Orchestrating GPU-hungry ML Pipelines with Airflow and Kubernetes

ML pipelines have a split personality: they need the choreography of a classical conductor (data collection, preprocessing, training, evaluation), and the brute force of a rock band on tour (GPU-backed training, bursty loads, multi-tenant stages). Over the last two years the tooling on both sides has matured: Apache Airflow made a major leap toward ML-friendly workflows, and Kubernetes + NVIDIA tooling unlocked new ways to share and time-slice scarce GPU capacity. This article walks through the practical architecture and trade-offs when managing ML pipelines with Airflow running on Kubernetes—focusing on how Airflow 3 and modern GPU tooling change the calculus for multi-tenant ML workloads. (airflow.apache.org)

Why this is timely

Two shifts are especially relevant:

• Airflow’s recent 3.x line formalized features that make event-driven and non-calendar ML workflows easier to author and run—opening up patterns like inference DAGs and hyperparameter sweeps that don’t map neatly to fixed intervals. (airflow.apache.org)
• NVIDIA and the wider Kubernetes ecosystem have advanced GPU scheduling: device plugins, the GPU Operator, and time-slicing/oversubscription affordments let clusters host higher GPU utilization and more predictable sharing across teams. That changes how pipelines can be scheduled and isolated on a shared cluster. (developer.nvidia.com)

Combined, these let platform teams treat ML workloads more like elastic, orchestrated services instead of monolithic training jobs—but the new flexibility requires careful design.

The fundamental tension

Three forces collide in ML orchestration:

• GPU scarcity: GPUs are expensive and often the bottleneck for training and large-batch inference.
• Bursty workloads: Model training and hyperparam sweeps arrive in waves, sometimes simultaneously across teams.
• Orchestration load: A highly concurrent Airflow deployment pushes the scheduler, API server, and metadata DB (and proxies like PgBouncer) into heavy usage patterns that are different from typical ETL workloads.

The result: naive designs either underutilize GPUs (wasting budget) or overload control plane components (causing flakiness). There are documented cases of Airflow 3 deployments hitting API server pressure and connection-pooling limits when concurrency grows—highlighting the need to design both orchestration and infra in tandem. (github.com)

A practical architecture pattern (conceptual)

Think of the cluster as a campus with specialized buildings:

• A “control center” where Airflow runs its API server, scheduler, and metadata components.
• Multiple “workshop” node pools: CPU pools for lightweight tasks, GPU pools for training, and perhaps a separate inference pool optimized for latency and smaller GPUs.
• A device-management layer (NVIDIA GPU Operator + device plugin) that announces GPU resources and can enable time-slicing or MIG-style partitioning where available.

Airflow on Kubernetes typically uses either the KubernetesExecutor or the KubernetesPodOperator. Both offload heavy compute into ephemeral Kubernetes pods rather than making the Airflow scheduler run work directly—this is central to keeping the control plane responsive. (airflow.apache.org)

Key placement controls are used to steer pods:

• nodeSelector/nodeAffinity to pick GPU node pools,
• taints/tolerations to keep non-GPU workloads off GPU nodes,
• pod resource requests/limits (and GPU requests like nvidia.com/gpu) so the scheduler understands actual needs. (kubernetes.io)

How modern GPU features change trade-offs

Two capabilities matter a lot:

• GPU time-slicing / shared GPU resources: NVIDIA’s device plugin and GPU Operator now support configurations enabling fractional or time-sliced access in many environments. This means multiple smaller workloads can safely share a physical GPU, improving utilization for many ML workloads that don’t fully saturate a GPU all the time. (developer.nvidia.com)
• Node-level isolation via taints/affinity: Treating GPU pools as schedulable-but-protected real estate prevents noisy neighbors from stealing CPU/network resources on GPU nodes while still letting platform tooling schedule GPU-hungry pods predictably. Managed Kubernetes providers often taint GPU node pools by default. (docs.cloud.google.com)

Together, you get more flexible sharing (time-slicing) and stronger placement guarantees (taints/affinity). But the extra flexibility also requires awareness in orchestration—task sizing, retry budgets, and admission control are more important than ever.

Practical considerations and patterns

Below are operational patterns and the reasons behind them.

• Keep heavy compute out of Airflow processes. Use KubernetesPodOperator or spawn pods via the KubernetesExecutor so GPU work runs in its own container lifecycle and image. This reduces memory/CPU pressure on Airflow components and improves reproducibility. (airflow.apache.org)

• Use dedicated node pools for GPU tasks. Mark them with a taint like nvidia.com/gpu=present:NoSchedule and rely on tolerations in GPU-task pods. Managed K8s providers often set similar taints automatically for GPU pools. This isolates GPU nodes from generic services. (docs.cloud.google.com)

• Tune resource requests and avoid over-large XCom payloads. The Airflow metadata DB (and PgBouncer in front of it) is frequently the first bottleneck at scale—keeping metadata small and offloading heavy artifacts (models, datasets) to object storage avoids DB saturation. There are community reports of Airflow 3 at high concurrency requiring PgBouncer and API autoscaling attention. (airflow.staged.apache.org)

• Prepare for control-plane load. Large bursts of pod creation from KubernetesExecutor or many short-lived pods from PodOperator can create API server pressure; batching pod creation, limiting creation rate, or using a worker pod pool can help. Monitoring API server metrics and PgBouncer connection pools is useful for capacity planning. (github.com)

• Leverage GPU Operator features intentionally. Time-slicing is excellent for inference or many light training jobs, but some workloads require exclusive GPUs for peak throughput. Label nodes by GPU capability (MIG, A100/A30, T4) and use node affinity to match workload needs to GPU capabilities. (docs.nvidia.com)

Example: a KubernetesPodOperator snippet (conceptual)

Below is a compact DAG task that runs a training job in a pod and requests one GPU (note: adapt images/registry and environment for your environment):

from airflow import DAG
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
from datetime import datetime

with DAG("train_example", start_date=datetime(2025,1,1), schedule_interval=None) as dag:
    train = KubernetesPodOperator(
        task_id="train",
        name="train-job",
        image="my-registry/ml-train:latest",
        cmds=["python","train.py"],
        resources={
            "limit_gpu": 1,  # platform convention; actual key: "nvidia.com/gpu" in pod spec
            "request_memory": "16Gi",
            "request_cpu": "4"
        },
        # nodeSelector / tolerations example (real pod spec uses nodeSelector/affinity/tolerations fields)
        # full_pod_spec or pod_template_file are recommended for advanced configs
    )

A corresponding Pod YAML would include a toleration for the GPU node taint and an explicit nvidia.com/gpu resource request; adapting these is a common integration point between Airflow DAG code and platform node pool labels.

Trade-offs and “gotchas”

• Shared GPUs can improve utilization, but not all model workloads are friendly to preemption or time-slicing. Large, memory-bound training jobs often still prefer exclusive devices. Characterize workloads (memory, compute, burstiness) before oversubscribing. (developer.nvidia.com)

• Airflow control-plane scaling grows nonlinearly with concurrency. The Airflow API server, scheduler, and DB will need attention as task concurrency increases—some teams observed PgBouncer and API-server OOMs at very high concurrent task counts. Autoscaling the API server and configuring connection pooling are common mitigations. (github.com)

• Debugging transient failures becomes a multi-system exercise: Airflow logs, Kubernetes events, node-level GPU health (via DCGM/Prometheus), and DB pool metrics all matter. Build observability for each layer.

A balanced closing note

Airflow and Kubernetes together offer a powerful control plane + execution plane for ML pipelines—especially now that Airflow 3 brings more ML-friendly workflow primitives and the Kubernetes GPU ecosystem supports fractional sharing and robust device management. The combination unlocks better GPU utilization and more elastic ML pipelines, but it also converts resource contention into a cross-system problem. Treat orchestration design like sound mixing: a few well-placed filters and buffers (node pools, taints, autoscalers, PgBouncer) can keep the signal clean; too many instruments without a conductor can quickly muddy the performance.

For teams moving from ETL-first workloads into heavy ML usage, the pragmatic takeaway is to view Airflow as the conductor and Kubernetes as the stage crew—design placement and sharing policies at the infra level, keep heavy compute in ephemeral pods, and monitor both orchestration and device layers closely. The recent releases and device-plugin advances give more options than before; choosing the right mix depends on whether your workloads need exclusivity, predictability, or raw throughput. (airflow.apache.org)

References (selected)

• Apache Airflow 3 announcement and release notes. (airflow.apache.org)
• Airflow KubernetesExecutor and KubernetesPodOperator documentation. (airflow.apache.org)
• NVIDIA blog and GPU Operator docs on GPU time-slicing and device plugin improvements. (developer.nvidia.com)
• Kubernetes taints/tolerations and node affinity docs (node pool isolation best practices). (kubernetes.io)
• Community scaling discussions and guidance on API server/PgBouncer concerns with highly concurrent Airflow workloads. (github.com)