Chapter 3: Model Versioning

Introduction

Model versioning is critical for production ML systems. It enables tracking, comparing, and rolling back models while maintaining a clear lineage from data to deployment.

Why Model Versioning?

The Model Management Challenge

Benefits of Model Versioning

Benefit	Description
Reproducibility	Recreate any model version
Lineage	Track data → model → predictions
Comparison	Compare performance across versions
Rollback	Quickly revert to previous versions
Compliance	Audit trail for regulated industries

Model Registry Concepts

Key Components

Model Lifecycle Stages

Stage	Description
None	Initial state, newly registered
Staging	Model is being tested/validated
Production	Model is serving live traffic
Archived	Model is deprecated, kept for reference

MLflow Model Registry

Setting Up MLflow

# Install MLflow
pip install mlflow

# Start MLflow server with backend store
mlflow server \
    --backend-store-uri postgresql://user:pass@localhost/mlflow \
    --default-artifact-root s3://mlflow-artifacts/ \
    --host 0.0.0.0 \
    --port 5000

Registering Models

# register_model.py
import mlflow
import mlflow.sklearn
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:
    # Train model
    model = RandomForestClassifier(n_estimators=100)
    model.fit(X_train, y_train)
    
    # Log parameters
    mlflow.log_param("n_estimators", 100)
    mlflow.log_param("max_depth", None)
    
    # Log metrics
    accuracy = model.score(X_test, y_test)
    mlflow.log_metric("accuracy", accuracy)
    
    # Log model
    mlflow.sklearn.log_model(
        model,
        "model",
        registered_model_name="customer-churn-model"
    )
    
    print(f"Model registered with run_id: {run.info.run_id}")

Managing Model Versions

# model_management.py
from mlflow.tracking import MlflowClient

client = MlflowClient()

# Get all versions of a model
versions = client.search_model_versions("name='customer-churn-model'")
for v in versions:
    print(f"Version {v.version}: {v.current_stage}")

# Transition model to staging
client.transition_model_version_stage(
    name="customer-churn-model",
    version=2,
    stage="Staging"
)

# Transition model to production
client.transition_model_version_stage(
    name="customer-churn-model",
    version=2,
    stage="Production",
    archive_existing_versions=True  # Archive current prod version
)

# Add description to model version
client.update_model_version(
    name="customer-churn-model",
    version=2,
    description="Improved feature engineering, AUC increased by 5%"
)

# Add tags
client.set_model_version_tag(
    name="customer-churn-model",
    version=2,
    key="validation_status",
    value="approved"
)

Loading Models from Registry

# load_model.py
import mlflow

# Load specific version
model_v1 = mlflow.sklearn.load_model(
    "models:/customer-churn-model/1"
)

# Load by stage
production_model = mlflow.sklearn.load_model(
    "models:/customer-churn-model/Production"
)

# Load staging model
staging_model = mlflow.sklearn.load_model(
    "models:/customer-churn-model/Staging"
)

# Make predictions
predictions = production_model.predict(X_new)

Model Metadata and Lineage

Comprehensive Model Card

# model_card.py
import json
from dataclasses import dataclass, asdict
from typing import List, Dict, Optional
from datetime import datetime

@dataclass
class ModelCard:
    """Structured model documentation following industry standards."""
    
    # Basic info
    model_name: str
    version: str
    description: str
    owner: str
    created_at: str
    
    # Training info
    training_data: Dict[str, str]
    features: List[str]
    target: str
    algorithm: str
    hyperparameters: Dict[str, any]
    
    # Performance
    metrics: Dict[str, float]
    evaluation_data: str
    
    # Intended use
    intended_use: str
    out_of_scope_use: str
    
    # Limitations and biases
    limitations: List[str]
    ethical_considerations: List[str]
    
    # Lineage
    parent_model: Optional[str] = None
    training_run_id: Optional[str] = None
    data_version: Optional[str] = None
    code_version: Optional[str] = None

# Create model card
model_card = ModelCard(
    model_name="customer-churn-model",
    version="2.1.0",
    description="Predicts customer churn probability within 30 days",
    owner="ml-team@company.com",
    created_at=datetime.now().isoformat(),
    
    training_data={
        "source": "s3://data-lake/customer-features/",
        "date_range": "2023-01-01 to 2024-01-01",
        "records": "1,500,000"
    },
    features=[
        "account_age_days",
        "total_purchases",
        "days_since_last_login",
        "support_tickets_90d",
        "nps_score"
    ],
    target="churned_within_30_days",
    algorithm="GradientBoostingClassifier",
    hyperparameters={
        "n_estimators": 200,
        "max_depth": 7,
        "learning_rate": 0.1,
        "min_samples_split": 10
    },
    
    metrics={
        "auc_roc": 0.89,
        "precision": 0.82,
        "recall": 0.75,
        "f1_score": 0.78
    },
    evaluation_data="s3://data-lake/test-set-2024-01/",
    
    intended_use="Identify at-risk customers for retention campaigns",
    out_of_scope_use="Credit scoring, hiring decisions",
    
    limitations=[
        "Performance may degrade for customers < 30 days old",
        "Not validated for enterprise tier customers",
        "Requires minimum 6 months of transaction history"
    ],
    ethical_considerations=[
        "Model should not be used to deny service",
        "Predictions should be human-reviewed before action"
    ],
    
    parent_model="customer-churn-model/1.0.0",
    training_run_id="abc123def456",
    data_version="data-v2.1.0",
    code_version="git-sha-789xyz"
)

# Save model card
with open("model_card.json", "w") as f:
    json.dump(asdict(model_card), f, indent=2)

Logging Model Lineage

# lineage_tracking.py
import mlflow
import git
import hashlib

def get_code_version():
    """Get current git commit hash."""
    repo = git.Repo(search_parent_directories=True)
    return repo.head.object.hexsha

def get_data_hash(data_path: str) -> str:
    """Calculate hash of data file."""
    with open(data_path, 'rb') as f:
        return hashlib.md5(f.read()).hexdigest()

def train_with_lineage(X_train, y_train, X_test, y_test, data_path: str):
    """Train model with complete lineage tracking."""
    
    with mlflow.start_run() as run:
        # Track code version
        code_version = get_code_version()
        mlflow.set_tag("code_version", code_version)
        mlflow.set_tag("git_repo", "https://github.com/company/ml-project")
        
        # Track data version
        data_hash = get_data_hash(data_path)
        mlflow.set_tag("data_hash", data_hash)
        mlflow.set_tag("data_path", data_path)
        
        # Track environment
        mlflow.set_tag("python_version", "3.10.0")
        mlflow.set_tag("sklearn_version", "1.3.0")
        
        # Track training configuration
        mlflow.log_param("training_samples", len(X_train))
        mlflow.log_param("test_samples", len(X_test))
        mlflow.log_param("features", list(X_train.columns))
        
        # Train model
        model = GradientBoostingClassifier(
            n_estimators=200,
            max_depth=7,
            learning_rate=0.1
        )
        model.fit(X_train, y_train)
        
        # Log metrics
        train_accuracy = model.score(X_train, y_train)
        test_accuracy = model.score(X_test, y_test)
        mlflow.log_metric("train_accuracy", train_accuracy)
        mlflow.log_metric("test_accuracy", test_accuracy)
        
        # Log model with signature
        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="customer-churn-model"
        )
        
        return run.info.run_id

Model Versioning Strategies

Semantic Versioning for Models

MAJOR.MINOR.PATCH

Examples:
- 1.0.0 → 1.0.1  # Patch: Bug fix, no retraining
- 1.0.0 → 1.1.0  # Minor: Retrained, same architecture
- 1.0.0 → 2.0.0  # Major: New architecture or breaking changes

Automated Version Bumping

# version_manager.py
import re
from enum import Enum
from mlflow.tracking import MlflowClient

class VersionBump(Enum):
    MAJOR = "major"
    MINOR = "minor"
    PATCH = "patch"

def get_next_version(current_version: str, bump: VersionBump) -> str:
    """Calculate next semantic version."""
    match = re.match(r"(\d+)\.(\d+)\.(\d+)", current_version)
    if not match:
        return "1.0.0"
    
    major, minor, patch = map(int, match.groups())
    
    if bump == VersionBump.MAJOR:
        return f"{major + 1}.0.0"
    elif bump == VersionBump.MINOR:
        return f"{major}.{minor + 1}.0"
    else:
        return f"{major}.{minor}.{patch + 1}"

def register_new_version(
    model_name: str,
    run_id: str,
    bump: VersionBump,
    description: str
):
    """Register model with proper versioning."""
    client = MlflowClient()
    
    # Get latest version
    versions = client.search_model_versions(f"name='{model_name}'")
    if versions:
        latest_version = max(v.version for v in versions)
        latest_tags = client.get_model_version(model_name, latest_version).tags
        current_semantic = latest_tags.get("semantic_version", "0.0.0")
    else:
        current_semantic = "0.0.0"
    
    # Calculate new semantic version
    new_semantic = get_next_version(current_semantic, bump)
    
    # Register model
    result = client.create_model_version(
        name=model_name,
        source=f"runs:/{run_id}/model",
        run_id=run_id,
        description=description
    )
    
    # Set semantic version tag
    client.set_model_version_tag(
        name=model_name,
        version=result.version,
        key="semantic_version",
        value=new_semantic
    )
    
    print(f"Registered {model_name} v{result.version} (semantic: {new_semantic})")
    return result

Model Approval Workflow

Automated Validation Pipeline

# model_validation.py
from typing import Dict, List, Tuple
import mlflow
from mlflow.tracking import MlflowClient

class ModelValidator:
    """Validate models before production deployment."""
    
    def __init__(self, model_name: str, version: int):
        self.client = MlflowClient()
        self.model_name = model_name
        self.version = version
        self.model_uri = f"models:/{model_name}/{version}"
        
    def validate_performance(
        self, 
        X_test, 
        y_test,
        thresholds: Dict[str, float]
    ) -> Tuple[bool, Dict[str, float]]:
        """Check if model meets performance thresholds."""
        model = mlflow.sklearn.load_model(self.model_uri)
        
        from sklearn.metrics import accuracy_score, roc_auc_score, f1_score
        
        predictions = model.predict(X_test)
        probabilities = model.predict_proba(X_test)[:, 1]
        
        metrics = {
            "accuracy": accuracy_score(y_test, predictions),
            "auc_roc": roc_auc_score(y_test, probabilities),
            "f1_score": f1_score(y_test, predictions)
        }
        
        # Check against thresholds
        passed = all(
            metrics.get(metric, 0) >= threshold
            for metric, threshold in thresholds.items()
        )
        
        return passed, metrics
    
    def validate_latency(
        self, 
        X_sample,
        max_latency_ms: float = 100
    ) -> Tuple[bool, float]:
        """Check if model meets latency requirements."""
        import time
        
        model = mlflow.sklearn.load_model(self.model_uri)
        
        # Warm up
        _ = model.predict(X_sample[:1])
        
        # Measure latency
        times = []
        for _ in range(100):
            start = time.time()
            _ = model.predict(X_sample[:1])
            times.append((time.time() - start) * 1000)
        
        p99_latency = sorted(times)[98]
        passed = p99_latency <= max_latency_ms
        
        return passed, p99_latency
    
    def validate_consistency(
        self,
        X_sample,
        production_model_uri: str
    ) -> Tuple[bool, float]:
        """Check predictions match production model."""
        new_model = mlflow.sklearn.load_model(self.model_uri)
        prod_model = mlflow.sklearn.load_model(production_model_uri)
        
        new_predictions = new_model.predict(X_sample)
        prod_predictions = prod_model.predict(X_sample)
        
        agreement = (new_predictions == prod_predictions).mean()
        passed = agreement >= 0.95  # At least 95% agreement
        
        return passed, agreement
    
    def run_all_validations(
        self,
        X_test,
        y_test,
        thresholds: Dict[str, float]
    ) -> Dict[str, any]:
        """Run all validations and return results."""
        results = {
            "model_name": self.model_name,
            "version": self.version,
            "validations": {}
        }
        
        # Performance validation
        perf_passed, perf_metrics = self.validate_performance(
            X_test, y_test, thresholds
        )
        results["validations"]["performance"] = {
            "passed": perf_passed,
            "metrics": perf_metrics
        }
        
        # Latency validation
        latency_passed, p99_latency = self.validate_latency(X_test[:100])
        results["validations"]["latency"] = {
            "passed": latency_passed,
            "p99_ms": p99_latency
        }
        
        # Overall result
        results["approved"] = all(
            v["passed"] for v in results["validations"].values()
        )
        
        # Update model tags
        status = "approved" if results["approved"] else "rejected"
        self.client.set_model_version_tag(
            name=self.model_name,
            version=self.version,
            key="validation_status",
            value=status
        )
        
        return results

Key Takeaways

Summary

1. Use a model registry for centralized model management
2. Track complete lineage - data, code, parameters, environment
3. Implement semantic versioning for clear version progression
4. Create model cards for documentation and compliance
5. Automate validation before production deployment
6. Maintain audit trails for regulated industries

Practical Exercise

1. Set up MLflow tracking server
2. Train and register multiple model versions
3. Create a model card for your best model
4. Implement model validation pipeline
5. Practice transitioning models between stages

---

Next: Module 2 Quiz