ML Orchestration
A DAG of data → features → train → evaluate → register → deploy, with scheduling, retries and backfills. The ML-specific hazards: a training step that is not idempotent, an evaluation gate, and artifact promotion as a step.
The problem, the obvious approach, and why it breaks
Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.
The orchestrator retried a failed training step and the registry now holds two candidates from the same run. What about ML steps makes the standard orchestration assumptions wrong?
We moved our retraining onto the data team's orchestrator. It handles our feature pipelines fine. Last week the training step timed out, the orchestrator retried it, and now the registry has two candidates with the same run id and different metrics. Separately, someone triggered a backfill for a feature table and it kicked off six retrains over the weekend, five of which deployed.
Orchestration is orchestration. The data team's DAG runner handles dependencies, scheduling, retries and backfills for the feature pipelines; add the training, evaluation and deployment steps as more tasks and inherit the defaults.
Data pipeline steps are written to be idempotent — rerun and overwrite the same partition. A training step is not: two runs with the same inputs produce two artifacts with different seeds and slightly different metrics, and each registers itself (Random Seeds, Reproducibility).
- Data pipeline steps are written to be idempotent — rerun and overwrite the same partition. A training step is not: two runs with the same inputs produce two artifacts with different seeds and slightly different metrics, and each registers itself (Random Seeds, Reproducibility).
- Backfills exist to recompute historical partitions. Wired naively, a feature backfill is an upstream change that triggers every downstream task, including training and deployment. Six retrains and five deployments over a weekend is the orchestrator doing exactly what it was told.
- Evaluation was a task that logged metrics and succeeded. It was not a gate: the deploy step depended on evaluation *completing*, not on evaluation *passing*.
What is being predicted, and from what data
This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.
- The model predicts demand per store per day. The orchestration target is that a retrain runs on the intended trigger, each step either completes exactly once or is safe to repeat, the evaluation gate is a step that can fail the DAG, and a backfill of upstream data does not silently redeploy models.
- A DAG: feature tables built nightly from sales and inventory; a training step that reads the latest feature snapshot; an evaluation step; a register step that writes to the registry; a deploy step that promotes and rolls out. Retries are inherited from the data pipeline defaults: three attempts, no idempotency key.
- The training step writes its artifact and registers it before evaluation "so the run is tracked". On timeout after registration, the retry trains again with a different seed and registers again.
How it actually works
Precisely enough to predict its behaviour — not a framework API.
- An orchestrator runs a DAG: tasks with dependencies, a schedule or trigger, retries with backoff, and a state store that records what ran. Data Engineering owns the depth of this — DAGs in Data Pipelines, Idempotent Data Pipelines, Retries in Pipelines, Backfills — and this lesson does not repeat it. What it adds is where ML steps break the orchestrator's assumptions.
- Three ML steps are different. Training is expensive, long and non-deterministic: a retry is a new experiment, not a repeat, so it needs an idempotency key (the run id plus the input hashes) and a rule that a retry resumes from a checkpoint rather than restarting (Checkpointing). Evaluation is a gate: its success state must encode pass or fail, and the downstream deploy must depend on pass. Promotion and deployment are side effects on production; they must be triggered by an explicit decision, never by an upstream data change.
- Backfills are the sharpest hazard. A backfill recomputes features for the past; it should not retrain unless someone decided the training set was wrong, and it should never deploy. The DAG needs a boundary between "data changed" and "model changed", with the model side triggered by policy, not by dependency.
The DAG, with the three steps that break the defaults
The diagram is the ordinary shape: data to features to training to evaluation to registry to deployment. The interesting parts are the edges that are deliberately absent — from a feature backfill into training — and the step that must encode a decision rather than a completion: evaluation.
Everything on the left half is a data pipeline and inherits Data Engineering's discipline. Everything on the right half touches an expensive non-deterministic computation or production, and inherits nothing by default.
A training step that can be retried
The orchestrator will retry. The question is what the step does when it runs a second time with the same inputs. The answer for a data step is "overwrite the partition". The answer for a training step has to be "look up the key; if a finished artifact exists, return it; if a checkpoint exists, resume; otherwise train" — and register only at the end, after evaluation.
The key must include every input that changes the result: the dataset snapshot hash, the feature definition version, the config, the code commit. A key on config alone skips legitimate retrains on new data; a key missing the code commit reuses an artifact from old code.
1import hashlib, json2 3def run_key(snapshot_sha, feature_version, config, code_commit):4 payload = json.dumps([snapshot_sha, feature_version, config, code_commit], sort_keys=True)5 return hashlib.sha256(payload.encode()).hexdigest()[:16]6 7def train_step(inputs, registry, store):8 key = run_key(inputs.snapshot_sha, inputs.feature_version, inputs.config, inputs.code_commit)9 # 1. a finished artifact for these exact inputs: the retry is a no-op10 existing = registry.find_candidate(run_key=key)11 if existing and existing.status == "evaluated":12 return existing13 # 2. a checkpoint from an interrupted attempt: resume, do not restart14 ckpt = store.latest_checkpoint(key)15 model = fit(inputs, resume_from=ckpt, seed=inputs.config["seed"],16 checkpoint_every=lambda m, step: store.save_checkpoint(key, m, step))17 # 3. write the artifact under the key; register only after evaluation passes18 store.save_artifact(key, model)19 return key # the evaluate step reads the artifact by key and decidesThe step never registers. Registration is the evaluate step's side effect on pass, so an interrupted training attempt leaves a checkpoint and nothing in the registry — which is what makes the retry safe.
Backfills, and the edge that must not exist
A feature backfill is routine for the data team: recompute last month's partitions because a bug was fixed. In a DAG where training depends on the feature table, the backfill is an upstream change and the orchestrator dutifully runs everything downstream, once per recomputed partition if the DAG is partitioned by date.
The fix is structural, not a flag. The data DAG publishes snapshots. The model DAG reads them when a retraining policy decides to — on a schedule, on a data-volume threshold, on a drift signal reviewed by a person. There is no dependency edge; there is a policy that reads state.
No change on the data side — a backfill, a late partition, a schema fix — can start a training run or a deployment without a retraining decision.
holds when The data DAG and the model DAG are separate; the model DAG's only trigger is the retraining policy; backfills are scoped to the data DAG.
breaks when Someone adds a convenience edge so that "new features retrain automatically"; a sensor task polls the feature table and fires on any change; the two DAGs are merged during a migration.
respond Remove the edge and roll back any deployments it caused to the last candidate that went through a decision; then write down why the edge is absent, because someone will add it again.
| Trigger | Symptom | Cause | Response |
|---|---|---|---|
| Training step times out; orchestrator retries | Two candidates, same run id, different metrics | Training is non-idempotent and registered before evaluation | Key the step on inputs; resume from checkpoint; register after evaluation only |
| Feature backfill over a month of partitions | Many retrains and deployments over a weekend | A dependency edge from feature partitions into the training task | Remove the edge; the model DAG is triggered by a retraining policy that reads snapshots |
| Candidate worse than incumbent | Deployed anyway; evaluation task shows success | Deploy depends on evaluation completing, not passing | Evaluation fails the DAG on a losing candidate; register and deploy depend on pass |
| Orchestrator restarts mid-deploy | Canary left at partial traffic with no rollback timer | Deployment treated as a task with no state of its own | Deployment is its own stateful rollout with a timer and rollback criteria; the orchestrator only starts it |
How to build it
Most important first.
- Give the training step an idempotency key from its inputs — dataset snapshot hash, feature definition version, config hash — and make the step check the registry for that key before training. A retry with the same key resumes from checkpoint or returns the existing artifact.
- Make evaluation a gate: the task fails the DAG when the candidate does not beat the incumbent on the agreed criteria, and the register step runs only on pass (Champion / Challenger, Promotion Is a Checklist, Not a Score). Register after evaluation, never before.
- Separate the data DAG from the model DAG. The feature pipeline's success publishes a snapshot; the model DAG is triggered by a retraining policy that reads snapshots (Retraining Strategies). Backfills on the data side do not propagate.
- Make deployment a manual or policy-gated step with its own rollout — canary and rollback are the deploy step's job, not the orchestrator's (Canary Rollout).
What to measure
Which number actually maps to the decision — and which numbers look relevant and are not.
- Number of candidates per run id in the registry: exactly one. More than one is a retry that trained twice.
- Number of deployments per week against the number of retraining decisions per week. They should match; a surplus of deployments is a dependency edge that should not exist.
- DAG success rate is not the number. A DAG that deployed five models over a weekend was fully successful.
What must stay true after deployment
The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.
- Every training step is keyed on its full inputs and checks the registry before training, so a retry cannot produce a second candidate for the same inputs.
- The deploy step depends on the evaluation gate's pass state, and no data-side task has a dependency edge into the model DAG.
- The checkpoint resume path has been exercised on purpose and the resumed run's metrics match an uninterrupted one within tolerance.
- Offline: kill a training step mid-run in staging and let the orchestrator retry it. Confirm one candidate in the registry, resumed from checkpoint.
- Online: run a feature backfill over a week of history and confirm no training or deployment task started. Then trigger the retraining policy by hand and confirm one run.
- Over time: audit the registry for run ids with multiple candidates and the deployment log for deployments without a matching retraining decision. Both should be empty.
What can go wrong
- The idempotency key is computed from the config but not the dataset hash, so a retrain on new data with the same config is skipped as "already done".
- The evaluation gate compares against the incumbent's recorded metric rather than re-evaluating the incumbent on the same fresh data, and the gate passes candidates that would lose a fair comparison.
- The checkpoint resume path is never exercised until a real timeout, and the resumed run fails on an optimizer-state mismatch at three in the morning.
- Idempotency keys and checkpoint resume add state to the training step and a registry lookup before every run; the step is no longer a pure script.
- Separating the data DAG from the model DAG means a feature fix does not automatically flow into a retrain, which is the point, and also a thing someone now has to remember to trigger.
- A real evaluation gate produces failed DAG runs that page someone, and the first reaction is usually to relax the gate.
- "Retries make the pipeline robust." Retries make idempotent steps robust. A retried training step is a second experiment that registers itself as if it were the first.
- "Evaluation ran, so the model was evaluated." The task completed. Whether the deploy step depended on the *result* is a separate question, and the default answer is no.
- "The backfill only touched feature tables." The DAG has edges. Everything downstream of a changed partition is fair game unless the edge into training is deliberately absent.
Where this applies
ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.
- GENERALThat training is non-idempotent, evaluation is a gate and deployment is a production side effect holds for every model family and every orchestrator; the mechanics of retries and backfills are the orchestrator's and belong to Data Engineering.
- FRAMEWORK-SPECIFICHow a task encodes pass/fail, how idempotency keys are expressed and whether backfills propagate by default vary by orchestrator; the design here is stated in terms every DAG runner can express, and the specifics belong in that tool's documentation.
- SCALE-SPECIFICA single model retrained monthly can be orchestrated by a cron job and a script with a registry check; the DAG, the gate as a task and the data/model boundary become necessary once several models share feature tables and retraining runs unattended.
Where the depth lives
This domain teaches the model and hands the rest off by name.