ML Model Versioning Done Right with MLflow

You wouldn't ship code without version control. Why ship models without it?

The Problem

Every ML team hits this wall:

• "Which model is in production?"
• "What hyperparameters did we use?"
• "Can we reproduce last month's results?"

MLflow Model Registry

python
import mlflow
from mlflow.tracking import MlflowClient

# Log your model with metadata
with mlflow.start_run():
    mlflow.log_params({
        "learning_rate": 0.001,
        "epochs": 100,
        "architecture": "transformer-base"
    })
    mlflow.log_metrics({"accuracy": 0.94, "f1": 0.91})
    mlflow.sklearn.log_model(model, "model")

Lifecycle Management

None → Staging → Production → Archived

python
client = MlflowClient()

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

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

Serving Models

bash
# Serve any registered model
mlflow models serve -m "models:/fraud-detector/Production" -p 5001

Or use the Python API:

python
model = mlflow.pyfunc.load_model("models:/fraud-detector/Production")
predictions = model.predict(data)

Best Practices

1. Tag everything: Model type, data version, training date
2. Automate promotion: CI/CD for models, not just code
3. Keep artifacts: Store training data snapshots alongside models

Our Setup

yaml
# docker-compose.yml
services:
  mlflow:
    image: mlflow:latest
    environment:
      - BACKEND_STORE_URI=postgresql://...
      - ARTIFACT_ROOT=s3://mlflow-artifacts/
    ports:
      - "5000:5000"

This saved us from three "which model is this?" incidents in one quarter.