AIP-104: Task Iteration (TI) and Dynamic Task Batching (DTB)

Status

Motivation

Airflow’s existing Dynamic Task Mapping (DTM) provides a powerful way to process collections by expanding tasks into multiple Task Instances. However, this model introduces significant overhead when applied to large datasets:

• Each item becomes a separate Task Instance
• The scheduler must create, track, and persist each instance
• The metadata database experiences increased load
• Worker slots are consumed per item

This approach works well for coarse-grained, independently trackable workloads, but becomes inefficient for:

• High-volume workloads (thousands to millions of items)
• Short lived workloads
• I/O-bound operations (e.g. API calls, SFTP downloads)
• Streaming or paginated data processing

In parallel, deferrable operators—often used with DTM for async workloads—introduce additional constraints:

• Dependence on triggerers
• Triggerers storing events in the metadata database
• Limited scalability compared to workers
• Inability to leverage custom XCom backends

With the introduction of native async support in Python tasks (AIP-98), a new execution pattern becomes possible:

Instead of distributing work across many tasks, process collections within a single task using concurrency (threads or async).

This leads to the introduction of Task Iteration (TI).

TI shifts iteration from the scheduler layer to the execution layer, eliminating scheduling overhead and enabling efficient in-task concurrency.

Additionally, a future pattern—Dynamic Task Batching (DTB)—combines DTM and TI to balance distribution and efficiency.

Considerations

Task Iteration introduces a fundamentally different execution model compared to Dynamic Task Mapping (DTM), which leads to several trade-offs and behavioural differences.

Retry semantics

• DTM:
  • Retries are handled per item (per Task Instance)
  • State is fully isolated and persisted by the scheduler
• TI:
  • Retries are also handled per item within the task execution
  • Progress is maintained across completed batches within the same Task Instance
  • However, if the worker dies during execution:
    • In-flight progress (current batch) may be lost
    • Retry state within that batch is not preserved
  • Previously completed batches do not need to be reprocessed

This makes Task Iteration partially stateful within a task execution, but not resilient to worker-level failures in the same way as DTM, unless we implement Task State thanks to AIP-103.

Observability

• DTM:
  • Full per-item observability in the Airflow UI
  • Each item is a separate Task Instance with its own logs and state
• TI:
  • Aggregated observability within a single Task Instance
  • Per-item visibility is not exposed in the UI
  • Reduced observability is a trade-off for significantly improved performance

Execution model

• DTM:
  • Work is distributed across multiple workers
  • Each item executes independently
• TI:
  • All items are processed within a single worker
  • Iteration happens in-process
  • Supports:
    • Sequential execution
    • Multi-threading (sync)
    • Async multiplexing (shared event loop)

Failure isolation

• DTM:
  • Failures are isolated per item
  • Only failed items are retried
• TI:
  • Failures are handled at the task level
  • Items within the current batch may be retried
  • In case of task failure, reprocessing scope depends on:
    • Completed batches (persisted)
    • In-flight batch (may be retried)

Dynamic Task Batching (DTB)

Dynamic Task Batching combines the strengths of both models:

• DTM distributes partitions across workers
• TI processes items within each partition

This provides:

• Reduced scheduler load (compared to full DTM)
• Improved failure isolation (compared to pure TI)
• High throughput within each partition

Summary of trade-offs

What change do you propose to make?

Introduce Task Iteration (TI) as a first-class execution model in Airflow.

This includes:

• Adding an .iterate() API to tasks/operators
• Allowing iteration over iterables (typically XCom results) within a single Task Instance
• Supporting multiple execution strategies within that task:
  • Sequential execution
  • Multi-threaded execution (for sync tasks)
  • Async multiplexing (shared event loop for async tasks)

TI enables:

• Per-item processing within a task
• Per-item retry handling inside the task execution
• Batch-based progress tracking within the task

Additionally, introduce Dynamic Task Batching (DTB) as an extension of Task Iteration (TI):

• Add a .batch(size=N) API to split an iterable into batches or chunks
• Combine with .iterate() to process each partition within a task

Example:

get_pokemon.batch(size=2).iterate(url=list_pokemon())

This pattern:

• Uses DTM implicitly to distribute partitions across workers
• Uses TI within each partition to process items efficiently

What problem does it solve?

TI and DTB together address several limitations of existing Airflow execution models:

1. Scheduler scalability limitations

DTM:

• One Task Instance per item
• High scheduler and database overhead

TI:

• Single Task Instance
• Minimal scheduler overhead

DTB:

• One Task Instance per partition (instead of per item)
• Reduces scheduler load from O(N) → O(N / batch_size)

2. Inefficient handling of high-volume I/O workloads

DTM and deferrable operators:

• Do not efficiently multiplex I/O
• Introduce overhead per item

TI:

• Enables async multiplexing within a task

DTB:

• Combines distribution + multiplexing
• Enables high throughput at scale

3. Triggerer bottlenecks

Deferrable operators:

• Depend on triggerers
• Store events in metadata DB

TI/DTB:

• Run entirely on workers
• Avoid triggerers
• Leverage worker scalability and custom XCom backends

4. Poor trade-off between scalability and retry granularity

• DTM → good retry granularity, high overhead
• TI → low overhead, coarse failure handling

DTB provides a balance:

• Retry per batch (coarse-grained)
• Efficient execution within each batch

Why is it needed?

TI introduces a third execution model in Airflow:

However, TI alone is not always sufficient:

• Too coarse-grained for large workloads (single task retry)
• Limited by single-worker execution

DTB is needed to bridge this gap:

• Combines horizontal scaling (DTM) with efficient iteration (TI)
• Provides a middle ground between:
  • Fine-grained (DTM)
  • Fully aggregated (TI)

This enables:

• Scalable high-throughput pipelines
• Efficient API pagination at scale
• Large dataset processing with controlled retry boundaries

Are there any downsides to this change?

Task Iteration downsides

• Reduced observability (no per-item UI visibility)
• Execution limited to a single worker
• Partial durability:
  • Completed batches are preserved
  • In-flight batch progress may be lost if worker dies (unless AIP-103 is used, which will be tested soon)
• Requires batching to avoid memory pressure

Dynamic Task Batching downsides

• More complex mental model (DTM + TI combined)
• Batch size tuning required:
  • Too small → behaves like DTM (high overhead)
  • Too large → behaves like TI (coarse retries)
• Retry granularity limited to partition level
• Still less observable than pure DTM

Which users are affected by the change?

• Users processing large datasets
• Users working with I/O-bound workloads (APIs, SFTP, databases)
• Users leveraging async hooks or async tasks
• Users experiencing scheduler or triggerer bottlenecks

DTB is particularly relevant for users who:

• Currently use large-scale Dynamic Task Mapping
• Need a balance between scalability and retry granularity

How are users affected by the change? (e.g. DB upgrade required?)

• No database schema changes required
• No breaking changes
• Fully backward compatible

New APIs are optional:

• .iterate() for TI
• .batch(size=N) for DTB

Existing DAGs continue to work unchanged.

What is the level of migration effort (manual and automated) needed for the users to adapt to the breaking changes? (especially in context of Airflow 3)

• None required (no breaking changes)
• Adoption is optional

DAG authors may:

• Replace .expand() (DTM) with .iterate() for TI
• Introduce .batch(size=N) to adopt DTB
• Refactor logic into task-level iteration
• Tune partition sizes based on workload

Example migration:

# Before (DTM)
task.expand(item=items)

# After (TI)
task.iterate(item=items)

# After (DTB)
task.batch(size=4).iterate(item=items)

Examples

Pokemon REST API with Task Iteration

from airflow.sdk import dag, task
from airflow.providers.http.hooks.http import HttpHook, HttpAsyncHook

from pendulum import datetime


@dag(
    start_date=datetime(2025, 1, 1),
    schedule=None,
    catchup=False,
)
def pokemon_iteration():
    @task
    def list_pokemon() -> list[str]:
        response = HttpHook(
            http_conn_id="pokeapi",
            method="GET",
        ).run(
            endpoint="api/v2/pokemon?limit=100",
        )

        return [
            pokemon["url"].replace("https://pokeapi.co/", "")
            for pokemon in response.json()["results"]
        ]

    @task(
        retries=3,
        task_concurrency=2,
        show_return_value_in_logs=False,
    )
    async def get_pokemon(url: str):
        async with HttpAsyncHook(
            http_conn_id="pokeapi",
            method="GET",
        ).session() as session:
            response = await session.run(endpoint=url)
            return await response.json()

    get_pokemon.iterate(
        url=list_pokemon(),
    )


pokemon_iteration()

Pokemon REST API with Dynamic Task Batching and Task Iteration

from airflow.sdk import dag, task
from airflow.providers.http.hooks.http import HttpHook, HttpAsyncHook

from pendulum import datetime


@dag(
    start_date=datetime(2025, 1, 1),
    schedule=None,
    catchup=False,
)
def pokemon_batched_iteration():
    @task
    def list_pokemon() -> list[str]:
        response = HttpHook(
            http_conn_id="pokeapi",
            method="GET",
        ).run(
            endpoint="api/v2/pokemon?limit=100",
        )

        return [
            pokemon["url"].replace("https://pokeapi.co/", "")
            for pokemon in response.json()["results"]
        ]

    @task(
        retries=3,
        task_concurrency=2,
        show_return_value_in_logs=False,
    )
    async def get_pokemon(url: str):
        async with HttpAsyncHook(
            http_conn_id="pokeapi",
            method="GET",
        ).session() as session:
            response = await session.run(endpoint=url)
            return await response.json()

    get_pokemon.batch(size=2).iterate(
        url=list_pokemon(),
    )


pokemon_batched_iteration()

Comparison

Other considerations?

What defines this AIP as "done"?

https://github.com/orgs/apache/projects/647/views/1

• AIP-104: Task Iteration and Dynamic Task Batching is merged
• Added asynchronous aget_hook method to BaseHook is merged
• Added deadlock detection in CommsDecoder is merged
• Add async XCom accessors for async tasks is merged
• Add a `POST /xcoms/{dag_id}/{run_id}/{task_id}/keys` execution API endpoint is merged
• Intermediate state is saved via AIP-103 Task State Management.