Model Versioning with MLflow: Practical Guide

The worst position to be in as an ML engineer: production is broken, the model that was working last week is gone, and you have no idea what changed. Model versioning prevents this. MLflow is the most widely used open-source solution.

Why Git Is Not Enough for Models

You version your code with Git. Why not models?

1. Model files are large - a fine-tuned transformer is 7GB+. Git LFS works but is painful.
2. Models depend on data, not just code - the same code on different data produces different models. You need to version both together.
3. Experiments are many-to-one - you run 50 experiments, pick 1. You want to query and compare all 50, not cherry-pick commit hashes.
4. Deployment has stages - staging, production, archived. Git branches do not map cleanly to these.

MLflow solves all four.

Core Concepts

Experiment → Runs → Artifacts
                ↓
         Model Registry → Versions → Stages (Staging / Production / Archived)

• Experiment: a collection of runs for one problem (e.g., "churn-model-v2")
• Run: one training job - captures params, metrics, artifacts
• Model Registry: named model entries with versioned artifacts and stage labels

Instrumenting Your Training Code

pythonCopy
import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
import numpy as np

# Point to your MLflow server (or use local ./mlruns)
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("churn-prediction")

def train_model(n_estimators, max_depth, X_train, y_train, X_val, y_val):
    with mlflow.start_run(run_name=f"rf-{n_estimators}-depth{max_depth}"):
        # Log hyperparameters
        mlflow.log_params({
            "n_estimators": n_estimators,
            "max_depth": max_depth,
            "model_type": "RandomForest"
        })

        # Train
        model = RandomForestClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            random_state=42
        )
        model.fit(X_train, y_train)

        # Evaluate and log metrics
        val_accuracy = model.score(X_val, y_val)
        val_auc = roc_auc_score(y_val, model.predict_proba(X_val)[:, 1])
        cv_scores = cross_val_score(model, X_train, y_train, cv=5)

        mlflow.log_metrics({
            "val_accuracy": val_accuracy,
            "val_auc": val_auc,
            "cv_mean": cv_scores.mean(),
            "cv_std": cv_scores.std()
        })

        # Log the model with input schema for validation
        from mlflow.models.signature import infer_signature
        signature = infer_signature(X_train, model.predict(X_train))

        mlflow.sklearn.log_model(
            model,
            "model",
            signature=signature,
            registered_model_name="churn-predictor"  # auto-registers
        )

        return model, val_auc

# Run a sweep
results = []
for n_est in [100, 200, 500]:
    for depth in [5, 10, None]:
        model, auc = train_model(n_est, depth, X_train, y_train, X_val, y_val)
        results.append((n_est, depth, auc))

Querying Experiments to Find the Best Run

pythonCopy
from mlflow.tracking import MlflowClient
import pandas as pd

client = MlflowClient()

# Get all runs for the experiment
experiment = client.get_experiment_by_name("churn-prediction")
runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    order_by=["metrics.val_auc DESC"],
    max_results=10
)

# Compare them
for run in runs:
    print(f"Run {run.info.run_id[:8]}: "
          f"AUC={run.data.metrics.get('val_auc', 'N/A'):.3f}, "
          f"n_estimators={run.data.params.get('n_estimators')}, "
          f"max_depth={run.data.params.get('max_depth')}")

Promoting Models Through Stages

pythonCopy
client = MlflowClient()

# Find the best run
best_run = runs[0]
run_id = best_run.info.run_id

# Register if not already registered
model_uri = f"runs:/{run_id}/model"
result = mlflow.register_model(model_uri, "churn-predictor")
version = result.version
print(f"Registered as version {version}")

# Transition to Staging for testing
client.transition_model_version_stage(
    name="churn-predictor",
    version=version,
    stage="Staging",
    archive_existing_versions=False
)

# After QA passes, promote to Production
client.transition_model_version_stage(
    name="churn-predictor",
    version=version,
    stage="Production",
    archive_existing_versions=True  # archive previous production version
)

# Add description for audit trail
client.update_model_version(
    name="churn-predictor",
    version=version,
    description="RF 200 trees, depth=10. AUC 0.871. Retrained on Q4 2025 data."
)

Loading Models by Stage in Your Serving Code

pythonCopy
import mlflow.sklearn

# Always load "Production" stage  -  update this string, not the code
model = mlflow.sklearn.load_model("models:/churn-predictor/Production")

# Or load by specific version for reproducibility in tests
model_v3 = mlflow.sklearn.load_model("models:/churn-predictor/3")

This decouples deployment from training: to deploy a new model, you change its stage in the registry. The serving code does not change. Rollback is transition_stage("Production" → version_n_minus_1).

Setting Up MLflow Locally

bashCopy
pip install mlflow

# Start the tracking server with a local SQLite backend
mlflow server \
  --backend-store-uri sqlite:///mlflow.db \
  --default-artifact-root ./mlruns \
  --host 0.0.0.0 \
  --port 5000

Open http://localhost:5000 for the UI.

For a team, run MLflow on a shared server with S3 for artifact storage:

bashCopy
mlflow server \
  --backend-store-uri postgresql://user:pass@host/mlflow_db \
  --default-artifact-root s3://your-bucket/mlflow-artifacts \
  --host 0.0.0.0 \
  --port 5000

What to Log Beyond Metrics

Most tutorials only log accuracy and loss. Log these too:

pythonCopy
with mlflow.start_run():
    # Dataset characteristics
    mlflow.log_params({
        "train_size": len(X_train),
        "val_size": len(X_val),
        "positive_rate_train": y_train.mean(),
        "feature_count": X_train.shape[1],
        "data_snapshot_date": "2025-01-15"
    })

    # Confusion matrix as artifact
    cm = confusion_matrix(y_val, model.predict(X_val))
    np.savetxt("confusion_matrix.txt", cm)
    mlflow.log_artifact("confusion_matrix.txt")

    # Feature importances
    importances = pd.DataFrame({
        'feature': feature_names,
        'importance': model.feature_importances_
    }).sort_values('importance', ascending=False)
    importances.to_csv("feature_importances.csv", index=False)
    mlflow.log_artifact("feature_importances.csv")

    # Training data fingerprint (not the data itself)
    import hashlib
    data_hash = hashlib.md5(X_train.tobytes()).hexdigest()
    mlflow.log_param("train_data_hash", data_hash)

The data hash lets you verify whether two runs used the same training set without storing the data twice.

The Versioning Workflow in Practice

1. Create experiment for the problem
2. Run training sweep → multiple runs logged automatically
3. Query runs, pick best by metric
4. Register winning run as new model version
5. Test the staged version against holdout set
6. If it passes, promote to Production → archive previous
7. If production breaks, roll back by promoting archived version

This loop is the foundation of systematic ML development. Everything else - CI/CD, feature stores, drift detection - plugs into it.

Common Mistakes

Logging metrics but not the dataset version or git commit hash. A logged metric without the code and data that produced it is nearly useless for reproducing or diagnosing a result. If you cannot identify which git commit and which dataset split produced a given run's numbers, you cannot audit it, compare it fairly to another run, or reproduce it six months later. Log git hash, dataset SHA or version identifier, and all hyperparameters as part of every run's metadata.

Using model registry stages informally without defined promotion criteria. Stage transitions (Staging → Production) mean nothing if anyone can promote a model at any time for any reason. Without explicit promotion criteria (minimum eval score, human sign-off, A/B test result), the registry stages become decorative labels rather than quality gates. Define and document the promotion criteria before your first model enters the registry.

Not saving the full experiment config alongside the artifact. Saving only the model artifact means you cannot reproduce the training environment - the library versions, preprocessing pipeline, feature engineering logic, and training script are all missing. Log the full config as a JSON artifact, pin dependencies in a requirements file, and include the training script itself alongside every model artifact.

What to Practice Next

• Create an MLflow experiment that logs git commit hash (via subprocess.check_output(['git', 'rev-parse', 'HEAD'])), dataset SHA, all hyperparameters, and at least three evaluation metrics; verify you can reproduce the exact run from the logged artifacts.
• Define written promotion criteria for moving a model from Staging to Production in your registry: minimum required improvements on two metrics, required sign-off process, and rollback procedure.
• Compare two MLflow runs for the same model family side-by-side in the MLflow UI; identify the hyperparameter differences and explain which change most likely drove the metric difference.