Version Models with MLflow

You trained a great model. Weeks later, you can't reproduce the exact pipeline that produced it, you don't know which dataset version fed it, or which hyperparameters gave you that 0.92 F1 score. Now you're stuck debugging a degraded model with zero visibility into what changed. This chaos is common in applied AI engineering — and versioning models with MLflow is the fix. It gives you a central registry where every model, its code version, parameters, metrics, and artifacts are tracked and retrievable on demand.

The problem this lesson solves

When you build models repeatedly (or in teams), manual versioning breaks down quickly. Naming files like model_final_v2_real_final.pkl is fragile and uninformative. You lose the link between a model and its training run, metrics, and even the exact code version that produced it. This leads to:

• Reproducibility failures: you can't recreate a model's exact metrics without knowing the data and parameter set.
• No lineage: no record of which model version is currently in production or why it changed.
• Poor collaboration: team members can't see or compare experiments, leading to confusion.
• Impossible rollback: when a new model performs worse, you can't easily switch back to a known-good one.

MLflow's Model Registry solves this by giving every model a unique version number, storing its metadata, and letting you manage its lifecycle (staging, production, archived) in one place.

Core concept / mental model

Think of MLflow Model Registry as Git for machine learning models. Where Git tracks code commits, MLflow tracks model versions. Each version is tied to a logged run — an immutable snapshot containing parameters, metrics, and artifacts.

Here's the core vocabulary:

• Run: a single execution of your training code.
• Experiment: a logical grouping of runs (e.g., 'churn-model').
• Model: a registered model is a named entity (e.g., 'ChurnPredictor') that can have multiple versions.
• Version: an immutable snapshot of the model plus its metadata.
• Stage: lifecycle status (Staging, Production, Archived).

Pro tip: A model version is immutable — you never change it. You always create a new version. This mirrors container image tags: you tag a build, you don't edit it.

The flow is: train → log to MLflow → register the model → assign a stage → later retrieve and promote.

How it works step by step

Versioning models with MLflow follows a predictable sequence. Here's the cause → effect chain:

1. Set up tracking: Initialize an MLflow tracking URI (local file or server).
2. Create an experiment: Group your runs under a logical name.
3. Log the run: Inside your training code, log parameters, metrics, and the model artifact.
4. Register the model: After the run, register the logged model as a named model.
5. Assign a stage: Mark the version as Staging, Production, or Archived.
6. Retrieve later: Pull any version by name and stage — each returns the exact model object.

Each step is deliberate: logging makes the run reproducible; registering gives a human-readable name; staging gives a lifecycle. Skip any step and you lose lineage.

Hands-on walkthrough

Let's implement a real example with a scikit-learn classifier. First, install MLflow:

pip install mlflow scikit-learn

Now, train and log a model with auto-logging:

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score

# 1. Set the tracking URI (local folder)
mlflow.set_tracking_uri("./mlruns")

# 2. Create an experiment
exp_name = "Churn-Prediction"
mlflow.set_experiment(exp_name)

# 3. Auto-log parameters, metrics, and model
data, target = make_classification(n_samples=2000, n_features=20, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2)

with mlflow.start_run():
    model = RandomForestClassifier(n_estimators=100, max_depth=5, random_state=42)
    model.fit(X_train, y_train)

    pred = model.predict(X_test)
    score = f1_score(y_test, pred)

    mlflow.log_metric("f1_score", score)
    mlflow.sklearn.log_model(model, "model")

    print(f"Logged run with F1: {score:.3f}")

This logs the run under the experiment, saving the model artifact. But we haven't registered it yet. Let's do that:

import mlflow

client = mlflow.tracking.MlflowClient()

# Find the last run ID from the experiment
exp = mlflow.get_experiment_by_name("Churn-Prediction")
runs = mlflow.search_runs([exp.experiment_id], order_by=["start_time desc"])
run_id = runs.iloc[0].run_id

# Register the model as a named model
result = mlflow.register_model(
    model_uri=f"runs:/{run_id}/model",
    name="ChurnPredictor"
)
print(f"Registered version {result.version}")

Now you can manage stages:

# Move version 1 to Production
client.transition_model_version_stage(
    name="ChurnPredictor",
    version=1,
    stage="Production"
)

# Later, load the production model
model = mlflow.pyfunc.load_model(
    model_uri="models:/ChurnPredictor/Production"
)
print(type(model))

Each run you log can get registered, and each registration creates a new version. That's versioning with MLflow in a nutshell.

Compare options / when to choose what

MLflow's model registry isn't the only versioning approach — but it's tightly integrated with experiment tracking. Here's a comparison:

Variations: - Use DVC (Data Version Control) if you need to version datasets alongside models. - Use Seldon Core or BentoML for serving, but they rely on MLflow for registry. - For deep learning, use MLflow's PyTorch/TensorFlow flavors instead of sklearn.

If you need reproducibility, auditability, and a model lifecycle, MLflow is the clear win.

Troubleshooting & edge cases

• Model not registered: Forgot to call mlflow.register_model. Check that your run finished and has a logged model.
• Stage transition fails: You must specify a valid stage — 'Staging', 'Production', or 'Archived'.
• Model load error: Loading a PMML versus pickle can throw. Use the correct flavor (mlflow.sklearn.load_model for sklearn, mlflow.pyfunc for generic)
• Version conflicts: If you register two models with the same name, you get consecutive versions. Verify you're using the right version.
• Tracking URI mismatch: If your client and server use different tracking URIs, you won't see runs. Keep them consistent.

Pro tip: For production, run a central MLflow tracking server (e.g., mlflow server --backend-store-uri postgresql://...) so all team members share the same registry.

What you learned & what's next

You've learned the core idea behind versioning models with MLflow: each model version is an immutable snapshot tied to a run, parameters, metrics, and artifacts. You can now register models, assign lifecycle stages, and load any version reliably. This tackles reproducibility and rollback in applied AI engineering.

Next in the track, you'll explore model serving and monitoring — taking a registered production version and exposing it via a REST API, then tracking its performance over time. You'll build on the registry you just created.