MLflow Model Registry

MLflow's Model Registry provides a centralized model store with versioning, stage transitions, and lineage tracking. It's the most widely-used open-source model registry.

Setup

# Install MLflow
pip install mlflow

# Start MLflow server with registry
mlflow server --backend-store-uri sqlite:///mlflow.db \
              --default-artifact-root ./mlartifacts \
              --host 0.0.0.0 --port 5000

Registering Models

From an Experiment Run

import mlflow
from sklearn.ensemble import RandomForestClassifier

# Set tracking URI
mlflow.set_tracking_uri("http://localhost:5000")

# Train and log model
with mlflow.start_run() as run:
    model = RandomForestClassifier(n_estimators=100)
    model.fit(X_train, y_train)

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

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

    # Register model
    mlflow.register_model(
        f"runs:/{run.info.run_id}/model",
        "fraud-detector"
    )

Direct Registration

from mlflow import MlflowClient

client = MlflowClient()

# Register from run
result = client.create_registered_model("fraud-detector")

# Add version
client.create_model_version(
    name="fraud-detector",
    source=f"runs:/{run_id}/model",
    run_id=run_id
)

Managing Model Stages

Stage Transitions

from mlflow import MlflowClient

client = MlflowClient()

# Transition to staging
client.transition_model_version_stage(
    name="fraud-detector",
    version=1,
    stage="Staging"
)

# After testing, promote to production
client.transition_model_version_stage(
    name="fraud-detector",
    version=1,
    stage="Production"
)

# Archive old production model
client.transition_model_version_stage(
    name="fraud-detector",
    version=0,
    stage="Archived"
)

Model Aliases (MLflow 2.x)

# Set alias for easier reference
client.set_registered_model_alias(
    name="fraud-detector",
    alias="champion",
    version=3
)

# Load by alias
model = mlflow.pyfunc.load_model("models:/fraud-detector@champion")

Loading Models

By Version

import mlflow

# Load specific version
model_v1 = mlflow.pyfunc.load_model("models:/fraud-detector/1")
model_v2 = mlflow.pyfunc.load_model("models:/fraud-detector/2")

By Stage

# Load production model
production_model = mlflow.pyfunc.load_model("models:/fraud-detector/Production")

# Load staging model
staging_model = mlflow.pyfunc.load_model("models:/fraud-detector/Staging")

By Alias

# Load champion model
champion_model = mlflow.pyfunc.load_model("models:/fraud-detector@champion")

Model Metadata

Adding Descriptions

client = MlflowClient()

# Model-level description
client.update_registered_model(
    name="fraud-detector",
    description="""
    Real-time fraud detection model for payment transactions.
    Owner: risk-team@company.com
    SLA: 99.9% uptime, <50ms p99 latency
    """
)

# Version-level description
client.update_model_version(
    name="fraud-detector",
    version=3,
    description="""
    v3: Added new features for device fingerprinting.
    Accuracy: 0.95 (+0.05 vs v2)
    Training data: 2025-01-01 to 2025-01-15
    """
)

Querying the Registry

List Models

client = MlflowClient()

# List all registered models
for model in client.search_registered_models():
    print(f"Model: {model.name}")
    for version in model.latest_versions:
        print(f"  Version {version.version}: {version.current_stage}")

Search Models

# Search by name
models = client.search_registered_models(filter_string="name LIKE 'fraud%'")

# Search versions by tag
versions = client.search_model_versions(
    filter_string="tags.approved_by = 'alice@company.com'"
)

Get Latest Production Model

def get_production_model(model_name: str):
    """Get the current production model."""
    client = MlflowClient()

    for version in client.search_model_versions(f"name='{model_name}'"):
        if version.current_stage == "Production":
            return mlflow.pyfunc.load_model(f"models:/{model_name}/Production")

    raise ValueError(f"No production model found for {model_name}")

CI/CD Integration

Automated Testing Before Promotion

def test_and_promote(model_name: str, version: int):
    """Test model and promote if passing."""
    client = MlflowClient()

    # Load staging model
    model = mlflow.pyfunc.load_model(f"models:/{model_name}/Staging")

    # Run tests
    predictions = model.predict(test_data)
    accuracy = calculate_accuracy(predictions, test_labels)

    if accuracy >= 0.90:
        # Promote to production
        client.transition_model_version_stage(
            name=model_name,
            version=version,
            stage="Production",
            archive_existing_versions=True
        )
        return True
    else:
        print(f"Model failed tests: accuracy={accuracy}")
        return False

GitHub Actions Example

# .github/workflows/model-promotion.yml
name: Model Promotion

on:
  workflow_dispatch:
    inputs:
      model_name:
        required: true
      version:
        required: true

jobs:
  promote:
    runs-on: ubuntu-latest
    steps:
      - name: Test and promote model
        env:
          MLFLOW_TRACKING_URI: ${{ secrets.MLFLOW_URI }}
        run: |
          python scripts/test_and_promote.py \
            --model ${{ inputs.model_name }} \
            --version ${{ inputs.version }}

Best Practices

Practice	Why
Use aliases	Decouple code from versions
Require descriptions	Document what changed
Automate transitions	Reduce human error
Test before promotion	Catch regressions
Archive, don't delete	Maintain audit trail

Key insight: MLflow Model Registry integrates seamlessly with MLflow tracking, giving you end-to-end lineage from experiment to production model.

Next, we'll explore model serving with BentoML. :::