:::tip 🎮 Interactive Playground Visualize this concept: Try the MLflow Tracking demo on the EngineersOfAI Playground - no code required. :::

Experiment Tracking

50 Scientists, One MLflow Instance

The MLflow tracking server had 14,000 runs. Nobody could find anything.

The company had grown from 5 to 50 data scientists over 18 months. Early on, the team had deployed MLflow with a shared tracking server and told everyone "log your experiments here." They had. With no naming conventions, no tagging strategy, and no organizational structure, the result was a garbage dump of experiments with names like run_1, final_final_v3, test_do_not_use, and johns_experiment_tuesday.

Experiments from six months ago were indistinguishable from experiments from last week. Models with identical hyperparameter names used completely different value ranges. Some runs had 3 metrics logged; others had 40. There was no way to tell which run produced the model currently serving in production. When a data scientist left the company, their experimental knowledge left with them - locked inside 2,000 runs with no documentation.

The platform team's response was not to replace MLflow. It was to build governance around it. After 6 weeks of work, the same 14,000 runs became navigable: every run had an owner, every experiment had a description, every production model traced back to a specific run. New experiments followed naming conventions enforced at the SDK level. Runs not linked to a model registry entry within 30 days were automatically archived.

The lesson: experiment tracking infrastructure is necessary but not sufficient. The governance structure - conventions, tooling, and automated enforcement - is what makes it useful at scale.

---

Why Experiment Tracking Exists

Before systematic experiment tracking, ML research was the Wild West. A researcher would train dozens of model variants, keep notes in a text file or Google Doc, and eventually lose track of what they'd tried. The "best model" was whichever one the researcher remembered being good, not whichever one actually was.

This approach fails in three critical ways. First, irreproducibility: without logging the exact hyperparameters, random seeds, data splits, and code commit used for training, the "best" experiment can never be reproduced. Second, lost knowledge: when a researcher moves to a new project, their experimental knowledge is gone. Newcomers repeat the same experiments. Third, collaboration overhead: two researchers working on the same problem duplicate experiments because there's no shared record of what's been tried.

Experiment tracking solves all three by creating a structured, searchable, permanent record of every training run - including the configuration used, the results achieved, and the artifacts produced.

---

Historical Context

The first ML experiment tracking tools were research lab notebook systems - literal paper notebooks where scientists recorded results. WANDB (Weights & Biases) was founded in 2018 and was among the first companies to build a polished, cloud-hosted tracking service specifically for deep learning. MLflow was open-sourced by Databricks in 2018, offering a self-hosted alternative. Sacred (2017) was an earlier Python library for experiment tracking, now largely superseded.

The real driver of experiment tracking adoption was the deep learning era: models became expensive enough to train that losing track of a promising experiment was a genuine loss of thousands of dollars of compute. When a single training run costs $5,000, you cannot afford not to have a record of what it produced.

---

MLflow Architecture

MLflow is the most widely deployed self-hosted experiment tracking system. Understanding its architecture is essential for deploying it at scale.

Production MLflow Deployment

# docker-compose.yml for production MLflow tracking server
version: '3.8'

services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.10.0
    ports:
      - "5000:5000"
    environment:
      # PostgreSQL backend for scalable metadata storage
      MLFLOW_BACKEND_STORE_URI: postgresql://mlflow:${POSTGRES_PASSWORD}@postgres:5432/mlflow
      # S3 for artifact storage (model files, plots, datasets)
      MLFLOW_DEFAULT_ARTIFACT_ROOT: s3://your-mlflow-artifacts/
      # Authentication (using nginx reverse proxy with basic auth)
      MLFLOW_AUTH_CONFIG_PATH: /etc/mlflow/auth_config.yaml
    command: >
      mlflow server
      --host 0.0.0.0
      --port 5000
      --backend-store-uri postgresql://mlflow:${POSTGRES_PASSWORD}@postgres:5432/mlflow
      --default-artifact-root s3://your-mlflow-artifacts/
      --workers 4
    volumes:
      - ./auth_config.yaml:/etc/mlflow/auth_config.yaml

  postgres:
    image: postgres:15
    environment:
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
      POSTGRES_DB: mlflow
    volumes:
      - postgres_data:/var/lib/postgresql/data

volumes:
  postgres_data:

:::tip Scale considerations For teams over 20 data scientists, run MLflow with at least 4 workers and connect to RDS/Cloud SQL instead of local Postgres. 14,000 runs is not a large dataset for Postgres - the bottleneck is typically the artifact listing queries on S3 when displaying the UI. :::

---

The Tracking SDK: What to Log

The most impactful decision in experiment tracking is what metadata to record. Log too little and runs are unreproducible. Log too much and the UI becomes noise.

The Minimum Required Set

import mlflow
import mlflow.pytorch
import subprocess
import json
from pathlib import Path

def create_training_run(
    experiment_name: str,
    run_name: str,
    tags: dict,
    hyperparams: dict,
    training_fn,
    dataset_path: str,
    model_save_path: str,
) -> mlflow.entities.Run:
    """
    Standard training run wrapper with complete metadata logging.
    Use this as the base for all training scripts.
    """
    mlflow.set_experiment(experiment_name)

    with mlflow.start_run(run_name=run_name, tags=tags) as run:

        # === REQUIRED: Configuration ===
        mlflow.log_params(hyperparams)

        # Code versioning - critical for reproducibility
        git_hash = subprocess.check_output(
            ["git", "rev-parse", "HEAD"]
        ).decode().strip()
        mlflow.set_tag("git_commit", git_hash)
        mlflow.set_tag("git_branch", subprocess.check_output(
            ["git", "rev-parse", "--abbrev-ref", "HEAD"]
        ).decode().strip())

        # Dataset versioning
        import hashlib
        dataset_hash = hashlib.md5(Path(dataset_path).read_bytes()).hexdigest()
        mlflow.set_tag("dataset_hash", dataset_hash)
        mlflow.set_tag("dataset_path", dataset_path)

        # Environment
        import platform, torch
        mlflow.set_tag("python_version", platform.python_version())
        mlflow.set_tag("torch_version", torch.__version__)
        mlflow.set_tag("cuda_version", torch.version.cuda or "cpu")

        # Required tags for governance
        required_tags = ["team", "project", "owner"]
        for tag in required_tags:
            if tag not in tags:
                raise ValueError(f"Missing required tag: {tag}")

        # === TRAINING ===
        model, metrics = training_fn(hyperparams)

        # === REQUIRED: Metrics ===
        mlflow.log_metrics(metrics)

        # === REQUIRED: Model artifact ===
        mlflow.pytorch.log_model(
            model,
            "model",
            registered_model_name=f"{experiment_name}-model",
        )

        # === OPTIONAL: Plots, configs ===
        mlflow.log_artifact("training_config.yaml")

        return run

Naming Conventions at Scale

Without naming conventions, MLflow becomes unusable at scale. Enforce these programmatically:

import re
from datetime import datetime

class ExperimentNamingConvention:
    """
    Enforces experiment and run naming conventions.
    Call validate() before creating any experiment or run.
    """

    # Format: {team}/{project}/{model-type}
    # Example: "recommendations/user-embedding/transformer"
    EXPERIMENT_PATTERN = re.compile(r'^[a-z][a-z0-9-]+/[a-z][a-z0-9-]+/[a-z][a-z0-9-]+$')

    # Format: {YYYY-MM-DD}_{description}_{optional-suffix}
    # Example: "2024-03-15_dropout-ablation_v2"
    RUN_PATTERN = re.compile(r'^\d{4}-\d{2}-\d{2}_[a-z][a-z0-9-]+(_[a-z0-9]+)?$')

    @classmethod
    def validate_experiment_name(cls, name: str) -> None:
        if not cls.EXPERIMENT_PATTERN.match(name):
            raise ValueError(
                f"Invalid experiment name: '{name}'\n"
                f"Required format: team/project/model-type\n"
                f"Example: recommendations/user-embedding/transformer"
            )

    @classmethod
    def validate_run_name(cls, name: str) -> None:
        if not cls.RUN_PATTERN.match(name):
            raise ValueError(
                f"Invalid run name: '{name}'\n"
                f"Required format: YYYY-MM-DD_description\n"
                f"Example: {datetime.now().strftime('%Y-%m-%d')}_dropout-ablation"
            )

    @classmethod
    def generate_run_name(cls, description: str) -> str:
        """Generate a valid run name from a description."""
        date = datetime.now().strftime('%Y-%m-%d')
        slug = re.sub(r'[^a-z0-9-]', '-', description.lower()).strip('-')
        return f"{date}_{slug}"

---

Weights & Biases: The Hosted Alternative

W&B offers most of the same capabilities as MLflow with a significantly better UI and easier setup. The key differences:

Feature	MLflow	W&B
Hosting	Self-hosted (or Databricks)	Cloud-hosted (or private cloud)
Setup time	1–2 days	30 minutes
UI quality	Good	Excellent
Sweeps	Manual (Optuna integration)	Built-in sweeps UI
Team features	Basic	Advanced (access control, reports)
Price	Free (infra cost only)	$50/seat/month
Data privacy	You control	W&B servers

# W&B integration - nearly identical API to MLflow
import wandb

def train_with_wandb(config: dict, tags: list[str]):
    """
    Training run with W&B tracking.
    The API surface is very similar to MLflow.
    """
    run = wandb.init(
        project="user-embedding",
        name=f"transformer-{config['hidden_dim']}d",
        config=config,
        tags=tags,
        notes="Testing dropout values 0.1-0.5",  # searchable description
    )

    model = build_model(config)

    for epoch in range(config["num_epochs"]):
        train_metrics = train_epoch(model, train_loader)
        val_metrics = validate(model, val_loader)

        # Log metrics - auto-plotted in W&B UI
        wandb.log({
            "epoch": epoch,
            "train/loss": train_metrics["loss"],
            "train/accuracy": train_metrics["accuracy"],
            "val/loss": val_metrics["loss"],
            "val/accuracy": val_metrics["accuracy"],
        })

    # Save model artifact
    artifact = wandb.Artifact("user-embedding-model", type="model")
    artifact.add_file("model.pt")
    run.log_artifact(artifact)

    wandb.finish()

W&B Sweeps: Hyperparameter Optimization

One area where W&B clearly leads MLflow is sweeps - automated hyperparameter optimization with a UI for visualization:

import wandb

sweep_config = {
    "method": "bayes",  # Bayesian optimization
    "metric": {"name": "val/accuracy", "goal": "maximize"},
    "parameters": {
        "learning_rate": {"distribution": "log_uniform_values", "min": 1e-5, "max": 1e-2},
        "dropout": {"values": [0.1, 0.2, 0.3, 0.4, 0.5]},
        "hidden_dim": {"values": [128, 256, 512]},
        "num_layers": {"min": 2, "max": 8},
    },
    "early_terminate": {
        "type": "hyperband",
        "min_iter": 3,
    },
}

def train_sweep():
    """Training function called by W&B sweep agent."""
    with wandb.init() as run:
        config = run.config
        model, metrics = train(config)
        wandb.log({"val/accuracy": metrics["val_accuracy"]})


sweep_id = wandb.sweep(sweep_config, project="user-embedding")
wandb.agent(sweep_id, function=train_sweep, count=50)  # run 50 trials

---

Organizing Experiments at Scale

The Hierarchy

MLflow's hierarchy: Experiments → Runs → Artifacts. A well-organized hierarchy makes search practical:

Experiments (one per team-project-model-type)
├── recommendations/user-embedding/transformer
│   ├── Run: 2024-03-15_baseline
│   ├── Run: 2024-03-16_dropout-0.2
│   └── Run: 2024-03-17_larger-context
└── recommendations/user-embedding/mlp
    ├── Run: 2024-03-10_baseline
    └── Run: 2024-03-11_wider-hidden

Governance Policies

import mlflow
from datetime import datetime, timedelta

class MLflowGovernancePolicy:
    """
    Automated governance for MLflow experiments.
    Run daily via cron job.
    """

    def __init__(self, client: mlflow.MlflowClient):
        self.client = client

    def archive_old_unregistered_runs(self, days: int = 30) -> list[str]:
        """
        Archive runs older than N days that aren't linked to a model registry entry.
        Returns list of archived run IDs.
        """
        cutoff_ms = (datetime.utcnow() - timedelta(days=days)).timestamp() * 1000

        experiments = self.client.search_experiments()
        archived = []

        for experiment in experiments:
            runs = self.client.search_runs(
                experiment_ids=[experiment.experiment_id],
                filter_string=f"attributes.start_time < {int(cutoff_ms)}",
                run_view_type=mlflow.entities.ViewType.ACTIVE_ONLY,
            )

            for run in runs:
                # Check if this run is linked to a registered model version
                model_versions = self.client.search_model_versions(
                    filter_string=f"run_id = '{run.info.run_id}'"
                )

                if not model_versions:
                    # No registered model linked - archive the run
                    self.client.set_tag(
                        run.info.run_id,
                        "archived_reason",
                        f"Not linked to model registry after {days} days"
                    )
                    self.client.delete_run(run.info.run_id)
                    archived.append(run.info.run_id)

        return archived

    def find_runs_missing_required_tags(
        self,
        required_tags: list[str] = None,
    ) -> list[dict]:
        """Find runs missing required metadata tags."""
        if required_tags is None:
            required_tags = ["team", "project", "owner", "git_commit"]

        non_compliant = []
        experiments = self.client.search_experiments()

        for experiment in experiments:
            runs = self.client.search_runs(
                experiment_ids=[experiment.experiment_id],
                run_view_type=mlflow.entities.ViewType.ACTIVE_ONLY,
            )

            for run in runs:
                existing_tags = set(run.data.tags.keys())
                missing = [t for t in required_tags if t not in existing_tags]

                if missing:
                    non_compliant.append({
                        "run_id": run.info.run_id,
                        "run_name": run.info.run_name,
                        "experiment": experiment.name,
                        "missing_tags": missing,
                    })

        return non_compliant

---

Reproducibility Requirements

A run is reproducible if, given the logged information, you can recreate the exact model. This requires:

1. Code version - git commit hash
2. Data version - hash or versioned dataset reference
3. Random seeds - all seeds (Python, NumPy, PyTorch) logged
4. Environment - Python version, library versions
5. Hyperparameters - every configurable parameter

import random
import numpy as np
import torch

def set_reproducible_seeds(seed: int) -> dict:
    """Set all random seeds and return them for logging."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # For full determinism on CUDA (slower but reproducible)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    return {
        "random_seed": seed,
        "numpy_seed": seed,
        "torch_seed": seed,
        "cuda_deterministic": True,
    }

def log_environment(run):
    """Log complete environment for reproducibility."""
    import pkg_resources
    import platform

    installed_packages = {
        pkg.project_name: pkg.version
        for pkg in pkg_resources.working_set
        if pkg.project_name in [
            "torch", "transformers", "numpy", "pandas",
            "scikit-learn", "mlflow", "datasets"
        ]
    }

    mlflow.log_dict(installed_packages, "requirements.json")
    mlflow.set_tag("platform", platform.platform())
    mlflow.set_tag("python_version", platform.python_version())

---

Common Mistakes

:::danger Logging metrics outside the training loop Logging only the final validation metric loses all information about training dynamics - loss curves, gradient norms, learning rate schedule. Log metrics every N steps during training. This is essential for diagnosing training instability (loss spikes), premature convergence, or learning rate issues. :::

:::warning Treating experiment names as throw-away Run names like "test1", "run_v2", "final" become meaningless in 2 weeks. Enforce naming conventions at the SDK level - raise an exception if the run name doesn't match the pattern. The 30 seconds of friction at run creation saves hours of confusion later. :::

:::danger Not logging the dataset version Training on slightly different data produces different models. If you don't log a hash or versioned reference to the training data, runs with the same hyperparameters but different data look identical in the UI. This is the most common source of "I can't reproduce this result." :::

:::warning Building experiment tracking from scratch Teams sometimes build custom experiment tracking because MLflow "doesn't do X." Almost always, the right answer is to extend MLflow (custom plugins, wrapper SDKs) rather than replace it. Building a complete tracking system from scratch takes months and almost never produces something better than existing OSS tools. :::

---

Interview Q&A

Q: What metadata should every ML experiment log?

A: I think of it in four categories. (1) Configuration: all hyperparameters, model architecture, training schedule - everything that determines the model's behavior. (2) Code version: git commit hash and branch - essential for reproducing the run months later. (3) Data version: hash of training data or a reference to a versioned dataset - different data means different results even with identical code and hyperparameters. (4) Environment: Python version, key library versions - training behavior can change across library versions. Beyond these minimums, I log training metrics at every epoch (not just the final value), random seeds for reproducibility, and hardware information (GPU type, distributed training config). In MLflow, parameters go in log_params(), metrics in log_metrics(), and everything else as tags.

Q: How do you handle experiment organization for a team of 50 data scientists?

A: Three-layer approach. First, naming conventions enforced at the SDK level - experiments named team/project/model-type and runs named YYYY-MM-DD_description. Non-conforming names raise exceptions at run creation, not as a polite suggestion. Second, required tags: team, project, owner, git_commit - enforced at the SDK wrapper level, not the MLflow default. Third, governance automation: a daily job archives runs older than 30 days that aren't linked to a model registry entry, and a weekly report goes to team leads showing their team's experiment hygiene score. The key insight: governance must be automated or it won't happen. Asking 50 scientists to "please follow the naming convention" gets you a 20% adoption rate. Making the SDK enforce it gets you 99%.

Q: When would you choose W&B over MLflow?

A: Team size is the primary driver. For teams under 20 data scientists, W&B is almost always worth the $50/seat/month - the UI is significantly better, sweeps are built-in, and the time saved on setup and maintenance exceeds the cost. For teams over 30, the seat cost ($18,000+/year) starts to approach the cost of a self-hosted MLflow instance with a dedicated maintainer. I'd also choose W&B when: the team does extensive hyperparameter optimization (W&B sweeps are much better), when data scientists are remote and need great collaboration features, or when the team is new to experiment tracking and needs quick adoption. I'd choose MLflow when: data privacy requires self-hosted storage, the team is already Databricks-based (MLflow is native), or at-scale cost matters more than UX polish.

Q: How do you connect experiment runs to production models?

A: Through the model registry. When a run produces a model good enough to potentially deploy, you register it: mlflow.register_model(f"runs:/{run_id}/model", "model-name"). This creates a model version in the registry with an explicit link back to the run ID. From that run ID, you can trace: the hyperparameters used, the training data hash, the git commit, and every metric from training. This chain - production model version → run ID → training config + data version + code version - is the model lineage. Without it, you can't answer "what training data went into the model currently serving traffic?" which is a question that comes up in every incident investigation and every compliance audit.

Q: What is the offline-online experiment gap and how do you manage it?

A: The offline-online gap is the difference between how a model performs on your holdout test set (offline) and how it performs when serving real users (online). It exists because: holdout data is historical and may not reflect current user behavior, users interact with recommendations in ways you can't simulate offline, and small differences in serving infrastructure can change predictions. To manage it: first, always A/B test models before full rollout - offline performance is a filtering criterion, not a deployment criterion. Second, track the offline-online correlation over time - if your test set accuracy has historically translated to 30% of online lift, use that as your conversion factor for estimating expected online gains. Third, investigate when the correlation breaks - it often signals data drift or that the holdout set has become unrepresentative.