4_model_registry.py

"""
MLflow Tutorial - Component 4: MLflow Model Registry

This module demonstrates the MLflow Model Registry component, which provides
centralized model store, versioning, and lifecycle management.

Key Features Demonstrated:
1. Registering models
2. Model versioning
3. Stage transitions (None -> Staging -> Production -> Archived)
4. Model annotations and descriptions
5. Loading models from registry
6. Comparing model versions
"""

import mlflow
import mlflow.sklearn
from mlflow.tracking import MlflowClient
import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.metrics import accuracy_score, f1_score
import time


def print_section(title):
    """Print a formatted section header."""
    print("\n" + "=" * 80)
    print(f"  {title}")
    print("=" * 80 + "\n")


def create_sample_data():
    """Create sample data."""
    X, y = make_classification(
        n_samples=1000,
        n_features=20,
        n_informative=15,
        n_classes=2,
        random_state=42
    )
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    return X_train, X_test, y_train, y_test


def register_first_model():
    """Create and register the first version of a model."""
    print_section("1. Registering First Model Version")
    
    mlflow.set_experiment("Model_Registry_Demo")
    
    X_train, X_test, y_train, y_test = create_sample_data()
    
    model_name = "classification_model"
    
    with mlflow.start_run(run_name="register_v1") as run:
        print(f"Run ID: {run.info.run_id}")
        
        # Train model
        print("\nTraining RandomForest model (v1)...")
        model = RandomForestClassifier(
            n_estimators=100,
            max_depth=5,
            random_state=42
        )
        model.fit(X_train, y_train)
        
        # Evaluate
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        f1 = f1_score(y_test, y_pred)
        
        print(f"Accuracy: {accuracy:.4f}")
        print(f"F1 Score: {f1:.4f}")
        
        # Log metrics
        mlflow.log_params({
            "n_estimators": 100,
            "max_depth": 5,
            "model_type": "RandomForest"
        })
        mlflow.log_metrics({
            "accuracy": accuracy,
            "f1_score": f1
        })
        
        # Log and register model
        from mlflow.models.signature import infer_signature
        signature = infer_signature(X_test, model.predict(X_test))
        
        model_info = mlflow.sklearn.log_model(
            model,
            "model",
            signature=signature,
            registered_model_name=model_name
        )
        
        print(f"\n✓ Model registered as: {model_name}")
        print(f"Version: 1")
        
        return run.info.run_id, model_name, accuracy


def register_improved_model(model_name):
    """Create and register an improved version of the model."""
    print_section("2. Registering Improved Model Version")
    
    mlflow.set_experiment("Model_Registry_Demo")
    
    X_train, X_test, y_train, y_test = create_sample_data()
    
    with mlflow.start_run(run_name="register_v2") as run:
        print(f"Run ID: {run.info.run_id}")
        
        # Train improved model
        print("\nTraining improved RandomForest model (v2)...")
        model = RandomForestClassifier(
            n_estimators=200,  # More trees
            max_depth=7,       # Deeper trees
            min_samples_split=5,
            random_state=42
        )
        model.fit(X_train, y_train)
        
        # Evaluate
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        f1 = f1_score(y_test, y_pred)
        
        print(f"Accuracy: {accuracy:.4f}")
        print(f"F1 Score: {f1:.4f}")
        
        # Log metrics
        mlflow.log_params({
            "n_estimators": 200,
            "max_depth": 7,
            "min_samples_split": 5,
            "model_type": "RandomForest"
        })
        mlflow.log_metrics({
            "accuracy": accuracy,
            "f1_score": f1
        })
        
        # Log and register model (this creates version 2)
        from mlflow.models.signature import infer_signature
        signature = infer_signature(X_test, model.predict(X_test))
        
        mlflow.sklearn.log_model(
            model,
            "model",
            signature=signature,
            registered_model_name=model_name
        )
        
        print(f"\n✓ New version registered for: {model_name}")
        print(f"Version: 2")
        
        return run.info.run_id, accuracy


def register_alternative_model(model_name):
    """Register a different algorithm as a new version."""
    print_section("3. Registering Alternative Model (Different Algorithm)")
    
    mlflow.set_experiment("Model_Registry_Demo")
    
    X_train, X_test, y_train, y_test = create_sample_data()
    
    with mlflow.start_run(run_name="register_v3_gbm") as run:
        print(f"Run ID: {run.info.run_id}")
        
        # Train GradientBoosting model
        print("\nTraining GradientBoosting model (v3)...")
        model = GradientBoostingClassifier(
            n_estimators=100,
            max_depth=5,
            learning_rate=0.1,
            random_state=42
        )
        model.fit(X_train, y_train)
        
        # Evaluate
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        f1 = f1_score(y_test, y_pred)
        
        print(f"Accuracy: {accuracy:.4f}")
        print(f"F1 Score: {f1:.4f}")
        
        # Log metrics
        mlflow.log_params({
            "n_estimators": 100,
            "max_depth": 5,
            "learning_rate": 0.1,
            "model_type": "GradientBoosting"
        })
        mlflow.log_metrics({
            "accuracy": accuracy,
            "f1_score": f1
        })
        
        # Log and register model
        from mlflow.models.signature import infer_signature
        signature = infer_signature(X_test, model.predict(X_test))
        
        mlflow.sklearn.log_model(
            model,
            "model",
            signature=signature,
            registered_model_name=model_name
        )
        
        print(f"\n✓ Alternative model registered for: {model_name}")
        print(f"Version: 3")
        
        return run.info.run_id, accuracy


def manage_model_stages(model_name):
    """Demonstrate model stage transitions."""
    print_section("4. Managing Model Lifecycle Stages")
    
    client = MlflowClient()
    
    print("Model lifecycle stages:")
    print("  None -> Staging -> Production -> Archived")
    print()
    
    # Get all versions
    versions = client.search_model_versions(f"name='{model_name}'")
    
    print(f"Current versions of '{model_name}':")
    for version in versions:
        print(f"  Version {version.version}: Stage = {version.current_stage}")
    
    # Transition version 1 to Staging
    print(f"\n→ Transitioning version 1 to Staging...")
    client.transition_model_version_stage(
        name=model_name,
        version=1,
        stage="Staging"
    )
    print("✓ Version 1 is now in Staging")
    
    # Transition version 2 to Production
    print(f"\n→ Transitioning version 2 to Production...")
    client.transition_model_version_stage(
        name=model_name,
        version=2,
        stage="Production"
    )
    print("✓ Version 2 is now in Production")
    
    # Later, transition version 1 to Archived
    print(f"\n→ Transitioning version 1 to Archived...")
    client.transition_model_version_stage(
        name=model_name,
        version=3,
        stage="Archived"
    )
    print("✓ Version 3 is now Archived")
    
    # Show updated stages
    print("\n" + "-" * 80)
    print("Updated model stages:")
    print("-" * 80)
    versions = client.search_model_versions(f"name='{model_name}'")
    for version in versions:
        print(f"  Version {version.version}: {version.current_stage}")
    
    print("\n✓ Model stage management demonstration complete")


def add_model_descriptions(model_name):
    """Add descriptions and tags to models."""
    print_section("5. Adding Model Descriptions and Tags")
    
    client = MlflowClient()
    
    # Update model description
    print("Adding description to registered model...")
    client.update_registered_model(
        name=model_name,
        description="Classification model for binary prediction. "
                   "Trained on synthetic data with 20 features. "
                   "Production model achieves >90% accuracy."
    )
    print("✓ Model description added")
    
    # Add descriptions to specific versions
    print("\nAdding descriptions to model versions...")
    
    client.update_model_version(
        name=model_name,
        version=1,
        description="Initial baseline model. RandomForest with 100 trees and max_depth=5."
    )
    print("✓ Version 1 description added")
    
    client.update_model_version(
        name=model_name,
        version=2,
        description="Improved model. RandomForest with 200 trees and max_depth=7. "
                   "Better performance than v1. Currently in Production."
    )
    print("✓ Version 2 description added")
    
    client.update_model_version(
        name=model_name,
        version=3,
        description="Alternative algorithm. GradientBoosting with 100 estimators. "
                   "Similar performance to v2 but different approach."
    )
    print("✓ Version 3 description added")
    
    # Add tags
    print("\nAdding tags to model versions...")
    client.set_model_version_tag(model_name, 1, "algorithm", "RandomForest")
    client.set_model_version_tag(model_name, 1, "status", "baseline")
    
    client.set_model_version_tag(model_name, 2, "algorithm", "RandomForest")
    client.set_model_version_tag(model_name, 2, "status", "production")
    client.set_model_version_tag(model_name, 2, "validated", "true")
    
    client.set_model_version_tag(model_name, 3, "algorithm", "GradientBoosting")
    client.set_model_version_tag(model_name, 3, "status", "experimental")
    
    print("✓ Tags added to all versions")
    
    # Display model information
    print("\n" + "-" * 80)
    print("Model Registry Information:")
    print("-" * 80)
    
    registered_model = client.get_registered_model(model_name)
    print(f"\nModel Name: {registered_model.name}")
    print(f"Description: {registered_model.description}")
    print(f"Creation Time: {registered_model.creation_timestamp}")
    print(f"Last Updated: {registered_model.last_updated_timestamp}")
    
    versions = client.search_model_versions(f"name='{model_name}'")
    for version in versions:
        print(f"\n  Version {version.version}:")
        print(f"    Stage: {version.current_stage}")
        print(f"    Description: {version.description}")
        print(f"    Tags: {version.tags}")


def load_models_from_registry(model_name):
    """Demonstrate loading models from the registry."""
    print_section("6. Loading Models from Registry")
    
    X_train, X_test, y_train, y_test = create_sample_data()
    
    # Method 1: Load by version
    print("Method 1: Load specific version")
    print("-" * 80)
    model_version_uri = f"models:/{model_name}/2"
    loaded_model_v2 = mlflow.sklearn.load_model(model_version_uri)
    predictions_v2 = loaded_model_v2.predict(X_test[:5])
    print(f"Model URI: {model_version_uri}")
    print(f"Predictions: {predictions_v2}")
    
    # Method 2: Load by stage
    print("\n\nMethod 2: Load by stage (Production)")
    print("-" * 80)
    model_stage_uri = f"models:/{model_name}/Production"
    loaded_model_prod = mlflow.pyfunc.load_model(model_stage_uri)
    predictions_prod = loaded_model_prod.predict(X_test[:5])
    print(f"Model URI: {model_stage_uri}")
    print(f"Predictions: {predictions_prod}")
    
    # Method 3: Load latest version
    print("\n\nMethod 3: Load latest version")
    print("-" * 80)
    client = MlflowClient()
    versions = client.search_model_versions(f"name='{model_name}'")
    latest_version = max([int(v.version) for v in versions])
    
    model_latest_uri = f"models:/{model_name}/{latest_version}"
    loaded_model_latest = mlflow.sklearn.load_model(model_latest_uri)
    predictions_latest = loaded_model_latest.predict(X_test[:5])
    print(f"Latest version: {latest_version}")
    print(f"Model URI: {model_latest_uri}")
    print(f"Predictions: {predictions_latest}")
    
    print("\n✓ Model loading demonstration complete")


def compare_model_versions(model_name):
    """Compare different model versions."""
    print_section("7. Comparing Model Versions")
    
    client = MlflowClient()
    
    # Get all versions
    versions = client.search_model_versions(f"name='{model_name}'")
    
    print(f"Comparison of all versions of '{model_name}':")
    print("=" * 80)
    print(f"{'Version':<10} {'Stage':<15} {'Algorithm':<20} {'Run ID':<35}")
    print("-" * 80)
    
    for version in sorted(versions, key=lambda v: int(v.version)):
        # Get run info to fetch metrics
        run = client.get_run(version.run_id)
        algorithm = run.data.params.get('model_type', 'Unknown')
        
        print(f"{version.version:<10} {version.current_stage:<15} "
              f"{algorithm:<20} {version.run_id:<35}")
    
    # Compare metrics
    print("\n" + "=" * 80)
    print("Performance Metrics:")
    print("=" * 80)
    print(f"{'Version':<10} {'Accuracy':<12} {'F1 Score':<12} {'Stage':<15}")
    print("-" * 80)
    
    for version in sorted(versions, key=lambda v: int(v.version)):
        run = client.get_run(version.run_id)
        accuracy = run.data.metrics.get('accuracy', 0)
        f1 = run.data.metrics.get('f1_score', 0)
        
        print(f"{version.version:<10} {accuracy:<12.4f} {f1:<12.4f} "
              f"{version.current_stage:<15}")
    
    print("\n✓ Model comparison complete")


def demonstrate_model_aliasing(model_name):
    """Demonstrate model aliasing feature."""
    print_section("8. Model Aliasing (Alternative to Stages)")
    
    client = MlflowClient()
    
    print("Model aliases provide a more flexible alternative to stages.")
    print("You can assign custom names to model versions.\n")
    
    # Set aliases
    print("Setting aliases...")
    client.set_registered_model_alias(model_name, "champion", 2)
    client.set_registered_model_alias(model_name, "challenger", 3)
    
    print("✓ Alias 'champion' set to version 2")
    print("✓ Alias 'challenger' set to version 3")
    
    # Load by alias
    print("\nLoading models by alias...")
    champion_uri = f"models:/{model_name}@champion"
    challenger_uri = f"models:/{model_name}@challenger"
    
    print(f"Champion model URI: {champion_uri}")
    print(f"Challenger model URI: {challenger_uri}")
    
    print("\n✓ Model aliasing demonstration complete")


def main():
    """Run all model registry examples."""
    print("\n" + "=" * 80)
    print("  MLflow Model Registry Tutorial - Component 4")
    print("=" * 80)
    print("\nThis tutorial demonstrates the MLflow Model Registry component.")
    print("The Model Registry provides centralized model versioning,")
    print("stage management, and deployment lifecycle tracking.\n")
    
    try:
        # Register models
        run_id_v1, model_name, acc_v1 = register_first_model()
        run_id_v2, acc_v2 = register_improved_model(model_name)
        run_id_v3, acc_v3 = register_alternative_model(model_name)
        
        # Manage lifecycle
        manage_model_stages(model_name)
        
        # Add metadata
        add_model_descriptions(model_name)
        
        # Load models
        load_models_from_registry(model_name)
        
        # Compare versions
        compare_model_versions(model_name)
        
        # Demonstrate aliasing
        demonstrate_model_aliasing(model_name)
        
        print_section("Tutorial Complete!")
        print("✓ All model registry examples completed successfully!")
        print("\nKey Takeaways:")
        print("1. Model Registry centralizes model storage and versioning")
        print("2. Multiple versions of the same model can coexist")
        print("3. Stages (None/Staging/Production/Archived) manage lifecycle")
        print("4. Descriptions and tags provide important metadata")
        print("5. Models can be loaded by version, stage, or alias")
        print("6. Registry enables model comparison and governance")
        print("7. Aliases provide flexible alternative to stages")
        
        print("\n" + "-" * 80)
        print("Summary:")
        print("-" * 80)
        print(f"Registered Model: {model_name}")
        print(f"Total Versions: 3")
        print(f"  Version 1 (Archived): Accuracy = {acc_v1:.4f}")
        print(f"  Version 2 (Production): Accuracy = {acc_v2:.4f}")
        print(f"  Version 3 (None): Accuracy = {acc_v3:.4f}")
        
        print("\n" + "-" * 80)
        print("To view in MLflow UI:")
        print("-" * 80)
        print("1. Start MLflow UI: mlflow ui")
        print("2. Visit: http://localhost:5000")
        print("3. Click on 'Models' in the top menu")
        print(f"4. Select '{model_name}'")
        
        print("\nCongratulations! You've completed all 4 MLflow components!")
        print("\nTo run all tutorials in sequence, use:")
        print("  $ python run_all.py")
        
    except Exception as e:
        print(f"\n❌ Error occurred: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()