Apache Airflow — ML Orchestration Skeleton

Interview Reference: ZenML-to-Airflow Translation

Goal: Demonstrate conceptual fluency in Airflow architecture and ML pipeline design. Sushant’s production stack is ZenML — use this document as the translation layer.

1. What Airflow Is (and Isn’t)

What Airflow is: Apache Airflow is an open-source workflow orchestration platform that lets you author, schedule, and monitor multi-step pipelines as DAGs (Directed Acyclic Graphs) of tasks — in pure Python. It does not move data or run compute; it tells other systems to do work and records what happened.

Airflow vs ZenML — analogous but not equivalent:

Are they analogous? At the structural level — yes. Both represent workflows as a directed graph of discrete units with dependency edges. At the purpose level — no. Airflow is a general orchestrator that happens to run ML pipelines; ZenML is an ML platform that happens to orchestrate. ZenML can even use Airflow as its backend execution engine.

What it solves: Multi-step ML workflows (load → clean → train → evaluate → promote) that need scheduling, dependency management, retry logic, and observability — things cron jobs and shell scripts cannot provide cleanly at scale.

2. Architecture — The Four Components

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

┌──────────────┐     parses DAGs      ┌──────────────────┐
│  Scheduler   │ ───────────────────▶ │  Metadata DB     │
│              │ ◀─── task states ─── │  (PostgreSQL)    │
└──────┬───────┘                      └──────────────────┘
       │ dispatches                          ▲
       ▼                                     │ reads
┌──────────────┐    executes tasks    ┌──────┴───────┐
│   Executor   │ ───────────────────▶ │  Worker(s)   │
└──────────────┘                      └──────────────┘
                                             ▲
┌──────────────┐    reads UI from DB          │
│  Web Server  │                      (runs actual code)
└──────────────┘

Key insight: The Scheduler never runs your code. Workers do. The Scheduler only decides when and in what order.

Kubernetes Primer — What an Interviewer Expects You to Know

Kubernetes (K8s) is a container orchestration platform — it manages the deployment, scaling, and lifecycle of containerised applications across a cluster of machines. Think of it as an operating system for a fleet of servers.

Core concepts:

Where Helm Charts come in:

Deploying Airflow (or any complex app) onto Kubernetes means writing dozens of YAML manifests — Deployments, Services, ConfigMaps, Secrets, PersistentVolumes. This is tedious and error-prone to manage manually.

Helm is the package manager for Kubernetes. A Helm Chart is a pre-packaged, parameterisable bundle of all the Kubernetes manifests needed to deploy an application.

1
2

Without Helm:  write 15+ YAML files manually → kubectl apply each one → repeat per environment
With Helm:     helm install airflow apache-airflow/airflow --set executor=KubernetesExecutor

Where Helm fits in the deployment chain:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

Docker image (your code)
      ↓ pushed to
ECR / DockerHub (image registry)
      ↓ referenced in
Helm Chart values.yaml
      ↓ deployed via
helm install / upgrade
      ↓ creates
Kubernetes resources (Pods, Services, ConfigMaps...)
      ↓ managed by
Kubernetes cluster

In practice: the official Apache Airflow Helm chart packages the Scheduler, Webserver, Workers, and PostgreSQL metadata DB as one deployable unit. Your team only needs to override values.yaml — executor type, image tag, resource requests — and helm upgrade handles the rest. No manual pod management.

How Kubernetes relates to Airflow:

With KubernetesExecutor, every Airflow task runs in its own dedicated pod — created when the task starts, deleted when it finishes. This means:

• A heavy training task gets a 16GB RAM pod; a lightweight logging task gets 512MB — no contention.
• Failed pods don’t affect other tasks — full isolation.
• The cluster auto-scales: if 20 tasks are queued, Kubernetes spins up 20 pods in parallel (subject to node capacity).

The interview framing:

“Kubernetes is the infrastructure layer beneath the orchestration layer. Airflow decides what to run and when; Kubernetes decides where to run it and ensures it gets the right resources. In a production ML platform, Airflow’s KubernetesExecutor bridges the two — each task becomes a pod spec, and Kubernetes handles placement, resource allocation, and cleanup.”

3. DAG — The Core Concept

A DAG is a Python file that defines a directed, acyclic graph of tasks. No loops — tasks flow forward only.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

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

@dag(
    dag_id="kpi_model_retraining",
    schedule="0 3 * * *",        # cron: 3am daily
    start_date=datetime(2024, 1, 1),
    catchup=False,                # ← critical in production
    max_active_runs=1,            # ← prevent concurrent training runs
    default_args={
        "retries": 2,
        "retry_delay": timedelta(minutes=5),
        "email_on_failure": True,
    },
)
def kpi_retraining():
    ...

Three parameters that matter most

catchup=False If you deploy a DAG with start_date 30 days ago and catchup=True (default), Airflow queues 30 simultaneous backfill runs — flooding your workers. Always set catchup=False unless backfilling is intentional.

max_active_runs=1 Prevents a second scheduled run from starting while the previous one is still running. Essential for training pipelines — you never want two jobs simultaneously writing to the same model registry entry.

schedule and logical date Airflow runs at the end of an interval, not the start. A daily DAG with schedule="@daily" and start_date=2024-01-01 triggers its first run on 2024-01-02 for the logical date 2024-01-01. Use {{ ds }} in templates to get data_interval_start as YYYY-MM-DD.

4. Tasks — The Unit of Work

A task = one step in the pipeline. Three ways to define one:

TaskFlow API (preferred — Airflow 2.x)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

from airflow.decorators import task

@task
def load_data(ds=None) -> str:           # ds injected from Airflow context
    path = fetch_from_vertica(ds)
    save_to_s3(path)
    return path                          # return value auto-pushed to XCom

@task
def train_model(data_path: str) -> dict: # argument = auto-pull from XCom
    df = pd.read_parquet(data_path)
    metrics = fit_and_save(df)
    return metrics

Classic Operator (for non-Python work)

1
2
3
4
5
6

from airflow.operators.bash import BashOperator

spark_job = BashOperator(
    task_id="run_spark_features",
    bash_command="spark-submit /jobs/features.py --date {{ ds }}",
)

Sensor (waits for external condition)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor

wait_for_data = S3KeySensor(
    task_id="wait_for_snapshot",
    bucket_name="data-lake",
    bucket_key="snapshots/{{ ds_nodash }}/",
    poke_interval=300,
    timeout=7200,
    mode="reschedule",    # release worker slot between polls — always use for long waits
)

Wiring dependencies

1
2
3
4
5
6
7
8

# Bitshift operators
extract >> transform >> train >> evaluate

# Fan-out (parallel)
train >> [evaluate_train, evaluate_valid]

# Fan-in (merge)
[load_table_a, load_table_b] >> join_step

5. XCom — Passing Data Between Tasks

XCom (Cross-Communication) stores values in the metadata DB. It is for metadata, not data.

1

Rule: XCom the PATH to data, never the data itself.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

@task
def train(data_path: str) -> dict:
    model_path = "s3://models/model_2024_01_01.pkl"
    fit_and_save(data_path, model_path)
    return {"model_path": model_path, "rmse": 0.12}   # ✓ small dict

# NOT this:
@task
def train_bad(df: pd.DataFrame) -> pd.DataFrame:      # ✗ DataFrames don't belong in XCom
    return df.transform(...)

Why: XCom is stored in PostgreSQL. A 100MB DataFrame serialised there kills your metadata DB. The pattern is: write large data to S3/GCS, XCom the path, downstream task reads from the path.

6. Complete ML Pipeline Skeleton

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

from airflow.decorators import dag, task
from airflow.operators.python import ShortCircuitOperator, BranchPythonOperator
from airflow.operators.empty import EmptyOperator
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from datetime import datetime, timedelta

@dag(
    dag_id="kpi_model_daily_retraining",
    schedule="0 3 * * *",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    max_active_runs=1,
    default_args={"retries": 2, "retry_delay": timedelta(minutes=5)},
)
def kpi_retraining():

    # 1. Wait for upstream data (sensor pattern)
    wait = S3KeySensor(
        task_id="wait_for_snapshot",
        bucket_name="data-lake",
        bucket_key="snapshots/{{ ds_nodash }}/",
        mode="reschedule",
        poke_interval=300,
        timeout=7200,
    )

    # 2. Quality gate — abort early if data is bad
    @task
    def validate(ds=None) -> str:
        count = get_row_count(ds)
        if count < 10_000:
            raise ValueError(f"Insufficient data: {count} rows")
        path = f"s3://staging/clean_{ds}.parquet"
        save_cleaned_data(ds, path)
        return path                            # XCom: path, not DataFrame

    # 3. Tune + Train (separate tasks so failures are isolated)
    @task
    def tune(data_path: str) -> dict:
        df = pd.read_parquet(data_path)
        return optuna_search(df)               # XCom: best params dict

    @task
    def train(data_path: str, best_params: dict, ds=None) -> dict:
        df = pd.read_parquet(data_path)
        model_path = f"s3://models/kpi_{ds}.pkl"
        fit_and_save(df, best_params, model_path)
        return {"model_path": model_path}

    # 4. Evaluate and promote
    @task
    def evaluate(model_info: dict, ds=None) -> dict:
        metrics = run_evaluation(model_info["model_path"], ds)
        mlflow.log_metrics(metrics)
        return metrics

    @task
    def promote(metrics: dict, model_info: dict):
        champion = get_champion_metrics()
        if metrics["rmse"] < champion["rmse"]:
            register_champion(model_info["model_path"])

    # 5. Always-run failure alert
    alert = PythonOperator(
        task_id="alert_on_failure",
        python_callable=send_pagerduty,
        trigger_rule="one_failed",             # runs even if upstream fails
    )

    # Wire
    path    = validate()
    params  = tune(path)
    model   = train(path, params)
    metrics = evaluate(model)
    promote(metrics, model)

    wait >> path
    [model, metrics] >> alert

dag_instance = kpi_retraining()

What this demonstrates:

• Sensor for upstream data dependency
• Quality gate with early abort
• Isolated steps (tune failure doesn’t lose cleaned data)
• XCom by path
• Always-on failure alerting via trigger_rule
• Implicit dependencies from TaskFlow call order

7. Key Patterns to Know

Branching

1
2
3
4
5
6
7
8

def choose_path(ds=None, **ctx) -> str:
    return "train_full" if get_row_count(ds) > 100_000 else "train_lite"

branch = BranchPythonOperator(task_id="branch", python_callable=choose_path)
merge  = EmptyOperator(task_id="merge", trigger_rule="none_failed_min_one_success")

branch >> [train_full, train_lite] >> merge
# Unchosen branch is SKIPPED, not FAILED — merge needs the trigger_rule override

Trigger Rules

Parallelism via expand (Airflow 2.3+)

1
2
3
4
5
6

@task
def train_kpi(config: dict) -> dict:
    return run_training(config)

configs = [{"goal": "ctr"}, {"goal": "ecpc"}, {"goal": "vcpm"}]
results = train_kpi.expand(config=configs)   # 3 parallel task instances at runtime

Executors — One-line summary each

KubernetesExecutor is the right answer for ML: heavy training tasks get 16GB RAM pods; lightweight monitoring tasks get 512MB pods — no resource contention.

8. Airflow vs ZenML — Your Translation Layer

The key difference to articulate

ZenML auto-versions every step’s input and output as a named artifact — you can reproduce any past run by loading the exact artifact versions. Airflow doesn’t version artifacts natively; you manage this by embedding dates in S3 paths (model_2024_01_01.pkl). MLflow or DVC fills that gap when using Airflow for ML.

9. Interview Talking Points

“What is Airflow and how does it work?”

“Airflow is a metadata-driven workflow orchestrator. The Scheduler parses DAG files and marks tasks ready when their dependencies are met. The Executor dispatches those tasks to Workers, which run the actual code. Everything is recorded in a PostgreSQL metadata database. Critically, Airflow doesn’t move data — it tells other systems to do work. That separation is what makes it composable: the same Airflow DAG can orchestrate Spark jobs, Python scripts, and dbt runs within one dependency graph.”

“Walk me through how you’d design a daily retraining pipeline.”

“I’d structure it as a DAG with catchup=False and max_active_runs=1 so concurrent runs can’t race on the model registry. First, an S3 sensor in reschedule mode waits for the upstream data snapshot — this decouples the training DAG from the data pipeline DAG. Then a validation task that aborts early if row count is below threshold. Then separate tune and train tasks — keeping them separate means a training failure doesn’t force hyperparameter search to re-run. Artifacts flow between tasks as S3 paths via XCom, never as DataFrames. A final alert task with trigger_rule='one_failed' fires PagerDuty if anything breaks.”

“How do you pass data between tasks?”

“Via XCom, but only for small metadata. XCom is backed by the metadata database — you don’t want a DataFrame going in there. The production pattern is: write data to S3, XCom the path, downstream task reads from S3. This also means you can restart a failed task from the middle of the pipeline without re-running the data loading step — the path is already in XCom.”

“How does this relate to your ZenML experience?”

“The mental model is identical — a directed graph of typed steps with dependency edges. The differences are in scheduling and artifact management. ZenML doesn’t have a built-in scheduler; we trigger via GitHub Actions CI. Airflow has a mature scheduler with cron support and backfill capability, which is a real advantage for time-series workflows. ZenML auto-versions artifacts; in Airflow you manage versioning yourself through path conventions or MLflow. If I were porting our KPI model training system to Airflow, it would be a DAG with task groups per KPI, expand() for parallelism across eight models, and the same MLflow logging calls we already use — the Python logic is identical, just the orchestration wrapper changes.”

“When would you use KubernetesExecutor?”

“Any ML platform where tasks have very different resource profiles. A data cleaning task needs 2 cores and 2GB RAM; a hyperparameter search needs 8 cores and 32GB. With CeleryExecutor, workers are sized for the heaviest task — most workers are wasteful most of the time. KubernetesExecutor creates a pod per task with its own resource spec, so resources are allocated exactly to what each task needs. The tradeoff is pod startup latency — 30-60 seconds — so for DAGs with many fast tasks, Celery is faster.”