Model Registry

Every example in this section continues with the well-known Adult Income dataset [1] used throughout the rest of this documentation, and with the Logistic Regression, Decision Tree, and Random Forest classifiers trained on it using the model_tuner library [2]. Click here to view the corresponding codebase for this workflow.

That page ends by loading the model objects back from disk with explicit paths. This works for one project on one machine. It stops working as soon as there are several tracking stores under one project, or the mlruns/ tree gets copied somewhere else, or you simply cannot remember which of eleven runs produced the number in the paper.

model_registry indexes every model in every MLflow store under a project root and addresses them by name. There are no hard-coded experiment ids, run ids, or paths anywhere in the module, and mlflow itself is not a dependency: the module reads the local FileStore tree (the meta.yaml files) directly from disk.

Two problems it handles that the plain MLflow client does not:

Several tracking roots under one project. mlruns/preprocessing and mlruns/models are both indexed, not just whichever one you point at.
Relocated artifact trees. The artifact_location recorded at training time often points into a directory that no longer exists, because mlruns/ trees get copied between machines and projects. Every MLflow artifact API then resolves to a missing path. Artifacts are instead rebased onto the store root actually found on disk.

Quick Start 

from model_metrics.model_registry import available, load, load_all, rank

available()                     # everything on disk, with logged metrics
rank("rf", metric="roc_auc")    # every random forest run, ordered
model_rf = load("rf_smote")     # pin the variant you actually reported
models = load_all()             # one model per variant

If nothing is found, the error explains why rather than just reporting an empty result. Start with diagnose() when that happens.

Configuration 

Every setting is optional and can be supplied by environment variable (before import) or by calling configure() (before the first lookup).

Variable	Meaning
`MODEL_REGISTRY_ROOT`	Directory to walk. Defaults to the working directory, walking up to the repository boundary to find an `mlruns/` tree.
`MODEL_REGISTRY_TARGET`	Outcome token stripped from artifact folder names.
`MODEL_REGISTRY_STORES`	Comma-separated store prefixes eligible to win a metric comparison.
`MODEL_REGISTRY_FILE`	Artifact filename. Defaults to `model.pkl`.
`MODEL_REGISTRY_TOLERANT`	Set to `0` to make scikit-learn version mismatches raise instead of warn.

Important

The default search root is the working directory, not the module’s own location. Installed in site-packages the module is nowhere near your project, and in a src/ layout its location resolves to src/ rather than the repository. The walk climbs until it finds an mlruns/ directory, which also lets the registry work when called from a notebooks/ subdirectory, and stops at .git so an unrelated mlruns/ higher up the filesystem is never picked up silently.

configure(root=None, target=None, model_file=None, loader=None, stores=None)

Override configuration at runtime and drop the cached index. Call before the first lookup, or after moving or regenerating mlruns/.

Parameters:

root (str or pathlib.Path, optional) – Directory to walk.
target (str, optional) – Outcome token to strip from artifact folder names.
model_file (str, optional) – Artifact filename to look for inside each run.
loader (callable, optional) – Custom callable path -> object replacing the default deserializer, for artifacts that need special handling.
stores (str or sequence of str, optional) – Store prefixes eligible to win a metric comparison.

Returns:

None.

Return type:

None

from model_metrics import model_registry

# Adult Income artifact folders are named <algo>_income
model_registry.configure(target="income")

refresh()

Drop the cached filesystem index and the model load cache. Call after writing new runs during the same session.

Returns:: None.
Return type:: None

Naming: algo vs variant 

Several runs commonly write to one artifact folder, so a folder name alone cannot address a run. The registry splits the two apart:

algo comes from the artifact folder: lr, dt, rf.
variant comes from the MLflow run name: lr_orig, rf_smote, rf_orig_no_sex.

variant is the addressable key. algo is what you group by.

Artifact folders are conventionally named <algo>_<TARGET>, so setting MODEL_REGISTRY_TARGET to the outcome name turns rf_income into rf and lets you group by algorithm rather than by outcome. The token is stripped wherever it appears, not only as a suffix, so ablation variants survive intact:

TARGET = "income"

rf_income           ->  rf
rf_income_no_sex    ->  rf_no_sex

This matters for the Adult Income examples specifically. A fairness workflow on that dataset usually produces ablation variants alongside the full models, dropping sex or race to see what the model loses. Those variants have to stay individually addressable while still grouping under their algorithm, which is what the wherever-it-appears rule buys.

Note

The default is an empty token, meaning no stripping, which is the safe behavior for a project whose naming convention is unknown.

Discovery 

available()

Every model found on disk, newest first, with its logged metrics.

Returns:: DataFrame with variant, algo, experiment, experiment_id, run_id, path (relative to the project root), and one column per logged metric.
Return type:: pandas.DataFrame

rank(name, metric='roc_auc', ascending=False, experiment=None)

Order every run matching name by a metric. For looking, not for loading: rank("rf") shows every random forest run so you can see which one is rf_smote.

Parameters:

name (str) – A variant, an algo, a run id, or a qualified experiment/variant key.
metric (str, optional) – Metric to sort on. Ignored if the runs did not log it.
ascending (bool, optional) – Sort direction. Defaults to False (best first for higher-is-better metrics).
experiment (str, optional) – Restrict to one experiment, by id or by name.

Returns:

DataFrame with variant, algo, experiment, run_id, and one column per logged metric.

Return type:

pandas.DataFrame

Raises:

LookupError – If nothing matches name.

variants()

Returns:: Sorted list of every addressable variant name.
Return type:: list[str]

algos()

Returns:: Sorted list of every algo name.
Return type:: list[str]

experiments()

Returns:: Sorted list of every experiment name.
Return type:: list[str]

metric_names()

Returns:: Sorted list of every metric key logged anywhere in the store.
Return type:: list[str]

is_available()

True if at least one loadable model was found, without raising. Use this to gate registry-dependent code paths in environments that may have no MLflow store at all.

Returns:: Whether the index is non-empty.
Return type:: bool

store_summary()

One row per tracking store on disk, with how many models it holds and whether it is currently eligible to win a metric comparison. Use it to see what STORES is excluding before trusting a best_per_algo() result.

Returns:: DataFrame with store, models, experiments, and eligible.
Return type:: pandas.DataFrame

diagnose()

Print the resolved configuration, active deserializer, YAML parser, any shims and repairs applied so far, the store summary, and the full inventory. Also runs as __main__.

Returns:: None.
Return type:: None

python -m model_metrics.model_registry

Notes

Failure messages are specific. Three distinct situations get three distinct explanations: the search root does not exist; no MLflow tracking store was found anywhere beneath it; stores were found but no run contains an artifact under the configured filename. Each names the setting to change.
Remote backends are out of scope. The module reads a local FileStore directory. Runs on a remote tracking server or a SQL backend need the mlflow client instead, and the error says so.

Stores: constraining what can win 

The index is always global. available(), rank(), load(), load_all(), and variants() see every model on disk, so nothing is ever hidden from you.

STORES constrains only the metric-based selectors, best_per_algo() and load_best_per_algo(), so that a superseded store cannot win a comparison it was never meant to be in.

from model_metrics.model_registry import set_stores, stores, best_per_algo

set_stores("mlruns/models")                     # only the live store
set_stores("mlruns/models", "mlruns/archive")   # two of them
set_stores()                                    # clear the constraint

best_per_algo(stores="mlruns/models")           # one call only

set_stores(*prefixes)

Restrict which stores may win a metric comparison.

Parameters:: prefixes (str) – Project-relative store prefixes. Accepts separate arguments, a single comma-separated string, or any iterable of either. Called with no arguments, clears the constraint.
Returns:: The new setting.
Return type:: tuple

stores()

Returns:: The store prefixes currently constraining metric-based selection.
Return type:: tuple

Important

Prefixes are matched on whole path segments. mlruns/models matches mlruns/models and mlruns/models/12345, but not mlruns/models_old/12345. A substring test would wrongly accept the latter, which is the exact confusion this setting exists to prevent.

Note

When a constraint excludes every run that logged the requested metric, the error lists the stores that do have it, so widening the constraint is a matter of copying a name rather than going spelunking.

Selection 

The default selection policy is "newest": when a name matches several runs, the latest start_time wins, and a one-line note reports which run was chosen and how.

Selecting by best test metric is available but is not the default, because picking the max-metric run is model selection on the evaluation set. Use it to inspect; pin by name to load.

resolve(name, experiment=None, policy='newest', metric='roc_auc')

Resolve one entry without loading it.

Parameters:

name (str) – A variant, an algo, a run id, or experiment/variant.
experiment (str, optional) – Restrict to one experiment, by id or by name.
policy (str, optional) – "newest" (default) or "best" (maximum metric).
metric (str, optional) – Metric used when policy="best".

Returns:

The resolved entry.

Return type:

ModelEntry

Raises:

LookupError –

If nothing matches; the error lists every available variant.
If policy="best" but no matching run logged metric; the error lists the metrics that are present.

load(name, experiment=None, policy='newest', metric='roc_auc')

Load one model. Cached. Aliased as load_model.

Parameters:

name (str) – A variant, an algo, a run id, or experiment/variant.
experiment (str, optional) – Restrict to one experiment.
policy (str, optional) – "newest" or "best".
metric (str, optional) – Metric used when policy="best".

Returns:

The deserialized model object.

Return type:

object

load_all(only=None, qualified=False)

One model per variant, not per algo, so runs sharing an artifact folder stay distinct.

Parameters:

only (list[str], optional) – Restrict to these variants, algos, or qualified keys.
qualified (bool, optional) – Key the result by "<experiment_name>/<variant>" so nothing collides across experiments. Defaults to False.

Returns:

Mapping of key to loaded model.

Return type:

dict

Loading the three Adult Income classifiers by name gives the same objects the rest of the documentation uses, without pinning a path:

from model_metrics.model_registry import load

model_lr = load("lr_orig")
model_dt = load("dt_orig")
model_rf = load("rf_orig")

resolve_metric(metric)

Map a loose metric name onto the key actually logged.

Parameters:

metric (str) – A metric name, exact or approximate.

Returns:

The logged key.

Return type:

str

Raises:

LookupError –

If the name is ambiguous; the error lists every candidate.
If nothing resembles it; the error lists every logged metric.

Note

resolve_metric normalizes spaces and hyphens to underscores, so average_precision finds test_average_precision, valid ap, or test Average Precision. Without that normalization a raw substring test never matches and the default metric argument to best_per_algo() raises on a perfectly ordinary store. Average-precision aliases (aucpr, pr_auc, ap, and others) are recognized. Ambiguity raises rather than silently ranking on the wrong split.

Metric-Based Selection 

best_per_algo(metric='average_precision', experiment=None, per_experiment=True, ascending=False, stores=None)

The winning run for each algo, ranked by metric. Inspect before loading.

Parameters:

metric (str, optional) – Metric to rank on. Resolved through resolve_metric.
experiment (str, optional) – Restrict to one experiment.
per_experiment (bool, optional) – Keep runs from different experiments in separate groups instead of letting them compete. Defaults to True.
ascending (bool, optional) – Set True for lower-is-better metrics.
stores (tuple, optional) – Override the global store constraint for this call only.

Returns:

DataFrame with algo, store, winner, experiment, the ranking metric, run_id, n_candidates, and the winner’s remaining metrics.

Return type:

pandas.DataFrame

Raises:

LookupError –

If no run logged the resolved metric.
If the store constraint excludes every run that did.

load_best_per_algo(metric='average_precision', experiment=None, per_experiment=True, ascending=False, qualified=False, stores=None)

Load the top run for each algo by metric.

Parameters:: qualified (bool, optional) – Key the result by "<experiment_name>/<variant>" instead of by algo. Defaults to False.
Returns:: Mapping of key to loaded model.
Return type:: dict

Important

If metric is computed on your test set, this is model selection on the test set and the winning score is optimistically biased. Fine for exploration. For anything you report, pin the variant by name, or select on validation data using the functions below.

Note

The n_candidates column says how many runs the winner beat under the same constraints. A winner with n_candidates=1 won by default, which is worth knowing before quoting it as the best of anything.

Validation-Set Selection 

The defensible path: choose champions on validation data, then report test metrics for those winners only.

score_candidates(data, scorer=None, name=None)

Score every indexed model against a held-out set.

Parameters:

data (tuple, dict, or callable) – One of three forms. A (X, y) tuple applies one matrix to every model. A dict keyed by algo, variant, or qualified key supplies a different matrix per model, which is what ablation variants need. A callable entry -> (X, y) covers anything more involved.
scorer (callable, optional) – Callable (y_true, y_score) -> float. Defaults to average_precision_score.
name (str, optional) – Restrict to models matching this variant, algo, or run id.

Returns:

DataFrame with variant, algo, experiment, score, run_id, and an error column for any model that failed.

Return type:

pandas.DataFrame

select_on_validation(data, scorer=None, per_experiment=True)

Champion of each family, chosen on the data you pass. Pass validation data. Then report test metrics for these winners only.

Parameters:

data (tuple, dict, or callable) – Same three forms accepted by score_candidates.
scorer (callable, optional) – Scoring callable. Defaults to average precision.
per_experiment (bool, optional) – Group by experiment as well as by algo. Defaults to True.

Returns:

One row per winning model, sorted by score.

Return type:

pandas.DataFrame

load_selected(selection, qualified=False)

Load the models named in a select_on_validation() frame.

Parameters:

selection (pandas.DataFrame) – The frame returned by select_on_validation.
qualified (bool, optional) – Key by "<experiment_name>/<variant>".

Returns:

Mapping of key to loaded model.

Return type:

dict

The dict form of data is what makes this work for the Adult Income fairness variants, where the ablated models were fitted on a narrower feature set and cannot consume the full validation matrix:

from model_metrics.model_registry import select_on_validation, load_selected

winners = select_on_validation((X_valid, y_valid))
models = load_selected(winners)

# ablation variants scored against their own matrices in the same sweep
winners = select_on_validation({
    "rf_orig": (X_valid, y_valid),
    "rf_no_sex": (X_valid.drop(columns=["sex"]), y_valid),
})

Note

A model that fails to load or score is captured per row with an error column rather than aborting the sweep, so one unloadable artifact does not cost you the other twenty scores.

Cross-Version Loading 

Artifacts are often written under a different scikit-learn than the one installed, and pickles reference private sklearn classes by module path. Those classes get renamed or removed between versions, so a plain load dies with something like:

AttributeError: Can't get attribute '_RemainderColsList' on
<module 'sklearn.compose._column_transformer'>

Two layers handle the drift.

Tolerant unpickling. Missing sklearn internals are synthesized rather than raising. Every load tries the configured deserializer, retries with the known shims installed, then falls back to an unpickler that synthesizes any remaining sklearn class on demand. Missing non-sklearn classes still raise, since synthesizing those would be guessing.

Post-load fitted-state repair. Shimming fixes unpickling but not fitted state: newer transform() code reads instance attributes that an older fit() never wrote. Each repair restores the attribute to the value that reproduces the original training-version behavior, not the current version’s.

repair_estimator(model, verbose=False)

Patch fitted attributes the installed scikit-learn expects but the pickle predates. Applied automatically on every load; exposed for models obtained some other way.

Parameters:

model (object) – A fitted estimator, pipeline, or wrapper. Walked recursively.
verbose (bool, optional) – Print each repair as it is applied.

Returns:

A description of each repair applied.

Return type:

list[str]

shimmed()

Returns:: The sklearn internals this process had to synthesize, if any.
Return type:: list[str]

repaired()

Returns:: The fitted-state attributes this process had to restore, if any.
Return type:: list[str]

backend()

Returns:: Which deserializer is in use: "model_tuner", "joblib", or "pickle".
Return type:: str

Notes

The two repairs that ship:
- SimpleImputer._fill_dtype is set to statistics_.dtype, which makes the newer astype() call a no-op and matches versions that did not cast at all. Gated on a probe of the installed sklearn: the attribute was introduced partway through the 1.x line, so on versions that never had it, its absence is normal and repairing anyway would fill repaired() with noise.
- ColumnTransformer remainder columns. A synthesized _RemainderColsList unpickles as an empty list, which under remainder="passthrough" silently drops every passthrough column. This is the failure most likely to bite an Adult Income pipeline, since those routinely one-hot the categorical columns and pass the numeric ones through. It surfaces far downstream as ValueError: Feature shape mismatch, expected: 25, got 15. The remainder set is fully determined by fitted state, so it is rebuilt exactly as sklearn computes it at fit time and written to all three places that must agree: _transformer_to_input_indices, _remainder, and transformers_. Repairing _remainder alone is not enough, because transform() reads transformers_ when fitted.
Deserializer selection:
- model_tuner’s loadObjects [2] is used when installed, so whatever it does at save time is mirrored at load time.
- Otherwise joblib (which reads plain pickles too), then pickle.
- Override with configure(loader=...) for artifacts needing special handling.

Important

None of this makes cross-version loading supported. It makes it survivable long enough to verify. Both layers warn when they fire, and the warning points at verify_entry(). The correct fix is to pin the training scikit-learn version; check the run’s artifacts/*/requirements.txt for what that was.

Verification 

This is the check that makes a cross-version load defensible. If the recomputed score matches the logged one, the estimator state survived whatever shimming and repair was needed. If it does not, do not use the model.

verify_entry(name, X, y, split='test', tol=1e-3, **resolve_kw)

Recompute ROC AUC, average precision, and Brier score for a loaded model and compare them against what was logged.

Parameters:

name (str) – The variant to verify.
X (array-like) – Feature matrix corresponding to split.
y (array-like) – True labels corresponding to split.
split (str, optional) – Which logged metrics to compare against, matched as a prefix on the logged key ("test", "valid", "train"). Pass None to compare against every logged metric.
tol (float, optional) – Absolute tolerance for the comparison. Defaults to 1e-3.
resolve_kw – Forwarded to resolve (experiment, policy, metric).

Returns:

DataFrame with logged_metric, logged, recomputed, delta, and ok.

Return type:

pandas.DataFrame

verify_all(X, y, metric='average_precision', split='test', tol=1e-3, **best_kw)

Verify every champion returned by best_per_algo() in one call.

Returns:: One row per model with algo, variant, n_compared, max_abs_delta, and a single ok.
Return type:: pandas.DataFrame

Important

split must correspond to the data you pass. Comparing a logged validation score against metrics recomputed on the test set will always disagree and means nothing. When no logged key matches the requested split, the function warns and returns an empty frame rather than silently reporting success.

Note

The default tolerance is 1e-3 rather than something tighter because MLflow commonly stores metrics rounded to three decimals. verify_all reports the worst absolute delta alongside the pass/fail, so a single bad load cannot hide behind an aggregate.

from model_metrics.model_registry import verify_entry, verify_all

verify_entry("rf_orig", X_test, y_test)                  # test split
verify_entry("rf_orig", X_valid, y_valid, split="valid")

verify_all(X_test, y_test)

The ModelEntry Record 

resolve() returns a frozen dataclass describing one indexed model. Most workflows never touch it directly, but it is what the DataFrame-returning functions are built from.

Field	Meaning
`algo`	From the artifact folder: `lr`, `dt`, `rf`, …
`variant`	From the MLflow run name: `rf_orig`, `rf_smote`, …
`path`	Local artifact path, verified to exist.
`run_id`	MLflow run id.
`experiment_id`	MLflow experiment id.
`experiment_name`	MLflow experiment name.
`store_root`	The tracking store this run belongs to.
`start_time`	Run start timestamp, used by the `"newest"` policy.
`metrics`	Mapping of logged metric key to value.
`key` (property)	`"<experiment_name>/<variant>"`, the fully qualified address.

Optional Dependencies 

The module degrades rather than requiring anything beyond pandas.

mlflow is not needed at all. The FileStore tree is read directly.
model_tuner [2], joblib, pickle: tried in that order for deserialization. backend() reports which is active.
PyYAML is used when present. Without it, a built-in flat-scalar parser covers every field the module actually reads (experiment_id, name, run_id, run_uuid, run_name, lifecycle_stage, start_time). MLflow’s meta.yaml files are flat key/value documents, so the fallback is not a compromise.