ML Experiment Tracking System

Log, compare and reproduce training runs

Design a service that records every model training run (its settings, metrics, files and environment) so researchers can compare runs, reproduce a result, and trace a production model back to how it was made.

Overview

An experiment tracker records what happened in each training run. A run is one execution of training code; an experiment groups related runs (MLflow's documentation describes it the same way). For each run the tracker keeps parameters (settings such as the learning rate, written once), metrics (numbers such as loss, logged many times against a step), tags, and artifacts (files such as model checkpoints and plots).

Researchers use it for three jobs: see a run's curves live, find and compare the runs that worked, and rebuild a result months later. Production teams add a fourth: given the model now serving traffic, find the run, code and data that produced it.

The two hard parts are opposite in shape. Metrics are billions of tiny points that must be written without slowing training and read back as curves in a second or two. Artifacts are few but huge, and without deduplication and retention they fill any budget.

Questions to ask first

How many runs at once? — About 2,000 runs training at the same time, most lasting about a day
How much do they log? — About 20 metrics each, about one point per metric per second
How large are checkpoints? — About 2 GB, about 10 saved per run
What must be kept? — The best checkpoint of every run, and forever anything registered as a model
May a tracker outage stop training? — Never; logging degrades, training continues
Hosted by us or per team? — One shared service for all teams

What you will be able to do

Design an SDK that never blocks or crashes training and still delivers every point exactly once.
Split the data across a metadata database, a columnar metric store and an object store, and say why.
Search runs by parameters and metrics, and draw 100,000-step curves quickly.
Store checkpoints once by content hash, and delete them safely under a retention policy.
List what a run must record to be reproducible, and trace a model back to its run.

Key concepts

Run and experiment
Parameters, metrics, tags
Asynchronous batched logging
Idempotent batches
Columnar metric store
Downsampling
Content-addressed artifacts
Retention and garbage collection
Reproducibility record
Model lineage

Functional requirements

Create a run inside an experiment and record its parameters, tags and environment
Log metrics against a step and a timestamp, many times a second, from many processes
Upload artifacts (checkpoints, plots, evaluation files) of any size
Search runs by parameters, tags, status and metric values; compare curves side by side
Keep artifacts under a retention policy; never delete one behind a registered model
Show the full reproducibility record and the lineage of a registered model

Non-functional requirements

Logging adds no blocking work to the training step; a tracker outage never stops training
Each logged point is stored exactly once even when the SDK retries
Comparing 20 runs of 100,000 steps draws in under 2 seconds (a target we chose)
Parameters and artifacts of a finished run never change
Metadata and metrics highly available; 99.9% for the API (a target we chose)

Interview questions

How do you log metrics without slowing down training?
Where do you store billions of metric points, and why not in PostgreSQL?
How do you draw a 100,000-step curve for 20 runs quickly?
How do you avoid storing the same checkpoint many times, and how do you delete safely?
What does it take to reproduce a run?
How do you trace a production model back to how it was trained?
How would you support a hyperparameter sweep of 200 runs?