ML Experiment Tracking skill for Codex

Done well means: any result in a report can be traced to one run that records the code commit, data version, config, environment and metrics; models are promoted by moving an alias after a recorded comparison; and tracking never blocks or crashes training.

Decide the tool

Situation	Default
Self-hosted, open source, model registry and serving hooks wanted	MLflow
Team already on W&B, rich dashboards, sweeps, reports	Weights & Biases
Both appear in the repo	Keep what AGENTS.md, requirements or CI already use
Network off (Codex sandbox, air-gapped)	MLflow with a local sqlite:/// or ./mlruns store, or WANDB_MODE=offline

Never put tokens (WANDB_API_KEY, MLflow credentials) in code or committed config; read them from the environment.

What every run must record

- [ ] git commit SHA and dirty flag
- [ ] full resolved config (hyperparameters, seeds, model architecture)
- [ ] dataset identifier + version or content hash, and split definition
- [ ] environment: python, key library versions, hardware
- [ ] metrics per step/epoch for train and validation, final test metrics once
- [ ] artifacts: best model, plots (confusion matrix, curves), eval predictions sample
- [ ] tags: purpose ("baseline", "ablation-lr"), owner, ticket

MLflow

import subprocess, mlflow

mlflow.set_tracking_uri("sqlite:///mlflow.db")       # or http://mlflow.internal:5000
mlflow.set_experiment("churn-model")

sha = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()
dirty = bool(subprocess.check_output(["git", "status", "--porcelain"], text=True).strip())

with mlflow.start_run(run_name="hgb-lr0.05") as run:
    mlflow.set_tags({"git_sha": sha, "git_dirty": str(dirty), "purpose": "baseline"})
    mlflow.log_params(cfg)                            # flat dict of scalars
    mlflow.log_param("data_version", data_hash)
    for epoch, m in enumerate(history):
        mlflow.log_metrics({"train_loss": m.train, "val_auc": m.val_auc}, step=epoch)
    mlflow.log_metric("test_auc", test_auc)
    mlflow.log_figure(fig, "plots/roc.png")
    mlflow.sklearn.log_model(pipe, name="model", input_example=X_val.head(5),
                             registered_model_name="churn-model")

• MLflow 3 uses name= for log_model; older versions use artifact_path=. Match the installed version.
• mlflow.autolog() covers sklearn, XGBoost, LightGBM, PyTorch Lightning and Transformers quickly; still add the git, data and purpose tags manually.
• Local UI: mlflow ui --backend-store-uri sqlite:///mlflow.db --port 5000.
• Team server: mlflow server --backend-store-uri postgresql://... --artifacts-destination s3://bucket/mlflow --host 0.0.0.0 behind an authenticating proxy; the open-source server has minimal built-in auth.

Registry and promotion

Stages (Staging/Production) are deprecated; use aliases.

from mlflow import MlflowClient
client = MlflowClient()
client.set_registered_model_alias("churn-model", "challenger", version=7)

# after comparison on the frozen holdout and user approval:
client.set_registered_model_alias("churn-model", "champion", version=7)

model = mlflow.pyfunc.load_model("models:/churn-model@champion")

Moving champion changes what production loads: confirm with the user and record the comparison run IDs in the version description before moving it. Rollback is moving the alias back to the previous version.

Weights & Biases

import wandb

run = wandb.init(project="churn-model", job_type="train", config=cfg,
                 tags=["baseline"], notes="HGB with target encoding")
for epoch, m in enumerate(history):
    run.log({"train/loss": m.train, "val/auc": m.val_auc, "epoch": epoch})
run.summary["test_auc"] = test_auc

data = wandb.Artifact("churn-data", type="dataset", metadata={"rows": len(df)})
data.add_file("data/train.parquet")
run.log_artifact(data)                     # versioned v0, v1, ... by content

model_art = wandb.Artifact("churn-model", type="model")
model_art.add_file("model.joblib")
run.log_artifact(model_art, aliases=["candidate"])
run.finish()

• W&B records the git commit automatically when run inside a repo; confirm in the run overview.
• WANDB_MODE=offline writes to ./wandb; sync later with wandb sync wandb/offline-run-* (needs network and approval). WANDB_MODE=disabled for tests.
• Use run.use_artifact("churn-data:v3") in training so lineage links data to model.
• Sweeps: wandb sweep sweep.yaml then wandb agent <id>; cap runs with count.

Framework hooks

Framework	Hook
Hugging Face Trainer / TRL	report_to="mlflow" or "wandb" in TrainingArguments
PyTorch Lightning	MLFlowLogger / WandbLogger
Keras	mlflow.autolog() or WandbMetricsLogger
Plain PyTorch	Explicit log_metrics every N steps (see $pytorch-training)

Comparing runs

runs = mlflow.search_runs(experiment_names=["churn-model"],
                          filter_string="tags.purpose = 'baseline' and metrics.val_auc > 0.8",
                          order_by=["metrics.val_auc DESC"])
print(runs[["run_id", "params.learning_rate", "metrics.val_auc", "tags.git_sha"]].head(10))

• Compare on the same split and data version; a filter on params.data_version avoids apples-to-oranges tables.
• Differences smaller than seed-to-seed variance are noise: run 3 seeds before claiming a win.
• Choose on validation, report test once for the chosen run.

Gotchas

• Logging inside the inner loop every step with .item() slows GPU training; log every N steps.
• Under DDP, only rank 0 should create the run and log, or you get N duplicate runs.
• log_params values are truncated strings; store large configs as an artifact (mlflow.log_dict(cfg, "config.json")).
• Tracking failures must not kill training: wrap remote logging in retries, or log offline and sync.
• Artifacts can contain customer data (prediction samples). Check where the artifact store lives and who can read it before logging raw rows.

Verification

• Open a run and reproduce it: check out the logged SHA, load the logged config and data version, retrain, and get the same metric within seed noise.
• models:/<name>@champion loads in a clean environment and predicts on input_example.

Report template

**Experiment:** churn-model (MLflow, sqlite store) | **Runs compared:** 14, data_version a3f9c1
**Winner:** run 8e21... hgb lr=0.05 depth=6: val AUC 0.871 +/- 0.004 (3 seeds), test AUC 0.866
**Registry:** churn-model v7 -> alias challenger; champion move awaiting approval
**Repro check:** retrained from logged SHA/config: val AUC 0.869