ML Pipeline Orchestration: Airflow, Kubeflow and Prefect Compared

Orchestrating an ML pipeline is not the same as scheduling ETL. In data engineering, a DAG that moves rows from A to B at 2 AM is usually deterministic. In ML, your pipeline has to handle data validation that fails non-deterministically, training jobs that run for 9 hours then OOM, models that need to be compared and conditionally promoted, and GPUs that are expensive to leave idle. I have migrated teams off cron scripts onto all three of the major orchestrators, and each time the trade-off was less about feature checklists and more about where you want to absorb complexity. This article compares Airflow, Kubeflow Pipelines, and Prefect on the dimensions that actually matter in production ML: authoring ergonomics, execution model, state and failure handling, and operational overhead.

Airflow's DAG-Centric Model: Proven Scheduling at the Cost of ML Flexibility

Airflow is the most mature option, and that maturity shows. Its scheduler is rock solid, its integration ecosystem is unmatched, and almost every data engineer already knows its mental model. A DAG is defined as a static graph at parse time, operators are instantiated, dependencies are declared with >>, and the scheduler figures out when to run what. For batch feature computation and daily retraining triggers, this is exactly what you want.

The friction appears when ML logic becomes dynamic. In my experience, the first sign that you are fighting Airflow is when you start using BranchPythonOperator or XCom to pass model metrics between tasks to decide whether to deploy. Airflow assumes tasks are largely independent and communicate via small metadata in its metadata database. Pushing a 500MB model artifact through XCom will break the database, and pulling branching logic into PythonOperators makes the DAG hard to test outside the scheduler. I have seen teams work around this by externalizing state to S3 and MLflow, which works but effectively means you are building your own orchestration layer on top of Airflow.

Where Airflow Fits Naturally in an ML Stack

Airflow excels when your ML pipeline is primarily a scheduled batch process that orchestrates other systems. For example, a nightly pipeline that runs Great Expectations checks, triggers a Spark job for feature generation, launches a training job on a separate Kubernetes cluster via KubernetesPodOperator, and then calls a model registry API to register the results. In that pattern, Airflow is the control plane, not the compute plane, and that separation keeps it stable. If your data team already runs Airflow for ETL, extending it to trigger ML jobs is often the lowest-risk path.

Authoring Patterns and Operational Realities

You will need to manage Airflow's two-phase execution: DAG parsing and task execution. DAG parsing happens every 30 seconds on the scheduler, so any heavy code at the top level of your DAG file (importing pandas, loading a Hugging Face tokenizer) will slow down the entire deployment. Keep DAG files lean and push heavy logic into operators. You will also need to choose an executor. CeleryExecutor with Redis/RabbitMQ is still common, but KubernetesExecutor or the newer CeleryKubernetesExecutor is far better for ML because each task can request its own CPU/memory/GPU profile. Upgrades are another consideration; I have found that pinning Airflow to a stable version and using the official Helm chart is the only sane way to run it. The PyTorch Documentation approach to versioned, reproducible environments is a good analogy for how to treat your Airflow images.

# airflow/dags/ml_training_dag.py
from airflow import DAG
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
from airflow.operators.python import BranchPythonOperator
from datetime import datetime, timedelta

def choose_deployment_branch(ti):
    # Pull metrics from prior task via XCom is small - just JSON, not artifacts
    metrics = ti.xcom_pull(task_ids="evaluate_model")
    return "register_model" if metrics["auc"] > 0.82 else "notify_no_promotion"

with DAG(
    dag_id="ml_retraining_pipeline",
    start_date=datetime(2024, 1, 1),
    schedule="@daily",
    catchup=False,
    default_args={"retries": 2, "retry_delay": timedelta(minutes=5)},
) as dag:
    
    validate = KubernetesPodOperator(
        task_id="validate_data",
        image="my-registry/ml-validation:1.4.0",
        cmds=["python", "validation/run_checks.py"],
        get_logs=True,
    )

    train = KubernetesPodOperator(
        task_id="train_model",
        image="my-registry/ml-training:1.4.0",
        cmds=["python", "train.py", "--config", "/configs/xgb.yaml"],
        container_resources={
            "request_memory": "16Gi",
            "request_cpu": "4",
            "limit_memory": "32Gi",
        },
    )

    evaluate = KubernetesPodOperator(
        task_id="evaluate_model",
        image="my-registry/ml-training:1.4.0",
        cmds=["python", "evaluate.py", "--output", "/airflow/xcom/return.json"],
        do_xcom_push=True,
    )

    branch = BranchPythonOperator(
        task_id="promotion_decision",
        python_callable=choose_deployment_branch,
    )

    validate >> train >> evaluate >> branch

Kubeflow Pipelines: Containerized Components and Kubernetes-Native Lineage

Kubeflow Pipelines (KFP) takes the opposite approach: everything is a container, everything runs on Kubernetes, and the pipeline is a first-class artifact with lineage. When I worked with a team that was already all-in on Kubernetes and needed to run distributed training with true GPU isolation, KFP was the obvious fit. Each component runs in its own pod with its own resource requests, artifacts are passed via typed inputs/outputs backed by MinIO/S3, and you get a visual graph in the Kubeflow UI that actually reflects the executed components, not just the intended DAG.

That isolation is powerful for reproducibility. You compile a pipeline from Python using the KFP SDK, and the compiler produces an Intermediate Representation (IR) YAML that is immutable. I have used this to enforce that the exact same pipeline IR that passed staging is what runs in production, which is critical for auditability. It also pairs naturally with artifact tracking; I typically wire KFP's artifact outputs directly into MLflow or the Kubeflow metadata store, which makes ML Model Versioning with MLflow and DVC: Tracking Experiments at Scale much more coherent compared to bolting versioning onto Airflow's XComs.

Component Authoring and the SDK Experience

The KFP developer experience has improved dramatically since V2. You can author lightweight Python components with the @dsl.component decorator for quick iteration, or build heavyweight containerized components from Docker images for training jobs that need CUDA, custom kernels, or large dependencies. My rule of thumb is: use lightweight components for orchestration logic (data slicing, metric comparison, conditional deployment) and full container components for anything that touches data at scale or requires a GPU. The compiler will handle input/output serialization, but you need to be explicit about types.

What catches teams off guard is that KFP is not just an orchestrator, it is a platform. You need a Kubernetes cluster, you need to manage Kubeflow's control plane (ml-pipeline API server, persistence agent, UI), and you need object storage configured correctly. If you do not already operate Kubernetes, the operational overhead is substantial. I have seen small teams spend more time debugging Istio and profile RBAC than building models. When it is already available, however, KFP gives you the cleanest story for elastic, GPU-aware scheduling.

# kfp_pipeline.py - KFP v2 style
from kfp import dsl, compiler

@dsl.component(
    base_image="python:3.10",
    packages_to_install=["pandas==2.1.0", "scikit-learn==1.3.0"]
)
def validate_dataset(
    input_path: str,
    output_report: dsl.Output[dsl.HTML],
    threshold: float = 0.05
):
    import pandas as pd
    df = pd.read_parquet(input_path)
    null_ratio = df.isnull().mean().max()
    drift_flag = null_ratio > threshold
    
    with open(output_report.path, "w") as f:
        f.write(f"<h1>Null ratio: {null_ratio:.4f}</h1>")
    
    if drift_flag:
        raise ValueError(f"Data quality check failed: null_ratio {null_ratio} > {threshold}")

@dsl.component(
    base_image="pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime"
)
def train_model(
    dataset_path: str,
    model_output: dsl.Output[dsl.Model],
    metrics_output: dsl.Output[dsl.Metrics]
):
    # Follow patterns from Hugging Face Documentation for training loops
    # https://huggingface.co/docs/transformers/training
    import torch
    # ... training logic using torch ...
    metrics_output.log_metric("auc", 0.847)
    # Model artifact is automatically uploaded to object storage

@dsl.pipeline(name="ml-training-kfp", description="Kubeflow training with validation")
def ml_pipeline(data_path: str = "s3://bucket/dataset/2024-10-01"):
    validate_task = validate_dataset(input_path=data_path)
    # KFP handles pod resource spec and caching
    train_task = train_model(dataset_path=data_path)
    train_task.after(validate_task)
    train_task.set_cpu_request("8").set_memory_request("32Gi").set_gpu_limit(1)

if __name__ == "__main__":
    compiler.Compiler().compile(ml_pipeline, "pipeline.yaml")

Prefect's Dynamic Flows: Pythonic Orchestration for Imperative ML Logic

Prefect 2.x is where I send teams that love Python and hate YAML. Unlike Airflow's static DAGs, Prefect flows are just Python functions decorated with @flow and @task. That sounds trivial, but it changes how you model ML workflows. You can use real if/else, for loops, and dynamic task mapping based on the output of previous tasks without creating a BranchOperator or a sensor. If your pipeline needs to train N models with different hyperparameters discovered at runtime, or conditionally skip evaluation based on a data quality score, Prefect expresses that naturally.

In my experience, Prefect is the fastest to get value from if you are starting from a set of notebooks or scripts. You can wrap an existing training script in a flow incrementally, add retries and caching, and get observability in Prefect Cloud or self-hosted server without rewriting everything into operators. It also handles the messy middle of ML orchestration well: stateful tasks, notifications, and result persistence. A task that fails due to a transient S3 read can retry with backoff, and results are persisted to storage so a rerun can resume from the last successful task.

Tasks, Flows, and Hybrid Execution

Prefect separates the workflow definition from the infrastructure. A flow defines what to run, and infrastructure blocks define where to run it. You can develop locally with a Process infrastructure and then deploy the same flow to KubernetesJob, ECS, or DockerContainer without changing code. For ML, this is critical because your validation task might need 2 CPUs while your training task needs an A100. With Prefect, you attach a different task runner and infrastructure to each flow deployment.

One detail I have learned to enforce: use persist_result and result storage for any artifact you care about, and never rely on return values alone for large objects. Prefect will serialize task results to S3/GCS by default, which keeps the orchestration backend light. When paired with a dedicated feature layer, this pattern complements Feature Store Design: Online-Offline Architecture and Feature Serving because your feature materialization tasks can cache results and avoid recomputing features if inputs have not changed.

# prefect_flow.py
from prefect import flow, task, get_run_logger
from prefect.cache_policies import INPUTS
from prefect.task_runners import ConcurrentTaskRunner

@task(retries=3, retry_delay_seconds=[10, 30, 60], persist_result=True, cache_policy=INPUTS)
def validate_and_slice(path: str) -> str:
    logger = get_run_logger()
    import pandas as pd
    df = pd.read_parquet(path)
    if df["label"].isnull().mean() > 0.1:
        raise ValueError("Label null ratio exceeds threshold - failing task for retry/alert")
    output = path.replace("raw", "validated")
    df.to_parquet(output)
    logger.info(f"Validated {len(df)} rows to {output}")
    return output

@task(cache_policy=INPUTS)
def train_variant(dataset_path: str, params: dict) -> dict:
    import mlflow
    # mlflow tracking pairs well with Prefect's run context
    with mlflow.start_run(nested=True):
        # ... train model with params ...
        mlflow.log_params(params)
        auc = 0.84  # placeholder for actual evaluation
        mlflow.log_metric("auc", auc)
    return {"params": params, "auc": auc, "artifact": f"s3://models/{params['lr']}"}

@flow(task_runner=ConcurrentTaskRunner(), log_prints=True)
def ml_training_flow(raw_path: str = "s3://bucket/raw/2024-10-01"):
    validated = validate_and_slice(raw_path)
    
    # Dynamic mapping - impossible in static DAG without workarounds
    param_grid = [{"lr": 0.01, "depth": 6}, {"lr": 0.05, "depth": 8}, {"lr": 0.1, "depth": 10}]
    results = train_variant.map(dataset_path=validated, params=param_grid)
    
    # Conditional logic is just Python
    best = max(results, key=lambda r: r["auc"])
    if best["auc"] > 0.82:
        print(f"Promoting model {best['artifact']} with AUC {best['auc']}")

if __name__ == "__main__":
    ml_training_flow()

Failure Domains and State Recovery: How Each Orchestrator Handles Broken ML Runs

ML pipelines fail differently than data pipelines. A data pipeline failure is often binary: a table is missing, a schema changed. An ML failure can be partial: the data validation task passes but warns about drift, the training job completes but the metric is 5 points lower than last week, or the GPU node is preempted 6 hours in. How the orchestrator models failure determines how quickly you can recover without wasting compute.

Airflow is conservative and persistent. Task instance state is stored in Postgres, retries are configured per task, and a failed task can be cleared and rerun from the UI. I have found this reliable for long-running training jobs where you want to manually intervene, but Airflow's retry model is task-centric, not workflow-centric. If a downstream evaluation task fails after a 4-hour training task succeeded, you do not want to rerun training. Airflow handles this by not rerunning upstream tasks, but you must ensure tasks are idempotent because clearing a task does not clean up its side effects like partially written S3 objects.

Kubeflow pipelines rely on Kubernetes primitives for recovery. Each component pod has its own retry policy, and the pipeline run has an overall TTL and retry strategy. Argo Workflows, which underpins KFP, will restart failed pods based on the retryStrategy. This works well for transient infra failures like node eviction, but I have learned to add explicit artifact garbage collection. If a training component writes a checkpoint to /tmp and fails, the pod is ephemeral and that checkpoint is lost unless you mounted a persistent volume or wrote to object storage incrementally.

Prefect has the most granular state model for ML use cases. Tasks have states like Completed, Failed, Cached, Crashed, and Cancelled, and flows can be resumed from a failed state with persisted results. In practice this means a failed hyperparameter sweep that crashed on variant 7 of 10 can resume from variant 7 instead of starting over. For teams running expensive sweeps or frequent retraining triggered by Model Monitoring in Production: Data Drift, Concept Drift and Alert Systems, that resume behavior directly saves budget and time.

Practical Recovery Patterns I Enforce

Regardless of orchestrator, I enforce three rules. First, make every task idempotent and write outputs to a deterministic, content-addressed path like s3://bucket/runs/{run_id}/model.pkl so reruns overwrite cleanly. Second, persist evaluation metrics and manifests outside the orchestrator metadata store, usually in MLflow or a database, so promotion logic does not depend on XCom or in-memory state. Third, set explicit timeouts and heartbeats for training tasks. A training job stuck at 99% GPU that never exits is worse than one that fails fast, because it blocks your quota. Prefect's task timeouts and Airflow's execution_timeout plus KFP's activeDeadlineSeconds all address this, but you must configure them.

Selecting Your Orchestrator: A Context-Driven Comparison for MLOps Teams

There is no universal winner. The right choice depends on team topology, existing infrastructure, and what your pipeline spends most of its time doing. If you are a platform team operating Kubernetes at scale with many ML engineers needing isolation and lineage, Kubeflow's container-native model pays off. If you are a data team with a mature Airflow deployment that now needs to trigger ML jobs, extending Airflow avoids new operational burden. If you are an ML-centric team iterating quickly on Python-defined workloads with dynamic branching and frequent rescheduling, Prefect reduces friction substantially.

In my experience, the migration hotspots are predictable. Teams on Airflow hit complexity when they need dynamic task generation based on data (e.g., training one model per customer segment discovered at runtime). Teams on Kubeflow struggle when they lack Kubernetes expertise to debug pipeline controller crashes or PVC binding issues. Teams on Prefect encounter gaps when they need Airflow's vast provider ecosystem or Kubeflow's strict artifact lineage guarantees for regulated workloads.

A practical way to decide is to run a one-week proof of concept with the same pipeline on two candidates and measure: time to first successful run, time to add a conditional promotion step, time to recover from an injected failure (kill a training pod mid-run), and operational tasks required to upgrade the platform. The table below distills the core trade-offs I use when advising teams.

Whatever you choose, keep the orchestrator thin. The orchestrator should schedule, retry, and observe, not contain your business logic. Put data validation in a shared library tested with Testing ML Systems: Data Tests, Model Tests and Integration Testing, put training in containerized jobs with pinned dependencies, and let the orchestrator decide when and where they run. That separation is what makes pipelines maintainable a year later when the team that wrote them has changed.

Frequently Asked Questions