AI Training Pipeline Architecture in Depth: run manifests, leakage checks, resumable training, checkpoint intervals and evaluation gates

The architecture: a run as a function

Think of a run as a pure function. Inputs: a data snapshot id, a code commit, a container image digest, a configuration and a random seed. Output: model weights, metrics and a lineage record. Everything in the architecture exists to keep that function honest, meaning the same inputs give the same model within known tolerance, and every output names its inputs.

An orchestrator (Airflow, Kubeflow Pipelines, Argo Workflows, Flyte or Metaflow are common) runs the steps as separate tasks with recorded inputs and outputs, so a failed evaluation can be rerun without retraining. An experiment tracker such as MLflow or Weights & Biases records parameters, curves and artifacts. A model registry stores the approved artifact with its lineage, as described in model registry architecture.

The run manifest

Write the run's identity down before the run starts, as a manifest the job reads instead of loose flags:

run_id: churn-2026-10-05-a
data:
  snapshot: s3://lake/churn/features@snapshot=8f3c1e        # immutable, never "latest"
  splits: {train: "signup < 2026-08-01", val: "2026-08", test: "2026-09"}
code:
  commit: 4b9e2d7
  image: registry.example.com/train@sha256:1c0f...           # digest, not a tag
config:
  model: gbdt_v3
  lr: 0.05
  max_steps: 40000
  global_batch: 4096
seed: 1234
eval:
  baseline_model: churn@prod
  gates: {auc_min_delta: -0.002, max_slice_drop: 0.01}

Two details carry most of the value. The data reference is a snapshot, never a path that can change, which is the subject of data versioning. And the image is pinned by digest, so a rebuilt image with a new library version cannot slip in under the same tag. Hash the manifest and store the hash with every metric and artifact.

Data validation and leakage

Validate before spending GPU hours. Check the schema and types, null rates, value ranges and category sets against the previous snapshot, and fail the run on large unexplained shifts. Check that labels are present and their rate is plausible. Then check for leakage, which produces the most convincing bad models:

• Split by entity or time, not by row. If one customer's rows land in both train and test, the model memorises customers and the test score is inflated. For anything forecast-like, split by time.
• Deduplicate across splits. Near-duplicate documents or images in train and test have the same effect.
• Check feature timestamps. A feature computed after the label's event, such as "number of support tickets this month" for a churn label in the same month, leaks the answer.

Freeze the test split. Every model is compared on the same rows, and nobody tunes against it; tuning happens on validation.

The input pipeline

The input pipeline decides whether expensive accelerators do any work. Measure it directly: time each step split into waiting for the next batch and compute. If the wait is more than a few percent of step time, the job is input-bound and adding GPUs makes it slower per dollar.

Fixes, roughly in order of payoff: store preprocessed data in large sequential shards (Parquet, WebDataset tar files or similar) instead of millions of small files; read with several loader worker processes; prefetch batches and copy to the GPU asynchronously from pinned host memory; move heavy decoding or augmentation either offline into preprocessing or onto the GPU. Each worker reads a disjoint set of shards, and the shuffle order is derived from the seed and epoch so it can be reproduced after a restart.

A training loop that resumes exactly

A resumable loop saves everything needed to continue as if nothing happened: model and optimizer state, the learning-rate scheduler, the mixed-precision gradient scaler, the step counter, the data position and the random number generator states. PyTorch, with distributed data parallel:

import os, random, numpy as np, torch, torch.distributed as dist

def save_ckpt(path, step, model, opt, sched, scaler, epoch):
    if dist.get_rank() != 0:
        return
    state = {
        "step": step, "epoch": epoch,
        "model": model.module.state_dict(),
        "opt": opt.state_dict(), "sched": sched.state_dict(), "scaler": scaler.state_dict(),
        "rng": {"py": random.getstate(), "np": np.random.get_state(),
                "torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state_all()},
    }
    tmp = path + ".tmp"
    torch.save(state, tmp)
    os.replace(tmp, path)            # atomic on one filesystem: never a half-written checkpoint

def train(cfg, model, opt, sched, scaler, dataset, ckpt_path):
    sampler = torch.utils.data.distributed.DistributedSampler(dataset, seed=cfg.seed)
    loader = torch.utils.data.DataLoader(dataset, batch_size=cfg.batch, sampler=sampler,
                                         num_workers=8, pin_memory=True)
    step, epoch = 0, 0
    if os.path.exists(ckpt_path):
        s = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        model.module.load_state_dict(s["model"]); opt.load_state_dict(s["opt"])
        sched.load_state_dict(s["sched"]); scaler.load_state_dict(s["scaler"])
        step, epoch = s["step"], s["epoch"]
        random.setstate(s["rng"]["py"]); np.random.set_state(s["rng"]["np"])
        torch.set_rng_state(s["rng"]["torch"]); torch.cuda.set_rng_state_all(s["rng"]["cuda"])
        # Only rank 0's generators are saved here; save one file per rank if every rank's RNG matters.
    while step < cfg.max_steps:
        sampler.set_epoch(epoch)       # different, reproducible shuffle per epoch
        skip = step % len(loader) if step else 0
        for i, (x, y) in enumerate(loader):
            if i < skip:
                continue               # fast-forward to the saved position
            with torch.autocast("cuda", dtype=torch.bfloat16):
                loss = model(x.cuda(non_blocking=True), y.cuda(non_blocking=True))
            opt.zero_grad(set_to_none=True)
            scaler.scale(loss).backward()
            scaler.step(opt); scaler.update(); sched.step()
            step += 1
            if not torch.isfinite(loss):
                raise RuntimeError(f"non-finite loss at step {step}")
            if step % cfg.ckpt_every == 0:
                save_ckpt(ckpt_path, step, model, opt, sched, scaler, epoch)
            if step >= cfg.max_steps:
                break
        epoch += 1

Three points. Writing to a temporary file and renaming means a crash mid-save leaves the previous checkpoint intact; on object storage, write a new key and update a small pointer object last. Skipping batches to the saved position is simple but rereads data; for long epochs, save the loader's position or use a resumable data loader instead. And with bfloat16 autocast the gradient scaler is not strictly required; it matters for float16, and keeping it in the checkpoint costs nothing.

Determinism and the noise floor

Bitwise reproducibility on GPUs is possible but costs speed. Seed Python, NumPy and PyTorch; call torch.use_deterministic_algorithms(True), which raises an error when an operation has no deterministic implementation; set the CUBLAS_WORKSPACE_CONFIG environment variable that PyTorch's documentation requires for deterministic cuBLAS; and fix the number of data workers. Even then, a different GPU model, driver or world size can change floating-point summation order.

Most teams aim for a weaker, testable property: reruns land within a measured tolerance. Run the same manifest three times with different seeds, record the spread of the main metric, and treat that spread as the noise floor. A candidate that beats the baseline by less than the noise floor has not beaten it.

Distributed failure and checkpoint frequency

In synchronous data-parallel training, every worker waits for every other at each gradient all-reduce, so one lost process stops the job. Launchers such as torchrun can restart workers and rejoin, but the model returns to the last checkpoint, so checkpoint frequency sets how much work a failure costs. Checkpoint too often and you spend the run writing; too rarely and failures waste hours. Young's approximation gives the interval that minimises expected waste: the square root of two times the checkpoint cost times the mean time between failures.

Worked example: a job on 32 GPUs writes a checkpoint in 90 seconds, and the fleet's history says such a job sees a failure about every 12 hours (43,200 seconds). The interval is the square root of 2 times 90 times 43,200, which is about 2,790 seconds, roughly 46 minutes. At that rate checkpointing costs about 3 percent of wall time and a failure costs on average half an interval, about 23 minutes, plus restart time. If the checkpoint is written asynchronously, so training continues while a copy uploads, the blocking cost falls and the optimal interval shrinks with it. Keep at least the last two checkpoints in case the newest is corrupt.

Evaluation gates and registration

Training ends with a candidate, not a model. The evaluation step loads the final weights in a separate task and compares them with the current production model on the frozen test set: the headline metric, calibration, and the same metrics per slice, such as region, customer tier or input length. Gates are written in the manifest and evaluated by code, so promotion is not a judgment made at the end of a long day. A typical rule: no worse than the baseline minus the noise floor overall, and no slice drops more than a set amount. A pass registers the model with its manifest hash, data snapshot, metrics and evaluation report; a fail blocks registration and reports why. Evaluation frameworks covers how to build the suites.

Worked through with the manifest above: three baseline reruns gave AUC 0.861, 0.863 and 0.862, so the noise floor is about 0.002, which is why the gate allows a delta of -0.002. The candidate scores 0.866 overall, a clear pass. Per slice, enterprise customers rise from 0.88 to 0.89, but customers in their first 30 days fall from 0.79 to 0.77, a drop of 0.02 against a limit of 0.01. The gate blocks the release even though the headline improved. Investigation finds a new tenure feature that is mostly missing for new accounts and was filled with zero; filling it with an explicit missing indicator restores the slice. Without slice gates that regression would have reached the customers most likely to churn.

Failure modes

Operations and trade-offs

Dashboard each run with step time, the fraction of it spent waiting for data, accelerator utilisation, throughput in samples or tokens per second, loss, gradient norm, checkpoint duration and time since the last good checkpoint. Alert on a stalled step counter, which is how a hung collective shows up, on non-finite loss and on throughput dropping below a fraction of the run's own median. Record cost per run alongside metrics so the team can see what an experiment cost.

Trade-offs: a heavyweight orchestrator adds setup but makes steps rerunnable and auditable; plain scripts are faster to start and fine until the second engineer joins. Synchronous checkpoints are simpler; asynchronous ones save time but need care to avoid copying a state that is still changing. Strict determinism helps debugging but slows training, so many teams enable it only in a debugging configuration.

What to do next

Apply this to your next run:

1. Write a run manifest with a data snapshot id, commit, image digest, config and seed, and hash it.
2. Add validation and leakage checks that fail the run before training starts.
3. Measure batch-wait time; fix the input pipeline if it exceeds a few percent.
4. Save full training state atomically and test resume by killing a run on purpose.
5. Compute your checkpoint interval from measured checkpoint cost and failure rate.
6. Run the baseline three times to measure the noise floor.
7. Encode evaluation gates, including per-slice gates, and register only passing models with lineage.