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.

Use Airflow when the ML workflow is one part of a larger data platform and you need a single scheduler for ingestion, transformation, training, and reporting.

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:

  1. Pin the container image by digest, not by tag.
  2. Record the code commit SHA as a run parameter or environment variable.
  3. Write every artifact to a path that includes the run ID or snapshot date.
  4. Capture hyperparameters as explicit parameters, not defaults buried in code.
  5. Emit a run manifest (framework, version, inputs, outputs) to a metadata store.
A common failure mode is a pipeline that succeeds but is not reproducible: the training step reads 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.
Orchestration Model Comparison How each framework moves data and state between steps Airflow extract_features train_model validate_model Artifacts via S3/GCS No native step cache Manual lineage Kubeflow Pipelines prepare_features container · Dataset out train container · Model out evaluate Typed artifacts + metadata Step cache on input hash Container isolation Metaflow start self.raw_path set train self.model stored end Built-in datastore Run versioning + resume Local → cloud scaling Platform-wide scheduling Kubernetes-native ML Data-science-owned flows

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.