2. Pipeline orchestration frameworks
A reproducible ML workflow is not just a script that trains a model. It is a directed acyclic graph (DAG) of tasks with explicit dependencies, versioned inputs, deterministic environments, and observable state at every step. Orchestration frameworks exist to make that graph executable, retryable, and auditable. This section compares the three frameworks most commonly used in production MLOps stacks — Apache Airflow, Kubeflow Pipelines (KFP), and Metaflow — and shows how to reason about which one fits a given workflow.
What Orchestration Actually Provides
Before comparing tools, it helps to separate the responsibilities an orchestrator owns from those owned by the surrounding platform. An orchestrator schedules tasks, resolves dependencies, manages retries and backoff, tracks run state, and passes artifacts or metadata between steps. It does not inherently provide compute isolation, container image building, or model registries — those come from the execution backend (Kubernetes, a batch scheduler, or a managed service) and adjacent systems.
For ML specifically, three properties matter more than in generic data engineering:
- Artifact lineage. A trained model must be traceable to the exact dataset snapshot, code commit, and hyperparameters that produced it.
- Environment reproducibility. Training and inference must run against pinned dependency versions, not whatever happens to be installed.
- Step-level caching. Expensive steps (feature engineering, embedding generation) should be skipped when their inputs are unchanged.
Airflow: General-Purpose DAG Scheduling
Airflow models workflows as Python-defined DAGs. Tasks are typically operators, and dependencies are declared with the bitshift syntax. It is the most widely deployed orchestrator and integrates with nearly every data system.
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
with DAG(
dag_id="churn_training",
start_date=datetime(2024, 1, 1),
schedule="@daily",
catchup=False,
tags=["ml", "churn"],
) as dag:
extract = PythonOperator(
task_id="extract_features",
python_callable=extract_features,
op_kwargs={"snapshot_date": "{{ ds }}"},
)
train = PythonOperator(
task_id="train_model",
python_callable=train_model,
op_kwargs={"run_date": "{{ ds }}"},
)
validate = PythonOperator(
task_id="validate_model",
python_callable=validate_model,
)
extract >> train >> validate
Airflow's strength is breadth: it can orchestrate Spark jobs, dbt models, Kubernetes pods, and ML training in one graph. Its weakness for ML is that artifact passing is not first-class. Data moves between tasks through external storage (S3, GCS, a database), and the DAG author is responsible for making those paths deterministic. Airflow also does not natively cache task outputs; re-running a DAG re-executes every task unless you build idempotency into each callable.
Kubeflow Pipelines: Container-Native ML Workflows
KFP compiles a Python pipeline definition into an Argo Workflows manifest that runs on Kubernetes. Every step is a container, which makes environment isolation explicit and reproducible by construction. Artifacts are typed and tracked by the KFP metadata store.
from kfp import dsl
from kfp.dsl import Input, Output, Dataset, Model
@dsl.component(base_image="python:3.11-slim", packages_to_install=["pandas", "scikit-learn"])
def train(features: Input[Dataset], model: Output[Model], n_estimators: int = 200):
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
import joblib
df = pd.read_parquet(features.path)
X, y = df.drop(columns=["label"]), df["label"]
clf = RandomForestClassifier(n_estimators=n_estimators, random_state=42)
clf.fit(X, y)
joblib.dump(clf, model.path)
@dsl.pipeline(name="churn-training", description="Train and register churn model")
def churn_pipeline(snapshot: str = "2024-06-01", n_estimators: int = 200):
prep = prepare_features(snapshot=snapshot)
trained = train(features=prep.outputs["features"], n_estimators=n_estimators)
evaluate(model=trained.outputs["model"])
Because each component declares its inputs and outputs, KFP can cache a step's result keyed on the component image digest plus input artifact hashes. If prepare_features is unchanged and the snapshot is the same, the training step reuses the cached feature artifact. This is the property Airflow lacks and the main reason ML teams adopt KFP.
The cost is operational: you need a Kubernetes cluster, and local iteration requires either a KFP cluster or a compatible local runner. Debugging a failed step means inspecting pod logs and the metadata store rather than a single scheduler log.
Metaflow: Human-Centric ML Workflows
Metaflow, originally developed at Netflix, treats the flow as a Python class with decorated steps. It emphasizes local development and seamless scaling to remote compute without rewriting the flow.
from metaflow import FlowSpec, step, Parameter, card
class ChurnFlow(FlowSpec):
snapshot = Parameter("snapshot", default="2024-06-01")
n_estimators = Parameter("n_estimators", default=200)
@step
def start(self):
self.raw_path = f"s3://features/churn/{self.snapshot}.parquet"
self.next(self.train)
@step
def train(self):
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
df = pd.read_parquet(self.raw_path)
X, y = df.drop(columns=["label"]), df["label"]
self.model = RandomForestClassifier(
n_estimators=self.n_estimators, random_state=42
).fit(X, y)
self.next(self.end)
@step
def end(self):
print(f"trained on {self.snapshot}")
if __name__ == "__main__":
ChurnFlow()
Run locally with python churn_flow.py run, then scale with python churn_flow.py run --with kubernetes. Metaflow automatically versions every run, stores artifacts in its datastore, and exposes them through a client API (Flow("ChurnFlow").latest_run.data.model). This makes lineage and reproducibility nearly free for the flow author.
Choosing Between Them
| Dimension | Airflow | Kubeflow Pipelines | Metaflow |
|---|---|---|---|
| Execution unit | Operator / task | Container component | Python step |
| Artifact passing | External storage (manual) | Typed artifacts + metadata store | Built-in datastore + client API |
| Step caching | Not native | Yes, keyed on inputs | Yes, via resume/run versioning |
| Local iteration | Local executor | Requires cluster or local runner | First-class local execution |
| Best fit | Cross-system data platform | Kubernetes-native ML platform | Data-science-owned ML flows |
The decision is rarely about raw capability. Airflow wins when the ML pipeline must coexist with ingestion and BI jobs under one scheduler. KFP wins when the organization already runs Kubernetes and needs strict container isolation and artifact lineage. Metaflow wins when data scientists own the flow end-to-end and want local-to-cloud scaling without platform tickets.
Reproducibility Checklist
Whichever framework you choose, the following practices determine whether a run can actually be reproduced six months later:
- Pin the container image by digest, not by tag.
- Record the code commit SHA as a run parameter or environment variable.
- Write every artifact to a path that includes the run ID or snapshot date.
- Capture hyperparameters as explicit parameters, not defaults buried in code.
- Emit a run manifest (framework, version, inputs, outputs) to a metadata store.
latest.parquet instead of a dated snapshot, so re-running the DAG tomorrow produces a different model with the same run ID. Always bind inputs to immutable references.The diagram above contrasts the three models along the dimension that most affects reproducibility: how state and artifacts cross step boundaries. Airflow leaves that to the author; KFP and Metaflow make it a framework concern. That single difference explains most real-world framework choices in ML platforms.