ml-experiment-tracking

ML Experiment Tracking

When to invoke

Use when instrumenting a training script for the first time, standardizing experiment tracking across a team, or deciding when a candidate model is genuinely better than the current production version.

Why tracking matters (in two sentences)

An experiment you can't reproduce is a result you can't trust. Tracking creates the immutable record — code version, data version, hyperparameters, environment, metrics — that makes "we improved accuracy by 4%" a verifiable claim rather than a notebook screenshot.

Step 1 — What to log (minimum set)

Category	Examples	Why
Git commit SHA	mlflow.log_param("git_sha", get_git_sha())	Pin the code version to the run
Dataset version	S3 URI + ETag, DVC hash, or dataset fingerprint	Pin the data; different data = different experiment
Hyperparameters	All of them — not just the ones you varied	Future you needs to reproduce the run
Environment	Python version, key package versions	Dependency drift changes results
Metrics (per epoch)	Train loss, val loss, val metric	Enables learning curve comparison, not just final score
Artifacts	Model checkpoint, confusion matrix, feature importances	Complete provenance

MLflow instrumentation template:

import mlflow

mlflow.set_experiment("churn-prediction-v2")

with mlflow.start_run(run_name=f"lgbm-{trial.number}"):
    mlflow.log_params({
        "git_sha": get_git_sha(),
        "dataset_version": dataset.version,
        "model_type": "lightgbm",
        "n_estimators": params["n_estimators"],
        "learning_rate": params["learning_rate"],
    })

    model = train(params, X_train, y_train)

    mlflow.log_metrics({
        "val_auc": val_auc,
        "val_pr_auc": val_pr_auc,
        "train_auc": train_auc,
    })

    mlflow.sklearn.log_model(model, "model",
        signature=mlflow.models.infer_signature(X_val, model.predict(X_val)),
        input_example=X_val.iloc[:3])

Step 2 — Run comparison workflow

When comparing two runs, follow this order to avoid cherry-picking:

1. Define the evaluation metric before comparing — choose the primary metric (e.g., val AUC) and the threshold (e.g., 0.5% improvement = meaningful) before looking at results.
2. Compare on the held-out test set only once — use val for tuning; the test set is touched once per candidate, not once per run.
3. Check for overfitting: train_metric >> val_metric is a warning sign even if val_metric is good.
4. Compare learning curves: a model that converges earlier with the same final val metric is strictly better.
5. Run the applied-statistics significance check (bootstrap or paired t-test) before declaring the new model "better."

Step 3 — Model registry promotion workflow

[Experiment runs] → [Staging registry] → [Champion vs Challenger test] → [Production registry]

MLflow model lifecycle:

client = mlflow.tracking.MlflowClient()

# Register from a run
result = mlflow.register_model(
    model_uri=f"runs:/{run_id}/model",
    name="churn-prediction"
)

# Promote to staging after validation
client.transition_model_version_stage(
    name="churn-prediction",
    version=result.version,
    stage="Staging",
    archive_existing_versions=False
)

# Promote to production only after:
# 1. Shadow/canary traffic test passes
# 2. Statistical significance check passes
client.transition_model_version_stage(
    name="churn-prediction",
    version=result.version,
    stage="Production",
    archive_existing_versions=True  # demote old version to Archived
)

Step 4 — Team conventions to standardize

Convention	Recommendation
Experiment naming	{team}-{problem}-v{n} e.g. risk-churn-v3
Run naming	Descriptive: lgbm-lr0.01-depth6 not run_42
Tags	Always tag env (dev/ci/prod) and triggered_by (manual/ci)
Test-set evaluation	One function, called once per candidate; no loops over test set
Model signature	Always log an input_example — serves as the schema contract for the serving endpoint

Step 5 — CI integration

On every PR that touches training/ or features/:

- name: Run training smoke test
  run: |
    python train.py \
      --dataset tests/fixtures/sample_dataset.parquet \
      --max-rows 1000 \
      --experiment ci-smoke \
      --run-name "pr-${{ github.sha }}"

- name: Assert min metric threshold
  run: |
    python scripts/assert_min_metric.py \
      --experiment ci-smoke \
      --metric val_auc \
      --min-value 0.70

The CI run uses a sample dataset; the full training run is triggered separately (not on every PR).

Pitfalls

• Logging only the final metric, not per-epoch metrics — you can't diagnose overfitting or compare convergence speed without the training curve.
• Sharing one experiment across all team members' runs — use per-developer experiment namespaces; shared experiments become a wall of unattributed runs.
• Not logging the dataset version — changing a feature, fixing a bug in preprocessing, and retraining without bumping the dataset version makes runs incomparable.
• Evaluating on the test set repeatedly — each evaluation on the test set is an implicit optimization step; the test set is a one-shot oracle, not a validation loop.
• Skipping the model signature — a model without a signature has no enforced input schema; a serving endpoint receiving mismatched features fails silently or produces garbage predictions.