Lab 01 — MLflow-style Tracking Store & Model Registry

Phase: 17 — MLOps: The ML Platform (Tracking, Registry, Drift, CI/CD) Difficulty: ⭐⭐⭐☆☆ (the data structures are simple; the state machine and lineage discipline are ⭐⭐⭐⭐⭐) Time: 3–4 hours

Every MLOps platform is, underneath the dashboards, two mechanisms: an experiment-tracking store that records what each training run did (params, metrics with full step history, tags, artifacts) and a model registry that turns the winning run into a versioned, staged, promotable, rollback-able model with lineage back to exactly what produced it. This lab builds both from first principles — a deterministic, in-memory MLflow — so the words run, param immutability, metric step history, best run, version, stage transition, archive-the-incumbent, and lineage stop being dashboard buttons and become code you own.

What you build

• Run — one trial: run_id, experiment, params, metrics (key → ordered (step, value) history), tags, artifacts, status (RUNNING/FINISHED/FAILED), parent_run_id for nested runs.
• TrackingStore — deterministic in-memory backend store: create_experiment, create_run (incrementing run-0001 ids — no uuid, no clock), log_param (immutable — re-logging a different value raises), log_metric (keeps the full step history), log_artifact/get_artifact, set_tag, end_run, get_run, search_runs (filter + order), best_run (max/min over a metric), compare_runs.
• ModelRegistry — versions, stages, lineage, backed by the store: register_model, create_model_version (auto-incrementing version bound to a run), transition_stage with a real state machine over {"None","Staging","Production", "Archived"} (illegal edges rejected; archive_existing=True demotes the current Production version to Archived when you promote a new one), get_latest_versions, get_model_lineage (version → run → params/metrics), set_version_tag.

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # against the reference (must be green)
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why params are immutable but tags and metrics are not — and why test_relog_param_different_value_raises is a correctness test, not pedantry.
• You can explain why metrics keep the full (step, value) history and why best_run uses the latest value (test_best_run_uses_latest_value_not_first).
• You can draw the stage state machine and say which edges are illegal and why Archived → Production directly is forbidden.
• You can explain what archive_existing=True guarantees (exactly one Production version) and why test_archive_existing_promotes_and_demotes_incumbent is the registry soul test.
• You can trace get_model_lineage from a version back to the params/metrics that produced it and say which 2 a.m. question that answers.

How this maps to the real stack

Limits of the miniature (say these in the interview): it is in-memory and single-process — a real tracking server has a persistent backend store (Postgres/MySQL/filesystem) plus a separate artifact store for blobs, and concurrent writers need transactional guarantees we don't model; our "lineage" is param/metric provenance only — true reproducibility also needs the code commit and the data/feature version (Git + DVC/Delta), which Lab 02/03 lean on; the stage state machine here is the classic 4-stage model, while MLflow ≥ 2.9 deprecates fixed stages in favor of model aliases (@champion, @challenger) and tags for exactly the flexibility our fixed edges lack; and a real registry adds access control / approvals / webhooks on transitions (the human gate before Production) that we leave as an extension.

Extensions (build these on real MLflow / your own platform)

• Run it on real MLflow. The same workflow, against a live server:

  import mlflow
  from mlflow import MlflowClient
  
  mlflow.set_experiment("fraud-classifier")
  with mlflow.start_run() as run:                 # a Run, auto-ended on exit
      mlflow.log_param("lr", 0.05)                # immutable
      for step, auc in enumerate([0.83, 0.88, 0.91]):
          mlflow.log_metric("val_auc", auc, step=step)   # step history
      mlflow.log_artifact("confusion_matrix.png")
      mlflow.sklearn.log_model(model, "model")    # writes an MLmodel flavor
  
  # register the winner and promote it, demoting the incumbent
  mv = mlflow.register_model(f"runs:/{run.info.run_id}/model", "fraud-model")
  MlflowClient().transition_model_version_stage(
      name="fraud-model", version=mv.version,
      stage="Production", archive_existing_versions=True,   # archive-the-incumbent
  )
  prod = MlflowClient().get_latest_versions("fraud-model", stages=["Production"])[0]
  src_run = MlflowClient().get_run(prod.run_id)   # lineage: params/metrics that made prod
  

• Add a real backend/artifact split: persist runs to SQLite and artifacts to a local directory; show that search_runs still works and survives a restart.

• Swap stages for aliases: add set_alias(name, alias, version) / get_by_alias and a blue/green champion/challenger pattern (this is where MLflow is going).

• Add an approval gate: require a set_version_tag(..., "approved_by", user) before a transition into Production is allowed — the human-in-the-loop release gate.

• Compare against W&B / Neptune: log the same sweep to one of them and note what their data model adds (system metrics, artifact lineage graphs, hosted UI) over this core.

Interview / resume

• Talking points: "Why are params immutable but metrics a time-series?" "Walk me through a model promotion to Production and how you'd roll back." "What does a model registry give you that a folder of .pkl files doesn't?" "How do you answer what model is in prod and what data trained it at 2 a.m.?"
• Resume bullet: Built an MLflow-style experiment-tracking store and model registry — deterministic run tracking with immutable params and full metric step history, search/best-run selection, and a versioned registry with a stage state machine (Staging→Production→Archived), atomic promote-and-archive, and run-level lineage — mapping each mechanism to the production MLflow API.