Experiment tracking and the model registry

The question “which run produced the model currently in production?” becomes unanswerable within about two weeks of not tracking. This is the cheapest discipline in ML and the most commonly skipped.

What to track

Every run, from the first experiment — not once things “get serious”, because you never notice the transition.

Data version and code version are the two that make a run reproducible. Metrics without them tell you a number was achieved but not how to achieve it again.

MLflow

The safe default answer: most portable, works across scikit-learn, PyTorch, XGBoost and HuggingFace, and recent versions added prompt logging, evaluation and tracing for GenAI workloads — so one tool covers both classical and LLM work.

import mlflow

mlflow.set_experiment("churn-prediction")

with mlflow.start_run(run_name="lgbm-baseline"):
    mlflow.log_params({"max_depth": 6, "learning_rate": 0.05})
    mlflow.set_tags({"data_version": DATA_HASH, "git_sha": GIT_SHA})

    model.fit(X_train, y_train)

    mlflow.log_metric("val_ap", average_precision_score(y_val, proba))
    for segment, idx in segments.items():                    # per-segment, not just aggregate
        mlflow.log_metric(f"val_ap_{segment}", average_precision_score(y[idx], proba[idx]))

    mlflow.sklearn.log_model(model, "model", signature=signature)

The signature — the expected input schema — is worth logging. It’s how you catch a serving-time schema mismatch at load rather than in production.

Alternatives: Weights & Biases (stronger UI and collaboration), Neptune, Comet, or the cloud-native ones (SageMaker, Vertex, Azure ML). MLflow is the portable default.

The registry

Tracking records experiments; the registry manages which model is where.

mlflow.register_model("runs:/<run_id>/model", "churn-predictor")
client.set_registered_model_alias("churn-predictor", "champion", version=7)

Aliases (champion, challenger) beat stage names, because they let you swap what production points at without touching the serving code.

What the registry gives you:

• Lineage — this model came from that run, that data, that commit.
• A promotion path — dev to staging to production, ideally gated by tests.
• Instant rollback — repoint the alias at the previous version.
• An audit trail — who promoted what, when. Increasingly a compliance requirement. See ../14_guardrails_safety/03_pii_privacy_and_compliance.md.

Serve by alias, never by version number. models:/churn-predictor@champion means rollback is a registry operation, not a deployment.

Reproducibility

# Seed everything
random.seed(42); np.random.seed(42); torch.manual_seed(42)

Full determinism on GPU is often unachievable — some kernels are non-deterministic by design. So:

• Report a mean over several seeds when comparing architectures. Run-to-run variance frequently exceeds the difference between two configurations.
• Pin library versions; a minor version bump can change results.
• Snapshot or hash the data.

The LLM variant

There’s no trained model, so the versioned artefacts change:

with mlflow.start_run():
    mlflow.log_params({
        "prompt_version": "v7",
        "model": "pinned-model-id",       # never an unpinned alias
        "retrieval_k": 5,
        "index_version": INDEX_VERSION,
    })
    results = run_eval_suite(eval_set_v3)
    mlflow.log_metrics({
        "faithfulness": results.faithfulness,
        "format_compliance": results.schema_pass_rate,
        "cost_per_query": results.mean_cost,
        "p95_latency_ms": results.p95,
    })

Prompts are code. They belong in version control, reviewed in pull requests, with the eval suite as their test. Storing prompts in a database where anyone can edit them live is how quality changes with no deployment and no record.

Cost and latency belong in the tracked metrics alongside quality — they’re part of whether a change is an improvement.

Interview angle

• “What do you track for an ML experiment?” — parameters, metrics including per-segment breakdowns, artifacts, and critically the data version and git commit. Without those last two a run is a number you can’t reproduce.
• “Which tool?” — MLflow as the portable default; it spans classical and GenAI workloads and recent versions handle prompt logging, evaluation and tracing. W&B if you want stronger collaboration UI.
• “What does a model registry give you beyond storage?” — lineage back to run, data and commit; a gated promotion path; instant rollback by repointing an alias; and an audit trail of who promoted what. Serve by alias so rollback isn’t a deploy.
• “How does this change for an LLM application?” — the versioned artefact is the prompt plus the pinned model ID, retrieval config and index version. Prompts are code: version-controlled, reviewed, with the eval suite as their tests.
• “Two training runs give different results with the same config. Bug?” — usually not. GPU non-determinism and seed variance are normal. Report a mean over several seeds, because run-to-run variance often exceeds the difference you’re trying to measure.
• “Why not store prompts in a database for live editing?” — you lose review, versioning and the link between a quality change and a deliberate action. Quality then shifts with no deployment and no record of why.