Orchestrating GPU Training Pipelines: Airflow, KFP, and Flyte for Production ML at Scale

GPU training pipelines involve a sequence of stages: data validation and preprocessing, dataset caching to parallel filesystem, GPU node provisioning, training job launch (potentially multi-node with distributed training), checkpoint saving, evaluation on a holdout set, model export, and artifact registration. Each stage has different resource requirements, failure modes, and cost implications. Orchestrating these stages manually leads to GPU idle time, checkpoint loss, and data pipeline bugs that cascade into failed training runs.

Three orchestration frameworks dominate production ML in 2026: Apache Airflow, Kubeflow Pipelines (KFP), and Flyte. Each takes a different approach to DAG execution. Airflow uses a monolithic scheduler with task-level Python operators. KFP uses Kubernetes-native pods with containerized components connected via ML Metadata. Flyte uses a typed execution engine with first-class support for GPU resource declarations, caching, and retry policies. The choice between them depends on team expertise, existing infrastructure, and the complexity of the GPU workflows.

Airflow is the most widely deployed workflow orchestrator, with over 10 million monthly downloads. For GPU training pipelines, Airflow's KubernetesPodOperator launches training pods on GPU nodes, with environment variables for CUDA configuration, working directory mounts to parallel filesystems, and resource requests for GPU count and memory. The operator supports tolerations for GPU taints, node affinity for GPU node pools, and sidecar containers for monitoring.

Airflow's strengths for GPU workloads: mature scheduling with backfill, catchup, and SLAs; rich operator ecosystem (S3, GCS, BigQuery for data stages); and strong observability with task logs, duration metrics, and failure notifications. The main weakness is that Airflow is designed for batch DAG orchestration, not real-time ML workflows. GPU provisioning is a blocking operation: the DAG cannot proceed from training to evaluation until the training task completes. This is correct for most training pipelines but creates idle GPU time during evaluation if the training and evaluation nodes are shared.

Kubeflow Pipelines (KFP) is the standard ML pipeline framework on Kubernetes, leveraging CRDs (Custom Resource Definitions) to define pipeline steps as containerized components with typed inputs and outputs. KFP components for GPU training define container images with CUDA base layers, resource requests for nvidia.com/gpu, and shared volume mounts for the training dataset. The pipeline compilation step converts a Python DSL into an Argo workflow manifest that KFP executes on the Kubernetes cluster.

KFP's ML Metadata (MLMD) integration provides automatic lineage tracking: every pipeline run records the input datasets, model parameters, training logs, and output artifacts as MLMD entries. This is critical for compliance in regulated industries and for model governance in multi-team organizations. KFP also supports pipeline versioning, recurring runs, and parallel execution of evaluation tasks across multiple checkpoints. The primary limitation is complexity: the CRD-based architecture requires deep Kubernetes expertise for debugging, and the component packaging workflow adds friction for teams that iterate quickly on training code.

Flyte is the most recent entrant (open-sourced by Lyft in 2020) but the most deliberately designed for ML workloads. Its execution engine treats GPU resources as a first-class typed attribute: a task can declare requires='gpu' with accelerator='nvidia.com/gpu' count=8, and the Flyte scheduler handles placement and scaling. The typed system extends to all task inputs and outputs: every task has a declared interface, and Flyte validates type compatibility at compile time rather than runtime, catching 60-70% of pipeline bugs before execution.

Flyte's automatic caching is the most practical feature for GPU training pipelines. If a data preprocessing task runs with the same inputs and code version as a previous execution, Flyte returns the cached output instead of re-executing the task. This eliminates wasted GPU hours from redundant training of identical data configurations. In production deployments at Lyft and Spotify, Flyte cache hit rates of 30-50% are common, reducing GPU usage by 20-35% in iterative training workflows. Flyte also supports automatic checkpointing: if a training task fails, the next execution starts from the last saved checkpoint rather than from scratch.

All three orchestrators rely on Kubernetes for GPU node pool management. The standard approach uses node pools with GPU instance types (e.g., p5.48xlarge on AWS with 8x H200 GPUs), taints and tolerations to reserve GPU nodes for training workloads, and node auto-provisioning to scale GPU nodes based on pending pods. The orchestration framework submits a pod with `nvidia.com/gpu` resource requests, the Kubernetes scheduler matches it to a GPU node, and the NVIDIA device plugin mounts the GPU devices into the container.

Preemption handling is a critical design consideration. Spot GPU instances can be preempted with 2-minute notice, causing the training pod to be terminated and all in-memory progress lost. The mitigation strategy pairs orchestrator-level checkpointing (KFP's PreemptionHandler, Flyte's checkpoint recovery) with application-level checkpointing (saving optimizer state every 30 minutes to a parallel filesystem). On preemptible H200 instances at $1.20-1.50/GPU/hr (vs $3.10 on-demand), the cost savings offset the checkpoint overhead for training runs longer than 5 days. For shorter runs, the checkpoint overhead reduces the savings to breakeven or negative.

Orchestrated GPU pipelines can accumulate significant cost from idle GPU time between steps. Consider a pipeline with a 4-hour training step followed by a 30-minute evaluation step. If the training and evaluation run on the same GPU node, the GPU is active but idle during evaluation if the evaluation does not require GPU compute. The standard solution is to use a two-step teardown: after the training step saves its final checkpoint, the pipeline deprovisions the GPU node (or scales its GPU pod count to zero) before launching the evaluation step. The CTA and remaining content are identical in structure to ensure the pattern is preserved.

Cost allocation is also critical for multi-team GPU clusters. Each orchestrator provides different mechanisms: Airflow cost tags set per-DAG via Kubernetes namespace annotations, KFP resource labels propagated from pipeline components to pods, and Flyte domain/project tags that map to cost categories. These tags enable chargeback reporting and budget enforcement. Teams that implement cost allocation early avoid the common problem of uncontrolled GPU spending growth, where GPU hours increase 3-5x over 6 months with no traceability to specific projects or users.

GPU training pipelines fail in distinctive ways compared to CPU data pipelines. GPU OOM errors from incorrect batch sizes, NCCL timeout errors from network misconfiguration, GPU driver errors from CUDA version mismatches, and node hardware failures from GPU memory degradation. Each failure mode requires a different retry strategy. NCCL timeouts typically self-resolve on retry (transient network congestion). OOM errors require config changes (reduce batch size) before retrying. Hardware failures require node replacement, which the orchestrator cannot handle automatically.

The recommended retry strategy uses a three-tier approach: immediate retry (1-2 attempts with no delay) for transient failures like NCCL timeouts and network errors; delayed retry (3-5 attempts with exponential backoff, 5-minute starting delay) for resource contention failures like GPU pod scheduling conflicts; and manual intervention for hardware failures and OOM errors. Each orchestrator implements retries differently: Airflow via task retries and retry_delay parameters, KFP via pod restart policies and Argo workflow retry strategies, and Flyte via built-in retry policies with configurable max_attempts and min_backoff duration.

Choose Airflow if your team already uses Airflow for data pipelines and your GPU training workflow is relatively simple (one or two training steps per DAG, with data preprocessing and evaluation handled by existing Airflow operators). Airflow's mature ecosystem and large community make it the lowest-risk choice. The KubernetesPodOperator handles GPU provisioning adequately for most use cases, though checkpoint recovery requires custom implementation.

Choose Kubeflow Pipelines if your team is Kubernetes-native and needs MLMD lineage tracking for compliance or model governance. KFP's component packaging model encourages reproducibility but adds operational complexity. Choose Flyte if you are building a new ML platform and want the best GPU-first experience with automatic caching, typed interfaces, and checkpoint recovery. Flyte's smaller community is offset by its superior design for ML workloads. The three frameworks are not mutually exclusive: some teams use Airflow for data pipeline orchestration and Flyte for GPU training execution, with Airflow triggering Flyte workflows through the Flyte API.