Experiment Tracking

Experiment Tracking

Overview

Experiment tracking enables reproducibility, comparison, and collaboration in ML development. It captures hyperparameters, metrics, artifacts, and model versions to ensure every experiment can be reproduced and compared.

When to Use This Reference

• Setting up MLflow for experiment tracking
• Implementing Weights & Biases integration
• Creating model registries and versioning
• Comparing experiments and selecting models
• Building custom tracking solutions

When NOT to Use

• Quick one-off experiments without reproducibility needs
• Simple scripts without hyperparameters
• Non-ML projects

MLflow Integration

Basic Experiment Tracking

import mlflow
from mlflow.tracking import MlflowClient
from pathlib import Path
import json

class MLflowTracker:
    """MLflow experiment tracking wrapper."""

    def __init__(
        self,
        experiment_name: str,
        tracking_uri: str = "http://localhost:5000",
        artifact_location: str = None,
    ):
        mlflow.set_tracking_uri(tracking_uri)

        # Create or get experiment
        experiment = mlflow.get_experiment_by_name(experiment_name)
        if experiment is None:
            self.experiment_id = mlflow.create_experiment(
                experiment_name,
                artifact_location=artifact_location,
            )
        else:
            self.experiment_id = experiment.experiment_id

        mlflow.set_experiment(experiment_name)
        self.client = MlflowClient()
        self.run = None

    def start_run(
        self,
        run_name: str = None,
        tags: dict = None,
        nested: bool = False,
    ) -> str:
        """Start a new MLflow run."""
        self.run = mlflow.start_run(
            run_name=run_name,
            experiment_id=self.experiment_id,
            nested=nested,
        )

        if tags:
            mlflow.set_tags(tags)

        return self.run.info.run_id

    def end_run(self, status: str = "FINISHED") -> None:
        """End the current run."""
        mlflow.end_run(status=status)
        self.run = None

    def log_params(self, params: dict) -> None:
        """Log hyperparameters."""
        mlflow.log_params(params)

    def log_metrics(self, metrics: dict, step: int = None) -> None:
        """Log metrics with optional step."""
        for key, value in metrics.items():
            mlflow.log_metric(key, value, step=step)

    def log_artifact(self, local_path: str, artifact_path: str = None) -> None:
        """Log file or directory as artifact."""
        mlflow.log_artifact(local_path, artifact_path)

    def log_model(
        self,
        model,
        artifact_path: str,
        registered_model_name: str = None,
        signature=None,
        input_example=None,
    ) -> str:
        """Log model with optional registration."""
        from mlflow.models import infer_signature

        if signature is None and input_example is not None:
            signature = infer_signature(input_example, model.predict(input_example))

        model_info = mlflow.sklearn.log_model(
            model,
            artifact_path=artifact_path,
            registered_model_name=registered_model_name,
            signature=signature,
            input_example=input_example,
        )

        return model_info.model_uri

# Usage example
def train_with_mlflow(
    model,
    X_train,
    y_train,
    X_val,
    y_val,
    params: dict,
):
    """Complete training run with MLflow tracking."""
    tracker = MLflowTracker("my_experiment")

    tracker.start_run(
        run_name=f"run_{params['model_type']}",
        tags={
            "model_type": params["model_type"],
            "dataset_version": "v1.0",
            "author": "ml-team",
        },
    )

    try:
        # Log parameters
        tracker.log_params(params)

        # Train model
        model.fit(X_train, y_train)

        # Evaluate and log metrics
        train_score = model.score(X_train, y_train)
        val_score = model.score(X_val, y_val)

        tracker.log_metrics({
            "train_accuracy": train_score,
            "val_accuracy": val_score,
        })

        # Log model
        model_uri = tracker.log_model(
            model,
            artifact_path="model",
            registered_model_name="my_model",
            input_example=X_train[:5],
        )

        tracker.end_run()
        return model_uri

    except Exception as e:
        tracker.end_run(status="FAILED")
        raise

PyTorch Model Logging

import mlflow.pytorch
import torch

def log_pytorch_model(
    model: torch.nn.Module,
    artifact_path: str,
    registered_model_name: str = None,
    sample_input: torch.Tensor = None,
) -> str:
    """Log PyTorch model with signature inference."""
    from mlflow.models import infer_signature

    # Create signature from sample input
    signature = None
    if sample_input is not None:
        model.eval()
        with torch.no_grad():
            sample_output = model(sample_input)

        signature = infer_signature(
            sample_input.numpy(),
            sample_output.numpy(),
        )

    model_info = mlflow.pytorch.log_model(
        model,
        artifact_path=artifact_path,
        registered_model_name=registered_model_name,
        signature=signature,
    )

    return model_info.model_uri

def load_pytorch_model(model_uri: str, device: str = "cpu") -> torch.nn.Module:
    """Load PyTorch model from MLflow."""
    model = mlflow.pytorch.load_model(model_uri, map_location=device)
    return model

Model Registry Operations

from mlflow.tracking import MlflowClient
from mlflow.entities.model_registry import ModelVersion

class ModelRegistry:
    """MLflow Model Registry wrapper."""

    def __init__(self, tracking_uri: str = "http://localhost:5000"):
        mlflow.set_tracking_uri(tracking_uri)
        self.client = MlflowClient()

    def register_model(
        self,
        model_uri: str,
        name: str,
        tags: dict = None,
        description: str = None,
    ) -> ModelVersion:
        """Register a new model version."""
        result = mlflow.register_model(model_uri, name)

        if tags:
            for key, value in tags.items():
                self.client.set_model_version_tag(name, result.version, key, value)

        if description:
            self.client.update_model_version(
                name,
                result.version,
                description=description,
            )

        return result

    def transition_model_stage(
        self,
        name: str,
        version: str,
        stage: str,
        archive_existing: bool = True,
    ) -> ModelVersion:
        """Transition model to new stage (Staging, Production, Archived)."""
        return self.client.transition_model_version_stage(
            name=name,
            version=version,
            stage=stage,
            archive_existing_versions=archive_existing,
        )

    def get_latest_version(
        self,
        name: str,
        stages: list[str] = None,
    ) -> list[ModelVersion]:
        """Get latest model versions by stage."""
        return self.client.get_latest_versions(name, stages=stages)

    def load_production_model(self, name: str) -> any:
        """Load the production model."""
        model_uri = f"models:/{name}/Production"
        return mlflow.pyfunc.load_model(model_uri)

    def compare_versions(
        self,
        name: str,
        version_a: str,
        version_b: str,
    ) -> dict:
        """Compare two model versions."""
        v_a = self.client.get_model_version(name, version_a)
        v_b = self.client.get_model_version(name, version_b)

        run_a = self.client.get_run(v_a.run_id)
        run_b = self.client.get_run(v_b.run_id)

        return {
            "version_a": {
                "version": version_a,
                "metrics": run_a.data.metrics,
                "params": run_a.data.params,
            },
            "version_b": {
                "version": version_b,
                "metrics": run_b.data.metrics,
                "params": run_b.data.params,
            },
        }

Weights & Biases Integration

Basic W&B Tracking

import wandb
from pathlib import Path

class WandbTracker:
    """Weights & Biases experiment tracking wrapper."""

    def __init__(
        self,
        project: str,
        entity: str = None,
        config: dict = None,
    ):
        self.project = project
        self.entity = entity
        self.config = config
        self.run = None

    def start_run(
        self,
        name: str = None,
        tags: list[str] = None,
        group: str = None,
        job_type: str = "train",
        resume: str = None,
    ) -> wandb.Run:
        """Initialize W&B run."""
        self.run = wandb.init(
            project=self.project,
            entity=self.entity,
            name=name,
            config=self.config,
            tags=tags,
            group=group,
            job_type=job_type,
            resume=resume,
        )
        return self.run

    def log(self, data: dict, step: int = None, commit: bool = True) -> None:
        """Log metrics and data."""
        wandb.log(data, step=step, commit=commit)

    def log_artifact(
        self,
        name: str,
        artifact_type: str,
        path: str,
        metadata: dict = None,
    ) -> wandb.Artifact:
        """Log artifact (model, dataset, etc.)."""
        artifact = wandb.Artifact(
            name=name,
            type=artifact_type,
            metadata=metadata,
        )

        if Path(path).is_dir():
            artifact.add_dir(path)
        else:
            artifact.add_file(path)

        self.run.log_artifact(artifact)
        return artifact

    def log_model(
        self,
        model_path: str,
        name: str,
        metadata: dict = None,
        aliases: list[str] = None,
    ) -> wandb.Artifact:
        """Log model artifact with aliases."""
        artifact = wandb.Artifact(
            name=name,
            type="model",
            metadata=metadata,
        )

        if Path(model_path).is_dir():
            artifact.add_dir(model_path)
        else:
            artifact.add_file(model_path)

        self.run.log_artifact(artifact, aliases=aliases or ["latest"])
        return artifact

    def watch_model(
        self,
        model,
        log: str = "all",
        log_freq: int = 100,
    ) -> None:
        """Watch model for gradient and parameter logging."""
        wandb.watch(model, log=log, log_freq=log_freq)

    def finish(self, exit_code: int = 0) -> None:
        """Finish the run."""
        wandb.finish(exit_code=exit_code)

# Usage with PyTorch
def train_with_wandb(
    model: torch.nn.Module,
    train_loader,
    val_loader,
    config: dict,
):
    """Training with W&B tracking."""
    tracker = WandbTracker(
        project="my-project",
        config=config,
    )

    tracker.start_run(
        name=f"experiment_{config['model_type']}",
        tags=["baseline", config["model_type"]],
        group="hyperparameter_search",
    )

    # Watch model gradients
    tracker.watch_model(model)

    for epoch in range(config["epochs"]):
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            # Training step
            loss = train_step(model, data, target)

            tracker.log({
                "train/loss": loss,
                "train/epoch": epoch,
            })

        # Validation
        val_metrics = evaluate(model, val_loader)
        tracker.log({
            "val/loss": val_metrics["loss"],
            "val/accuracy": val_metrics["accuracy"],
            "epoch": epoch,
        })

    # Save and log model
    torch.save(model.state_dict(), "model.pt")
    tracker.log_model(
        "model.pt",
        name="trained_model",
        metadata={"accuracy": val_metrics["accuracy"]},
        aliases=["latest", "best"],
    )

    tracker.finish()

W&B Sweeps for Hyperparameter Tuning

import wandb

sweep_config = {
    "method": "bayes",  # bayes, grid, random
    "metric": {
        "name": "val/loss",
        "goal": "minimize",
    },
    "parameters": {
        "learning_rate": {
            "distribution": "log_uniform_values",
            "min": 1e-5,
            "max": 1e-2,
        },
        "batch_size": {
            "values": [16, 32, 64, 128],
        },
        "hidden_size": {
            "values": [128, 256, 512],
        },
        "dropout": {
            "distribution": "uniform",
            "min": 0.1,
            "max": 0.5,
        },
    },
    "early_terminate": {
        "type": "hyperband",
        "min_iter": 3,
    },
}

def sweep_train():
    """Training function for sweep."""
    with wandb.init() as run:
        config = wandb.config

        model = build_model(
            hidden_size=config.hidden_size,
            dropout=config.dropout,
        )

        optimizer = torch.optim.Adam(
            model.parameters(),
            lr=config.learning_rate,
        )

        train_loader = DataLoader(train_dataset, batch_size=config.batch_size)

        for epoch in range(10):
            loss = train_epoch(model, train_loader, optimizer)
            val_loss = evaluate(model, val_loader)

            wandb.log({
                "train/loss": loss,
                "val/loss": val_loss,
                "epoch": epoch,
            })

# Run sweep
sweep_id = wandb.sweep(sweep_config, project="my-project")
wandb.agent(sweep_id, function=sweep_train, count=50)

Custom Experiment Tracking

Lightweight Tracker

import json
from datetime import datetime
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import Optional
import hashlib
import uuid

@dataclass
class Experiment:
    """Experiment metadata and results."""
    experiment_id: str
    name: str
    params: dict
    metrics: dict = field(default_factory=dict)
    artifacts: list = field(default_factory=list)
    tags: dict = field(default_factory=dict)
    start_time: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    end_time: Optional[str] = None
    status: str = "running"

    def to_dict(self) -> dict:
        return asdict(self)

class SimpleTracker:
    """Lightweight file-based experiment tracker."""

    def __init__(self, experiments_dir: str = "./experiments"):
        self.experiments_dir = Path(experiments_dir)
        self.experiments_dir.mkdir(parents=True, exist_ok=True)
        self.current_experiment: Optional[Experiment] = None

    def start_experiment(
        self,
        name: str,
        params: dict,
        tags: dict = None,
    ) -> Experiment:
        """Start a new experiment."""
        experiment_id = str(uuid.uuid4())[:8]

        self.current_experiment = Experiment(
            experiment_id=experiment_id,
            name=name,
            params=params,
            tags=tags or {},
        )

        # Create experiment directory
        exp_dir = self.experiments_dir / experiment_id
        exp_dir.mkdir(exist_ok=True)

        self._save_experiment()
        return self.current_experiment

    def log_metrics(self, metrics: dict, step: int = None) -> None:
        """Log metrics to current experiment."""
        if self.current_experiment is None:
            raise ValueError("No active experiment")

        for key, value in metrics.items():
            if key not in self.current_experiment.metrics:
                self.current_experiment.metrics[key] = []

            self.current_experiment.metrics[key].append({
                "value": value,
                "step": step,
                "timestamp": datetime.utcnow().isoformat(),
            })

        self._save_experiment()

    def log_artifact(self, path: str, name: str = None) -> str:
        """Copy artifact to experiment directory."""
        if self.current_experiment is None:
            raise ValueError("No active experiment")

        import shutil

        source = Path(path)
        artifact_name = name or source.name
        exp_dir = self.experiments_dir / self.current_experiment.experiment_id
        dest = exp_dir / "artifacts" / artifact_name

        dest.parent.mkdir(parents=True, exist_ok=True)

        if source.is_dir():
            shutil.copytree(source, dest)
        else:
            shutil.copy2(source, dest)

        self.current_experiment.artifacts.append(str(dest))
        self._save_experiment()

        return str(dest)

    def end_experiment(self, status: str = "completed") -> None:
        """End current experiment."""
        if self.current_experiment is None:
            return

        self.current_experiment.status = status
        self.current_experiment.end_time = datetime.utcnow().isoformat()
        self._save_experiment()
        self.current_experiment = None

    def _save_experiment(self) -> None:
        """Save experiment to JSON file."""
        if self.current_experiment is None:
            return

        exp_dir = self.experiments_dir / self.current_experiment.experiment_id
        with open(exp_dir / "experiment.json", "w") as f:
            json.dump(self.current_experiment.to_dict(), f, indent=2)

    def load_experiment(self, experiment_id: str) -> Experiment:
        """Load experiment by ID."""
        exp_file = self.experiments_dir / experiment_id / "experiment.json"
        with open(exp_file) as f:
            data = json.load(f)
        return Experiment(**data)

    def list_experiments(self, tags: dict = None) -> list[Experiment]:
        """List all experiments, optionally filtered by tags."""
        experiments = []

        for exp_dir in self.experiments_dir.iterdir():
            if not exp_dir.is_dir():
                continue

            exp_file = exp_dir / "experiment.json"
            if not exp_file.exists():
                continue

            exp = self.load_experiment(exp_dir.name)

            if tags:
                if not all(exp.tags.get(k) == v for k, v in tags.items()):
                    continue

            experiments.append(exp)

        return sorted(experiments, key=lambda x: x.start_time, reverse=True)

    def compare_experiments(self, experiment_ids: list[str]) -> dict:
        """Compare metrics across experiments."""
        comparison = {}

        for exp_id in experiment_ids:
            exp = self.load_experiment(exp_id)
            comparison[exp_id] = {
                "name": exp.name,
                "params": exp.params,
                "final_metrics": {
                    k: v[-1]["value"] if v else None
                    for k, v in exp.metrics.items()
                },
            }

        return comparison

Experiment Comparison and Analysis

Metrics Comparison

import pandas as pd
import matplotlib.pyplot as plt
from mlflow.tracking import MlflowClient

def compare_runs(
    experiment_name: str,
    metric_keys: list[str],
    n_runs: int = 10,
) -> pd.DataFrame:
    """Compare recent runs in an experiment."""
    client = MlflowClient()
    experiment = client.get_experiment_by_name(experiment_name)

    runs = client.search_runs(
        experiment_ids=[experiment.experiment_id],
        order_by=["start_time DESC"],
        max_results=n_runs,
    )

    data = []
    for run in runs:
        row = {
            "run_id": run.info.run_id,
            "run_name": run.info.run_name,
            "status": run.info.status,
            "start_time": run.info.start_time,
        }
        row.update(run.data.params)
        row.update({k: run.data.metrics.get(k) for k in metric_keys})
        data.append(row)

    return pd.DataFrame(data)

def plot_metric_comparison(
    runs_df: pd.DataFrame,
    metric: str,
    group_by: str = None,
) -> plt.Figure:
    """Plot metric comparison across runs."""
    fig, ax = plt.subplots(figsize=(10, 6))

    if group_by:
        for group, group_df in runs_df.groupby(group_by):
            ax.bar(group_df["run_name"], group_df[metric], label=str(group))
        ax.legend(title=group_by)
    else:
        ax.bar(runs_df["run_name"], runs_df[metric])

    ax.set_xlabel("Run")
    ax.set_ylabel(metric)
    ax.set_title(f"Comparison of {metric}")
    plt.xticks(rotation=45, ha="right")
    plt.tight_layout()

    return fig

Best Practices

What to Track

# Always track:
REQUIRED_PARAMS = [
    "learning_rate",
    "batch_size",
    "epochs",
    "model_architecture",
    "optimizer",
    "random_seed",
    "dataset_version",
]

REQUIRED_METRICS = [
    "train_loss",
    "val_loss",
    "train_accuracy",
    "val_accuracy",
]

REQUIRED_ARTIFACTS = [
    "model_checkpoint",
    "training_config",
    "requirements.txt",
]

# Recommended tags
RECOMMENDED_TAGS = {
    "author": "username",
    "environment": "dev|staging|prod",
    "model_type": "classification|regression|etc",
    "dataset": "dataset_name",
    "git_commit": "commit_hash",
}

Experiment Naming Conventions

# Good naming patterns
run_name = f"{model_type}_{dataset}_{timestamp}"
run_name = f"exp_{experiment_number:03d}_{description}"
run_name = f"{feature_flag}_{ablation_type}_{seed}"

# Organize with groups and tags
tags = {
    "project": "recommendation_engine",
    "sprint": "sprint_42",
    "hypothesis": "larger_embedding_helps",
}

Related References

• training-pipelines.md - Integrating tracking with training
• model-validation.md - Validating tracked models
• pipeline-orchestration.md - Tracking in automated pipelines

Cross-Reference Skills

• DevOps Engineer - MLflow server deployment
• Data Engineer - Artifact storage integration