Experiment Tracking with MLflow

Experiment tracking with MLflow gives an ML team a searchable record of what was tried, why it was tried, and which fitted artifact came out of it. In this lesson, the outcome is practical: you will know how MLflow records runs, how parameters, metrics, tags, artifacts, experiments, and model registry entries relate to each other, and how to troubleshoot common tracking mistakes before they corrupt a comparison.

In the MLOps course, this lesson sits in the experiments and artifacts section because tracking is the bridge between local model development and controlled promotion. A model should not reach deployment because someone remembers that a notebook cell looked good. It should reach deployment because a specific run points to its code revision, data snapshot, configuration, metrics, and artifacts.

What MLflow Tracking Stores

MLflow Tracking organizes work around an experiment and a run. An experiment is a named container, such as fraud-risk-baseline. A run is one execution inside that container. During a run, client code writes structured records to a tracking server or local tracking store. The most common records are params, metrics, tags, and artifacts.

A parameter is an input choice: learning rate, feature set name, random seed, embedding model, or training window. Parameters are stored as strings and are intended to be stable for a run. A metric is a numeric observation: validation AUC, loss at step 20, latency, precision, or calibration error. Metrics can be logged repeatedly with steps and timestamps, which is why a loss curve can be reconstructed. A tag is metadata used for search, ownership, lineage, and review state. An artifact is a file or directory: model binaries, plots, feature importance files, evaluation reports, schemas, or serialized preprocessors.

Internally, the MLflow client sends tracking operations to a backend store and an artifact store. The backend store holds entities and metadata, often in a local file layout during development and in a database for shared use. The artifact store holds larger files, often on local disk, object storage, or another remote store. This split matters: a run can appear in the UI while its artifact upload failed, and a registry entry can point to a model source that later becomes unreadable if storage permissions or retention rules are wrong.

API Anatomy

The usual API shape is small but important. mlflow.set_tracking_uri() chooses the tracking backend. mlflow.set_experiment() selects or creates the experiment. mlflow.start_run() opens a run context, and logging calls inside that context attach evidence to that run. mlflow.log_param(), mlflow.log_metric(), mlflow.set_tag(), and mlflow.log_artifact() write the main evidence types. MlflowClient is the lower-level API for searching runs, reading model versions, and automating registry workflows.

The design choice is whether the tracking store is personal, team-shared, or production-governed. A local file store is fast for a tutorial and works without infrastructure, but it is poor for collaboration. A shared tracking server centralizes comparisons and registry workflows, but it needs authentication, backups, storage policy, and conventions for experiment naming. Object storage scales artifacts better than a database filesystem, but every training and deployment identity must be able to read the exact artifact paths it needs.

Example 1: Record One Run

This first example logs one run to a temporary local tracking store. It records the input choices, final metric, a lineage tag, and a small artifact. The deterministic output confirms that an artifact was written and that the run completed.

import tempfile
from pathlib import Path

import mlflow

tracking_dir = tempfile.mkdtemp(prefix="mlruns-")
mlflow.set_tracking_uri(Path(tracking_dir).as_uri())
mlflow.set_experiment("lesson-mlflow-tracking")

with tempfile.TemporaryDirectory() as work_dir:
    report = Path(work_dir) / "evaluation.txt"
    report.write_text("auc=0.91\nprecision=0.84\n", encoding="utf-8")

    with mlflow.start_run(run_name="baseline-logistic") as run:
        mlflow.log_param("model_family", "logistic_regression")
        mlflow.log_param("training_window", "2026-01")
        mlflow.log_metric("validation_auc", 0.91)
        mlflow.set_tag("data_snapshot", "transactions-v5")
        mlflow.log_artifact(str(report), artifact_path="reports")
        run_id = run.info.run_id

print("logged", len(run_id) > 0)

The expected output is logged True. In a real workflow, the data_snapshot tag should identify an immutable dataset or table version, not a vague filename that can be overwritten.

Example 2: Compare Candidate Runs

Experiment tracking is useful when comparisons are reproducible. This example creates three runs, then searches them by metric. The key detail is that the comparison is made from stored run metadata, not from values still living in notebook memory.

import tempfile
from pathlib import Path

import mlflow
from mlflow.tracking import MlflowClient

tracking_dir = tempfile.mkdtemp(prefix="mlruns-")
mlflow.set_tracking_uri(Path(tracking_dir).as_uri())
experiment = mlflow.set_experiment("lesson-mlflow-selection")

candidates = [
    ("small-tree", 3, 0.873),
    ("medium-tree", 7, 0.902),
    ("large-tree", 13, 0.897),
]

for name, max_depth, auc in candidates:
    with mlflow.start_run(run_name=name):
        mlflow.log_param("algorithm", "decision_tree")
        mlflow.log_param("max_depth", max_depth)
        mlflow.log_metric("validation_auc", auc)

client = MlflowClient()
runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    order_by=["metrics.validation_auc DESC"],
    max_results=1,
)
best = runs[0]
print(best.data.params["max_depth"], best.data.metrics["validation_auc"])

The expected output is 7 0.902. This illustrates a common promotion gate: choose the best run according to a declared metric. In production, you would add secondary checks such as minimum recall, fairness slices, inference latency, and a comparison against the currently deployed model.

Example 3: Query by Tags and Review State

Tags turn runs into searchable evidence. Here, a team records data source and review state, then selects only reviewed candidates trained on an approved dataset snapshot.

import tempfile
from pathlib import Path

import mlflow
from mlflow.tracking import MlflowClient

tracking_dir = tempfile.mkdtemp(prefix="mlruns-")
mlflow.set_tracking_uri(Path(tracking_dir).as_uri())
experiment = mlflow.set_experiment("lesson-mlflow-tags")

rows = [
    ("draft", "transactions-v4", 0.915),
    ("approved", "transactions-v4", 0.918),
    ("approved", "transactions-v3", 0.921),
]

for review_state, snapshot, auc in rows:
    with mlflow.start_run():
        mlflow.set_tag("review_state", review_state)
        mlflow.set_tag("data_snapshot", snapshot)
        mlflow.log_metric("validation_auc", auc)

client = MlflowClient()
selected = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    filter_string="tags.review_state = 'approved' and tags.data_snapshot = 'transactions-v4'",
    order_by=["metrics.validation_auc DESC"],
)
print(len(selected), selected[0].data.metrics["validation_auc"])

The expected output is 1 0.918. The highest metric overall was rejected because it used the wrong data snapshot. That is exactly the point of tracking: the best numerical score is not automatically the best operational candidate.

Registry and Artifact Decisions

MLflow can also register models, assigning model names and versions that point back to artifact sources. The registry is valuable when deployment systems need a controlled reference such as CreditRiskModel version 12 instead of a copied file path. The trade-off is process overhead. A registry without stage rules, reviewers, and artifact retention still leaves teams guessing which model should be deployed.

For artifacts, store enough to reproduce and inspect the model: the serialized model, preprocessing objects, dependency information, evaluation reports, and input schema. Avoid logging raw training data unless there is a strong governance reason and the storage system is approved for that data. Metrics should be comparable across runs, so define the evaluation dataset and metric calculation once and reuse it. If one run logs macro F1 on a holdout set and another logs weighted F1 on cross-validation folds, the tracking UI will compare numbers that do not mean the same thing.

Failure Modes and Troubleshooting

Symptom: a run appears in MLflow, but the model or report artifact is missing. Cause: the artifact store path was unavailable, credentials expired, or the process exited before upload completed. Diagnose: open the run artifact tab, check training logs for artifact upload errors, and verify that the training identity can write to the configured artifact URI. Correct: fix storage permissions, retry the run, and make artifact existence part of the promotion check.

Symptom: searches return no runs even though the UI shows them. Cause: the code is pointed at a different tracking URI or experiment name than the UI. Diagnose: print mlflow.get_tracking_uri(), list experiments with MlflowClient, and compare experiment IDs. Correct: set the tracking URI from configuration and fail startup when it is absent or unexpected.

Symptom: the best run cannot be reproduced. Cause: the run logged hyperparameters and metrics but not the data snapshot, code revision, environment, or feature pipeline version. Diagnose: inspect tags and artifacts for lineage fields, then try to rebuild the dataset and environment from the recorded values. Correct: add mandatory tags and artifacts for lineage, and reject runs that omit them.

Symptom: the UI contains many runs with inconsistent metric names such as auc, val_auc, and validation_auc. Cause: teams logged free-form names without a metric contract. Diagnose: search run data by experiment and count metric keys. Correct: publish a metric naming convention and wrap logging in a small project helper.

Security, Reliability, and Performance

Tracking systems often receive sensitive metadata. Do not log secrets, access tokens, personal identifiers, or full raw records as params, tags, or artifacts. Use bounded identifiers that point to governed systems. For shared servers, require authentication and give training jobs write access only to their experiments and artifact locations. Deployment jobs usually need read access to approved model artifacts, not permission to mutate past runs.

Reliability depends on immutability and backups. Runs used for promotion should be treated as evidence. If artifact retention deletes model files after 30 days while registry versions keep pointing to them, rollback will fail. For performance, avoid logging thousands of tiny artifacts or per-row metrics. Aggregate metrics, log curves at sensible intervals, and keep large profiling outputs in compressed files.

Hands-On Lab: Track and Select a Candidate

Prerequisites: Python, MLflow installed in the active environment, and a writable temporary directory. No external tracking server is required.

  1. Create a new Python file or notebook cell using Example 2.
  2. Run it once and confirm the output is 7 0.902.
  3. Change the metric for large-tree to 0.925 and run again in a fresh temporary tracking directory.
  4. Verify that the selected depth changes to 13.
  5. Add a tag named review_state to each run, mark large-tree as draft, and update the search filter to select only approved runs.
  6. Verify that medium-tree is selected again even though large-tree has the highest metric.

Cleanup: because the examples use tempfile.mkdtemp(), the tracking directory path can be removed after inspection. For a persistent local store, delete only the lesson experiment directory, not a shared mlruns directory used by other work.

Assessment

  1. You find two runs with the same validation AUC, but one has no data snapshot tag. Which run is safer to promote, and what evidence would you require before deciding?
  2. A deployment job can read the registry entry but fails to load the model artifact. Explain the likely split between backend store and artifact store that caused this.
  3. Design a filter that excludes draft runs and candidates trained on deprecated data. Which tags must be mandatory?
  4. Why is logging validation_auc from two different evaluation datasets misleading, even if MLflow stores both numbers correctly?
  5. What should a project logging helper enforce before allowing mlflow.end_run() to complete successfully?

Summary

MLflow experiment tracking records model development as durable evidence. Runs capture parameters, metrics, tags, and artifacts; experiments group those runs; the backend store and artifact store serve different roles; and the registry can turn selected artifacts into deployable versions. Good MLOps practice is to log lineage, compare candidates with declared gates, protect artifact access, and troubleshoot tracking gaps before a model is promoted.