ML Experiment Tracking with MLflow

Training a model once is easy. Training dozens of variants with different hyperparameters, datasets, and preprocessing steps — and remembering which combination produced the best results — is where teams lose track. MLflow solves this by providing a structured system for logging experiments, comparing results, and managing model versions.

This guide covers practical usage from first log call to production model registry.

Setting Up MLflow

Install MLflow and start the tracking server:

pip install mlflow

# Start the UI (stores data locally in ./mlruns)
mlflow ui --port 5000

For team use, point MLflow at a shared backend:

mlflow server \
    --backend-store-uri postgresql://user:pass@host/mlflow \
    --default-artifact-root s3://my-bucket/mlflow-artifacts \
    --port 5000

The backend store holds run metadata (parameters, metrics). The artifact store holds files (models, plots, data samples).

Logging Your First Experiment

import mlflow
import mlflow.sklearn
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score

# Set the tracking server URI
mlflow.set_tracking_uri("http://localhost:5000")

# Create or get an experiment
mlflow.set_experiment("iris-classification")

# Load data
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Train and log
with mlflow.start_run(run_name="random-forest-baseline"):
    # Log parameters
    params = {"n_estimators": 100, "max_depth": 5, "random_state": 42}
    mlflow.log_params(params)

    # Train
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)

    # Evaluate and log metrics
    y_pred = model.predict(X_test)
    mlflow.log_metric("accuracy", accuracy_score(y_test, y_pred))
    mlflow.log_metric("f1_macro", f1_score(y_test, y_pred, average="macro"))

    # Log the model
    mlflow.sklearn.log_model(model, "model")

    print(f"Run ID: {mlflow.active_run().info.run_id}")

Every start_run block creates a new run under the experiment. Parameters, metrics, and the model artifact are all linked to this run and visible in the MLflow UI.

Hyperparameter Sweep with Logging

Run multiple configurations and log each one. MLflow makes it trivial to compare them afterward.

from itertools import product

param_grid = {
    "n_estimators": [50, 100, 200],
    "max_depth": [3, 5, 10, None],
    "min_samples_split": [2, 5],
}

def sweep_params(X_train, X_test, y_train, y_test):
    mlflow.set_experiment("iris-hyperparameter-sweep")

    best_accuracy = 0
    best_run_id = None

    for n_est, depth, min_split in product(
        param_grid["n_estimators"],
        param_grid["max_depth"],
        param_grid["min_samples_split"],
    ):
        with mlflow.start_run():
            params = {
                "n_estimators": n_est,
                "max_depth": depth if depth else "None",
                "min_samples_split": min_split,
            }
            mlflow.log_params(params)

            model = RandomForestClassifier(
                n_estimators=n_est,
                max_depth=depth,
                min_samples_split=min_split,
                random_state=42,
            )
            model.fit(X_train, y_train)
            y_pred = model.predict(X_test)

            accuracy = accuracy_score(y_test, y_pred)
            f1 = f1_score(y_test, y_pred, average="macro")

            mlflow.log_metric("accuracy", accuracy)
            mlflow.log_metric("f1_macro", f1)
            mlflow.sklearn.log_model(model, "model")

            if accuracy > best_accuracy:
                best_accuracy = accuracy
                best_run_id = mlflow.active_run().info.run_id

    print(f"Best run: {best_run_id} with accuracy {best_accuracy:.4f}")
    return best_run_id

Logging Artifacts

Artifacts are files associated with a run: plots, data samples, configuration files, or any other file you want to preserve.

import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

def log_confusion_matrix(y_true, y_pred, class_names):
    """Create and log a confusion matrix plot."""
    cm = confusion_matrix(y_true, y_pred)
    disp = ConfusionMatrixDisplay(cm, display_labels=class_names)
    fig, ax = plt.subplots(figsize=(8, 6))
    disp.plot(ax=ax)
    plt.title("Confusion Matrix")
    plt.tight_layout()

    # Save and log
    plt.savefig("confusion_matrix.png", dpi=150)
    mlflow.log_artifact("confusion_matrix.png")
    plt.close()

def log_data_sample(X_test, y_test, n_samples=10):
    """Log a sample of test data for reproducibility."""
    import pandas as pd
    sample = pd.DataFrame(X_test[:n_samples])
    sample["target"] = y_test[:n_samples]
    sample.to_csv("test_sample.csv", index=False)
    mlflow.log_artifact("test_sample.csv")

Comparing Runs Programmatically

The MLflow UI provides visual comparison, but you can also query runs programmatically.

def find_best_run(experiment_name: str, metric: str = "accuracy") -> dict:
    """Find the best run in an experiment by a given metric."""
    experiment = mlflow.get_experiment_by_name(experiment_name)
    if not experiment:
        raise ValueError(f"Experiment '{experiment_name}' not found")

    runs = mlflow.search_runs(
        experiment_ids=[experiment.experiment_id],
        order_by=[f"metrics.{metric} DESC"],
        max_results=5,
    )

    if runs.empty:
        raise ValueError("No runs found")

    best = runs.iloc[0]
    print(f"Best run: {best['run_id']}")
    print(f"  {metric}: {best[f'metrics.{metric}']:.4f}")

    # Print all parameters for the best run
    param_cols = [c for c in runs.columns if c.startswith("params.")]
    for col in param_cols:
        param_name = col.replace("params.", "")
        print(f"  {param_name}: {best[col]}")

    return best.to_dict()

Model Registry

The Model Registry manages model versions and their deployment stages. After finding the best run, register the model for production use.

def register_best_model(
    experiment_name: str,
    registry_name: str = "iris-classifier",
):
    """Register the best model from an experiment."""
    best = find_best_run(experiment_name)
    run_id = best["run_id"]

    # Register the model
    model_uri = f"runs:/{run_id}/model"
    result = mlflow.register_model(model_uri, registry_name)

    print(f"Registered model: {result.name} version {result.version}")
    return result

def load_production_model(registry_name: str):
    """Load the latest version of a registered model."""
    model_uri = f"models:/{registry_name}/latest"
    model = mlflow.sklearn.load_model(model_uri)
    return model

The registry supports version aliases. Tag a specific version as “production” and your serving code always loads the right model without hardcoding version numbers.

Custom Metrics Over Time

For training loops (deep learning, iterative algorithms), log metrics at each step to see learning curves.

def train_with_step_logging(X_train, y_train, X_val, y_val):
    """Log metrics at each training iteration."""
    from sklearn.ensemble import GradientBoostingClassifier

    mlflow.set_experiment("gradient-boosting-curves")

    with mlflow.start_run(run_name="gbm-staged"):
        n_estimators = 200
        params = {"n_estimators": n_estimators, "learning_rate": 0.1}
        mlflow.log_params(params)

        model = GradientBoostingClassifier(**params, random_state=42)
        model.fit(X_train, y_train)

        # Log metrics at each boosting stage
        for i, y_pred in enumerate(model.staged_predict(X_val)):
            acc = accuracy_score(y_val, y_pred)
            mlflow.log_metric("val_accuracy", acc, step=i)

        # Final metrics
        final_pred = model.predict(X_val)
        mlflow.log_metric("final_accuracy", accuracy_score(y_val, final_pred))
        mlflow.sklearn.log_model(model, "model")

Step-logged metrics appear as line charts in the MLflow UI, making it easy to spot overfitting (validation metric peaking then declining) or underfitting (metric still climbing when training ended).

Organizing Experiments

Keep experiments organized with naming conventions and tags:

with mlflow.start_run(run_name="rf-v2-larger-dataset") as run:
    # Tags for filtering and organization
    mlflow.set_tags({
        "team": "ml-platform",
        "dataset_version": "v2.1",
        "purpose": "production-candidate",
        "framework": "sklearn",
    })

    # ... training code ...

Search with tags to find specific runs:

runs = mlflow.search_runs(
    experiment_ids=["1"],
    filter_string="tags.purpose = 'production-candidate'",
)

Summary

MLflow provides three capabilities that every ML team needs: experiment tracking to log and compare model runs, an artifact store to preserve models and plots alongside their metadata, and a model registry to manage versions for deployment. Start by adding mlflow.log_params and mlflow.log_metric to your training scripts. Once you have runs to compare, use the UI or programmatic search to find the best configurations. Register production-ready models in the registry so your serving infrastructure can load them by name instead of file path.