Experiment Tracking Setup

Chapter 1

Model Registry Workflow

How to version models, promote through stages, and roll back safely.

The Problem

After 50 training runs you have dozens of model files scattered across directories. When someone asks "which model is serving in production?" you grep through logs and hope for the best.

A model registry solves this by:

Assigning version numbers to every model artifact
Tracking which version is in each lifecycle stage
Maintaining an audit trail of every promotion/demotion
Making rollback a single command

Lifecycle Stages

  ┌──────┐     ┌─────────┐     ┌────────────┐     ┌──────────┐
  │ NONE │ ──▶ │ STAGING │ ──▶ │ PRODUCTION │ ──▶ │ ARCHIVED │
  └──────┘     └─────────┘     └────────────┘     └──────────┘
                    ▲                                    │
                    └────────────────────────────────────┘
                              (rollback path)

NONE: Freshly registered, not yet evaluated
STAGING: Under evaluation (shadow traffic, A/B test, manual review)
PRODUCTION: Actively serving predictions
ARCHIVED: Previously in production, kept for audit and rollback

Rule: Only ONE version can be PRODUCTION at a time per model name. Promoting a new version auto-archives the current one.

Step-by-Step Workflow

1. Register a model after training

python

from src.model_registry import ModelRegistry

registry = ModelRegistry(base_dir="./model_registry")

# Create the model container (once per model name)
registry.create_model("churn-predictor", description="Customer churn binary classifier")

# Register a version from a training run
v1 = registry.register_version(
    model_name="churn-predictor",
    run_id="abc123def456",
    metrics={"accuracy": 0.87, "f1": 0.82, "auc": 0.91},
    params={"model": "logistic_regression", "C": 1.0},
    artifact_path="./models/churn_lr_v1.pkl",
    description="Baseline: L2-regularized logistic regression",
)

2. Promote to staging for evaluation

python

from src.model_registry import ModelStage

registry.transition_stage(
    "churn-predictor", v1.version, ModelStage.STAGING,
    user="ml-engineer",
    reason="Passed offline evaluation, sending to shadow traffic test",
)

3. Promote to production

python

registry.transition_stage(
    "churn-predictor", v1.version, ModelStage.PRODUCTION,
    user="ml-lead",
    reason="Shadow traffic showed no regression vs current baseline",
)

4. Train a better model and replace

python

v2 = registry.register_version(
    model_name="churn-predictor",
    run_id="789ghi012jkl",
    metrics={"accuracy": 0.92, "f1": 0.89, "auc": 0.95},
    params={"model": "xgboost", "max_depth": 6, "lr": 0.01},
    artifact_path="./models/churn_xgb_v2.pkl",
    description="XGBoost with tuned hyperparameters, +5% accuracy",
)

# Stage it
registry.transition_stage(
    "churn-predictor", v2.version, ModelStage.STAGING,
    user="ml-engineer",
    reason="Outperformed v1 in offline eval by 5% accuracy",
)

# Promote to production — v1 is automatically archived
registry.transition_stage(
    "churn-predictor", v2.version, ModelStage.PRODUCTION,
    user="ml-lead",
    reason="A/B test: +4% conversion lift over v1",
)

5. Roll back if something goes wrong

python

# v2 is degrading in production — roll back to v1
rolled_back = registry.rollback(
    "churn-predictor", to_version=v1.version,
    user="on-call-engineer",
    reason="v2 showing 15% increase in false positives since Tuesday",
)

The rollback method handles the stage transitions automatically:

1. Archives the current production version (v2)

2. Moves the target version (v1) from ARCHIVED → STAGING → PRODUCTION

3. Records audit entries for every stage change

Querying the Registry

Find the current production model

python

prod = registry.get_production_version("churn-predictor")
if prod:
    print(f"Production: v{prod.version}, accuracy={prod.metrics['accuracy']}")
    print(f"Artifact: {prod.artifact_path}")

List all versions

python

for v in registry.list_versions("churn-predictor"):
    print(f"  v{v.version} [{v.stage}] — {v.description}")

Filter by stage

python

staging = registry.list_versions("churn-predictor", stage=ModelStage.STAGING)
archived = registry.list_versions("churn-predictor", stage=ModelStage.ARCHIVED)

Generate a Markdown report

python

report = registry.version_history_markdown("churn-predictor")
print(report)

Promotion Policies

Use these as guidelines for when to promote:

Gate	NONE → STAGING	STAGING → PRODUCTION
Offline metrics	Beat current prod on key metric	N/A (already passed)
Data validation	Training data passes quality checks	N/A
Shadow traffic	N/A	No regression vs current prod
A/B test	N/A	Statistically significant improvement
Latency check	N/A	p99 latency within SLA
Bias audit	N/A	No disparate impact on protected groups
Approval	ML engineer	ML lead or product owner

Integration with Experiment Tracking

Combine the registry with the experiment logger for a complete workflow:

python

from src.experiment_logger import ExperimentLogger
from src.model_registry import ModelRegistry, ModelStage

exp = ExperimentLogger(backend="stdlib")
registry = ModelRegistry()

# Train and track
with exp.run("churn-model", run_name="xgb-v3") as run_id:
    exp.log_params({"model": "xgboost", "max_depth": 8})
    # ... training code ...
    exp.log_metric("accuracy", 0.94)
    exp.log_metric("f1", 0.91)

# Register the model version using the run_id
v3 = registry.register_version(
    "churn-predictor",
    run_id=run_id,
    metrics={"accuracy": 0.94, "f1": 0.91},
    params={"model": "xgboost", "max_depth": 8},
    description="XGBoost v3, deeper trees",
)

# Now promote through stages as described above

Storage Layout

The local registry stores everything as JSON:

model_registry/
├── models/
│   ├── churn-predictor.json       # Model metadata
│   └── fraud-detector.json
└── versions/
    ├── churn-predictor/
    │   ├── v1.json                # Version metadata + audit trail
    │   ├── v2.json
    │   └── v3.json
    └── fraud-detector/
        └── v1.json

For production, migrate to MLflow Model Registry or a managed service (SageMaker, Vertex AI) — the workflow logic stays the same.

Chapter 2

Setup Guide

Get experiment tracking running in under 10 minutes.

Option 1: Local-Only (Zero Dependencies)

The stdlib tracker works immediately with no external services. Every run is saved as a JSON file.

bash

# No installation needed beyond Python 3.10+
python src/stdlib_tracker.py

This creates an experiments/ directory:

experiments/
├── runs/
│   └── abc123def456.json    # One JSON per run
├── artifacts/               # Model files, plots, etc.
└── index.json               # Experiment → run_id mapping

Using the ExperimentLogger wrapper

python

from src.experiment_logger import ExperimentLogger

exp = ExperimentLogger(backend="stdlib")

with exp.run("my-first-experiment", run_name="baseline") as run_id:
    exp.log_params({"model": "random_forest", "n_estimators": 100})

    for epoch in range(10):
        loss = 1.0 / (epoch + 1)
        exp.log_metric("loss", loss, step=epoch)

    exp.set_tag("status", "baseline")

print(f"Run saved: {run_id}")

Comparing runs

After running a few experiments:

python

from src.run_comparison import load_runs_from_directory, generate_leaderboard

runs = load_runs_from_directory("./experiments/runs")
lb = generate_leaderboard(runs, metric="loss", ascending=True)
print(lb.to_markdown())

Option 2: MLflow

Install

bash

pip install mlflow

Start a local tracking server

bash

mlflow server \
  --backend-store-uri sqlite:///mlflow.db \
  --default-artifact-root ./mlartifacts \
  --host 0.0.0.0 --port 5000

Configure the ExperimentLogger

python

import os
os.environ["MLFLOW_TRACKING_URI"] = "http://localhost:5000"

from src.experiment_logger import ExperimentLogger

exp = ExperimentLogger(backend="mlflow")
# or backend="auto" — will detect MLflow from the env var

Production MLflow

For team use, deploy MLflow on a VM or container:

1. Use PostgreSQL as the backend store (not SQLite)

2. Use S3/GCS/Azure Blob for artifact storage

3. Put a reverse proxy (nginx) in front for auth

4. See configs/mlflow_config.yaml for all settings

Option 3: Weights & Biases

Install and authenticate

bash

pip install wandb
wandb login
# Paste your API key from https://wandb.ai/authorize

Configure the ExperimentLogger

bash

export WANDB_API_KEY=YOUR_API_KEY_HERE
export WANDB_PROJECT=my-ml-project
export WANDB_ENTITY=acme-ml-team

python

from src.experiment_logger import ExperimentLogger

exp = ExperimentLogger(backend="wandb", project="my-ml-project")

W&B sweep example

bash

wandb sweep configs/wandb_config.yaml
wandb agent <sweep-id>

See configs/wandb_config.yaml for sweep configuration.

Switching Backends

The ExperimentLogger uses the same API regardless of backend. To migrate from stdlib to MLflow:

1. Keep your training code unchanged

2. Set MLFLOW_TRACKING_URI as an environment variable

3. Change backend="stdlib" to backend="mlflow" (or use "auto")

python

# This code works with ALL backends — no changes needed
exp = ExperimentLogger(backend="auto")

with exp.run("churn-model") as run_id:
    exp.log_params({"lr": 0.01})
    exp.log_metric("accuracy", 0.95)

Directory Structure

experiment-tracking-setup/
├── src/
│   ├── __init__.py              # Package init
│   ├── stdlib_tracker.py        # Zero-dependency JSON tracker
│   ├── experiment_logger.py     # Unified logging facade
│   ├── run_comparison.py        # Leaderboard & comparison tools
│   ├── model_registry.py        # Model version & stage management
│   └── artifact_manager.py      # Content-addressed artifact storage
├── configs/
│   ├── mlflow_config.yaml       # MLflow server configuration
│   ├── wandb_config.yaml        # W&B project configuration
│   └── tracking_config.yaml     # ExperimentLogger settings
├── tests/
│   ├── test_experiment_logger.py
│   └── test_run_comparison.py
└── guides/
    ├── setup-guide.md           # (this file)
    └── registry-workflow.md     # Model promotion workflow

Experiment Tracking Setup v1.0.0 — Free Preview

Contents

Model Registry Workflow

The Problem

Lifecycle Stages

Step-by-Step Workflow

1. Register a model after training

2. Promote to staging for evaluation

3. Promote to production

4. Train a better model and replace

5. Roll back if something goes wrong

Querying the Registry

Find the current production model

List all versions

Filter by stage

Generate a Markdown report

Promotion Policies

Integration with Experiment Tracking

Storage Layout

Setup Guide

Option 1: Local-Only (Zero Dependencies)

Using the ExperimentLogger wrapper

Comparing runs

Option 2: MLflow

Install

Start a local tracking server

Configure the ExperimentLogger

Production MLflow

Option 3: Weights & Biases

Install and authenticate

Configure the ExperimentLogger

W&B sweep example

Switching Backends

Directory Structure