Machine Learning Pipelines: From Data Ingestion to Model Deployment

A comprehensive intermediate-level guide to designing, building, and operating production-grade machine learning pipelines from raw data to deployed, monitored models.

Table of Contents

• Chapter 1: Foundations of ML Pipelines and MLOps
• Chapter 2: Data Ingestion: Sources, Formats, and Patterns
• Chapter 3: Data Validation, Cleaning, and Quality
• Chapter 4: Feature Engineering and Feature Stores
• Chapter 5: Data and Pipeline Versioning
• Chapter 6: Pipeline Orchestration Frameworks
• Chapter 7: Model Training Infrastructure and Distributed Training
• Chapter 8: Experiment Tracking and Hyperparameter Tuning
• Chapter 9: Model Evaluation, Validation, and Testing
• Chapter 10: Model Packaging, Registry, and Versioning
• Chapter 11: Model Deployment Patterns: Batch, Online, and Edge
• Chapter 12: Serving Infrastructure: Latency, Throughput, and Scalability
• Chapter 13: Monitoring, CI/CD, and Production Operations

Chapter 1: Foundations of ML Pipelines and MLOps

Learning Objectives

• Define an ML pipeline and distinguish it from a one-off training script or exploratory notebook.
• Explain the MLOps maturity model and articulate the gap between research-grade ML and production-grade ML.
• Identify the major stages of an end-to-end ML pipeline and describe how they interconnect through a directed acyclic graph (DAG).
• Compare ML pipelines to traditional software CI/CD pipelines along the dimensions of artifacts, testing, triggers, and feedback loops.
• Recognize the canonical failure modes — technical debt, train-serve skew, silent data dependency changes — that motivate disciplined pipeline engineering.

What is an ML Pipeline?

From Recipe to Restaurant

Imagine a home cook who develops a fantastic chocolate cake recipe. Standing in their own kitchen, they measure flour by feel, taste the batter as they go, and pull the cake from the oven when “it smells right.” The result is delicious — once. Now imagine that cake needs to be produced 10,000 times a day, at consistent quality, across forty franchise locations, by staff who have never met the original cook. The recipe alone is no longer enough. The franchise needs measured ingredients in standardized packages, calibrated ovens, written procedures, quality-control inspections, a supply chain that delivers fresh inputs, and a feedback system that catches it when an oven starts running hot.

This is exactly the difference between a one-off training script and an ML pipeline. A pipeline is the industrialized version of an experimental workflow — the set of automated, orchestrated, repeatable steps that take data from its source and turn it into a deployed, monitored, retrainable model.

More formally, an ML pipeline is “a reusable, orchestrated DAG of stages (ingest → validate → feature-engineer → train → evaluate → validate → deploy → monitor), as opposed to a one-off training script that bundles all logic ad hoc in a notebook” [Source: https://eitt.academy/knowledge-base/mlops-in-practice-from-jupyter-to-production/]. Each stage is a clearly defined unit with known inputs and outputs; the orchestrator wires the stages together, runs them on schedule or on demand, retries on transient failure, and records what happened for later inspection.

Pipeline vs. Script vs. Notebook

It helps to put these three artifacts side by side. A notebook is an interactive document — the data scientist’s laboratory bench. A script is a packaged, command-line-runnable program. A pipeline is a directed graph of such programs (or containers) executed by an orchestrator with versioning, retries, and metadata.

Notebooks are “optimized for exploration and prototyping” but tend “to produce code that is linear, stateful, and interleaved with analysis and visualization, rather than modular, testable, and production-ready” [Source: https://www.ekascloud.com/our-blog/from-notebooks-to-production-the-hard-truth-about-deploying-ml/3598]. A script is a step up — it can be scheduled with cron and accept parameters — but it is still a monolith. Once you ask “Which data did we train on last Tuesday?” or “Why did this morning’s run fail at the feature step?” you find yourself needing the affordances that only a pipeline provides.

Why Production ML Requires Pipelines

Why not just take a well-written script, slap a cron job on it, and call it done? Because production ML carries a unique combination of properties that ad hoc scripts cannot handle gracefully:

1. Data is alive. Inputs change continuously, often silently. A column that contained “USD” yesterday might contain “USD,EUR” today because an upstream service added multi-currency support without telling anyone.
2. Models degrade. Even if code never changes, model behavior decays as the world changes — what statisticians call concept drift. Industry experience suggests that many ML incidents trace not to model code changes but to upstream data changes [Source: https://super.ai/blog/7-costly-surprises-of-machine-learning-part-four].
3. Reproducibility is non-trivial. Recreating yesterday’s model requires recreating yesterday’s data snapshot, hyperparameters, random seeds, code version, and library versions — all simultaneously.
4. Train-serve skew lurks everywhere. The same preprocessing must run identically in training (often offline, batch) and inference (often online, single-record). A mismatch creates a silent, hard-to-diagnose performance regression.
5. Stakeholders are heterogeneous. Data scientists, data engineers, ML engineers, DevOps, security, compliance, and business owners all touch the system.

The orchestrated pipeline is the technical artifact that lets a team manage all five at once.

Stakeholders Along the Pipeline

Unlike traditional software, where a typical project sits comfortably between developers and operators, ML pipelines span an unusually wide range of disciplines. The MLOps literature emphasizes that success “depends on cross-functional collaboration among data scientists, data engineers, DevOps engineers, ML engineers, software developers, and business stakeholders” [Source: https://www.missioncloud.com/blog/10-mlops-best-practices-every-team-should-be-using].

A useful analogy: the ML pipeline is a factory floor. The data engineer is the parts supplier; the data scientist is the design engineer; the ML engineer is the manufacturing engineer; DevOps runs the building; the business owner sets quality targets; compliance is the safety inspector. Each role hands off well-defined artifacts to the next stage.

Key Takeaway: An ML pipeline is a versioned, orchestrated graph of stages that turns raw data into a deployed, monitored, retrainable model — fundamentally different from a notebook (exploratory) or a script (single-purpose automation). Pipelines exist because production ML has live data, decaying models, heterogeneous stakeholders, and rigorous reproducibility needs that no single script can satisfy.

The MLOps Discipline

Origins: DevOps Meets the Data World

To understand MLOps, start with the practice it generalizes. DevOps emerged in the late 2000s as a cultural and technical movement that broke down the wall between developers (who wrote software) and operators (who ran it). Its central insight: if you automate the path from a developer’s commit to a running production service — with continuous integration (CI), continuous delivery (CD), monitoring, and infrastructure-as-code — you can ship faster and more reliably at the same time.

DataOps applied a similar philosophy to data engineering: treat data pipelines as products, version them, test them, monitor them. MLOps is the synthesis: it “applies DevOps-style automation, testing, and monitoring to the entire machine learning lifecycle so that models can be reliably taken from experimentation into production and kept working over time” [Source: https://aws.amazon.com/what-is/mlops/].

Across major vendors, MLOps is defined in nearly identical terms — “a set of practices that unify ML development (Dev) and deployment/operations (Ops) to automate and standardize ML workflows” [Source: https://www.ibm.com/think/topics/mlops]. Crucially, it extends beyond classic DevOps by treating data, models, and experimental configurations as first-class objects to be versioned, tested, deployed, and monitored [Source: https://www.databricks.com/blog/what-is-mlops].

The Google MLOps Maturity Model: L0, L1, L2

Google’s widely-cited maturity model gives teams a vocabulary for “how production-ready is our ML?” It defines three levels.

[Source: https://www.databricks.com/blog/what-is-mlops]

Level 0 is what most teams start with: a data scientist runs a notebook on their laptop, exports a .pkl file, and hands it to an engineer who copies it into a service. There is no orchestrator, no automated testing, no drift detection. This may be acceptable for a low-stakes, infrequently-retrained model (think: an annual demand-forecasting exercise), but for anything customer-facing it is fragile.

Level 1 is where many teams aspire. The training workflow itself is automated — refactored from a notebook into Python modules (ingest.py, features.py, train.py, evaluate.py) and wired together as an Airflow DAG that runs on a schedule [Source: https://www.youtube.com/watch?v=7Xjrp9j9bLw]. The pipeline can be re-run with new data on demand. However, the software lifecycle around that pipeline — testing changes to feature code, deploying new pipeline versions — is still largely manual.

Level 2 is full production-grade ML: every change to pipeline code triggers automated tests (including data tests and “train-on-sample” smoke tests), passing tests trigger automated deployment, and the system continuously monitors production for drift. When drift crosses thresholds, a retraining pipeline kicks off automatically, and if the new model passes evaluation gates, it is promoted with safe deployment patterns like canary releases [Source: https://www.databricks.com/blog/what-is-mlops].

Worked Example: A Churn Model Climbing the Maturity Curve

Let’s follow a churn-prediction model through all three levels.

At Level 0, the data scientist Anna builds the model in a notebook. Each quarter she:

1. SQL-exports last quarter’s customer data to CSV.
2. Runs all twenty-seven cells of churn_v3_final_FINAL.ipynb.
3. Emails model.pkl to engineer Bob.
4. Bob copies it onto the production server and restarts the API.

If predictions look off the next morning, nobody can tell whether it was a code change, a data change, or just statistical noise. Re-creating the exact model in two months will be nearly impossible.

At Level 1, Anna and Bob refactor. The notebook becomes a four-task Airflow DAG: ingest → features → train → evaluate. Every nightly run logs metrics to MLflow, and the trained model lands in a registry tagged with the data snapshot and code commit. Anna reviews the dashboard weekly; if metrics look good, she asks Bob to deploy. Errors now manifest as task-level Airflow failures — “feature step failed because column region was renamed to state” — instead of a wall of red text inside a notebook.

At Level 2, every pull request to the pipeline repository triggers a CI run that executes unit tests on the feature code, a data validation pass on a sample dataset, and a “train-on-sample” smoke test confirming that training converges and metrics fall within expected bounds. If CI passes, a CD pipeline builds a new Docker image and rolls it out to the orchestrator. In production, an Evidently dashboard tracks input distributions; when the Population Stability Index for any top-10 feature exceeds 0.2, a retraining pipeline fires. The candidate model is shadow-deployed for 48 hours, then canary-released at 5% → 25% → 100% as long as A/B metrics stay healthy [Source: https://www.qwak.com/post/shadow-deployment-vs-canary-release-of-machine-learning-models]. Rollback is one button click.

The point is not that every team must reach Level 2 immediately. The maturity model is a roadmap, not a checklist. Many production systems are healthy and stable at Level 1; only systems where data changes rapidly, stakes are high, or many models must be maintained truly demand Level 2.

Figure 1.2: Google MLOps maturity model (L0 to L2)

flowchart TD
    L0["Level 0: Manual ML<br/>Notebooks, .pkl handoffs<br/>No orchestration, no tests"]
    L1["Level 1: Pipeline Automation<br/>Orchestrated DAG (Airflow/KFP)<br/>Scheduled retraining, basic monitoring"]
    L2["Level 2: CI/CD + Continuous Training<br/>Automated tests, canary deploys<br/>Drift-triggered retraining"]
    L0 ==>|"Refactor notebook<br/>into DAG"| L1
    L1 ==>|"Add CI/CD,<br/>drift monitoring,<br/>auto-retraining"| L2
    L2 -.->|Roadmap, not checklist| L1

Common Failure Modes

Why does MLOps as a discipline even exist? Because ML systems fail in distinctive, often invisible ways. The most chronic failure modes include:

• Train-serve skew. Training preprocesses data one way (offline, batch, pandas); serving preprocesses it another way (online, single record, a different library). Predictions degrade silently. Feature stores exist primarily to eliminate this category of bug [Source: https://www.snowflake.com/en/fundamentals/feature-store/].
• Silent data dependency changes. An upstream team changes the semantics of a column; your model’s accuracy drops, but no alarm rings because no code changed [Source: https://super.ai/blog/7-costly-surprises-of-machine-learning-part-four].
• Hidden technical debt. ML systems “have a special capacity for incurring technical debt, often in hidden forms that compound over time” — Sculley et al. describe ML as a “high-interest credit card” of debt [Source: https://papers.neurips.cc/paper/5656-hidden-technical-debt-in-machine-learning-systems.pdf].
• Feedback loops. A recommender that only shows users content from genres they have already engaged with eventually trains on data that reflects only its own choices — a self-reinforcing rut [Source: https://super.ai/blog/7-costly-surprises-of-machine-learning-part-four].
• Reproducibility gaps. Without versioned data, code, environment, and seeds, “which model is in production?” becomes a research project.

The Tooling Landscape

The MLOps tooling ecosystem has exploded. While specific tools come and go, the categories have stabilized.

[Source: https://montecarlo.ai/blog-ml-orchestration-tools/] [Source: https://www.zenml.io/blog/mlflow-vs-weights-and-biases]

A common pattern is to choose one tool from each category and integrate them, often through a cloud platform (Vertex AI on Google, SageMaker on AWS) that bundles many roles. An empirical study of serving frameworks found that TensorFlow Serving and TorchServe outperform general-purpose alternatives like BentoML, MLServer, and MLflow on latency for deep learning workloads, while general-purpose frameworks excel at flexibility and heterogeneous workloads [Source: https://arxiv.org/html/2411.10337v1].

Key Takeaway: MLOps is DevOps generalized to encompass data and models as first-class artifacts. Google’s three-level maturity model — manual workflows (L0), automated pipelines (L1), and full CI/CD + Continuous Training (L2) — gives teams a roadmap. The discipline exists because ML systems exhibit unique failure modes (train-serve skew, silent data dependency changes, feedback loops, hidden technical debt) that traditional software practices were never designed to catch.

Anatomy of an End-to-End Pipeline

The Canonical Stages

Strip away the vocabulary differences between vendors and frameworks, and almost every ML pipeline contains the same eight stages. We will visit each in much more depth in later chapters; this section gives the bird’s-eye view.

1. Data ingestion — get raw data into a stable, queryable form (data lake, warehouse, message stream).
2. Data validation — check schema, distributions, ranges, and quality before spending compute on training.
3. Data preparation / feature engineering — transform validated data into model-ready features.
4. Model training — fit a model on prepared data, often including hyperparameter optimization.
5. Model evaluation — measure performance on held-out data, with slicing across subpopulations.
6. Model validation (gating) — decide if the new model is good enough to ship.
7. Model deployment / serving — package the approved model and expose it for predictions.
8. Monitoring — observe data, model, and system metrics in production; close the feedback loop.

In DAG form, the linear backbone looks like this:

ingest → validate_data → prepare_features → train_model
       → evaluate_model → validate_model → deploy_model → monitor

Figure 1.1: End-to-end ML pipeline stages with feedback loop

flowchart LR
    A[Ingest] --> B[Validate Data]
    B --> C[Prepare Features]
    C --> D[Train Model]
    D --> E[Evaluate Model]
    E --> F{Validate Model}
    F -->|Pass| G[Deploy Model]
    F -.->|Fail| D
    G --> H[Monitor]
    H -.->|Drift / Decay| A

But the edges fan out and feed back. The same ingested data flows to monitoring (for baseline distributions); the same feature engineering logic must be packaged and shipped alongside the model (to prevent train-serve skew); and monitoring can trigger a brand-new run of the entire pipeline (continuous training).

Stage-by-Stage in Brief

[Source: https://developers.google.com/machine-learning/crash-course/production-ml-systems/deployment-testing]

A concrete example helps. Suppose we are training a TFX-style fraud detection pipeline. The TFX components map almost 1:1 to our canonical stages: ExampleGen (ingest), StatisticsGen + SchemaGen + ExampleValidator (validation), Transform (feature engineering), Trainer (training), Evaluator (evaluation), ModelValidator (gating), and Pusher (deployment) [Source: https://developers.google.com/machine-learning/crash-course/production-ml-systems/deployment-testing]. Whether your stack is TFX, Kubeflow, or hand-rolled Airflow, the shapes are the same.

DAG Orchestration: Why a Directed Acyclic Graph?

The DAG is not arbitrary jargon. It is the most natural data structure for representing “these tasks have these dependencies.” Three properties matter:

1. Directed — edges have arrows. The train task can only start after prepare_features has finished, not before.
2. Acyclic — no cycles. You cannot have train → evaluate → train as a single DAG, because that would never terminate. (Retraining loops are implemented at a higher level: the monitoring pipeline triggers a new run of the training pipeline.)
3. Graph — many edges, not just one chain. Two tasks with no dependency can run in parallel.

Each node is a task — typically a containerized program. Each edge is a data or metadata dependency: “task B needs an artifact that task A produces.” The orchestrator handles the rest: scheduling, retries, logging, and emitting metadata to a lineage store.

Figure 1.3: DAG orchestration with parallel branches and conditional gating

flowchart TD
    Ingest((Ingest)) --> Validate[Validate Data]
    Validate --> Features[Prepare Features]
    Features --> Train[Train Model]
    Features --> Baseline[Train Baseline]
    Train --> Eval[Evaluate Candidate]
    Baseline --> Eval
    Eval --> Gate{Metrics > baseline?}
    Gate -->|Yes| Deploy[Deploy]
    Gate -.->|No| Stop((Stop))

Here is a minimal Kubeflow Pipelines example showing how the DAG is wired:

import kfp
from kfp import dsl

@dsl.component
def ingest_op(output_path: str) -> str:
    # write data to output_path
    return output_path

@dsl.component
def validate_data_op(data_path: str) -> str:
    # run validation, return report path
    return "gs://.../validation_report.json"

@dsl.component
def train_model_op(data_path: str, report_path: str) -> str:
    # train, save model
    return "gs://.../model"

@dsl.pipeline
def ml_pipeline():
    ingest = ingest_op(output_path="gs://.../raw")
    validate = validate_data_op(data_path=ingest.output)
    train = train_model_op(data_path=ingest.output,
                           report_path=validate.output)

[Source: https://montecarlo.ai/blog-ml-orchestration-tools/]

Notice that the DAG is implicit — it is built by tracking which outputs flow into which inputs. The KFP backend constructs a graph where validate depends on ingest, and train depends on both. The same pattern applies in Airflow (with explicit >> operators), TFX (via component output wiring), and Vertex AI Pipelines (built on KFP).

The DAG is structural backbone of ML pipelines: edges encode artifact/metadata dependencies (not time), enabling parallelism, retries, lineage tracking, and conditional branching such as “only deploy if validation gates pass.”

Triggers: What Starts a Pipeline?

Traditional CI/CD pipelines have one trigger: a git push. ML pipelines have several.

In a Level 2 system, several of these triggers can fire the same pipeline. The orchestrator records why each run was triggered — invaluable forensic information when a particular model version misbehaves in production.

Artifacts and Metadata: The Pipeline’s Memory

Every pipeline run produces artifacts. In ML, these are far broader than software binaries: code, raw and processed data, feature definitions, model weights and checkpoints, hyperparameters, evaluation reports, and governance documents — each requiring versioning and lineage. H2O.ai defines an ML artifact as any output created by the training process, including fully trained models, checkpoints, and intermediate files [Source: https://h2o.ai/wiki/artifacts/].

Equally important is the metadata that links these artifacts: which data snapshot produced which features, which features fed which trained model, which model achieved which metrics, which model is currently deployed and serving which fraction of traffic. This web of relationships is the lineage graph, and storing it in a queryable form (e.g., ML Metadata in TFX, Vertex ML Metadata, or a custom database) is what enables debugging questions like “What changed between yesterday’s model and today’s?”

A real-world analogy: lineage is the chain-of-custody document in a forensics lab. Without it, you have evidence; with it, you have evidence you can defend in court.

Key Takeaway: An end-to-end ML pipeline is a DAG of eight canonical stages — ingest, validate, prepare features, train, evaluate, validate model, deploy, monitor — orchestrated to produce and link artifacts (data, models, metrics) with full lineage. The DAG structure enables parallelism, retries, and conditional logic, while the rich set of triggers (code, schedule, data, drift, performance) sets ML pipelines apart from one-trigger software CI/CD.

ML Pipelines vs Software CI/CD

A Synthesis Comparison

We have already touched on most of the differences; this section gives them a unified treatment. The fundamental shift is from systems whose behavior is fully specified by code to systems whose behavior emerges from data and learning algorithms.

[Source: https://valohai.com/cicd-for-machine-learning/] [Source: https://www.wwt.com/blog/mlops-cicd-ct-whats-continuous-training]

ML pipelines do not replace CI/CD — they superimpose new layers on top of it. Every good ML pipeline still needs a healthy software CI/CD foundation. What changes is that the foundation must be augmented to handle data and learned behavior as first-class concerns.

Data as a First-Class Artifact

In traditional CI/CD, the build is a pure function of code: same source plus same toolchain yields the same binary, bit-for-bit. In ML, the function is impure: same training code plus same hyperparameters but different data yields a different model.

This forces data into the artifact universe. Weights & Biases and adjacent literature emphasize that “version control for datasets and ML models is as essential as for source code, providing traceability, reproducibility, rollback, debugging support, and collaboration” [Source: https://wandb.ai/site/articles/intro-to-mlops-data-and-model-versioning/]. Tools such as DVC (Data Version Control) integrate dataset versioning into Git workflows, storing pointers and metadata in the repository while data lives in cloud storage [Source: https://github.com/treeverse/dvc].

A useful analogy: in software engineering, code is the source and the binary is the build output. In ML, both code and data are the source, and the model is the build output. If you do not version both inputs, you cannot reproduce the output.

Continuous Training Alongside CI/CD

In MLOps, two intertwined loops characterize the continuous process [Source: https://valohai.com/cicd-for-machine-learning/]:

• The software CI/CD loop is triggered by code commits and builds/tests/deploys pipeline code, training code, and serving code.
• The model CT loop is triggered by data signals (new data, drift, performance degradation) and reruns the training pipeline using the current codebase and fresh data, producing new model artifacts that are evaluated and potentially deployed.

These loops share infrastructure — the same orchestrator, the same registries, often the same tests — but they have different triggers, different cadences, and different stakeholders.

Figure 1.4: Two intertwined loops - CI/CD plus Continuous Training

flowchart TD
    subgraph CICD["Software CI/CD Loop"]
        Commit[Git Commit] --> Build[Build & Test]
        Build --> DeployCode[Deploy Pipeline Code]
    end
    subgraph CT["Continuous Training Loop"]
        Signal["Data Signal:<br/>new data / drift / decay"] --> Retrain[Retrain on Fresh Data]
        Retrain --> EvalCT[Evaluate Model]
        EvalCT --> DeployModel[Deploy Model Artifact]
    end
    DeployCode -.->|Updates pipeline used by| Retrain
    DeployModel -.->|Production metrics feed| Signal

This is what people mean when they say MLOps adds an axis to DevOps. Traditional CD asks: “Is the new code safe to deploy?” CT asks an additional question: “Has the world changed enough that we need a new model?”

Testing: From Deterministic to Probabilistic

Software tests are nearly always deterministic. assert add(2, 3) == 5 either passes or fails, and the answer never changes. ML testing is fundamentally probabilistic.

Consider the testing pyramid in each world:

A particularly important difference: A/B testing. In traditional software, an A/B test is one technique among many; in ML, it is often the primary way to validate a new model against the incumbent in production-like conditions. Teams define an Overall Evaluation Criterion (OEC) such as click-through rate, choose a minimum detectable effect size and acceptable error rates, compute the required sample size, and let the experiment run until it has the statistical power to make a confident call [Source: https://mlops.community/blog/the-what-why-and-how-of-a-b-testing-in-ml].

Reproducibility: Why Models Are Not Just Code

Reproducing a software build requires: source code, build toolchain, dependency versions. That is hard but tractable.

Reproducing an ML training run requires all of that, plus:

• The exact data snapshot (partition, version, or hash).
• The random seed for splits, initialization, and any stochastic optimization.
• The exact preprocessing pipeline, including any state (e.g., a normalizer’s fitted mean and standard deviation).
• The exact hardware (GPUs can have nondeterministic kernel orderings).
• The exact library versions (NumPy, scikit-learn, TensorFlow — minor versions of any of these can change numerical results).

Even with perfect rigor, exact reproducibility may be impossible because of nondeterminism in GPU kernels or in distributed training. The practical goal is therefore statistical reproducibility — getting models that are equivalent in behavior given the same inputs, not necessarily bit-for-bit identical.

This is one reason ML pipelines lean heavily on container images. A container freezes the environment; combined with versioned data and seeded code, it gets you the closest thing to a software build’s reproducibility guarantee.

Hidden Technical Debt in ML

Sculley et al.’s paper “Hidden Technical Debt in Machine Learning Systems” is required reading for the field. They liken ML to a “high-interest credit card” of technical debt — it enables rapid development of complex systems, but the resulting systems can be fragile and expensive to maintain in the long run [Source: https://papers.neurips.cc/paper/5656-hidden-technical-debt-in-machine-learning-systems.pdf].

They enumerate several ML-specific debt categories that traditional software does not have direct analogues for:

Data dependency debt is uniquely dangerous because upstream data schemas and semantics can change silently without any code change, so data lineage tracking and schema validation are essential for production ML [Source: https://datahub.com/blog/data-lineage-for-ml/].

Why Models Are Not Just Code: A Final Synthesis

If we had to compress the difference between traditional CI/CD and ML pipelines into a single insight, it would be this: in software, behavior is determined by code; in ML, behavior is determined by code interacting with data. Every consequence flows from that. Data must be versioned. Tests must be probabilistic. Triggers must include data signals, not just code commits. Monitoring must include data and model metrics, not just availability. Deployment must include shadow and canary stages to manage probabilistic risk. Governance must include model cards and audit trails because the model’s behavior is not transparent from the source.

This shift from code-determined to data-determined behavior is not a minor extension. It rewires every part of the development lifecycle — and motivates everything we will build in the rest of this book.

Key Takeaway: ML pipelines do not replace software CI/CD; they superimpose data and model concerns on top of it. The result is a system with two intertwined loops (CI/CD plus Continuous Training), broader artifacts (data and models alongside code), probabilistic rather than deterministic tests, and a uniquely insidious set of technical debt modes. The single defining shift is that ML behavior is determined by code and data — and almost every difference between ML pipelines and software pipelines flows from that fact.

Chapter Summary

An ML pipeline is the production-grade industrialization of an experimental ML workflow: a versioned, orchestrated graph of stages that transforms raw data into a deployed, monitored, retrainable model. It differs from a notebook (which is built for exploration) and a script (which automates a single task) by virtue of its modular stages, lineage tracking, automated retries, and ability to coordinate across the many stakeholders — data engineers, data scientists, ML engineers, DevOps, business owners, and compliance — that production ML inevitably requires.

MLOps is the discipline that makes such pipelines work. It generalizes DevOps by treating data, models, and experimental configurations as first-class objects alongside code. Google’s three-level maturity model gives teams a vocabulary for assessing themselves: Level 0 is manual ML; Level 1 automates the training pipeline itself; Level 2 layers full software CI/CD plus Continuous Training on top, with code- and data-driven triggers, automated testing, and monitoring-driven retraining. The discipline exists because ML systems exhibit failure modes — train-serve skew, silent data dependency changes, feedback loops, and the rich taxonomy of hidden technical debt described by Sculley et al. — that traditional software practices were never designed to catch.

The canonical end-to-end pipeline is a DAG of eight stages: ingest, validate, prepare features, train, evaluate, validate model, deploy, monitor. Orchestrators like Airflow, Kubeflow Pipelines, TFX, and Vertex AI Pipelines render this graph executable, manage retries and metadata, and link artifacts via lineage. Compared to software CI/CD, ML pipelines have broader artifacts (data and models alongside code), more diverse triggers (code, schedule, data drift, performance), probabilistic rather than binary tests, and an additional continuous process — Continuous Training — that retrains models in response to changing data even when no code has changed. The single load-bearing insight that ties it all together is that ML behavior is determined by code interacting with data, and almost every distinguishing feature of ML pipelines flows from that fact. The remaining chapters of this book will take each stage of the canonical pipeline in turn, showing how to build, test, deploy, and operate it with the rigor that production ML demands.

Key Terms

Chapter 2: Data Ingestion: Sources, Formats, and Patterns

Every machine learning system begins with the same problem: getting data from the place it was created to the place a model can learn from it. This sounds mundane, but ingestion is where most pipeline failures originate. A subtle schema change in an upstream microservice, a Kafka consumer that double-counts on retry, a Parquet file partitioned by the wrong column - any of these can silently degrade model quality for weeks before anyone notices. This chapter examines the sources from which ML pipelines draw data, the two dominant ingestion paradigms (batch and streaming), the file formats used to persist training data, and the reliability patterns that keep ingestion correct under failure.

If Chapter 1 framed the ML pipeline as a factory, ingestion is the loading dock. The choice of trucks (batch vs. streaming), the shape of the crates (file formats), and the receiving procedures (idempotency, schema validation) determine whether the rest of the factory can actually operate.

Section 1: Data Sources for ML

ML pipelines rarely consume data from a single tidy source. A production fraud-detection model might pull transaction events from Kafka, account state from a Postgres database via change data capture, merchant metadata from a nightly S3 export, and risk-list updates from a third-party REST API. Each source has a native shape, a natural cadence, and a preferred ingestion mechanism. Understanding these properties up front prevents the common mistake of forcing every source through the same ingestion path.

Figure 2.1: Heterogeneous data sources feeding an ML ingestion layer

flowchart LR
    A[Postgres OLTP] -->|CDC via Debezium| E[Ingestion Layer]
    B[S3 Data Lake] -->|Batch read| E
    C[Kafka Clickstream] -->|Stream consumer| E
    D[Third-party REST API] -->|Scheduled pull| E
    E --> F[(Feature Store)]
    E --> G[(Training Lake)]

Relational Databases and Change Data Capture

Operational systems - order management, user accounts, billing, CRM - typically live in relational databases. These systems are the source of truth for entity state: who the user is, what their current balance is, which tickets are open. For ML, this state matters because it often dominates the feature vector.

Two strategies exist for pulling data out of an OLTP database. The simplest is a periodic full-table or incremental query, scheduled by Airflow or a similar orchestrator and executed by Spark or a SQL engine. This works for small tables and tolerant latency budgets, but it scales poorly: nightly scans of a billion-row orders table waste compute when only a small percentage of rows changed, and they put load on a database that is also serving production traffic.

Change Data Capture (CDC) solves this by reading the database’s transaction log directly. Tools like Debezium tail the MySQL binlog or the Postgres write-ahead log (WAL), translate each insert, update, and delete into a structured event, and publish those events to Kafka [Source: https://www.databricks.com/blog/what-feature-store-complete-guide-ml-feature-engineering]. Downstream consumers then process row-level changes in near real time without ever issuing a query against the production database.

CDC is particularly valuable for ML because it preserves source-of-truth semantics while enabling incremental updates. Slowly changing dimensions - user profile attributes, KYC flags, account status - flow as event streams that can be materialized into both an offline feature store (for training history) and an online store (for serving) [Source: https://arize.com/blog/feature-store/]. The tradeoff is operational complexity: log access permissions, schema evolution when DBAs add columns, and the need for periodic full re-syncs to recover from gaps.

Object Storage, Data Lakes, and Lakehouses

The second great reservoir of ML data is object storage - S3, Google Cloud Storage, Azure Data Lake Storage - typically organized as a data lake. Files arrive from batch ETL jobs, third-party data providers, log shippers, or older data warehouses that periodically export. A data lake is permissive: any team can drop any file in any format. A lakehouse adds the missing structure on top - a table format like Delta Lake, Apache Iceberg, or Apache Hudi that gives a collection of Parquet files transactional semantics, schema evolution, and time travel [Source: https://www.databricks.com/blog/what-feature-store-complete-guide-ml-feature-engineering].

For ML, the lakehouse is usually the canonical home of historical training data. A typical layered design uses bronze (raw ingested events), silver (cleaned, deduplicated), and gold (feature-ready aggregates) zones. Models train on the gold layer; reproducibility comes from snapshotting the table at a known version.

Event Streams: Kafka and Kinesis

User clicks, IoT telemetry, application logs, ad impressions - any high-volume continuous event flow lands naturally in an event streaming platform. Apache Kafka is the de facto standard in self-managed and on-prem deployments; AWS Kinesis, Google Pub/Sub, and Azure Event Hubs are the managed equivalents. These systems treat data as an append-only log partitioned across brokers, with durable replication and the ability for many consumers to read independently at their own pace.

For ML, event streams matter for two reasons. First, they are the substrate for real-time features: a recommendation system computing “items viewed in the last five minutes” reads directly from a clickstream topic [Source: https://aws.amazon.com/blogs/machine-learning/use-streaming-ingestion-with-amazon-sagemaker-feature-store-and-amazon-msk-to-make-ml-backed-decisions-in-near-real-time/]. Second, they serve as the durable buffer between source systems and downstream processors - a fraud pipeline can fall behind for an hour during a deploy without losing data, because Kafka retains messages for days [Source: https://www.youtube.com/watch?v=WvdLydIAD44].

APIs, Logs, and Sensors

Beyond databases and streams, ML pipelines often ingest from external HTTP APIs (weather, FX rates, third-party risk scores), application log files (typically shipped via Fluentd, Logstash, or a cloud agent), and IoT sensor feeds (frequently MQTT before crossing into Kafka). Each requires its own connector: API ingestion needs rate-limit awareness and retry logic; logs need parsing and timestamp normalization; sensor data needs gap detection and clock-skew handling.

Key Takeaway: Real ML pipelines blend multiple source types - OLTP databases via CDC, lakes via batch reads, event streams via Kafka, and APIs via scheduled pulls. Match the ingestion mechanism to the source’s native cadence rather than forcing one paradigm on all sources.

Section 2: Batch vs. Streaming Ingestion

Once you know where data lives, the next question is how often to move it. The answer divides cleanly into two paradigms - batch and streaming - and a handful of hybrid architectures that combine them. The decision is driven by one variable above all: how stale can features be before model quality suffers?

Batch Patterns

Batch ingestion moves data in discrete, scheduled bulk loads - hourly, nightly, weekly. The pattern is mature: an orchestrator like Airflow, Prefect, or Dagster triggers a Spark or SQL job on a cron schedule; the job reads a delta from the source, transforms it, and writes to a destination table [Source: https://www.databricks.com/blog/what-feature-store-complete-guide-ml-feature-engineering]. Latency is measured in minutes to hours.

Batch fits ML in three common situations. First, building large historical training sets from a lake or warehouse, where you need weeks or months of data and tolerate that the model is trained on slightly stale data [Source: https://arize.com/blog/feature-store/]. Second, computing aggregated offline features - 30-day customer spend, 7-day click count, lifetime value - that change slowly and are too expensive to recompute on every event [Source: https://datavidhya.com/learn/de-system-design/question-breakdowns/feature-store-ml/]. Third, backfills and re-training, where a new feature definition needs to be applied to historical data going back months.

The strengths of batch are simplicity, debuggability, and economy of scale. A failed Spark job can be re-run on a fixed input. Throughput on a well-tuned cluster is enormous. The weakness is staleness: a feature computed at 2 a.m. is six hours old by 8 a.m.

Streaming with Kafka and Flink

Streaming ingestion processes events continuously as they arrive. A producer publishes to Kafka or Kinesis; a stream processor - Apache Flink, Spark Structured Streaming, or Kafka Streams - consumes events, applies transformations (filtering, joining, windowed aggregations), and writes results to an online store [Source: https://www.youtube.com/watch?v=WvdLydIAD44]. End-to-end latency targets are typically seconds, sometimes tens of milliseconds.

Streaming is the right choice when models must react to events within the same session: real-time fraud scoring on the most recent card swipe, recommendation features built from clicks in the current visit, dynamic pricing that responds to current traffic. Online feature stores - Redis, DynamoDB, Cassandra, Aerospike - hold the latest feature value per entity and serve it to the inference layer at single-digit-millisecond latency [Source: https://aerospike.com/blog/feature-store/].

Streaming pays for low latency with complexity. Out-of-order events, late arrivals, exactly-once semantics, stateful joins across hours of history, and rolling restarts all require careful design. Debugging is harder because you cannot easily “re-run yesterday’s job” - you must replay the source log.

Figure 2.2: Batch vs streaming ingestion paths

flowchart TD
    S[Source Systems] --> B{Latency Budget?}
    B -->|Minutes to hours| BATCH[Batch Path]
    B -->|Seconds| STREAM[Streaming Path]
    BATCH --> AF[Airflow Schedule] --> SP[Spark Job] --> OFF[(Offline Store)]
    STREAM --> KF[Kafka Topic] --> FL[Flink Processor] --> ON[(Online Store)]
    OFF --> M[Model Training]
    ON --> I[Model Inference]

Lambda and Kappa Architectures

Most production ML stacks combine batch and streaming, and two named architectures describe how. Lambda architecture maintains separate batch and speed layers. The batch layer computes accurate historical features and writes them to an offline store; the speed layer computes approximate recent features from a stream and writes them to an online store. At serving time, the two are merged [Source: https://datavidhya.com/learn/de-system-design/question-breakdowns/feature-store-ml/]. The advantage is that each layer uses its optimal tool. The disadvantage is two code paths - one in Spark SQL, one in Flink - which can drift apart in subtle ways.

Kappa architecture takes a different approach. A single streaming pipeline is the source of truth; both historical and real-time processing run through the same code, with history reconstructed by replaying the log. Feature logic is implemented exactly once. The cost is that replaying a year of log to backfill a new feature is expensive and operationally tricky, and many organizations already rely on warehouses, making pure Kappa hard to adopt.

Figure 2.3: Lambda architecture with batch, speed, and serving layers

flowchart TD
    SRC[Source Events] --> BL[Batch Layer]
    SRC --> SL[Speed Layer]
    BL -->|Spark on lake| OFF[(Offline Feature Store)]
    SL -->|Flink on stream| ON[(Online Feature Store)]
    OFF --> SV[Serving Layer]
    ON --> SV
    SV --> APP[ML Application]

In practice, vendor feature stores (Feast, Tecton, SageMaker Feature Store) implement a Lambda-like dual store but mitigate the drift problem by letting users declare feature logic once in a DSL that generates both the batch and streaming pipelines [Source: https://www.qwak.com/post/top-ml-feature-stores].

When to Choose Each

A useful rule of thumb: if the business decision is offline (a weekly model refresh, a quarterly risk report), use batch. If the business decision is per-user, per-event, and the user is waiting, use streaming. If the source of truth is an OLTP database whose changes drive predictions, layer CDC on top of streaming to keep entity state fresh [Source: https://chalk.ai/blog/what-is-a-feature-store].

Key Takeaway: Latency requirements drive the batch-vs-streaming choice. Use batch for historical, aggregate, slow-moving features; use streaming for fresh, per-event features; use Lambda to combine them, with shared feature definitions to avoid logic drift.

Section 3: File Formats for ML Data

The file format you choose for training data is not just a serialization detail - it determines I/O throughput, storage cost, schema-evolution flexibility, and how easily different teams can share the data. ML workloads have particular characteristics (wide rows, repeated reads of the same dataset, mixed read patterns across frameworks) that interact with format choice in non-obvious ways.

Row-based Formats: CSV, JSON, Avro

Row-based formats store each record contiguously. The simplest are CSV and JSON, ubiquitous because every tool can read them, but inefficient: text encoding wastes space, parsing is CPU-heavy, and there is no native schema enforcement. They are fine for ad-hoc small datasets and human inspection; they are wrong for production ML.

Apache Avro is the serious row-based format. Records are encoded as compact binary with the schema stored separately (typically in a Confluent Schema Registry), enabling field-by-field deserialization without a parser per record [Source: https://www.youtube.com/watch?v=yQ2IibGvU9U]. Avro’s defining strength is schema evolution: the writer’s schema and the reader’s schema are reconciled at read time, supporting added fields with defaults, renamed fields via aliases, and other compatible changes. This makes Avro the canonical format for Kafka topics in production - event schemas evolve over years, and Avro plus the registry guarantees consumers do not break when producers add a field.

For ML, Avro is the right format at the raw ingest layer (bronze in lakehouse parlance) but a poor choice for analytical or training reads, because reading any subset of columns still requires loading the entire row.

Columnar Formats: Parquet and ORC

Columnar formats store all values of a single column contiguously, which is transformative for analytical workloads. Apache Parquet and Apache ORC are the two production columnar formats, and they share the same key advantages: high compression (similar values pack well together), predicate pushdown (skip whole row groups based on column statistics), and column pruning (read only the columns the query needs) [Source: https://www.youtube.com/watch?v=yQ2IibGvU9U].

Parquet is the dominant choice in modern lakehouses. It integrates natively with Spark (vectorized reads, predicate pushdown), Trino, Presto, Snowflake external tables, and the entire PyArrow ecosystem. For ML, Parquet is almost always the right format for the silver and gold layers - feature tables, training datasets, and the offline tier of a feature store [Source: https://www.databricks.com/blog/what-feature-store-complete-guide-ml-feature-engineering].

ORC offers similar properties and is slightly more optimized for Hive-centric workloads. In legacy Hadoop stacks, ORC is the natural choice; in greenfield cloud lakehouses, Parquet wins on ecosystem support.

Schema evolution on Parquet and ORC is workable for additive changes (adding a column with a default) but tricky for drops or type changes. In practice, ML teams delegate schema evolution to the table format layered on top: Delta Lake, Iceberg, or Hudi tracks versioned schemas, supports MERGE INTO operations, and enables time travel for training reproducibility [Source: https://www.databricks.com/blog/what-feature-store-complete-guide-ml-feature-engineering].

TFRecord and Petastorm

TFRecord is TensorFlow’s native training format: a sequence of length-prefixed protobuf messages, typically tf.train.Example records, optionally gzip-compressed at the file level. It is row-based and optimized for one specific access pattern: sequential reads with prefetching, shuffling, and interleaving via tf.data.TFRecordDataset [Source: https://www.youtube.com/watch?v=yQ2IibGvU9U]. When the training input pipeline is the bottleneck - typically when GPUs would otherwise sit idle waiting for data - TFRecord can deliver more stable and higher throughput than reading Parquet through a Python adapter.

The cost of TFRecord is significant. The “schema” is defined in your parsing code, not the file, so adding or removing features means updating every reader. Spark integration is awkward, requiring custom input formats. Cross-framework reuse - the same data feeding a PyTorch model and a TensorFlow model - is painful. Most teams should use TFRecord only as a derived training artifact materialized from Parquet for a specific TensorFlow job at scale.

Petastorm (originally from Uber) bridges Parquet and deep-learning frameworks, exposing a Parquet dataset as a streaming PyTorch or TensorFlow dataset with sharding, shuffling, and tensor conversion. For mixed-framework shops on a Parquet lake, Petastorm or similar libraries (NVIDIA DALI, Ray Data) remove much of the motivation for TFRecord.

Compression Tradeoffs

All four formats support multiple compression codecs - Snappy, ZSTD, Gzip, Zlib, LZ4 - with the same general tradeoff: stronger compression (Gzip, ZSTD high level) reduces storage and network cost but raises CPU cost on every read. Snappy and LZ4 are the common defaults for ML training data because read CPU often matters more than storage. ZSTD has emerged as a strong middle ground, offering compression close to Gzip with decompression speed close to Snappy.

Format Comparison

Key Takeaway: Use Avro at the event layer for schema evolution, Parquet (with a table format like Delta or Iceberg) as the canonical lake and offline feature store format, and TFRecord only as a derived artifact for TensorFlow training when measured I/O is the bottleneck.

Section 4: Ingestion Reliability

A pipeline that ingests data correctly 99% of the time is not 1% wrong - it is broken. The 1% manifests as silent feature drift, training-serving skew, or missing labels that degrade model accuracy in ways that are nearly impossible to debug after the fact. Reliable ingestion rests on four pillars: idempotency, schema management, backpressure handling, and lineage.

Idempotency and Exactly-Once Semantics

Distributed systems fail. Network calls time out, brokers restart, consumers get rebalanced, sinks return 5xx. Every reliable ingestion pipeline must assume retries will happen and ensure that processing the same message twice produces the same result as processing it once. This property is idempotency, and it is the practical foundation for what Kafka calls “exactly-once” semantics.

Within Kafka itself, idempotence is configurable. Setting enable.idempotence=true on a producer assigns it a producer ID and tracks sequence numbers per partition, so a retry after an in-flight failure does not duplicate the message [Source: https://aws.amazon.com/blogs/machine-learning/use-streaming-ingestion-with-amazon-sagemaker-feature-store-and-amazon-msk-to-make-ml-backed-decisions-in-near-real-time/]. For read-process-write pipelines that stay inside Kafka (topic to topic), transactional producers go further: transactional.id, initTransactions(), beginTransaction(), sendOffsetsToTransaction(), and commitTransaction() make the output records and the consumer offset commit atomic. Consumers in read_committed mode see only committed transactions.

Figure 2.4: Kafka exactly-once flow across producer, broker, and consumer

sequenceDiagram
    participant P as Producer
    participant B as Kafka Broker
    participant C as Consumer
    P->>B: initTransactions(transactional.id)
    P->>B: beginTransaction()
    P->>B: send(record, seq#)
    B-->>P: ack (dedup via PID+seq)
    P->>B: sendOffsetsToTransaction()
    P->>B: commitTransaction()
    B->>C: deliver (read_committed)
    C->>C: process exactly once

Outside Kafka, transactional guarantees do not extend - a feature store or a data lake cannot participate in a Kafka transaction. The practical pattern is at-least-once delivery from Kafka combined with idempotent writes at the sink. Every event carries a stable identifier (a UUID assigned upstream, or a hash of entity ID and event time), and the sink performs upserts keyed by that identifier. Writing the same event twice overwrites the same row with the same value; the duplicate is invisible downstream.

For lake ingestion, the idiom is MERGE INTO on a Delta, Iceberg, or Hudi table keyed by event ID. For online feature stores, it is SET keyed by (entity_id, feature_name) with a write-time check that ignores updates with timestamps older than the current value, preventing out-of-order events from overwriting fresh data with stale data. Compacted Kafka topics provide a third option: keyed by entity, the broker retains only the latest value per key, achieving deduplication at the storage layer.

A simple analogy: idempotency is like a hotel reservation confirmation number. If your booking app crashes and you click “Reserve” again, the hotel uses the confirmation number to recognize the duplicate and charges you once, not twice. Stable event IDs play the same role for ML ingestion.

Schema Evolution

Source schemas change. A microservice adds a field. A column type widens from int32 to int64. A nullable field becomes required after a backfill. If ingestion treats every change as a breaking change, every minor source update halts the pipeline.

The pattern is a schema registry. Confluent Schema Registry (and its compatible alternatives) stores Avro, Protobuf, or JSON Schema definitions keyed by topic and subject, and enforces compatibility rules - backward, forward, or full - on every new version [Source: https://www.databricks.com/blog/what-feature-store-complete-guide-ml-feature-engineering]. Producers register the schema before publishing; consumers fetch the writer’s schema by ID and reconcile it with their reader’s schema at deserialization.

The practical rules: add new fields with defaults (backward compatible), avoid renames (use aliases if you must), never drop required fields, never change a field’s type incompatibly. ML pipelines should always include stable identifiers - entity ID, event ID, event timestamp - as required fields, because these are the keys on which idempotency depends.

On the lake side, table formats (Delta, Iceberg, Hudi) extend schema evolution to columnar files. Add column, drop column, rename column, and reorder are all supported as table-level operations that update metadata without rewriting historical data.

Backpressure, Retries, and Dead Letter Queues

A pipeline that ingests faster than its sink can absorb does not just slow down - it falls over. Memory fills, GC pauses extend, consumers get evicted, and the lag chart turns into a wall. Backpressure is the mechanism by which a slow downstream consumer signals an upstream producer to slow down.

In Kafka consumers, backpressure is largely manual: tune max.poll.records, max.partition.fetch.bytes, and fetch.max.bytes so each poll returns only as much as the processor can handle before the next poll deadline. In stream processing frameworks (Flink, Spark Structured Streaming, Kafka Streams), backpressure is automatic: the framework measures sink throughput and slows source reads to match.

When the sink itself fails on a specific record - bad schema, business-rule violation, missing reference data - the answer is a dead letter queue (DLQ). The failing message is published to a side topic with error metadata; the main pipeline continues processing. Operators triage the DLQ separately, often replaying records back to the main topic after fixing the underlying issue. For transient errors (sink 5xx, rate-limit responses), exponential backoff with jitter prevents thundering-herd retries that turn a brief outage into a sustained one.

Lineage

The final pillar is lineage: knowing, for any given training row or feature value, exactly which source events produced it, through which transformations, at which versions. Lineage matters for ML for three reasons. First, debugging: when a model degrades, lineage lets you trace a suspicious feature back to its source. Second, compliance: regulated industries (finance, healthcare) need to reproduce any prediction with the exact inputs used. Third, reproducibility: retraining last quarter’s model requires reconstructing last quarter’s training data, including schema versions and source snapshots.

Lineage is captured at multiple layers. Table formats (Delta, Iceberg) record version histories and the operations that produced each version. Orchestrators (Airflow, Dagster) record job runs and their inputs and outputs. Dedicated lineage tools (OpenLineage, Marquez, DataHub) aggregate signals across the stack into a graph. For ML specifically, MLflow and feature stores like Feast record the dataset versions and feature definitions used in each training run, closing the loop from data to model.

Key Takeaway: Reliable ingestion combines idempotent producers and sinks (keyed by stable event IDs), schema registries with compatibility enforcement, backpressure-aware consumers with DLQs for poison messages, and end-to-end lineage so any model output can be traced back to its source events.

Chapter Summary

This chapter mapped the front end of the ML pipeline. Data sources fall into four broad categories - relational databases (best ingested via CDC for incremental changes), object stores and lakehouses (the canonical home of historical training data), event streams like Kafka and Kinesis (the substrate of real-time features), and APIs, logs, and sensors (each with its own connector idioms). Mixing these sources is normal; matching each to its natural ingestion mechanism is essential.

Ingestion comes in two main flavors. Batch processing moves data periodically via Airflow plus Spark, suiting historical training and slow aggregated features with minute-to-hour latency. Streaming processing uses Kafka or Kinesis plus Flink or Spark Structured Streaming, delivering sub-second freshness for online inference. Lambda architecture combines both with separate batch and speed layers, while Kappa unifies on a single streaming pipeline. Most production feature stores adopt a Lambda-like dual store but mitigate code drift by declaring feature logic once.

File format choice shapes I/O performance, storage cost, and schema flexibility. Avro is row-based with strong schema evolution, ideal for Kafka topics and raw bronze layers. Parquet and ORC are columnar with excellent compression and predicate pushdown, making them the right choice for offline feature stores and gold training tables. TFRecord is TensorFlow-native and high-throughput for sequential reads but is rarely the canonical format; it is best used as a derived artifact for specific training jobs.

Finally, reliability rests on four pillars. Idempotency, achieved through Kafka idempotent producers, transactional semantics where applicable, and stable event IDs that drive sink-side upserts, ensures that retries do not corrupt data. Schema registries enforce compatibility so source evolution does not break ingestion. Backpressure controls and dead letter queues keep the pipeline stable under load and isolate poison messages. Lineage closes the loop, enabling debugging, compliance, and reproducibility from any model output back to its source events.

With ingestion sorted, the pipeline has data flowing in. The next chapter takes up what happens to that data: cleaning, validation, feature engineering, and the construction of a feature store that bridges training and serving.

Key Terms

Chapter 3: Data Validation, Cleaning, and Quality

In machine learning, the model you ship is only as trustworthy as the data that feeds it. A common analogy is that data is the fuel of an ML system — but unlike gasoline, data is rarely refined to a single specification. It arrives noisy, partial, misformatted, occasionally adversarial, and almost always changing over time. A pipeline that ingests this data without validation, cleaning, and ongoing quality measurement is like an aircraft engine running on whatever liquid happens to be in the tank: it may run for a while, but the failure mode is catastrophic and silent until it isn’t.

This chapter introduces the practices and tools that transform raw ingested data into trustworthy ML inputs. We begin with the dimensions that define “quality,” then move to schema and statistical validation frameworks, then to active cleaning strategies for missing values, outliers, and label noise, and finally to the detection of drift, skew, and anomalies that emerge once a model is in production. By the end, you should be able to design data-quality SLAs that hold a pipeline to a contract, not a hope.

Data Quality Dimensions

Data quality is multidimensional. Treating it as a single binary (“clean” vs. “dirty”) obscures the very tradeoffs an ML engineer must reason about. Most practitioners decompose data quality into five dimensions: completeness, accuracy, consistency, timeliness, and uniqueness. Each maps to a distinct class of failure in downstream ML.

Completeness, Accuracy, Consistency, Timeliness, Uniqueness

Completeness measures the fraction of expected values that are actually present. A churn dataset where 30% of total_charges values are null has a completeness problem in that column. Completeness failures often arise from upstream pipeline bugs: a renamed column, a dropped join key, a partial backfill.

Accuracy asks whether a value reflects ground truth. An age of -7 or a country code of XX is inaccurate by definition. Subtler accuracy issues — a timestamp recorded in the wrong timezone, a price denominated in cents instead of dollars — can hide for months and silently corrupt model training.

Consistency measures whether the same fact is represented the same way everywhere. If user_id = 42 has country = "US" in one table and country = "United States" in another, you have a consistency problem. Consistency issues are especially dangerous across the train/serve boundary, because the same upstream entity may be presented to the model in two incompatible forms.

Timeliness measures the lag between when an event occurred in the real world and when it became available to the pipeline. A fraud model trained on transactions delayed by 48 hours will systematically underweight rapidly emerging attack patterns.

Uniqueness measures whether each entity appears exactly as often as expected. Duplicate rows inflate certain classes during training and produce biased loss estimates; missing primary-key uniqueness breaks downstream joins.

How Each Affects ML

Each dimension maps to a distinct model-level failure:

A useful analogy: if data quality were a five-legged stool, removing any one leg destabilizes the whole model. You can ship to production missing any single dimension only if you compensate explicitly elsewhere.

Quality Scoring

Mature ML platforms compute a data quality score for each batch, often as a weighted aggregate of dimensional scores. A typical scheme might be:

• Completeness score: 1 − (mean null rate across critical features)
• Accuracy score: 1 − (fraction of rows failing range/domain checks)
• Consistency score: 1 − (fraction of cross-source mismatches)
• Timeliness score: 1 − min(1, lag / SLA_lag)
• Uniqueness score: 1 − (duplicate rate)

Scores are persisted to a time-series store and dashboarded alongside model metrics, so a sudden drop in any dimension is visible long before a customer-impacting prediction error occurs [Source: https://pub.towardsai.net/codequeries-answering-semantic-queries-over-code-944a93c302ee].

Cost of Poor Quality

The cost of poor quality compounds at every stage of the ML lifecycle. A bad row costs perhaps a millisecond of compute during ingestion, a few seconds of an analyst’s time during EDA, hours of debugging when a training run produces puzzling results, days when a model is retrained on the corruption, and potentially weeks of business impact when the model misclassifies in production. The “1-10-100” rule from data management — that a defect costs $1 to prevent at source, $10 to remediate downstream, and $100 once it reaches a customer — translates almost directly to ML.

Key Takeaway: Data quality is not a single property but five orthogonal dimensions — completeness, accuracy, consistency, timeliness, and uniqueness — each with distinct failure modes in downstream ML. Score each dimension explicitly so degradation is measurable, not merely felt.

Schema and Statistical Validation

Once you have decided what quality means, you need machinery to enforce it. Two complementary categories of tools dominate the field: schema-based validators (which enforce structural and statistical expectations on every batch) and statistical anomaly detectors (which compare incoming distributions against a baseline). Modern frameworks blend both, but their design philosophies and ecosystems differ.

TFDV (TensorFlow Data Validation)

TensorFlow Data Validation (TFDV) is the data-quality component of the TFX stack. It is built on Apache Beam, which lets it scale to terabyte-scale datasets on Dataflow, Spark, or Flink runners [Source: https://www.oreilly.com/content/question-answering-with-tensorflow/]. A typical TFDV workflow has three steps: compute statistics over a reference dataset, infer a schema, and validate every subsequent batch against that schema.

import tensorflow_data_validation as tfdv

# Step 1: generate statistics from training data
stats = tfdv.generate_statistics_from_csv(data_location='train.csv')

# Step 2: infer a schema
schema = tfdv.infer_schema(stats)

# Step 3: validate a new batch
eval_stats = tfdv.generate_statistics_from_csv('batch.csv')
anomalies = tfdv.validate_statistics(eval_stats, schema)
tfdv.display_anomalies(anomalies)

The schema is a protobuf describing feature types (INT, FLOAT, STRING, BYTES), presence (required vs. optional), domains (allowed categorical values or numeric ranges), and structure for nested features. TFDV produces structured Anomalies objects: each anomaly is tied to a feature, a reason, and a severity, so it can be wired directly into a TFX gating step that blocks training when severity exceeds a threshold [Source: https://www.youtube.com/watch?v=tpCFfeUEGs8].

Crucially, TFDV automates training-serving skew detection and drift detection as first-class operations. You hand it two stats artifacts — training and serving, or yesterday and today — and it emits per-feature comparisons. This is the principal reason TFDV remains popular for TensorFlow-centric platforms even though its ecosystem has narrowed: skew detection in ML pipelines is the problem TFDV was built for.

Figure 3.1: Data validation pipeline — raw inputs flow through schema and statistical checks before reaching the cleaned, model-ready dataset.

flowchart LR
    A[Raw Batch] --> B[Schema Check<br/>types, presence, domains]
    B -->|pass| C[Statistical Check<br/>ranges, distributions]
    B -->|fail| X[Anomaly Report]
    C -->|pass| D[Cleaned Dataset]
    C -->|fail| X
    X --> E[Gate / Alert]

Great Expectations

Great Expectations (GE) takes a different philosophical stance. Instead of inferring a statistical baseline and detecting deviations, GE asks the team to write expectations — human-readable assertions about the data, organized into Expectation Suites.

import great_expectations as ge
import pandas as pd

df = pd.read_csv("train.csv")
ge_df = ge.from_pandas(df)

ge_df.expect_column_values_to_not_be_null("user_id")
ge_df.expect_column_values_to_be_between("age", min_value=0, max_value=120)
ge_df.expect_column_mean_to_be_between("click_count", min_value=0, max_value=10)
result = ge_df.validate()

Every expectation that fails is, in GE’s model, an anomaly. The suite can be auto-profiled from a reference dataset, then refined by humans — a workflow analogous to generating unit tests via a coverage tool, then editing them by hand. GE renders results as browsable Data Docs: HTML reports that double as living documentation of the data contract [Source: https://github.com/Bhanupriya-art/INT426-Coursera-Answers]. The framework integrates first-class with Airflow via the GreatExpectationsOperator, making it the natural choice for warehouse- and dbt-centric data stacks.

Pandera and pydantic

For lighter-weight validation embedded in Python services, pandera (DataFrame schemas) and pydantic (Pythonic data models) provide expressive, in-process validation. Pandera schemas validate Pandas/Polars DataFrames inline:

import pandera as pa
from pandera.typing import Series

class TransactionSchema(pa.DataFrameModel):
    user_id: Series[int] = pa.Field(ge=1)
    amount: Series[float] = pa.Field(ge=0, le=1_000_000)
    country: Series[str] = pa.Field(isin=["US", "CA", "UK", "DE"])

TransactionSchema.validate(df)

pydantic shines at request-payload validation in FastAPI services that wrap models for online inference, ensuring that what arrives at the model matches what the model was trained on. Both are typically used as the last mile of validation inside the application — TFDV/GE catch the upstream problems, pandera/pydantic catch the request-level ones.

Expectations as Code

The unifying idea across all four frameworks is expectations as code: data contracts that are version-controlled, reviewed in pull requests, executed in CI, and deployed with the pipeline. This is the data-quality analog of infrastructure-as-code. The benefits are the same: reproducibility, auditability, and a shared source of truth between data engineers, ML engineers, and analysts.

In practice, mature ML platforms use a hybrid: Great Expectations upstream in the warehouse to enforce business rules, TFDV downstream for ML-specific feature and skew checks, and pandera/pydantic at service boundaries.

Key Takeaway: TFDV automates statistical schema inference and drift/skew detection for TFX pipelines, while Great Expectations encodes human-readable assertions integrated with Airflow and data warehouses. Treat expectations as code — version them, review them, and run them in CI.

Cleaning Strategies

Validation tells you what is wrong; cleaning decides what to do about it. The three perennial cleaning problems in ML are missing values, outliers, and label noise. Each requires a distinct strategy, and each has subtle failure modes that can silently bias the model.

Missing Values: Drop, Impute, Flag

Before choosing a treatment for missingness, diagnose the mechanism. The literature distinguishes three:

• MCAR (Missing Completely At Random): missingness is unrelated to any data. Dropping rows is safe but wasteful.
• MAR (Missing At Random): missingness depends on other observed variables. Imputation works but should condition on those variables.
• MNAR (Missing Not At Random): missingness depends on the unobserved value itself — e.g., income missing more often for high earners. In this regime, missingness is informative, and a “missingness indicator” feature often boosts performance.

Figure 3.2: Imputation decision tree — choose a strategy by missingness mechanism, rate, and feature importance.

flowchart TD
    A[Missing values detected] --> B{Mechanism?}
    B -->|MCAR| C{Missing rate < 5%?}
    C -->|Yes| D[Drop rows]
    C -->|No| E[Mean/Median impute]
    B -->|MAR| F[Conditional impute<br/>KNN or MICE]
    B -->|MNAR| G[Sentinel + missingness<br/>indicator feature]
    F --> H{Feature critical?}
    G --> H
    E --> H
    H -->|Yes| I[Model-based imputation]
    H -->|No| J[Keep simple imputer]

Common imputation strategies, in increasing order of complexity:

• Drop: remove rows with missing values. Acceptable only when missing rate is low (<1–5%) and MCAR.
• Mean/Median/Mode: the workhorses. Median is more robust to skew and outliers; mode is used for categoricals. Shrinks feature variance and biases correlations, so pair with a missingness indicator for important features.
• Constant + indicator: fill with a sentinel (0, -1, "__MISSING__") and add a boolean column flagging the imputation. Tree-based models such as XGBoost, LightGBM, and CatBoost handle this exceptionally well — many of them natively respect a NaN.
• KNN imputation: impute using the average of the k nearest training rows. Captures local multivariate structure but requires scaling and scales poorly past ~100k rows.
• MICE / IterativeImputer: chained equations that model each missing column as a function of the others. Statistically principled but compute-heavy.
• Model-based imputation: train a dedicated model (random forest, autoencoder) per high-value missing feature. Best when a feature is critical but frequently missing.

A non-negotiable rule: fit imputers only on training data, then apply the fitted artifact to validation, test, and production. Fitting on the full dataset leaks information into the test set and produces optimistic estimates [Source: https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/prompt-engineering].

from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.ensemble import RandomForestClassifier

numeric_features = ["age", "income", "balance"]
categorical_features = ["country", "device_type"]

numeric_transformer = Pipeline(steps=[
    ("imputer", SimpleImputer(strategy="median", add_indicator=True)),
    ("scaler", StandardScaler())
])

categorical_transformer = Pipeline(steps=[
    ("imputer", SimpleImputer(strategy="most_frequent")),
    ("encoder", OneHotEncoder(handle_unknown="ignore"))
])

preprocess = ColumnTransformer(transformers=[
    ("num", numeric_transformer, numeric_features),
    ("cat", categorical_transformer, categorical_features),
])

clf = Pipeline(steps=[("preprocess", preprocess),
                     ("model", RandomForestClassifier(n_estimators=200, random_state=0))])
clf.fit(X_train, y_train)

The serialized pipeline now applies the same imputation logic in both training and inference — eliminating an entire category of training-serving skew.

Outliers

Outliers come in three flavors: data errors (sensor glitches, parsing bugs), rare but valid cases (legitimate high-value customers), and distribution shifts (a new population the model has never seen). They demand different responses.

Detection methods range from simple to multivariate:

• IQR rule: flag values outside [Q1 − 1.5·IQR, Q3 + 1.5·IQR]. Robust, univariate, interpretable.
• Z-score: flag |z| > 3. Assumes approximate normality; mean and σ are themselves affected by outliers. Use median + MAD for robustness.
• Isolation Forest: randomly partitions feature space; points isolated in few splits are anomalies. Handles mixed distributions; scales to large data; key hyperparameter is contamination.
• DBSCAN: density-based clustering. Points unassigned to any cluster are noise. Sensitive to eps and min_samples.
• Local Outlier Factor (LOF): compares local density of a point to its neighbors. Good for heterogeneous data with local anomalies.

Treatment options:

• Fix or correct known data errors.
• Winsorize: cap values at the 1st/99th percentile. Preserves the row, kills the tail.
• Remove: only when you are confident the rows are not part of the population you care about.
• Down-weight: for loss functions that support instance weights.
• Use robust models: trees, Huber loss, quantile regression.

Critical guidance: do not auto-delete every flagged outlier. In fraud detection or rare-disease prediction, the “outliers” are the positive class.

Deduplication and Normalization

Deduplication seems trivial but is surprisingly subtle. Exact deduplication on a primary key is fast. Near-duplicate detection — same person with two email addresses, same product with two SKUs — typically requires fuzzy matching (Jaro-Winkler, MinHash, embedding similarity). Duplicates inflate training-set size without adding information, biasing the model toward the duplicated examples and producing optimistic cross-validation scores.

Normalization standardizes representations: lowercasing strings, collapsing whitespace, mapping "USA"/"U.S.A."/"United States" to a canonical "US", parsing dates to ISO 8601. Inconsistent normalization is one of the most common sources of train-serve skew: the training pipeline lowercases country codes, the serving path does not, and the model silently encounters unseen categories.

Label Noise

Label noise is the most damaging form of data corruption because it directly corrupts the learning signal. If 10% of your labels are wrong, ~90% is the ceiling for accuracy on that noisy test set — and worse, the model will memorize the errors.

The dominant modern technique is confident learning, implemented in the Cleanlab library. The workflow:

1. Train a baseline model and obtain out-of-sample predicted probabilities for every training example, typically via 5-fold cross-validation.
2. Pass labels and probabilities to find_label_issues, which estimates the joint distribution of noisy vs. true labels and flags examples where the model is confidently disagreeing with the assigned label.
3. Route flagged examples for human review, drop them, or use CleanLearning to retrain with noise-aware reweighting.

from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_predict
from cleanlab.filter import find_label_issues

base_clf = RandomForestClassifier(n_estimators=300, random_state=0)
probs = cross_val_predict(base_clf, X, y, cv=5, method="predict_proba")
label_issues = find_label_issues(labels=y, probs=probs)

X_clean, y_clean = X[~label_issues], y[~label_issues]
base_clf.fit(X_clean, y_clean)

Complementary techniques include label smoothing (softens hard targets), co-teaching (two networks teach each other from their respective low-loss examples), and per-example loss tracking (persistently high-loss examples are often mislabeled). In production, the single highest-leverage action is usually establishing a human-in-the-loop review queue for high-confidence disagreements between model and label.

Key Takeaway: Choose missing-value strategies by mechanism (MCAR/MAR/MNAR), fit imputers only on training data, distinguish data-error outliers from rare-but-valid cases before deleting anything, and treat label noise as a first-class data-quality problem — cleaning labels often beats tuning the model.

Drift, Skew, and Anomalies

A model trained on a clean snapshot still degrades in production, because the world moves. Detecting that movement is the work of drift detection: comparing live distributions against a training reference and raising alarms when they diverge enough to matter.

Training-Serving Skew

Training-serving skew is a systematic difference between what a model saw in training and what it sees at serving time, caused by something other than natural distribution shift — typically a pipeline bug. Examples include: a feature engineered with a SQL query in training but a Python function in serving (and the two implementations disagree on edge cases); a categorical encoder fitted on training categories that silently maps unseen serving categories to 0; a normalization step accidentally re-fitted on each serving batch.

The defining property of skew is that it is fixable by engineering. The cure is structural: use a single serialized preprocessing pipeline for both paths, validate serving inputs against the same schema used in training, and run continuous comparisons between training and serving feature distributions. Feature stores exist largely to make this hard problem easy by maintaining a single source of truth for feature values across train and serve.

Figure 3.3: Training-serving skew — divergent feature pipelines silently corrupt predictions; a shared serialized transformer eliminates the gap.

flowchart TD
    R[(Raw Source)] --> T1[Training Pipeline<br/>SQL feature query]
    R --> S1[Serving Pipeline<br/>Python feature fn]
    T1 --> T2[Fitted Encoders]
    S1 --> S2[Ad-hoc Encoders]
    T2 --> M1[Model Training]
    S2 --> M2[Online Inference]
    M1 -.skew.- M2
    R --> P[Serialized Preprocessing<br/>Pipeline / Feature Store]
    P --> M1
    P --> M2
    style P fill:#1f6feb,color:#fff

Concept Drift vs. Data Drift

The two forms of true distributional change are conceptually distinct:

• Data drift (covariate drift): P(X) changes. The kinds of inputs the model sees in production differ from training — perhaps the user demographic shifted, or a new geography came online.
• Concept drift: P(Y | X) changes. The relationship between inputs and outputs has moved. A spam filter trained on 2023 spam will degrade as spammers adopt new tactics in 2025, even if the surface features look identical.

Data drift is detectable from inputs alone. Concept drift requires either labels (often delayed in production) or proxies such as prediction distribution shifts, performance estimation, or shifts in model confidence.

KS, PSI, JS, and Other Divergence Measures

Drift detection comes down to two-sample distribution tests and divergence measures comparing a reference (training, or a stable baseline window) to a current production sample. The major methods, their use cases, and typical thresholds:

Population Stability Index (PSI) deserves special attention because it is the de-facto industry standard in credit risk and many production ML platforms. For each bin i with training proportion p_i and production proportion q_i, PSI = Σ (q_i − p_i) · ln(q_i / p_i). The heuristic thresholds — <0.1 stable, 0.1–0.25 moderate drift, ≥0.25 significant drift — come from decades of scorecard practice and are configurable per feature based on business criticality [Source: https://www.youtube.com/watch?v=KuzEm1VhJYE].

The KS test complements PSI with statistical significance. KS compares empirical CDFs of two continuous samples and outputs a p-value; teams typically require both p < 0.05 and a minimum effect size (D > 0.1–0.2) because at high N the test rejects on imperceptibly small differences [Source: https://pmc.ncbi.nlm.nih.gov/articles/PMC4905616/].

Jensen-Shannon divergence is the symmetric, bounded sibling of KL. JS(P||Q) = ½ KL(P||M) + ½ KL(Q||M) where M = ½(P+Q). Because it stays finite even when supports differ, it is far safer than raw KL for production dashboards.

MMD is the heavy artillery: a kernel-based two-sample test that operates in a reproducing kernel Hilbert space (RKHS), making it ideal for multivariate drift detection on raw feature vectors, image embeddings, or text representations [Source: https://www.emergentmind.com/topics/maximum-mean-discrepancy-mmd]. The Alibi Detect library implements MMDDrift as one of its core detectors. Thresholds are derived analytically from concentration bounds or empirically from permutation tests [Source: https://arxiv.org/html/2205.12706v3].

Two further practical heuristics:

• Statistical vs. practical significance: a p-value answers “is the difference real?”; an effect-size threshold (PSI, Wasserstein) answers “is the difference big enough to matter?”. Use both.
• Calibrate thresholds on healthy data. Universal textbook numbers are starting points. Compute the distribution of each drift metric over historical no-drift periods and alert at, say, the 95th or 99th percentile of that baseline [Source: https://www.r-bloggers.com/2022/01/universal-estimation-with-maximum-mean-discrepancy-mmd/].

Alerting and SLAs

Drift detection without alerting is theater. A production drift monitor should:

1. Run continuously, comparing the latest window (hour/day/batch) to a fixed training reference.
2. Compute multiple metrics per feature — at minimum, PSI plus one significance test.
3. Aggregate per-feature signals into an overall drift share (fraction of features above threshold).
4. Route alerts by severity: warn on moderate drift (investigate), page on significant drift (intervene).
5. Tie drift alerts to operational playbooks: retraining triggers, fallback model activation, kill-switch deployment.

A useful data quality SLA template for a production ML system:

When an SLA is violated, the response should be policy-driven: block scoring for high-severity failures, route to a fallback model for moderate ones, file a ticket for everything. The point is that the response is automatic — humans should be informed, not interrupted, when the system behaves as designed [Source: https://www.anthropic.com/engineering/effective-context-engineering-for-ai-agents].

Figure 3.4: Drift detection and SLA monitoring loop — reference and live distributions feed statistical tests; severity routes to playbooks.

flowchart TD
    REF[Reference Distribution<br/>training baseline] --> CMP[Drift Tests<br/>PSI, KS, JS, MMD]
    LIVE[Live Distribution<br/>rolling window] --> CMP
    CMP --> AGG[Drift Share<br/>aggregate per-feature signals]
    AGG --> SEV{Severity?}
    SEV -->|None| LOG[Log metrics]
    SEV -->|Moderate| WARN[Warn + ticket<br/>investigate]
    SEV -->|Significant| PAGE[Page on-call<br/>fallback / retrain / kill-switch]
    LOG --> LIVE
    WARN --> LIVE
    PAGE --> LIVE

Key Takeaway: Distinguish skew (engineering bug) from drift (world changing) and concept drift (P(Y|X) shift) from data drift (P(X) shift). Combine a significance test (KS, χ², MMD) with an effect-size metric (PSI, Wasserstein, JS) per feature, calibrate thresholds on historical healthy data, and codify the response as data-quality SLAs.

Chapter Summary

Data quality is the silent variable that decides whether an ML system delivers value or accumulates technical debt. We decomposed it into five orthogonal dimensions — completeness, accuracy, consistency, timeliness, uniqueness — and showed how each maps to a distinct ML failure mode. Quality scores per dimension, persisted to time-series storage, transform “feeling bad about the data” into an engineering signal.

We surveyed schema and statistical validation frameworks. TFDV automates schema inference, drift detection, and training-serving skew detection at Apache Beam scale; it is the natural choice inside TFX. Great Expectations encodes human-readable assertions into version-controlled Expectation Suites and renders them as Data Docs; it is the natural choice for Airflow and warehouse-centric stacks. Pandera and pydantic cover DataFrame and request-payload validation inside Python services. Most mature platforms use a hybrid: GE upstream, TFDV downstream, pandera/pydantic at service boundaries.

For cleaning, we covered missing-value treatments (drop, simple impute, KNN, MICE, model-based) keyed to the missingness mechanism (MCAR/MAR/MNAR); outlier detection via IQR, Z-score, Isolation Forest, DBSCAN, and LOF, with the critical warning not to auto-delete what may be the positive class; deduplication and normalization as defenses against train-serve skew; and label noise mitigation via confident learning and Cleanlab.

Finally we examined drift, skew, and anomalies in production. We separated training-serving skew (an engineering bug fixable by serialized preprocessing pipelines and feature stores) from concept drift (P(Y|X) change) and data drift (P(X) change). We compared the major detection methods — KS, PSI, JS, KL, χ², Wasserstein, MMD — by data type and use case, and emphasized combining statistical significance with practical effect size, calibrated on historical baselines. We closed with a data-quality SLA template that translates all of the above into a contract a pipeline can be held to.

The unifying theme: every data-quality property worth caring about should be measured continuously, expressed as code, versioned in Git, validated in CI, and gated in production. Hoping the data is clean is not a strategy; making cleanliness observable is.

Key Terms

Chapter 4: Feature Engineering and Feature Stores

If raw data is the crude oil of machine learning, features are the refined fuel that engines actually burn. A model can only ever be as good as the signal encoded in its inputs, and the discipline of crafting, transforming, storing, and serving those inputs is what we call feature engineering. This chapter takes you from the bedrock transformations applied to tabular, text, and time-series data, through the architectural pattern that prevents your training data from silently disagreeing with your production data: the feature store.

By the end of the chapter you will be able to apply the core encoding, scaling, and windowing transforms used across modern ML systems; articulate the train/serve consistency problem and how a feature store solves it; compare Feast, Tecton, SageMaker Feature Store, and adjacent platforms with a clear sense of when each fits; and design a feature pipeline with versioning, point-in-time correctness, and proper materialization schedules.

Section 1: Feature Engineering Fundamentals

Feature engineering is the bridge between raw events and the matrix of numbers a model consumes. A useful mental analogy is cooking: raw ingredients (data) rarely go straight onto the plate. They are washed, chopped, marinated, and balanced. Likewise, before a learning algorithm can extract patterns, raw fields must be scaled, encoded, aggregated, and shaped into a representation the model can use. The most effective production work focuses on simple, robust transformations first, then adds complexity (embeddings, advanced time-series features) only where they clearly improve business metrics and can be reliably maintained [Source: https://pmc.ncbi.nlm.nih.gov/articles/PMC9904526/].

Scaling, Encoding, and Binning of Numeric Features

Most learning algorithms care about the scale and distribution of numeric inputs. Linear models, neural networks, k-means, and PCA are all sensitive to feature magnitude, while tree-based models such as XGBoost and LightGBM are largely scale-invariant. A few canonical numeric transforms apply across the board:

• Standardization (z-score): shifts each feature to mean 0, standard deviation 1. Default for linear models, neural networks, and any distance-based method.
• Min-max scaling: rescales to a [0, 1] range. Useful when features have natural bounds (probabilities, percentages) or when downstream code expects bounded inputs.
• Robust scaling: uses the median and interquartile range instead of mean and standard deviation. More stable under heavy tails and outliers [Source: https://pmc.ncbi.nlm.nih.gov/articles/PMC9904526/].
• Binning / discretization: converts a continuous variable into ordinal buckets (for example, age into decades). Useful for non-linear effects in linear models, and as a defense against extreme outliers.
• Log and Box-Cox transforms: compress skewed positive distributions (prices, counts, sales) so the model sees a more Gaussian-shaped input.

The most important production rule is one of discipline: fit scalers and binners only on the training partition, persist the parameters (mean, std, bin edges) as part of your model artifact, and apply the identical transformation online via a shared feature store or portable pipeline.

Categorical Encoding

Categorical features are where naive choices break models. Choosing the wrong encoding for a 5-million-cardinality user ID column will either explode your feature matrix or leak labels into training. The right answer depends on cardinality, model family, and your tolerance for retraining encoders.

For high-cardinality features such as IDs, URLs, or zip codes, three strategies dominate in production. Target encoding replaces each category with an aggregated target statistic such as conversion rate, but must be computed on out-of-fold data with smoothing toward a global mean for rare categories, plus optional noise injection during training [Source: https://pmc.ncbi.nlm.nih.gov/articles/PMC9904526/]. The hashing trick maps category strings into a fixed number of buckets via a hash function, which gives bounded dimensionality and trivial handling of new categories at the cost of occasional collisions. Learned embeddings train a small lookup table jointly with the model, common in deep CTR and recommender architectures; a useful starting heuristic is d ~ min(50, sqrt(cardinality)) for embedding dimension, with L2 regularization and dropout to prevent overfitting on rare IDs [Source: https://pmc.ncbi.nlm.nih.gov/articles/PMC9904526/].

Text: Bag-of-Words, TF-IDF, and Embeddings

Text needs to be turned into vectors before a model can touch it, and the ladder of techniques goes from cheap and interpretable to expensive and semantically rich.

• Bag-of-Words (BoW): counts of tokens or n-grams. Fast, easy, ignores word order [Source: https://machinelearningmastery.com/a-gentle-introduction-to-word-embedding-and-text-vectorization/].
• TF-IDF (Term Frequency-Inverse Document Frequency): weights tokens by how common they are in a document relative to the corpus, emphasizing words that are frequent in a document but rare overall. TF-IDF is a remarkably strong baseline for classification, FAQ retrieval, and search [Source: https://healthark.ai/keyword-extraction-using-regex-tf-idf-and-bert-a-comprehensive-approach/].
• Static embeddings (Word2Vec, GloVe, Doc2Vec): dense vectors (50-300 dims) where semantically similar words are geometrically close [Source: https://pub.towardsai.net/from-words-to-vectors-exploring-text-embeddings-af64ee798759].
• Contextual embeddings (BERT, DistilBERT, modern transformers): the same word gets different vectors depending on context. Highest accuracy, highest cost [Source: https://ai.plainenglish.io/bert-vs-tf-idf-embeddings-in-an-enterprise-chatbot-782ba161ffcf].

Production tips for TF-IDF: cap the vocabulary at 20k-100k tokens, use n-grams up to 2 or 3, apply L2 normalization, and optionally reduce dimensionality with truncated SVD if latency or memory is tight. Freeze the vocabulary at training time and handle out-of-vocabulary terms via an “unknown” bucket or hashing [Source: https://machinelearningmastery.com/a-gentle-introduction-to-word-embedding-and-text-vectorization/].

Time-Series: Lags, Rolling Windows, and Seasonality

Time-series problems are typically reshaped into supervised learning by extracting features over past windows. The three workhorses are:

• Lag features: past values of the series itself or its covariates - for example x_{t-1}, x_{t-7}, x_{t-30} to capture daily, weekly, and monthly dependencies.
• Rolling-window statistics: mean, median, min, max, standard deviation, skewness, exponential moving averages, rolling counts and rates. These capture local trend and volatility.
• Seasonality and calendar features: hour of day, day of week, month, week of year, holiday flags. Cyclic variables should be encoded with sin/cos transforms so January and December are close in feature space. Fourier features (sine/cosine terms at chosen frequencies) provide an explicit handle on periodic patterns for regression-style models.

The single most important rule in time-series feature engineering is no future leakage: every feature window must end strictly before the prediction time, and validation must use time-based splits rather than random splits.

Figure 4.1: Time-series feature construction with lag and rolling window features ending strictly before the prediction time t.

flowchart TD
    A[Raw time series x_t] --> B[Lag features]
    A --> C[Rolling window stats]
    A --> D[Calendar / cyclic features]
    B --> B1[x_t-1, x_t-7, x_t-30]
    C --> C1[mean, std, min, max over window W]
    C --> C2[EMA, rolling counts]
    D --> D1[hour, day_of_week, month]
    D --> D2[sin/cos cyclic encoding]
    B1 --> E[Feature vector at time t]
    C1 --> E
    C2 --> E
    D1 --> E
    D2 --> E
    E --> F{Window ends strictly before t?}
    F -->|Yes| G[Safe to train / serve]
    F -->|No| H[Future leakage - reject]

Key Takeaway: Start with simple, robust transforms - z-score scaling, one-hot or hashing for categoricals, TF-IDF for text, and backward-looking lag/window aggregates for time series - and reach for embeddings, target encoding, or BERT only when the offline gains justify the operational cost.

Section 2: The Feature Store Pattern

Once a team starts shipping more than one model, a structural problem emerges. The same “average spend over 30 days” feature gets written three times: once in a Snowflake query for the training set, once in a Python microservice for online inference, and once in a Spark job for batch scoring. Each implementation drifts from the others. Models that look great on the offline test set degrade silently in production. This is the problem a feature store is built to eliminate.

Why Feature Stores Exist

In organizations without a feature store, data scientists typically write feature code separately for training in Spark, SQL, or notebooks over a warehouse, and for serving in Python or Java microservices with ad-hoc caches. The consequences are predictable and painful:

• Train-serve skew: features computed differently offline versus online.
• Feature duplication: the same business concept reimplemented across teams.
• Hard-to-reproduce historical training datasets.
• Ad-hoc caches and brittle pipelines that no one wants to own.

A feature store solves these by providing one shared system for feature definition, computation, storage, and serving so that training and online inference use the same logic and data, with correct time semantics and low operational overhead.

Figure 4.2: Feature store architecture - shared definitions feed offline and online stores from one source of truth.

flowchart LR
    DS1[Warehouse / Lake] --> FE[Feature Engineering]
    DS2[Kafka / Kinesis streams] --> FE
    DS3[Operational DBs] --> FE
    FE --> REG[Feature Registry / Catalog]
    REG --> OFF[(Offline Store<br/>S3/Parquet, BigQuery,<br/>Snowflake)]
    REG --> ON[(Online Store<br/>Redis, DynamoDB)]
    OFF --> TRAIN[Training<br/>point-in-time joins]
    ON --> SERVE[Online Serving<br/>millisecond reads]
    TRAIN --> MODEL[Model Artifact]
    MODEL --> SERVE

              +-----------------------------+
              |   Feature Registry / Catalog|
              |   (definitions, owners,     |
              |    metadata, versions)      |
              +--------------+--------------+
                             |
        +--------------------+--------------------+
        |                                         |
+-------v--------+                       +--------v--------+
| Offline Store  |   <- materialization  |  Online Store   |
| (S3/Parquet,   |   ------------------> |  (Redis,        |
|  BigQuery,     |                       |   DynamoDB)     |
|  Snowflake)    |                       |  low-latency KV |
+-------+--------+                       +--------+--------+
        |                                         |
        v                                         v
   Training data                          Online prediction
   (point-in-time joins)                  (millisecond reads)

Online versus Offline Stores

Feature stores split storage into two cooperating layers because training and serving have fundamentally different access patterns.

The offline store is optimized for large historical datasets, cheap storage, and analytical queries. It backs training set generation, backfills, and experimentation. Common offline stores include S3 with Parquet, BigQuery, Snowflake, Redshift, and Delta Lake.

The online store is optimized for low-latency, key-value random access by entity ID such as user_id or account_id. It backs real-time inference from API services and model servers. Common online stores include Redis, DynamoDB, and managed key-value services. Read latency is measured in single-digit milliseconds.

Materialization is the process that moves computed features from source systems into the offline and/or online stores. Three variants dominate:

• Batch materialization - an hourly or nightly job from the warehouse to the online database.
• Streaming materialization - Kafka or Kinesis events flowing through Flink or Spark Streaming into the online store.
• On-demand features - computed at request time from raw inputs or other features (for example, the time_since_last_login value can be computed at predict time as now - last_login_at).

Feature Registry and Metadata

The registry is the central catalog of feature definitions, owners, lineage, tags, and versions. It is what turns a feature store from a database into a governable system. A healthy registry answers questions like:

• Who owns customer_lifetime_value_90d?
• Which models currently consume it?
• What version was used when model v3.2 was trained?
• When was it last successfully materialized?
• What are the upstream data sources?

In Feast, the registry is typically a file on S3/GCS or a small SQL database; in Tecton, it is a rich first-class service with UI, lineage, and access control; in SageMaker Feature Store, it is a set of Feature Group definitions plus IAM-controlled metadata.

Point-in-Time Joins

The single most subtle and most important capability of a feature store is the point-in-time join, also called an AS-OF join or time-travel join. It ensures that, for each training example, you only join in feature values that were available as of the prediction time. This is what prevents label leakage and aligns offline training features with what the model will see online.

Conceptually, given an entity e, a prediction time t_p, and a feature table F(e, t) of feature values across times, a point-in-time join returns the row F(e, t*) where t* <= t_p and t* is the latest such timestamp. Two timestamps matter:

• Event timestamp: when the business event actually happened.
• Created / ingestion timestamp: when the event became visible in the system.

Tracking both lets the store correctly exclude late-arriving corrections that would not have been known at prediction time.

Figure 4.3: Point-in-time join - only feature rows whose event and created timestamps precede the label time are eligible.

sequenceDiagram
    participant E as Entity Timeline
    participant F as Feature Store
    participant J as PIT Join
    participant T as Training Row
    Note over E: t1: balance=1200 (event)
    Note over E: t2: balance=850 (event)
    Note over E: t_p: label_time = 2023-03-31
    Note over E: t3: balance=100 (after t_p)
    E->>F: write feature rows with<br/>event_ts + created_ts
    J->>F: find latest row where<br/>event_ts <= t_p AND<br/>created_ts <= t_p
    F-->>J: returns t2 row (balance=850)
    Note over J: t3 row excluded -<br/>not yet known at t_p
    J->>T: balance=850, label=1

Worked example: point-in-time join for credit risk. Suppose we are training a model that predicts whether a customer will default within 90 days. The label table looks like this:

We have a feature table of daily balance snapshots:

A naive join would pick the latest balance for each customer regardless of time, leaking the post-default 100 balance for customer 7 and the post-prediction 5,100 for customer 9. Offline metrics would look fantastic and production would crater.

A point-in-time join with the constraint feature.event_timestamp <= label.label_time AND feature.created_timestamp <= label.label_time produces:

For customer 7 the join picks the 2023-03-30 row (the most recent event with created_timestamp <= 2023-03-31). For customer 9 it picks the 2023-04-10 row. Both selections faithfully emulate “what we knew at prediction time” [Source: https://developers.openai.com/cookbook/examples/gpt4-1_prompting_guide].

Key Takeaway: A feature store gives you one definition for each feature, served from an offline store for training and an online store for inference, with point-in-time joins that ensure training data reflects only what the model would have known when it predicted.

Section 3: Feature Store Implementations

There is no single “right” feature store; the right one depends on your cloud, your team, and your latency budget. The three reference implementations are Feast (open source), Tecton (commercial end-to-end), and Amazon SageMaker Feature Store (managed AWS), with Databricks Feature Store, Hopsworks, and Vertex AI Feature Store as common alternatives.

Feast

Feast is an open-source feature store originally created at Gojek and incubated with Tecton contributors. It focuses on being a feature serving and registry layer on top of your existing data infra, rather than a turnkey platform. You bring the warehouse, you bring the orchestration, you bring the online KV store; Feast wires them together with a consistent Python SDK [Source: https://natesnewsletter.substack.com/p/context-windows-are-a-lie-the-myth].

Architecturally, Feast offers:

• A pluggable registry stored as a file on S3/GCS or in a small SQL database.
• A pluggable offline store - BigQuery, Snowflake, Redshift, Hive, local Parquet.
• A pluggable online store - Redis, DynamoDB, Postgres.
• A feast materialize / materialize-incremental command that reads from the offline store and writes the latest values to the online store. You schedule it from Airflow, Dagster, or cron.

A typical Feast stack on GCP might look like this: a nightly job computes avg_spend_30d into a partitioned BigQuery table; every 15 minutes feast materialize-incremental reads new rows by timestamp and writes them into Redis keyed by customer_id; a prediction API on Kubernetes calls the Feast SDK to fetch features for a customer ID and passes them into an XGBoost model.

Feast is the right choice for teams with strong platform engineers who value flexibility and open source, and for organizations that are multi-cloud or want to avoid vendor lock-in. Its limitations are operational: you build and run the orchestration, the streaming jobs, the access control, and the UI yourself.

Tecton, Hopsworks, and Databricks

Tecton is a commercial feature platform built by engineers from Uber’s Michelangelo team. It provides the full end-to-end stack: declarative feature definitions in Python, managed compute orchestration (Spark/Flink), online and offline storage, monitoring, governance UI, and serving APIs. It is available as SaaS or in a customer-managed VPC depending on plan.

Tecton’s distinguishing capability is first-class support for real-time streaming features and complex pipelines. For example, a fraud detection use case might define num_transactions_5m and avg_amount_1h as sliding-window aggregations over a Kafka stream. Tecton’s streaming jobs compute these continuously and write to both offline (Delta Lake on S3) and online (DynamoDB) stores. A fraud API queries Tecton’s online serving API in ~10-20 ms, then calls a model for the decision. Data scientists generate point-in-time correct training sets directly from the same Feature Service.

Hopsworks is another commercial/open-core feature platform with strong support for both feature storage and end-to-end MLOps, including model registry and serving. It is particularly popular in EU organizations and on-prem deployments.

Databricks Feature Store (with Unity Catalog) is the natural choice for teams already on Databricks. It integrates tightly with Delta Lake, MLflow, and Unity Catalog governance, and supports both batch and online (via Databricks Online Tables) serving.

DIY Redis + Parquet

For early-stage teams, a credible feature store can be built from a warehouse plus Redis plus a couple of Airflow jobs:

• Warehouse (Snowflake, BigQuery, Redshift) as the offline store.
• Redis or DynamoDB as the online store.
• A YAML or dbt-driven registry of feature definitions.
• Airflow / Dagster DAGs to compute features and write them to both stores.
• A thin Python client that fetches features by entity key.

This is effectively a lightweight Feast-clone, and it is often the right starting point before adopting a full platform. The risk to manage is governance: as the catalog grows, you will increasingly want lineage, access control, and a UI - precisely what Feast or Tecton add.

Vertex AI and SageMaker Feature Store

Amazon SageMaker Feature Store is a managed AWS service. Features are grouped into Feature Groups, each defining a record identifier, an event time, and whether online and/or offline storage is enabled. The offline store lives in S3 as Parquet (queryable via Athena, EMR, or SageMaker Processing), and the online store is a managed DynamoDB-backed key-value layer. Ingestion comes from your own ETL - Glue, EMR, Lambda, Kinesis, or SageMaker Processing - which calls the Feature Store API directly. When both online and offline are enabled, a single write populates both. There is no separate “materialize” command in the Feast/Tecton sense.

Google Vertex AI Feature Store is the GCP analog: a managed offline and online store, deeply integrated with BigQuery and Vertex AI Pipelines.

Here is the side-by-side comparison most teams need:

Key Takeaway: Pick Feast when flexibility and OSS matter and you have platform engineers; pick Tecton when streaming features, governance, and reduced internal glue justify the cost; pick SageMaker Feature Store when you are already on AWS and want a managed option; consider a DIY warehouse + Redis stack to start.

Section 4: Pipelining Features

A feature store is only as good as the pipelines that feed it. This section covers materialization schedules, feature versioning, train/serve skew prevention, and the special case of real-time features on streams.

Materialization and Refresh Schedules

Each feature has its own freshness requirement, which in turn dictates how often it must be materialized. A useful framework is to classify features by SLA:

For batch features, the most important design choice is incremental materialization. A feature like purchases_30d should not be recomputed from scratch every hour; instead the pipeline should read only rows whose timestamps changed since the last successful run and update those entity keys in the online store. Feast offers this via materialize-incremental; Tecton handles it internally.

Figure 4.4: Feature materialization pipeline - incremental compute fans out to both stores keyed by entity ID.

flowchart LR
    SRC[Source Tables<br/>events, transactions] --> CDC[Detect new rows<br/>since last run]
    CDC --> COMP[Feature Computation<br/>SQL / Spark / Flink]
    COMP --> OFF[(Offline Store<br/>partitioned by date)]
    COMP --> ON[(Online Store<br/>keyed by entity_id)]
    OFF --> BACKFILL[Backfills /<br/>historical training]
    ON --> SERVE[Low-latency<br/>inference reads]
    SCHED[Scheduler<br/>Airflow / Dagster] -.triggers.-> CDC

Versioning

Features evolve. A definition change - a new outlier filter, a different smoothing constant in target encoding, a renamed column - changes the meaning of the feature in subtle ways. Without versioning, the model that consumed purchases_30d last month and the model consuming it today might disagree on what it means.

Practical versioning strategies include:

• Suffix the feature name on breaking changes: purchases_30d_v2. Old models keep using purchases_30d_v1; new models opt in.
• Tag the FeatureView with a semantic version in the registry and pin model artifacts to specific versions.
• Pin training datasets to a registry snapshot: record the registry commit hash or version at training time, so any historical training set can be reproduced.

The model registry (covered in later chapters) should record which feature versions a model consumed, much like a Pipfile.lock for ML.

Train/Serve Skew Prevention

Train-serve skew is the systematic mismatch between features used in training (offline) and features computed at serving time (online or batch inference). It comes from:

• Different code paths - Python offline versus ad-hoc Java/Go online.
• Different data sources or freshness - daily batch table offline versus streaming online.
• Missing time-travel logic offline - training on more complete data than you can serve.

The structural fix is to use the same FeatureView definition for both historical feature generation (offline) and online or batch serving. Feast’s get_historical_features(entity_df, features) performs the AS-OF join automatically, and get_online_features(entity_rows) serves the same definitions from the online store. Tecton’s Feature Services play the same role.

Operationally, also enforce:

1. One source of truth for transformations - a single FeatureView, never reimplemented in serving code.
2. Tracked event and created timestamps for every feature row, so backfills don’t sneak into training [Source: https://developers.openai.com/cookbook/examples/gpt4-1_prompting_guide].
3. Backward-looking windows ending strictly before prediction time - never “next 30 days” as a feature, only as a label.
4. Time-based train/validation splits rather than random splits, so offline evaluation matches production deployment order.
5. Monitoring of feature distributions in production versus training to catch drift early [Source: https://www.youtube.com/watch?v=MqqKT6etxpQ].

When offline AUC is much higher than online AUC and there are no obvious deployment bugs, the culprit is almost always one of: a missing time filter on a join, a feature table without an event_timestamp being treated as a static snapshot, or a label timestamp accidentally used as a prediction timestamp.

Figure 4.5: Train-serve skew prevention - one FeatureView feeds both paths, with monitoring closing the loop.

flowchart TD
    DEF[One FeatureView Definition<br/>transformation + window + timestamps]
    DEF --> OFFP[Offline path:<br/>get_historical_features]
    DEF --> ONP[Online path:<br/>get_online_features]
    OFFP --> TRAIN[Training Dataset<br/>time-based split]
    ONP --> PRED[Prediction Service]
    TRAIN --> MODEL[Trained Model]
    MODEL --> PRED
    PRED --> MON[Production Monitoring<br/>distributions + freshness]
    TRAIN --> BASE[Training Baseline]
    BASE --> DRIFT{Drift detected?}
    MON --> DRIFT
    DRIFT -->|Yes| ALERT[Alert / retrain / fix definition]
    DRIFT -->|No| OK[Continue serving]
    ALERT -.update.-> DEF

Real-Time Features on Streams

Real-time features turn fast-moving event streams into low-latency signals like “number of transactions in the last 5 minutes” or “average click-through rate over the last hour.” Architecturally:

   Kafka / Kinesis  ---->  Flink / Spark Streaming  ---->  Online Store (Redis/DynamoDB)
        |                          |                              |
        v                          v                              v
    Raw events           Sliding-window aggregates         Read in <20ms by model server
                                  |
                                  +----------> Offline Store (Delta/Parquet)
                                                (same values written for training)

The critical design constraint is that the same windowed aggregation logic must produce both offline (training) and online (serving) values. In Tecton this is enforced by the platform: a batch_feature_view or stream_feature_view definition writes to both stores from one specification. In Feast or a DIY system, you typically run a streaming job that updates the online store and lands raw events into the offline store, then run a backfill that recomputes the same aggregations historically.

A canonical fraud-detection example: a Kafka topic of card transactions feeds a Flink job that maintains tumbling and sliding windows of count and sum(amount) for each (user_id, card_id). The aggregates are written to DynamoDB on every update. A backfill job replays the same windowing logic across historical transactions to fill the offline store. The prediction service reads from DynamoDB in ~5-10 ms per request, and data scientists generate training data using point-in-time joins against the offline store. Because the windowing code is one definition, offline and online stay consistent by construction.

Key Takeaway: Match materialization cadence to feature SLA, version FeatureViews semantically, prevent skew by sharing one definition across offline and online paths, and for streaming features keep the same windowing logic on both sides so training and serving stay aligned by construction.

Chapter Summary

Feature engineering and feature stores together convert raw data into the disciplined, reproducible inputs that production ML demands. The chapter began with the bedrock transforms: z-score, min-max, and robust scaling for numerics; one-hot, target, hashing, and learned embeddings for categoricals; TF-IDF and contextual embeddings for text; and lag, rolling, and Fourier features for time series - always with the rule to start simple and add complexity only when it pays.

We then introduced the feature store pattern, an architectural response to the fact that ad-hoc feature code drifts between offline and online environments. A feature store unifies a feature registry, an offline store, an online store, and materialization pipelines under one definition, and uses point-in-time joins to ensure training data only ever reflects information that would have been available at the moment of prediction. We worked through a credit-risk AS-OF join to make the semantics concrete, then compared the leading implementations: Feast for OSS flexibility, Tecton for end-to-end commercial governance and streaming, SageMaker Feature Store for AWS-native managed simplicity, and adjacent options on Databricks and Vertex AI. Finally, we covered the operational discipline that makes pipelines reliable: materialization schedules calibrated to feature SLA, semantic feature versioning pinned to models, structural prevention of train-serve skew through a single shared definition, and streaming-feature pipelines that emit the same windowed aggregates to both offline and online stores.

The architectural payoff is large: when features are defined once, served from two purpose-built stores, and time-traveled through point-in-time joins, the train/serve skew that haunts so many production models disappears at the platform layer. Models built on this foundation are easier to reproduce, easier to debug, and far less likely to regress silently when raw data shifts beneath them.

Key Terms

Chapter 5: Data and Pipeline Versioning

In software engineering, “it works on my machine” is a punchline. In machine learning, it is a crisis. An ML system is not just code; it is the marriage of code, data, environment, and stochastic processes, each of which can drift independently and silently invalidate yesterday’s results. A model that achieved 92% accuracy on Tuesday may produce 87% on Wednesday because someone re-uploaded the training CSV with two extra rows, or because the CUDA driver was patched, or because a random seed was never set in the first place. Worse, when regulators or auditors come knocking, “I think we used roughly this data” is not a defensible answer.

This chapter is about engineering discipline that makes ML reproducible, auditable, and safe to evolve. We will examine the four dimensions of reproducibility, survey the modern toolchain for versioning large datasets (DVC, lakeFS, Delta Lake), discuss how to version pipelines and environments together, and close with how lineage systems like OpenLineage and Marquez stitch the entire graph together for debugging and compliance.

Section 1: Reproducibility in ML

Why git alone is not enough

Git is the canonical tool for code versioning, but ML workflows have at least three artifacts that git handles badly or not at all: large binary datasets, ephemeral execution environments, and stochastic state. A git repository can faithfully record that train.py changed between commits, but if data/train.csv is a 40 GB file that you .gitignored, or if it lives in an S3 bucket that someone overwrote last week, your commit history is an illusion. You can roll back the code, but you cannot roll back the world the code ran in. Reproducible ML requires versioning the entire causal chain that produced a model, not just the source files. [Source: https://doc.dvc.org/start]

The four dimensions of ML reproducibility

Practitioners typically decompose ML reproducibility into four orthogonal dimensions. Each has its own failure modes, its own tooling, and its own conventions. The table below summarizes them.

A reproducible experiment is one where a tuple of (git commit, data version, container image digest, seed/config) uniquely defines the run, and re-running that tuple yields the same result within a documented tolerance. [Source: https://doc.dvc.org/start]

Figure 5.1: Four dimensions of ML reproducibility and their toolchains

graph TD
    R[Reproducible ML Run]
    R --> C[Code Dimension]
    R --> D[Data Dimension]
    R --> E[Environment Dimension]
    R --> S[Randomness Dimension]
    C --> C1[Git commit SHA]
    C --> C2[Pipeline-as-code: Airflow, Kubeflow, dvc.yaml]
    D --> D1[DVC hashes]
    D --> D2[lakeFS commits]
    D --> D3[Delta Lake table version]
    E --> E1[Container image digest]
    E --> E2[pip-tools / Poetry / conda-lock]
    E --> E3[CUDA + driver version]
    S --> S1[Framework seed APIs]
    S --> S2[Deterministic algorithm flags]
    S --> S3[Per-rank seed offsets]
    style R fill:#1f4068,stroke:#58a6ff,color:#fff
    style C fill:#2d4a6b,stroke:#58a6ff,color:#fff
    style D fill:#2d4a6b,stroke:#58a6ff,color:#fff
    style E fill:#2d4a6b,stroke:#58a6ff,color:#fff
    style S fill:#2d4a6b,stroke:#58a6ff,color:#fff

Reproducibility levels

It helps to distinguish degrees of reproducibility, because the engineering cost rises sharply as you tighten the bar:

• Bitwise reproducibility: Identical floating-point outputs across runs. Achievable only with identical hardware, drivers, deterministic kernels, and disabled autotuning. Often required in regulated domains.
• Numerical reproducibility: Outputs match within a small numerical tolerance (e.g., loss within 1e-4, accuracy within 0.1%). Feasible across the same GPU family with deterministic flags set.
• Statistical reproducibility: Distributions of outcomes match, but individual runs differ slightly. Acceptable for research and for many production scenarios with mixed-precision training.
• Conceptual reproducibility: Same methodology, same dataset description, but reimplemented. The bar for academic publication, not for production rollback.

Pick the lowest acceptable level for each pipeline and tool accordingly; chasing bitwise determinism in a Spark-based feature job is usually wasted effort.

Determinism in distributed training

Distributed training adds another axis of variance. Floating-point addition is not associative, which means a parallel reduction across 32 GPUs may sum gradients in a slightly different order on each run. CUDA kernels that use atomic operations are inherently non-deterministic. cuDNN’s autotuner (torch.backends.cudnn.benchmark = True) chooses the fastest convolution algorithm for the current input shape, but the choice can vary between runs, producing different numerical paths. [Source: https://news.ycombinator.com/item?id=44095189]

The mitigation playbook in PyTorch looks like this:

import os, random, numpy as np, torch

def set_seed(seed: int):
    os.environ["PYTHONHASHSEED"] = str(seed)
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
    random.seed(seed); np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.use_deterministic_algorithms(True)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

For DistributedDataParallel jobs, derive each rank’s seed as base_seed + rank so workers diverge predictably rather than coincidentally. DataLoaders need a worker_init_fn that re-seeds NumPy and Python’s random per worker, plus an explicit torch.Generator for shuffling. TensorFlow 2.9+ offers a comparable surface via TF_DETERMINISTIC_OPS=1 and tf.config.experimental.enable_op_determinism(True). None of these flags guarantees bitwise reproducibility across different GPU architectures - an A100 and an H100 will produce slightly different floats no matter what you set. [Source: https://developers.openai.com/cookbook/examples/gpt4-1_prompting_guide]

Analogy: Think of reproducibility like a recipe. The git repo is the cookbook; data versioning is the pantry inventory; the container image is the kitchen; the random seed is the chef’s mood. If any one of them changes silently, you cannot promise the same cake twice.

Key Takeaway: ML reproducibility is a four-dimensional problem - code, data, environment, randomness - and git alone covers only one of those dimensions. A run is reproducible only when the entire tuple (commit, data hash, image digest, seed) is recorded and restorable.

Section 2: Data Versioning Tools

Large datasets do not fit in git, and even when they do, git’s text-oriented diffing makes binary diffs useless. The 2020s gave us three dominant patterns for versioning data at scale: a git-like project tool (DVC), a branchable layer over object storage (lakeFS), and a transactional table format (Delta Lake). They are not strict competitors; they often coexist.

DVC: git-like at the project level

DVC (“Data Version Control”) models itself explicitly on git. You run dvc add data/raw/images/ and DVC computes a content hash of the directory, moves the content into a content-addressable cache, and writes a tiny .dvc metadata file that you commit to git. The actual bytes live in a separate DVC remote - typically S3, GCS, Azure Blob, or a shared filesystem - which you push to with dvc push and pull from with dvc pull. Because the .dvc file contains the hash, the git history of .dvc files is the git history of your data. [Source: https://doc.dvc.org/start]

A typical reproducibility workflow:

git checkout v1.3-paper-submission
dvc pull          # downloads the exact data hashes for this commit
dvc repro         # re-runs the pipeline defined in dvc.yaml

DVC also defines pipelines via dvc.yaml, with dvc.lock recording the input and output hashes of each stage. Re-running dvc repro only re-executes stages whose inputs changed - the same incremental rebuild logic as make, applied to ML stages. [Source: https://doc.dvc.org/start]

Figure 5.2: DVC workflow - git tracks metadata, remote stores the bytes

flowchart LR
    A[Developer workspace<br/>data/raw/images/] -->|dvc add| B[Content-addressable<br/>local cache]
    B -->|writes hash| C[.dvc metadata file]
    C -->|git commit| D[(Git repository<br/>code + .dvc files)]
    B -->|dvc push| E[(DVC Remote<br/>S3 / GCS / Azure)]
    D -->|git checkout v1.3| F[Reproducer workspace]
    F -->|dvc pull| E
    E -->|fetches by hash| F
    F -->|dvc repro| G[Re-executed pipeline<br/>identical outputs]
    style D fill:#1f4068,stroke:#58a6ff,color:#fff
    style E fill:#1f4068,stroke:#58a6ff,color:#fff
    style G fill:#2d6b4a,stroke:#58a6ff,color:#fff

DVC’s sweet spot is a single-team repository with datasets in the tens to low hundreds of gigabytes. It struggles when datasets reach multi-terabyte scale because dvc checkout may need to materialize entire directories on the local filesystem.

lakeFS: branchable over object storage

lakeFS sits at a different layer. Instead of versioning a project, it versions an entire object store bucket. It exposes S3 (or GCS, or Azure Blob) through a versioning layer that supports branches, commits, and merges - operations that should feel familiar to any git user, but applied to potentially petabytes of objects. [Source: https://www.youtube.com/watch?v=efnw2QvlhZM]

Branches are cheap because lakeFS uses copy-on-write semantics: creating an experiment-202506 branch from main does not duplicate any data. Objects are copied only when they are modified on the branch. This makes it cheap and safe to run “what-if” experiments on a fork of production data, then merge results back or throw them away.

Figure 5.3: lakeFS branching model over object storage

flowchart TD
    M[(main branch<br/>production data)]
    M -->|branch zero-copy| F[feature-eng-v4]
    M -->|branch zero-copy| E[experiment-202506]
    F -->|Spark writes new features| F1[s3://repo/feature-eng-v4/features/]
    E -->|exploratory rewrites| E1[s3://repo/experiment-202506/labels/]
    F -->|model wins -> merge| M
    E -->|model loses -> delete| X[discarded, no storage cost]
    M -->|commit ID| L[Lineage: model trained at<br/>s3://repo@commitId/...]
    style M fill:#1f4068,stroke:#58a6ff,color:#fff
    style F fill:#2d4a6b,stroke:#58a6ff,color:#fff
    style E fill:#2d4a6b,stroke:#58a6ff,color:#fff
    style X fill:#6b2d2d,stroke:#ff8b8b,color:#fff

A practical lakeFS-anchored workflow:

1. Data engineer branches main to feature-eng-v4 for a new feature pipeline.
2. Spark jobs run against s3://repo/feature-eng-v4/..., writing experimental feature tables.
3. ML team trains a model referencing the branch ref, logs the commit ID alongside MLflow metrics.
4. If the model wins, merge the branch into main; if not, delete it.

lakeFS shines in centralized data lakes used by multiple teams. It does not replace experiment tracking, and it does not give you SQL/ACID semantics on tables - for that, you layer Delta Lake (or Iceberg) on top of lakeFS-backed paths.

Delta Lake and Iceberg: ACID time travel for tables

Delta Lake is a transactional table format. Underneath, it is just Parquet files in object storage, but a _delta_log/ directory records every transaction (add file, remove file, update schema) as a JSON or checkpoint entry. Engines like Spark, Flink, and Trino read that log to determine which files constitute a given table version, giving you ACID guarantees for MERGE, UPDATE, and DELETE operations on data lakes that historically had none. [Source: https://www.youtube.com/watch?v=8I0jMEs470o]

Figure 5.4: Delta Lake time-travel architecture - transaction log over Parquet

flowchart LR
    Q[Query engine<br/>Spark / Trino / Flink]
    Q -->|VERSION AS OF 37| L[_delta_log/]
    L --> L1[00000.json: add file A,B]
    L --> L2[00001.json: add file C]
    L --> L3[00037.json: remove B, add D]
    L --> L4[00050.checkpoint.parquet]
    L1 -.resolves to.-> P
    L2 -.resolves to.-> P
    L3 -.resolves to.-> P
    P[Parquet files in object storage]
    P --> PA[part-A.parquet]
    P --> PB[part-B.parquet]
    P --> PC[part-C.parquet]
    P --> PD[part-D.parquet]
    Q -->|reads only files in version 37| PA
    Q --> PC
    Q --> PD
    style Q fill:#1f4068,stroke:#58a6ff,color:#fff
    style L fill:#2d4a6b,stroke:#58a6ff,color:#fff
    style P fill:#2d4a6b,stroke:#58a6ff,color:#fff

Time travel is the headline reproducibility feature:

-- Read the exact features table the model was trained on
SELECT * FROM features.user_features VERSION AS OF 37;
SELECT * FROM features.user_features TIMESTAMP AS OF '2026-05-01 09:00:00';

Logging the version number alongside an MLflow run becomes a complete, queryable record of the training data. Delta Lake 2024-2025 features include Change Data Feed (CDF) for incremental retraining (read only the rows that changed since the last training run) and Delta UniForm plus Delta Kernel, which let non-Spark engines and Iceberg-aware tools read Delta tables, reducing format lock-in. [Source: https://www.cliffsnotes.com/study-notes/28411172]

Apache Iceberg is the closest competitor: similar ACID + time travel guarantees, different metadata layout (manifest lists rather than a _delta_log), and historically stronger multi-engine support. The two formats are converging functionally; choose based on your engine ecosystem (Databricks-heavy shops gravitate to Delta; Trino/Snowflake/AWS-heavy shops often pick Iceberg).

Tool comparison by scale and use case

Use the table below as a starting decision matrix, not a verdict. The three tools are often combined.

In practice, large organizations commonly stack them: lakeFS provides bucket-level branching, Delta Lake provides ACID tables inside lakeFS, and DVC manages project-local slices, models, and configs in each ML repo - with each tool’s version pointer logged in the experiment tracker so a model can be traced to a lakeFS commit, a Delta version, and a DVC hash. [Source: https://www.youtube.com/watch?v=efnw2QvlhZM]

Analogy: DVC is a household pantry inventory - granular, but only for one family. lakeFS is the warehouse’s branch system - cheap copies of the entire inventory for testing layouts. Delta Lake is a ledger for a specific shelf - every transaction recorded so you can replay the shelf’s state at any moment.

Key Takeaway: DVC, lakeFS, and Delta Lake operate at different layers - project, bucket, table - and the right choice depends on whether your reproducibility problem is per-project, lake-wide, or table-centric. Many mature organizations layer all three.

Section 3: Pipeline and Environment Versioning

Versioning data and code is necessary but not sufficient. The pipeline that wires them together, and the environment in which it executes, must also be captured.

Docker as the environment unit of versioning

A container image is the standard packaging format for an ML environment: it encodes the base OS, CUDA and cuDNN versions, Python interpreter, system libraries (libjpeg, libsndfile, libglib), and all Python dependencies into a single immutable artifact identified by a content-addressable digest (a SHA-256 hash). Two engineers running docker pull myorg/ml@sha256:abc123... are guaranteed to execute against byte-identical environments. [Source: https://news.ycombinator.com/item?id=44095189]

A practical Dockerfile for a GPU training job:

FROM nvidia/cuda:12.1.0-cudnn9-runtime-ubuntu22.04
RUN apt-get update && apt-get install -y \
    git python3 python3-pip && \
    rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY requirements.lock /app/
RUN pip install --no-cache-dir -r requirements.lock
COPY . /app
CMD ["python3", "train.py", "--config", "config/exp1.yaml"]

Two discipline rules to apply religiously:

1. Log the image digest, not the tag. Tags like :latest or even :v1.3 can be re-pointed. The SHA digest cannot.
2. Treat the image as immutable. If you need to change anything - even one dependency patch - rebuild and re-tag. Never docker exec your way into reproducing fixes.

Locking dependencies: pip-tools, Poetry, conda-lock

The Dockerfile is only as reproducible as the dependency resolution that produced it. RUN pip install -r requirements.txt is a moving target - transitive dependencies can shift on every rebuild. Lockfiles solve this:

• pip-tools (pip-compile) produces a fully pinned requirements.txt from a higher-level requirements.in, including transitive dependencies and hashes.
• Poetry does the same for projects defined in pyproject.toml, producing a poetry.lock with exact versions and content hashes.
• conda-lock generates platform-specific lockfiles from environment.yml, capturing the entire conda solve including non-Python packages.

In the Dockerfile above, copying requirements.lock (the resolved lockfile) instead of requirements.txt ensures the image is byte-deterministic given the same base image. Combined with a digest-pinned base image (FROM nvidia/cuda@sha256:...), the entire build is reproducible.

Pipeline-as-code

Pipeline-as-code is the principle that the orchestration topology - which step runs, in what order, with what inputs and outputs - lives in version-controlled source files, not in a UI someone clicked on six months ago. The pipeline definition becomes a first-class artifact that can be reviewed, diffed, branched, and rolled back like any other code. [Source: https://doc.dvc.org/start]

Concrete forms across the modern ecosystem:

• Airflow DAGs: Python files defining tasks and dependencies; deployed via GitOps.
• Kubeflow Pipelines / Argo Workflows: YAML or Python-generated YAML describing containerized steps on Kubernetes.
• DVC pipelines: dvc.yaml defines stages with input/output hashing.
• Dagster, Prefect, Metaflow: Python-native, with built-in lineage and run-tracking primitives.

A representative dvc.yaml snippet:

stages:
  prepare_features:
    cmd: python src/features.py --input data/raw --output data/features
    deps:
      - data/raw
      - src/features.py
    outs:
      - data/features
  train:
    cmd: python src/train.py --features data/features --model models/churn.pt
    deps:
      - data/features
      - src/train.py
    outs:
      - models/churn.pt
    metrics:
      - metrics.json:
          cache: false

When inputs change, dvc repro re-executes only the affected stages, and the dvc.lock file records the exact input/output hashes for every stage in the run.

Compute environment snapshots

Beyond the container image, a fully reproducible run also depends on runtime configuration: the number of GPUs, the GPU model, the driver version, the values of OMP_NUM_THREADS and CUDA_VISIBLE_DEVICES, and the cloud instance type. Mature teams log these as runtime metadata alongside the experiment:

{
  "image_digest": "sha256:abc123...",
  "git_sha": "f4a2c8e",
  "data_dvc_hash": "md5:9f1ec4...",
  "gpu_model": "NVIDIA A100-SXM4-80GB",
  "gpu_count": 4,
  "cuda_version": "12.1",
  "driver_version": "535.104.05",
  "seed": 1234,
  "deterministic_flags": ["torch.use_deterministic_algorithms=True"]
}

This metadata bundle is what makes a run not just reproducible in principle but auditable in practice. Six months from now, when accuracy regression hits production, this JSON tells you the exact configuration to recreate.

Analogy: Pipeline-as-code is to ML what infrastructure-as-code is to operations: stop describing what you want by clicking, start declaring it in a file that survives staff turnover.

Key Takeaway: Reproducible environments require digest-pinned container images built from locked dependency files, with pipelines defined as code and runtime metadata logged for every run. The image digest plus the pipeline commit plus the data version is the closure of an experiment.

Section 4: Lineage and Provenance

So far we have versioned the ingredients. Data lineage versions the recipe in motion: which job, running which code, on which inputs, produced which outputs, and when. In MLOps, lineage is the graph that ties raw sources to features to models to predictions. It is the difference between “this model was trained on user data” and “this exact training run, with this commit SHA, on these specific Delta table versions, produced this model artifact.” [Source: https://www.youtube.com/watch?v=MqqKT6etxpQ]

OpenLineage: the open standard

OpenLineage is a tool-agnostic JSON specification for emitting lineage events. Its data model centers on four entities:

• Job: A logical unit of work (an Airflow task, a Spark job, a dbt model), identified by name + namespace.
• Run: A single execution of a Job, identified by a runId (typically a UUID).
• Dataset: A logical dataset read or written (name + namespace); can represent tables, files, streams, or ML model artifacts.
• Event: JSON messages of type START, COMPLETE, or FAIL emitted during a run, listing inputs, outputs, and extensible facets. [Source: https://www.youtube.com/watch?v=MqqKT6etxpQ]

Facets are pluggable schemas that carry the rich metadata: schema (columns and types), columnLineage (per-column upstream mappings), dataQualityMetrics (row counts, null rates), errorMessage (stacktrace on failure), nominalTime, parent (for nested workflows), sourceCode, and sql. You can also define custom ML-specific facets for hyperparameters, training metrics, or feature store references. [Source: https://www.youtube.com/watch?v=MqqKT6etxpQ]

Integrations: Airflow, Spark, dbt

The point of an open standard is that you do not emit lineage by hand. Integrations hook into orchestrators and engines and emit events automatically:

• Airflow: The OpenLineage provider hooks into task lifecycle callbacks. It inspects operators - PostgresOperator, BigQueryInsertJobOperator, SparkSubmitOperator - to infer inputs and outputs, emits START when a task begins, and emits COMPLETE or FAIL with error facets when it ends.
• Spark: A listener jar attaches to Spark’s logical and physical plans. As DataFrames are read from and written to tables and files, the listener emits jobs and datasets along with schema and (where derivable) column-level lineage.
• dbt: A plugin reads dbt’s manifest.json and run_results.json to map every dbt model to an OpenLineage job, with rich column-level lineage derived directly from compiled SQL. [Source: https://www.youtube.com/watch?v=MqqKT6etxpQ]

Marquez: storage and visualization

Marquez is the reference open-source implementation that ingests OpenLineage events, stores them, and exposes a UI. It maintains a time-aware graph of datasets, jobs, and runs, so you can ask “what did the lineage graph look like on May 3rd?” and get an answer. The UI offers a dataset view (upstream/downstream jobs, recent runs, row counts), a job view (input and output datasets with run history), and global lineage navigation that lets you expand multiple hops in either direction. Where columnLineage facets are present, Marquez can visualize per-column dependencies - critical for fairness audits and PII tracking. [Source: https://www.youtube.com/watch?v=MqqKT6etxpQ]

End-to-end: dataset → features → model → predictions

A canonical ML lineage chain looks like this:

1. Airflow ingest_events writes raw.events from Kafka.
2. dbt model stg_events reads raw.events, writes stg.events, with column-level lineage from SQL.
3. Spark job user_features_job reads stg.events, writes features.user_features.
4. Airflow train_churn_model reads features.user_features, writes models.churn_model:v1.3 with run facets capturing hyperparameters, git SHA, Delta table version, and metrics.
5. Airflow batch_inference reads models.churn_model:v1.3 and features.user_features, writes predictions.churn_scores.

In Marquez, you can start from predictions.churn_scores, walk upstream through the model artifact, the feature table, the dbt staging model, and finally to raw.events and the ingestion job that wrote it. Every node carries its run history, schema, and facets. The same graph can be walked downstream from any column: drop events.event_type and Marquez shows you the exact set of features, models, and predictions that depend on it. [Source: https://www.youtube.com/watch?v=MqqKT6etxpQ]

Figure 5.5: End-to-end data lineage graph - dataset to predictions

graph LR
    K[(Kafka<br/>events stream)]
    K -->|Airflow: ingest_events| RE[raw.events]
    RE -->|dbt: stg_events| SE[stg.events]
    SE -->|Spark: user_features_job| UF[features.user_features<br/>Delta v=42]
    UF -->|Airflow: train_churn_model| MC[models.churn_model:v1.3<br/>git SHA + image digest]
    UF -->|Airflow: batch_inference| CP[predictions.churn_scores]
    MC -->|Airflow: batch_inference| CP
    CP -->|consumed by| APP[Downstream apps<br/>CRM, dashboards]
    style K fill:#1f4068,stroke:#58a6ff,color:#fff
    style MC fill:#2d6b4a,stroke:#58a6ff,color:#fff
    style CP fill:#2d6b4a,stroke:#58a6ff,color:#fff
    style APP fill:#4a2d6b,stroke:#58a6ff,color:#fff

GDPR and the EU AI Act: lineage as auditability

Regulatory pressure is the operational reason lineage moves from “nice to have” to “must have.” The EU’s General Data Protection Regulation grants users a right to be forgotten and requires controllers to demonstrate which datasets contain a subject’s data and how that data has been processed. The EU AI Act, phased in through 2025-2026, requires providers of high-risk AI systems to maintain detailed technical documentation of training datasets, data governance procedures, and traceability of decisions. [Source: https://arxiv.org/html/2603.20576v1]

With OpenLineage + Marquez, common compliance questions become graph queries:

• “Show all models trained on PII-tagged datasets.” → Filter datasets by a sensitivity: PII facet, walk downstream to model artifacts.
• “User X has revoked consent. Which models need retraining?” → Identify datasets containing user X, walk forward to all downstream models.
• “Prove this credit-scoring model was not trained on the prohibited applicant_race column.” → Use column-level lineage from the model’s training inputs back to source columns.

Without lineage, these questions become weeks of forensic SQL across logs. With lineage, they are dashboard clicks.

Debugging and impact analysis

Beyond compliance, lineage pays for itself in incident response. A typical scenario: the churn model’s AUC drops from 0.80 to 0.72 overnight. The lineage walk:

1. Open the latest train_churn_model run in Marquez.
2. Inspect run facets: git SHA, data versions, hyperparameters - unchanged.
3. Walk upstream to features.user_features - the latest write run shows row count down 30% and events_30d mostly NULL.
4. Walk upstream to stg.events - the upstream dbt model’s run facets reveal a schema-mismatch error.
5. Fix the source schema mapping, re-run the chain, and verify the whole downstream graph turns green.

The same graph supports impact analysis for planned changes: before deprecating a feature column, walk downstream to enumerate every model that consumes it and notify the owners.

Analogy: Lineage is the camera roll of your data. Versioning tells you what happened; lineage tells you who did it, when, and what they touched.

Key Takeaway: Lineage closes the reproducibility loop by recording the runtime graph of jobs, runs, and datasets. OpenLineage standardizes the event format; Marquez stores and visualizes the graph; together they convert “we think we know what trained this model” into an auditable, queryable record.

Chapter Summary

Reproducibility in ML is engineered, not assumed. We started by decomposing it into four dimensions - code, data, environment, randomness - and argued that git, by itself, covers only one. Each dimension has its own toolchain and its own failure mode, and a run is reproducible only when all four are pinned simultaneously.

For data, we surveyed three dominant patterns. DVC brings git-like ergonomics to project-scoped datasets, storing hashes in git and bytes in cloud remotes. lakeFS lifts that model up to entire object stores, offering zero-copy branches over petabyte-scale data lakes. Delta Lake (and its sibling Iceberg) provides ACID transactions and time-travel queries at the table level, integrating tightly with Spark and lakehouse architectures. The three coexist more often than they compete.

For pipelines and environments, we treated the container image as the unit of environment versioning, with digest-pinned images built from locked dependency files. Pipeline-as-code moves orchestration topology from UIs into version-controlled files, making the wiring as auditable as the steps. A complete experiment closure is (git SHA, data hash, image digest, seed, runtime config).

Finally, lineage turned the static version-pinning story into a dynamic graph. OpenLineage defines an open event format for jobs, runs, and datasets; Marquez ingests and visualizes those events; integrations with Airflow, Spark, and dbt capture lineage automatically without manual instrumentation. The resulting graph supports debugging, impact analysis, and regulatory auditability for GDPR and the EU AI Act.

The discipline is unglamorous but cumulative. Each pinned dimension is a future incident you do not need to investigate, each lineage event is a regulatory question you will not have to research, and each lockfile is a “works on my machine” conversation you do not have to have.

Key Terms

Chapter 6: Pipeline Orchestration Frameworks

By the end of Chapter 5, you had a working understanding of how features and training code can be packaged into reusable steps. But knowing what the steps are is only half the problem. In production, somebody (or something) has to wake up every night at 02:00 UTC, run the feature build for yesterday’s partition, wait for it to finish, kick off training for three customer segments in parallel, only deploy the new model if its evaluation metrics beat the previous one, and retry the whole thing tomorrow if the warehouse was flaky. That “somebody” is the orchestrator, and this chapter is about how to choose and operate one for ML.

Think of an orchestrator the way air traffic control thinks about aircraft. Each plane (task) has its own engines, fuel, and pilots; ATC does not fly the plane, but it sequences takeoffs, prevents collisions, reroutes around storms, and decides whether the flight is allowed to land at all. Likewise, an ML orchestrator does not train your model; it decides when training runs, what it depends on, what to do when it fails, and how to backfill last week’s data after you fix a bug.

Section 1: Orchestration Concepts

Sub-topic 1.1: DAGs, Tasks, Operators, and Executors

Every modern orchestrator is built on the same four abstractions, even when the vocabulary differs. A Directed Acyclic Graph (DAG) describes the dependency structure of work: nodes are units of computation, edges are “must run before” relationships, and the graph has no cycles, which guarantees the schedule can finish. A task (sometimes called a step, op, or component) is a single unit of work, typically a Python function, a SQL query, or a containerized job. An operator is a reusable template for a class of tasks — Airflow’s BashOperator, PythonOperator, and KubernetesPodOperator are canonical examples. An executor is the engine that actually runs the tasks: a process pool on a single machine, a Celery cluster, or a Kubernetes scheduler [Source: https://eng.lyft.com/orchestrating-data-pipelines-at-lyft-comparing-flyte-and-airflow-72c40d143aad].

The DAG is the plan; the executor is the labor. Confusing the two is the most common newcomer mistake: a “scaled” orchestrator usually means you scaled the executor, not the scheduler. The scheduler — the brain that scans DAG definitions, decides what is ready, and dispatches it — is almost always the bottleneck before the executor is, especially when DAGs have thousands of small tasks [Source: https://eng.lyft.com/orchestrating-data-pipelines-at-lyft-comparing-flyte-and-airflow-72c40d143aad].

Figure 6.1: Example DAG structure with parallel fan-out and downstream join.

flowchart TD
    A[ingest_raw] --> B[validate_schema]
    B --> C[build_features]
    C --> D1[train_retail]
    C --> D2[train_smb]
    C --> D3[train_enterprise]
    D1 --> E[evaluate_all]
    D2 --> E
    D3 --> E
    E --> F[register_winners]

Sub-topic 1.2: Imperative vs Declarative

ML practitioners coming from notebooks usually find imperative workflows comfortable: you write Python that does things, and decorators like @task or @step turn the function calls into a graph at runtime. Prefect and Metaflow exemplify this. Declarative workflows, by contrast, ask you to describe the desired graph and let the system schedule it; Argo Workflows expressed as Kubernetes YAML is the purest example. Airflow sits in between — Python files that read declaratively but execute imperatively at parse time [Source: https://www.zenml.io/blog/flyte-vs-airflow].

The trade-off is the usual one between expressiveness and analyzability. Imperative DAGs are easy to write but harder to inspect statically (the graph may depend on runtime branches). Declarative DAGs are easier for the platform team to validate, secure, and reason about, but feel verbose to data scientists.

Sub-topic 1.3: Materialization vs Orchestration

A subtle but important conceptual split has emerged. Orchestration-first systems (Airflow, Prefect, Argo) think in terms of tasks: did this job run, and did it succeed? Materialization-first systems (Dagster, and to a degree Flyte, Kubeflow Pipelines, ZenML) think in terms of assets: is this dataset, feature table, or model up-to-date relative to its inputs? [Source: https://www.union.ai/blog-post/we-compared-the-data-models-of-every-major-ai-orchestrator].

The distinction matters for ML because the canonical question is rarely “did training run?” — it is “is this model version fresh given the latest features and the latest code?” Asset-aware systems can answer that natively; task-aware systems push the question into external tools (MLflow, a model registry, a custom database).

Sub-topic 1.4: Triggers — Cron, Sensors, Events

Pipelines need a reason to start. Three patterns dominate:

• Cron schedules fire on calendar time (“every day at 02:00”). Simple, universal, and what most teams reach for first.
• Sensors poll an external system until a condition is true (the file landed in S3, the upstream Airflow DAG finished, the queue has 10k messages). They are convenient but can monopolize worker slots; modern Airflow offers deferrable sensors that release the slot and resume when a trigger fires [Source: https://www.prompts.ai/blog/tools-orchestrating-machine-learning-workflows.html].
• Events push a trigger from outside (Pub/Sub message, cloud-storage notification, webhook). Event-driven is generally preferred over polling at scale [Source: https://www.prompts.ai/blog/tools-orchestrating-machine-learning-workflows.html].

Key Takeaway: Every orchestrator decomposes work into a DAG of tasks executed by some executor and triggered by cron, sensor, or event. The deeper choice is paradigm: task-centric “did it run?” vs asset-centric “is it fresh?” — this single decision shapes how you will model lineage and incremental recomputation for the next several years.

Section 2: General-Purpose Orchestrators

General-purpose orchestrators were not designed for ML; they were designed for ETL, reporting, and arbitrary batch jobs. That history is both their strength (mature, battle-tested, huge ecosystem) and their weakness (they treat models and datasets as opaque side effects).

Sub-topic 2.1: Apache Airflow Architecture

Airflow is the de facto standard for data engineering. Its architecture has four moving parts: a metadata database (Postgres or MySQL) that stores DAG definitions, run history, and task state; a scheduler that scans DAG files and dispatches ready tasks; a webserver that renders the UI; and one or more executors (Local, Celery, Kubernetes, or LocalKubernetes) that actually run the tasks [Source: https://eng.lyft.com/orchestrating-data-pipelines-at-lyft-comparing-flyte-and-airflow-72c40d143aad].

Figure 6.2: Apache Airflow architecture — scheduler, executor, workers, webserver, and metadata DB.

flowchart TD
    subgraph Authors[DAG Authors]
        DAGs[DAG Files<br/>Python]
    end
    subgraph Control[Control Plane]
        Scheduler[Scheduler]
        Webserver[Webserver / UI]
    end
    subgraph State[State]
        MetaDB[(Metadata DB<br/>Postgres / MySQL)]
    end
    subgraph Compute[Compute Plane]
        Executor[Executor<br/>Celery / K8s / Local]
        W1[Worker 1]
        W2[Worker 2]
        W3[Worker N]
    end
    DAGs --> Scheduler
    Scheduler <--> MetaDB
    Webserver <--> MetaDB
    Scheduler --> Executor
    Executor --> W1
    Executor --> W2
    Executor --> W3
    W1 --> MetaDB
    W2 --> MetaDB
    W3 --> MetaDB

Airflow 2.x added the TaskFlow API, datasets as first-class citizens, a stable scheduler, a REST API, and DAG versioning, dramatically modernizing the experience [Source: https://www.union.ai/blog-post/we-compared-the-data-models-of-every-major-ai-orchestrator]. Airflow’s ML story is mostly through operators: KubernetesPodOperator for containerized training, DatabricksRunNowOperator, SageMakerTrainingOperator, and the like. Airflow does the orchestration; the operators delegate the actual work [Source: https://www.zenml.io/blog/flyte-vs-airflow].

The classical weakness for ML is short, fine-grained tasks. Airflow’s scheduler was tuned for a few hundred long ETL jobs per DAG, not the ten-thousand-step hyperparameter sweep [Source: https://eng.lyft.com/orchestrating-data-pipelines-at-lyft-comparing-flyte-and-airflow-72c40d143aad].

Sub-topic 2.2: Prefect 2.x and Dagster

Prefect 2 (“Orion”) is the most Pythonic of the major orchestrators. A flow is a normal Python function decorated with @flow; tasks are functions decorated with @task. The control plane (Prefect Cloud or self-hosted server) is lightweight; workers pull flow runs from work pools and execute them on whatever infrastructure you point them at. Prefect is well-suited to “I want to wrap my existing ML training script with a few decorators and get observability” [Source: https://www.prompts.ai/blog/tools-orchestrating-machine-learning-workflows.html].

Dagster takes the radical position that the unit of orchestration should be the asset, not the task. A @asset declaration says “this dataset/model exists, and it depends on these other assets.” The runtime then materializes assets in the right order, exposes partitions and freshness in the UI, and runs backfills by selecting partition ranges [Source: https://www.union.ai/blog-post/we-compared-the-data-models-of-every-major-ai-orchestrator]. For teams already using dbt for the warehouse layer and wanting that same lineage philosophy for features and models, Dagster is a strong fit.

Sub-topic 2.3: Argo Workflows on Kubernetes

Argo Workflows is the Kubernetes-native workflow engine: workflows are CRDs (Custom Resource Definitions), every step is a pod, and the controller reconciles desired state with cluster reality. It is declarative YAML, scales with the cluster, and is the substrate that Kubeflow Pipelines actually compiles to under the hood [Source: https://asya.sh/docs/comparisons/as-ml-pipeline-tool/]. Argo is rarely used directly by data scientists — it is too low-level — but it is the right answer when you want a thin, K8s-native orchestration layer that other systems can compile into.

Sub-topic 2.4: Strengths and Weaknesses for ML

Key Takeaway: General-purpose orchestrators are mature and scalable, but with the exception of Dagster they treat ML artifacts as opaque side effects. If you choose one, plan to pair it with an external model registry and experiment tracker — the orchestrator alone will not give you lineage.

Section 3: ML-Native Orchestrators

ML-native orchestrators start from a different premise: pipelines are sequences of typed ML steps that produce versioned artifacts (datasets, models, metrics), and the platform should track those artifacts as first-class entities.

Sub-topic 3.1: Kubeflow Pipelines

Kubeflow Pipelines (KFP) v2 is the canonical open-source K8s-native ML orchestrator. You write pipelines in a Python DSL that compiles to a Kubernetes workflow (Argo or Tekton, depending on the install) [Source: https://www.zenml.io/blog/metaflow-vs-kubeflow]. Every step is a containerized component with typed inputs and outputs; artifact lineage is tracked in ML Metadata (MLMD), a service that records executions, contexts, and artifacts so the UI can show “this model came from these features, which came from this dataset” [Source: https://asya.sh/docs/comparisons/as-ml-pipeline-tool/].

The price of admission is operational. Running Kubeflow means running a Kubernetes cluster, the KFP control plane, MLMD, MinIO or another object store, and ideally Istio. Vertex AI Pipelines is Google’s managed offering that speaks the KFP SDK, letting teams skip most of the platform work [Source: https://www.zenml.io/blog/metaflow-vs-kubeflow].

Figure 6.3: Kubeflow Pipelines architecture on Kubernetes.

flowchart TD
    SDK[KFP Python SDK] -->|compile| Spec[Pipeline Spec<br/>YAML / IR]
    Spec --> API[KFP API Server]
    API --> Argo[Argo / Tekton<br/>Workflow Controller]
    API <--> MLMD[(ML Metadata<br/>MLMD)]
    Argo --> P1[Step Pod 1]
    Argo --> P2[Step Pod 2]
    Argo --> P3[Step Pod N]
    P1 --> Store[(Artifact Store<br/>MinIO / GCS / S3)]
    P2 --> Store
    P3 --> Store
    P1 --> MLMD
    P2 --> MLMD
    P3 --> MLMD
    UI[KFP UI] <--> API
    UI <--> MLMD

Sub-topic 3.2: Metaflow

Metaflow was open-sourced by Netflix and is unapologetically optimized for data scientists. A flow is a Python class with @step methods; you run it locally with python flow.py run, and the same code can scale out by adding @batch or @kubernetes to a step [Source: https://www.zenml.io/blog/metaflow-vs-kubeflow]. Anything assigned to self inside a step becomes a versioned artifact, automatically persisted to S3 (or another datastore) and queryable by run ID via the Metaflow client.

Metaflow’s superpower is local-to-cloud transparency: a data scientist iterating in a notebook can run the same flow against a 10-row sample locally and a 10-million-row partition on AWS Batch with no code change [Source: https://mlops.community/blog/zenml-vs-flyte-vs-metaflow]. The deliberate trade-off is that Metaflow does less platform-level governance than Flyte or KFP.

Sub-topic 3.3: ZenML and Flyte

Flyte originated at Lyft for ML at scale and is the strongest “K8s-native, typed, reproducible” option. Tasks are Python functions with type hints; Flyte uses those hints to serialize/deserialize artifacts, to cache task outputs based on content addressing, and to validate the workflow graph at compile time [Source: https://eng.lyft.com/orchestrating-data-pipelines-at-lyft-comparing-flyte-and-airflow-72c40d143aad]. Resource specifications (GPU count, memory, accelerator type) are decorator arguments. Flyte’s multi-tenant, multi-project design makes it appealing for centralized ML platforms.

ZenML is a meta-orchestrator: rather than executing pipelines itself, it compiles ML-centric definitions into the backend you already have (Airflow, KFP, Kubernetes, AWS Step Functions, Vertex). On top of that it provides typed Artifact classes (Dataset, Model, Evaluation), a central metadata store, and stack abstractions for swapping in different artifact stores or experiment trackers without rewriting your code [Source: https://mlops.community/blog/zenml-vs-flyte-vs-metaflow]. ZenML is attractive when you do not want to bet the farm on a single execution engine.

Sub-topic 3.4: Vertex AI and SageMaker Pipelines

The cloud vendors offer managed equivalents. Vertex AI Pipelines (GCP) runs KFP-compatible pipelines without the Kubernetes ops burden, integrates with Vertex Metadata for lineage, and bills per pipeline-second. SageMaker Pipelines (AWS) provides a CI/CD-flavored DSL that integrates tightly with SageMaker Training, Processing, and Model Registry; it is the path of least resistance for teams already on SageMaker. Both trade flexibility for managed convenience.

Sub-topic 3.5: The Big Comparison

Key Takeaway: ML-native orchestrators trade operational simplicity for first-class artifacts, typed interfaces, and built-in caching. KFP and Flyte are the K8s-heavy ML platforms; Metaflow optimizes for developer ergonomics; ZenML lets you swap backends; the cloud vendors offer managed escape hatches.

Section 4: Operationalizing Pipelines

Choosing an orchestrator is the easy half. Operating one well — so that it survives bad data, flaky infrastructure, half-written deployments, and frantic backfills — requires a small set of disciplines that are essentially the same across every tool.

Sub-topic 4.1: A Worked Example — A Daily Training DAG

Before diving into operational mechanics, consider a concrete pipeline. It builds features for a date partition, validates them, conditionally trains one model per customer segment, evaluates each, and registers the winners. Here is the same DAG expressed in Prefect:

from prefect import flow, task
from prefect.tasks import task_input_hash
from datetime import timedelta, datetime

@task(retries=3, retry_delay_seconds=30, retry_jitter_factor=0.2,
      cache_key_fn=task_input_hash, cache_expiration=timedelta(hours=24))
def build_features(date: str) -> str:
    # idempotent: writes to s3://features/date=YYYY-MM-DD/
    return f"s3://features/date={date}/"

@task(retries=2)
def validate_features(features_uri: str) -> float:
    # returns a data-quality score in [0, 1]
    return 0.97

@task(retries=2, timeout_seconds=3600)
def train_segment(features_uri: str, segment: str, model_type: str) -> dict:
    # writes to s3://models/{model_type}/{segment}/{run_id}/
    return {"segment": segment, "auc": 0.84, "model_uri": "..."}

@task
def register_if_better(result: dict, baseline_auc: float) -> bool:
    return result["auc"] > baseline_auc

@flow(name="daily-training")
def daily_training(date: str = None,
                   model_type: str = "xgboost",
                   segments: list[str] = ("retail", "smb", "enterprise"),
                   baseline_auc: float = 0.80):
    date = date or datetime.utcnow().strftime("%Y-%m-%d")
    features = build_features(date)
    dq = validate_features(features)

    if dq < 0.95:
        # conditional execution: skip training on bad data
        return {"status": "skipped", "reason": "data-quality"}

    # dynamic fan-out: one training task per segment
    results = train_segment.map(features, segments, model_type)
    promotions = [register_if_better(r, baseline_auc) for r in results]
    return promotions

Notice four things. (1) Retries with jitter are declared on the decorator, not embedded in code. (2) The feature build is cached by its input hash — re-running the flow with the same date skips recomputation. (3) Conditional execution is a plain Python if; Prefect treats it as a visible branch in the run graph [Source: https://www.prompts.ai/blog/tools-orchestrating-machine-learning-workflows.html]. (4) The training step fans out dynamically across segments via .map(), a pattern that has direct equivalents in Airflow (.expand()), Dagster (DynamicOutput), Flyte (@dynamic), and KFP (dsl.ParallelFor) [Source: https://mlops.community/blog/zenml-vs-flyte-vs-metaflow].

Sub-topic 4.2: Retries, Timeouts, and Idempotency

Retries are how pipelines absorb the natural flakiness of distributed systems. The default playbook: 3-5 retries with exponential backoff capped at 30-60 minutes, plus jitter to avoid thundering-herd patterns [Source: https://www.prompts.ai/blog/tools-orchestrating-machine-learning-workflows.html]. Airflow expresses this with retries, retry_delay, retry_exponential_backoff, and max_retry_delay; Prefect adds retry_jitter_factor; Dagster uses RetryPolicy; Flyte uses @task(retries=..., retry_delay=...); KFP delegates to the underlying Argo/Vertex retryStrategy [Source: https://eng.lyft.com/orchestrating-data-pipelines-at-lyft-comparing-flyte-and-airflow-72c40d143aad].

Retries are dangerous without idempotency. If build_features("2025-01-15") writes a row per record and you retry it three times, you have triple-counted rows unless the write is partitioned and overwriting. The canonical patterns are:

• Partitioned writes keyed by date or dataset_id (INSERT OVERWRITE PARTITION date='2025-01-15', or writing to s3://features/date=2025-01-15/ and overwriting the prefix) [Source: https://www.prompts.ai/blog/tools-orchestrating-machine-learning-workflows.html].
• Versioned artifact paths for models, e.g. s3://models/xgboost/retail/{run_id}/model.pkl, with downstream consumers loading by run_id rather than “latest”.
• Deduplication keys on warehouse writes — a business key plus a run_id lets MERGE/UPSERT handle replays cleanly.

Treat the orchestrator’s state as ephemeral and the data system’s state as the source of truth. If a task succeeds but the orchestrator crashes before recording it, the rerun should detect the existing output and skip the work, not duplicate it.

Figure 6.4: Retry-with-backoff state machine for a single task.

stateDiagram-v2
    [*] --> Pending
    Pending --> Running: scheduler dispatches
    Running --> Success: exit 0
    Running --> Failed: exit != 0 / timeout
    Failed --> Backoff: attempts < max
    Backoff --> Running: wait = base * 2^n + jitter
    Failed --> DeadLetter: attempts >= max
    Success --> [*]
    DeadLetter --> [*]: alert on-call

Timeouts are the safety net against runaway jobs. Every long-running task — training, large queries, sensors — should have a timeout shorter than the next scheduled run, otherwise a stuck task will silently delay every downstream pipeline.

Sub-topic 4.3: Backfills and Catch-Up

A backfill runs a pipeline for a historical range, usually because you fixed a bug, ingested missing data, or rebuilt features under a new schema. The mechanics vary widely:

[Source: https://www.union.ai/blog-post/we-compared-the-data-models-of-every-major-ai-orchestrator]

Two operational rules: (1) Limit concurrency during backfills — a year of daily features fired off in parallel will overload the warehouse and cost a fortune. (2) Tag backfill runs distinctly (e.g., a run-name suffix or a label) so monitoring, cost dashboards, and alerts can distinguish them from normal scheduled runs.

Figure 6.5: Backfill execution timeline with bounded concurrency interleaved with scheduled runs.

sequenceDiagram
    participant Op as Operator
    participant Sched as Scheduler
    participant Pool as Backfill Pool (max=3)
    participant Prod as Prod Pool
    Op->>Sched: backfill 2025-01-01..2025-01-07
    Sched->>Pool: enqueue 7 partition runs
    par Bounded fan-out
        Pool->>Pool: run 2025-01-01
        Pool->>Pool: run 2025-01-02
        Pool->>Pool: run 2025-01-03
    end
    Note over Prod: 02:00 UTC scheduled run unaffected
    Sched->>Prod: daily-training (today)
    Pool->>Pool: run 2025-01-04
    Pool->>Pool: run 2025-01-05
    Pool->>Pool: run 2025-01-06
    Pool->>Pool: run 2025-01-07
    Pool-->>Op: backfill complete

Sub-topic 4.4: Parameterization

Parameterize everything that changes between runs: date, env (prod/stage/dev), model_name, segment, and hyperparameter presets. The right interface depends on the tool:

• Airflow: dag_run.conf for ad-hoc JSON payloads; params for typed defaults; Jinja templating ({{ ds }}) for time variables.
• Prefect: function arguments on the @flow are the parameters; deployments define defaults, runs override.
• Dagster: config_schema on jobs/ops; partition keys for time-dimensional params.
• Flyte: Python type-hinted workflow inputs; LaunchPlans provide defaults and schedules.
• KFP: pipeline parameters in the DSL; UI/SDK supplies values per run.
• Metaflow: @parameter decorators; CLI arguments work naturally.

[Source: https://www.prompts.ai/blog/tools-orchestrating-machine-learning-workflows.html]

Good practice: type-hint every parameter, validate it at the start of the flow, and tag runs with their parameter values so you can filter “all daily-training runs with model_type=xgboost in env=prod” in the UI.

Sub-topic 4.5: Resource Management and Queueing

The final operational dimension is sharing finite compute fairly. Two-GPU training jobs and 2-CPU feature builds should not contend for the same pool; backfills should not starve scheduled production runs.

• Work pools / queues: Prefect work pools, Airflow queues, Dagster runQueueDaemon — partition workers by purpose (“gpu”, “default”, “backfill”) and route tasks accordingly.
• Resource specs per task: KFP and Flyte let each step declare CPU/GPU/memory; the K8s scheduler bin-packs them across nodes. Airflow expresses the same via KubernetesPodOperator resources [Source: https://eng.lyft.com/orchestrating-data-pipelines-at-lyft-comparing-flyte-and-airflow-72c40d143aad].
• Concurrency limits: per-DAG, per-task, and per-pool. A daily-training flow with max_concurrency=1 cannot start a new run while yesterday’s is still going — a simple guard against pile-ups.
• Priorities: most orchestrators let you give production runs higher priority than ad-hoc experiments; use it.

Key Takeaway: Reliable pipelines are built from four disciplines: idempotent tasks, declarative retries with jittered backoff and timeouts, explicit parameterization tagged on runs, and resource isolation via pools and concurrency limits. Get these right and the choice of orchestrator becomes a matter of taste; get them wrong and no orchestrator will save you.

Chapter Summary

Pipeline orchestration is the connective tissue that turns isolated ML steps into a reliable production system. We started from the four primitives every orchestrator shares — DAGs, tasks, operators, and executors — and noted that the deeper paradigm split is between task-centric orchestration (Airflow, Prefect, Argo) and asset/artifact-centric orchestration (Dagster, KFP, Flyte, Metaflow, ZenML). The former asks “did the job run?”; the latter asks “is this dataset, feature table, or model fresh?” — a question much closer to what ML practitioners actually need to answer.

Among general-purpose tools, Airflow remains the safe default for organizations already standardized on it, Prefect offers the most Pythonic ergonomics, Dagster is uniquely strong for data-aware lineage, and Argo Workflows sits beneath everything Kubernetes-native. Among ML-native tools, Kubeflow Pipelines and Flyte dominate the K8s-first, GPU-heavy ML platform space; Metaflow wins on ML-developer productivity; ZenML acts as a portable meta-orchestrator; and the cloud vendors offer managed escape hatches via Vertex AI Pipelines and SageMaker Pipelines.

The operational layer — retries with jittered exponential backoff, idempotent partitioned writes, explicit backfill mechanisms, type-hinted parameterization, and resource isolation through work pools — is largely framework-agnostic. Master those disciplines and you can change orchestrators in a quarter; ignore them and no tool will rescue you. The worked Prefect DAG in Section 4 illustrates the canonical shape: parameterize the date, build features idempotently with retries and caching, validate data quality, branch conditionally, fan out training across segments dynamically, and register only the winners. Every framework in this chapter expresses some version of that pattern; their differences are mostly about how much they help with versioning, lineage, and resource control along the way.

Key Terms

Chapter 7: Model Training Infrastructure and Distributed Training

Training a modern deep learning model is, in many ways, less about clever algorithms and more about choreographing thousands of arithmetic units, gigabytes of memory, and miles of high-speed cabling. The model is the recipe; the infrastructure is the kitchen. If the kitchen is small, badly lit, or has only one oven, even the best recipe will take days to bake. If the kitchen is well designed, with parallel ovens and assistants who pass ingredients efficiently, you can produce a Michelin-star transformer in a fraction of the time.

This chapter walks through how that kitchen is built. We start with the hardware: GPUs, TPUs, CPUs, and the cluster managers that schedule them. Then we look at how a single model is split across many devices using data, model, and pipeline parallelism. Next, we survey the frameworks (PyTorch DDP and FSDP, Horovod, DeepSpeed, Ray Train) that turn those strategies into a few dozen lines of code. Finally, we discuss how to keep the electricity bill manageable using mixed precision, gradient checkpointing, spot instances, and elastic training.

By the end, you should be able to look at a training job description (“we need to fine-tune a 70B-parameter LLM on 1 trillion tokens”) and sketch a plausible infrastructure plan, including hardware choices, parallelism strategy, framework, and cost controls.

Section 1: The Training Compute Landscape

Before we can distribute work, we need to understand what kinds of compute units exist, how their memory is organized, and how they are aggregated into clusters. The choices here set the ceiling on everything else: communication bandwidth, peak FLOPs, and the price tag of every experiment.

GPUs and the Memory Hierarchy

Graphics processing units (GPUs) dominate deep learning training because they trade a few latency-optimized CPU cores for thousands of throughput-optimized arithmetic units. NVIDIA’s data center GPUs - the V100, A100, and H100 generations - are the workhorses of most production training stacks. The A100 (40 or 80 GB of HBM2e memory) introduced the BF16 numerical format with hardware acceleration, while the H100 added FP8 support and roughly 3x the matrix-math throughput per watt for transformer workloads [Source: https://arxiv.org/html/2407.02883v3].

GPU memory has a steep hierarchy that influences every design decision. At the bottom are registers and shared memory inside each streaming multiprocessor (kilobytes, single-digit nanosecond latency). Above that sits L2 cache (tens of megabytes), then HBM - high-bandwidth memory - which is the “RAM” of the GPU (40-80 GB on data center cards, with bandwidth in the terabytes per second). Beyond the device lives host CPU memory (often hundreds of gigabytes, but reached over the PCIe bus with much lower bandwidth and higher latency) and finally NVMe storage. A useful analogy is a chef’s workstation: registers are the knife in your hand, HBM is the cutting board, host RAM is the pantry across the room, and NVMe is the warehouse downtown.

This hierarchy matters because every byte that has to move from HBM to registers (or, worse, from host to device) costs time. Activations, parameters, and gradients all compete for HBM, and when they spill out, you either OOM or pay a heavy bandwidth tax to offload them.

Inside a node, GPUs talk to each other over NVLink and NVSwitch - proprietary high-speed interconnects providing hundreds of gigabytes per second of bidirectional bandwidth between cards. Across nodes, clusters use InfiniBand or high-speed Ethernet (100-400 Gbps) [Source: https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide]. The dramatic difference between intra-node and inter-node bandwidth is the single most important fact about distributed training topology, and we will return to it repeatedly.

TPUs and Other Accelerators

Google’s Tensor Processing Units (TPUs) are an alternative class of accelerator built specifically for dense linear algebra. Where a GPU is a general-purpose throughput machine that happens to be good at matmuls, a TPU is essentially a giant systolic array - a 2D grid of multiply-accumulate units - wrapped in surrounding control logic. TPUs use BF16 natively and excel at large batched transformer workloads, particularly when paired with Google’s XLA compiler and JAX. The trade-off is ecosystem: PyTorch on TPU works, but most of the cutting-edge training recipes assume CUDA and NVIDIA hardware.

Other accelerators include AWS Trainium and Inferentia, Cerebras’s wafer-scale chips, Graphcore IPUs, and Habana Gaudi. They have niches - sometimes price/performance, sometimes specialized memory patterns - but as of writing, NVIDIA GPUs and Google TPUs account for the overwhelming majority of large-scale training.

When CPUs Are Still the Right Answer

It is tempting to assume every training job needs a GPU. They do not. Tabular gradient-boosted models (XGBoost, LightGBM), classical scikit-learn pipelines, small NLP fine-tunes, and lightweight recommendation models often train faster and certainly cheaper on a beefy multi-socket CPU node. CPUs also dominate feature engineering, hyperparameter search orchestration, and lightweight inference. The rule of thumb is simple: if your model is small enough that data movement dominates compute, a CPU is fine; if it is dense linear algebra over millions or billions of parameters, you need an accelerator.

Cluster Managers: Kubernetes, Slurm, and Ray

A bare GPU is useless without a scheduler. Three families dominate ML training clusters:

• Kubernetes, often with the Kubeflow Training Operator, Volcano, or KubeRay. K8s shines for hybrid workloads - training, inference, and microservices all sharing infrastructure - and integrates naturally with cloud-native storage and networking.
• Slurm, the workload manager that has run HPC clusters for two decades. Slurm is excellent at gang scheduling (allocating N GPUs across M nodes atomically), MPI workloads, and traditional supercomputing patterns. Most academic supercomputers and many industrial pretraining clusters use Slurm.
• Ray, an open-source distributed framework that has expanded into a full ML platform via Ray Train, Ray Tune, and Ray Serve. Ray is appealing because it presents a single Python API for distributed compute, hyperparameter search, and serving.

In practice, larger organizations often layer these: Slurm or Kubernetes at the bottom managing raw GPU allocations, with Ray or a Kubeflow PyTorchJob on top to launch the actual training script.

Key Takeaway: The training kitchen is built from accelerators (GPUs, TPUs, CPUs) wired together through a steep memory and bandwidth hierarchy, with cluster managers like K8s, Slurm, and Ray handling allocation; understanding NVLink-vs-network bandwidth is the foundation for every distributed training decision.

Section 2: Distributed Training Strategies

Once you have hardware, you need a strategy for splitting the model across it. There are three fundamental axes - data, model/tensor, and pipeline parallelism - that can be combined into the “3D parallelism” used to train modern large language models.

Data Parallelism: Replicate the Model, Split the Batch

Data parallelism is the simplest and most common strategy. Every GPU holds a complete copy of the model, the global batch is divided into per-GPU mini-batches, each GPU runs forward and backward on its share, and then gradients are all-reduced (averaged) across all GPUs so every replica updates identically [Source: https://arxiv.org/html/2407.02883v3].

The analogy is a study group: every student has the same textbook, each works on a different set of practice problems, and at the end they pool their answers to agree on the corrections. The communication overhead is one big synchronization per step, and efficient implementations overlap that synchronization with the backward pass so it is mostly hidden.

Figure 7.1: Data parallelism - replicated model, sharded batch, all-reduced gradients

flowchart TD
    Batch[Global Mini-Batch] --> Split{Split Across N GPUs}
    Split --> S1[Shard 1]
    Split --> S2[Shard 2]
    Split --> S3[Shard 3]
    Split --> S4[Shard 4]
    S1 --> G1[GPU 1: Full Model Replica]
    S2 --> G2[GPU 2: Full Model Replica]
    S3 --> G3[GPU 3: Full Model Replica]
    S4 --> G4[GPU 4: Full Model Replica]
    G1 --> AR[NCCL All-Reduce: Average Gradients]
    G2 --> AR
    G3 --> AR
    G4 --> AR
    AR --> U1[Identical Optimizer Step on Every Replica]

Data parallelism works beautifully when the full model (parameters + optimizer state + activations for a reasonable batch) fits on a single GPU. For Adam optimizers, the optimizer state alone is roughly 2-3x the parameter size, which sneaks up on practitioners trying to scale moderate models. Within those limits, DP scales nearly linearly to dozens of GPUs, particularly when interconnect is fast.

Model and Tensor Parallelism: Split the Model Itself

When the model is too large for one GPU - even with all the memory tricks we will discuss later - you have to physically partition it. Two flavors:

• Layer-wise model parallelism assigns different layers to different GPUs. In a 24-layer transformer with two GPUs, layers 0-11 might live on GPU 0 and layers 12-23 on GPU 1. Activations flow forward across the boundary, gradients flow back.
• Tensor parallelism splits the weight matrices inside a layer. For a giant linear layer y = xW where W is too big for one GPU, you partition W column-wise into [W_1, W_2, …, W_n], compute partial outputs on each GPU, and combine them via an all-gather or all-reduce [Source: https://arxiv.org/html/2407.02883v3]. Transformer implementations like Megatron-LM split attention heads and MLP intermediate dimensions across GPUs this way.

Tensor parallelism dramatically reduces per-GPU memory for huge layers but introduces fine-grained communication: every layer’s forward and backward triggers a collective. This is fine when GPUs sit on the same NVLink island, but disastrous across slow inter-node links. The standard rule is: keep tensor parallelism within a node.

Pipeline Parallelism: Assembly Line for Layers

Pipeline parallelism divides the model into sequential stages and pushes microbatches through them like an assembly line. With four stages and eight microbatches, while stage 1 processes microbatch 2, stage 2 can process microbatch 1; the goal is to keep every stage busy simultaneously.

The challenge is pipeline bubbles - idle time at the start (filling the pipeline) and the end (draining it). The standard remedy is to use many more microbatches M than stages S (M >> S) and to use a smarter schedule like 1F1B (one-forward-one-backward) instead of the simpler GPipe pattern. The trade-off: more microbatches means more in-flight activation memory, partially offset by activation checkpointing.

ZeRO Sharding and FSDP

Between “pure data parallel” and “split the model” lives a third option: shard the model state across data-parallel ranks. Microsoft’s Zero Redundancy Optimizer (ZeRO) was the first to systematize this. ZeRO has three stages:

• ZeRO-1 shards the optimizer states (the Adam moments) across ranks.
• ZeRO-2 also shards the gradients.
• ZeRO-3 shards parameters too. No GPU ever holds the entire model at rest.

PyTorch’s Fully Sharded Data Parallel (FSDP) is the native implementation of ZeRO-3-style sharding. The communication pattern is elegant: before each layer’s forward pass, FSDP issues an all-gather to reassemble that layer’s parameters across ranks; after backward, it issues a reduce-scatter so each rank only stores its shard of the gradient. The optimizer step then updates only the local shard.

Figure 7.2: ZeRO sharding stages - progressive partitioning of training state

graph TD
    Base[Baseline DDP: Every Rank Holds Full State] --> Z1
    subgraph Z1[ZeRO-1]
        Z1A[Params: Replicated]
        Z1B[Gradients: Replicated]
        Z1C[Optimizer States: SHARDED]
    end
    Z1 --> Z2
    subgraph Z2[ZeRO-2]
        Z2A[Params: Replicated]
        Z2B[Gradients: SHARDED]
        Z2C[Optimizer States: SHARDED]
    end
    Z2 --> Z3
    subgraph Z3[ZeRO-3 / FSDP]
        Z3A[Params: SHARDED]
        Z3B[Gradients: SHARDED]
        Z3C[Optimizer States: SHARDED]
    end
    Z3 --> Off[+ CPU/NVMe Offload: Push State Beyond GPU RAM]

NCCL All-Reduce and Ring All-Reduce

The communication backbone for nearly every GPU collective on NVIDIA hardware is NCCL (NVIDIA Collective Communications Library). NCCL implements all-reduce, all-gather, reduce-scatter, and broadcast primitives, choosing internally between ring, tree, and hierarchical algorithms depending on topology.

Ring all-reduce - popularized by Horovod - arranges N GPUs in a logical ring. The tensor is split into N chunks; in a scatter-reduce phase, chunks circulate around the ring with each GPU adding its local contribution, and in an all-gather phase, the reduced chunks circulate again so every GPU ends up with the full result. Each GPU transmits roughly 2(N-1)/N times the tensor size, which is bandwidth-optimal but has latency growing linearly in N. For small models on many ranks, this latency becomes painful, which is one reason large models lean on the all-gather/reduce-scatter pattern of FSDP rather than naive all-reduce [Source: https://arxiv.org/html/2407.02883v3].

Figure 7.3: NCCL ring all-reduce - bandwidth-optimal gradient aggregation around a logical ring

flowchart LR
    G0[GPU 0<br/>chunk A] -->|send chunk| G1[GPU 1<br/>chunk B]
    G1 -->|send chunk| G2[GPU 2<br/>chunk C]
    G2 -->|send chunk| G3[GPU 3<br/>chunk D]
    G3 -->|send chunk| G0
    G0 -.scatter-reduce: each rank accumulates one chunk.- G1
    G1 -.all-gather: reduced chunks circulate again.- G2

Comparing Parallelism Strategies

For models with tens or hundreds of billions of parameters, none of these is sufficient alone. The industry standard is 3D parallelism: tensor parallel within a node to shard huge layers, pipeline parallel across nodes to spread depth, and data parallel across replicas for throughput [Source: https://arxiv.org/html/2407.02883v3].

Key Takeaway: Data parallelism scales batch size, tensor parallelism scales layer width, pipeline parallelism scales depth, and ZeRO/FSDP sharding scales total state - real LLM training combines all of them in a 3D mesh tuned to the cluster’s bandwidth topology.

Section 3: Frameworks and Tools for Distributed Training

Strategies are conceptual; frameworks are what you actually import. Four (really four-and-a-half) dominate the open-source landscape: PyTorch DDP, PyTorch FSDP, Horovod, and DeepSpeed, with Ray Train as an orchestration layer that wraps them.

PyTorch DistributedDataParallel (DDP)

DDP is the workhorse. It is the canonical PyTorch implementation of data parallelism: each rank holds a full copy of parameters, gradients, and optimizer state, and gradients are bucketed and all-reduced via NCCL during the backward pass [Source: https://docs.pytorch.org/docs/stable/elastic/run.html]. The typical setup is small enough to read in one breath:

import torch, os
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

model = MyModel().cuda(local_rank)
model = DDP(model, device_ids=[local_rank])
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

DDP is launched with torchrun (formerly torch.distributed.launch), which spawns one process per GPU and orchestrates the rendezvous between nodes. For most models that fit comfortably on a single GPU and scale up to 8-16 GPUs, DDP is the right tool. It is simple to reason about, has excellent NCCL integration, and surfaces failures clearly.

PyTorch Fully Sharded Data Parallel (FSDP)

FSDP is what you switch to when the model no longer fits. It implements ZeRO-3-style full sharding of parameters, gradients, and optimizer states. The communication pattern - all-gather before each layer, reduce-scatter after - means more total NCCL calls than DDP, but each rank’s memory footprint drops by roughly 1/N for large enough models.

The tricky part is the auto-wrap policy: FSDP must decide which submodules to wrap as independently-sharded units. Wrapping every transformer block is the standard recipe; wrapping too coarsely defeats the memory savings, while wrapping too finely produces excessive communication. Mixed precision, backward prefetching, and CPU offload are configurable per-wrap. FSDP integrates with torch.distributed.checkpoint for sharded checkpoint I/O, which becomes essential when checkpoints span dozens of nodes.

Horovod

Horovod is a framework-agnostic ring-allreduce library originally from Uber. It works with PyTorch, TensorFlow, and MXNet, which made it the default in heterogeneous shops circa 2018-2020 [Source: https://arxiv.org/html/2407.02883v3]. Each rank holds a full model copy and optimizer state, and gradients are aggregated via ring-allreduce (NCCL or MPI as backend). Setup looks like:

import horovod.torch as hvd
hvd.init()
torch.cuda.set_device(hvd.local_rank())

optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
hvd.broadcast_parameters(model.state_dict(), root_rank=0)

Horovod still ships and works, but in PyTorch-only environments DDP has eclipsed it for the same use cases, and FSDP/DeepSpeed have eclipsed it for memory-bound workloads. The remaining sweet spot is multi-framework clusters or MPI-heavy environments where Horovod’s horovodrun launcher integrates cleanly with existing tooling.

DeepSpeed and Megatron-LM

Microsoft’s DeepSpeed is the “batteries-included” stack for very large training. At its core is the ZeRO optimizer (stages 1, 2, 3), but DeepSpeed adds:

• CPU and NVMe offload for parameters, gradients, and optimizer states - letting you push models beyond total GPU RAM into host memory or fast SSDs.
• Activation checkpointing integrated with the optimizer.
• Pipeline parallelism with multiple scheduling strategies.
• Tensor parallelism via integration with NVIDIA’s Megatron-LM (the combined stack is often called Megatron-DeepSpeed) [Source: https://arxiv.org/html/2407.02883v3].

The trade-off is configuration overhead. DeepSpeed jobs are driven by a JSON config file with dozens of knobs, and migrating between ZeRO stages or enabling offload requires careful tuning of bucket sizes, communication overlap, and CPU-side optimizer kernels. When you genuinely need to train a 70B+ parameter model with limited GPU RAM, DeepSpeed remains a top choice.

Ray Train

Ray Train is not really a distributed training algorithm; it is an orchestration and abstraction layer. You write a single training function that uses DDP, FSDP, or DeepSpeed internally, then ask Ray to run it on a cluster: Ray handles scaling, fault tolerance, hyperparameter search integration via Ray Tune, and checkpoint management to durable storage. Ray Train is increasingly popular in cloud-native and Kubernetes environments because it abstracts away the launcher (torchrun vs deepspeed vs horovodrun) behind a uniform Python API.

Framework Comparison

A useful migration path: start with DDP, switch to FSDP when you hit OOM, and reach for DeepSpeed when you need CPU/NVMe offload, pipeline parallel, or tightly integrated Megatron-LM tensor parallelism.

Key Takeaway: PyTorch DDP is the default for small-to-medium models, FSDP brings ZeRO-3 sharding into pure PyTorch for LLM-scale jobs, DeepSpeed adds offload and 3D parallelism for the largest models, Horovod serves multi-framework clusters, and Ray Train provides uniform orchestration over all of them.

Section 4: Cost and Throughput Optimization

Distributed training is expensive enough that small percentage savings translate to real money. Four levers - mixed precision, gradient checkpointing/accumulation, spot instances with elastic launchers, and profiling - account for most of the wins.

Mixed Precision and BF16

The fastest way to halve your training bill is often to switch from FP32 to mixed precision. Modern GPUs have tensor cores that perform matmuls in 16-bit precision at 4-8x the throughput of FP32. Two flavors matter:

• FP16 (half precision) has only 5 exponent bits, so gradients can underflow. Automatic Mixed Precision (AMP) with a GradScaler multiplies the loss by a large factor before backward so gradients stay in range, then unscales before the optimizer step.
• BF16 (brain float) uses 8 exponent bits - the same range as FP32 - at the cost of fewer mantissa bits. BF16 typically does not need dynamic loss scaling, making it simpler. A100 and H100 GPUs have first-class BF16 support [Source: https://docs.pytorch.org/docs/stable/elastic/run.html].

A canonical AMP training loop:

scaler = torch.cuda.amp.GradScaler(enabled=use_fp16)
for inputs, targets in loader:
    with torch.cuda.amp.autocast(dtype=torch.bfloat16 if use_bf16 else torch.float16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

Mixed precision typically yields 1.5-2x speedup over FP32 on A100/H100 when the workload is compute-bound, and reduces activation memory enough to fit larger batches. The H100 takes this further with FP8 training, where supported, for additional throughput on transformer workloads.

Figure 7.4: Mixed-precision training flow with loss scaling

flowchart TD
    M[FP32 Master Weights] --> Cast[Cast to FP16/BF16]
    Cast --> F[Forward Pass in FP16/BF16<br/>Tensor Cores]
    F --> L[Loss in FP32]
    L --> Scale{FP16?}
    Scale -->|Yes| LS[Multiply Loss by Scale Factor]
    Scale -->|No, BF16| B[Backward Pass]
    LS --> B
    B --> G[FP16/BF16 Gradients]
    G --> Unscale{FP16?}
    Unscale -->|Yes| US[Unscale Gradients<br/>Check for Inf/NaN]
    Unscale -->|No| Opt[Optimizer Step on FP32 Master]
    US --> Opt
    Opt --> M

Gradient Checkpointing and Accumulation

Two more memory-throughput trades:

Gradient checkpointing (also called activation recomputation) drops most intermediate activations during the forward pass and recomputes them during backward. This typically reduces activation memory by 30-50% at the cost of 20-40% extra compute. For memory-bound jobs - say, fitting a large transformer on T4 GPUs - this is the difference between training and OOM.

Gradient accumulation simulates a larger global batch by running accum_steps micro-batches before each optimizer update:

accum_steps = 4
optimizer.zero_grad()
for step, (x, y) in enumerate(loader):
    with autocast():
        loss = model_step(x, y) / accum_steps
    scaler.scale(loss).backward()
    if (step + 1) % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

This is particularly valuable when spot capacity forces you to drop from 8 nodes to 4: accumulation lets you maintain the same effective global batch (and thus convergence behavior) without raising per-GPU memory.

Spot, Preemptible, and Elastic Training

Cloud spot or preemptible instances offer 50-80% discounts over on-demand pricing in exchange for the right to reclaim them with little warning [Source: https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide]. Used carelessly, spot is a recipe for losing days of training to a single preemption. Used well, it is the single largest cost lever available.

The key is elastic, fault-tolerant launch via torchrun:

torchrun \
  --nnodes=2:8 \
  --nproc-per-node=8 \
  --rdzv-backend=c10d \
  --rdzv-endpoint=$HOST:29400 \
  --max-restarts=5 \
  train.py ...

The --nnodes=2:8 flag means the job can run with anywhere from 2 to 8 nodes; if nodes are preempted, torchrun re-rendezvouses and resumes with the remaining members [Source: https://docs.pytorch.org/docs/stable/elastic/run.html]. Combined with --max-restarts, it treats spot preemption as a routine membership change rather than a fatal error.

For this to work, the training script must be stateless across restarts:

• A single main() entry point (no mp.spawn), since torchrun spawns the processes itself.
• On startup, attempt to load the latest snapshot from durable storage (S3, GCS, or networked NFS) before initializing the process group [Source: https://pytorch-cn.com/tutorials/beginner/ddp_series_fault_tolerance.html].
• Snapshots must include model state, optimizer state, AMP GradScaler state, LR scheduler, global step/epoch, and RNG seeds.
• Use DistributedSampler and call sampler.set_epoch(epoch) so resharding is deterministic.

Checkpoint frequency depends on cost and risk tolerance: every 5-15 minutes is typical, more often for very expensive steps or volatile spot pools.

Figure 7.5: Spot preemption and elastic checkpoint recovery lifecycle

sequenceDiagram
    participant Cloud as Cloud Spot Pool
    participant TR as torchrun (elastic)
    participant Job as Training Workers
    participant S3 as Durable Storage (S3/GCS)
    TR->>Job: Rendezvous N nodes, start training
    Job->>S3: Snapshot every 5-15 min (model+opt+scaler+RNG)
    Cloud-->>Job: 2-minute preemption notice on Node K
    Job->>S3: Emergency snapshot (termination handler)
    Cloud-->>Job: Node K reclaimed
    TR->>TR: Detect membership change, re-rendezvous
    TR->>Job: Resume with N-1 nodes (within --nnodes=min:max)
    Job->>S3: Load latest snapshot
    Job->>Job: Resume from saved global step + RNG state
    Cloud-->>TR: New spot capacity available
    TR->>Job: Re-rendezvous, scale back up to N

On AWS, Spot Fleets or EC2 Auto Scaling groups with capacity-optimized allocation reduce interruption rate, and the 2-minute termination notice can trigger an emergency checkpoint via a sidecar or systemd hook. On GCP, preemptible/Spot VMs with managed instance groups play the same role.

Profiling: Find Your Bottleneck

Optimization without measurement is wishful thinking. Three tools dominate:

• PyTorch Profiler (torch.profiler) gives operator-level timings, memory snapshots, and a Chrome trace viewer. It is the first stop for finding slow ops, unfused kernels, or pathological Python overhead.
• NVIDIA Nsight Systems profiles at the CUDA stream level: NCCL collectives, kernel overlap, GPU utilization, and host-device synchronization. This is what you reach for when you suspect communication is dominating step time.
• DCGM and dcgm-exporter plus Grafana dashboards expose cluster-wide GPU utilization, memory, and SM occupancy in real time.

The numbers to chase are: GPU SM utilization > 70%, NCCL time < 30% of step time, per-step time stable across ranks (a slow rank stalls the all-reduce). When you see ranks finishing at very different times, suspect data loader stragglers or unbalanced FSDP wrap policies.

GPU Selection by Workload

The dominant cost metric is not $/hour but cost per training token or sample. An H100 at 4x the hourly rate of a V100 may still be cheaper per token if the workload fully utilizes it with BF16 or FP8.

Putting It Together

A minimal cost-efficient training recipe:

1. Launch via torchrun with elastic mode (--nnodes=min:max --max-restarts=N).
2. Save robust snapshots (model + optimizer + AMP scaler + scheduler + step + RNG) to durable storage every 5-15 minutes.
3. Enable BF16 AMP on A100/H100, FP16 elsewhere.
4. Add gradient checkpointing when memory-bound; use gradient accumulation to maintain global batch under elastic shrinkage.
5. Use spot instances with capacity-optimized allocation; install a termination-notice handler that forces immediate snapshots.
6. Pick GPUs based on cost per token measured on a short pilot, not list price.

Key Takeaway: Mixed precision, gradient checkpointing, and gradient accumulation maximize FLOPs per GPU-hour, while elastic torchrun launches with robust snapshotting let you run on cheap spot capacity - the combination routinely cuts training cost by 60-80% without sacrificing convergence.

Chapter Summary

Modern ML training is an exercise in coordinated infrastructure. The hardware foundation - NVIDIA A100/H100 GPUs (or TPUs, or CPUs for lighter workloads), wired together with NVLink within nodes and InfiniBand or fast Ethernet across them, scheduled by Kubernetes, Slurm, or Ray - sets the ceiling on what is possible. The steep memory and bandwidth hierarchy of GPUs (registers, L2, HBM, host RAM, NVMe) drives every distribution decision.

Three parallelism strategies map work onto that hardware. Data parallelism replicates the model and splits the batch, requiring only that the model fits on one GPU; it scales well to dozens of GPUs and is the simplest to implement. Tensor (model) parallelism shards weight matrices within layers across GPUs, enabling very wide layers but introducing fine-grained intra-node communication. Pipeline parallelism splits the model into sequential stages and feeds microbatches through them, scaling depth at the cost of pipeline bubbles. ZeRO-style sharding and PyTorch FSDP shard parameters, gradients, and optimizer states across data-parallel ranks, reducing per-GPU memory dramatically. Frontier LLM training combines all three into 3D parallelism.

Frameworks turn these strategies into code. PyTorch DDP is the default for small-to-medium models; FSDP brings ZeRO-3 sharding to pure PyTorch for LLM-scale jobs; DeepSpeed adds CPU/NVMe offload, pipeline parallelism, and Megatron-LM tensor parallelism for the largest models; Horovod still serves multi-framework or MPI-centric clusters; and Ray Train wraps everything in a cloud-native orchestration layer. NCCL handles the actual GPU-to-GPU collectives, with all-reduce, all-gather, and reduce-scatter as the building blocks.

Cost optimization rests on four levers: mixed precision (BF16 on A100/H100, FP16 elsewhere) delivers 1.5-2x speedups; gradient checkpointing and accumulation trade compute for memory and let you maintain global batch under elastic shrinkage; spot instances with torchrun elastic launch and robust snapshotting cut compute cost by 50-80% with manageable restart overhead; and profiling with PyTorch Profiler, Nsight Systems, and DCGM identifies whether you are compute-bound, memory-bound, or communication-bound so optimizations target the actual bottleneck. The metric that matters is cost per training token or sample, not list price per GPU-hour.

Key Terms

Chapter 8: Experiment Tracking and Hyperparameter Tuning

Building a model is essentially an exercise in disciplined optimism: you try something, measure what happens, and decide whether to keep it. The catch is that “something” rarely means a single change. A single training run encompasses a dataset version, a preprocessing pipeline, a feature set, an architecture, a learning-rate schedule, regularization choices, a random seed, a library version, and the specific revision of the code that wove them together. Lose track of any one of those, and the model’s results stop being scientific findings and start being folklore.

This chapter shows how to turn ad hoc model development into a tracked, comparable, and reproducible workflow. We will look at why experiment tracking matters before pipelines exist, survey the four most influential tracking platforms (MLflow, Weights & Biases, Neptune, Comet), study modern hyperparameter optimization (HPO) algorithms in depth (Bayesian, ASHA, PBT), tour the distributed HPO frameworks that run them at scale (Optuna, Ray Tune, Katib, Vizier), and finally codify the bridge from notebook experimentation to production pipelines.

Why Track Experiments?

Experiment tracking is the discipline of recording, for every model training run, the inputs (data version, code commit, hyperparameters, environment), the outputs (metrics, artifacts, predictions, plots), and the context (who, when, why) so that any past result can be understood, compared, and reproduced. Without it, the modeling workflow degrades into a fog of half-remembered notebook cells.

The Lost Notebook Problem

Anyone who has worked in a Jupyter-driven ML team recognizes the pattern: a data scientist runs ten variants of a model in a notebook, saves a model_final_v3_actually_final.pkl somewhere on a workstation, posts a screenshot of the validation AUC to Slack, and moves on. Three months later, someone asks: “Which preprocessing did that model use? Which features? What random seed? Was that the run with min_samples_leaf=5 or min_samples_leaf=15?” Nobody knows. The notebook has been edited a hundred times since then, the workstation has been reimaged, and the AUC screenshot does not record the data partition that produced it.

The cost of this is not just embarrassment. It is concretely (1) wasted compute, because the team must re-run search to recover a known-good result; (2) silent regressions, because a “reproduction” attempt is actually a new experiment with subtly different settings; and (3) blocked collaboration, because nobody else can build on a result they cannot inspect. Modern tracking platforms exist primarily to make this entire failure mode impossible by writing every run to a durable, queryable backend at the moment the run happens [Source: https://mlflow.org/docs/latest/python_api/mlflow.genai.html].

Reproducibility and Comparison

A tracked run is a row in a database that ties together a git commit, a configuration object, a dataset hash, a Python environment specification, a set of metrics over time, and a folder of artifacts. When two runs differ on one metric, that database lets you compute the symmetric difference of everything else to find the cause. Without that, “why did accuracy drop?” is an exploratory archaeology project; with it, it is a SQL query. Comparison views, parallel coordinate plots, and metric-vs-step charts all rest on this same machinery.

Figure 8.1: Experiment tracking flow from code to registry

flowchart LR
    A[Code Commit] --> B[Training Run]
    C[Data Version] --> B
    D[Hyperparameters] --> B
    B --> E[Log Params]
    B --> F[Log Metrics]
    B --> G[Log Artifacts]
    E --> H[(Tracking Backend)]
    F --> H
    G --> H
    H --> I[Compare and Select]
    I --> J[Model Registry]

Reproducibility means more than re-running the same code. It means that the same code, on the same data, in the same environment, with the same seeds, produces the same numbers. Tracking systems contribute to this in two ways: they record the inputs precisely enough that you could re-create them, and they store the model artifact and its environment together so that even years later you can re-instantiate the exact training context [Source: https://mlflow.org/docs/latest/python_api/mlflow.metrics.html].

Auditability and Compliance

In regulated domains (finance, healthcare, government, hiring), models do not get to exist as folklore. Auditors need to answer: “Which version of which model produced the score that denied this loan? What data trained it? Who approved its promotion? Were the validation metrics within thresholds at promotion time?” These questions are almost trivially answerable when each run is timestamped, signed by a user, linked to a code commit, and tied to a registered model version. They are nearly impossible to answer otherwise.

The same audit trail also serves internal governance: model risk teams, security reviewers, and ML engineers all benefit from knowing exactly what changed between versions. MLflow’s design explicitly separates the act of logging a model from registering it, with stage transitions (Staging, Production, Archived) that constitute an auditable promotion workflow [Source: https://home.mlops.community/public/videos/mlflow-leading-open-source].

Foundation for the Model Registry

Experiment tracking and model registries are two halves of one system. A tracker records the noisy reality of dozens of runs per day, including failures, sweeps, and ablations. A registry records the small subset of those runs that the team decided to standardize on, with explicit versions and lifecycle stages. The bridge between them is a record of provenance: when version 7 of the fraud_classifier model is promoted to Production, the registry should be able to point back to the exact run, code, data, and metrics that produced it.

This bidirectional link is what turns “we have a model in production” into “we have an accountable, debuggable, replaceable model in production.” Chapter 9 will dive into registries in detail; for now it is enough to note that without disciplined tracking, registries become decorative.

Key Takeaway: Experiment tracking exists to make every model result auditable, reproducible, and comparable by writing the full context (code, data, hyperparameters, environment, metrics, artifacts) of every run to a durable backend. It is the prerequisite for both safe iteration and any meaningful model registry.

Experiment Tracking Tools

Four platforms dominate production experiment tracking: MLflow, Weights & Biases (W&B), Neptune, and Comet. All four cover the core capabilities (log parameters, metrics, artifacts, code, and environment) but they make different trade-offs along the axes that actually matter when a team commits to one: self-hosting vs SaaS, model registry maturity, collaboration features, visualization quality, and scalability [Source: https://home.mlops.community/public/videos/mlflow-leading-open-source].

MLflow Tracking

MLflow is the de facto open-source standard for ML lifecycle management. Its design covers four pillars: Tracking (runs, params, metrics, artifacts), Projects (reproducible packaging), Models (a generic model format), and the Model Registry (versions, stages, lineage). For tracking, you start a run with mlflow.start_run(), then call mlflow.log_param, mlflow.log_metric, and mlflow.log_artifact to record everything about that run; autologging hooks into frameworks like scikit-learn, PyTorch, and XGBoost to record metrics and parameters automatically [Source: https://mlflow.org/docs/latest/python_api/mlflow.genai.html].

Architecturally, MLflow Tracking is a server (FastAPI + a UI) backed by a relational database for metadata (Postgres, MySQL, SQLite) and a pluggable object store for artifacts (S3, GCS, Azure Blob, NFS). Because every piece is open and replaceable, regulated organizations can run it entirely inside their VPC, integrate it with internal auth (LDAP, OIDC), and treat it as the system of record for models. The trade-off is operational: you run the database, the storage, the server, and the upgrades.

Figure 8.2: MLflow tracking server architecture

flowchart TD
    A[Training Client<br/>mlflow.log_*] --> B[MLflow Tracking Server<br/>FastAPI + UI]
    C[Notebook Client] --> B
    D[Pipeline Client] --> B
    B --> E[(Metadata DB<br/>Postgres / MySQL / SQLite)]
    B --> F[(Artifact Store<br/>S3 / GCS / Azure Blob / NFS)]
    E --> G[Run Metadata<br/>params, metrics, tags]
    F --> H[Models, Plots,<br/>Datasets, Logs]
    B --> I[Model Registry<br/>Versions and Stages]

MLflow’s UI is functional but spartan compared to W&B’s. You get run lists, hyperparameter comparison, metric charts, and the registry view, but no built-in reports, comments, or dashboards. Many teams pair MLflow with notebook-driven analyses or BI tools for richer reporting.

Weights & Biases (W&B)

W&B is SaaS-first and visualization-strongest. A W&B run captures the same building blocks (params, metrics, artifacts) but the platform’s identity lives in its UI: interactive metric panels, system metrics (GPU/CPU/memory), gradient histograms, image and audio media, custom panels, and shareable Reports that let stakeholders see exactly what a researcher saw. Integrations exist as one-line callbacks for PyTorch Lightning, TensorFlow/Keras, scikit-learn, HuggingFace, XGBoost, RL frameworks, and most things in between.

The W&B Artifacts system handles versioned datasets and models with explicit lineage graphs; combined with their Model Registry, this can serve as both registry and tracker for many teams. The strongest single feature for hyperparameter work is W&B Sweeps, which orchestrates HPO (grid, random, Bayesian) and visualizes the search space as the sweep progresses. The principal trade-offs are commercial: pricing is per-seat above the free tier, self-hosting is available only on enterprise plans, and storing all metadata externally is a non-starter for some compliance contexts.

Neptune

Neptune positions itself as a metadata store for ML, emphasizing structured logging, custom fields, and fast search across very large numbers of runs and projects. If a team’s pain point is “we have 50,000 experiments across 12 teams and we need to query them like a database,” Neptune is built for that: schema-friendly tagging, hierarchical metadata, and a UI tuned for filtering and comparison rather than dashboard storytelling.

Neptune supports both SaaS and self-hosted/VPC deployments, integrates with the usual frameworks (PyTorch, TensorFlow, scikit-learn, Kedro, Airflow), and has a strong story for organizing many projects consistently. Its model registry is present but less central than MLflow’s; teams that use Neptune commonly pair it with MLflow or an internal registry for the production model lifecycle.

Comet

Comet is a balanced choice for mid-size teams that want a hosted, polished tracking experience without committing to all of W&B’s price point or all of MLflow’s operations. It supports the standard logging surface (hyperparameters, metrics, artifacts, code, environment), has a model registry with versions and stage transitions, and offers a useful online/offline mode where runs can be cached locally on a spot instance or in a restricted network and synced later. SaaS and self-hosted/VPC deployment options exist.

Self-Hosted vs SaaS

The single biggest organizational choice in tracking is self-hosted vs SaaS. Self-hosted (MLflow OSS, Neptune VPC, Comet on-prem, W&B Enterprise on-prem) means data and metadata stay inside your perimeter, which is mandatory in regulated environments and often desirable for very sensitive data. The cost is operational ownership: database backups, HA, scaling, SSO integration, upgrades. SaaS (W&B default, Neptune cloud, Comet cloud, MLflow on Databricks/Azure ML) shifts that burden to a vendor and tends to scale more transparently as teams grow, at the cost of data residency and lock-in considerations.

Comparison Table

[Source: https://mlflow.org/docs/latest/python_api/mlflow.genai.html] [Source: https://home.mlops.community/public/videos/mlflow-leading-open-source]

Key Takeaway: Pick MLflow when you need open-source control, strong registry, and infra integration; W&B when collaboration and visualization speed matter most; Neptune when you treat experiments as searchable records across many teams; Comet when you want a polished hosted experience with a private-cloud option. All four cover the basics, so choose on the axes that match your organization, not on a feature checklist.

Hyperparameter Tuning

A model’s hyperparameters are the knobs the optimizer cannot turn for you: learning rate, depth, width, dropout, regularization strength, batch size, kernel choice, tree depth, number of estimators. Their values can swing validation performance by more than the architecture itself, and finding good ones is itself an optimization problem - one usually treated as a black-box search over a high-dimensional, mixed (continuous/discrete/categorical) space where each evaluation is expensive.

Grid and Random Search

The simplest strategies treat hyperparameter space as something to be enumerated. Grid search picks a discrete value set for each hyperparameter and evaluates the Cartesian product. It is deterministic, trivially parallel, and easy to reason about, but it scales exponentially with the number of hyperparameters and wastes trials on dimensions that do not matter. Random search samples configurations from per-hyperparameter distributions; it is also trivially parallel and, given the same budget, usually finds better configurations than grid search because it allocates samples non-redundantly across important dimensions. Random search is the standard baseline that any smarter method should beat [Source: https://natesnewsletter.substack.com/p/context-windows-are-a-lie-the-myth].

Neither method learns: trial number 100 is sampled exactly the same way as trial number 1. When training is cheap, that is fine. When training takes hours and the search space has even moderate dimensionality, it becomes wasteful.

Bayesian Optimization

Bayesian optimization (BO) treats hyperparameter search as the problem of finding the maximum of an unknown function f(lambda) (validation performance as a function of hyperparameters) using as few evaluations as possible. It fits a surrogate model of f (Gaussian Process, random forest as in SMAC, or Tree-structured Parzen Estimator as in HyperOpt/Optuna) to the past evaluations, then uses an acquisition function (Expected Improvement, Upper Confidence Bound, Probability of Improvement) to decide where to sample next, balancing exploration (try uncertain regions) against exploitation (refine promising regions).

The strength of BO is sample efficiency: when each training run is expensive and the search space is small-to-moderate (roughly up to 20-30 dimensions), BO consistently beats random search at the same budget. The weaknesses are equally consistent: surrogate fitting struggles in high dimensions, BO is awkward with categorical/conditional parameters, parallelization beyond ~10-20 workers gives diminishing returns (because the surrogate cannot incorporate results before launching more trials), and classic BO treats each evaluation as a single scalar - it ignores intermediate learning curves [Source: https://platform.claude.com/docs/en/build-with-claude/prompt-engineering/claude-prompting-best-practices].

Figure 8.3: Bayesian optimization loop

flowchart TD
    A[Observed Trials<br/>lambda, performance] --> B[Fit Surrogate Model<br/>GP / TPE / RF]
    B --> C[Evaluate Acquisition Function<br/>EI / UCB / PI]
    C --> D[Select Next Hyperparameter<br/>lambda*]
    D --> E[Train Model<br/>and Evaluate]
    E --> F[Record Performance]
    F --> A
    F --> G{Budget<br/>exhausted?}
    G -->|No| B
    G -->|Yes| H[Return Best Config]

Hyperband and ASHA

Hyperband and ASHA take a different angle: instead of modeling the objective, they allocate compute adaptively by aggressively early-stopping bad runs. The unit of adaptation is a resource that can be incrementally increased (epochs, training time, training-set size, image resolution).

The core algorithm is Successive Halving: launch many configurations cheaply (small resource), evaluate, keep the top fraction (e.g., the best 1/eta), increase the resource for survivors, and repeat until a few configurations have trained to full budget. Hyperband wraps this in multiple “brackets” with different starting (n, r) trade-offs, hedging against not knowing how predictive early performance is. ASHA (Asynchronous Successive Halving) is the practical parallel variant: trials are promoted or stopped at each “rung” as soon as they finish, with no global synchronization. This scales to hundreds or thousands of workers, tolerates heterogeneous runtimes, and handles preemption gracefully.

Figure 8.4: ASHA rung-based early stopping

flowchart TD
    A[Rung 0: 27 trials<br/>1 epoch each] --> B{Top 1/3<br/>by metric}
    B -->|Promote 9| C[Rung 1: 9 trials<br/>3 epochs each]
    B -->|Stop 18| X1[Pruned]
    C --> D{Top 1/3<br/>by metric}
    D -->|Promote 3| E[Rung 2: 3 trials<br/>9 epochs each]
    D -->|Stop 6| X2[Pruned]
    E --> F{Top 1/3<br/>by metric}
    F -->|Promote 1| G[Rung 3: 1 trial<br/>27 epochs - full budget]
    F -->|Stop 2| X3[Pruned]
    G --> H[Best Configuration]

ASHA shines when early performance is reasonably predictive of final performance (typical for most deep learning) and when you have substantial parallel compute. Its weakness is exactly the inverse: models that “learn late” or have non-monotonic curves get cut prematurely, and it is no more sample-efficient than random search at picking which configurations to try - it only decides which to stop.

BOHB (Bayesian Optimization + Hyperband) and similar hybrids combine the two ideas: use a TPE-like surrogate to sample better configurations and Hyperband to early-stop. In practice BOHB is often the strongest default for large DL tuning workloads [Source: https://platform.claude.com/docs/en/build-with-claude/prompt-engineering/claude-prompting-best-practices].

Population-Based Training

Population-Based Training (PBT) is conceptually different: it optimizes hyperparameters and model weights jointly over training time, borrowing from evolutionary algorithms. A population of N models trains in parallel with different hyperparameters. At periodic “exploit/explore” steps, low-performing members exploit by copying weights and hyperparameters from better-performing peers, then explore by perturbing the copied hyperparameters. The output is not a single best hyperparameter vector but a schedule of hyperparameters over training - which is often what you actually want for things like learning rate, dropout, and entropy regularization.

Figure 8.5: PBT exploit/explore cycle

sequenceDiagram
    participant W1 as Worker 1 (low perf)
    participant W2 as Worker 2 (top perf)
    participant Sched as PBT Scheduler
    participant Store as Checkpoint Store

    W1->>Sched: Report metric @ step T
    W2->>Sched: Report metric @ step T
    Sched->>Sched: Rank population
    Sched-->>W1: Exploit: copy from W2
    W2->>Store: Save weights + hparams
    Store-->>W1: Load W2 checkpoint
    Sched-->>W1: Explore: perturb hparams
    W1->>W1: Resume training with new schedule
    W2->>W2: Continue training
    Note over W1,W2: Repeat every K steps

PBT excels in deep reinforcement learning, large supervised models with long training horizons, and any setting where the best hyperparameters change over the course of training. Its costs are real: substantial compute (you train many models for full durations), infrastructure complexity (frequent checkpointing and weight copying between workers), and orchestration overhead. PBT is overkill when models are small or runs are short.

HPO Algorithm Comparison

Tools: Optuna, Ray Tune, Katib, Vizier

Three open-source frameworks dominate distributed HPO in practice, with Google’s Vizier as the influential research/internal precursor.

Optuna is a Python-native library built around a Study (an optimization run) containing Trials (individual evaluations). Its design separates Samplers (which propose configurations: TPE, CMA-ES, random, Gaussian Process) from Pruners (which decide whether to stop a trial early: median, ASHA, percentile). Optuna’s ask-and-tell API lets external systems propose, evaluate, and report trials without Optuna controlling execution - making it composable with Kubernetes jobs, Airflow tasks, or any external scheduler. Distributed coordination happens through a shared storage backend (SQLite, MySQL, Postgres), with a gRPC storage proxy that fronts the database for high-throughput, thousands-of-worker scenarios [Source: https://www.youtube.com/watch?v=tVskbekONlw].

Ray Tune is built on the Ray distributed runtime: each trial is a Ray task or actor, scheduled by Ray’s resource-aware scheduler with explicit CPU/GPU requirements (including fractional GPUs like gpus_per_trial=0.25). Tune integrates ASHAScheduler, PBT, BOHB-style algorithms, and a wide selection of search algorithms (HyperOpt, Optuna, BayesOpt, AxSearch). Because Tune inherits Ray’s primitives, it scales naturally to multi-node GPU clusters, handles fault tolerance through Ray’s actor model, and integrates cleanly with MLflow for tracking, so that Tune drives the search while MLflow records every trial’s metrics and artifacts. Ray clusters typically run on Kubernetes via KubeRay or directly on cloud VMs.

Kubeflow Katib is the Kubernetes-native option. Its architecture is built from CRDs: an Experiment defines the search space, objective, algorithm, and max trial counts; a Suggestion is a hyperparameter set proposed by an algorithm service; a Trial wraps a user training workload (a TFJob, PyTorchJob, MPIJob, or generic Kubernetes Job). Because Trials are arbitrary container workloads, Katib is language- and framework-agnostic - any container that emits metrics through logs or files can be tuned. Search algorithms are themselves gRPC services packaged as Docker images, so adding new algorithms is a matter of writing and registering a container. Katib handles parallelism with parallelTrialCount and maxTrialCount, leverages Kubernetes for fault tolerance (failed Pods get restarted, failures count toward Experiment-level termination), and integrates with Kubeflow Pipelines for end-to-end workflows.

Google Vizier is the internal Google service that pioneered much of the modern HPO ecosystem: TPE-style search, transfer learning across studies, and large-scale parallel trial management. Its public-facing descendant is the open-source library of the same name, which serves as both a research platform and a production-grade HPO service.

Key Takeaway: Match the algorithm to the regime: grid/random for cheap models, Bayesian optimization for expensive runs at modest parallelism, ASHA/BOHB for large DL with many workers and meaningful early signals, PBT for long runs where hyperparameter schedules matter. Match the tool to your platform: Optuna for Python-native flexibility, Ray Tune for Ray clusters with rich scheduling, Katib for Kubernetes-native AutoML.

From Experiment to Pipeline

A good HPO sweep does not end with a winning configuration; it ends with a winning recipe that can be re-run reliably as part of a production pipeline. This handoff is where many ML projects quietly break, because the artifacts that justified a model in a notebook are not the same artifacts that retrain it on a schedule.

Codifying Winning Hyperparameters

The first job is to stop letting hyperparameters live in someone’s head, in a Slack message, or in a notebook cell. Best practice is to commit the winning configuration to version control as a structured config file (YAML, JSON, or a Hydra/Pydantic config object) under a path like configs/fraud_classifier/v3.yaml. The training pipeline reads this file - never inline literals - and the same file is logged as an MLflow parameter or W&B config at the start of every training run.

This has several downstream benefits. Code review now meaningfully covers hyperparameter changes. Git history records who changed learning_rate from 3e-4 to 1e-4 and when. Different environments (staging, production) can pin different config files. And the artifact that “wins” an HPO sweep is no longer a one-off notebook cell but a literal file that gets committed.

Avoiding Overfit to the Validation Set

Aggressive hyperparameter search is, statistically, a form of multiple-comparisons testing against the validation set. Run a thousand configurations and the best one will look better than its true generalization warrants - sometimes substantially. Three practices defend against this:

First, hold out a true test set that no HPO sweep ever sees. The validation set is for HPO; the test set is used exactly once, at promotion time, on the configuration the team intends to ship. Second, use nested cross-validation or rolling-origin evaluation for time series, so that hyperparameter selection happens in an inner loop and final evaluation happens in an outer loop on data the HPO never touched. Third, prefer configurations near the top of the leaderboard rather than the literal best - a configuration that is robustly excellent across cross-validation folds is more trustworthy than one that wins narrowly on one fold.

The HPO platform helps here only insofar as it makes these comparisons visible. Tracked sweeps let you ask “how many configurations were within 0.1% of the winner?” and “how stable was the winner across folds?” - questions that turn into instant queries when every trial is in a backend.

Reproducing Deterministically

Once a winning configuration is committed, the pipeline must be able to reproduce it. This requires more than the same code and config; it requires the same data version, the same library versions, the same random seeds, and ideally the same hardware-level determinism. In practice this means:

• The training pipeline pins library versions in a requirements.txt, poetry.lock, or container image, and that image (or hash) is logged with the run.
• The dataset is referenced by version - a DVC tag, a Delta Lake version, a feature store snapshot ID, an S3 object hash - not a mutable path like s3://data/train.csv.
• Random seeds are set centrally (Python random, NumPy, PyTorch/TensorFlow CPU and GPU) and logged. Where determinism matters more than throughput, set torch.use_deterministic_algorithms(True) or the TF equivalent.
• The MLflow run logs the code commit hash, the container image, the data version, and the config file path. From any registered model version you can answer “what would it take to retrain this exactly?” by reading the run metadata [Source: https://learn.microsoft.com/en-us/azure/databricks/mlflow3/genai/tracing/collect-user-feedback/].

Determinism is not always achievable bit-for-bit (some GPU kernels are non-deterministic, some libraries use system entropy), but the pipeline should at least be deterministic up to numerical noise and exactly deterministic in everything you can control.

Linking to the Model Registry

The last step in the experiment-to-pipeline bridge is the registry handoff. After a winning configuration is codified and a retraining pipeline run produces a candidate model, the pipeline should register that model in the registry with rich metadata: the source run ID, the git commit, the data version, the configuration file path, the evaluation metrics, and any validation reports. The registry then promotes the model through stages (None -> Staging -> Production -> Archived) through an explicit, auditable transition, ideally gated by automated tests and an approval workflow rather than by a human clicking a button without checks [Source: https://home.mlops.community/public/videos/mlflow-leading-open-source].

The registry-to-tracking link runs in both directions. From a registry version you can navigate to its source run and see every metric and artifact. From a tracking run you can see which registry versions, if any, were ever created from it. That bidirectional traceability is exactly what made the audit, debugging, and reproducibility stories of this chapter possible - and it is also what Chapter 9 will build on as we look at the model registry as a system in its own right.

Figure 8.6: Experiment-to-pipeline promotion

flowchart LR
    A[HPO Sweep<br/>Notebook] --> B[Winning Config]
    B --> C[Commit config YAML<br/>to git]
    C --> D[Training Pipeline<br/>reads config]
    D --> E[Tracked Run<br/>pinned data + env + seed]
    E --> F[Test-set Evaluation]
    F --> G{Pass<br/>thresholds?}
    G -->|Yes| H[Register Model<br/>with provenance]
    G -->|No| A
    H --> I[Staging]
    I --> J[Production]

Key Takeaway: Winning a sweep is not the end; codifying the winner is. Commit hyperparameters to version control, hold out a true test set the sweep never sees, pin data and environment for deterministic reproduction, and complete the loop by registering the model with full provenance back to its source run.

Chapter Summary

Experiment tracking turns ML development from folklore into engineering. By recording the full context of every training run - code commit, data version, hyperparameters, environment, metrics, and artifacts - tracking platforms enable reproducibility, comparison, audit, and the model registry workflows that production ML depends on. MLflow leads on open-source flexibility and registry maturity, Weights & Biases on collaboration and visualization, Neptune on structured metadata at scale, and Comet on hybrid SaaS/private deployment. The choice is organizational more than technical.

Hyperparameter tuning is itself a learnable problem. Grid and random search are the cheap baselines. Bayesian optimization wins at sample efficiency in moderate-dimensional spaces with expensive evaluations. Hyperband and ASHA win at compute efficiency when partial training predicts full training and you have many parallel workers. Population-Based Training wins when hyperparameters should be schedules rather than fixed values. Optuna, Ray Tune, and Kubeflow Katib implement these algorithms at scale on Python, Ray, and Kubernetes respectively, with Google Vizier as the influential ancestor.

The bridge from notebook to pipeline closes the loop. A winning sweep result becomes a committed config file; a true test set defends against overfitting to validation; pinned data, environment, and seeds guarantee reproducibility; and a model registry handoff with full provenance turns one tracked run into a versioned, auditable production model. The next chapter takes that registry as a system and shows how to design, govern, and operate it.

Key Terms

Chapter 9: Model Evaluation, Validation, and Testing

A model that scores 95% accuracy on a held-out test set can still be a disaster in production. It might be 95% accurate because the positive class is only 2% of the data and the model predicts “negative” for everything. It might leak information from the future. It might work brilliantly for one demographic and fail catastrophically for another. It might break the moment a user adds a typo or rephrases a sentence. Aggregate metrics are necessary but rarely sufficient — they tell you the average, not the failure modes.

This chapter treats evaluation as a multi-layered discipline. You will learn how to choose metrics that approximate your real business objective, how to construct validation splits that resist leakage, how to disaggregate performance across slices and fairness criteria, and how to write behavioral and invariance tests that catch failures aggregate metrics will never see. Together these layers form the pre-deployment gate that decides whether a model is ready to face real users.

Choosing Metrics

A framework for picking the right metric

Before you compute a single number, answer five questions: What decision does the model’s output drive? What are the relative costs of different errors? What kind of output does the model produce — label, probability, real value, ranked list? What constraints apply (capacity, latency, regulation)? And who consumes the metric — engineers, product managers, or executives? Only then should you pick a primary metric tied to the business objective, supplement it with secondary metrics that monitor trade-offs, and validate that offline improvements correlate with online KPIs through A/B testing [Source: https://developers.openai.com/cookbook/examples/gpt4-1_prompting_guide].

This framework matters because metrics are surrogates for utility. Optimizing AUC, RMSE, or NDCG is not the goal; minimizing fraud loss, stockout cost, or improving conversion is. The metric is just a tractable approximation of that goal — and a bad approximation produces a model that wins on the leaderboard and loses in production [Source: https://developers.google.com/machine-learning/guides/rules-of-ml].

Classification metrics

Classification metrics live and die by the confusion matrix. Precision (TP / (TP + FP)) answers “of items I flagged, how many were real?” — use it when false positives are costly, such as bothering users with marketing SMS or wrongly blocking transactions. Recall (TP / (TP + FN)) answers “of all real positives, how many did I catch?” — use it when missing a positive is costly, such as fraud, cancer screening, or safety alerts. The F1 score is the harmonic mean of the two, giving a single number that balances both but hiding the underlying trade-off [Source: https://developers.openai.com/cookbook/examples/gpt4-1_prompting_guide].

Figure 9.1: Confusion matrix anatomy and derived metrics

flowchart TD
    subgroup_actual["Actual class"]
    subgroup_pred["Predicted class"]

    A["Actual: Positive"] --> TP["True Positive (TP)<br/>Predicted: Positive"]
    A --> FN["False Negative (FN)<br/>Predicted: Negative"]
    B["Actual: Negative"] --> FP["False Positive (FP)<br/>Predicted: Positive"]
    B --> TN["True Negative (TN)<br/>Predicted: Negative"]

    TP --> P["Precision = TP / (TP + FP)"]
    FP --> P
    TP --> R["Recall = TP / (TP + FN)"]
    FN --> R
    P --> F1["F1 = 2·P·R / (P + R)"]
    R --> F1

AUC-ROC measures threshold-independent ranking quality: the probability that a randomly chosen positive scores higher than a randomly chosen negative. It is useful for comparing models early in development, but on heavily imbalanced data (e.g., 0.1% positives) it can look optimistic while precision on the positive class is terrible. PR-AUC and precision-recall curves are far more informative under severe imbalance. Log-loss, by contrast, evaluates the quality of predicted probabilities themselves and is the right metric when calibration matters — for example, when downstream business logic multiplies probabilities by dollar values to compute expected loss [Source: https://natesnewsletter.substack.com/p/context-windows-are-a-lie-the-myth].

Consider a fraud-detection model where the business goal is to minimize fraud loss while limiting customer friction. Accuracy is useless (predicting “not fraud” gives 99.9% accuracy on its own). The right workflow: examine the precision-recall curve, pick a threshold that meets a target precision of 90% (alerts must mostly be real fraud), then among thresholds meeting that constraint maximize recall. Track F1 as a summary, but use precision@k and recall@k where k matches your investigation capacity, and tie everything back to dollar fraud prevented per investigator-hour.

Regression metrics

Regression metrics differ mainly in how they treat the magnitude of errors. RMSE squares errors before averaging, so a few large misses dominate — appropriate when large errors are disproportionately harmful (energy demand forecasting, where a big underestimate causes blackouts). MAE averages absolute errors, giving each one linear weight; it is robust to outliers and easy to explain (“our predictions are off by 12 minutes on average”). MAPE expresses error as a percentage, which stakeholders love for cross-scale comparisons (revenue forecasting across markets of different sizes), but it breaks when actual values approach zero and overweights small targets [Source: https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/prompt-engineering].

Quantile loss matters when over- and under-prediction have different costs. Inventory forecasting is the canonical example: stockouts (lost sales, lost customers) often cost far more than overstock (holding cost, markdowns). A model trained with quantile loss at the 90th percentile produces predictions deliberately biased upward, accepting slightly worse RMSE in exchange for fewer stockouts. The optimal regression metric is rarely the one that minimizes squared error — it is the one whose loss function mirrors the business cost function.

Ranking metrics

Ranking metrics evaluate ordered lists, not individual predictions. NDCG@k (Normalized Discounted Cumulative Gain at k) sums graded relevance scores discounted by rank position and normalizes by the ideal ordering; it handles graded relevance (“highly relevant,” “somewhat relevant,” “irrelevant”) and is the standard for web search, recommenders, and feed ranking. MAP (Mean Average Precision) works for binary relevance when multiple items per query can be relevant — legal document retrieval is a classic case. MRR (Mean Reciprocal Rank) computes 1/rank of the first relevant item and is appropriate when users stop reading after the first good answer, as in FAQ retrieval or question answering [Source: http://susandumais.com/CHI2012-12-tailanswers-chi2012.pdf].

Hit-rate@k — the fraction of sessions where the target item appears in the top k — is a blunt but business-friendly metric for recommenders. It maps directly to “did we surface what the user wanted?” Always specify the @k cutoff that reflects the actual user-facing position; NDCG@1000 is meaningless if users only see the top 10.

Aligning metrics to KPIs

The metric selection table below summarizes when each family applies. The critical move is mapping metrics back to KPIs: the primary metric is your offline proxy, but the gold-standard validation is an A/B test showing the proxy moves with revenue, fraud loss, CTR, or whatever the business actually cares about [Source: https://www.anthropic.com/engineering/effective-context-engineering-for-ai-agents].

Key Takeaway: Choose the metric whose mathematical structure mirrors your business cost function — and validate that offline gains translate to online KPIs before declaring victory.

Validation Strategies

Train/val/test pitfalls

The textbook recipe — split data into train, validation, and test sets — is correct but constantly misapplied. The most common pitfall is using the test set repeatedly during development, which silently turns it into a second validation set and inflates your final estimate. The discipline is to lock the test set away and touch it only once, at the end, after all modeling decisions are frozen. Anything else is sample-size contamination dressed up as rigor [Source: https://developers.google.com/machine-learning/guides/rules-of_ml].

Other common errors include splitting before deduplicating (near-duplicate examples land in both train and test), splitting tabular rows when the natural unit is a user or session (allowing the model to memorize per-user patterns), and shuffling time-series data so the model trains on the future to predict the past. Each of these inflates offline metrics relative to production performance, often dramatically.

Cross-validation strategies

When data is limited, k-fold cross-validation reuses every example for both training and validation by rotating folds. Standard k-fold splits randomly, which works for i.i.d. tabular data but breaks for imbalanced or temporal datasets. Stratified k-fold preserves the class distribution within each fold and is mandatory for imbalanced classification — without it, one fold might have zero positives. Group k-fold ensures all rows belonging to the same entity (user, hospital, document) stay in the same fold, preventing leakage when entities are the unit of generalization [Source: https://developers.google.com/machine-learning/guides/rules-of-ml].

Time-series validation is its own discipline. Random splits leak the future into the training set; the model learns from tomorrow to predict today. Instead use forward-chaining (also called expanding-window) cross-validation: train on weeks 1-4, validate on week 5; train on weeks 1-5, validate on week 6; and so on. This mimics deployment, where the model is always predicting forward from a fixed point. Purged k-fold goes further by excluding a buffer of examples around the validation window to handle slowly-resolving labels (e.g., a 30-day default flag means today’s training label depends on next month’s outcome).

Figure 9.2: Forward-chaining (expanding-window) time-series cross-validation

flowchart TD
    F1["Fold 1<br/>Train: W1-W4 → Validate: W5"]
    F2["Fold 2<br/>Train: W1-W5 → Validate: W6"]
    F3["Fold 3<br/>Train: W1-W6 → Validate: W7"]
    F4["Fold 4<br/>Train: W1-W7 → Validate: W8"]

    F1 --> F2 --> F3 --> F4
    F4 --> AGG["Aggregate per-fold metrics<br/>(mean, variance across folds)"]

    style F1 fill:#1a3d5c,stroke:#58a6ff,color:#fff
    style F2 fill:#1a3d5c,stroke:#58a6ff,color:#fff
    style F3 fill:#1a3d5c,stroke:#58a6ff,color:#fff
    style F4 fill:#1a3d5c,stroke:#58a6ff,color:#fff
    style AGG fill:#0d3b2e,stroke:#58d68d,color:#fff

Time-series cross-validation (forward chaining):

Week:     1   2   3   4   5   6   7   8
Fold 1:  [T] [T] [T] [V]
Fold 2:  [T] [T] [T] [T] [V]
Fold 3:  [T] [T] [T] [T] [T] [V]
Fold 4:  [T] [T] [T] [T] [T] [T] [V]

T = train, V = validate

Hold-out and golden datasets

Beyond rotating cross-validation folds, mature teams maintain a permanent “golden dataset” — a curated, manually verified, slowly-evolving evaluation set that represents the canonical problem. Every model release runs against the golden dataset, producing a stable baseline that can be tracked across months and architectures. Golden datasets typically include edge cases, regression scenarios from past production bugs, and adversarial examples — they intentionally over-sample the hard tail rather than mirror the i.i.d. distribution.

A second hold-out, often called a “shadow” or “online evaluation” set, is collected from recent production traffic and refreshed periodically. This set catches distribution shift in a way frozen golden datasets cannot. Together the two answer different questions: “does the model still handle the cases we explicitly care about?” and “is the world drifting away from what the model was trained on?”

Data leakage detection

Data leakage is the silent killer of offline evaluations. It occurs whenever the training data contains information that would not be available at prediction time. Classic patterns: a target-derived feature (a “fraud_score” column that was computed using the fraud label), temporal leakage (using next week’s price in this week’s training row), entity leakage (the same user appearing in both train and test), and preprocessing leakage (computing normalization statistics over train+test combined before splitting) [Source: https://developers.google.com/machine-learning/guides/rules-of-ml].

Figure 9.3: Four common data leakage paths into the training set

flowchart TD
    LABELS["Ground-truth labels"]
    FUTURE["Future observations"]
    USERS["User / entity identity"]
    STATS["Train+test combined statistics"]

    LABELS -->|"Target-derived feature<br/>(e.g., fraud_score column built from label)"| TRAIN["Training set"]
    FUTURE -->|"Temporal leakage<br/>(next week's price as today's feature)"| TRAIN
    USERS -->|"Entity leakage<br/>(same user in train and test)"| TRAIN
    STATS -->|"Preprocessing leakage<br/>(scaler fit on full dataset)"| TRAIN

    TRAIN --> METRIC["Inflated offline metric"]
    METRIC --> PROD["Production performance collapses"]

    style TRAIN fill:#5c1a1a,stroke:#ff6b6b,color:#fff
    style METRIC fill:#5c1a1a,stroke:#ff6b6b,color:#fff
    style PROD fill:#5c1a1a,stroke:#ff6b6b,color:#fff

Detection techniques include: training a model with each feature individually and looking for suspiciously high single-feature AUC; comparing offline metrics to online metrics from a similar prior model and investigating large gaps; running adversarial validation (train a classifier to distinguish train from test rows — if it succeeds, your split is broken); and a feature audit that asks for every feature, “would this value actually be known at the time of the decision in production?” A leakage-free pipeline performs all feature engineering, including imputation and scaling, on the training fold only and applies the fitted transformers to validation and test data.

Key Takeaway: A validation strategy is only as good as its resistance to leakage and its fidelity to deployment conditions — split by entity, respect time, lock the test set, and audit every feature for hindsight bias.

Slice-Based and Fairness Evaluation

Why aggregate metrics hide failures

A model with 92% global accuracy can be 97% accurate on one slice and 71% on another. If that 71% slice is “new customers in emerging markets” or “users over 65,” the aggregate metric is actively misleading you. Slice-based evaluation decomposes overall performance into per-subgroup metrics so that worst-case rather than average behavior becomes visible. This is a direct application of the principle of looking for patterns in measured errors and quantifying undesirable behavior before changing the model [Source: https://developers.google.com/machine-learning/guides/rules-of-ml].

The slices that matter are domain-specific: by protected attribute (gender, race, age), by business segment (new vs. returning customers, geography, product category), by data characteristics (record length, language, image resolution), and by intersection of these. Intersectional slices routinely surface failures invisible in single-attribute analysis — Black women may have lower recall than either Black men or White women, and you will not see it if you only slice on gender or race alone [Source: https://news.ycombinator.com/item?id=44095189].

Figure 9.4: Slice-based evaluation workflow with disparity check

flowchart LR
    PRED["Predictions<br/>(y_pred, y_true)"]
    ATTR["Sensitive / segment<br/>attributes"]
    PRED --> SLICE["Group by slice<br/>(gender, age, market,<br/>intersections)"]
    ATTR --> SLICE
    SLICE --> METRICS["Per-group metrics<br/>(accuracy, recall, FPR,<br/>selection rate)"]
    METRICS --> WORST["Worst-case group<br/>min / max disparity"]
    METRICS --> DI["Disparate-impact ratio<br/>(unprivileged / privileged)"]
    WORST --> GATE{"Meets thresholds?"}
    DI --> GATE
    GATE -->|"Yes"| PASS["Pass slice gate"]
    GATE -->|"No"| MITIGATE["Mitigation:<br/>reweight / constrain / threshold"]

Subgroup performance with Fairlearn

Fairlearn’s MetricFrame is the workhorse for slice-based evaluation in Python. Given true labels, predictions, and a sensitive attribute, it computes any metric per group and overall in one call:

from fairlearn.metrics import MetricFrame, selection_rate
from sklearn.metrics import accuracy_score, recall_score

mf = MetricFrame(
    metrics={"accuracy": accuracy_score,
             "recall": recall_score,
             "selection_rate": selection_rate},
    y_true=y_true, y_pred=y_pred,
    sensitive_features=df["gender"],
)
mf.by_group      # per-group metrics
mf.difference()  # max minus min across groups
mf.group_min()   # worst-case group performance

Passing a DataFrame as sensitive_features produces intersectional slices indexed by combinations of attributes. The .difference() and .group_min() accessors immediately surface the worst-case subgroup, which is far more actionable than the global average. Common pitfalls include misaligned indices between y_true, y_pred, and sensitive features (causing ValueError: y_true and sensitive_features must have the same length) and passing nested lists instead of pandas Series [Source: https://dev.to/thebitforge/common-coding-mistakes-at-every-level-and-how-to-fix-them-4cgb].

Subgroup analysis with Aequitas

Aequitas takes a more audit-oriented approach. It expects a DataFrame with columns named score, label_value, and any number of attribute columns. The standard flow runs Preprocessor to normalize types, then Group.get_crosstabs to produce per-group ppr (predicted positive rate), tpr, fpr, fnr, and pprev (prevalence), and finally Fairness.get_fairness to compute disparity ratios versus a chosen reference group. The output is a table of ppr_disparity, tpr_disparity, and fpr_disparity values, each flagged when they cross a configurable threshold such as the 80% rule [Source: https://news.ycombinator.com/item?id=44095189].

The most common Aequitas mistake is forgetting to rename your columns to score and label_value — the library silently produces zeroes or confusing errors otherwise. The second most common is using probability scores when binary predictions are needed (or vice versa) for a given metric.

Fairness metrics compared

Three fairness criteria dominate practice: demographic parity, equal opportunity, and equalized odds. They are not interchangeable, and except in degenerate cases you cannot satisfy all three simultaneously.

The mathematical fact that demographic parity and equalized odds are incompatible when base rates differ across groups means choosing a fairness criterion is a policy decision, not a technical one. Document it, justify it, and have stakeholders sign off [Source: https://developers.google.com/machine-learning/guides/rules-of-ml].

Mitigation strategies and their trade-offs

Once disparities are detected, mitigation falls into three buckets. Pre-processing reweights or resamples training data to balance representation. In-processing adds fairness constraints to the training objective — Fairlearn’s ExponentiatedGradient reduction enforces demographic parity or equalized odds by reweighting examples during training. Post-processing adjusts decision thresholds per group to equalize chosen metrics after training.

Each carries trade-offs. Pre-processing is simple but loses information. In-processing produces principled models but requires retraining and may sacrifice accuracy globally. Post-processing is fast and reversible but uses group membership at decision time, which may be legally prohibited in domains like lending or hiring under disparate-treatment doctrine. There is no free lunch — improving worst-group recall typically costs aggregate accuracy or precision, and the size of that cost should be measured and reported alongside fairness gains.

Small slices deserve special care. A 23-row subgroup with 15 errors will show 65% error rate by sampling variance alone; bootstrapping confidence intervals and enforcing minimum sample sizes (often 100-500 depending on the metric) prevent the team from chasing noise. Conversely, aggregating tiny slices to make numbers look better can hide real harms — privacy and statistical-reliability concerns have to be balanced against transparency.

Key Takeaway: Evaluate every model on the slices that matter, choose a single fairness criterion as a policy decision, and report both worst-case subgroup performance and the accuracy cost of any mitigation.

Behavioral and Robustness Testing

Why aggregate accuracy is blind

Even after slice-based evaluation, aggregate metrics tell you nothing about whether the model handles negation, typos, paraphrases, demographic substitutions, or numerical reasoning. A sentiment model with 93% accuracy might consistently get “not good” wrong; a question-answering model might confidently change its answer when “he” becomes “she.” Behavioral testing — popularized by Ribeiro et al.’s 2020 ACL paper “Beyond Accuracy: Behavioral Testing of NLP Models with CheckList” — treats an ML model like a piece of software, probing specific capabilities with unit-test-style assertions.

The shift in mindset is from “Model A has 92% accuracy, Model B has 93%” to “Model B is better overall but fails badly on negation and demographic invariance; Model A is more robust there.” Stakeholders gain a behavioral specification of the model — a list of capabilities you have explicitly checked — that parallels requirements in traditional software engineering.

The CheckList test types

CheckList organizes tests along two axes: linguistic or reasoning capabilities (negation, coreference, intensifiers, fairness across demographics) and test types that probe those capabilities differently. Three test types dominate.

MFTs are unit tests for atomic capabilities — “the model must handle negation” becomes “predict negative on a battery of ‘not good’ sentences.” INVs are metamorphic tests — applying a transformation that should not change the answer and asserting the answer does not change. Common invariances include synonym substitution, typo injection, gender or name swapping (for fairness), and adding irrelevant filler clauses. DIRs are directional metamorphic tests — adding “very” should make positive sentences more positive, adding slurs should increase toxicity scores, removing ambiguity should increase model confidence on the correct answer [Source: https://www.promptingguide.ai/introduction/examples].

Building test suites at scale

CheckList scales because tests are generated from templates and lexicons rather than written by hand. A template like "The {adj} {noun} was {sentiment_adj}." combined with lexicons of adjectives, nouns, and sentiment words produces thousands of MFTs in seconds. Transformation functions (add_negation, swap_gender, introduce_typos, paraphrase) turn each base case into many INV and DIR variants. Hand-written tests cover the irreducibly weird cases; templates cover the bulk.

The output of running a CheckList suite is a capability x test-type matrix of pass rates: “Negation MFT: 68% pass; Spelling INV: 55% pass; Gender-swap INV: 92% pass; Intensifier DIR: 80% pass.” This is far more actionable than a single accuracy number — each failed capability points to specific data augmentation, architectural choices, or guardrails that might fix it.

Adversarial and stress testing

Behavioral tests probe capabilities the team knows about. Adversarial testing probes the ones the team does not. Adversarial example generators (e.g., TextAttack-style perturbations for NLP, projected gradient descent for vision) automatically search for inputs that flip predictions while remaining semantically equivalent or visually indistinguishable. Stress tests feed the model deliberately noisy, out-of-distribution, or rare inputs to characterize its degradation curve — what happens at low light, with code-switched language, with unfamiliar entity names, with adversarial typos.

Adversarial and stress tests usually integrate cleanly into a CheckList suite as additional INV and DIR cases. The key discipline is to distinguish robustness (graceful degradation on rare-but-natural inputs) from adversarial robustness (resistance to actively malicious inputs); they require different test distributions and different mitigations.

Shadow evaluation and pre-deployment gates

The final pre-deployment layer is shadow evaluation — running the new model on live production traffic in parallel with the existing model, logging both predictions, and comparing them without affecting users. Shadow runs reveal three things offline evaluation cannot: real-world input distribution (which often differs from any held-out set), real-world latency and resource consumption, and disagreement patterns between old and new model. A model that disagrees with production on 8% of cases for ostensibly harmless reasons may still cause user-visible surprises after launch.

Figure 9.5: Shadow evaluation architecture

flowchart LR
    USER["Production traffic"] --> ROUTER["Request router"]
    ROUTER --> LIVE["Live model<br/>(serves user)"]
    ROUTER -.->|"mirror copy"| SHADOW["Shadow model<br/>(no user impact)"]

    LIVE --> RESP["Response to user"]
    LIVE --> LOG["Prediction log"]
    SHADOW --> LOG

    LOG --> COMPARE["Compare:<br/>agreement rate, latency,<br/>distribution drift"]
    COMPARE --> REPORT["Shadow eval report<br/>(gate input)"]

    style LIVE fill:#1a3d5c,stroke:#58a6ff,color:#fff
    style SHADOW fill:#3d2a1a,stroke:#f0a020,color:#fff
    style COMPARE fill:#0d3b2e,stroke:#58d68d,color:#fff

A mature pre-deployment gate combines all the layers in this chapter into a release checklist:

1. Aggregate metrics meet the primary-metric threshold on the golden dataset.
2. Per-slice metrics meet minimum-group performance thresholds.
3. Fairness metrics (demographic parity diff, equalized odds diff, or chosen criterion) stay within agreed bounds.
4. Behavioral test suite pass rates meet thresholds per capability — with hard failures on safety-critical MFTs blocking release.
5. Shadow evaluation shows acceptable agreement and no latency regression on production traffic.
6. Online A/B test shows the offline-primary metric correlates with the business KPI on a small slice of real users.

Only when every gate passes does the model graduate to full production. This is the closest analog the ML world has to traditional software release engineering — and like in software, the cost of building the gates pays off the first time one of them catches a regression that would have shipped to users.

Key Takeaway: Behavioral, adversarial, and shadow tests turn evaluation from a one-number summary into a release checklist that catches the specific failure modes aggregate metrics will always miss.

Chapter Summary

Evaluation is the discipline that converts a trained model into a justified deployment decision. It is layered: aggregate metrics on a held-out set tell you the average story, validation splits guard against leakage and over-fitting, slice-based evaluation surfaces systematic per-subgroup failures, fairness analysis quantifies group-level disparities, and behavioral tests probe specific capabilities that aggregate metrics cannot see.

Choosing metrics begins with the business cost structure: classification metrics depend on the relative cost of false positives and false negatives, regression metrics depend on whether large errors dominate or cancel, and ranking metrics must always specify the @k cutoff that matches user-facing position. Validation strategies must respect time, entity boundaries, and class balance — and must compute every preprocessing statistic on the training fold alone. Slice-based and fairness evaluation use tools like Fairlearn’s MetricFrame and Aequitas to disaggregate performance, expose worst-case subgroups, and quantify disparities against criteria like demographic parity, equal opportunity, equalized odds, and the 80% disparate-impact rule — none of which can be simultaneously satisfied when base rates differ across groups, making the choice a policy decision. Behavioral testing in the CheckList tradition reframes evaluation as a matrix of capabilities and test types (MFT, INV, DIR), turning what was a single accuracy number into a software-style regression test suite.

The release gate that combines all these layers — golden-dataset metrics, slice thresholds, fairness bounds, behavioral pass rates, shadow agreement, and A/B-validated KPI correlation — is what separates models that ship cleanly from models that ship and break.

Key Terms

Chapter 10: Model Packaging, Registry, and Versioning

A trained model in a notebook is like a finished symphony performance recorded only as a memory in the conductor’s head: vivid, complete, and totally useless to anyone else. To move from “the model works on my machine” to “the model serves a million predictions a day with documented provenance,” teams must convert that memory into a portable artifact, file it in a registry, and stamp it with a version that links back to the exact code and data that produced it. This chapter walks through the four pillars of that work: packaging formats, registries, containerization, and a versioning strategy that ties everything together.

Think of the chapter as the chain of custody for a model. Packaging is the evidence bag, the registry is the property room, the container is the courier vehicle, and versioning is the case number that lets investigators trace any artifact back to its origin.

Section 1: Model Packaging Formats

Why Pickle and joblib Are Dangerous

The first instinct for most data scientists is pickle.dump(model, f) or joblib.dump(model, "model.pkl"). It works for scikit-learn, it is one line of code, and it round-trips to disk. The catch is that pickle is, by design, a Turing-complete bytecode for reconstructing arbitrary Python objects. Loading a pickle file from an untrusted source executes whatever code the file tells the interpreter to execute, which makes pickle a remote code execution vector dressed up as a serialization format. Even when the file is trusted, pickles are brittle: they encode references to specific class paths and library versions, so upgrading scikit-learn or moving from Python 3.10 to 3.11 can break deserialization without warning.

For internal experimentation pickle is fine. For anything that crosses a security boundary, an environment boundary, or a multi-year compatibility horizon, you want a format whose contract is “data, not code.” That is exactly the niche that ONNX, TorchScript, SavedModel, GGUF, and safetensors occupy [Source: https://mlflow.org/docs/latest/python_api/mlflow.entities.html].

ONNX: The Cross-Framework Intermediate Representation

ONNX (Open Neural Network Exchange) is a protocol-buffer description of a directed acyclic computation graph, plus a versioned set of operators called an opset. The spec is deliberately independent of any single framework, so a PyTorch model exported to ONNX can be loaded by ONNX Runtime in C++, by TensorRT on an NVIDIA GPU, or by OpenVINO on an Intel CPU without any of those runtimes needing to know that PyTorch exists [Source: https://developers.openai.com/cookbook/examples/gpt4-1_prompting_guide].

The analogy to keep in mind is a PDF. A Word document is tied to Word, a Pages document is tied to Pages, but a PDF is the print-ready exchange format that any reader can render. ONNX plays the same role for neural networks. The trade-off is the same too: like PDF, ONNX is good at preserving the structure of a finished artifact and less good at preserving the editable, dynamic behavior of the original.

In practice, you export from PyTorch with torch.onnx.export (classic) or torch.onnx.dynamo_export (PyTorch 2.x), and from TensorFlow with tf2onnx. Common failure modes are unsupported operators (“Operator X is not supported in opset Y”), control flow that depends on non-tensor Python values, and dynamic shapes that were not declared at export time [Source: http://susandumais.com/CHI2012-12-tailanswers-chi2012.pdf].

Figure 10.1: ONNX as the cross-framework intermediate representation between training frameworks and inference runtimes

flowchart LR
    PT[PyTorch Model] -->|torch.onnx.export| ONNX[(ONNX Graph<br/>+ Opset)]
    TF[TensorFlow Model] -->|tf2onnx| ONNX
    SK[scikit-learn] -->|skl2onnx| ONNX
    ONNX --> ORT[ONNX Runtime<br/>CPU/GPU]
    ONNX --> TRT[TensorRT<br/>NVIDIA GPU]
    ONNX --> OV[OpenVINO<br/>Intel CPU/NPU]
    ONNX --> TRI[Triton<br/>ONNX Backend]

TorchScript and TensorFlow SavedModel: Framework-Native Formats

TorchScript is PyTorch’s answer to “how do I run my model without the Python interpreter.” It compiles a restricted subset of Python plus tensor operations into a JIT graph that can be loaded by LibTorch in C++ or PyTorch Mobile on a phone. You produce it via torch.jit.script (preferred when the model has control flow) or torch.jit.trace (records the graph from example inputs, but only captures the branches actually exercised) [Source: https://platform.claude.com/docs/en/build-with-claude/prompt-engineering/claude-prompting-best-practices].

TensorFlow SavedModel is not a single file but a directory containing one or more MetaGraphs (protobufs), a variables checkpoint, optional assets such as vocabulary files, and named serving signatures like serving_default. That structure is what TF Serving, TFX, TF Lite, and Vertex AI all consume. Because the format can carry multiple graphs (train, serve, eval) and tightly integrates with tf.function-decorated callables, it is the natural canonical format for TensorFlow-end-to-end shops [Source: https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/prompt-engineering].

A key 2024-2025 trend is that PyTorch’s compiler innovation has shifted toward torch.export and torch.compile. TorchScript is still production-grade, but new large bets on cross-framework workflows tend to go through ONNX rather than doubling down on TorchScript [Source: https://www.publichealth.columbia.edu/research/population-health-methods/content-analysis].

GGUF and safetensors for LLMs

Two newer formats matter specifically for large language models. safetensors is a flat, mmap-friendly tensor container that, like ONNX, encodes data not code; loading a safetensors file cannot execute arbitrary Python, which is the exact property that makes it a safer drop-in for the legacy pytorch_model.bin pickles that dominate Hugging Face. It also enables zero-copy loads from disk, so massive checkpoints come up faster.

GGUF is the format behind the llama.cpp ecosystem. It bundles weights, tokenizer, metadata, and chat templates into a single quantization-aware file optimized for CPU and edge inference. The mental model is that safetensors is the secure cousin of the bare weights file, while GGUF is the all-in-one cartridge that a local LLM runtime can plug in and play without any framework dependency.

Key Takeaway: Pickle is convenient but executes arbitrary code on load; production packaging belongs in framework-neutral formats (ONNX, safetensors, GGUF) or in tightly integrated framework-native formats (TorchScript, SavedModel) chosen to match the deployment runtime.

Section 2: The Model Registry

Artifacts, Metadata, and Lineage

A registry is to a model what a library card catalog is to a book: not the content itself, but the index that makes the content findable, attributable, and governable. A registry entry typically bundles three things: the artifact (or a pointer to it in object storage), the metadata (training metrics, hyperparameters, data schema, signature), and the lineage (which run produced it, which dataset version it consumed, which git commit of the training code). Without those three, a model is just a .bin file in a bucket [Source: https://mlflow.org/docs/latest/python_api/mlflow.entities.html].

Lineage is the part most teams underinvest in. The question “which training run produced the model currently serving 95% of production traffic” should be answerable in one click. If it requires archaeology in Slack, the registry is failing at its job [Source: https://www.youtube.com/watch?v=daBTYQP23-A].

Stages: Staging, Production, Archived

The classic MLflow lifecycle defines four explicit stages: None, Staging, Production, and Archived. A version starts in None after registration, moves to Staging for integration testing, gets promoted to Production once it passes business and quality gates, and is Archived when superseded. The MLflow Python API expresses this with a single call:

client.transition_model_version_stage(
    name="churn_model",
    version="3",
    stage="Production",
    archive_existing_versions=True,
)

The archive_existing_versions=True flag is the small but vital touch that prevents two versions from claiming to be “Production” at the same time [Source: https://www.youtube.com/watch?v=6ngxBkx05Fs].

Figure 10.2: MLflow model registry lifecycle states and the transitions between them

stateDiagram-v2
    [*] --> None: register_model()
    None --> Staging: transition(Staging)
    Staging --> Production: transition(Production,<br/>archive_existing=True)
    Staging --> Archived: superseded
    Production --> Archived: new version promoted
    Production --> Staging: rollback for re-eval
    Archived --> Staging: re-promote for rollback
    Archived --> [*]

Newer MLflow versions (>=1.30) add aliases such as @prod or @champion, which are mutable pointers to a specific version. Aliases decouple downstream serving code from raw version numbers: a serving container that loads models:/churn_model@prod keeps working when you re-point prod to v8, with no redeployment needed [Source: https://home.mlops.community/public/videos/mlflow-leading-open-source].

MLflow, Vertex AI, and SageMaker Registries

The three dominant registries in 2024-2025 share the core idea of versioned model entries but diverge sharply on governance, lineage, and cloud integration.

MLflow is the lightweight, vendor-neutral option. The trade-off is that you build governance yourself: a typical pattern is a GitHub pull request that toggles an alias, combined with MLflow tags like approved_by and approval_ticket [Source: https://docs.databricks.com/aws/en/mlflow3/genai/eval-monitor/code-based-scorer-examples].

Vertex AI is the registry of choice when you live inside Google Cloud. It does not enforce stages, but it makes promotion natural by combining version aliases, labels, and deployment targets such as a prod-endpoint with 10% canary traffic. Vertex’s killer feature is end-to-end lineage: a single graph that shows “this prod model v5 was trained by pipeline X, which used dataset Y and code version Z” [Source: https://mlflow.org/docs/latest/genai/serving/responses-agent/].

SageMaker leans hardest into governance. ModelApprovalStatus is an explicit field on every Model Package, and the typical CI/CD pattern is: pipeline trains and registers v7 as PendingManualApproval, an approver in the ML-Governance IAM role reviews the Model Card and metrics, sets the status to Approved, and a deployment pipeline detects the state change and updates the production endpoint [Source: https://www.youtube.com/watch?v=bDflB17YUNc]. That explicit gate is exactly what regulated industries need to satisfy auditors.

Approval Workflows and RBAC

Across all three registries, the operational rule is the same: promotion requires stronger permissions than read access. A junior engineer’s notebook should be able to register experimental versions but not promote anything to production. The mechanics differ - IAM policies in SageMaker, Vertex IAM bound to service accounts in GCP, MLflow’s auth plugins or external RBAC in OSS deployments - but the principle is universal. The registry is the chokepoint where governance attaches, and a registry without RBAC is a sticky note labeled “production” [Source: https://www.certlibrary.com/exam/Certified%20Machine%20Learning%20Professional].

Key Takeaway: Registries centralize versioned artifacts, metadata, and lineage; choose MLflow for vendor neutrality, Vertex for GCP-native lineage, and SageMaker for the strongest built-in approval and audit story, and always gate promotion with RBAC.

Section 3: Containerization for Serving

Model-Serving Docker Images and Multi-Stage Builds

A serving image is a contract that says “given an HTTP or gRPC request, I will load this exact artifact in this exact runtime and return a prediction.” The cleanest way to build that contract is a multi-stage Dockerfile that strictly separates the heavy build-time tools from the lean runtime [Source: https://www.blacksmith.sh/blog/understanding-multi-stage-docker-builds].

The pattern is straightforward: a builder stage based on a CUDA -devel image installs compilers, Python build dependencies, and any custom CUDA kernels, producing wheels, .so files, and exported model artifacts. A runtime stage based on nvcr.io/nvidia/tritonserver:<version>-py3 or a CUDA -runtime image then COPY --from=builder only the produced artifacts. The result is an image that contains the model and serving binary but no compiler, no header files, and no apt cache [Source: https://cycle.io/learn/multi-stage-builds].

FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS builder
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential cmake protobuf-compiler \
    && rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip wheel --no-cache-dir -r requirements.txt -w /wheels

FROM nvcr.io/nvidia/tritonserver:24.05-py3 AS runtime
COPY --from=builder /wheels /wheels
RUN pip install --no-cache-dir /wheels/*.whl && rm -rf /wheels
RUN useradd -r -u 10001 triton
USER triton
COPY --chown=triton:triton models/ /models/
EXPOSE 8000 8001 8002
CMD ["tritonserver", "--model-repository=/models"]

Two cache-friendly habits make a real difference: copy requirements.txt before the rest of the source so dependency layers are reused across code edits, and group apt-get install with rm -rf /var/lib/apt/lists/* in the same RUN so the cache never lingers in a lower layer [Source: https://docs.docker.com/build/building/best-practices/].

Figure 10.3: Multi-stage container build pipeline from source to scanned registry image

flowchart LR
    SRC[Source + Dockerfile<br/>+ requirements.txt] --> BUILD[Builder Stage<br/>CUDA -devel<br/>compilers, wheels]
    BUILD -->|COPY --from=builder| RT[Runtime Stage<br/>tritonserver -py3<br/>non-root user]
    RT --> IMG[Image Layers]
    IMG --> SCAN[Vulnerability Scan<br/>Trivy / Grype / Scout]
    SCAN -->|pass| REG[(Container Registry<br/>tagged + digest)]
    SCAN -->|critical CVE| FAIL[Fail CI]
    REG --> K8S[Kubernetes /<br/>Serving Cluster]

TorchServe, TF Serving, and Triton Base Images

Each major serving stack publishes a vendor-supported base image, and starting there saves weeks of CUDA debugging.

• NVIDIA Triton Inference Server is the multi-backend powerhouse: a single Triton image can serve ONNX, TorchScript, TensorFlow SavedModel, TensorRT engines, and Python backends from the same process. The official nvcr.io/nvidia/tritonserver:<version>-py3 images bundle matched CUDA, cuDNN, TensorRT, and Triton binaries.
• TorchServe typically starts from pytorch/torchserve or a custom pytorch/pytorch image with TorchServe installed; you must align the PyTorch CUDA build (e.g., +cu121) with the base CUDA image to avoid silent fallback to CPU.
• TensorFlow Serving is published by Google as a small C++ server with both CPU and GPU variants.

Triton’s multi-backend nature is the reason many organizations adopt a “default ONNX, native exceptions” pattern: most models export to ONNX and serve through the Triton ONNX backend, while the handful of models that resist clean ONNX export get served natively through the PyTorch or TensorFlow backends in the same server [Source: https://nickjanetakis.com/blog/shrink-your-docker-images-by-50-percent-with-multi-stage-builds].

Image Size and Cold-Start

Image size is not just a storage concern - it is a cold-start concern. Every additional gigabyte is a gigabyte that must be pulled to a fresh Kubernetes node before the first prediction can be served, and on a scale-from-zero autoscaler that latency is in the user-facing critical path [Source: https://depot.dev/blog/docker-multi-stage-builds].

The levers are well-rehearsed: multi-stage builds, runtime-only CUDA images, --no-cache-dir for pip, avoiding shells/editors/curl in runtime images, stripping debug symbols where safe, and mounting large model repositories as volumes from object storage rather than baking every version into the image [Source: https://www.harness.io/blog/how-to-create-multi-stage-docker-builds-with-harness-continuous-delivery]. For Triton fleets that host dozens of models, the volume-mount pattern is especially powerful: one slim Triton image plus per-model artifacts in S3 or GCS, instead of one fat image per model.

The analogy to keep in mind is luggage. A runtime image should be a carry-on with exactly the runtime, the model, and a non-root user, not a steamer trunk that also contains the compiler, the test suite, and three Python interpreters “just in case.”

Pinning CUDA, Drivers, and ABI

GPU containers are uniquely fragile because the host kernel driver, the CUDA user-space libraries inside the container, and the framework’s CUDA build (+cu121, +cu118, etc.) must all be compatible. The standard discipline is:

1. Do not bake driver components into the image. The host provides the kernel driver; the container provides the CUDA runtime.
2. Run containers with the NVIDIA Container Runtime (--gpus=all on Docker, nvidia.com/gpu resources in Kubernetes).
3. Pin explicit versions everywhere: base image tag, Triton/TorchServe version, CUDA, Python, and PyTorch/TF build. Never use latest in production.
4. Rebuild regularly with docker build --pull to pick up base-image security updates [Source: https://www.youtube.com/watch?v=ajetvJmBvFo].

Runtime security hardening completes the picture: run as a dedicated non-root user (useradd -r -u 10001 triton; USER triton), set readOnlyRootFilesystem: true in Kubernetes, drop all Linux capabilities you do not need, never use --privileged, expose only the required ports (Triton’s HTTP/gRPC/metrics on 8000/8001/8002, TorchServe’s inference/management APIs), and inject secrets via environment variables or secret managers rather than baking them into layers. Pair this with continuous vulnerability scanning (Trivy, Grype, Anchore, Docker Scout) that fails CI on critical CVEs.

Key Takeaway: Multi-stage Docker builds on vendor base images give you small, secure, cold-start-friendly serving containers; rely on host GPU drivers, pin every version, run as non-root, and scan images on every build.

Section 4: Versioning Strategy

Semantic Versioning for Models

Semantic versioning (MAJOR.MINOR.PATCH) was designed for libraries, but it adapts naturally to models if you interpret the three numbers in terms of consumers of predictions rather than callers of an API.

• MAJOR bumps when the model’s interface changes in a breaking way: a new input feature, a different output schema, a re-encoded label space. Consumers must update integration code.
• MINOR bumps when the model’s behavior changes in a meaningful but compatible way: a retrain on fresh data, a new architecture that produces the same shape and similar quality. Consumers can adopt without code changes but should re-evaluate downstream metrics.
• PATCH bumps for fixes that should be transparent: a bug fix in preprocessing that does not measurably change predictions, an optimization that preserves outputs.

This discipline is what lets a downstream team pin churn_model >=2.3,<3.0 and trust that they will pick up improvements without breaking their pipeline.

Linking Model to Code SHA to Dataset Version

A model version is meaningful only as a tuple of three things: the model artifact, the exact code commit that trained it, and the exact dataset snapshot that fed it. If any of the three is missing, the model is not reproducible, which means you cannot debug regressions or satisfy auditors.

The practical pattern stitches together three tools:

• Code versioning via Git; record the commit SHA in the model’s metadata.
• Data versioning via DVC, LakeFS, Delta Lake time travel, or simply an immutable dataset hash; record the dataset version in the model’s metadata.
• Model artifact versioning via the registry (MLflow version number, Vertex Model version, SageMaker Model Package).

MLflow makes this concrete because runs log code version (Git SHA) and any parameter you choose to record, and each registered model version stores the run_id pointer. You can therefore drill from “production v7” back to “trained by run abc123, code SHA 7f3a9c1, dataset hash sha256:e8b...” in a single query [Source: https://mlflow.org/docs/latest/python_api/mlflow.entities.html]. Vertex Lineage and SageMaker Lineage provide richer graph-based equivalents.

The chain is only as strong as its weakest link. A model artifact with no dataset hash is an orphan; a Git SHA with no model hash is a thought experiment. The registry is where you bolt the chain together [Source: https://www.youtube.com/watch?v=daBTYQP23-A].

Figure 10.4: Linkage between a model version and the code, data, training run, and runtime that produced it

graph LR
    MV[Model Version<br/>v7] --> CODE[Git Commit SHA<br/>7f3a9c1]
    MV --> DATA[Dataset Hash<br/>sha256:e8b...]
    MV --> RUN[Training Run<br/>run_id abc123]
    MV --> IMG[Container Image<br/>tag + digest]
    RUN --> METRICS[Metrics &<br/>Hyperparams]
    DATA --> DVC[DVC / LakeFS /<br/>Delta time-travel]
    CODE --> REPO[Git Repository]
    IMG --> REG[(Container<br/>Registry)]

Promotion Criteria

Promotion from staging to production should be a checklist, not a hunch. A defensible promotion gate typically includes:

1. Offline quality: holdout metrics meet or exceed the current production version by a defined margin on the same evaluation set.
2. Subgroup fairness: performance does not regress on protected or business-critical slices.
3. Latency and throughput: p95 latency under the SLA, throughput at expected QPS within infrastructure budget.
4. Shadow or canary results: shadow traffic or a small canary cohort shows no regression on online metrics.
5. Governance artifacts: Model Card (or equivalent), risk assessment, and approval ticket are attached.
6. Lineage completeness: code SHA, dataset version, and training run are all linked.

In MLflow this often manifests as a CI job that runs the checklist and only then calls transition_model_version_stage(stage="Production", archive_existing_versions=True). In Vertex it manifests as a pipeline step that conditionally moves the prod alias. In SageMaker it manifests as an update_model_package(ModelApprovalStatus="Approved") call gated behind an IAM-protected human approver [Source: https://www.youtube.com/watch?v=6ngxBkx05Fs].

Figure 10.5: Model promotion workflow from data scientist commit through CI checks and SRE approval to production traffic

sequenceDiagram
    participant DS as Data Scientist
    participant CI as CI Pipeline
    participant REG as Model Registry
    participant SRE as SRE / Approver
    participant PROD as Prod Endpoint
    DS->>CI: Push training code + config
    CI->>CI: Train, evaluate, log run
    CI->>REG: Register version v7 (None)
    CI->>REG: Run quality / latency / fairness gates
    REG->>SRE: Notify PendingApproval
    SRE->>REG: Review Model Card + lineage
    SRE->>REG: Approve / move @prod alias
    REG->>PROD: Serving container reloads model
    PROD-->>DS: Live predictions on v7

Rollback Strategies

Rollback is the dual of promotion and should be just as cheap. Three patterns dominate:

1. Re-point the alias. In MLflow or Vertex, move the @prod alias back to the previous version. The serving container picks up the change at next model load, with no redeployment.
2. Re-approve a previous Model Package. In SageMaker, the previous version remains Approved; the rollback is to update the endpoint configuration to reference it again. CloudTrail captures the action for audit.
3. Blue/green and traffic splitting. Keep the previous version warm on a fraction of traffic during canary deployment. If the new version regresses, shift 100% of traffic back instantly without any artifact movement [Source: https://www.certlibrary.com/exam/Certified%20Machine%20Learning%20Professional].

The non-negotiable property is that rollback must not require retraining. If your only path to undo a bad model is “re-run the training pipeline on yesterday’s data,” your registry is doing the property-room job poorly. Every previous production version should remain available, indexed, and one alias-flip away from being live again.

The analogy here is firefighting. Promotion is the building inspector’s signed certificate of occupancy; rollback is the sprinkler system you hope you never use but test every quarter. Both belong in the design, not improvised at 3 a.m.

Key Takeaway: Treat model versioning as a triple of code SHA, dataset version, and artifact ID linked through the registry, promote only against an explicit checklist, and design rollback as a one-command alias flip rather than a retraining exercise.

Chapter Summary

Packaging a model for production is the act of converting an in-memory experiment into a portable, auditable artifact and giving it a permanent address. The portable artifact comes from choosing the right format: ONNX when you need to cross framework or hardware boundaries, TorchScript or SavedModel when you stay inside a single ecosystem, safetensors for safe LLM weight loading, and GGUF for self-contained edge LLM cartridges. Pickle and joblib should be reserved for ephemeral experiments because they execute arbitrary code on load.

The permanent address comes from a registry. MLflow gives you a lightweight, vendor-neutral catalog with built-in lifecycle stages and aliases. Vertex AI gives you deep GCP-native lineage and a flexible alias-driven promotion model. SageMaker gives you the strongest first-class governance with explicit approval status and Model Cards. All three converge on the same essentials: versioned artifacts, attached metadata, traceable lineage, and RBAC-gated promotion.

Serving turns the registry entry into a running service through containers. Multi-stage Docker builds on vendor base images such as nvcr.io/nvidia/tritonserver produce small, fast-starting, hardened images that run as non-root, rely on host GPU drivers, expose only required ports, and are continuously scanned for vulnerabilities. Triton’s multi-backend nature in particular enables a “default ONNX, native exceptions” pattern that simplifies fleets.

Versioning ties everything together. Semantic versioning gives consumers a contract. Linking model version to code SHA and dataset version gives auditors a reproducibility chain. Promotion gates turn lifecycle transitions into checklist-driven events, and alias-based rollback makes recovery a single command rather than a retraining incident. Done right, the entire stack - artifact, registry, container, version - acts as a chain of custody that the next engineer, the next regulator, and the next on-call rotation can all trust.

Key Terms

Chapter 11: Model Deployment Patterns: Batch, Online, and Edge

A trained model creates no value until it reaches a prediction surface — a nightly scoring table, a user-facing API, a stream operator, or a phone in someone’s pocket. The choice of deployment pattern is one of the highest-leverage decisions in an ML system. It dictates latency budgets, cloud bills, on-call complexity, and how quickly you can iterate. A team that picks the wrong pattern often discovers the mistake only after months of operational pain: a recommender that reloads embeddings every request, a fraud system that runs nightly when it should run per-transaction, or a vision model bundled into a mobile app that drains battery in twenty minutes.

This chapter develops a practical decision framework across four dimensions. First, we distinguish the four core inference patterns — batch, online, streaming, and embedded — and show how their trade-offs along latency, throughput, cost, and complexity map to real use cases like recommendations, fraud detection, and content moderation. Second, we cover safe rollout strategies — shadow, canary, A/B, and blue-green — and the ML-specific monitoring that makes them trustworthy. Third, we explore edge and mobile deployment, where quantization, pruning, and distillation determine whether a model is usable at all. Finally, we survey the serving-platform landscape, from self-hosted KServe and Seldon to managed and serverless options, so you can pick the right substrate for your scale and team.

Inference Patterns

Batch Offline Inference

Batch inference runs the model on large groups of inputs at scheduled intervals — hourly, nightly, or whenever a pipeline upstream writes new data. Inputs are read from a warehouse or object store, predictions are written back to a table or cache, and downstream applications read the precomputed scores when they need them. The trigger is a scheduler — Airflow, Argo, Prefect, or cron — not a user request, so the model has no real-time relationship with the consumer of its predictions [Source: https://blog.codinghorror.com/the-problem-with-logging/].

The classic batch stack uses Spark, Beam, Dask, or plain Python on Kubernetes/EMR/Databricks, with predictions landing in BigQuery, Snowflake, S3, or a relational database. Latency between data arrival and prediction availability ranges from minutes to hours, but throughput is enormous — billions of records per job are routine — and cost per prediction is the lowest of any pattern because work batches efficiently and can run on off-peak or spot capacity [Source: https://pub.towardsai.net/how-i-cut-my-llm-costs-by-80-without-sacrificing-quality-85f8505eec96]. Think of batch like a printing press: setup is expensive, but once it’s running, each page is nearly free.

Batch fits when staleness is acceptable. A retailer that scores customer lifetime value nightly can serve a CRM team perfectly well; a streaming service that precomputes “top-N items per user” between 2 a.m. and 4 a.m. delivers fast recommendations because the application just reads a Redis key at request time. Batch also dominates re-scoring and backfill workloads: when a new model ships, the easiest way to populate scores for every historical user is a one-off batch job.

Key Takeaway: Batch inference offers the lowest cost and highest throughput in exchange for stale predictions; choose it whenever downstream applications can tolerate minutes-to-hours of latency.

Synchronous Online Inference

Online inference runs the model per user request, synchronously, behind a REST or gRPC endpoint. A client — web, mobile, or backend service — calls the model service, which fetches features, runs the model, and returns a response in the same HTTP round trip. Typical latency budgets sit between 1 and 200 milliseconds at the 95th percentile, and the serving fleet must be provisioned for peak QPS rather than average load [Source: http://susandumais.com/CHI2012-12-tailanswers-chi2012.pdf].

The online stack centers on a serving runtime — FastAPI, Flask, gRPC servers, or specialized servers like TensorFlow Serving, TorchServe, BentoML, and NVIDIA Triton — fronted by a load balancer or API gateway. Features come from a feature store (Feast, Tecton, or a custom service) or a low-latency cache such as Redis or DynamoDB. Because the model is in the critical path of every request, cost per prediction is the highest of the patterns: capacity must be ready for peak traffic, GPUs may sit idle to guarantee p99 latency, and autoscaling overhead is real.

Online is the only choice when the user is waiting on the response. Fraud scoring during card authorization, search ranking, autocomplete, chatbot replies, and ad targeting all require synchronous predictions. The engineering discipline online inference demands — autoscaling, circuit breakers, request hedging, careful feature retrieval — is significant, but it is the price of putting a model in the user’s critical path.

Key Takeaway: Online inference is unavoidable when a human is waiting on the answer; budget for always-on capacity, strict latency engineering, and far higher cost per prediction than batch.

Asynchronous and Streaming Inference

Streaming inference sits between batch and online. The model runs continuously on a flow of events arriving through Kafka, Kinesis, Pub/Sub, or Pulsar, with frameworks like Apache Flink, Spark Structured Streaming, Kafka Streams, or Beam orchestrating the work [Source: https://news.ycombinator.com/item?id=43499862]. Predictions are themselves another stream, written back to a topic or feature store for consumers to subscribe to. Latency is near-real-time — tens of milliseconds to a few seconds — and throughput can reach millions of events per second, but operational complexity is the highest of the patterns: checkpointing, exactly-once semantics, backpressure, and stateful windowing all have to be handled.

Streaming shines when freshness is critical and decisions are time-ordered. A trending-content algorithm that needs to react within seconds to a viral video, a fraud system that maintains a rolling 5-minute transaction count per card, a live-chat moderator that flags abusive messages within a second — all are natural fits. Streaming also powers the feature side of hybrid systems: a Flink job continuously updates “items viewed in last 10 minutes” features that an online model reads at request time.

A lighter cousin of streaming is async inference: clients send a request, the system enqueues it, a worker runs the model, and the client either polls for the result or receives a callback. Async relaxes the synchronous latency contract and reduces peak-capacity needs — useful for slow models (e.g., document summarization) where users can wait a few seconds but the system would buckle under synchronous load.

Key Takeaway: Streaming inference delivers fresh predictions over continuous event flows at the cost of significant infrastructure complexity; use it when missing time-sensitive signals carries real business cost.

Embedded and On-Device Inference

Embedded inference runs the model directly on the device generating the data: a phone, a camera, an industrial sensor, a car. There is no network call to a server. The model ships with the application binary or is downloaded over the air, and predictions happen inside the user’s hardware [Source: https://discuss.pytorch.org/t/mobile-deployment-best-practice/96197].

On-device deployment unlocks three properties that no server-side pattern can match. Privacy improves because raw input — a photo, a heart-rate reading, a voice clip — never leaves the device. Latency drops to single-digit milliseconds because there is no round trip; this is essential for AR, real-time camera effects, and offline voice assistants. Reliability improves in poor-connectivity environments: a tractor scoring crop health in a field cannot wait for a 4G signal. The trade-offs are equally real: model size, RAM, battery, and thermal limits cap what’s possible, OTA model updates are a separate engineering problem, and you give up the easy observability of server-side logs.

The following table summarizes the four patterns side by side and forms the decision-making backbone of this section.

Most production systems combine these patterns. A typical recommender precomputes candidate sets in batch nightly, updates a “recently viewed” feature via streaming, and runs final ranking with online inference at page load — a Lambda-like architecture that pushes heavy work offline and reserves online for the cheap final step.

Figure 11.1: Inference pattern comparison across latency, trigger, and canonical use cases

graph TD
    A[Inference Patterns] --> B[Batch Offline]
    A --> C[Online Synchronous]
    A --> D[Streaming]
    A --> E[Embedded/Edge]
    B --> B1[Trigger: Scheduler<br/>Latency: Minutes-Hours<br/>Cost: Low]
    B --> B2[Use: Nightly CLV<br/>Precomputed Recs<br/>Backfills]
    C --> C1[Trigger: Per Request<br/>Latency: 1-200ms p95<br/>Cost: High]
    C --> C2[Use: Fraud Auth<br/>Search Ranking<br/>Autocomplete]
    D --> D1[Trigger: Event Flow<br/>Latency: Sub-second<br/>Cost: Medium-High]
    D --> D2[Use: Trending Content<br/>Live Moderation<br/>Rolling Features]
    E --> E1[Trigger: Local App<br/>Latency: Single-digit ms<br/>Cost: None at runtime]
    E --> E2[Use: AR Filters<br/>Offline Voice<br/>On-device Vision]

Key Takeaway: No production system uses a single pattern; the engineering art is choosing which work belongs in batch, online, streaming, or on-device — and combining them to maximize freshness while minimizing cost.

Safe Rollout Strategies

Shipping a new model is riskier than shipping new code. Models can degrade silently — no 5xx errors, no stack traces, just worse predictions. Labels and business outcomes arrive with delay, so you may not know a rollout is bad for hours or days. And rolling back means reverting not just code but the model artifact, the feature pipeline that fed it, and the configuration that wired them together [Source: https://erichorvitz.com/tail_answers.pdf]. The four rollout strategies below — shadow, canary, A/B, and blue-green — exist to manage this risk. Mature teams use them in sequence rather than in isolation.

Shadow and Mirror Traffic

Shadow mode (sometimes called dark launch or mirror mode) tests a new model under real production traffic without exposing users to its predictions. The current production model continues to serve responses; the new candidate receives a mirrored copy of the same requests, runs inference, and logs predictions for offline comparison — but its outputs never reach the user [Source: https://www.cliffsnotes.com/study-notes/28411172].

The implementation is mechanically simple at the routing layer. In Istio, a VirtualService adds a mirror directive that duplicates traffic to a second backend. Seldon Core exposes a “shadow predictor” inside a SeldonDeployment. KServe combines a separate InferenceService with mesh-level mirroring. The harder problem is preventing side effects: if the new model writes to a database, increments counters, or calls external APIs, those must be disabled or routed to isolated targets — otherwise “shadow” silently affects production.

Figure 11.2: Shadow deployment with mirrored traffic and offline comparison

flowchart LR
    U[User Request] --> GW[API Gateway / Mesh]
    GW -->|Primary Response| LIVE[Live Production Model v1]
    LIVE -->|Returned to user| U
    GW -.->|Mirrored Copy| SHADOW[Shadow Candidate Model v2]
    SHADOW -.->|Predictions Only| LOG[(Prediction Log)]
    LIVE -->|Predictions| LOG
    LOG --> CMP{Offline Comparator}
    CMP -->|Score Deltas<br/>Drift<br/>Latency| EVAL[Evaluation Report]
    SHADOW -.->|Side effects DISABLED| EXT[External APIs / DB Writes]

Shadow mode is the right time to ask three questions. Does the candidate handle real input distributions without schema errors or NaN explosions? Does its latency and resource profile fit the production envelope? And on offline metrics computed from shadow logs — score deltas, distribution shifts, eventual label-based performance — does it at least match the incumbent? Rollback is trivial: stop mirroring, since users never saw anything. The cost is doubled inference compute during the shadow window, which is worth it for high-stakes models.

Key Takeaway: Shadow mode is the cheapest insurance against silent model regressions; mirror full production traffic, disable side effects, and only promote candidates that match or beat incumbents on real-data offline metrics.

Canary and Progressive Rollout

A canary release sends a small slice of real traffic — typically 1 to 5% — to the new model and watches metrics in near real time. If signals are healthy, traffic ramps to 10%, 25%, 50%, and finally 100% over hours or days; if anything degrades, traffic instantly reverts to the incumbent [Source: https://www.marks4sure.com/sy0-701-comptia-securityp-exam-questions.html]. Canary is operational risk mitigation, not a statistical experiment — its job is to catch catastrophic failures before they reach everyone.

The traffic-splitting machinery is well-established. Istio’s VirtualService supports weighted routing (weight: 95 to v1, weight: 5 to v2) that can be adjusted via config or a rollout controller like Argo Rollouts or Flagger. Seldon Core lets you declare multiple predictors with traffic percentages in a single SeldonDeployment. KServe exposes a canaryTraffic field on its InferenceService that controls the split with one number.

Figure 11.3: Canary rollout state machine with progressive traffic ramps and rollback

stateDiagram-v2
    [*] --> Shadow0: Deploy candidate
    Shadow0 --> Canary1: Pass shadow checks
    Canary1 --> Canary5: Healthy at 1%
    Canary5 --> Canary25: Healthy at 5%
    Canary25 --> Canary50: Healthy at 25%
    Canary50 --> Promoted100: Healthy at 50%
    Promoted100 --> [*]: Old version retired
    Canary1 --> Rollback: SLO breach
    Canary5 --> Rollback: SLO breach
    Canary25 --> Rollback: Drift / metric drop
    Canary50 --> Rollback: Drift / metric drop
    Rollback --> [*]: Traffic reverts to v1

Monitoring during canary spans three classes: system metrics (p50/p95/p99 latency, error rates, pod restarts), model quality (CTR, conversion, AUC once labels arrive), and data/drift (feature distribution shifts between variants, training–serving skew). Because labels often lag, early canary decisions rely on proxy metrics — short-term engagement, add-to-cart rates — rather than the eventual metric you ultimately care about. Two ML-specific pitfalls deserve attention: ensure canary traffic is representative (don’t accidentally route only one region or one device type), and consider sticky assignment so users don’t see flipping behavior between requests.

Key Takeaway: Canaries catch operational failures fast by ramping traffic gradually with automated rollback triggers; always keep the previous model hot and capable of taking 100% of traffic instantly.

A/B Tests and Multi-Armed Bandits

A/B testing is a randomized statistical experiment, not a rollout. A canary asks “is the new model breaking?”; an A/B test asks “is the new model genuinely better on business metrics?” A typical setup splits traffic 50/50 (or some experiment-specific ratio) with sticky per-user assignment — once a user is in variant B, they stay there for the duration — and runs for a predetermined window with predefined primary and secondary metrics [Source: https://developers.openai.com/cookbook/examples/gpt-5/gpt-5-2_prompting_guide].

Implementation typically separates the assignment layer from the routing layer. The application or an experiment service uses consistent hashing (variant = hash(user_id, experiment_id) mod 100) to pick a variant and attaches a header like X-Experiment-Variant. The mesh — Istio, Seldon, or KServe — routes on that header. Sample size and duration are computed up front from the minimum detectable effect (MDE) and statistical power, and analysis uses two-sample tests or Bayesian methods depending on team convention.

A/B for ML has wrinkles unique to the domain. For ranking and recommendation tasks, evaluate at the list level (NDCG, MAP) rather than per item. Watch for leakage — shared caches or feature stores that contaminate one variant with another’s predictions. And recognize that for systems with exploration policies (contextual bandits, RL), the i.i.d. assumption underlying classical A/B may not hold; multi-armed bandit algorithms — epsilon-greedy, Thompson sampling, upper-confidence-bound — dynamically shift traffic toward the better-performing variant while still exploring, which is more sample-efficient but harder to analyze.

Key Takeaway: Canary protects against regression; A/B and bandits quantify improvement. Use canary first to confirm safety, then run a proper experiment with sticky assignment and pre-registered metrics before declaring victory.

Blue-Green and Rollback

Blue-green deployment maintains two complete environments — blue (current production) and green (new version) — and switches all traffic at once when ready. Green is built in full: new model, updated feature transformations, possibly new candidate generators and post-processing services. While blue continues to serve, green runs in parallel taking shadow or small canary traffic for validation. When confidence is sufficient, a single configuration change at the gateway or service mesh — flipping an Istio DestinationRule, updating a load balancer target — sends 100% of traffic to green. Rollback is the same single change in reverse [Source: https://onlinelibrary.wiley.com/doi/pdf/10.1002/9781119480280.app].

Blue-green is the right tool for major changes that go together: a new architecture, a new feature store schema, a redesigned ranking stack. It is the safest path when you want a hard cutover with a clean revert. The ML-specific challenges center on state and schemas. If the model updates online (incremental learning), blue and green diverge, and you must plan how state migrates. If green depends on a new feature store schema (user_features_v2), the rollback path requires v1 still computable — usually achieved with versioned feature views and dual-write windows. Batch-fed features need both pipelines running in parallel before switchover.

Across all four strategies, ML rollback discipline goes beyond the routing layer. Keep the previous model “hot” — fully deployed and capable of taking full traffic. Version features and schemas explicitly. Log enough request, prediction, and metadata to recompute metrics offline. And wire guardrails — error-rate, latency, and proxy-metric thresholds — directly into the rollout controller so automatic reversion triggers without human delay.

Key Takeaway: Combine strategies in sequence — offline validation, then shadow, then small canary, then A/B at 50/50, then full blue-green cutover — keeping the previous version deployable at every stage.

Edge and Mobile Deployment

Edge deployment trades server-side flexibility for privacy, latency, and offline operation. The constraint set flips: instead of optimizing for throughput on a cluster of GPUs, you are optimizing for milliseconds and milliwatts on a phone CPU or microcontroller. Three compression techniques — quantization, pruning, and distillation — and four frameworks — TFLite, Core ML, ONNX Runtime, PyTorch Mobile — form the working vocabulary.

Edge Frameworks: TFLite, Core ML, ONNX Runtime, PyTorch Mobile

TensorFlow Lite dominates Android and microcontroller deployment. It offers mature post-training and quantization-aware training flows, full int8 and float16 support, and hardware delegates that route to NNAPI, GPU, or Edge TPUs [Source: https://dzone.com/articles/edge-ai-tensorflow-lite-vs-onnx-runtime-vs-pytorch]. TensorFlow Lite for Microcontrollers (TFLM) extends the stack to devices with kilobytes of RAM. The pain point is the conversion path: models trained in PyTorch or other frameworks must round-trip through ONNX or custom converters, sometimes losing ops along the way.

Core ML is Apple’s native runtime, the only way to fully exploit the Neural Engine on iPhones, iPads, Macs, and Apple Watches. The coremltools package converts from TensorFlow, PyTorch (via TorchScript or ONNX), and other sources, automatically partitioning work across the Neural Engine, GPU, and CPU based on op support and power profile [Source: http://www.ml-illustrated.com/2020/06/15/deploy-pytorch-sound-classification-model-via-coreml.html]. It is Apple-specific by design, so cross-platform teams usually keep an ONNX intermediate.

ONNX Runtime (ORT) is the cross-platform option. A single ONNX model can run on Android (via NNAPI or XNNPACK execution providers), iOS (via the Core ML execution provider), desktop, and server [Source: https://onnxruntime.ai/docs/execution-providers/CoreML-ExecutionProvider.html]. Quantization tooling supports both dynamic and static int8. The cost is a bit more glue code and careful attention to opset versions — quantization requires opset 10 or higher, and some advanced operators need execution-provider-specific handling [Source: https://devblogs.microsoft.com/xamarin/machine-learning-in-xamarin-forms-with-onnx-runtime/].

PyTorch Mobile and ExecuTorch target PyTorch-first teams that want minimal conversion friction. TorchScript models run directly on mobile, keeping training and inference code aligned. ExecuTorch — PyTorch’s newer on-device runtime — targets mobile and embedded with more optimized backends including GPU and NPU acceleration on Android. Historically PyTorch Mobile has been less lean than TFLite for tiny devices, and many production teams still convert to TFLite or ONNX for final optimization [Source: https://huggingface.co/blog/tugrulkaya/running-large-transformer-models-on-mobile].

INT8 and INT4 Quantization Plus Pruning

Quantization reduces numeric precision — typically from float32 to int8 — shrinking models roughly 4× and accelerating inference on hardware with integer kernels (NNAPI, the Apple Neural Engine, Edge TPUs, XNNPACK) [Source: https://fs-eire.github.io/onnxruntime/docs/execution-providers/CoreML-ExecutionProvider.html]. Two flavors matter in practice. Post-Training Quantization (PTQ) quantizes an already-trained float model using a small calibration set; it requires no retraining, takes minutes, and typically delivers 2–4× size and latency wins. Its weakness is accuracy: small models, transformers, and highly non-linear architectures can lose meaningful accuracy under aggressive PTQ. Quantization-Aware Training (QAT) inserts quantization stubs into the graph during training so the model learns weights that are robust to int8 arithmetic. QAT preserves accuracy at lower bit widths but requires a training pipeline, additional engineering, and longer iteration time.

INT4 and mixed-precision quantization push further — 8× size reduction over float32 — and have become essential for running large language models on phones, often combined with techniques like GPTQ or AWQ. The recommended discipline is to start with PTQ, measure the accuracy gap on task-specific and fairness metrics (not just top-1), and escalate to QAT only when PTQ falls short.

Pruning is complementary. Magnitude-based (unstructured) pruning zeroes out small weights, reducing parameter count and model size but rarely speeding up inference on mobile because most edge runtimes lack optimized sparse kernels. Structured pruning removes whole channels, attention heads, or transformer blocks, directly shrinking the computational graph and delivering real latency wins — at the cost of larger accuracy risk and an architecture-change retraining cycle.

Knowledge Distillation

Knowledge distillation trains a small student model to mimic a large teacher. The student learns from the teacher’s soft probability outputs (logits) — often with a temperature parameter to smooth the distribution — combined with the original hard labels. Distillation often beats direct compression of the large model because the student can be designed from scratch to fit the device budget rather than retrofitted from a too-big architecture. It is the dominant compression path for shipping transformers to mobile: TinyBERT, DistilBERT, MobileBERT, and similar distilled models power on-device search, autocorrect, and voice. Once a distilled student is trained, you can apply pruning and quantization on top for additional gains.

OTA Model Updates

Models shipped on devices age quickly. Drift, new content, new attack patterns, and bug fixes all demand updates without forcing users to download a new app. Over-the-air (OTA) model updates decouple the model artifact from the application binary: the app downloads the latest model from a CDN or model server, verifies signatures, swaps it in atomically, and falls back to the previous version if the new one fails health checks. Best practices include staged rollouts (small device percentage first), differential updates to minimize bandwidth, A/B comparison between model versions on the device, and clear telemetry — inference latency, prediction distributions, occasional sampled outputs — flowing back to detect regressions in the wild. The recommended compression order for any edge model is: start from a mobile-architected baseline (MobileNet, EfficientNet-Lite, distilled transformers), apply structured pruning, distill if needed, then PTQ first and QAT if accuracy demands. Always profile on the actual target device with the final framework — desktop benchmarks systematically mislead.

Figure 11.4: Edge deployment pipeline from trained model to on-device OTA delivery

flowchart LR
    T[Trained Float32 Model] --> D[Knowledge Distillation<br/>Teacher to Student]
    D --> P[Structured Pruning<br/>Remove Channels/Heads]
    P --> Q{Quantization}
    Q -->|PTQ first| Q1[INT8 / INT4 Weights]
    Q -->|QAT if accuracy gap| Q2[Quantization-Aware Trained]
    Q1 --> CV[Framework Conversion]
    Q2 --> CV
    CV --> A[TFLite<br/>Android/MCU]
    CV --> B[Core ML<br/>Apple Neural Engine]
    CV --> C[ONNX Runtime<br/>Cross-platform]
    A --> CDN[(Signed Model CDN)]
    B --> CDN
    C --> CDN
    CDN --> OTA[OTA Staged Rollout<br/>1% to 100%]
    OTA --> DEV[On-Device<br/>Atomic Swap + Fallback]
    DEV --> TEL[Telemetry: latency,<br/>distributions, energy]
    TEL -.->|Drift signal| T

Key Takeaway: Edge deployment is a compression problem: distill, prune, then quantize against the target device’s actual hardware accelerators, ship updates OTA with signed staged rollouts, and measure energy per task, not just per-inference latency.

Serving Platforms

The serving platform is the substrate that turns a model artifact into a live prediction service: handling autoscaling, traffic splitting, monitoring, multi-model packing, and the integration glue between a model registry and the user. The landscape splits into four broad categories — self-hosted, managed, serverless, and multi-model/multi-tenant — with very different operational profiles.

Self-Hosted: KServe, BentoML, Seldon

Self-hosted serving runs on your Kubernetes cluster (or equivalent), giving you full control over hardware, configuration, and routing. Three platforms dominate.

KServe (formerly KFServing) is the Kubernetes-native, serverless-style ML serving framework. Its central abstraction is the InferenceService CRD, which packages predictor, transformer, and explainer components, supports the Open Inference Protocol across many backends (TF Serving, TorchServe, Triton, scikit-learn), and integrates with Knative for scale-to-zero and Istio for routing. The canaryTraffic field on InferenceService makes a percentage-based canary a one-line change. KServe is the default choice for teams already invested in Kubernetes and Knative.

BentoML focuses on the developer-experience end of the stack. You declare a service in Python using Pythonic decorators, package the model and its dependencies into a “bento” (a reproducible artifact), and deploy that bento to Docker, Kubernetes, or BentoML’s own cloud. Strengths include first-class Python serving, easy multi-model composition, and an opinionated workflow that shortens the path from notebook to production. The trade-off is less granular Kubernetes-native control compared to KServe.

Seldon Core is the ML-aware Kubernetes platform. Its SeldonDeployment CRD natively understands predictors, shadow predictors, A/B routing, ensembles, explainers, and outlier detectors as first-class concepts, integrating tightly with Istio. Seldon is the most ML-feature-rich of the three, with built-in support for the rollout strategies covered earlier in this chapter, but its surface area is larger and operating it well requires familiarity with both Kubernetes and ML serving concerns.

Managed: SageMaker, Vertex AI, Azure ML

Managed serving platforms — Amazon SageMaker, Google Vertex AI, and Azure ML — handle the infrastructure layer for you. You upload a model artifact, declare an endpoint configuration, and the cloud provider runs the autoscaling, load balancing, health checks, and (often) the canary deployment. SageMaker endpoints support multi-model endpoints, serverless inference, and asynchronous inference modes. Vertex AI offers similar features with tight integration into Google’s data stack. Azure ML provides managed online endpoints with built-in blue-green and traffic-split semantics.

The appeal of managed serving is operational simplification: you trade flexibility for not running Kubernetes. The drawbacks are vendor lock-in, sometimes opaque pricing at high QPS, and limits on customizing the request path (custom preprocessing, complex routing). Managed services are often the right starting point for small teams that want to ship quickly and migrate to self-hosted later if scale or cost demands it.

Serverless: Lambda, Cloud Run, Functions

Serverless serving — AWS Lambda, Google Cloud Run, Azure Functions — provisions compute per request, charges by execution time, and scales to zero when idle. For low-QPS or spiky workloads it is dramatically cheaper than always-on serving. The constraints are model size limits (a few hundred MB for Lambda containers), cold-start latency (seconds for large models loading from scratch), no native GPU support on most platforms (Cloud Run supports GPUs in limited regions), and a request-response model that may not suit batched or streaming inference. Cloud Run is the most ML-friendly of the three because it supports containers up to several GB and offers concurrency-per-instance, letting one container handle many concurrent requests and amortize the cost of loading a model.

Multi-Model and Multi-Tenant Serving

Most teams eventually outgrow the “one model per pod” model. A team that ships 200 personalized recommender variants — one per merchant, country, or experiment — cannot afford one pod per variant. Multi-model serving packs many models into a single serving process and loads them on demand: SageMaker Multi-Model Endpoints, Triton’s model repository with on-demand loading, and BentoML’s multi-runner all support this pattern. The trade-offs are cache management (which models stay warm?), per-request memory pressure when a cold model loads, and noisier-neighbor effects when one model spikes resource use.

Multi-tenant serving generalizes the same idea across users or organizations sharing a serving infrastructure, with careful attention to isolation, quotas, and authentication. SaaS ML products almost always need multi-tenant patterns; pricing per inference and isolation guarantees become first-class concerns alongside the usual latency and throughput.

Key Takeaway: Match the serving platform to scale and team — managed for fast starts, self-hosted (KServe/BentoML/Seldon) for control and ML-aware features, serverless for low-QPS workloads, and multi-model when you ship many variants and per-model dedicated capacity becomes uneconomic.

Chapter Summary

Deployment is where ML systems meet reality, and the choice of pattern shapes everything that follows. Batch inference, the cheapest and simplest pattern, fits anywhere downstream tolerates minutes-to-hours staleness — nightly CLV scoring, precomputed recommendations, backfills. Online inference is the only choice when a human waits on the response, paying for always-on capacity and strict latency engineering in exchange for synchronous predictions in fraud, search, and ranking. Streaming inference threads the middle, processing event flows in near-real-time for trending detection, live moderation, and feature freshening, at the cost of significant operational complexity. Embedded inference moves the model into the user’s device, unlocking privacy, sub-10ms latency, and offline operation while forcing a compression-and-OTA discipline that server-side patterns avoid.

Rollout strategy matters as much as inference pattern. Shadow mode mirrors production traffic to validate candidates without user exposure. Canary releases ramp traffic gradually with automated rollback to catch regressions. A/B tests measure genuine business uplift with sticky randomized assignment and pre-registered metrics. Blue-green swaps entire environments for major bundled changes. The mature pattern composes them: offline validation, shadow, small canary, A/B at 50/50, and blue-green cutover — with the old version kept hot at every stage. ML-specific monitoring across system, model-quality, and drift metrics is what makes any of this safe; silent model regressions don’t show up as 5xx errors.

Edge and mobile deployment is a compression problem. Quantization shrinks models 4× (int8) or 8× (int4) with PTQ as the starting point and QAT as the rescue for accuracy-sensitive cases. Structured pruning removes whole channels for real latency wins. Knowledge distillation produces small students that often beat compressed teachers. Framework choice maps to platform: TFLite for Android and microcontrollers, Core ML for Apple, ONNX Runtime for cross-platform, PyTorch Mobile and ExecuTorch for PyTorch-first teams. OTA model updates decouple model lifecycle from app releases.

Finally, the serving platform — self-hosted (KServe, BentoML, Seldon), managed (SageMaker, Vertex AI, Azure ML), serverless (Lambda, Cloud Run), or multi-model — is the substrate that ties everything together. The right choice depends on scale, team operational maturity, and how many models you ship. Start simple, measure relentlessly, and let cost and latency push you toward more sophisticated patterns only when the data demands it.

Key Terms

Chapter 12: Serving Infrastructure: Latency, Throughput, and Scalability

Once a model has been trained, validated, registered, and deployed behind an API, the real engineering challenge begins. Production serving is the discipline of delivering predictions fast enough, often enough, and cheaply enough to satisfy a service-level objective (SLO) — typically expressed as something like “99% of requests must return within 80 ms” or “the endpoint must sustain 5,000 queries per second at peak.” This chapter moves from the fundamentals of latency measurement, through the optimization techniques that compress per-request work, into the scaling strategies that horizontally expand capacity, and finally into the advanced topologies used to serve ensembles and large language models. Think of serving infrastructure as a freeway system: latency is how long any one car takes to reach its destination, throughput is how many cars per hour the freeway carries, and scaling is the difference between a single-lane country road and a twelve-lane expressway with on-ramps that materialize when traffic builds.

Section 1: Latency and Throughput Fundamentals

Tail Latency and Service-Level Objectives

The single most important habit a serving engineer must develop is to stop reasoning about average latency. Averages hide the catastrophes. If a model returns in 20 ms for 99 requests and 5,000 ms for the hundredth, the average looks like a healthy 70 ms — but one out of every hundred users just watched a loading spinner for five full seconds. Real serving systems are evaluated on percentiles: p50 (median, the typical experience), p95 (the bad-day experience), and p99 (the worst-case experience that still happens hundreds of times per hour at scale) [Source: https://erichorvitz.com/tail_answers.pdf].

A useful analogy: latency percentiles are like restaurant wait times. The median customer might wait 8 minutes, but if your p99 is 45 minutes, one in every hundred parties walks out furious — and they tell ten friends. A service-level objective (SLO) makes this concrete: “p99 < 200 ms for the recommendation endpoint” is a contract between the ML platform team and downstream consumers. When the SLO is breached, error budgets are burned, on-call engineers are paged, and the autoscaler should already be expanding capacity.

Tail latency originates from sources that average measurements simply cannot see: garbage collection pauses, cold caches, head-of-line blocking when a long request stalls others behind it, kernel scheduler jitter, and rare large input shapes that fall outside an optimization profile [Source: https://erichorvitz.com/tail_answers.pdf]. Engineering for p99 is fundamentally about reducing variance, not just reducing the mean.

Key Takeaway: Serving SLOs are written in terms of tail percentiles (p95, p99) because averages mask the rare-but-frequent bad experiences that define user perception; engineering for p99 means engineering for variance reduction, not just average speed.

Throughput Versus Latency Tradeoffs

Throughput (queries per second, QPS, or tokens per second for LLMs) and latency are tightly coupled but not the same thing. A model can have low latency at low load and still collapse under high concurrency, or it can have moderate single-request latency but excellent throughput because the hardware is being used efficiently. The relationship is governed by queueing theory: as utilization approaches 100%, queue depth grows nonlinearly and latency explodes. A widely cited operational heuristic is to keep GPU utilization under 60–70% under normal traffic; above that threshold, even small bursts of incoming requests cause queueing delays that inflate p99 dramatically [Source: https://www.youtube.com/watch?v=MqqKT6etxpQ].