MLOps Experiment Tracking

Real-World Scenario

A team runs 50 experiments tuning a gradient boosting model — different feature sets, hyperparameters, and preprocessing. Without tracking, they lose track of what combination produced the best result. MLflow records every run: who ran it, what parameters, what metrics, which dataset version. The best model is promoted to the registry with a single API call.

Basic Experiment Tracking

import mlflow
import mlflow.sklearn
from sklearn.datasets import make_classification
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, f1_score, accuracy_score
import numpy as np

# Point to local MLflow (or remote server)
mlflow.set_tracking_uri("http://localhost:5000")  # or "mlruns" for local file
mlflow.set_experiment("churn-prediction")         # creates experiment if not exists

X, y = make_classification(n_samples=5000, n_features=20, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)


def run_experiment(
    n_estimators: int,
    max_depth: int,
    learning_rate: float,
    run_name: str = None,
) -> str:
    with mlflow.start_run(run_name=run_name) as run:
        # Log all hyperparameters
        mlflow.log_params({
            "n_estimators":  n_estimators,
            "max_depth":     max_depth,
            "learning_rate": learning_rate,
            "model_type":    "HistGradientBoosting",
        })

        # Log dataset info
        mlflow.log_params({
            "train_size": len(X_train),
            "test_size":  len(X_test),
            "n_features": X.shape[1],
            "pos_rate":   round(float(y_train.mean()), 4),
        })

        model = HistGradientBoostingClassifier(
            max_iter=n_estimators,
            max_depth=max_depth,
            learning_rate=learning_rate,
            validation_fraction=0.1,
            n_iter_no_change=10,    # early stopping
            random_state=42,
        )
        model.fit(X_train, y_train)

        # Log metrics
        y_pred  = model.predict(X_test)
        y_prob  = model.predict_proba(X_test)[:, 1]

        metrics = {
            "auc":      round(roc_auc_score(y_test, y_prob), 4),
            "f1":       round(f1_score(y_test, y_pred), 4),
            "accuracy": round(accuracy_score(y_test, y_pred), 4),
            "n_iter":   model.n_iter_,   # actual iterations (early stopping)
        }
        mlflow.log_metrics(metrics)

        # Log the trained model
        mlflow.sklearn.log_model(
            model,
            artifact_path="model",
            registered_model_name="churn_classifier",  # adds to registry
        )

        print(f"Run: {run.info.run_id[:8]}  AUC={metrics['auc']}  F1={metrics['f1']}")
        return run.info.run_id


# Run multiple experiments
run_id_1 = run_experiment(200, 5, 0.10, "baseline")
run_id_2 = run_experiment(500, 3, 0.05, "more-trees-shallow")
run_id_3 = run_experiment(200, 8, 0.20, "deep-fast")

Comparing Runs Programmatically

import mlflow
import pandas as pd

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

# Search all runs in the experiment
runs = mlflow.search_runs(
    experiment_names=["churn-prediction"],
    order_by=["metrics.auc DESC"],
    max_results=20,
)

# Focus on key columns
cols = ["run_id", "tags.mlflow.runName",
        "params.n_estimators", "params.learning_rate",
        "metrics.auc", "metrics.f1", "metrics.n_iter"]

report = runs[cols].rename(columns={
    "tags.mlflow.runName":  "name",
    "params.n_estimators":  "n_trees",
    "params.learning_rate": "lr",
    "metrics.auc":          "auc",
    "metrics.f1":           "f1",
    "metrics.n_iter":       "iters",
})
report["run_id"] = report["run_id"].str[:8]
print(report.to_string(index=False))

# Best run
best = runs.iloc[0]
print(f"\nBest run: {best['tags.mlflow.runName']} — AUC={best['metrics.auc']:.4f}")

Logging Metrics Per Epoch

import mlflow
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score
import numpy as np

mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("neural-net-training")

X, y = make_classification(n_samples=3000, n_features=20, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

with mlflow.start_run(run_name="epoch-tracking-demo"):
    mlflow.log_params({"model": "SGDClassifier", "lr": 0.01, "epochs": 50})

    from sklearn.linear_model import SGDClassifier
    model = SGDClassifier(loss="log_loss", learning_rate="constant", eta0=0.01,
                          warm_start=True, random_state=42)

    best_auc = 0
    for epoch in range(1, 51):
        model.max_iter = epoch
        model.fit(X_train, y_train)

        train_auc = roc_auc_score(y_train, model.predict_proba(X_train)[:, 1])
        val_auc   = roc_auc_score(y_test,  model.predict_proba(X_test)[:, 1])

        # step= makes metrics plottable as a time series in the MLflow UI
        mlflow.log_metric("train_auc", train_auc, step=epoch)
        mlflow.log_metric("val_auc",   val_auc,   step=epoch)

        if val_auc > best_auc:
            best_auc = val_auc
            mlflow.log_metric("best_val_auc", best_auc, step=epoch)

    mlflow.log_metric("final_val_auc", best_auc)
    print(f"Best val AUC: {best_auc:.4f}")

Model Registry and Promotion

import mlflow
from mlflow.tracking import MlflowClient

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


def get_latest_run_id(experiment_name: str, metric: str = "auc") -> str:
    runs = mlflow.search_runs(
        experiment_names=[experiment_name],
        order_by=[f"metrics.{metric} DESC"],
        max_results=1,
    )
    return runs.iloc[0]["run_id"]


def register_and_promote(run_id: str, model_name: str, min_auc: float = 0.85) -> bool:
    # Check if this run meets the quality gate
    run = client.get_run(run_id)
    auc = run.data.metrics.get("auc", 0)

    if auc < min_auc:
        print(f"Run {run_id[:8]} failed quality gate: AUC={auc:.4f} < {min_auc}")
        return False

    # Register (creates a new version)
    model_uri = f"runs:/{run_id}/model"
    mv = mlflow.register_model(model_uri, model_name)
    print(f"Registered {model_name} v{mv.version}  AUC={auc:.4f}")

    # Transition to staging
    client.transition_model_version_stage(
        name=model_name, version=mv.version, stage="Staging"
    )
    print(f"v{mv.version} → Staging")

    # Add tags for traceability
    client.set_model_version_tag(model_name, mv.version, "auc", str(auc))
    client.set_model_version_tag(model_name, mv.version, "run_id", run_id[:8])
    return True


def list_registry(model_name: str):
    print(f"\n{'Version':<10} {'Stage':<15} {'AUC':<8} {'Run'}")
    print("-" * 45)
    for mv in client.search_model_versions(f"name='{model_name}'"):
        tags = mv.tags
        print(f"{mv.version:<10} {mv.current_stage:<15} "
              f"{tags.get('auc', 'N/A'):<8} {tags.get('run_id', 'N/A')}")


# Get best run and register it
best_run = get_latest_run_id("churn-prediction")
success  = register_and_promote(best_run, MODEL_NAME)
list_registry(MODEL_NAME)


# Load the production model anywhere
def load_production_model(model_name: str):
    model_uri = f"models:/{model_name}/Production"
    return mlflow.sklearn.load_model(model_uri)

Reproducible Training Script

#!/usr/bin/env python
# train.py — designed to be called from CLI or CI/CD
import argparse
import mlflow
import mlflow.sklearn
from sklearn.datasets import make_classification
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score
import os

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--n-estimators",  type=int,   default=200)
    parser.add_argument("--max-depth",     type=int,   default=5)
    parser.add_argument("--learning-rate", type=float, default=0.1)
    parser.add_argument("--experiment",    type=str,   default="churn-prediction")
    args = parser.parse_args()

    mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI", "mlruns"))
    mlflow.set_experiment(args.experiment)

    X, y = make_classification(n_samples=5000, n_features=20, random_state=42)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    with mlflow.start_run():
        mlflow.log_params(vars(args))
        mlflow.set_tag("git_commit", os.popen("git rev-parse --short HEAD").read().strip())

        model = HistGradientBoostingClassifier(
            max_iter=args.n_estimators,
            max_depth=args.max_depth,
            learning_rate=args.learning_rate,
            random_state=42,
        )
        model.fit(X_train, y_train)

        auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
        mlflow.log_metric("auc", auc)
        mlflow.sklearn.log_model(model, "model")
        print(f"AUC: {auc:.4f}")

if __name__ == "__main__":
    main()

# Usage:
# python train.py --n-estimators 300 --learning-rate 0.05
# MLFLOW_TRACKING_URI=http://mlflow.internal python train.py ...