Lab 03 — ML Pipelines: a DAG Engine with Artifacts, Caching & Conditions

Phase 25 · Lab 03 · Phase README · Warmup

The problem

One training job is a script. A production ML system is a pipeline: preprocess → train → evaluate → if the metrics clear the bar → register — re-run on a schedule, on new data, on every code change. All three clouds converged on the same abstraction for this (SageMaker Pipelines, Vertex AI Pipelines / Kubeflow, Azure ML pipelines), and it has exactly four load-bearing ideas:

1. A DAG of steps. Each step declares what it consumes and produces; the edges are derived from those declarations, not drawn by hand. The engine topologically sorts and executes, and a cycle is a definition error caught before anything runs.
2. Artifacts. Steps don't share memory; they pass named artifacts (in real platforms: S3/GCS URIs) — which is what makes steps independently retryable, cacheable, and auditable.
3. Caching. If a step's inputs and parameters haven't changed, don't run it again — replay its recorded outputs. This is why re-running a 6-hour pipeline after fixing the last step takes minutes, and why a cache key must cover everything that determines the output.
4. Conditions. "Register the model only if accuracy beats the threshold" is a first-class branch node, and the untaken branch is skipped, not failed — with skips cascading downstream.

This lab builds that engine. Every step function is a pure injected callable, so topological order, artifact flow, cache hits/misses, and both branch outcomes are all exact assertions.

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

• topological_order respects every dependency, is deterministic (name-order tie-break), and raises CycleError for a 2-cycle and a self-cycle; an unknown input reference raises ValidationError before anything runs.
• Artifacts flow: a 3-step linear pipeline computes ((start + 1) * 10) - 3 end to end via the artifact store; a step that omits a declared output fails the run.
• Caching: an identical second run re-executes nothing (all steps cached=True, outputs replayed); changing any input value re-executes; compute_cache_key differs on step name, input values, and params — and matches on identical ones.
• Conditions: both branches tested — the taken branch runs, the untaken one is Skipped, and a step downstream of a skipped step is skipped too (never failed).
• Fan-out/fan-in: two parallel branches after one parent; the join sees both outputs.
• All 24 tests pass under both lab and solution.

How this maps to the real stack

• SageMaker Pipelines: TrainingStep, ProcessingStep, ConditionStep, ModelStep / RegisterModel, with step properties (e.g. a training step's model-artifact S3 URI) wired into downstream steps — that property-reference wiring is exactly our "step.output" reference strings, and SageMaker derives the DAG from them just like _deps_of does. SageMaker step caching keys on the step arguments and, on a hit, reuses the prior execution's outputs.
• Vertex AI Pipelines / Kubeflow Pipelines (KFP): @dsl.component functions whose typed inputs and outputs define the DAG, dsl.If/dsl.Condition for branches, and execution caching that keys on the component spec + inputs and skips re-execution on a hit — our cache is that mechanism minus the container image in the key (our step name stands in for it). Vertex tracks every artifact in ML Metadata, which is the production version of our run.artifacts dict plus lineage.
• Azure ML pipelines: @pipeline / @command_component definitions, inputs/outputs binding steps together, and is_deterministic-based reuse of previous results — the same three ideas with Azure vocabulary.
• The condition-gated register step in build_train_evaluate_register_pipeline is the shape SageMaker's own MLOps examples use: evaluate writes a metrics artifact, a ConditionStep compares it to a threshold, and only the passing branch calls RegisterModel — Lab 01's registry is what that step writes into, closing the loop across this phase's labs.

Limits of the miniature. Real steps are containers on remote compute with retries, timeouts, and per-step resources; ours are in-process function calls. Real artifacts are object-store URIs with metadata lineage; ours are in-memory values keyed by name. Real cache keys include the container image digest and command — renaming our step changes the key, but changing its code does not, which is precisely the "stale cache" foot-gun the Warmup warns about (KFP keys on the component spec, so a code change that doesn't change the spec can also serve stale results). Real engines run ready steps in parallel; ours runs them sequentially in deterministic order so tests can assert on it.

Extensions (your own machine)

• Add retry-with-limit per step (a step may raise transiently; re-run it up to max_retries times before failing the run) with an injected failure schedule.
• Add parallel execution with a worker pool over the ready set, and prove the artifact results are identical to sequential execution (determinism under concurrency).
• Persist the artifact store and cache to disk (JSON) so a re-run in a new process still gets cache hits — the first step toward real object-store-backed artifacts.
• Add lineage queries: given an artifact, walk back through the steps and inputs that produced it (Vertex ML Metadata's core query), reusing Lab 01's registry lineage idea.

Interview / resume signal

"Built a miniature ML-pipeline engine in the SageMaker Pipelines / KFP mold: a DAG derived from declared artifact references, deterministic topological execution with cycle detection, content-keyed step caching (hit ⇒ replay, miss ⇒ re-run), first-class conditional branches with cascading skips, and fan-out/fan-in — culminating in the canonical train→evaluate→gate→register pipeline with both branches under test."