Lab 04 · AI-native Airflow

The lab runs a five-task pipeline — extract, validate, classify, load, summarize — end-to-end on your laptop with airflow dags test, then flips one environment variable to inject a malformed record and walks the triage loop an on-call engineer would run: capture the traceback, feed it to the copilot with DAG context, get ranked causes and a fix, verify. Three production patterns in one artifact:

• AI-assisted DAG authoringyou verify by execution.
• LLM steps as first-class pipeline taskswith graceful degradation.
• AI-in-the-incident-loopwith the human still owning the fix.

Time & tier

~120 minutes hands-on · Core tier (maps to module M4 · AI-Native Orchestration, Docs & Cost).

Setup

Local Airflow, no cloud, $0. Airflow 3 installs cleanly into a venv; the Astro CLI (astro dev start) is an equally good containerized route if you prefer it, and Dagster or Prefect users can port the task graph one-to-one.

mkdir airflow-lab && cd airflow-lab && mkdir -p dags data
python3 -m venv .venv && source .venv/bin/activate
pip install "apache-airflow==3.*" duckdb anthropic pydantic
export AIRFLOW_HOME=$PWD/airflow_home
export AIRFLOW__CORE__DAGS_FOLDER=$PWD/dags
export AIRFLOW__CORE__LOAD_EXAMPLES=False
export LAB_DIR=$PWD
airflow db migrate

Save make_events.py — the raw feed your DAG ingests, with a switchable poison pill:

"""Generate a day of raw sensor alert events for the DAG to process."""
import csv, random, os
random.seed(21)
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
os.makedirs(OUT, exist_ok=True)
kinds = [("temp spike to {v}C on line {l}", "THERMAL"), ("humidity {v}% exceeds spec on line {l}", "MOISTURE"),
         ("battery at {v}% on sensor {l}", "POWER"), ("no heartbeat for {v} min from {l}", "CONNECTIVITY")]
rows = []
for i in range(1, 201):
    tmpl, _ = random.choice(kinds)
    rows.append([i, tmpl.format(v=random.randint(5, 99), l=f"L{random.randint(1,8)}"),
                 f"2025-07-14 {random.randint(0,23):02d}:{random.randint(0,59):02d}:00"])
if os.environ.get("INJECT_FAILURE") == "1":
    rows.append([201, None, "not-a-timestamp"])   # malformed row for the triage exercise
with open(f"{OUT}/alerts.csv", "w", newline="") as f:
    w = csv.writer(f); w.writerow(["alert_id", "message", "raised_at"]); w.writerows(rows)
print(f"wrote {len(rows)} alerts (INJECT_FAILURE={os.environ.get('INJECT_FAILURE','0')})")

And the DAG, dags/alert_triage_dag.py. Draft it with your assistant first if you like (the prompt that works: "Airflow 3 TaskFlow DAG, five tasks extract→validate→classify→load→summarize, CSV in, DuckDB out, retries 1, no schedule") — then reconcile against this reference, which is the version the rest of the lab assumes:

"""alert_triage: ingest raw sensor alerts, classify each one, load to DuckDB, publish a summary.

Offline-first: classification uses deterministic rules by default. If ANTHROPIC_API_KEY is set,
the classify task upgrades itself to an LLM classifier with structured outputs.
"""
from __future__ import annotations
import csv, os
from datetime import datetime, timedelta
from airflow.sdk import dag, task

LAB_DIR = os.environ.get("LAB_DIR", os.path.expanduser("~/airflow-lab"))
DATA = os.path.join(LAB_DIR, "data", "alerts.csv")
DB = os.path.join(LAB_DIR, "data", "alerts.duckdb")
CATEGORIES = ["THERMAL", "MOISTURE", "POWER", "CONNECTIVITY", "OTHER"]

def classify_rules(message: str) -> tuple[str, float]:
    m = message.lower()
    if "temp" in m or "thermal" in m: return "THERMAL", 0.9
    if "humidity" in m or "moist" in m: return "MOISTURE", 0.9
    if "battery" in m or "power" in m: return "POWER", 0.9
    if "heartbeat" in m or "offline" in m: return "CONNECTIVITY", 0.9
    return "OTHER", 0.3

def classify_llm(messages: list[str]) -> list[tuple[str, float]]:
    """Optional Claude path -- one batched call, structured output, validated labels."""
    import anthropic
    from pydantic import BaseModel
    class Labels(BaseModel):
        labels: list[str]
        confidences: list[float]
    client = anthropic.Anthropic()
    numbered = "\n".join(f"{i+1}. {m}" for i, m in enumerate(messages))
    response = client.messages.parse(
        model="claude-opus-4-8",
        max_tokens=16000,
        thinking={"type": "adaptive"},
        messages=[{"role": "user", "content":
            f"Classify each alert into exactly one of {CATEGORIES}. "
            f"Return one label and one confidence (0-1) per alert, in order.\n{numbered}"}],
        output_format=Labels,
    )
    out = response.parsed_output
    return [(lab if lab in CATEGORIES else "OTHER", conf)
            for lab, conf in zip(out.labels, out.confidences)]

@dag(
    dag_id="alert_triage",
    start_date=datetime(2025, 7, 1),
    schedule=None,               # trigger manually while developing
    catchup=False,
    default_args={"retries": 1, "retry_delay": timedelta(minutes=1)},
    tags=["lab", "ai-at-work"],
)
def alert_triage():
    @task
    def extract() -> list[dict]:
        with open(DATA, newline="") as f:
            return list(csv.DictReader(f))

    @task
    def validate(rows: list[dict]) -> list[dict]:
        bad = [r for r in rows if not r["message"] or len(r["raised_at"]) != 19]
        if bad:
            raise ValueError(f"{len(bad)} malformed rows, first: {bad[0]}")
        return rows

    @task
    def classify(rows: list[dict]) -> list[dict]:
        if os.environ.get("ANTHROPIC_API_KEY"):
            results = classify_llm([r["message"] for r in rows])
        else:
            results = [classify_rules(r["message"]) for r in rows]
        return [{**r, "category": c, "confidence": conf}
                for r, (c, conf) in zip(rows, results)]

    @task
    def load(rows: list[dict]) -> int:
        import duckdb
        con = duckdb.connect(DB)
        con.sql("""create table if not exists alert_facts (
                   alert_id int, message varchar, raised_at timestamp,
                   category varchar, confidence double)""")
        con.sql("delete from alert_facts")   # idempotent re-run for the lab
        con.executemany("insert into alert_facts values (?,?,?,?,?)",
                        [(r["alert_id"], r["message"], r["raised_at"], r["category"], r["confidence"]) for r in rows])
        return len(rows)

    @task
    def summarize(n: int):
        import duckdb
        print(f"loaded {n} alerts")
        print(duckdb.connect(DB).sql("select category, count(*) from alert_facts group by 1 order by 2 desc"))

    summarize(load(classify(validate(extract()))))

alert_triage()

Guided steps

1. 1.

   Parse before anything. python make_events.py (expect wrote 200 alerts), then python dags/alert_triage_dag.py. A clean exit means the DAG file imports — the cheapest gate in orchestration and the one AI-generated DAGs fail most (hallucinated imports, Airflow-2-isms like schedule_interval, decorators on nested functions). If your own assistant-drafted version fails here, diff it against the reference and log what the model got wrong: that list is your personal review checklist for generated DAGs.

2. 2.

   Review the DAG like a senior engineer, with your assistant. Paste the file in and ask: "Review this Airflow 3 DAG for production readiness: XComs carrying full datasets, retry semantics on non-idempotent tasks, missing SLAs/alerts, schedule choices." The review should surface the lab's deliberate simplifications — 200 rows through XCom is fine for a lab but a real feed needs a storage-backed handoff (object store paths or Airflow assets); retries: 1 on load is only safe because the delete-then-insert makes it idempotent. Write the three best findings into review-notes.md; keep or reject each with a reason.

3. 3.

   Run it for real. airflow dags test alert_triage executes the whole graph in-process without the scheduler. Expect all five tasks to succeed and the summary to print a category table (THERMAL/MOISTURE/POWER/CONNECTIVITY roughly balanced, OTHER near zero). This is your green baseline — screenshot or copy the tail into baseline-run.txt.

4. 4.

   Understand the LLM task's design before you scale it. Three choices in classify are the production pattern, and you should be able to defend each: it batches (one call for 200 records, not 200 calls — two orders of magnitude on latency and cost); it validates (structured output parsed against a schema, labels checked against the allowed set, junk demoted to OTHER — LLM output is untrusted input); and it degrades (no key → deterministic rules, the DAG never hard-depends on an external API being up). In production you would reach for the Airflow Common AI provider or the astronomer airflow-ai-sdk, which wrap exactly this shape in dedicated operators with retry/cost controls — note in review-notes.md which parts of classify_llm those operators would replace. If you have an ANTHROPIC_API_KEY, export it and rerun step 3: the classifier upgrades in place and the summary counts should barely move (the rules are a decent proxy — that agreement check is itself an eval).

5. 5.

   Inject the failure. INJECT_FAILURE=1 python make_events.py (writes 201 rows, one malformed), then airflow dags test alert_triage again. The run fails at validate with ValueError: 1 malformed rows, first: {'alert_id': '201', 'message': None, ...} and downstream tasks never run. Capture the traceback block from the terminal into incident/logs.txt — resist the urge to just fix it; the exercise is the loop, not the bug.

6. 6.

   Build the triage prompt template. Save incident/triage_prompt.md — a reusable template with four context slots, because a triage copilot is only as good as what you feed it:

   Copy

   You are an Airflow on-call triage copilot. Diagnose the failure below.
   
   ## DAG context
   dag_id: {dag_id} -- purpose: {one_line_purpose}
   task order: {task_list}; failed task: {task_id}; prior runs green until now.
   
   ## Failure log
   {traceback}
   
   ## Recent changes
   {deploys_or_data_changes_or_none_known}
   
   ## Task source
   {failed_task_code}
   
   Respond with: (1) probable root cause, ranked, max 3, each with confidence;
   (2) the exact evidence in the log supporting each; (3) is this a data problem,
   a code problem, or an infra problem; (4) immediate remediation vs proper fix;
   (5) what to check before re-running so retries do not re-fail.
   

   Fill the slots from your artifacts and run it through your assistant (Claude, ChatGPT, or a local model). A good response pins it as a data problem (upstream feed shipped a null-message row), evidence being the None message in the ValueError, and separates remediation (quarantine the bad row, rerun) from the proper fix (validate should divert bad rows to a reject table instead of failing the run — a dead-letter pattern). Save the full exchange as incident/transcript.md.

7. 7.

   Close the loop like on-call would. Implement the copilot's remediation: change validate to split good and bad rows, write rejects to data/rejects.csv with a reason column, log a warning, and only fail if more than 5% of rows are bad. Rerun with the poison pill still in place: the DAG now goes green, rejects.csv holds row 201, and the summary still reports 200 loaded alerts. Append a three-line postmortem to incident/transcript.md: what failed, why the fix is the right trade-off (availability over strictness at 0.5% bad rows, with the threshold making mass corruption still fail loudly), and what monitor would catch a rising reject rate.

Verification

• •
  python dags/alert_triage_dag.py exits cleanly (DAG imports).
• •
  airflow dags test alert_triage on clean data: all five tasks succeed, summary prints the category table, data/alerts.duckdb contains 200 rows in alert_facts.
• •
  With INJECT_FAILURE=1 and the original validate: the run fails at validate with the ValueError naming 1 malformed row; the traceback is captured in incident/logs.txt.
• •
  With the dead-letter fix: the poisoned run goes green, data/rejects.csv contains exactly the malformed row, and 200 rows load.
• •
  incident/triage_prompt.md is fully generic (slots, no hardcoded specifics) and incident/transcript.md shows the solved incident plus the postmortem.

Completion artifact

The airflow-lab/ repo: DAG, generator, review notes, baseline run, and the incident/ folder (logs, prompt template, solved transcript, postmortem). Self-grading rubric (pass = 8/10):

• DAG quality (0–2): parses, runs green, tasks idempotent where retried = 2.
• AI-review judgment (0–2): three findings logged with keep/reject reasoning, XCom-size and retry-semantics issues among them = 2.
• LLM-task design (0–3): batching, output validation against the label set, and offline degradation all present and explained = 3.
• Triage loop (0–2): prompt template reusable; transcript shows ranked causes with log evidence; data-vs-code-vs-infra call correct = 2.
• Fix quality (0–1): dead-letter implementation with a thresholded failure mode, not a silent swallow = 1.

Stuck?

• •
  airflow db migrate errors about existing tables: delete airflow_home/ and re-export the env vars in a fresh shell.
• •
  dags test says dag not found: AIRFLOW__CORE__DAGS_FOLDER is not pointing at your dags/ dir, or the file has a syntax error — rerun step 1.
• •
  from airflow.sdk import dag, task fails: you have Airflow 2.x — the lab needs 3.x (pip show apache-airflow).
• •
  The LLM path errors on import: anthropic and pydantic must be in the same venv Airflow runs from; or just unset ANTHROPIC_API_KEY and use the rules path.

Client-safe variant

Pipeline logs and DAG source frequently embed client identifiers — connection names, bucket paths, schema names — so triage prompts are a real leak vector. Route the review and triage steps through the client-approved tenancy (Amazon Bedrock or Azure OpenAI) or a local model via Ollama, and add a scrub pass to the prompt template: replace connection IDs, hostnames, and bucket names with placeholders before the log leaves the client boundary (the template's slot structure makes this mechanical). The DAG itself, airflow dags test, and the dead-letter fix are all local and LLM-free, and the in-DAG classifier's rules fallback means the pipeline runs with zero external calls when the engagement demands it.

Answer key (don't read before attempting)

1. 1Injected failure:make_events.py under INJECT_FAILURE=1 appends row 201 with a None message and not-a-timestamp raised_at; validate raises ValueError: 1 malformed rows and downstream tasks are skipped. It is a data incident (upstream feed), not code or infra.
2. 2Why retries cannot save it: the input file is unchanged on retry, so the same row fails deterministically — the copilot should say "do not just re-run"; retry only helps transient faults.
3. 3Proper fix: dead-letter the bad rows (rejects file/table with a reason), warn, and fail only above a badness threshold (lab uses 5%) so a mass-corruption event still stops the pipeline; add a reject-rate monitor so quarantining does not become silent data loss.