Lab 01 — Managed Training Jobs, HPO & the Model Registry

Phase 25 · Lab 01 · Phase README · Warmup

The problem

Strip the branding off SageMaker, Vertex AI, and Azure ML and the training half of all three is the same spine, and it is the part an interviewer will make you draw:

1. You describe a training job — a container/script, hyperparameters, and a compute request — and hand it to a managed compute pool. The platform schedules it onto instances, runs it, and reports a lifecycle: Pending → InProgress → Completed/Failed.
2. Because GPU instances are expensive, you run on cheap, interruptible spot/preemptible capacity and rely on checkpointing to survive an interruption — a resumed job must produce the same model it would have without the interruption.
3. To find good hyperparameters you launch a tuning job that fans out many trials and keeps the best by an objective metric.
4. The winning artifact lands in a model registry: versioned, staged (None/Staging/Production/Archived), and traceable by lineage back to the exact job and hyperparameters that produced it.

This lab builds that whole spine as an offline, deterministic miniature. The one non-deterministic, expensive thing — training — is injected as a pure per-epoch function, so the whole job becomes a pure function of its hyperparameters and every assertion is reproducible.

What you build

File map

Run it

pytest test_lab.py -v                       # your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # the reference (green)
python solution.py                          # worked example

Success criteria

• A TrainingJob starts Pending and the runner drives it to Completed with a TrainingResult (metrics + artifact hash + epochs), or Failed if train_step raises.
• interrupt_before_epoch puts the job in Stopped at the last durable checkpoint, and resume finishes it — producing an artifact hash identical to an uninterrupted run.
• TrainingCluster places a Pending job on a free instance; a Stopped job keeps its instance; placing over capacity raises CapacityError.
• HyperparameterTuner fans out trials (grid exhaustive, random seed-deterministic, bayesian deterministic) and selects the true best trial by a Minimize/Maximize objective.
• ModelRegistry increments versions, enforces the stage state machine (no direct None → Production), auto-archives the incumbent on promotion, records lineage, and answers get_latest_production.
• All 30 tests pass under both lab and solution.

How this maps to the real stack

• The training job is SageMaker's CreateTrainingJob / Vertex AI's CustomJob / Azure ML's command job. You give it a container or an entry script, hyperparameters, an instance_type + instance_count, and input/output data channels; the platform provisions the cluster, runs your code, ships logs/metrics to the platform's metric store, and writes the model artifact to S3/GCS/Blob. Our train_step injection is the seam that stands in for "your training code" so the lifecycle is testable without a GPU. The real lifecycle strings are essentially ours (InProgress, Completed, Failed, Stopped).
• Spot/managed checkpointing is real and load-bearing: SageMaker Managed Spot Training (and Vertex/Azure preemptible VMs) can cut training cost by a large margin, but only if your job checkpoints to the output path and can resume — SageMaker re-supplies the last checkpoint on restart. Our checkpoint-every-N-epochs + resume-from-checkpoint, and the assertion that the resumed artifact is byte-identical, is exactly that contract in miniature.
• Hyperparameter tuning is SageMaker's Automatic Model Tuning (HyperParameterTuningJob) / Vertex AI Vizier / Azure ML sweep jobs. Real ones support grid, random, and Bayesian strategies, early stopping, and parallel trials; ours implements grid, seeded random, and a simulated (exploitation-only) Bayesian search. The mechanism to internalize is that tuning is a search loop that launches training jobs and reads back an objective metric — not a distinct black box.
• The model registry is SageMaker Model Registry (model package groups + versions + ModelApprovalStatus), the MLflow Model Registry (which uses the exact None/Staging/ Production/Archived stage vocabulary we model), and Vertex AI Model Registry / Azure ML registered models. The state machine, the "one production version" invariant via auto-archive-on-promote, and lineage back to the source run are the real value: at 2 a.m. you must be able to answer "what is in production and which job/params produced it" from the registry alone.

Limits of the miniature. Real training is stochastic (data shuffling, dropout, non-deterministic GPU kernels) and reproducibility takes real effort (fixed seeds, deterministic ops); our train_step is pure by construction. Real checkpoints are multi-GB tensor snapshots to object storage; ours is a small state dict. Real tuning runs trials in parallel across a fleet; ours runs them sequentially so the test suite can assert order and selection. SageMaker Model Registry's lifecycle is actually an approval-status workflow (PendingManualApproval/Approved/ Rejected) wired into CI/CD, whereas MLflow uses the stage vocabulary we chose — both ideas appear in the Warmup.

Extensions (your own machine)

• Add distributed training: split train_step across simulated workers (data-parallel: average a per-shard metric; model-parallel: partition the state) and prove the averaged result matches single-worker training.
• Add early stopping to the tuner: stop a trial (and free its instance) once its objective is provably worse than the best-so-far after k epochs.
• Add a warm-pool to TrainingCluster (keep an instance hot for ttl ticks after a job finishes, using an injected tick clock) and measure the scheduling-latency win.
• Wire the registry to an approval-status workflow instead of stages, and gate Production on an injected "CI passed" signal — the SageMaker Model Registry shape.

Interview / resume signal

"Built a miniature managed-training platform: a training-job lifecycle on a finite compute pool, spot-interruption survival via checkpoint/resume proven to yield a byte-identical artifact, a hyperparameter-tuning job that fans out grid/random/Bayesian trials and selects the best by objective, and a model registry with a stage state machine (auto-archiving the incumbent on promotion) and full lineage back to the source job and hyperparameters."