Experiment tracking (MLflow, W&B)

An Experiment is a named container that groups related Runs, typically for one modeling problem or initiative. A Run is a single execution of model code that records params, metrics, tags, and artifacts. Runs always belong to exactly one Experiment, and the Experiment provides a namespace for comparing Runs side by side.

log_param records a hyperparameter or config value that is fixed for the run and stored as a string; it cannot be updated to a different value after being set. log_metric records a numeric value that can change over time and supports a step argument, so MLflow stores a time series you can plot. Using metrics for evolving values (loss per epoch) and params for static config keeps the UI and comparisons correct.

mlflow.log_param("lr", 0.001)
for epoch in range(n):
    mlflow.log_metric("loss", loss, step=epoch)

autolog automatically captures params, metrics, model artifacts, and signatures for supported libraries (sklearn, XGBoost, PyTorch Lightning, etc.) without explicit log calls. A common pitfall is that it can silently log nothing or partial data if the library version is unsupported, or it can conflict with manual logging by duplicating or overriding values. It must be called before model training begins.

import mlflow
mlflow.autolog()
model.fit(X_train, y_train)  # params/metrics logged automatically

The Model Registry is a centralized store for versioned models with lineage back to the producing run. Each registered model has incrementing versions, and (in the legacy stage API) versions can be assigned stages like None, Staging, Production, and Archived. Newer MLflow recommends aliases and tags instead of fixed stages for more flexible promotion workflows.

Stages were a fixed enum (Staging/Production/Archived) that allowed only one version per stage and didn't model real-world workflows like champion/challenger or multi-environment routing. Aliases are mutable, named references (e.g., @champion, @prod-eu) that you can point at any version, enabling flexible, deterministic loading via models:/name@alias. Aliases plus tags decouple promotion semantics from a rigid lifecycle.

client.set_registered_model_alias("fraud", "champion", version=7)
model = mlflow.pyfunc.load_model("models:/fraud@champion")

A flavor is a convention describing how a saved model can be loaded and used by a particular tool (e.g., sklearn, pytorch, onnx). A single MLflow model can be saved with multiple flavors. The python_function (pyfunc) flavor is a universal interface so any tool can load and serve the model with a generic predict() regardless of the underlying framework.

A signature declares the expected input and output schema (column names, types, tensor shapes) of a model. At serving/inference time MLflow uses it to validate and coerce incoming data, catching schema mismatches early and preventing silent wrong predictions. It is auto-inferred via infer_signature or set explicitly when logging the model.

from mlflow.models import infer_signature
sig = infer_signature(X_train, model.predict(X_train))
mlflow.sklearn.log_model(model, "model", signature=sig)

MLflow Tracking uses a tracking server (REST/UI), a backend store for metadata like params/metrics/tags (a database or file store), an artifact store for large files like models and plots (S3, GCS, Azure Blob, local), and the client API. Separating the backend store (DB) from the artifact store (object storage) is the recommended production pattern for scalability.

The backend store (a relational DB) is optimized for fast queries over structured metadata like metrics and params, while artifacts are large binary blobs better suited to scalable object storage like S3. Mixing them (e.g., file-based backend) doesn't scale for concurrent users and large models. Separation also lets clients upload/download artifacts directly to/from object storage, reducing load on the tracking server.

A Project is a workspace that groups related Runs (and reports/artifacts). A Run is a single training/eval execution logging config and metrics via wandb.log. A Sweep is a hyperparameter search that orchestrates many Runs within a project according to a search strategy (grid, random, Bayesian).

W&B sweeps support grid, random, and bayes (Bayesian optimization) methods. Grid is exhaustive but explodes combinatorially; random samples efficiently in high dimensions. Bayesian builds a surrogate model of the objective to focus sampling on promising regions, which is most valuable when each run is expensive and the search space is moderate.

method: bayes
metric:
  name: val_loss
  goal: minimize
parameters:
  lr:
    distribution: log_uniform_values
    min: 1e-5
    max: 1e-1

Hyperband (configured under early_terminate in a sweep) monitors a metric at bracket milestones and kills underperforming runs early so compute is reallocated to promising configurations. This dramatically improves search efficiency by not wasting full budgets on clearly losing trials. You specify min_iter, eta, and the metric to monitor.

early_terminate:
  type: hyperband
  min_iter: 3
  eta: 2

Artifacts are versioned, deduplicated objects (datasets, models, checkpoints) stored with content-addressed hashing. When a run logs (output) or uses (input) an artifact, W&B records the dependency, building a lineage DAG linking data to models to evaluations. Versions are immutable and can carry aliases like 'latest' or 'production' for stable references.

art = wandb.Artifact("train-data", type="dataset")
art.add_file("data.csv")
run.log_artifact(art)
# downstream run
used = run.use_artifact("train-data:latest")

By default initialize and log only from the rank-0 (main) process so a single run captures aggregated metrics. Alternatively, use grouped runs where each process is its own run sharing a group name, which W&B aggregates in the UI. The key is to avoid multiple processes writing to the same run object concurrently, which corrupts step ordering.

if rank == 0:
    wandb.init(project="ddp")
# or grouped:
wandb.init(group="experiment-A", name=f"rank-{rank}")

step defines the x-axis position of a metric point, creating a monotonic time series for plotting. In W&B, logged steps must be non-decreasing per run; logging a smaller step than the current one is dropped or causes warnings. To log multiple metrics against a custom axis, use define_metric (W&B) or pass an explicit step in MLflow's log_metric.

wandb.define_metric("val_loss", step_metric="epoch")
wandb.log({"epoch": e, "val_loss": vl})

Calling mlflow.start_run(nested=True) inside an active run creates a child run linked to its parent. This is ideal for hyperparameter tuning, cross-validation folds, or pipeline stages where you want a parent run for the overall job and child runs for each trial or fold. The UI lets you expand the parent to compare children.

with mlflow.start_run() as parent:
    for params in grid:
        with mlflow.start_run(nested=True):
            mlflow.log_params(params)

MLFLOW_TRACKING_URI (or mlflow.set_tracking_uri) sets where runs are logged. Valid targets include a local ./mlruns directory, a file path, a SQLAlchemy DB URI (e.g., sqlite/postgres) for the backend store, or an http(s) URL pointing at a remote tracking server. Databricks workspaces use the special 'databricks' scheme.

mlflow.set_tracking_uri("http://mlflow.internal:5000")

Use the context manager (with mlflow.start_run():) so MLflow sets the run status to FINISHED on success and FAILED on exception automatically. If you manage runs manually, wrap in try/except and call mlflow.end_run(status='FAILED'). Leftover RUNNING runs usually mean the process was killed (SIGKILL/OOM) without cleanup.

try:
    with mlflow.start_run():
        train()
except Exception:
    mlflow.end_run(status="FAILED")
    raise

Use wandb.Table for structured tabular data and rich media columns, and wandb.plot helpers (or wandb.Image) for visualizations; these become interactive in the dashboard. Tables support querying, filtering, and joining across runs, and can hold images, audio, or molecule objects per row. This is preferred over logging static PNGs because it stays interactive and queryable.

tbl = wandb.Table(columns=["img","pred","label"])
tbl.add_data(wandb.Image(x), pred, label)
wandb.log({"samples": tbl})

With must, wandb.init requires the specified run id to already exist and errors otherwise—used when you know you're continuing. With allow, it resumes if the id exists or creates a new run if not. With never, it always starts fresh and errors if the id already exists. Correct resume handling preserves step continuity for long or preempted jobs.

wandb.init(id=run_id, resume="allow")

Use the runs table with column filters, sorting by metric, and grouping by config values, plus parallel-coordinates plots to see hyperparameter-to-metric relationships at scale. Tag or name runs systematically and use search queries to slice cohorts. Both MLflow (search_runs with filter strings) and W&B (filters, parallel coords, parameter importance) support programmatic and UI-based comparison.

runs = mlflow.search_runs(
  experiment_ids=[id],
  filter_string="metrics.f1 > 0.9 and params.model = 'xgb'",
  order_by=["metrics.f1 DESC"])

mlflow.evaluate runs a standardized evaluation on a model and dataset, automatically computing task-appropriate metrics (e.g., accuracy/F1 for classifiers, RMSE/MAE for regressors) and producing diagnostic artifacts like confusion matrices, SHAP explanations, and per-example results. It can also enforce promotion gates by passing validation_thresholds (MetricThreshold objects), optionally compared against a baseline_model. This standardizes evaluation and logs everything to the run.

from mlflow.models import MetricThreshold
mlflow.evaluate(
    model_uri, data, targets="y", model_type="classifier",
    validation_thresholds={
        "accuracy_score": MetricThreshold(threshold=0.9, greater_is_better=True)
    })

Capturing exact dependency versions ensures the model can be reloaded and served reproducibly, avoiding 'works on my machine' drift. MLflow saves a conda.yaml, python_env.yaml, and requirements.txt alongside the logged model so the runtime can recreate the environment. This is essential for reproducibility and for serving the model in a clean container later.

MLflow auto-records the git commit hash and source file as run tags (mlflow.source.git.commit) when run from a repo, linking results to code state. W&B can save the actual code (save_code=True) and capture the git state and diff, storing an uncommitted-changes patch. Both aim to tie metrics to exact code, but W&B can persist the literal source and diff, not just the commit reference.

wandb.init(project="p", save_code=True)

They capture nested traces of LLM/agent executions—prompts, responses, tool calls, latencies, token counts, and costs—so you can debug chains and compare prompt/model variants. MLflow Tracing (mlflow.openai.autolog, @mlflow.trace) and W&B Weave log spans and support evaluation over datasets with LLM-as-judge scorers. This extends classic metric tracking to non-deterministic generative pipelines.

import mlflow
mlflow.openai.autolog()

@mlflow.trace
def agent(q): ...  # spans captured automatically

Log a content hash or version identifier of the data (and a pointer/URI) rather than the raw file inline, and use a versioned artifact system so runs reference an immutable dataset version. MLflow offers mlflow.data with dataset logging (log_input) that records schema, source, and a digest; W&B uses versioned dataset Artifacts. This guarantees you can map any metric back to the exact data it was computed on.

ds = mlflow.data.from_pandas(df, source="s3://bucket/v3.parquet")
mlflow.log_input(ds, context="training")

The W&B Registry lets you promote specific model Artifact versions into a registered, named collection with aliases (e.g., production, staging) and protected lifecycle controls. Because models are artifacts, the registry inherits full lineage: you can trace a registered version back to the run, code, and dataset that produced it. This provides governance and a single source of truth for deployment-candidate models.

An MLproject file declares entry points, parameters, and the execution environment (conda/docker/virtualenv) so a run can be launched consistently with mlflow run, locally or remotely (e.g., on Databricks or Kubernetes). It packages how to run the code, complementing tracking which records what happened. This makes experiments portable and re-runnable from a git URI.

mlflow run https://github.com/org/repo -P lr=0.01 --experiment-name exp1

Frequent synchronous logging adds network/IO latency and can balloon stored time-series, slowing both training and the UI. Mitigate by batching logs, reducing logging frequency (e.g., log every N steps or aggregate per epoch), and relying on the client's async/background upload (W&B uploads in a separate process; MLflow can batch via log_metrics/log_batch). Sample or downsample noisy per-step metrics for storage efficiency.

mlflow.log_batch(run_id, metrics=[Metric("loss", v, ts, step) for ...])

Params capture inputs that define the experiment's configuration (learning rate, model type), while tags are mutable metadata for organization, filtering, and governance (e.g., team='nlp', ticket='JIRA-123', mlflow.user, run_purpose='smoke-test'). Tags can be added or changed after a run completes; params are typically fixed. Using tags consistently enables powerful search and grouping without polluting the param space.

mlflow.set_tag("run_purpose", "smoke-test")

For W&B, set WANDB_MODE=offline (or wandb.init(mode='offline')) to write runs to disk, then sync later with 'wandb sync'. MLflow naturally supports a local backend (file or local DB) with a local artifact path, requiring no external service. Both let you collect runs locally and later push to a central server when connectivity is available.

WANDB_MODE=offline python train.py
wandb sync ./wandb/offline-run-*

You need authentication, role-based access control over projects/experiments and the model registry, audit logs of who promoted or deleted models, and secrets management so credentials aren't logged as params. Sensitive data (PII) must be kept out of logged artifacts/metrics. Managed offerings (W&B, Databricks MLflow) provide SSO, RBAC, and team scoping; self-hosted MLflow needs an auth proxy or its built-in basic-auth plugin.

The tracker faithfully records what happened—the discrepancy almost always stems from non-determinism, not the logger: GPU/cuDNN nondeterministic kernels, unset library seeds, data-loader worker shuffling, mixed precision, or different hardware/library versions. The fix is to enforce determinism (seed all RNGs, torch.use_deterministic_algorithms, fixed cuDNN flags) and log the environment so divergence is explainable. Tracking surfaces the difference; it doesn't cause it.

torch.manual_seed(s); np.random.seed(s)
torch.use_deterministic_algorithms(True)

Pass secrets via environment variables or a secrets manager and never log them as params, tags, or in the config dict that autolog/config capture serializes. Use W&B's WANDB_API_KEY env (kept out of run config) and avoid putting tokens in argparse namespaces that get auto-logged. Review autolog and config-capture behavior, and scrub or exclude sensitive keys before logging.