Skip to content

Model Versioning

Model versioning ensures you can track, manage, and reproduce ML models. This sheet covers version control, model registry, stage transitions, rollback, and governance.

Table of Contents

Installation / Setup

pip install mlflow

Core Syntax

Version Control Strategies

Strategy Description When to Use
Git + DVC Code in git, data/models in DVC Full version control
MLflow Registry Centralized model management Team collaboration
Artifact Storage S3/GCS for model files Simple versioning
Database Store metadata in DB Custom governance

Model Stages

None → Staging → Production → Archived

None:     New model, not yet reviewed
Staging:  Under testing/evaluation
Production: Live, serving predictions
Archived: Deprecated, no longer in use

Common Operations

MLflow Model Registry

import mlflow

# Register model from run
model_uri = "runs:/run_id/model"
mlflow.register_model(model_uri, "my_model")

# Register with description
mlflow.register_model(
    model_uri,
    "my_model",
    description="Churn prediction model v1"
)

# Get model versions
client = mlflow.tracking.MlflowClient()
versions = client.get_latest_versions("my_model")
for v in versions:
    print(f"Version {v.version}: stage={v.current_stage}")

Stage Transitions

client = mlflow.tracking.MlflowClient()

# Transition to staging
client.transition_model_version_stage(
    name="my_model",
    version=1,
    stage="staging"
)

# Transition to production
client.transition_model_version_stage(
    name="my_model",
    version=1,
    stage="production"
)

# Archive old version
client.transition_model_version_stage(
    name="my_model",
    version=1,
    stage="archived"
)

Model Metadata

# Update model description
client.update_model_version(
    name="my_model",
    version=1,
    description="Best model from experiment exp-123"
)

# Add tags
client.set_model_version_tag(
    name="my_model",
    version=1,
    key="team",
    value="ml-team"
)

# Get model info
model_info = client.get_model_version("my_model", 1)
print(model_info.tags)
print(model_info.description)

Loading Models

import mlflow

# Load from registry by stage
model = mlflow.pyfunc.load_model("models:/my_model/production")

# Load from registry by version
model = mlflow.pyfunc.load_model("models:/my_model/1")

# Load from run
model = mlflow.pyfunc.load_model("runs:/run_id/model")

DVC Model Versioning

# Track model file
dvc add model.pkl
git add model.pkl.dvc .gitignore
git commit -m "Add model v1"

# Update model
cp new_model.pkl model.pkl
dvc add model.pkl
git add model.pkl.dvc
git commit -m "Update model v2"

# Checkout specific version
git checkout v1
dvc checkout
# model.pkl is now v1

Practical Examples

Example 1 — Full versioning workflow

import mlflow
from mlflow.tracking import MlflowClient
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# 1. Train and log model
mlflow.set_experiment("churn_prediction")

with mlflow.start_run(run_name="v1"):
    model = RandomForestClassifier(n_estimators=100)
    model.fit(X_train, y_train)

    accuracy = accuracy_score(y_test, model.predict(X_test))

    mlflow.log_param("n_estimators", 100)
    mlflow.log_metric("accuracy", accuracy)
    mlflow.sklearn.log_model(model, "model")

    run_id = mlflow.active_run().info.run_id

# 2. Register model
client = MlflowClient()
model_uri = f"runs:/{run_id}/model"
client.create_model_version("churn_model", model_uri)

# 3. Transition to staging
client.transition_model_version_stage("churn_model", 1, "staging")

# 4. Run validation
# ... validation code ...

# 5. Transition to production
client.transition_model_version_stage("churn_model", 1, "production")

# 6. Add description
client.update_model_version(
    "churn_model", 1,
    description="Production model: accuracy=0.95"
)

Example 2 — Rollback

import mlflow
from mlflow.tracking import MlflowClient

client = MlflowClient()

# Get all versions
versions = client.get_latest_versions("churn_model", stages=["production"])

# Find previous production version
if len(versions) > 1:
    previous_version = versions[1]

    # Rollback
    client.transition_model_version_stage(
        "churn_model",
        previous_version.version,
        "production"
    )

    print(f"Rolled back to version {previous_version.version}")

Example 3 — Version comparison

import mlflow
from mlflow.tracking import MlflowClient

client = MlflowClient()

# Get model versions
versions = client.get_latest_versions("churn_model")

# Compare versions
for v in versions:
    run = client.get_run(v.run_id)
    print(f"Version {v.version}:")
    print(f"  Stage: {v.current_stage}")
    print(f"  Metrics: {run.data.metrics}")
    print(f"  Description: {v.description}")

Common Patterns

Pattern: version tagging

import mlflow

# Tag runs with version
with mlflow.start_run():
    mlflow.set_tag("version", "1.0.0")
    mlflow.set_tag("git_commit", "abc123")
    mlflow.set_tag("dataset_version", "v2")

Pattern: A/B testing versions

import mlflow
import random

def get_version(user_id):
    # Deterministic routing
    if hash(user_id) % 100 < 10:  # 10% to new version
        return "2"
    return "1"

def predict(user_id, data):
    version = get_version(user_id)
    model = mlflow.pyfunc.load_model(f"models:/churn_model/{version}")
    return model.predict(data)

Pattern: model lineage

# Track complete lineage
with mlflow.start_run() as run:
    # Log data version
    mlflow.log_param("data_version", "v2")

    # Log code version
    import subprocess
    git_commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode().strip()
    mlflow.log_param("git_commit", git_commit)

    # Log model
    mlflow.sklearn.log_model(model, "model")

    # Log metrics
    mlflow.log_metric("accuracy", accuracy)

Best Practices

  • Use model registry — centralized version management.
  • Tag versions — add metadata for traceability.
  • Use stage transitions — None → Staging → Production.
  • Keep previous versions — enable rollback.
  • Document changes — add descriptions to versions.
  • Test before promoting — validate in staging first.
  • Archive old versions — clean up production.
  • Track lineage — data version → code version → model version.

Common Mistakes

Mistake Fix
No version control Use MLflow or DVC
No stage transitions Use None → Staging → Production
No rollback plan Keep previous versions
No documentation Add descriptions to versions
No lineage tracking Track data → code → model

Troubleshooting

Symptom Cause / Fix
Model not found Check model name; check version/stage
Stage transition fails Check permissions; check model exists
Cannot rollback Check version exists; check stage
Version conflict Use unique version numbers
Metadata missing Add tags and descriptions

Official Documentation

Last Updated

2026-09-06 — Initial version.