1_tracking.py

"""
MLflow Tutorial - Component 1: MLflow Tracking

This module demonstrates the MLflow Tracking component, which is used to log
parameters, metrics, and artifacts during ML experiments.

Key Features Demonstrated:
1. Creating and managing experiments
2. Logging parameters and metrics
3. Logging artifacts (plots, models, files)
4. Nested runs for hyperparameter tuning
5. Tags and run organization
"""

import mlflow
import mlflow.sklearn
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-GUI backend for matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.metrics import confusion_matrix, classification_report
import os
import json


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


def create_sample_data():
    """Create a sample classification dataset."""
    print("Creating sample classification dataset...")
    X, y = make_classification(
        n_samples=1000,
        n_features=20,
        n_informative=15,
        n_redundant=5,
        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
    )
    print(f"Train set: {X_train.shape}, Test set: {X_test.shape}")
    return X_train, X_test, y_train, y_test


def plot_confusion_matrix(y_true, y_pred, save_path="confusion_matrix.png"):
    """Create and save a confusion matrix plot."""
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.title('Confusion Matrix')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()
    return save_path


def plot_feature_importance(model, save_path="feature_importance.png"):
    """Create and save a feature importance plot."""
    importances = model.feature_importances_
    indices = np.argsort(importances)[::-1][:10]  # Top 10 features
    
    plt.figure(figsize=(10, 6))
    plt.bar(range(len(indices)), importances[indices])
    plt.title('Top 10 Feature Importances')
    plt.xlabel('Feature Index')
    plt.ylabel('Importance')
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()
    return save_path


def basic_tracking_example():
    """Demonstrate basic MLflow tracking."""
    print_section("1. Basic Tracking Example")
    
    # Create experiment
    experiment_name = "Basic_Tracking_Demo"
    mlflow.set_experiment(experiment_name)
    print(f"Created/Set experiment: {experiment_name}")
    
    # Prepare data
    X_train, X_test, y_train, y_test = create_sample_data()
    
    # Start MLflow run
    with mlflow.start_run(run_name="basic_rf_model") as run:
        print(f"\nStarted run: {run.info.run_id}")
        
        # Log parameters
        params = {
            "n_estimators": 100,
            "max_depth": 5,
            "min_samples_split": 2,
            "random_state": 42
        }
        mlflow.log_params(params)
        print("Logged parameters:", params)
        
        # Train model
        print("\nTraining model...")
        model = RandomForestClassifier(**params)
        model.fit(X_train, y_train)
        
        # Make predictions
        y_pred = model.predict(X_test)
        
        # Log metrics
        metrics = {
            "accuracy": accuracy_score(y_test, y_pred),
            "precision": precision_score(y_test, y_pred),
            "recall": recall_score(y_test, y_pred),
            "f1_score": f1_score(y_test, y_pred)
        }
        mlflow.log_metrics(metrics)
        print("Logged metrics:", metrics)
        
        # Log tags
        mlflow.set_tag("model_type", "RandomForest")
        mlflow.set_tag("dataset", "synthetic_classification")
        mlflow.set_tag("tutorial", "component_1")
        print("Logged tags")
        
        # Create and log artifacts
        os.makedirs("temp_artifacts", exist_ok=True)
        
        # Confusion matrix
        cm_path = plot_confusion_matrix(y_test, y_pred, "temp_artifacts/confusion_matrix.png")
        mlflow.log_artifact(cm_path)
        print(f"Logged artifact: {cm_path}")
        
        # Feature importance
        fi_path = plot_feature_importance(model, "temp_artifacts/feature_importance.png")
        mlflow.log_artifact(fi_path)
        print(f"Logged artifact: {fi_path}")
        
        # Classification report
        report = classification_report(y_test, y_pred)
        report_path = "temp_artifacts/classification_report.txt"
        with open(report_path, 'w') as f:
            f.write(report)
        mlflow.log_artifact(report_path)
        print(f"Logged artifact: {report_path}")
        
        # Log model with signature and input example
        from mlflow.models.signature import infer_signature
        signature = infer_signature(X_test, model.predict(X_test))
        mlflow.sklearn.log_model(
            model, 
            "model",
            signature=signature,
            input_example=X_test[:5]
        )
        print("Logged model with signature and input example")
        
        print(f"\n✓ Basic tracking completed. Run ID: {run.info.run_id}")
    
    # Cleanup
    import shutil
    if os.path.exists("temp_artifacts"):
        shutil.rmtree("temp_artifacts")


def hyperparameter_tuning_example():
    """Demonstrate nested runs for hyperparameter tuning."""
    print_section("2. Hyperparameter Tuning with Nested Runs")
    
    # Create experiment
    experiment_name = "Hyperparameter_Tuning_Demo"
    mlflow.set_experiment(experiment_name)
    
    # Prepare data
    X_train, X_test, y_train, y_test = create_sample_data()
    
    # Define hyperparameter grid
    param_grid = {
        "n_estimators": [50, 100, 200],
        "max_depth": [3, 5, 7],
        "min_samples_split": [2, 5, 10]
    }
    
    # Parent run for the tuning process
    with mlflow.start_run(run_name="hyperparameter_tuning_parent") as parent_run:
        print(f"Started parent run: {parent_run.info.run_id}\n")
        
        best_score = 0
        best_params = None
        
        # Try different hyperparameter combinations
        run_count = 0
        for n_est in param_grid["n_estimators"]:
            for max_d in param_grid["max_depth"]:
                for min_split in param_grid["min_samples_split"]:
                    run_count += 1
                    
                    # Child run for each combination
                    with mlflow.start_run(run_name=f"run_{run_count}", nested=True):
                        params = {
                            "n_estimators": n_est,
                            "max_depth": max_d,
                            "min_samples_split": min_split,
                            "random_state": 42
                        }
                        
                        # Train model
                        model = RandomForestClassifier(**params)
                        model.fit(X_train, y_train)
                        y_pred = model.predict(X_test)
                        
                        # Calculate metrics
                        accuracy = accuracy_score(y_test, y_pred)
                        f1 = f1_score(y_test, y_pred)
                        
                        # Log parameters and metrics
                        mlflow.log_params(params)
                        mlflow.log_metrics({
                            "accuracy": accuracy,
                            "f1_score": f1
                        })
                        
                        print(f"Run {run_count}: n_est={n_est}, max_d={max_d}, "
                              f"min_split={min_split} -> Accuracy: {accuracy:.4f}")
                        
                        # Track best model
                        if accuracy > best_score:
                            best_score = accuracy
                            best_params = params
        
        # Log best parameters to parent run
        mlflow.log_params({f"best_{k}": v for k, v in best_params.items()})
        mlflow.log_metric("best_accuracy", best_score)
        mlflow.set_tag("total_runs", run_count)
        
        print(f"\n✓ Hyperparameter tuning completed.")
        print(f"Best parameters: {best_params}")
        print(f"Best accuracy: {best_score:.4f}")


def batch_logging_example():
    """Demonstrate batch logging of metrics over epochs/iterations."""
    print_section("3. Batch Logging Example (Simulated Training)")
    
    experiment_name = "Batch_Logging_Demo"
    mlflow.set_experiment(experiment_name)
    
    with mlflow.start_run(run_name="iterative_training") as run:
        print(f"Started run: {run.info.run_id}\n")
        
        # Simulate training over epochs
        n_epochs = 20
        print(f"Simulating {n_epochs} epochs of training...")
        
        for epoch in range(n_epochs):
            # Simulate metrics that improve over time
            train_loss = 1.0 - (epoch / n_epochs) * 0.8 + np.random.random() * 0.1
            val_loss = 1.0 - (epoch / n_epochs) * 0.7 + np.random.random() * 0.15
            train_acc = (epoch / n_epochs) * 0.9 + np.random.random() * 0.05
            val_acc = (epoch / n_epochs) * 0.85 + np.random.random() * 0.05
            
            # Log metrics for this epoch
            mlflow.log_metrics({
                "train_loss": train_loss,
                "val_loss": val_loss,
                "train_accuracy": train_acc,
                "val_accuracy": val_acc
            }, step=epoch)
            
            if epoch % 5 == 0:
                print(f"Epoch {epoch}: train_loss={train_loss:.4f}, "
                      f"val_loss={val_loss:.4f}, train_acc={train_acc:.4f}, "
                      f"val_acc={val_acc:.4f}")
        
        # Log final parameters
        mlflow.log_params({
            "epochs": n_epochs,
            "learning_rate": 0.001,
            "batch_size": 32
        })
        
        print(f"\n✓ Batch logging completed over {n_epochs} epochs")


def custom_artifacts_example():
    """Demonstrate logging various types of artifacts."""
    print_section("4. Custom Artifacts Example")
    
    experiment_name = "Custom_Artifacts_Demo"
    mlflow.set_experiment(experiment_name)
    
    with mlflow.start_run(run_name="custom_artifacts") as run:
        print(f"Started run: {run.info.run_id}\n")
        
        os.makedirs("temp_artifacts", exist_ok=True)
        
        # 1. Log a dictionary as JSON
        config = {
            "model_architecture": "RandomForest",
            "preprocessing": {
                "normalization": "StandardScaler",
                "feature_selection": "SelectKBest"
            },
            "deployment": {
                "platform": "AWS",
                "instance_type": "ml.m5.large"
            }
        }
        config_path = "temp_artifacts/config.json"
        with open(config_path, 'w') as f:
            json.dump(config, f, indent=2)
        mlflow.log_artifact(config_path)
        print(f"✓ Logged JSON config")
        
        # 2. Log a text file with experiment notes
        notes = """
        Experiment Notes:
        - Dataset: Synthetic classification data
        - Purpose: Demonstrate MLflow tracking capabilities
        - Key findings: Model performs well with default parameters
        - Next steps: Try deep learning approach
        """
        notes_path = "temp_artifacts/experiment_notes.txt"
        with open(notes_path, 'w') as f:
            f.write(notes)
        mlflow.log_artifact(notes_path)
        print(f"✓ Logged experiment notes")
        
        # 3. Log multiple plots
        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        
        # Plot 1: Random data distribution
        axes[0, 0].hist(np.random.randn(1000), bins=30)
        axes[0, 0].set_title('Data Distribution')
        
        # Plot 2: Training curve
        epochs = range(50)
        loss = [1.0 / (1 + x/10) for x in epochs]
        axes[0, 1].plot(epochs, loss)
        axes[0, 1].set_title('Training Loss')
        axes[0, 1].set_xlabel('Epoch')
        axes[0, 1].set_ylabel('Loss')
        
        # Plot 3: Scatter plot
        x = np.random.randn(100)
        y = 2 * x + np.random.randn(100)
        axes[1, 0].scatter(x, y, alpha=0.6)
        axes[1, 0].set_title('Feature Correlation')
        
        # Plot 4: Bar chart
        categories = ['A', 'B', 'C', 'D', 'E']
        values = [23, 45, 56, 78, 32]
        axes[1, 1].bar(categories, values)
        axes[1, 1].set_title('Category Performance')
        
        plt.tight_layout()
        plots_path = "temp_artifacts/analysis_plots.png"
        plt.savefig(plots_path)
        plt.close()
        mlflow.log_artifact(plots_path)
        print(f"✓ Logged analysis plots")
        
        # 4. Log a directory of artifacts
        data_dir = "temp_artifacts/data_samples"
        os.makedirs(data_dir, exist_ok=True)
        for i in range(3):
            with open(f"{data_dir}/sample_{i}.txt", 'w') as f:
                f.write(f"Sample data file {i}\n")
        mlflow.log_artifacts(data_dir, artifact_path="data_samples")
        print(f"✓ Logged directory of data samples")
        
        print(f"\n✓ Custom artifacts logged successfully")
    
    # Cleanup
    import shutil
    if os.path.exists("temp_artifacts"):
        shutil.rmtree("temp_artifacts")


def main():
    """Run all tracking examples."""
    print("\n" + "=" * 80)
    print("  MLflow Tracking Tutorial - Component 1")
    print("=" * 80)
    print("\nThis tutorial demonstrates the MLflow Tracking component.")
    print("Make sure to start the MLflow UI in a separate terminal:")
    print("  $ mlflow ui")
    print("Then visit http://localhost:5000 to view the experiments.\n")
    
    try:
        # Run all examples
        basic_tracking_example()
        hyperparameter_tuning_example()
        batch_logging_example()
        custom_artifacts_example()
        
        print_section("Tutorial Complete!")
        print("✓ All tracking examples completed successfully!")
        print("\nKey Takeaways:")
        print("1. MLflow Tracking logs parameters, metrics, and artifacts")
        print("2. Experiments organize related runs")
        print("3. Nested runs are useful for hyperparameter tuning")
        print("4. Metrics can be logged at different steps/epochs")
        print("5. Artifacts can be files, plots, models, or directories")
        print("\nNext: Run Component 2 (MLflow Projects) with:")
        print("  $ python 2_projects.py")
        
    except Exception as e:
        print(f"\n❌ Error occurred: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()