Training pipelines and orchestration

The big idea

The first version of a model is usually trained by a person running notebook cells in order. That works once. The tenth time, someone skips a cell, the data has changed, the person is on holiday, or the upstream table is half empty and nobody notices until the new model is serving. A training pipelineAn automated, repeatable workflow that turns raw data into a validated, registered model: extract, validate, build features, train, evaluate, register. is that notebook rewritten as a sequence of separate tasks that a machine runs the same way every time, with checks that stop it before a bad model reaches users.

Think of a bakery's overnight production. Nobody improvises: flour is checked when it arrives, dough is mixed, proofed and baked in a fixed order, the first loaf is tasted before the batch goes out, and if the flour is bad, no bread ships and yesterday's stock stays on the shelf. The orchestrator is the shift manager who starts the run, enforces the order, retries the oven that tripped, and records what happened.

How it works

Step 1: The pipeline is a DAG

A pipeline is a DAGDirected acyclic graph: tasks joined by one-way dependencies with no cycles, so there is always a valid order to run them in.: tasks connected by dependencies, with no cycles. Splitting the work into tasks buys four things. Each task can be retried on its own when it fails for a transient reason. Independent tasks can run in parallel. Outputs of finished tasks can be cached, so a rerun after a training failure does not extract the data again. And each task's inputs, outputs and logs are recorded, which links straight into experiment tracking ( Experiment tracking and reproducibility ).

Step 2: The orchestrator

An orchestratorA system that runs DAGs: it decides when a run starts, executes tasks in dependency order on workers, retries failures, and records the state and logs of every task. runs DAGs for you. Two common ones:

• Apache Airflow is a general workflow scheduler. A Dag (Airflow's spelling) bundles the tasks, their dependencies and the schedule; dependencies are written with >>, as in extract >> validate >> train. The schedule argument takes a preset such as "@daily" or a cron expression such as "0 0 * * *", and default arguments such as retries apply to every task.
• Kubeflow Pipelines (KFP) is built for ML workflows on Kubernetes. Each component execution runs in its own container, so the training step can ask for a GPU while the validation step runs on a small CPU. Pipelines are written in Python with @dsl.component and @dsl.pipeline and compiled to an intermediate-representation YAML file. KFP can cache results to skip redundant executions.

from kfp import compiler, dsl

@dsl.component
def train(rows: int) -> float:                       # each component runs in its own container
    return 0.88                                      # ... real training code here

@dsl.pipeline
def training_pipeline(rows: int = 100000):
    train(rows=rows)

compiler.Compiler().compile(training_pipeline, package_path="pipeline.yaml")

Managed cloud services (for example the pipeline services of the big cloud ML platforms) and other open-source tools follow the same model of tasks, dependencies and runs.

Step 3: Scheduled or triggered retraining

• Scheduled: retrain every night or every week. Simple and predictable; costs compute even when nothing changed, and reacts to change only at the next slot.
• Triggered by data: run when new data lands. Airflow can schedule a Dag on assets: a producer task declares an asset as an output, and a consumer Dag with schedule=[asset] runs after the producer task succeeds.
• Triggered by monitoring: run when drift or a performance drop is detected ( ML system architecture overview sketched this loop; drift detection gets its own lesson, Monitoring, drift and retraining ).

Many teams combine them: a regular schedule as a floor, plus triggers for urgent cases.

Step 4: Validation gates

An automated pipeline will eventually run on bad data, so it needs automated checks that can say no.

• Data validation before training: schema (columns, types), null rates, row counts, value ranges and distributions against the previous run.
• Model validation after training: evaluate on a fixed holdout and compare with the champion, the model currently serving. TFX's Evaluator is a well-known example: it "compares new models against a baseline (such as the currently serving model)", and only if the new model's metrics meet the criteria relative to the baseline (for example, AUC is not lower) is it "blessed", the signal that it can be pushed to production.

A failed gate is a success for the pipeline: it stops, alerts someone, and the champion keeps serving.

Going deeper

Code: an orchestrator in 40 lines

from graphlib import TopologicalSorter
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
rng = np.random.default_rng(5)
w_true = np.array([1.5, -1.0, 0.8])
def batch(n, bug=None):                                     # one day of upstream data
    X = rng.normal(size=(n, 3)); y = (X @ w_true + rng.logistic(size=n) > 0).astype(int)
    if bug == "nulls": X[rng.random(n) < 0.4, 1] = np.nan    # a broken join upstream
    if bug == "labels": y = rng.integers(0, 2, n)            # a silent label bug: no NaNs, right schema
    return X, y
holdout = batch(3000); state = {"champion": None, "auc": 0.0, "flaky": 1}
# The DAG: task -> the tasks it depends on. Each task reads and writes a shared context.
dag = {"extract": [], "validate_data": ["extract"], "train": ["validate_data"],
       "evaluate": ["train"], "gate": ["evaluate"], "register": ["gate"]}
def extract(c):
    if c["day"] == 3 and state["flaky"]: state["flaky"] = 0; raise TimeoutError("source timed out")
    c["X"], c["y"] = batch(2000, {5: "nulls", 6: "labels"}.get(c["day"]))
def validate_data(c):
    null_rate = np.isnan(c["X"]).mean()
    if null_rate > 0.01: raise ValueError(f"null rate {null_rate:.1%} > 1%")
def train(c): c["model"] = LogisticRegression().fit(c["X"], c["y"])
def evaluate(c): c["auc"] = roc_auc_score(holdout[1], c["model"].predict_proba(holdout[0])[:, 1])
def gate(c):
    if c["auc"] < state["auc"] - 0.005: raise ValueError(f"AUC {c['auc']:.3f} < champion {state['auc']:.3f}")
def register(c): state["champion"], state["auc"] = c["model"], c["auc"]
tasks = {f.__name__: f for f in [extract, validate_data, train, evaluate, gate, register]}

def run_pipeline(day, retries=1):                          # a tiny orchestrator: order, retries, stop on failure
    c = {"day": day}
    for name in TopologicalSorter(dag).static_order():
        for attempt in range(retries + 1):
            try: tasks[name](c); break
            except TimeoutError as e: print(f"  day {day}: {name} attempt {attempt + 1} failed ({e}), retrying")
            except ValueError as e: return f"stopped at {name}: {e}"
    return f"registered new champion, holdout AUC {c['auc']:.3f}"

for day in range(1, 8):                                    # scheduled: once a day
    print(f"day {day}: {run_pipeline(day)}   serving AUC {state['auc']:.3f}")

day 1: registered new champion, holdout AUC 0.877   serving AUC 0.877
day 2: registered new champion, holdout AUC 0.876   serving AUC 0.876
  day 3: extract attempt 1 failed (source timed out), retrying
day 3: registered new champion, holdout AUC 0.877   serving AUC 0.877
day 4: registered new champion, holdout AUC 0.877   serving AUC 0.877
day 5: stopped at validate_data: null rate 12.9% > 1%   serving AUC 0.877
day 6: stopped at gate: AUC 0.532 < champion 0.877   serving AUC 0.877
day 7: registered new champion, holdout AUC 0.877   serving AUC 0.877

The code defines the DAG as a dictionary of dependencies, orders it with Python's graphlib.TopologicalSorter, and runs it once a day for seven days, with one retry for transient errors and a stop on validation errors.

• Days 1, 2, 4 and 7 run cleanly; each new model has holdout AUC 0.876 to 0.877 and is registered (the gate allows a drop of at most 0.005).
• Day 3: the extract task times out, the orchestrator retries it once, and the run completes. Nobody was paged.
• Day 5: a broken upstream join leaves 12.9 percent of all values missing. The data gate stops the run before training; the day 4 model keeps serving.
• Day 6 is the dangerous one: a label bug with the right schema and no missing values. Data validation cannot see it, but the model trained on it has holdout AUC 0.532 against the champion's 0.877, and the model gate rejects it. Without that gate, day 6 would have shipped a model that is barely better than a coin flip.

Trade-offs

Common mistakes

• Putting the whole notebook into one giant task, so any failure reruns everything and nothing is cached.
• Retrying everything, including validation failures; retries are for transient errors, not for bad data.
• No model gate, so a pipeline faithfully trains and ships a broken model at 3 a.m.
• Comparing the new model with the champion on different evaluation data, so the gate compares apples with oranges.
• Retraining on a schedule without monitoring, and never asking whether the retrains help.

Where you will see this

Pipelines like this sit behind most production models. Uber's Michelangelo post, from ML system architecture overview , notes that many teams use its API and workflow tools to schedule regular retraining of their models. TFX, Google's open-source pipeline framework, pairs the Evaluator's blessing with a Pusher component that relies on blessings from validation components to decide whether to push a model to its deployment target. Teams that already run their data warehouse jobs in Airflow can add training as one more Dag, while teams on Kubernetes often choose Kubeflow Pipelines.

Check yourself

Recap