Orchestrating ML pipelines with Airflow

Generated by Claude through GigAI 0.1.11.10 on 2026-10-11.

Concept: Airflow · also covers Workflow orchestration · Step 8 of 35 in MLOps / ML platform engineer · version 1

Practice this

On this page

Airflow models a pipeline as a DAG of tasks defined in Python, and a scheduler decides when each task runs based on a schedule, dependencies and state. For ML platforms it is usually the control plane that triggers heavy work elsewhere (Spark, Kubernetes pods, training services) rather than the place where the compute happens. Running it well depends on idempotent tasks, careful DAG authoring and an executor that fits the workload.

Read more

How it is done in practice

A typical ML platform team runs Airflow on Kubernetes from the official Helm chart or uses a managed offering, with DAGs delivered from a Git repository through CI. Data scientists get a template DAG with standard stages: validate inputs, build features, train, evaluate, register. Each heavy stage runs as a pod or an external job, and Airflow only tracks state and passes references such as a storage path or a model version. The team sets default retries, timeouts and failure callbacks through shared default arguments, uses pools to protect scarce resources such as GPUs and warehouse slots, and watches scheduler health, parse time and task queue latency. Retraining is scheduled by time or triggered by upstream data updates, and backfills are run with concurrency limits after feature logic changes.

Technologies

Apache AirflowTaskFlow APIKubernetesExecutorCeleryExecutorKubernetesPodOperatorAmazon MWAAGoogle Cloud ComposerAstronomerPostgreSQLHelm

Diagram

flowchart LR
  Git[DAG repo] --> DP[DAG processor]
  DP --> DB[(Metadata DB)]
  S[Scheduler] --> DB
  UI[Web UI and API] --> DB
  S --> E[Executor]
  E --> P1[Task pod: features]
  E --> P2[Task pod: train on GPU]
  E --> P3[Task pod: evaluate]
  P1 --> OS[(Object storage)]
  P2 --> OS
  P3 --> REG[Model registry]
flowchart LR
  Git[DAG repo] --> DP[DAG processor]
  DP --> DB[(Metadata DB)]
  S[Scheduler] --> DB
  UI[Web UI and API] --> DB
  S --> E[Executor]
  E --> P1[Task pod: features]
  E --> P2[Task pod: train on GPU]
  E --> P3[Task pod: evaluate]
  P1 --> OS[(Object storage)]
  P2 --> OS
  P3 --> REG[Model registry]

Code

python
from datetime import datetime, timedelta
from airflow.decorators import dag, task

@dag(schedule="@daily", start_date=datetime(2025, 1, 1), catchup=False,
     default_args={"retries": 2, "retry_delay": timedelta(minutes=5)})
def train_daily():
    @task
    def build_features(data_interval_start=None) -> str:
        # write to a partition keyed by the interval, overwrite on rerun
        return f"features/dt={data_interval_start:%Y-%m-%d}/"

    @task
    def train(path: str) -> str:
        print("training on", path)
        return "model-version-ref"

    train(build_features())

train_daily()

Subtopics

The DAG model: tasks, operators and dependencies

Core concepts

A DAG is a Python-defined collection of tasks with directed dependencies and no cycles. It describes order and conditions of execution, not what the tasks do internally. A task is an instance of an operator (a prebuilt template such as a Bash, Python or Kubernetes pod operator), a sensor (waits for an external condition) or a function decorated with the TaskFlow task decorator. Dependencies are declared with the bitshift operators or set_upstream and set_downstream, and TaskFlow infers them from function calls. Each scheduled execution creates a DAG run, and each task in it becomes a task instance with its own state: queued, running, success, failed, up for retry or skipped.

Control flow

Beyond linear chains, Airflow supports branching (pick one downstream path at runtime), trigger rules (run a task when all upstream succeeded, when any failed, or always, which is useful for cleanup and alerts), task groups for visual and logical grouping, and dynamic task mapping, which expands one task into many instances at runtime from a list, for example one training task per region or hyperparameter set. Small values pass between tasks through XCom; anything large should be written to object storage and passed by reference.

Scheduling, data intervals, catchup and backfill

How runs are timed

A DAG has a start date and a schedule, which can be a cron expression, a preset such as daily, a timedelta, a custom timetable, or a data-aware trigger that fires when an upstream dataset or asset is updated. Each run carries a logical date and a data interval describing the slice of data it is responsible for. A daily run typically processes the previous complete day, so the run for an interval starts after that interval ends. Tasks should read the interval from the runtime context or templates instead of calling the current time, so that a rerun months later processes the same slice.

Catchup and backfill

With catchup enabled, the scheduler creates a run for every interval between the start date and now that has not yet run; disabled, it only schedules the latest. Backfill deliberately reruns a historical range, for example after fixing a feature bug that requires recomputing ninety days of features. Both only work safely when tasks are partitioned by interval and idempotent. Limits such as max active runs and pools keep a backfill from flooding a warehouse or a GPU queue.

Idempotent tasks, retries and failure handling

Tasks as transactions

Airflow's guidance is to treat each task like a database transaction: it either completes fully or leaves nothing partial behind, and rerunning it gives the same result. In practice that means writing to a partition keyed by the data interval and overwriting it (or using upsert and merge) instead of appending, never relying on local disk between tasks because tasks may run on different workers, and avoiding non-deterministic inputs such as the current time. Retries are set per task with a count, a delay and optional exponential backoff. Timeouts, SLA-style alerts and on-failure callbacks turn a silent stall into a page or a chat message. For ML, training tasks should checkpoint to remote storage so a retry resumes instead of restarting.

Architecture and executors: running tasks on Kubernetes

Components

An Airflow deployment has a scheduler that decides what should run, a DAG processor that parses DAG files, a metadata database (usually Postgres) holding all state, a web UI and API server, and an executor that determines where tasks run. The local executor runs tasks as subprocesses beside the scheduler. The Celery executor distributes them to a fixed pool of workers through a broker. The Kubernetes executor launches one pod per task, giving isolation and scale to zero at the cost of pod startup latency.

The pattern for ML

Independent of the executor, the Kubernetes pod operator runs a task as a pod from any image with its own resource requests, node selectors and tolerations. This keeps ML dependencies and GPU requirements out of the Airflow image. Platform teams commonly keep Airflow workers light and delegate: submit a Spark job, start a training job on a cluster or managed service, then poll or use a deferrable operator, which releases the worker slot while waiting and resumes through the triggerer process.

DAG authoring practices and testing

Keep parsing cheap

DAG files are parsed repeatedly by the DAG processor, so top-level code must be fast. Do not query databases, call APIs or read Airflow Variables at module level, and move heavy imports inside task functions. Read Variables and Connections through templates or inside the task at runtime. Keep secrets in Connections or a secrets backend, not in code. Generate similar DAGs from configuration with care, since each generated DAG adds parse time and scheduler load.

Testing and delivery

A baseline CI test loads every DAG file and fails on import errors or cycles. Unit tests cover the Python callables in isolation, and a single DAG can be executed locally with the test command or dag.test for a debug run without a scheduler. DAGs ship through Git: synced into the deployment or baked into the image, promoted across staging and production, with the Airflow version and provider packages pinned so upgrades are deliberate.

Sources

Related lessons

Practice this
Feature stores and feature pipelines: offline/online consistencyML-native orchestrators: Kubeflow Pipelines, Flyte, Prefect and Metaflow