The Model Registry

After tuning, you have a winning model. But where does it live? On someone's laptop? In a folder called final? A model needs a proper home, where every version is stored, described, and approved before it reaches users. That home is the model registry. In this lesson, we use the one in M L flow.

• A home for every model version
• Approved before it ships
• Applications ask by name, not file

What a registry gives you

A registry gives four things. A catalogue, with one entry for each model, and a numbered version for every update. Lineage, each version links back to the run, the code, and the data that produced it. Stages or aliases, labels like champion, or staging, that say which version is in use. And governance, who approved a version, and when. It turns models from loose files into managed products.

The registry in code

Here is the registry in code. We train two versions of a model, with different depths. For each, we log the model to a run, and register it under one name, churn dash model. The registry creates version one, and version two. We list the versions, to confirm. Now the key move. We give version two the alias, champion. Applications load the model by that alias, not by a number. And we ask the champion for predictions. To promote a new model, you only move the alias.

import mlflow
from mlflow import MlflowClient
from sklearn.ensemble import RandomForestClassifier
from churn import make_data

mlflow.set_tracking_uri("sqlite:///mlflow.db")
df = make_data(2000, seed=1)
X, y = df.drop(columns="churn"), df["churn"]

for depth in (3, 8):
    with mlflow.start_run():
        model = RandomForestClassifier(max_depth=depth, random_state=0).fit(X, y)
        info = mlflow.sklearn.log_model(model, name="model", serialization_format="cloudpickle")
        mlflow.register_model(info.model_uri, "churn-model")

client = MlflowClient()
for v in client.search_model_versions("name='churn-model'"):
    print("version", v.version, "created")

client.set_registered_model_alias("churn-model", "champion", 2)
champion = mlflow.pyfunc.load_model("models:/churn-model@champion")
print("champion is version", client.get_model_version_by_alias("churn-model", "champion").version)
print("prediction:", champion.predict(X.head(3)).tolist())

version 2 created
version 1 created
champion is version 2
prediction: [0, 0, 0]

import numpy as np, pandas as pd

def make_data(n, seed, shifted=False):
    rng = np.random.default_rng(seed)
    tenure = rng.integers(1, 72, n)
    charges = rng.normal(65 + (28 if shifted else 0), 18, n).clip(20, 140).round(2)
    calls = rng.poisson(1.5 + (1.5 if shifted else 0), n)
    contract = rng.choice([0, 1, 2], n, p=[.55, .25, .20] if not shifted else [.30, .30, .40])
    logit = -1.0 - 0.04 * tenure + 0.025 * (charges - 65) + 0.45 * calls - 0.8 * contract
    if shifted:
        logit = -0.2 - 0.01 * tenure - 0.01 * (charges - 65) + 0.2 * calls - 0.2 * contract
    churn = (rng.random(n) < 1 / (1 + np.exp(-logit))).astype(int)
    return pd.DataFrame({"tenure": tenure, "charges": charges, "calls": calls, "contract": contract, "churn": churn})

A promotion workflow

In a healthy workflow, promotion is deliberate. A new model version is registered, after training. Automatic tests run on it, as we will see next. If they pass, it receives the staging alias, and is tried in a safe environment. A person, or an automatic gate, then compares it with the current champion. If it is better, the champion alias moves to the new version. And the old version stays in the registry, so rolling back is just moving the alias back.

1. Registera new numbered version
2. Testautomatic checks pass
3. Stagetried in a safe place
4. Promotemove the champion alias

Summary