04 - Experiment Tracking with MLflow: Setup, Logging, Comparison

# Install MLflow (2.x/3.x)
pip install mlflow

# With extras for specific integrations
pip install mlflow[extras]      # scikit-learn, XGBoost, LightGBM
pip install mlflow[databricks]  # Databricks integration
pip install mlflow[genai]       # GenAI and LLM tools (MLflow 3+)

# Verify installation
mlflow --version
# mlflow, version 2.x.x or 3.x.x

# Launch the local UI (uses ./mlruns as storage)
mlflow ui
# UI available at http://localhost:5000

# Start with SQLite database
mlflow server \
  --backend-store-uri sqlite:///mlflow.db \
  --default-artifact-root ./mlruns \
  --host 0.0.0.0 \
  --port 5000

# ==================== docker-compose.yml ====================
version: '3.8'

services:
  postgres:
    image: postgres:15-alpine
    environment:
      POSTGRES_DB: mlflow
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
    volumes:
      - postgres_data:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U mlflow"]
      interval: 10s
      timeout: 5s
      retries: 5

  minio:
    image: minio/minio:latest
    command: server /data --console-address ":9001"
    environment:
      MINIO_ROOT_USER: ${MINIO_ACCESS_KEY}
      MINIO_ROOT_PASSWORD: ${MINIO_SECRET_KEY}
    volumes:
      - minio_data:/data
    ports:
      - "9000:9000"
      - "9001:9001"

  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.19.0
    depends_on:
      postgres:
        condition: service_healthy
    environment:
      MLFLOW_BACKEND_STORE_URI: postgresql://mlflow:${POSTGRES_PASSWORD}@postgres:5432/mlflow
      MLFLOW_ARTIFACT_ROOT: s3://mlflow-artifacts/
      AWS_ACCESS_KEY_ID: ${MINIO_ACCESS_KEY}
      AWS_SECRET_ACCESS_KEY: ${MINIO_SECRET_KEY}
      MLFLOW_S3_ENDPOINT_URL: http://minio:9000
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:${POSTGRES_PASSWORD}@postgres:5432/mlflow
      --default-artifact-root s3://mlflow-artifacts/
      --host 0.0.0.0
      --port 5000
    ports:
      - "5000:5000"

volumes:
  postgres_data:
  minio_data:

import mlflow
import mlflow.sklearn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import (
    accuracy_score, f1_score, roc_auc_score,
    confusion_matrix, classification_report, RocCurveDisplay
)
from sklearn.model_selection import train_test_split
import os


# ==================== MLflow Configuration ====================

# Connect to the tracking server (local or remote)
mlflow.set_tracking_uri("http://localhost:5000")

# Create or use an existing experiment
experiment_name = "churn-prediction-gbm"
mlflow.set_experiment(experiment_name)

experiment = mlflow.get_experiment_by_name(experiment_name)
print(f"Experiment ID: {experiment.experiment_id}")


# ==================== Training with Full Tracking ====================

def train_churn_model(X_train, X_val, y_train, y_val, params: dict) -> str:
    """
    Train a GBM for churn prediction with complete MLflow tracking.
    Returns the MLflow run_id.
    """

    with mlflow.start_run(
        run_name=f"gbm-lr{params['learning_rate']}-depth{params['max_depth']}"
    ) as run:

        # ---- 1. TAGS: run metadata ----
        mlflow.set_tags({
            "team": "ml-engineering",
            "project": "churn-prediction",
            "dataset_version": "v2.1",
            "git_commit": os.popen("git rev-parse HEAD").read().strip(),
            "environment": "dev",
        })

        # ---- 2. PARAMS: hyperparameters and config ----
        mlflow.log_params(params)
        mlflow.log_params({
            "train_size": len(X_train),
            "val_size": len(X_val),
            "n_features": X_train.shape[1],
            "target_positive_rate": float(y_train.mean()),
        })

        # ---- 3. TRAINING ----
        model = GradientBoostingClassifier(**params)
        model.fit(X_train, y_train)

        # ---- 4. STEP-BY-STEP METRICS during training ----
        train_scores = list(model.staged_predict(X_train))
        val_scores = list(model.staged_predict(X_val))

        for step, (tr_pred, val_pred) in enumerate(zip(train_scores, val_scores)):
            train_acc = accuracy_score(y_train, tr_pred)
            val_acc = accuracy_score(y_val, val_pred)
            mlflow.log_metrics({
                "train_accuracy_step": train_acc,
                "val_accuracy_step": val_acc,
            }, step=step)

        # ---- 5. FINAL METRICS ----
        y_pred = model.predict(X_val)
        y_prob = model.predict_proba(X_val)[:, 1]

        final_metrics = {
            "accuracy": accuracy_score(y_val, y_pred),
            "f1_score": f1_score(y_val, y_pred),
            "auc_roc": roc_auc_score(y_val, y_prob),
        }
        mlflow.log_metrics(final_metrics)

        # ---- 6. ARTIFACTS: model file and reports ----

        # Confusion matrix as an image artifact
        fig, ax = plt.subplots(figsize=(6, 5))
        cm = confusion_matrix(y_val, y_pred)
        im = ax.imshow(cm, interpolation='nearest', cmap='Blues')
        ax.set_title('Confusion Matrix - Churn Prediction')
        ax.set_xlabel('Predicted')
        ax.set_ylabel('True')
        plt.colorbar(im)
        plt.tight_layout()
        mlflow.log_figure(fig, "confusion_matrix.png")
        plt.close()

        # Classification report as a text artifact
        report = classification_report(y_val, y_pred, target_names=["No Churn", "Churn"])
        mlflow.log_text(report, "classification_report.txt")

        # Feature importance as a JSON table
        feature_imp = pd.DataFrame({
            "feature": X_train.columns.tolist(),
            "importance": model.feature_importances_
        }).sort_values("importance", ascending=False)
        mlflow.log_table(feature_imp.to_dict(orient="list"), "feature_importance.json")

        # ---- 7. LOG MODEL: save with signature and input example ----
        input_example = X_val.head(3)
        signature = mlflow.models.infer_signature(X_val, y_pred)

        mlflow.sklearn.log_model(
            sk_model=model,
            artifact_path="model",
            signature=signature,
            input_example=input_example,
            registered_model_name="churn-gbm-model",  # Auto-register in Registry
        )

        print(f"Run ID: {run.info.run_id}")
        print(f"Accuracy: {final_metrics['accuracy']:.4f}")
        print(f"AUC-ROC: {final_metrics['auc_roc']:.4f}")

        return run.info.run_id

import mlflow
import optuna
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import cross_val_score


def hyperparameter_search(X_train, y_train, n_trials: int = 50) -> str:
    """
    Hyperparameter search with Optuna + MLflow nested runs.
    Parent run: contains the search summary.
    Child runs: each Optuna trial is a child MLflow run.
    """

    mlflow.set_tracking_uri("http://localhost:5000")
    mlflow.set_experiment("churn-hparam-search")

    with mlflow.start_run(run_name="optuna-search-v1") as parent_run:
        mlflow.set_tag("search_method", "optuna-tpe")
        mlflow.log_param("n_trials", n_trials)
        mlflow.log_param("optimization_metric", "auc_roc")

        def objective(trial) -> float:
            params = {
                "n_estimators": trial.suggest_int("n_estimators", 100, 1000),
                "learning_rate": trial.suggest_float("learning_rate", 0.001, 0.3, log=True),
                "max_depth": trial.suggest_int("max_depth", 2, 8),
                "subsample": trial.suggest_float("subsample", 0.5, 1.0),
            }

            # Each trial gets its own child MLflow run
            with mlflow.start_run(
                run_name=f"trial-{trial.number}",
                nested=True  # marks this as a child run
            ):
                mlflow.log_params(params)

                model = GradientBoostingClassifier(**params, random_state=42)
                scores = cross_val_score(model, X_train, y_train, cv=3, scoring="roc_auc")
                auc_mean = scores.mean()
                auc_std = scores.std()

                mlflow.log_metrics({
                    "cv_auc_mean": auc_mean,
                    "cv_auc_std": auc_std,
                })

            return auc_mean

        study = optuna.create_study(
            direction="maximize",
            sampler=optuna.samplers.TPESampler(seed=42)
        )
        study.optimize(objective, n_trials=n_trials)

        best_trial = study.best_trial
        mlflow.log_params({f"best_{k}": v for k, v in best_trial.params.items()})
        mlflow.log_metric("best_auc", best_trial.value)

        print(f"Best AUC: {best_trial.value:.4f}")
        print(f"Best params: {best_trial.params}")

        return parent_run.info.run_id

import mlflow
import mlflow.sklearn
import mlflow.xgboost

# ==================== scikit-learn Autologging ====================

mlflow.sklearn.autolog(
    log_input_examples=True,
    log_model_signatures=True,
    log_models=True,
    log_datasets=False,
    max_tuning_runs=100,
    exclusive=False,
)

mlflow.set_experiment("autolog-demo")

from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV

with mlflow.start_run(run_name="rf-gridsearch-autolog"):
    param_grid = {
        "n_estimators": [100, 300],
        "max_depth": [4, 6, 8],
        "min_samples_split": [10, 20],
    }

    model = GridSearchCV(
        RandomForestClassifier(random_state=42),
        param_grid,
        cv=3,
        scoring="roc_auc",
        n_jobs=-1
    )
    model.fit(X_train, y_train)
    # MLflow automatically logged:
    # - All RandomForest parameters
    # - cv=3, scoring, n_jobs
    # - Best params from GridSearchCV
    # - Cross-validation score
    # - The model as an artifact


# ==================== XGBoost Autologging ====================

mlflow.xgboost.autolog(
    importance_types=["gain", "weight"],
    log_model_signatures=True,
)

import xgboost as xgb

with mlflow.start_run(run_name="xgb-autolog"):
    dtrain = xgb.DMatrix(X_train, label=y_train)
    dval = xgb.DMatrix(X_val, label=y_val)

    params = {
        "objective": "binary:logistic",
        "eval_metric": ["logloss", "auc"],
        "learning_rate": 0.05,
        "max_depth": 6,
    }

    # Autolog tracks all metrics per boosting round
    booster = xgb.train(
        params,
        dtrain,
        num_boost_round=500,
        evals=[(dtrain, "train"), (dval, "val")],
        early_stopping_rounds=50,
        verbose_eval=False,
    )

import mlflow
from mlflow.tracking import MlflowClient


mlflow.set_tracking_uri("http://localhost:5000")
client = MlflowClient()

MODEL_NAME = "churn-gbm-model"


# ==================== 1. REGISTER A MODEL ====================

# Method A: during log_model (most common)
with mlflow.start_run() as run:
    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        registered_model_name=MODEL_NAME,
    )
    # Automatically registered as version 1 with stage "None"


# Method B: from an existing run via URI
run_id = "abc123def456"
version = mlflow.register_model(
    model_uri=f"runs:/{run_id}/model",
    name=MODEL_NAME,
    tags={"team": "ml-eng", "algorithm": "gbm"}
)
print(f"Registered: {MODEL_NAME} version {version.version}")


# ==================== 2. MANAGE VERSIONS AND STAGES ====================

# Add a description to the version
client.update_model_version(
    name=MODEL_NAME,
    version=version.version,
    description=(
        "GBM for churn prediction v2.1. "
        "Accuracy: 0.9423, AUC-ROC: 0.9567. "
        "Trained on 2024-01 to 2025-01 dataset, 45k samples."
    )
)

# Promote to Staging (after internal validation)
client.transition_model_version_stage(
    name=MODEL_NAME,
    version=version.version,
    stage="Staging",
    archive_existing_versions=False,
)
print(f"Model v{version.version} promoted to Staging")

# Add traceability tags
client.set_model_version_tag(
    name=MODEL_NAME,
    version=version.version,
    key="validated_by",
    value="alice.rossi@company.com"
)

# Promote to Production (after team approval)
client.transition_model_version_stage(
    name=MODEL_NAME,
    version=version.version,
    stage="Production",
    archive_existing_versions=True,  # Archive the previous Production version
)
print(f"Model v{version.version} is now in production!")


# ==================== 3. LOAD THE PRODUCTION MODEL ====================

def load_production_model(model_name: str):
    """Always loads the Production version from the registry."""
    model_uri = f"models:/{model_name}/Production"
    return mlflow.sklearn.load_model(model_uri)


# In an inference script or serving layer
model = load_production_model(MODEL_NAME)
predictions = model.predict(new_data)


# ==================== 4. QUERY THE REGISTRY ====================

# List all versions of a model
versions = client.search_model_versions(f"name='{MODEL_NAME}'")
for v in versions:
    print(f"v{v.version} | Stage: {v.current_stage} | Run: {v.run_id[:8]}...")

# Get only Production versions
prod_versions = client.get_latest_versions(MODEL_NAME, stages=["Production"])
if prod_versions:
    latest_prod = prod_versions[0]
    run_data = client.get_run(latest_prod.run_id).data
    print(f"Production AUC-ROC: {run_data.metrics.get('auc_roc', 'N/A')}")

# ==================== Serving via CLI ====================

# Serve the latest Production version of the model
mlflow models serve \
  --model-uri "models:/churn-gbm-model/Production" \
  --host 0.0.0.0 \
  --port 8080 \
  --env-manager conda

# Serve a specific run
mlflow models serve \
  --model-uri "runs:/abc123def456/model" \
  --port 8080

# With Docker (recommended for production)
mlflow models build-docker \
  --model-uri "models:/churn-gbm-model/Production" \
  --name "churn-model-server" \
  --enable-mlserver

docker run -p 8080:8080 churn-model-server

# ==================== Test the Server ====================

import requests
import json
import pandas as pd

test_data = pd.DataFrame({
    "feature_0": [0.5, -1.2],
    "feature_1": [1.3, 0.8],
})

payload = {
    "dataframe_split": {
        "columns": test_data.columns.tolist(),
        "data": test_data.values.tolist()
    }
}

response = requests.post(
    "http://localhost:8080/invocations",
    headers={"Content-Type": "application/json"},
    data=json.dumps(payload)
)
print(f"Predictions: {response.json()}")
# {"predictions": [0, 1]}

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import mlflow
import pandas as pd
import numpy as np
import threading
import time
import logging
from typing import List

logger = logging.getLogger(__name__)
app = FastAPI(title="Churn Prediction API", version="2.0.0")


class MLflowModelManager:
    """
    Manages loading and auto-refresh of the model
    from the MLflow Model Registry.
    """

    def __init__(self, model_name: str, tracking_uri: str, refresh_interval: int = 300):
        self.model_name = model_name
        self._model = None
        self._model_version = None
        self._lock = threading.Lock()

        mlflow.set_tracking_uri(tracking_uri)
        self._load_model()
        self._start_refresh_thread()

    def _load_model(self) -> None:
        try:
            model_uri = f"models:/{self.model_name}/Production"
            new_model = mlflow.sklearn.load_model(model_uri)

            client = mlflow.MlflowClient()
            versions = client.get_latest_versions(self.model_name, stages=["Production"])
            version = versions[0].version if versions else "unknown"

            with self._lock:
                self._model = new_model
                self._model_version = version
            logger.info(f"Model {self.model_name} v{version} loaded from registry")
        except Exception as e:
            logger.error(f"Failed to load model: {e}")

    def _start_refresh_thread(self) -> None:
        def refresh_loop():
            while True:
                time.sleep(300)  # Check every 5 minutes
                self._load_model()
        thread = threading.Thread(target=refresh_loop, daemon=True)
        thread.start()

    def predict(self, features: pd.DataFrame) -> np.ndarray:
        with self._lock:
            if self._model is None:
                raise RuntimeError("Model not available")
            return self._model.predict(features)

    def predict_proba(self, features: pd.DataFrame) -> np.ndarray:
        with self._lock:
            return self._model.predict_proba(features)[:, 1]

    @property
    def model_version(self) -> str:
        return self._model_version or "unknown"


model_manager = MLflowModelManager(
    model_name="churn-gbm-model",
    tracking_uri="http://mlflow-server:5000",
)


class PredictionRequest(BaseModel):
    features: List[List[float]]
    feature_names: List[str]

class PredictionResponse(BaseModel):
    predictions: List[int]
    probabilities: List[float]
    model_version: str


@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest) -> PredictionResponse:
    try:
        df = pd.DataFrame(request.features, columns=request.feature_names)
        predictions = model_manager.predict(df).tolist()
        probabilities = model_manager.predict_proba(df).tolist()
        return PredictionResponse(
            predictions=predictions,
            probabilities=probabilities,
            model_version=model_manager.model_version,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health_check():
    return {"status": "healthy", "model_version": model_manager.model_version}

import mlflow
import pandas as pd

mlflow.set_tracking_uri("http://localhost:5000")

# ==================== Search Runs with Filters ====================

# Find all runs with AUC-ROC > 0.92 in your experiment
runs = mlflow.search_runs(
    experiment_names=["churn-prediction-gbm"],
    filter_string="metrics.auc_roc > 0.92 and tags.environment = 'dev'",
    order_by=["metrics.auc_roc DESC"],
    max_results=20,
)

print(runs[["run_id", "metrics.auc_roc", "metrics.f1_score",
            "params.learning_rate", "params.max_depth"]].head())


# ==================== Find and Load the Best Run ====================

def get_best_run(experiment_name: str, metric: str = "metrics.auc_roc") -> dict:
    runs = mlflow.search_runs(
        experiment_names=[experiment_name],
        filter_string="status = 'FINISHED'",
        order_by=[f"{metric} DESC"],
        max_results=1,
    )
    if runs.empty:
        raise ValueError(f"No runs found in {experiment_name}")

    best = runs.iloc[0]
    return {
        "run_id": best["run_id"],
        "auc_roc": best.get("metrics.auc_roc"),
        "accuracy": best.get("metrics.accuracy"),
    }


best_run = get_best_run("churn-prediction-gbm")
print(f"Best run: {best_run['run_id']}, AUC-ROC: {best_run['auc_roc']:.4f}")

# Load the model from the best run
model = mlflow.sklearn.load_model(f"runs:/{best_run['run_id']}/model")


# ==================== Export Step-by-Step Metric History ====================

from mlflow.tracking import MlflowClient

client = MlflowClient()
run_id = "abc123def456"
metric_history = client.get_metric_history(run_id, "val_accuracy_step")
steps = [m.step for m in metric_history]
values = [m.value for m in metric_history]

accuracy_over_time = pd.DataFrame({"step": steps, "val_accuracy": values})
print(f"Training steps logged: {len(accuracy_over_time)}")

# Recommended naming schema for experiments and runs

# EXPERIMENT: [team]-[project]-[type]
"ml-eng-churn-prediction-gbm"
"ml-eng-churn-prediction-neural-net"
"research-recommender-collaborative"

# RUN NAME: [algorithm]-[key-param]-[date]
"gbm-lr0.05-depth6-2025-11-15"
"xgb-v2-autofeat-2025-11-20"

# Always use tags for structured metadata
mlflow.set_tags({
    "team":            "ml-engineering",
    "project":         "churn-prediction",
    "environment":     "dev",           # dev | staging | prod
    "dataset_version": "v2.1",
    "git_branch":      git_branch,
    "git_commit":      git_commit[:8],
    "triggered_by":    "ci-cd",         # manual | ci-cd | scheduled
    "approved_by":     "",              # Filled before promotion
})

# .github/workflows/ml-experiment.yml
name: ML Experiment + Model Promotion

on:
  push:
    branches: [main]
    paths:
      - 'src/models/**'
      - 'params.yaml'

jobs:
  train-and-evaluate:
    runs-on: ubuntu-latest
    env:
      MLFLOW_TRACKING_URI: ${{ secrets.MLFLOW_TRACKING_URI }}
      AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
      AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}

    steps:
      - uses: actions/checkout@v4

      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'

      - name: Install dependencies
        run: pip install -r requirements.txt

      - name: Run training
        id: training
        run: |
          RUN_ID=$(python src/models/train.py --output-run-id)
          echo "run_id=$RUN_ID" >> $GITHUB_OUTPUT

      - name: Evaluate metrics
        run: |
          python scripts/ci_evaluate.py \
            --run-id ${{ steps.training.outputs.run_id }} \
            --min-auc 0.92 \
            --min-accuracy 0.90

      - name: Promote to Staging
        if: success()
        run: |
          python scripts/promote_model.py \
            --run-id ${{ steps.training.outputs.run_id }} \
            --model-name churn-gbm-model \
            --target-stage Staging