AIP-104: Task Iteration

[dabla]

---

Was generative AI tooling used to co-author this PR?

• [ x ] Yes (please specify the tool below)

Claude Code (Fable 5.1).

Description

This PR is the initial implementation of Iterable Tasks (IT), as discussed in the devlist and building upon the foundations of AIP-104. (Originally prototyped as "Dynamic Task Iteration"; renamed to Iterable Tasks following review feedback to avoid confusion with Dynamic Task Mapping.)

For further context on the use cases and performance benefits of IT, see this Medium Article and the new dynamic-task-mapping-vs-iteration.rst doc added in this PR, which compares IT with Dynamic Task Mapping (DTM) and Dynamic Task Batching in depth.

The XCom Database Constraint Challenge

While porting our internal "monkey-patched" version of IT (used since Airflow 2.x) to the core, I've identified a significant technical hurdle regarding XCom handling.

Around Airflow 2.10/2.11, a change was introduced to the database constraints for the XCom table. Specifically:

• Current State: The DB prevents creating indexed XComs (map_index >= 0) unless a corresponding mapped TaskInstance exists in the task_instance table.
• The Conflict: IT is designed to process multiple indexed XComs within a single Task Instance. Because there is no 1-to-1 mapping of a sub-task index to a physical TI row, the DB constraint blocks the insertion of these results.

The drawback is that XComs wouldn't automatically be removed from the database when a TaskInstance is deleted, which is the purpose of that constraint. So appending the index to the XCom key would be a good enough solution for IT, but not for DTM.

Current implementation in this PR

• XComIterable (airflow.sdk.bases.xcom) appends the sub-task index directly to the XCom key (return_value_<index>) to bypass the constraint, and exposes the results as a lazy Sequence (__len__/__getitem__/__iter__) so a downstream task can consume them the same way it would consume an .expand() result.
• XComIterable.flatten() returns a FlattenedXComIterable that lazily expands nested iterables (e.g. a list of pages into a single stream of items) without ever materializing the flattened stream in memory — __len__/__getitem__/__iter__ all speak consistently in flattened items.
• Progress and crash recovery no longer rely on XCom alone. IterableOperator tracks per-sub-task progress in the AIP-103 Task State Store rather than XCom, and participates in Airflow's standard retry mechanism: if the IterableOperator's task instance is retried (or manually cleared before it finishes), already-succeeded sub-tasks are skipped and only pending/failed ones re-run. XCom is only (re)written per index once a sub-task succeeds; once every index has succeeded, checkpoints are dropped so a subsequent manual clear reruns all indices from scratch rather than replaying stale results.
• Outlet asset/inlet events and on_kill propagation are handled per sub-task: a checkpointed sub-task replays its recorded outlet events on the following attempt instead of losing them, and killing the IterableOperator's task instance propagates to any sub-tasks still in flight.

I believe the cleanest long-term path is still to add a dedicated route in the Execution API that retrieves multiple XComs for a single TaskInstance by a list of keys in one round trip, so XComIterable.__getitem__/slicing don't need one request per element. I have a PR open to address this, intentionally split out of this PR.

This was also discussed in the devcall, see 2026-06-04 Dev Call Minutes.

AIP-104 itself was discussed again in the latest devcall, where the concerns raised there have also been addressed: 2026-09-10 Dev call Minutes.

Examples

The examples below assume an HTTP connection named pokeapi pointing to https://pokeapi.co.

Task Iteration

This example fetches a list of Pokémon from the PokéAPI and then uses Iterable Tasks (IT) to retrieve the details of each Pokémon. A single task instance processes all Pokémon URLs.

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()

Comparison

Pattern	Task Instances	Work Per Task
get_pokemon.expand(url=urls)	100	1 Pokémon
get_pokemon.iterate(url=urls)	1	100 Pokémon

This demonstrates how Task Iteration can significantly reduce TaskInstance creation overhead. Task Spreading (running one iteration over exactly N TaskInstances with .batch(size=N).iterate(), to be renamed .spread()) is split out into #73688.

Notable design points addressed since the initial draft

• .iterate()'s dict-argument semantics now match .expand(): passing a dict value forwards (key, value) pairs to each sub-task instead of bare keys.
• IterableOperator.task_type and .operator_name both forward to the wrapped operator (including @task-decorated callables with a custom_operator_name), so sub-tasks report the correct type in the UI/API instead of always showing MappedOperator/IterableOperator.
• XComIterable.flatten() moved out to Add XComIterable.flatten() to read an iterated task's pages as one sequence #73807, stacked on this PR, so this PR stays about running a task over its input.
• on_kill() propagates to in-flight sub-tasks, and outlet/asset events recorded by a sub-task that already succeeded are replayed from its checkpoint on a later retry instead of being lost.
• Deferred operators, reschedule-mode sensors, TriggerDagRunOperator, and ShortCircuitOperator-style downstream skipping are explicitly rejected inside IterableOperator with actionable errors rather than being silently mishandled — see the class docstring for the full list of current limitations.
• multiple_outputs is ignored for iterated tasks, explicitly at the IterableOperator level. A @task with a Mapping return annotation infers multiple_outputs=True, but the value the runner pushes for an iterated task is the XComIterable aggregate rather than a dict, so honouring the flag made the runner reject the result after every sub-task had already succeeded. Each sub-task's return value is pushed whole as return_value_<index>; keys are not fanned out into separate XComs the way .expand() does. Documented on the class and in the Task SDK docs, and pinned by a runner-level regression test for a dict-returning task under .iterate().

Per-iteration keys: XComs and task state

Every iteration of an iterated task runs under the same task instance (same dag id, task id, run id and map index). Anything an iteration writes into a per-task-instance store therefore competes with its siblings for the same key, and with the async executor the winner is whichever iteration finishes last. Two stores are affected, and both now apply the same rule: a key written from inside an iteration carries that iteration's index.

• XComs. IndexedTaskInstance.xcom_push/axcom_push suffix the key with _<index>. That is what makes return_value_<index> and XComIterable work, and it applies to any key an operator pushes from execute, including the keys of a multiple_outputs dict. Pulls are not suffixed: ti.xcom_pull(task_ids="upstream") reaches the upstream's XCom untouched.
• Task state store (the AIP-103 store, context["task_state_store"]). This was a gap: an iteration that stored a watermark or a cursor with task_state_store.set("last_offset", ...) shared that key with every sibling. IndexedTaskStateStoreAccessor closes it: IndexedTaskInstance.task_state_store is the parent's accessor seen through the index, suffixing keys on get/set/delete and their async twins, and the sub-task's context carries the same object, so an operator does not need to know it is being iterated. clear() is refused inside an iteration, since it would wipe the siblings' state and the operator's own checkpoints; an iteration deletes its own keys instead.
• The operator's checkpoints are separate. IterableOperator records per-index progress in the parent's store under _iterable_<index> and _iterable_completed, written through the parent's accessor, so they are never double-suffixed and never collide with user keys.

IndexedTaskRunner (formerly TaskExecutor, renamed because it read like one of Airflow's executors) builds the context an iteration runs against: a copy of the parent's context with the iteration's own task instance, its indexed state store view and its own outlet events. The operator binds it from the with statement, runs the operator inside that block, and records the outcome (checkpoint, XCom push, outlet-event merge) after it, so on_kill and the failure callbacks apply to the operator's execution only.

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

[kaxil]

Thanks for working on this — excited to see DTI taking shape for Airflow 3.2. I've gone through the full diff and have feedback on the implementation, some are bugs that would crash at runtime, others are design choices worth iterating on.

A few high-level things:

1. No tests. ~700 lines of new production code with zero test coverage. We need tests for IterableOperator, TaskExecutor, MappedTaskInstance, HybridExecutor, XComIterable, DecoratedDeferredAsyncOperator, and the iterate/iterate_kwargs methods — covering success, failure, retry, deferral, and edge cases.

2. Worker resilience. Since DTI runs N sub-tasks inside a single worker process, we need to think through what happens when that worker dies mid-execution — the scheduler has no record of which sub-tasks completed. Worth documenting the expected behavior and trade-offs here (and whether we want to add checkpointing later).

3. Thread safety. Several shared mutable structures (context dict, os.environ) are accessed concurrently from multiple threads without synchronization. This needs to be addressed before merge.

Inline comments below with specifics.

[dabla]

Thanks for working on this — DTI is an interesting concept and I can see the use case. I've gone through the full diff and have a number of concerns, some are bugs that would crash at runtime, others are architectural questions worth discussing before this goes further.

A few high-level things:

1. No tests. ~700 lines of new production code with zero test coverage. We need tests for IterableOperator, TaskExecutor, MappedTaskInstance, HybridExecutor, XComIterable, DecoratedDeferredAsyncOperator, and the iterate/iterate_kwargs methods — covering success, failure, retry, deferral, and edge cases.

Thanks for pointing this out. As mentioned earlier on Slack, this PR is currently intended as an initial draft to demonstrate the concept and gather early architectural feedback.

I agree that proper test coverage is essential before this can move forward. The plan is to add unit tests covering the components you mentioned (IterableOperator, TaskExecutor, MappedTaskInstance, HybridExecutor, XComIterable, DecoratedDeferredAsyncOperator, and the iterate/iterate_kwargs APIs), including scenarios for success, retries, failures, deferral, and edge cases.

Once we converge on the architectural direction, I will add the corresponding test suite.

2. Architectural concern. This builds a mini-executor inside an operator — running N tasks in threads with in-memory XCom, custom retry logic, and sleep()-based retry delays. The scheduler has no visibility into sub-task states, so if the worker dies mid-execution there's no record of which sub-tasks completed. This feels like it needs broader design discussion (probably an AIP) before merging, since it fundamentally changes how task execution works.

I agree this is an important architectural concern and worth discussing further.

The goal of this prototype is to explore a trade-off between observability and scheduling overhead, @ashb and @potiuk mentioned the same remark before. If we try to preserve the same visibility and lifecycle guarantees as Dynamic Task Mapping, we essentially end up re-implementing DTM semantics, which brings back the same scheduler overhead that this approach is trying to avoid.

This proposal intentionally explores a different point in that trade-off space: executing iterations within a single task while allowing controlled parallelism. That does mean the scheduler has indeed less visibility (but also less load) into the internal execution units.

3. Thread safety. Several shared mutable structures (context dict, os.environ) are accessed concurrently from multiple threads without synchronization.

Good point — thread safety needs to be handled carefully here.

Regarding the task context, my understanding is that operators already receive a per-task context instance, but you're right that when running iterations concurrently we should avoid sharing mutable structures across threads. One possible approach would be to create a shallow or deep copy of the context for each execution unit to ensure isolation.

If you have concerns about specific structures (e.g., os.environ or others), I'm happy to address them and introduce appropriate synchronization or isolation mechanisms where needed.

Stale review

[kaxil]

@uranusjr You should also review this PR since it touches several important modules :)

[dabla]

@uranusjr @kaxil In our patched Airflow installation I had to register the XComIterable manually with serde for serialization. How do I make sure it’s automatically registered with serde?

[dabla]

Level with main again (13 commits merged below the PR commits; tip 21b1794). The four stacked PRs are restacked, one commit each.

[kaxil]

Can you squash commits please?

[dabla]

Can you squash commits please?

Squashed commit like you asked, all current commits are related to each comment of the last review round.

[kaxil]

Only run() goes to the pool for a sync operator. IndexedTaskRunner.__exit__ stays on the loop thread, and so do the item's on_success_callback / on_skipped_callback.

So if a notifier reads a connection or variable while a sibling item is inside an async SDK call like aset_state (comms.py:283-290), it gets DeadlockImminentError. _run_task_state_change_callbacks only catches Exception, so that gets out. The item did its work, gets checkpointed UP_FOR_RETRY anyway, and the task fails with no retry. The error message blames an async sub-task.

I know we went over DeadlockImminentError in the earlier thread, but that was user code calling sync APIs from aexecute. Here the user wrote an ordinary sync operator and we're the ones who put it on a thread.

on_kill has the same problem. The SIGTERM handler and _run_tasks call it on the loop thread, and its except Exception means one failure skips the rest of the operators.

Could the sync item's enter/exit run inside the worker thread as well, and on_kill catch BaseException per operator?

[kaxil]

On SIGTERM the runner calls this once and the in-flight items are killed, but nothing stops the iteration from pulling more. Each killed item comes back as a failure, _fill_pending (executor.py:229-237) submits the next item into the freed slot, and new items, with their remote jobs, keep starting until the supervisor escalates to SIGKILL. Those never get an on_kill. A stop flag set here, checked by _run_tasks before it submits anything else, would make the kill stick.

[kaxil]

ti, task_instance, task_state_store and outlet_events are swapped, but task isn't, so inside execute and in the item callbacks context["task"] is the IterableOperator. Under .expand() it is the item's own unmapped operator (context_update_for_unmapped sets it, and the rendering in _create_task already gets that). Adding "task": self.task_instance.task here would match.

[kaxil]

This turns a string or other scalar from an upstream XCom into a one-item input. .expand() never sees those, because the upstream's _push_xcom_if_needed raises UnmappableXComTypePushed when it has a mapped dependant. An IterableOperator isn't a MappedOperator, though, so iter_mapped_dependants doesn't find it and that check never fires. An upstream that returns a JSON string makes .iterate() succeed over one wrong item. Raising here when a resolved XComArg isn't is_mappable_value would match .expand().

[kaxil]

Items that emit to the same asset end up as one event: extra.update keeps whatever the last item to finish wrote. So three items that each set extra["file"] produce one event carrying the last file, where .expand() produces three, and a consumer reading triggering_asset_events sees all three. The task instance can only send one event per asset, so this probably can't change, but the comparison table should say so. Users porting a per-file Metadata pattern from .expand() will hit it.

[kaxil]

With K failed items and a retry_policy, the policy is evaluated K times in _failure_for_the_runner, once more here and once in the runner's _evaluate_retry_policy. For a policy that calls a model, like common.ai's LLMRetryPolicy, that's K+2 calls per failed attempt, and since the answers can differ between calls, the callbacks can announce a retry the runner then doesn't take, which is what the callback change set out to prevent. Could _failure_for_the_runner keep the decision for the exception it hands over, and this reuse that instead of evaluating again?

[kaxil]

This presents iteration as a way to keep large results out of the metadata database through a custom XCom backend. But each item's full result is also written to its checkpoint (indexed_task_state.result in _run_task), which goes to the task_state_store table unless a state_store_backend is configured, and stays there for the 30-day retention. Worth saying here and pointing at state_store_backend, so someone iterating over large payloads doesn't just move them from the XCom table to the state store table.

[kaxil]

Small one: this docstring and the one at line 1791 refer to me instead of describing the behaviour. The test name and docstring should stand on their own, for example "a failed item's retry callback waits until no sibling rules out the retry".