Experiment Tracking Tools

Three weeks ago one run hit ROC-AUC 0.917. The notebook scrolled out of memory, a Slack message says the best one was Tuesday-ish, and the model file in the shared drive has no name worth trusting. You need the configuration, the code, and the weights - and none of the three is recoverable. That is the problem experiment tracking solves.

MLflow's model is simple: experiments contain runs, and a run is one record with four slots. Params holds hyperparameters, one scalar per key. Metrics holds scores, logged as a series so per-epoch values become curves. Tags holds searchable key-value pairs: git SHA, owner, status. Artifacts holds files: model binaries, plots, feature lists, prediction samples.

import os
import mlflow, mlflow.sklearn
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import cross_val_score

mlflow.set_experiment("fraud-risk")
with mlflow.start_run(run_name="gbm-d3-lr0.1-data-v7") as run:
    params = {"n_estimators": 200, "max_depth": 3, "learning_rate": 0.1}
    clf = GradientBoostingClassifier(**params)
    scores = cross_val_score(clf, X, y, cv=5, scoring="roc_auc")

    mlflow.log_params(params)
    mlflow.log_metric("roc_auc_mean", scores.mean())
    mlflow.log_metric("roc_auc_std", scores.std())
    mlflow.log_tag("git_sha", os.popen("git rev-parse HEAD").read().strip())
    mlflow.log_dict({"features": feature_cols}, "feature_list.json")
    mlflow.log_artifact("confusion_matrix.png")
    mlflow.sklearn.log_model(clf, artifact_path="model")
    print(run.info.run_id)          # the id binding all of the above together

Out of the box mlflow ui serves the local ./mlruns directory (files plus SQLite) - enough for one engineer on one machine. Teams point the backend at PostgreSQL and the artifact store at S3 or GCS so runs survive laptops and are visible to everyone. The recovery story then plays out end to end: open the UI, filter to the experiment, sort by roc_auc_mean, click the top run, and read the params, the git SHA, and the feature list; download the model artifact and retrain at that SHA with that config. The run id is what binds the weights to the configuration that produced them - a model file without one is an anonymous matrix.

Autologging (mlflow.sklearn.autolog(), mlflow.pytorch.autolog()) captures params, metrics, and models without instrumenting each call, which is the right default in a notebook. It is incomplete for real projects: autologging cannot know your business metric or your dataset version, so log those explicitly. For command-line reproduction, an MLproject file pins the entry point and parameters so mlflow run . -P max_depth=5 reruns the experiment as a command rather than as archaeology.

Two habits make the stored runs actually comparable. First, log every metric with the same name and unit across runs - roc_auc_mean computed by the same 5-fold protocol every time; a run that logged accuracy while its neighbors logged roc_auc cannot appear in the same comparison view, and silently mixing protocols makes the comparison a lie even when the names match. Second, use nested or child runs for the outer structure: one parent run per configuration with a child per fold keeps the fold-level scores under the run they belong to, and the parent carries the aggregated mean and standard deviation. The UI then shows both the spread and the summary without you re-deriving anything.

Teams usually hit one scaling question: where do the artifacts live when the laptop closes. MLflow separates the tracking backend (an MLflow server holding metadata in SQL) from the artifact store (S3, Azure Blob, GCS, or a shared filesystem), so a team of ten points at one server and one bucket, and every run anyone logs is visible to everyone within seconds - the single biggest upgrade over emailing notebooks. Runs can also carry status and lifecycle tags (FINISHED, FAILED, candidate), and filtering to successful candidates before a model-selection meeting turns three weeks of scattered output into a ten-row table you can defend.

MLflow's record-and-compare is table-shaped and retrospective. Weights & Biases is built around the moment training is running: streaming charts, auto-logged system metrics, and team-wide comparison without anyone exporting CSVs.

import wandb
wandb.init(project="fraud-risk", config={"n_estimators": 200, "max_depth": 3,
                                         "learning_rate": 0.1, "data": "v7"},
           tags=["gbm", "baseline"], notes="after feature cleanup, no threshold tuning")

for epoch in range(10):
    train_one_epoch()
    wandb.log({"train_loss": loss, "val_auc": auc, "lr": lr}, step=epoch)

wandb.log({"confusion": wandb.plot.confusion_matrix(
    preds=probs >= 0.4, y_true=y_val, class_names=["legit", "fraud"])})
wandb.finish()

While the run executes you see the loss, the learning rate, GPU memory and utilization, and any custom series you log. That is a debugging tool, not just a record: divergence you would have discovered at the end shows up at step 400, and a dead learning rate is visible the moment it stops decaying. Hyperparameter search becomes a first-class object rather than a folder of scripts:

sweep_config = {
    "method": "bayes",
    "metric": {"name": "val_auc", "goal": "maximize"},
    "parameters": {
        "max_depth": {"values": [2, 3, 5, 8]},
        "learning_rate": {"distribution": "log_uniform_values", "min": 0.01, "max": 0.3},
    },
}
sweep_id = wandb.sweep(sweep_config, project="fraud-risk")
wandb.agent(sweep_id, function=train_and_log, count=20)

Each trial lands as a child run under one sweep, and the parallel-coordinates view shows which knob actually moved the metric - direct preparation for the search strategies in Grid and Random Search. Two details decide whether sweeps produce answers or noise. The metric you maximize must be computed on a validation fold never touched by the trial's own training, or the optimizer chases overfit; and the search space should encode what you learned from Bias-Variance and Error Analysis - if the capacity curve says the model is on the right arm of the U, narrow max_depth downward rather than re-searching the whole range every quarter.

Weights & Biases also auto-logs system metrics - GPU memory, utilization, CPU, disk, network - which answers a class of questions that have nothing to do with the model. A run that OOMs at epoch 40 shows memory climbing against the limit; a run that trains at half the expected throughput shows a starved GPU; two runs differing only in batch size reveal which one actually used the hardware. For teams without an ML engineer who lives in nvidia-smi, these charts are the difference between the training is slow and the DataLoader is the bottleneck.

Data lineage closes the loop that MLflow's artifacts open. Log a dataset as a versioned artifact and every run records which version it consumed:

wandb.log_artifact("data/processed/v7", name="training-set", type="dataset")

Choosing between the tools is a team question, not a technology question:

Many organizations run both: W&B for the experiment experience, MLflow for the registry that gates deployment. Neither rescues sloppy habits. Four hundred unnamed runs with empty notes is a dashboard-shaped mess; name runs by configuration (gbm_d3_lr010_datav7), tag status (baseline, candidate, rejected-bad-labels), and write a one-line hypothesis or conclusion in notes when the run ends. The note is what you will search for in six weeks; the chart will not remind you what you were testing.

A run with metrics but no provenance is a rumor. The difference between an experiment record and folklore is four columns, and each one exists because of a specific way runs become unreproducible.

import hashlib, platform, random, subprocess

def run_fingerprint(data_path: str, seed: int = 42) -> dict:
    random.seed(seed)
    with open(data_path, "rb") as fh:
        data_sha = hashlib.sha256(fh.read()).hexdigest()[:12]
    sha = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode().strip()[:12]
    dirty = bool(subprocess.check_output(["git", "status", "--porcelain"]).decode().strip())
    return {"git_sha": sha, "dirty": dirty, "data_sha": data_sha,
            "python": platform.python_version(), "seed": seed}

The data column is the one teams skip, and it is the most common source of impossible results: two runs with identical params and different scores are usually different data - a join order changed, a filter was fixed, a snapshot was refreshed. Hashing the input makes the difference visible in one glance, and logging the row count plus max event date is a cheap human-readable complement. The dirty flag deserves the same severity: a run recorded from modified, uncommitted code has no identity at all, and CI should refuse to log it rather than let it masquerade as a release candidate.

Seeds close the loop without guaranteeing perfection. GPU reductions and DataLoader shuffling reorder floating-point operations, so bitwise reproducibility needs torch.use_deterministic_algorithms(True) and costs throughput. The professional habit is to state the noise floor explicitly - reruns vary by +/- 0.003 - so a 0.3-point difference between two configs is understood as nothing.

The habit that ties all four columns together is a naming and tagging scheme that makes queries possible: data-v7 as a tag, candidate and baseline as status tags, one notes line per run. That turns a thousand runs into a searchable experiment history, and it is the raw material from which the reviewer-facing artifact is built in Model Cards and Documentation. The curves and gaps you decided what to log in the first place come from Bias-Variance and Error Analysis: training and validation series per epoch are exactly the metrics worth streaming into the run record.

A realistic failure makes the four columns concrete. Two engineers report conflicting AUCs for the same config: 0.917 and 0.884. With all four columns logged, the diff takes minutes - same git SHA, same seed, same environment; the data hashes differ (one snapshot predates the deduplication fix), and the lower number came from the pre-dedup file. Without the columns, that discrepancy becomes a two-day investigation conducted in meetings, or worse, an unexplained regression that gets shipped to production as a rollback. The value of provenance is not felt when runs agree; it is felt the week they disagree.

A practical CI pattern: the training job computes run_fingerprint() at startup, refuses to run on a dirty tree in release mode, attaches the fingerprint as tags, and uploads the frozen requirements file as an artifact. Add a post-training assertion that the logged data hash matches the hash of the file the job actually read - pipelines get reordered, caches get stale, and a check that costs one line keeps a stale snapshot from masquerading as fresh data. These habits cost about twenty lines of code per project and are the difference between an experiment history and an experiment rumor.