{"title":"ML Pipelines","description":"","content":"Getting a model into production, and keeping it there, runs as a chain of steps: extract data, validate it, compute features, train, evaluate, register, deploy, and monitor.\n\nIf those steps run as independent scripts, the failure modes are predictable: training starts before features finish, evaluation reads the wrong dataset, a model is deployed without its matching feature schema, or a retry writes a second conflicting artifact. A production ML pipeline exists to remove that ambiguity.\n\nThe pipeline should make execution order, inputs, outputs, checks, and lineage explicit.\n\n---\n\n# What an ML Pipeline Owns\n\nAn ML pipeline is usually represented as a directed acyclic graph. Each node is a step, and each edge is a dependency. The trigger that matters is data, not the clock: a step runs once its required inputs exist, pass their checks, and come from the expected upstream step, rather than simply because a scheduled time arrived.\n\nA data pipeline moves raw data into validated tables, features, and training snapshots. A training pipeline turns a dataset into a model artifact. An end-to-end ML pipeline connects both with evaluation, model registration, deployment, and rollback hooks.\n\n\n```mermaid\nflowchart TD\n A[Data
Ingestion]:::primary --> B[Data
Validation]:::teal\n B -->|Pass| C[Feature
Computation]:::primary\n B -->|Fail| X[Block Run
Alert Owner]:::red\n C --> D[Training
Snapshot]:::teal\n D --> E[Model
Training]:::orange\n E --> F[Offline
Evaluation]:::teal\n F -->|Pass| G[Model
Registry]:::primary\n F -->|Fail| Y[Keep Champion
Store Report]:::red\n G --> H[Shadow or
Canary Deploy]:::orange\n H -->|Pass| I[Promote
Production]:::green\n H -->|Fail| J[Rollback
Investigate]:::red\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef teal fill:#38d9a9,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n classDef red fill:#ff8787,stroke:#000,color:#000\n```\n\n\nThe deployment branch leans on two safety patterns. A shadow deploy runs the new model on live traffic but throws its predictions away, so you can compare them against the champion without affecting users. A canary deploy routes a small slice of real traffic, say 1 to 5 percent, to the new model and watches its metrics before widening the rollout.\n\nTwo gates matter most in interviews and in production:\n\n- **Data quality gate:** blocks training when the input data violates schema, freshness, completeness, or label-quality expectations.\n- **Model quality gate:** blocks deployment when the challenger fails to match the champion on primary metrics, critical slices, calibration, latency, or safety constraints.\n\nSkip these two gates and all the automation buys you is bad releases shipped on a schedule.\n\n\n| Stage | What the Pipeline Must Know | Common Failure |\n|---|---|---|\n| Ingestion | Expected partitions, source freshness, row counts | Training reads an incomplete daily partition |\n| Validation | Schema, constraints, drift checks, label checks | New enum value maps to unknown without an alert |\n| Feature computation | Feature version, point-in-time logic, materialization status | Training and serving use different null handling |\n| Training snapshot | Dataset version, time range, joins, exclusions | Snapshot contains future information or duplicate labels |\n| Training | Code version, config, container image, resources, seed policy | Retry creates a second model with the same version name |\n| Evaluation | Champion reference, metric thresholds, slices, test set version | Aggregate AUC improves while a key segment regresses |\n| Registry | Model artifact, metadata, aliases/tags, schema, lineage | Model is deployable but not traceable |\n| Deployment | Traffic split, rollback trigger, compatibility checks | Canary passes latency but breaks feature compatibility |\n\n\nIn an interview, the orchestrator you reach for matters far less than showing that every boundary between steps carries an explicit contract.\n\n---\n\n# Contracts Between Steps\n\nA pipeline step should declare what it consumes, what it produces, and what proves it succeeded. \"The task finished\" is not enough.\n\nFor example, data validation should not return only `passed=True`. It should produce a validation artifact: schema version, row count, missing-value rates, freshness, distribution summaries, failed checks, warnings, and the exact data snapshot it inspected. Downstream steps should consume that artifact rather than assuming the upstream table is safe.\n\nThe same applies to model training. A training step should output the model artifact, training config, feature schema, data version, evaluation dataset pointer, metrics, and environment metadata. Leave that metadata out and the model registry degrades into a folder of opaque files that nobody can trace, debug, or safely roll back.\n\n\n| Artifact | Producer | Consumer | Why It Matters |\n|---|---|---|---|\n| Validation report | Data validation | Feature computation, training | Proves the dataset passed the expected checks |\n| Training snapshot ID | Snapshot builder | Training, evaluation | Makes the run reproducible |\n| Feature schema | Feature computation | Training, serving, registry | Prevents serving a model with incompatible inputs |\n| Model artifact | Training | Evaluation, registry, deployment | The deployable output |\n| Evaluation report | Evaluation | Registry, deployment gate | Explains why the model was accepted or rejected |\n| Lineage metadata | Every step | Debugging, audit, rollback | Answers what produced the production model |\n\n\nTyped artifacts make these contracts enforceable. In Kubeflow Pipelines or TFX, artifacts are first-class objects. In Airflow, Dagster, Prefect, or Argo, teams often store large artifacts in object storage and pass references plus metadata through the orchestrator. The exact mechanism is less important than the discipline: every downstream step should know exactly which upstream output it is using.\n\n---\n\n# Choosing an Orchestrator\n\nML orchestration has no single default answer. Airflow is still common because many companies already use it for data workflows. Kubeflow Pipelines and Argo Workflows fit Kubernetes-native teams. Dagster and Prefect are common when teams want Python-native orchestration with stronger local development ergonomics. Managed platforms such as Vertex AI Pipelines, SageMaker Pipelines, and Azure ML Pipelines reduce operational burden at the cost of cloud coupling.\n\n\n| Tool Family | Strength | Trade-off | Good Fit |\n|---|---|---|---|\n| Airflow | Mature scheduling, large ecosystem, strong data workflow adoption | ML artifact handling and caching need extra work | Teams already running Airflow for data pipelines |\n| Dagster / Prefect | Python-native developer experience, asset-aware workflows, easier local iteration | Smaller installed base than Airflow in older data orgs | Small to mid-sized ML platform teams |\n| Kubeflow Pipelines | Containerized ML workflows, artifacts, metadata, Kubernetes execution | Kubernetes and KFP operational complexity | ML platform teams already standardized on Kubernetes |\n| Argo Workflows | General Kubernetes workflow engine, strong container isolation | ML metadata and registry integration are mostly custom | Platform teams that want a lower-level workflow primitive |\n| Managed cloud pipelines | Integrated storage, registry, IAM, compute, and deployment | Vendor coupling and less portability | Teams optimizing for speed over infrastructure control |\n\n\nContainer isolation is a real design point. A training step may need CUDA, PyTorch, and large system libraries. A validation step may need only pandas and a schema checker. Running each step in its own container prevents dependency conflicts and gives the scheduler clear CPU, memory, and GPU requirements.\n\nThe practical decision tree:\n\n\n```mermaid\nflowchart TD\n A[Pick Pipeline
Orchestrator]:::primary --> B{Existing
platform?}:::yellow\n B -->|Airflow already runs data jobs| C[Use Airflow
plus ML metadata]:::teal\n B -->|Kubernetes-first platform| D{Need ML-specific
pipeline UI?}:::yellow\n D -->|Yes| E[Kubeflow
Pipelines]:::orange\n D -->|No| F[Argo or
Dagster on K8s]:::teal\n B -->|Cloud ML platform standard| G[Managed
pipeline service]:::orange\n B -->|Small team, few models| H[Keep it simple:
scheduled jobs + registry]:::green\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef teal fill:#38d9a9,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n classDef yellow fill:#ffd43b,stroke:#000,color:#000\n```\n\n\nIn an interview, do not over-prescribe Kubeflow because it sounds \"MLOps-y.\" If the system has one weekly retrained model, Airflow or a managed training job may be enough. If the system has dozens of models, GPU workloads, experiment lineage, and frequent retraining, a more specialized platform becomes easier to justify.\n\n---\n\n# Pipeline Definition Example\n\nThe point of pipeline code is to make dependencies and artifacts explicit. A minimal Kubeflow-style pipeline looks like this:\n\n\n```python\nfrom kfp import dsl\n\n@dsl.component(base_image=\"python:3.11\", packages_to_install=[\"pandas\", \"pyarrow\"])\ndef validate_data(data_uri: str) -> str:\n import pandas as pd\n\n df = pd.read_parquet(data_uri)\n if len(df) < 100_000:\n raise ValueError(f\"Too few rows: {len(df)}\")\n if df[\"label\"].isna().mean() > 0.001:\n raise ValueError(\"Label missing rate is above threshold\")\n\n return data_uri\n\n@dsl.component(base_image=\"python:3.11\", packages_to_install=[\"pandas\", \"scikit-learn\", \"joblib\"])\ndef train_model(validated_data_uri: str) -> str:\n import joblib\n import pandas as pd\n from sklearn.ensemble import HistGradientBoostingClassifier\n\n df = pd.read_parquet(validated_data_uri)\n X = df.drop(columns=[\"label\"])\n y = df[\"label\"]\n\n model = HistGradientBoostingClassifier(max_iter=200).fit(X, y)\n model_uri = \"/tmp/model.joblib\"\n joblib.dump(model, model_uri)\n return model_uri\n\n@dsl.pipeline(name=\"fraud-training\")\ndef fraud_training_pipeline(data_uri: str):\n validated = validate_data(data_uri=data_uri)\n train_model(validated_data_uri=validated.output)\n```\n\n\nThe training step receives its input from the validation step. That creates a real dependency. If validation fails, training does not run.\n\nFor Airflow, the same principle applies even if artifacts are object-store paths and metadata records instead of typed KFP artifacts. The important part is that dependencies are data-aware, not just time-aware.\n\n---\n\n# Artifact Passing, Caching, and Idempotency\n\nLarge artifacts should not move through the orchestrator itself. Datasets, embeddings, and model checkpoints belong in object storage or a model registry. The pipeline should pass stable references and metadata.\n\n\n| Pattern | Use For | Watch Out For |\n|---|---|---|\n| Object storage URI | Large datasets, model weights, embeddings | Path naming, retention, permissions, cleanup |\n| Model registry ID | Deployable model artifacts | Missing data/schema lineage |\n| Experiment run ID | Metrics, params, plots, training diagnostics | Tracking server becomes an operational dependency |\n| Pipeline-native artifact | Small reports, configs, metadata | Portability across orchestrators |\n\n\nCaching saves real money, but only when the cache key captures the right inputs. A cache key should reflect the input data version, code version, parameters, container image, and relevant environment. If feature computation reads \"latest\" from a table and the cache key only sees the table name, the pipeline can reuse stale results.\n\nKubeflow Pipelines v2 turns on task caching by default, so you switch it off per task when a step is non-deterministic. Argo Workflows supports memoization with an explicitly configured cache key. Airflow has no built-in step cache, so teams add their own skip logic or external metadata. Whatever the tool, cache only deterministic steps. Do not cache steps that depend on wall-clock time, mutable \"latest\" paths, random sampling without a seed, or external services with changing behavior.\n\nIdempotency is the other non-negotiable property. Retrying a failed training step should not create two production candidates with the same version name or append duplicate rows into a training table. Use run IDs, atomic writes, and \"write to temp then commit\" patterns.\n\n---\n\n# Testing ML Pipelines\n\nPipeline tests should verify wiring and contracts. They should not try to prove model quality.\n\n\n| Test Level | What It Checks | Runtime | Example |\n|---|---|---|---|\n| Unit | A single transformation or validator | Seconds | Feature function returns expected schema and values |\n| Integration | Boundary between two steps | Minutes | Validation output can be consumed by feature computation |\n| End-to-end smoke | Full DAG on tiny data | 10-30 minutes | Pipeline trains a tiny model and writes all artifacts |\n| Release rehearsal | Production-like run without promotion | Hours | Full training and evaluation before major launch |\n\n\nA useful unit test for feature computation is small and boring:\n\n\n```python\nimport pandas as pd\n\nfrom pipeline.features import compute_user_features\n\ndef test_compute_user_features_schema():\n events = pd.DataFrame({\n \"user_id\": [1, 1, 2],\n \"event_type\": [\"view\", \"click\", \"view\"],\n \"duration_sec\": [30, 45, 60],\n })\n\n features = compute_user_features(events)\n\n assert set([\"user_id\", \"avg_duration_sec\", \"click_rate\"]).issubset(features.columns)\n assert features[\"click_rate\"].between(0, 1).all()\n```\n\n\nEnd-to-end tests should use a tiny deterministic dataset and a cheap model. Running the full production model in CI usually creates slow, flaky tests that teams learn to ignore. Use production-scale training for scheduled release rehearsals or pre-launch validation, not every pull request.\n\n---\n\n# Versioning and Lineage\n\nWhen a production model regresses, the instinct is to ask who changed the model. The more useful question is what changed anywhere in the system that produced it: code, data, features, config, or environment.\n\nYou need lineage from a prediction back to:\n\n- model version or alias,\n- training pipeline run ID,\n- pipeline code commit,\n- data snapshot or table version,\n- feature definitions and schema,\n- training configuration,\n- container image digest,\n- evaluation report,\n- deployment configuration.\n\n\n```mermaid\nflowchart TB\n A[Production
Prediction]:::primary --> B[Model Alias:
champion]:::primary\n B --> C[Model Version
v23]:::orange\n C --> D[Pipeline Run
#847]:::teal\n D --> E[Data Snapshot
2026-05-20]:::green\n D --> F[Git Commit
a3f7b2c]:::green\n D --> G[Feature Schema
v12]:::green\n D --> H[Container Digest
sha256:8d4e]:::green\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n classDef teal fill:#38d9a9,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n```\n\n\nRegistries increasingly use flexible aliases and tags rather than only fixed stages such as `Staging` and `Production`. The concept is the same: deployment should refer to an audited model version, and the registry should make it clear which version is the current champion, which versions are challengers, and why each version moved or did not move forward.\n\nThe minimum viable lineage rule: every registered model should carry the pipeline run ID, data version, feature schema version, git commit, and evaluation report link. That small set answers most debugging and interview follow-up questions.\n\n---\n\n# Quiz","pageType":"ml-system-design"}

ML Pipelines

Ashish Pratap Singh

Get Premium