Model Registry: Managing the ML Model Lifecycle

Training a model produces a file. Running a model in production needs much more: knowing which version is live, what it was trained on, whether it was checked before release, and how to go back if it misbehaves. A model registry provides that. It turns trained models into managed releases with versions, names and history.

This guide walks through a whole lifecycle in MLflow's registry, with real data, and explains each step.

The building blocks

MLflow's documentation describes aliases as the deployment mechanism: point champion at a version, load models:/MyModel@champion, and change what production serves by reassigning the alias (MLflow: model registry). Aliases and tags can be managed in the UI or through the API (MLflow: registry workflows).

A lifecycle, end to end

The script uses the UCI Bike Sharing data, daily rental counts for 2011 and 2012 under CC BY 4.0 (UCI: Bike Sharing). The production model was trained on 2011 alone. A retrained candidate adds the first half of 2012. Both are judged on July to December 2012, which neither has seen.

import mlflow
import pandas as pd
from mlflow import MlflowClient
from mlflow.models import infer_signature
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error

mlflow.set_tracking_uri("sqlite:///registry.db")
mlflow.set_experiment("bike-demand")
client = MlflowClient()
NAME = "bike-demand"

day = pd.read_csv("data/day.csv")                        # UCI Bike Sharing (CC BY 4.0)
features = ["season", "yr", "mnth", "holiday", "weekday", "workingday", "weathersit", "temp", "atemp", "hum", "windspeed"]
holdout = day[day.dteday >= "2012-07-01"]                 # data neither model trained on


def train_and_register(until, note):
    data = day[day.dteday < until]
    with mlflow.start_run(run_name=f"trained-to-{until}"):
        model = RandomForestRegressor(n_estimators=200, random_state=42).fit(data[features], data.cnt)
        info = mlflow.sklearn.log_model(model, name="model", input_example=data[features].head(2),
                                        signature=infer_signature(data[features], model.predict(data[features])),
                                        skops_trusted_types=["sklearn.tree._tree.Tree"])
    version = mlflow.register_model(info.model_uri, NAME).version
    client.set_model_version_tag(NAME, version, "trained_on", f"{data.dteday.min()} to {data.dteday.max()}")
    client.update_model_version(NAME, version, description=note)
    return version


def holdout_mae(alias_or_version):
    model = mlflow.pyfunc.load_model(f"models:/{NAME}/{alias_or_version}" if str(alias_or_version).isdigit()
                                     else f"models:/{NAME}@{alias_or_version}")
    return mean_absolute_error(holdout.cnt, model.predict(holdout[features]))


def aliases():
    return {a: v for a, v in client.get_registered_model(NAME).aliases.items()}


# 1. Production today: a model trained on 2011 only
v1 = train_and_register("2012-01-01", "first production model")
client.set_registered_model_alias(NAME, "champion", v1)
print("step 1, aliases:", aliases())

# 2. A new candidate arrives and is registered as the challenger, not yet serving anyone
v2 = train_and_register("2012-07-01", "retrained with the first half of 2012")
client.set_registered_model_alias(NAME, "challenger", v2)
print("step 2, aliases:", aliases())

# 3. The gate: promote only if the challenger beats the champion clearly on held-out data
champion_mae, challenger_mae = holdout_mae("champion"), holdout_mae("challenger")
improvement = 1 - challenger_mae / champion_mae
print(f"step 3, holdout MAE: champion {champion_mae:,.0f}, challenger {challenger_mae:,.0f}, improvement {improvement:.0%}")
if improvement >= 0.05:
    client.set_registered_model_alias(NAME, "previous", v1)       # keep a way back
    client.set_registered_model_alias(NAME, "champion", v2)
    client.delete_registered_model_alias(NAME, "challenger")
    client.set_model_version_tag(NAME, v2, "validation", f"passed: holdout MAE {challenger_mae:,.0f}")
print("        aliases after the gate:", aliases())

# 4. Something goes wrong in production: roll back by moving one alias
client.set_registered_model_alias(NAME, "champion", aliases()["previous"])
print("step 4, rollback, aliases:", aliases())
print("        the application now loads version", client.get_model_version_by_alias(NAME, "champion").version)

print("\nversion history:")
for mv in sorted(client.search_model_versions(f"name='{NAME}'"), key=lambda m: int(m.version)):
    print(f"   v{mv.version}: {mv.description} | trained_on {mv.tags.get('trained_on')} | {mv.tags.get('validation', 'no validation tag')}")

step 1, aliases: {'champion': 1}
step 2, aliases: {'challenger': 2, 'champion': 1}
step 3, holdout MAE: champion 2,149, challenger 885, improvement 59%
        aliases after the gate: {'champion': 2, 'previous': 1}
step 4, rollback, aliases: {'champion': 1, 'previous': 1}
        the application now loads version 1

version history:
   v1: first production model | trained_on 2011-01-01 to 2011-12-31 | no validation tag
   v2: retrained with the first half of 2012 | trained_on 2011-01-01 to 2012-06-30 | passed: holdout MAE 885

What happened at each step

• Step 1: version 1 was registered and given the champion alias. Applications loading @champion got it.
• Step 2: version 2 arrived as the challenger. It was registered, but nothing served it yet. Registering a model is not deploying it.
• Step 3: the gate compared both on held-out data. The champion was off by 2,149 rentals a day on average; the challenger by 885, a 59% improvement. It passed the 5% threshold, so the challenger became champion, version 1 became previous, and the result was written into a tag.
• Step 4: a rollback moved champion back to version 1 in one call, and the application immediately loaded version 1 again.
• The history explains itself. Each version carries its description, its training dates and, where it was validated, the result.

The size of the improvement is itself a lesson. Ridership grew strongly in 2012, and the model trained only on 2011 had never seen that level of demand. Models trained on the past go stale as the world moves, which is why retraining and the gate exist. See model drift.

What a real promotion gate checks

The script's gate checks a single number. Google's MLOps guidance lists what offline validation should cover before a model is promoted (Google Cloud: MLOps):

1. Metrics on a test set the model has not seen.
2. A comparison with the current production model, so a new model must actually be better.
3. Consistency across segments, such as regions or customer groups, where an overall gain can hide a local loss.
4. Compatibility with the serving infrastructure and its API, so the model loads and accepts the same inputs.

The script covers the first two. For the bike model, a segment check would compare errors on working days and weekends separately, and on each season, because a model can improve on average while getting worse for one of them. Automate these, record the results as tags, and require a person's approval for high-stakes models.

Who may move the champion

Because one alias decides what production serves, moving it is a release, and it deserves the controls a release gets.

• Separate registering from promoting. Anyone training models can register versions. Only an automated gate or a named reviewer should move the champion.
• Record every change: who moved an alias, when, from which version to which, and why. A tag on the version and a note in the change log are enough to start.
• Require approval for high-stakes models, such as credit, pricing or medical decisions, with the validation results attached.
• Restrict permissions on a shared server, so a notebook cannot repoint production by accident.

Rollback is a feature

Every release plan needs a way back. With aliases it is cheap: keep a previous alias pointing at the last good version, and moving champion back is instant. Rollback still needs a trigger. Monitoring has to notice the problem first. See monitoring ML models in production.

For the tracking that feeds the registry, see the MLflow tutorial. For automating promotion inside a pipeline, see CI/CD for machine learning.

Common mistakes

• Treating registration as deployment.
• Loading models by file path or run ID in applications instead of an alias.
• Promoting on a single overall metric without checking segments.
• No previous alias, so rollback means a scramble.
• Undocumented versions, with no record of training data or validation.

The model registry and model versioning are part of the Experiment Tracking and Model Management module of the MLOps course.

Sources

• Model Registry (MLflow documentation) — MLflow
• Model Registry Workflows (MLflow documentation) — MLflow
• MLOps: Continuous delivery and automation pipelines in machine learning (Google Cloud) — Google Cloud
• Bike Sharing dataset (UCI Machine Learning Repository) — UCI Machine Learning Repository