MLflow: the model registry

Tracking tells you which run was best. But which model is in production right now? Which version do you roll back to when the new one misbehaves? How does your serving code fetch "the current model" without hard-coding a file path? That's the job of the Model Registry — a versioned, named shelf for the models you actually ship.

Install: pip install mlflow pandas (the registry's pyfunc models use pandas). Needs a database backend — we use a local SQLite file.

The problem it solves

Without a registry, "deploy the model" means copying a file and hoping everyone uses the right one. The registry replaces that with named, versioned models and aliases like @production and @staging, so your serving code asks for "models:/sentiment@production" and always gets the right version — and you can promote or roll back by moving an alias, not redeploying code.

Step 1 — wrap the model as a pyfunc

The registry works best with a pyfunc — MLflow's universal model interface. Wrapping our custom model in it means any MLflow-aware system can load and run it the same way, regardless of framework. From registry.py:

import mlflow.pyfunc

class SentimentPyfunc(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        self.model = SentimentModel.load(context.artifacts["model_path"])

    def predict(self, context, model_input, params=None):
        return self.model.predict_proba(model_input["text"].tolist())

load_context restores the model from its artifact; predict defines the universal inference call. That's the adapter that makes our NumPy model speak fluent MLflow.

Step 2 — log and register in one call

mlflow.set_tracking_uri("sqlite:///mlflow.db")     # DB backend enables the registry
with mlflow.start_run() as run:
    mlflow.pyfunc.log_model(
        name="model",
        python_model=SentimentPyfunc(),
        artifacts={"model_path": "model.json"},
        registered_model_name="sentiment",          # <-- registers it
    )

The registered_model_name argument is the magic word: it logs the model and creates (or adds a new version to) a registered model called sentiment.

Step 3 — promote with an alias

A registered model accumulates versions (v1, v2, v3…). You point an alias at the one that's live:

client = mlflow.MlflowClient()
version = client.get_latest_versions("sentiment")[0].version
client.set_registered_model_alias("sentiment", "production", version)

Don't be confused: aliases vs. the old "stages." Older MLflow used fixed stages (Staging, Production, Archived). Modern MLflow (2.9+) replaced these with flexible aliases — you can have @production, @champion, @canary, whatever you need. If you see transition_model_version_stage in old tutorials, the modern equivalent is set_registered_model_alias.

Step 4 — load by alias (what serving does)

Your API never hard-codes a path or version. It asks for the alias:

model = mlflow.pyfunc.load_model("models:/sentiment@production")
out = model.predict(pd.DataFrame({"text": ["this is great", "this is terrible"]}))

Running the whole thing

cd code/sentiment
python registry.py

Output:

logged + registered 'sentiment' from run 6fc7412d
set alias 'production' -> sentiment v1
loaded model predictions: [0.977, 0.024]

It registered version 1, pointed @production at it, then loaded by that alias and predicted — 0.977 for the positive sentence, 0.024 for the negative. Run it again and you'll get v2, with @production moved to the new version — that's a deploy. Point the alias back to v1 and you've rolled back, instantly, without touching serving code.

The registry in the MLflow UI

mlflow ui --backend-store-uri sqlite:///mlflow.db --port 5000

The Models tab now shows sentiment with all its versions, their aliases, which run produced each, and its metrics. This is the team's source of truth for "what's in production and how did it get there."

How this looks at a real company

  • A CI pipeline trains a candidate, logs it, and registers a new version.
  • An automated gate (or a human) checks it beats the current @production on a holdout, then moves the alias — often via a canary: @production for 95% of traffic, @canary (the new version) for 5%, watched before full promotion.
  • Serving code loads models:/sentiment@production and is blissfully unaware of version numbers.
  • An incident? Move the alias back to the previous version. Rollback in seconds.

Managed platforms (SageMaker Model Registry, Vertex Model Registry) offer the same concept; the idea — named, versioned models with promotable aliases — transfers everywhere.

The takeaway

The Model Registry turns a pile of model files into named, versioned, promotable models. Wrap the model as a pyfunc, register a version, point a @production alias at it, and load by alias in serving — so deploys and rollbacks are just moving an alias, never editing code. Tracking found the best model; the registry ships it. Now let's actually expose it to the world with an API. 👉