Model Versioning & Experiment Tracking with MLflow

Here's a bug we've been shipping since part one. Every time train.py runs, it overwrites model/model.joblib. Train a new model and the old one is gone — not archived, gone. If the new one turns out worse, there is no command that gets the old one back. We've been running production on a single mutable file that any training run can silently clobber.

"Just commit it to Git," says the reflex. But we buried that idea back in part one: Git versions a 47MB binary blob it can't diff, and even if it stored it perfectly, it still couldn't answer the questions that actually matter. Which hyperparameters produced this model? What was its error on the test set? Is it better or worse than the one currently serving traffic? Git tracks bytes. What we need is a system that tracks experiments.

That's MLflow. Two pieces of it: experiment tracking, which records the params, metrics, and model of every run so nothing is ever lost or unattributable; and the Model Registry, which turns those runs into versioned, promotable models. Together they finally answer the question part four left dangling — a green pipeline proves the tests passed, but MLflow is how you prove the model is any good.

Experiment tracking: never lose a run again

Instead of overwriting a file, each training run logs itself to MLflow: its parameters, its metrics, and the model artifact. Wrapping our part-one training in tracking is a handful of lines — mlflow.start_run() opens a run, and you log into it.

# registry/train_mlflow.py
import mlflow, mlflow.sklearn
from sklearn.ensemble import RandomForestRegressor
# ... same dataset + split as part one ...

mlflow.set_experiment("house-price")

for params in [
    {"n_estimators": 200, "max_depth": 12, "random_state": 42},
    {"n_estimators": 300, "max_depth": 16, "random_state": 42},
]:
    with mlflow.start_run() as run:
        model = RandomForestRegressor(**params).fit(X_train, y_train)
        preds = model.predict(X_test)
        mae = mean_absolute_error(y_test, preds)
        r2 = r2_score(y_test, preds)

        mlflow.log_params(params)
        mlflow.log_metrics({"mae": mae, "r2": r2})
        # artifact_path works on mlflow 2.x and 3.x; the newer name= is 3.x-only.
        mlflow.sklearn.log_model(model, artifact_path="model",
                                 registered_model_name="house-price")
        print(f"run {run.info.run_id[:8]}  mae={mae:.3f}  r2={r2:.3f}")

Run it against a local tracking store and you get two runs, both recorded, neither destroying the other:

run 87ebd3eb  mae=0.346  r2=0.792     # 200 trees, depth 12
run feb67576  mae=0.330  r2=0.804     # 300 trees, depth 16

These are the real numbers from our dataset. The second config — more trees, deeper — is genuinely better: MAE drops from 0.346 to 0.330, about $1,600 less average error per prediction. In the old world we'd have overwritten the first model with the second and hoped. Now both are on disk, tagged with exactly the params that produced them, and comparable side by side. Run mlflow ui and you get a sortable table of every run you've ever done — the experiment-tracking equivalent of finally turning on the lights.

The Model Registry: ECR for models

Logging runs is history. The Registry is control. Because we passed registered_model_name="house-price", each logged model also lands in the registry as a version — v1, v2, v3 — and this is exactly the ECR-for-models mental model from part one. Instead of pushing image:sha and promoting a tag through staging to production, you register model versions and move a named pointer onto whichever one earned it.

Modern MLflow does that promotion with aliases, not the old Staging/Production "stages" (those are deprecated). An alias is a movable, named pointer — think a Docker tag, or a Git branch that only ever points at one commit. We'll use one alias, production, and the rule for moving it is the champion/challenger gate we introduced all the way back in part one.

Promotion is a decision, not a default

This is the heart of it. A newer model is not automatically a better model, so "promote" can't mean "promote the latest." It means: the challenger only takes the production alias if it beats the current champion on a metric that matters. Lower MAE is fewer dollars of error, so that's our gate.

# registry/promote.py
from mlflow import MlflowClient
client = MlflowClient()

# Newest registered version is the challenger.
versions = client.search_model_versions("name='house-price'")
challenger = max(versions, key=lambda v: int(v.version))
challenger_mae = mae_of(client, challenger.version)

# Current champion, if any (None the very first time).
try:
    champion = client.get_model_version_by_alias("house-price", "production")
    champion_mae = mae_of(client, champion.version)
except Exception:
    champion, champion_mae = None, float("inf")

if challenger_mae < champion_mae:
    client.set_registered_model_alias("house-price", "production", challenger.version)
    print(f"PROMOTED v{challenger.version} to @production")
else:
    print("kept champion; challenger did not beat it")

One run of train_mlflow.py trains both configs, so the registry already holds v1 (MAE 0.346) and v2 (MAE 0.330). promote.py takes the newest, v2, as the challenger; with no champion yet it wins the production alias by default — and it deserves to, since 0.330 is the lower error. The gate earns its keep on the next round: if you later train a worse config — say a v3 that comes back at 0.351 — the script keeps the champion and refuses to promote it, no matter how new it is. That refusal is the entire point. It's the model-quality gate part four couldn't give you, and it drops straight into the CI pipeline: run the tests, and run this.

What this changes about serving

Here's the payoff that ties back to part two. Our API currently loads a hard-coded file path, model/model.joblib. With a registry, it instead loads by alias:

# the serving change, conceptually
import mlflow.sklearn
model = mlflow.sklearn.load_model("models:/house-price@production")

The service no longer knows or cares which version number it's running. It asks for "whatever is production right now." So promoting v2 over v1 becomes a pointer move in the registry — the next time a pod starts (or refreshes), it picks up the new champion with no image rebuild and no code change. Deploying a better model stops being a deploy at all. It becomes a promotion, and the "which model version is live" question from part four now has a precise answer: whatever @production points at.

In production you'd run MLflow with a real backend — a Postgres database for the metadata and S3 for the artifacts — so the registry outlives any one machine. We stand that up as part of the infrastructure in part seven. Locally, a sqlite file and a folder are enough to see the whole loop work end to end.

What's next

We can now track every experiment, keep every version, and promote only models that earn it — the three-part reproducibility (code, data, config) from part one finally has a home. What's still missing is somewhere real to run all this. Our container runs on a laptop; the registry runs on a laptop. In part six we deploy the API to Kubernetes — Deployments, Services, Ingress, ConfigMaps, Secrets, and autoscaling — and the readiness probe we've been quietly building since part two finally meets the orchestrator it was made for. The MLflow scripts are in the series repo under registry/. See you in part six.