Orchestrate LLM Calls as First-Class Tasks in Your DAGs

Tier: Core · Workflow area: Orchestration, ops & cost

LLM calls deserve the same production treatment as any external dependency: typed operators, backoff, dead-letter handling, and cost caps. This playbook wires LLM work into Airflow 3 as first-class tasks using the Common AI provider / airflow-ai-sdk patterns, with the batch API as the default for anything bulk.

When to reach for this

• •
  A scheduled pipeline needs LLM steps — classify support tickets nightly, summarize documents on landing, enrich records, generate embeddings.
• •
  An existing hand-rolled LLM call in a DAG keeps failing on rate limits or silently doubling spend on retries.
• •
  You're building the ingestion side of a RAG system and the enrichment/embedding stages need orchestration-grade reliability.

For interactive/low-latency LLM features, this is the wrong layer — orchestrators are for scheduled and event-driven batch work.

The workflow

1. 1.

   Pick the integration layer. Airflow's Common AI provider (built on Pydantic AI, 20+ model providers) and airflow-ai-sdk give you @task.llm-style tasks and typed outputs instead of hand-rolled clients — provider-agnostic, so switching models is config, not code. Where the data already lives in a warehouse, in-SQL functions (Snowflake Cortex AISQL, BigQuery AI.GENERATE family) can replace the external call entirely; weigh governance and cost, not novelty.

2. 2.

   Shape LLM work into idempotent, batched tasks. The unit of work is a partition (this hour's tickets), not a row. Each task: reads its partition, skips rows already processed (check output table by content hash), calls the model over the remainder, writes results keyed by input hash. That single design decision makes retries and backfills safe.

3. 3.

   Default to the batch API for bulk stages. Anything that can tolerate minutes-to-hours latency — which is nearly everything in a DAG — should use provider batch endpoints at ~50% discount. Pattern: task A submits the batch and pushes the batch id via XCom; a deferrable sensor task polls until it ends (deferrable, so it doesn't hold a worker slot); task B fetches results, keyed by custom_id (never by order), and dead-letters failed items:

   Copy

   @task
   def submit(partition: str) -> str:
       reqs = build_requests(partition)     # custom_id = input row hash
       return client.messages.batches.create(requests=reqs).id
   
   @task
   def collect(batch_id: str) -> None:
       for r in client.messages.batches.results(batch_id):
           if r.result.type == "succeeded":
               write_row(r.custom_id, parse(r.result.message))
           else:
               dead_letter(r.custom_id, r.result.type)
   

4. 4.

   Handle the failure taxonomy explicitly. Rate limits (429): retry with exponential backoff and honor retry-after — Airflow task retries with backoff cover this if per-call retries are bounded. Malformed outputs: enforce structured outputs so parse failures are impossible or loud, and route irreparable rows to a dead-letter table rather than failing the whole partition. Provider outages: the task fails, Airflow retries later — this is why LLM calls belong in the orchestrator.

5. 5.

   Cap and meter cost per run. Before submitting, estimate tokens (count tokens on a sample × row count) and fail fast if the estimate exceeds the DAG-run budget parameter. After collecting, write actual usage (input/output tokens, model, cost) to a metering table keyed by dag_id/run_id — this table is what pb-26's cost-engineering work reads.

6. 6.

   Alert on quality drift, not just failure. A classification task that "succeeds" while returning 40% unknown is broken. Add a lightweight check task downstream: distribution of output labels vs. trailing average, dead-letter rate threshold, empty-output rate. Wire it to the same alerting as any data-quality check.

Verify before you trust it

• •
  Re-running a partition produces zero duplicate outputs and zero re-billed rows (the hash-skip works).
• •
  Kill the polling task mid-batch and re-run: it resumes on the same batch id instead of resubmitting (no double spend).
• •
  The dead-letter table receives deliberately malformed inputs; the run still succeeds and the alert fires at your threshold.
• •
  The metering table's per-run cost matches the provider console within a few percent.
• •
  The cost cap actually trips: run with a tiny budget parameter and confirm fail-fast before submission.

Pitfalls

• Retry-driven double billing. A task that submits a batch and fails after submission will resubmit on retry unless the batch id is persisted first. Persist, then poll.
• Worker-slot starvation. Poke-mode polling of a 2-hour batch holds a slot the whole time. Use deferrable sensors/triggers for the wait.
• Row-per-task explosion. Dynamic task mapping over 100k rows melts the scheduler. Map over partitions, loop over rows inside the task.
• Unpinned model versions.model="latest" means silent behavior change mid-history. Pin model IDs in DAG params and change them via PR.
• XCom-ing large payloads. Pass object-store paths and batch ids through XCom, never the actual documents or results.

Tools

• Airflow Common AI provider / airflow-ai-sdktyped LLM and agent operators, 20+ providers, the idiomatic Airflow 3 integration layer (OSS).
• Anthropic batch API + structured outputsthe bulk-work default: ~50% cheaper, schema-enforced outputs that survive orchestration parsing.
• Snowflake Cortex AISQL / BigQuery AI.GENERATEthe in-warehouse alternative when data gravity or governance says the text shouldn't leave the platform.

Client-safe variant

The pipeline pattern is generic; the payload is client data — so provider choice is the client's call, not yours. Confirm which model endpoints are approved (their cloud's hosted models, a self-hosted vLLM, or a signed provider agreement) and parameterize the DAG so the provider is a connection, not code. Keep prompts free of client identifiers where the approved endpoint is external, and route all metering to the client's own cost tables. The dead-letter and metering tables make good handover artifacts: they let the client's team operate the pipeline — and audit its spend — without you.