Experiment Tracking & Model Registry: The Version Control for ML

Imagine deploying software without Git. No record of what changed, who changed it, or why. No ability to revert to a working state. That's the situation at ML teams without experiment tracking and a model registry.

Typical failures without experiment tracking:

• 'Which experiment produced the model that's currently in production?' Nobody knows — the Jupyter notebook was overwritten.
• 'Production accuracy dropped today. Did we deploy a new model yesterday?' No record. The engineer who trained it is on vacation.
• 'Let's revert to last month's model.' Which one? The weights file in S3 doesn't have metadata. What dataset was it trained on?
• 'Why does the model score user_id=12345 as high-risk?' No feature values were logged at prediction time. Cannot debug.

Experiment tracking solves: 'which code, data, and hyperparameters produced this model, and what were the results?'

Model registry solves: 'which model is in production right now, what's its lineage, and how do I roll it back?'

They are complementary: tracking manages the experiment phase, the registry manages the production lifecycle.

MLflow is the most widely adopted open-source experiment tracking and model registry system. Four components:

MLflow Tracking: Logs parameters (hyperparameters, data paths), metrics (loss curves, AUC at each epoch), artifacts (model weights, confusion matrix plots, feature importance charts). Every training run creates a unique run_id. Runs are organized into experiments (one per model type or project). Access via UI (browser-based) or Python API.

import mlflow

with mlflow.start_run(run_name='fraud_model_xgb_v3') as run:
    mlflow.log_params({'n_estimators': 500, 'learning_rate': 0.05, 'max_depth': 6})
    mlflow.log_metric('val_auc_pr', 0.847, step=epoch)
    mlflow.log_metric('val_recall_at_80pct_precision', 0.63)
    mlflow.xgboost.log_model(model, 'model',
        signature=mlflow.models.infer_signature(X_train, predictions),
        input_example=X_train[:5]
    )
    run_id = run.info.run_id

MLflow Projects: Package ML code in a reproducible format (conda.yaml or Docker). mlflow run executes a project in an isolated environment. Solves 'it works on my machine' problems.

MLflow Models: A standard format for packaging ML models from any framework (PyTorch, sklearn, XGBoost). Includes: the model artifact, a MLmodel file (metadata), a conda.yaml (dependencies), input/output signature (expected feature names and types). The signature is critical: it validates that the serving infrastructure provides exactly the features the model expects.

MLflow Model Registry: Versioned catalog of models. A registered model has named versions (version 1, 2, 3...). Each version has a Stage: None → Staging → Production → Archived. Only one version should be in Production at a time. APIs to transition versions between stages.

Weights & Biases (W&B) offers richer visualization and team collaboration features than MLflow, at the cost of being a hosted SaaS product (with a self-hosted option).

Where W&B exceeds MLflow:

Media logging: Log images, audio, video, 3D point clouds, and custom visualizations directly in the UI. Essential for computer vision and audio ML teams.

Sweeps: Built-in hyperparameter optimization that automatically runs multiple experiments with different configurations (grid search, Bayesian optimization, random search) and visualizes the hyperparameter-performance landscape.

Team features: Comments on runs, run comparison UI, shared dashboards. Better for teams with multiple ML engineers collaborating on the same model.

W&B Artifacts: Versioned data and model artifacts with lineage tracking. Visualize the entire lineage graph: which dataset version + which code commit + which hyperparameters → this model version.

Practical choice: MLflow for self-hosted, privacy-sensitive environments (enterprise, regulated industries). W&B for research teams, computer vision/NLP teams that need rich media logging, teams that prioritize collaboration UI. Many teams use both: W&B for experiment tracking (better UI), MLflow for model registry (open-source, self-hosted, integrates with Spark/Databricks).

A model in the registry without complete metadata is a black box. When debugging a production regression, you need to answer: what changed? Without metadata, you can't.

Required metadata for every registered model version:

1. Training dataset reference: S3 path or BigQuery table + partition, with the dataset creation timestamp and git hash of the data processing code. 'Trained on s3://ml-data/fraud/training/2024-01-15/v3/' is sufficient. This enables reproducibility.

2. Feature schema: The exact list of features, their types, and expected value ranges. {user_age: float, [18, 100]}, {merchant_category: string, [food, retail, travel, ...]}, ...}. The MLflow model signature enforces this at serving time.

3. Offline evaluation metrics: AUC-PR, Recall@80%Precision, Calibration ECE, and performance on each evaluation slice (by merchant category, geography, transaction size). Stored as key-value pairs on the model version.

4. Code version: git commit hash of the training code. 'This model was produced by commit abc1234 of the fraud-model repository.' Reproducibility requires pointing back to the exact code.

5. Run ID: link back to the MLflow/W&B experiment run where all training curves, hyperparameters, and intermediate metrics are logged.

6. Promoter and timestamp: who promoted this model to Production and when. Audit trail for compliance.

7. Serving requirements: expected input format, latency budget, GPU/CPU requirement. 'This model expects a feature vector of dimension 128 at float32. It requires a GPU with > 4GB VRAM for < 10ms P99 latency at batch=1.'