dabla commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r4236744455
########## task-sdk/docs/mapped-tasks-vs-iterable-tasks.rst: ########## @@ -0,0 +1,507 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _sdk-mapped-tasks-vs-iterable-tasks: + +Mapped tasks vs iterable tasks +============================== + +.. versionadded:: 3.4.0 + +Airflow provides two complementary ways to process collections of data: + +- **Mapped tasks** distribute work **across multiple workers**. + Each item becomes a separate Task Instance that can run on a different worker, + giving you horizontal scalability and per-item observability. + +- **Iterable Tasks (IT)** improves concurrency **within a single task**. + All items are processed inside one Task Instance on one worker, eliminating + scheduling overhead and — when combined with async operators — enabling true + I/O multiplexing through a shared event loop. + +In short: **mapping spreads load across workers; IT speeds up work within one worker.** + +While both approaches allow you to apply an operation over a collection, +they differ significantly in execution model, scheduler impact, and observability. +This page explains the trade-offs and when to use each. + +Real-World Motivation +--------------------- + +Consider a workflow that downloads ~17,000 XML files from an SFTP server and loads +them into a data warehouse. Community benchmarks compare a mapped operator with a loop +written by hand inside a single ``@task``, the pattern IT runs for you (none of these rows +uses ``iterate()`` itself): + +.. list-table:: + :header-rows: 1 + + * - Approach + - Execution Time + * - Mapped ``SFTPOperator`` + - 3 h 25 m + * - Sync ``@task`` with ``SFTPHook`` (sequential loop) + - 1 h 21 m + * - Async ``@task`` with ``SFTPHookAsync`` (concurrent loop) + - 8 m 29 s + * - Async ``@task`` with ``SFTPHookAsync`` and connection pooling + - 3 m 32 s + +The ~60× improvement comes from eliminating per-item scheduling overhead and +sharing a single event loop for concurrent I/O. This is the kind of workload +IT is built for: many small, I/O-bound operations processed within one task. + +Mapped tasks +------------ + +Task mapping allows you to expand a single task definition into multiple +Task Instances (TIs). + +For more details, see :ref:`task mapping <sdk-dynamic-task-mapping>`. + +Key characteristics: + +- Each item in the iterable creates a separate Task Instance. +- The scheduler is responsible for creating and managing all mapped tasks. +- Tasks can run in parallel across multiple worker slots. +- Fine-grained retry, logging, and observability per item. +- Well suited for workloads where each item should be independently scheduled and tracked. + +The following example fetches Pokémon data from a REST API. Each Pokémon becomes +a separate Task Instance, individually scheduled, retried, and visible in the UI: + +.. code-block:: python + + from datetime import datetime + + from airflow.providers.http.operators.http import HttpOperator + from airflow.sdk import DAG, task + + with DAG(dag_id="dtm-http-pokemon-example", start_date=datetime(2026, 1, 1)): + list_pokemon_task = HttpOperator( + task_id="list_pokemon", + http_conn_id="pokeapi", + method="GET", + endpoint="api/v2/pokemon?limit=100", + response_filter=lambda response: [ + pokemon["url"].replace("https://pokeapi.co/";, "") for pokemon in response.json()["results"] + ], + log_response=False, + ) + + get_pokemon_task = HttpOperator.partial( + task_id="get_pokemon", + http_conn_id="pokeapi", + method="GET", + ).expand(endpoint=list_pokemon_task.output) + + list_pokemon_task >> get_pokemon_task + + +With 100 Pokémon the scheduler creates 100 Task Instances, each occupying +a worker slot. This is fine for small lists, but for thousands of items the +scheduler and database overhead becomes significant. + +Iterable Tasks (IT) +---------------------------- + +Iterable Tasks allows you to iterate over an iterable (typically an XCom result) +*within a single Task Instance*, applying an operator multiple times without creating +separate Task Instances. + +This means that iteration happens inside the task execution itself rather than at the +scheduler level. + +Key characteristics: + +- A single Task Instance processes all items in the iterable. +- No task expansion; the scheduler manages only one task. +- Lower scheduler overhead compared to mapping. +- Iterations share the same execution context (e.g., memory, event loop). +- Particularly well suited for async operators and high-throughput workloads. + +The same Pokémon fetching problem can be solved with IT. Here, a single Task +Instance processes all Pokémon concurrently using the sync +:class:`~airflow.providers.http.operators.http.HttpOperator`: + +.. code-block:: python + + from datetime import datetime + + from airflow.providers.http.operators.http import HttpOperator + from airflow.sdk import DAG, task + + with DAG(dag_id="it-http-pokemon-example", start_date=datetime(2026, 1, 1)): + list_pokemon_task = HttpOperator( + task_id="list_pokemon", + http_conn_id="pokeapi", + method="GET", + endpoint="api/v2/pokemon?limit=100", + response_filter=lambda response: [ + pokemon["url"].replace("https://pokeapi.co/";, "") for pokemon in response.json()["results"] + ], + log_response=False, + ) + + get_pokemon_task = HttpOperator.partial( + task_id="get_pokemon", + http_conn_id="pokeapi", + method="GET", + ).iterate(endpoint=list_pokemon_task.output) + + list_pokemon_task >> get_pokemon_task + + +The scheduler only manages a single task. With a sync operator, iterations run in +a pool of up to ``task_concurrency`` threads, so blocking I/O such as these HTTP +requests overlaps up to that many at a time. Async operators scale further: a +coroutine waiting on I/O costs far less than a thread, so ``task_concurrency`` can +be set much higher. CPU-bound Python code speeds up with neither, since the +iterations share one process and its GIL. + +To **multiplex** many I/O-bound operations on one event loop, use an async task with +:class:`~airflow.providers.http.hooks.http.HttpAsyncHook`: + +.. code-block:: python + + from datetime import datetime + + from airflow.providers.http.hooks.http import HttpAsyncHook, HttpHook + from airflow.sdk import dag, task + + + @dag( + dag_id="it-async-http-pokemon-example", + start_date=datetime(2026, 1, 1), + ) + def it_async_http_pokemon_example(): + @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(), + ) + + + it_async_http_pokemon_example() + + +When ``iterate()`` is used with an async task, all iterations share the same +event loop, enabling true multiplexing of I/O-bound operations without any +manual concurrency management by the DAG author. For a handful of items the +difference is negligible, but for hundreds or thousands of items the +concurrent approach is dramatically faster — see the +:ref:`benchmarks above <sdk-mapped-tasks-vs-iterable-tasks>`. + +.. note:: + + ``multiple_outputs`` is ignored by ``iterate()``. Each iteration's return value is pushed + whole as ``return_value_<index>`` and the task's own return value is the lazy sequence over + them, so a ``dict`` return annotation on the task does not fan its keys out into separate + XComs the way it does with ``expand()``. Every key an iteration pushes or stores carries its + index the same way: ``ti.xcom_push("foo", v)`` in iteration 2 lands under ``foo_2``, and so + does ``task_state_store.set("foo", v)``, so iterations never overwrite each other's values. + Reading them back follows the same rule: ``ti.xcom_pull(key="foo")`` in iteration 2, or with + its own ``task_ids``, reads ``foo_2``, as ``task_state_store.get("foo")`` does; a pull from + another task keeps its key. To read another iteration's value, use ``XCom.get_one`` + (``XCom.aget_one`` in an async operator) with the suffixed key, or read every value downstream + through the task's lazy sequence. + This holds in the iteration's own thread or coroutine. A ``threading.Thread`` the task starts, + or ``loop.run_in_executor()``, does not inherit it: ``get_current_context()`` there returns + the task's own context, whose ``ti`` and ``task_state_store`` add no index, so the iterations' + keys overwrite each other. Use ``context["ti"]`` passed to the task, or start helpers with + ``asyncio.to_thread()`` or ``contextvars.copy_context().run()``, which carry the iteration's + context over. + +.. warning:: + + Inputs are shared between iterations. All iterations run in one process, so a value handed + to several of them is the same object in each: every value passed through ``partial()``, and + with ``iterate(a=..., b=...)`` every element of ``a`` and of ``b``, which the cross product + combines more than once. With ``expand()`` each task instance runs in its own process and + gets its own copy. Treat inputs as read-only, or copy what the task changes in place. + +Why Iterable Tasks? +--------------------------- + +IT is designed to address limitations of task mapping in specific scenarios: + +- **Scheduler scalability**: + Mapping creates one Task Instance per item, which can put pressure on the scheduler + for very large datasets. IT avoids this by keeping execution within a single task. + +- **Async multiplexing**: + With Python-native async support in Airflow 3.2, IT allows multiple + operations to share the same event loop within a single Task Instance. + This enables efficient multiplexing of I/O-bound workloads. + +- **Lower overhead**: + No need to serialize, schedule, and track thousands of Task Instances. + +- **Triggerer and deferrable-operator bottleneck**: + Deferrable operators delegate async work to triggerers, which store yielded + events directly in the Airflow metadata database. Unlike workers, triggerers + cannot leverage a custom XCom backend to offload large payloads. This makes + triggerers a bottleneck for sustained high-load async execution or workloads + that return large results. Mapping deferrable operators + amplifies the problem further. IT sidesteps triggerers entirely — iterations + execute on workers, which scale more effectively and support custom XCom + backends. + + A custom XCom backend is not the whole story for IT, though: each item's result + is also written to its checkpoint in the task state store, so that a retry can + replay it instead of running the item again. Those checkpoints land in the + ``task_state_store`` table of the metadata database and stay there for the + store's retention (``[state_store] default_retention_days``, 30 days unless + configured) unless a ``[state_store] state_store_backend`` is configured, in Review Comment: Right, `state_store_backend` lives under `[workers]`; both the docs page and the class docstring say so now (5f7a4db3cc). Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting ########## task-sdk/src/airflow/sdk/bases/decorator.py: ########## @@ -726,8 +743,69 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> XComArg: start_trigger_args=self.operator_class.start_trigger_args, start_from_trigger=self.operator_class.start_from_trigger, returns_dag_result=self.returns_dag_result, + register_with_dag=register_with_dag, + ) + + def iterate(self, **map_kwargs: OperatorExpandArgument) -> XComArg: + """ + Iterate the task over ``map_kwargs`` inside a single task instance. + + The counterpart of :meth:`expand` for Iterable Tasks: the same inputs, but processed by one + :class:`~airflow.sdk.definitions.iterableoperator.IterableOperator` instead of one task + instance per item. + """ + if self.kwargs.get("trigger_rule") == TriggerRule.ALWAYS and any( + [isinstance(expanded, XComArg) for expanded in map_kwargs.values()] + ): + raise ValueError( + "Task-generated iterating within a task using 'iterate' is not allowed with trigger rule 'always'." + ) + if not map_kwargs: + raise TypeError("no arguments to iterate against") + self._validate_arg_names("iterate", map_kwargs) Review Comment: Agreed. Both validators now run the `is_mappable` check for `iterate` as well, with the method named in the error, so `.iterate(x=5)`, `x="abc"` and `x=None` fail at parse time as `.expand()` does, and `Source.from_argument` refuses a scalar instead of wrapping it; the "checked at parse time" comment holds again (730ef65ef8). Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting ########## task-sdk/src/airflow/sdk/definitions/iterableoperator.py: ########## @@ -0,0 +1,1263 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +import hashlib +import json +import os +import threading +import time +import warnings +from collections.abc import AsyncIterable, AsyncIterator, Callable, Iterable, Mapping, Sequence +from functools import partial +from typing import TYPE_CHECKING, Any + +try: + # Python 3.11+ + BaseExceptionGroup +except NameError: + from exceptiongroup import BaseExceptionGroup + +from airflow.sdk import BaseXCom, TaskInstanceState, TriggerRule +from airflow.sdk.bases.operator import BaseAsyncOperator, BaseOperator, event_loop +from airflow.sdk.bases.skipmixin import SkipMixin +from airflow.sdk.bases.xcom import XComIterable +from airflow.sdk.definitions.retry_policy import RetryAction, RetryDecision +from airflow.sdk.definitions.xcom_arg import XComArg +from airflow.sdk.exceptions import ( + AirflowFailException, + AirflowRescheduleException, + AirflowSensorTimeout, + AirflowSkipException, + AirflowTaskTerminated, + AirflowTaskTimeout, + DagRunTriggerException, + DownstreamTasksSkipped, + TaskDeferred, +) +from airflow.sdk.execution_time.comms import DeadlockImminentError +from airflow.sdk.execution_time.context import context_update_for_unmapped +from airflow.sdk.execution_time.executor import AsyncAwareExecutor +from airflow.sdk.execution_time.task_runner import ( + IndexedTaskInstance, + IndexedTaskRunner, + IndexedTaskState, + _push_xcom_if_needed, +) +from airflow.sdk.serde import serialize + +if TYPE_CHECKING: + import jinja2 + + from airflow.sdk.definitions._internal.expandinput import ExpandInput, Resolved + from airflow.sdk.definitions.context import Context + from airflow.sdk.definitions.mappedoperator import MappedOperator + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + from airflow.sdk.types import Logger + + +# The trigger rules under which one skipped upstream task instance skips a task, whatever the other +# upstream task instances did (see TriggerRuleDep). A skipped iteration has the same effect on a +# downstream task with one of these rules as a skipped mapped task instance would. +SKIPPED_WITH_A_SKIPPED_UPSTREAM = frozenset( + {TriggerRule.ALL_SUCCESS, TriggerRule.NONE_SKIPPED, TriggerRule.ALL_DONE_MIN_ONE_SUCCESS} +) + + +# Raised by an indexed task, these fail the task without a retry, as the runner does for a task that raises +# them itself (see _run_task_and_map_outcome and _handle_handler_failure). +FAIL_WITHOUT_RETRY = (AirflowFailException, AirflowSensorTimeout, AirflowTaskTerminated) + + +# How strongly a retry policy decision speaks for the task when several indexed tasks failed: one the +# policy says must not be retried fails the task, one it says to retry makes it retry on its terms. +class Checkpoints: + """ + Decide whether one attempt of an IterableOperator may resume from its per-index checkpoints. + + Checkpoints are only consulted from the second attempt onwards. A completion marker left by a + previous fully successful run means this attempt follows a manual clear (which raises + ``max_tries`` but does not reset ``try_number``), so every index must run again: the stale + ``SUCCESS`` checkpoints are ignored and overwritten. On entry the marker is replaced by the + attempt the rerun starts at, so that a crash during the rerun resumes from the checkpoints + written since, and only from those: an index the crashed rerun did not reach still holds its + checkpoint from before the clear, which must not be replayed. The marker is written again when + the block exits without an exception, and when every iteration skipped: the task is then + ``SKIPPED``, a final state, and a clear of it must run every iteration again rather than + replay the skips, as a cleared mapped task instance would. An attempt that fails, even with + some iterations skipped, writes no marker, so a retry or a clear after it resumes: iterations + that succeeded or skipped keep that outcome and only the others run again. + + The checkpoints themselves are never deleted: one marker write costs the same whatever the indexed task + count, and the store is scoped to the parent task instance, so state a sub-task stored for itself + is never touched. They expire with the store's default retention (``[state_store] + default_retention_days``, 30 days unless configured, 0 disables expiry), which also bounds how + long a task that exhausted its retries keeps them. Keeping them until then is intended: a manual + clear of such a task resumes from the checkpoints instead of re-running every index, and the + marker is what tells a clear-after-success apart from that. + """ + + # Same namespace as IndexedTaskState.build_key, for the same reason. + COMPLETION_KEY = "_iterable_completed" + + def __init__(self, context: Context) -> None: + self._store = context["task_state_store"] + self._try_number = context["ti"].try_number + self.trust_checkpoints = False + # The attempt from which checkpoints may be resumed; older ones predate a manual clear. + self.since = 0 + + def __enter__(self) -> Checkpoints: + if self._try_number > 1: + marker = self._store.get(self.COMPLETION_KEY) + if isinstance(marker, Mapping) and marker.get("completed"): + self._store.set(self.COMPLETION_KEY, {"completed": False, "since": self._try_number}) + else: + self.trust_checkpoints = True + since = marker.get("since") if isinstance(marker, Mapping) else None + if isinstance(since, int): + self.since = since + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + if exc_type is None or issubclass(exc_type, AirflowSkipException): + self._store.set(self.COMPLETION_KEY, {"completed": True, "try_number": self._try_number}) + + +class IndexedTaskInstanceNotStarted(Exception): + """ + The outcome of an indexed task instance pulled before a kill and reached after it. + + Never raised: ``IterableOperator._run_task`` returns it as the instance's outcome, so that + :meth:`IndexedTaskOutcomes.record` counts it apart from the failures. It ran no code, wrote + no checkpoint and fires no callback; the next attempt runs it. + """ + + +class IterationState: + """ + The state of one run of an iterated task, kept apart from the operator's configuration. + + It holds what the run needs to remember while it is going: the sub-operators in flight and + those already killed, so that :meth:`IterableOperator.on_kill` reaches each one once; the + stop flag a kill sets, which the executor consults before starting the next indexed task; + and the resolved input. What the indexed tasks ended with is :class:`IndexedTaskOutcomes`. + + A deep copy of the operator (``dag.partial_subset``, ``prepare_for_execution``) is another Review Comment: You are right, `prepare_for_execution` is a `copy.copy`, and the probe showed the next attempt starting stopped. `IterableOperator.__copy__` now gives the copy a fresh `IterationState` (and `IterationState.__copy__` matches `__deepcopy__`); the test kills one prepared copy and checks the next one and the Dag's operator are untouched (b4ada44135). Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
