kaxil commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4262230324
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -950,6 +985,623 @@ def mark_success_url(self) -> str:
return self.log_url
+class CheckpointDeserializationError(Exception):
+ """An indexed task's checkpoint holds a value serde refuses to restore,
such as a class off the allow-list."""
+
+
+@dataclass
+class IndexedTaskState:
+ status: TaskInstanceState
+ result: Any | None = None
+ # Outlet asset events the indexed task recorded on a previous successful
attempt. Outlet events
+ # only reach the server on the parent's success payload
(_handle_current_task_success), so an
+ # attempt that fails never registers what its succeeded indexed tasks
emitted. An indexed task skipped
+ # on retry (because it already succeeded) never re-executes, so it never
re-emits into the
+ # fresh OutletEventAccessors created for the new attempt; persisting a
snapshot here lets
+ # IterableOperator._run_indexed_task replay it, which is the only way
those events survive, and it
+ # cannot double-emit because the failed attempt sent nothing.
+ outlet_events: list[dict[str, Any]] | None = None
+ # Digest of the input the indexed task ran with. An index only means the
same work while its input
+ # is the same, so a checkpoint is honoured on a later attempt only when
this still matches:
+ # after the upstream was cleared and produced other values, the indexed
task runs again.
+ fingerprint: str | None = None
+ # The attempt that wrote the checkpoint. After a manual clear only the
checkpoints written since
+ # are resumed from (see Checkpoints), which this tells apart from those
left by the run before.
+ try_number: int = 0
+ # The other XComs the indexed task pushed, by their key before the index
suffix. The runner deletes
+ # every XCom before an attempt, so an indexed task skipped on retry has
them pushed again from here,
+ # next to its result, as a mapped task instance that succeeded keeps its
own.
+ xcoms: dict[str, Any] | None = None
+
+ def record_outlet_events(self, accessors: OutletEventAccessors) -> None:
+ """
+ Keep a JSON-safe snapshot of the outlet asset events the indexed task
recorded.
+
+ Persisted with the checkpoint so a later attempt can
:meth:`replay_outlet_events` when
+ the indexed task is skipped because it already succeeded.
+ """
+ events: list[dict[str, Any]] = []
+ for _asset_or_alias, accessor in accessors.items():
+ if isinstance(accessor.key, AssetUniqueKey):
+ events.append(
+ {
+ "kind": "asset",
+ "name": accessor.key.name,
+ "uri": accessor.key.uri,
+ "extra": accessor.extra,
+ "partition_keys": sorted(accessor.partition_keys),
+ }
+ )
+ for alias_event in accessor.asset_alias_events:
+ events.append(
+ {
+ "kind": "asset_alias",
+ "source_alias_name": alias_event.source_alias_name,
+ "dest_asset_key": {
+ "name": alias_event.dest_asset_key.name,
+ "uri": alias_event.dest_asset_key.uri,
+ },
+ "dest_asset_extra": alias_event.dest_asset_extra,
+ "extra": alias_event.extra,
+ }
+ )
+ if events:
+ self.outlet_events = events
+
+ def replay_outlet_events(self, target: OutletEventAccessorsProtocol) ->
None:
+ """
+ Re-emit into ``target`` the events the indexed task recorded on a
previous attempt.
+
+ An indexed task skipped on retry (because it already succeeded) never
re-executes, so it
+ never re-emits into the fresh ``OutletEventAccessors`` created for the
new attempt. The
+ failed attempt sent nothing to the server either (outlet events travel
only on the success
+ payload), so replaying cannot emit an event twice; without it the
events would be lost.
+ They land in the parent's accessors as
:meth:`IndexedTaskRunner.merge_outlet_events_into`
+ lands the live ones: one event per asset per task instance.
+ """
+ for event in self.outlet_events or ():
+ if event["kind"] == "asset":
+ accessor = target[Asset(name=event["name"], uri=event["uri"])]
+ accessor.extra.update(event["extra"])
+ if event["partition_keys"]:
+ accessor.add_partitions(event["partition_keys"])
+ else:
+
target[AssetAlias(name=event["source_alias_name"])].asset_alias_events.append(
+ AssetAliasEvent(
+ source_alias_name=event["source_alias_name"],
+
dest_asset_key=AssetUniqueKey(**event["dest_asset_key"]),
+ dest_asset_extra=event["dest_asset_extra"],
+ extra=event["extra"],
+ )
+ )
+
+ @staticmethod
+ def build_key(index: int) -> str:
+ # The task state store is already scoped to the parent task instance
(dag, run, task and
+ # map index), so the key carries no identity, only a namespace that
keeps the operator's own
+ # entries apart from anything user code stores from inside an indexed
task.
+ return f"_iterable_{index}"
+
+ def serialize(self) -> dict[str, Any]:
+ # The checkpoint travels to the supervisor as a JsonValue, which
rejects anything that is not
+ # plain JSON (tuples, datetimes, models, ...). Serde turns those into
JSON-compatible
+ # structures and restores them on read, exactly as XCom does with the
same result.
+ data: dict[str, Any] = {"status": self.status.value}
+ if self.result is not None:
+ data["result"] = serde_serialize(self.result)
+ if self.outlet_events:
+ data["outlet_events"] = self.outlet_events
+ if self.fingerprint:
+ data["fingerprint"] = self.fingerprint
+ if self.try_number:
+ data["try_number"] = self.try_number
+ if self.xcoms:
+ data["xcoms"] = {key: serde_serialize(value) for key, value in
self.xcoms.items()}
+ return data
+
+ @classmethod
+ def deserialize(cls, raw: Any) -> IndexedTaskState | None:
+ """
+ Restore a checkpoint, or None when there is none.
+
+ A result or XCom value serde refuses to restore, a class off the
allow-list
+ (``ImportError``) or one it has no deserializer for (``TypeError``),
raises
+ :class:`CheckpointDeserializationError`: refused once, it is refused
on every retry.
+ """
+ if not isinstance(raw, Mapping):
+ return None
+ try:
+ result = serde_deserialize(raw.get("result"))
+ xcoms = (
+ {key: serde_deserialize(value) for key, value in
raw["xcoms"].items()}
+ if raw.get("xcoms")
+ else None
+ )
+ except (ImportError, TypeError) as e:
Review Comment:
A pydantic result goes through `model_validate`, which raises
`ValidationError` (a `ValueError`) when the model changed between attempts,
e.g. a new required field. That gets past this `except`, so the item that
succeeded fails every retry and every clear again until the checkpoint expires.
`raw` is already fetched, so nothing transient can fail in these two calls and
catching `Exception` here looks safe. The test only covers the `ImportError`
path; parametrizing it with a changed pydantic model and a class whose
constructor raises `TypeError` would cover the rest.
##########
task-sdk/src/airflow/sdk/definitions/_internal/contextmanager.py:
##########
@@ -38,8 +39,17 @@
# the `get_current_context` function.
_CURRENT_CONTEXT: list[Context] = []
+# The contexts of the iterations of an iterated task, pushed on top of the
task's own context.
+# Iterations run concurrently in one process, in worker threads and asyncio
tasks, so each one's
+# context lives in a ContextVar and is seen by its own execution only. The
task's context stays in
+# the module-level list above, where any thread of the process finds it,
including one the task
+# started itself, whose ContextVars start out empty.
+_INDEXED_CONTEXT: ContextVar[tuple[Context, ...]] =
ContextVar("_indexed_context", default=())
+
def _get_current_context() -> Context:
+ if indexed := _INDEXED_CONTEXT.get():
Review Comment:
The PR removes `test_get_current_context_without_context`, which was the
only test of the "no context was found" error below, and it still passes if put
back. Could it stay, with a case showing an indexed context wins and that
resetting it falls back to `_CURRENT_CONTEXT`?
##########
airflow-core/tests/unit/models/test_taskinstance.py:
##########
@@ -4132,6 +4133,350 @@ def show(a, b):
dag_maker.run_ti(ti.task_id, map_index=ti.map_index,
dag_run=dag_run, session=session)
assert outputs == [(2, 5), (2, 10), (4, 5), (4, 10), (8, 5), (8, 10)]
+ def test_iterate_literal_cross_product(self, dag_maker, session):
+ """Test an iterated task with literal cross product args properly."""
+ outputs = []
+
+ with dag_maker(dag_id="product_same_types", session=session,
serialized=True) as dag:
+
+ @dag.task
+ def show(a, b):
+ outputs.append((a, b))
+
+ show.iterate(a=[2, 4, 8], b=[5, 10])
+
+ dag_run = dag_maker.create_dagrun()
+
+ show_task = dag.get_task("show")
+ with pytest.raises(NotMapped):
+ show_task.get_parse_time_mapped_ti_count()
+ with pytest.raises(NotMapped):
+ expand_mapped_task_instances(show_task, dag_run.run_id,
session=session)
+
+ tis = session.scalars(
+ select(TaskInstance)
+ .where(
+ TaskInstance.dag_id == dag.dag_id,
+ TaskInstance.task_id == "show",
+ TaskInstance.run_id == dag_run.run_id,
+ )
+ .order_by(TaskInstance.map_index)
+ ).all()
+ for ti in tis:
+ ti.refresh_from_task(show_task)
+ dag_maker.run_ti(ti.task_id, map_index=ti.map_index,
dag_run=dag_run, session=session)
+ assert outputs == [(2, 5), (2, 10), (4, 5), (4, 10), (8, 5), (8, 10)]
Review Comment:
The six sync items append from the executor's worker threads, so this exact
order isn't guaranteed the way it was for the `.expand()` version above. The
sibling `.iterate()` tests compare `sorted(...)`; same here?
##########
task-sdk/docs/mapped-tasks-vs-iterable-tasks.rst:
##########
@@ -0,0 +1,592 @@
+ .. 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.
+
+Inside that one Task Instance, the *iterated task*, the operator runs once per
item of the
+iterable. Each of those runs is an iteration, or *indexed task*: it has the
item's index, its own
+view of the context and its own XComs and checkpoint, but no Task Instance of
its own in the
+scheduler.
+
+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>`.
+
+.. warning::
+
+ An async task under ``iterate()`` should make async SDK calls. The
iterations share the task's
+ channel to the supervisor and its event loop: a synchronous call made on
the loop, by the task,
+ a hook it uses or a sync callback, works, but it first waits for another
iteration's call in
+ flight and then holds up every iteration for its round trip. A sync task
is not concerned: its
+ iterations run in worker threads, where a synchronous call waits its turn
without holding up
+ the others.
+
+ .. list-table::
+ :header-rows: 1
+
+ * - Instead of
+ - Use in an async task
+ * - ``Variable.get`` / ``Variable.set``
+ - ``await Variable.aget`` / ``await Variable.aset``
+ * - ``BaseHook.get_connection``
+ - ``await BaseHook.aget_connection``
+ * - ``BaseHook.get_hook``
+ - ``await BaseHook.aget_hook``
+ * - ``ti.xcom_pull`` / ``ti.xcom_push``
+ - ``await ti.axcom_pull`` / ``await ti.axcom_push``
+ * - ``XCom.get_one``
+ - ``await XCom.aget_one``
+
+.. warning::
+
+ An async task that raises ``AirflowTaskTimeout`` itself is taken for the
task's own
+ ``execution_timeout`` running out. Async iterations run on the main
thread, where the signal of
+ that timeout lands, so the two cannot be told apart: the iteration stops,
the iterations still
+ running are killed, and the task ends as on its own timeout, retried when
it has retries left.
+ A sync task is not concerned: an ``AirflowTaskTimeout`` raised in an
iteration's worker thread
+ is that iteration's own failure, and the others go on. For a time limit of
the operator's own,
+ such as one converted from ``asyncio.timeout``, raise another exception,
for example
+ ``TimeoutError``.
+
+.. 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.
+
+.. note::
+
+ The order of the arguments decides how often an upstream is read.
``iterate(a=..., b=...)``
+ combines them as a cross product, the last one varying fastest, and runs
the iterations in
+ that order: each element of ``a`` is read once for the run of iterations
that use it, while
+ ``b`` is read for every iteration. Reading an upstream's XCom is a round
trip to the
+ supervisor per element, so put the longest upstream first and literal
lists last:
+ ``iterate(url=upstream.output, retries=[1, 2, 3])`` reads each URL once,
where
+ ``iterate(retries=[1, 2, 3], url=upstream.output)`` reads each of them
three times. No more
+ than the element in use is kept per argument, so a large upstream is never
held whole.
Review Comment:
This holds when `upstream` is mapped or iterated, where each element is its
own read. For an unmapped upstream that returns a list, like `list_pokemon()`
above, `PlainXComArg.aresolve` pulls the whole list in one call and keeps it in
memory, so the argument order changes nothing there. Could the note say it
applies to mapped or iterated upstreams?
##########
task-sdk/src/airflow/sdk/definitions/iterableoperator.py:
##########
@@ -0,0 +1,1507 @@
+#
+# 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
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable,
Iterable, Mapping, Sequence
+from dataclasses import dataclass
+from functools import cached_property, partial, wraps
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from airflow.sdk import BaseXCom, TaskInstanceState, TriggerRule
+from airflow.sdk.bases.operator import 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.context import context_update_for_unmapped
+from airflow.sdk.execution_time.executor import AsyncAwareExecutor
+from airflow.sdk.execution_time.task_runner import (
+ CheckpointDeserializationError,
+ 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 outcome an indexed task's step returns, kept as it is by
keeps_failure_to_its_indexed_task.
+Outcome = TypeVar("Outcome")
+
+# 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)
+
+
+def keeps_failure_to_its_indexed_task(
+ step: Callable[..., Awaitable[Outcome]],
+) -> Callable[..., Awaitable[Outcome | tuple[IndexedTaskInstance, None,
BaseException]]]:
+ """
+ Make what an indexed task's step raises that indexed task's outcome, not
the end of the task.
+
+ ``step`` is a coroutine method of :class:`IterableOperator` taking the
indexed task as its first
+ argument after ``self``. Whatever it raises, a ``BaseException`` included,
is returned as
+ ``(task, None, raised)``: the siblings go on, the run counts it with the
other outcomes, and a
+ non-``Exception`` still fails the task at once (see
``IterableOperator._fail_fast_for``). A
+ cancellation and the parent's ``execution_timeout`` are not the step's
outcome and go on unchanged.
+ """
+
+ @wraps(step)
+ async def kept(
+ self: IterableOperator, task: IndexedTaskInstance, *args: Any,
**kwargs: Any
+ ) -> Outcome | tuple[IndexedTaskInstance, None, BaseException]:
+ try:
+ return await step(self, task, *args, **kwargs)
+ except (asyncio.CancelledError, AirflowTaskTimeout):
+ raise
+ except BaseException as e:
+ return task, None, e
+
+ return kept
+
+
+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 attempt whose iterations all finished means this attempt follows
a manual clear when
+ ``max_tries`` changed since: a clear sets it to ``try_number + retries``,
a retry leaves it
+ alone, and it can go down when the retries were lowered, so any change
counts. Every index then
+ runs again: the stale ``SUCCESS`` checkpoints are ignored and overwritten.
With ``max_tries``
+ unchanged, the iterations finished but the task failed afterwards (the
worker died, or
+ ``post_execute`` or the result push failed), and this retry resumes from
the checkpoints that
+ attempt could trust; only a clear that leaves ``max_tries`` exactly as it
was (the retries
+ lowered by as many attempts as ran since) is taken for such a retry. A
marker without
+ ``max_tries`` counts as a clear. On entry after a clear 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 an
indexed 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._max_tries = context["ti"].max_tries
+ 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")
+ and marker.get("max_tries") != self._max_tries
+ ):
+ self.since = self._try_number
+ self._store.set(self.COMPLETION_KEY, {"completed": False,
"since": self.since})
+ 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,
+ "max_tries": self._max_tries,
+ "since": self.since,
+ },
+ )
+
+
+class IndexedTaskInstanceNotStarted(Exception):
+ """
+ The outcome of an indexed task instance pulled before a kill and reached
after it.
+
+ Never raised: ``IterableOperator._run_indexed_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.
+ """
+
+
+@dataclass(frozen=True)
+class UnreadIndex:
+ """
+ An index whose input could not be read or rendered, in place of its
indexed task instance.
+
+ A ``.map()`` callable runs, and template fields are rendered, when an
index is read: what
+ they raise is that index's outcome, as it is the mapped task instance's
under ``.expand()``.
+ A skip leaves the index out of the result, anything else fails it, and the
siblings go on.
+ No operator was created for it, so unlike a mapped task instance it fires
no callback; nor
+ does it write a checkpoint, so the next attempt reads its input again.
+ """
+
+ task_id: str
+ index: int
+
+
+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 indexed
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 copy of the operator (``copy.copy`` in ``prepare_for_execution``, the
deep copy of
+ ``dag.partial_subset``) is another task with nothing in flight and no kill
pending, so copying
+ the state gives a fresh one. The
+ operator renews it once a run ended, not when one starts: a rerun in the
same process starts
+ clean, while a kill that arrives before the run sets the stop flag on the
state the run uses.
+ """
+
+ def __init__(self) -> None:
+ # Keyed by identity: BaseOperator equality compares fields such as
task_id, which every
+ # indexed operator of one iterated task shares, so a set would hold
one of them at most.
+ self._in_flight: dict[int, BaseOperator] = {}
+ self._lock = threading.Lock()
+ # The runner calls on_kill() again after the execution timeout that
made _run_tasks call
+ # it first, and an indexed operator is killed once.
+ self._killed: set[int] = set()
+ # A plain flag and a plain list, not an Event and not guarded by the
lock: on_kill() runs in
+ # the runner's SIGTERM handler, on the main thread, between two
bytecodes of whatever the
+ # loop thread was doing, register() or unregister() under the lock
included. A lock the
+ # handler takes that its own thread holds would hang it, so
request_stop() and start_kill()
+ # take none; take_in_flight() is called from the kill thread instead.
+ self._stop_requested = False
+ # The threads on_kill() started to kill the indexed operators in
flight; see await_kill.
+ self._kill_threads: list[threading.Thread] = []
+ #: The input resolved for this task instance, once ``aresolve``
returned.
+ self.resolved: Resolved | None = None
+ #: Whether the operator the task iterates pushes its results, as built
for an indexed task:
+ #: an operator may set ``do_xcom_push`` in its own ``__init__``, not
only through
+ #: ``partial()``. The same for every indexed task; True until one is
built.
+ self.do_xcom_push = True
+
+ def __copy__(self) -> IterationState:
+ return IterationState()
+
+ def __deepcopy__(self, memo: dict[int, Any]) -> IterationState:
+ return IterationState()
+
+ def register(self, operator: BaseOperator) -> None:
+ """Note that ``operator`` is executing, so a kill reaches it."""
+ with self._lock:
+ self._in_flight[id(operator)] = operator
+
+ def unregister(self, operator: BaseOperator) -> None:
+ """Note that ``operator`` is done, one way or another."""
+ with self._lock:
+ self._in_flight.pop(id(operator), None)
+
+ def __contains__(self, operator: object) -> bool:
+ with self._lock:
+ return id(operator) in self._in_flight
+
+ def take_in_flight(self) -> list[BaseOperator]:
+ """Return the indexed operators in flight that were not handed out
before, and mark them killed."""
+ with self._lock:
+ operators = [op for key, op in self._in_flight.items() if key not
in self._killed]
+ self._killed.update(map(id, operators))
+ return operators
+
+ def request_stop(self) -> None:
+ """Ask the iteration to start nothing else; see
:meth:`stop_requested`. Takes no lock."""
+ self._stop_requested = True
+
+ def start_kill(self, kill: Callable[[], None]) -> None:
+ """
+ Run ``kill`` in a thread of its own, kept so that :meth:`await_kill`
can wait for it.
+
+ Takes no lock, so it can be called from a signal handler; ``kill``
takes the indexed operators
+ in flight itself, with :meth:`take_in_flight`, on its thread.
+ """
+ thread = threading.Thread(target=kill,
name="iterable-operator-on-kill", daemon=True)
+ self._kill_threads.append(thread)
+ thread.start()
+
+ def await_kill(self, timeout: float) -> None:
+ """
+ Wait up to ``timeout`` seconds for the threads :meth:`start_kill`
started.
+
+ Called by ``IterableOperator._run_tasks`` once the loop closed and
before the run
+ concludes, so the task does not end, and the process with it, while an
indexed operator's
+ ``on_kill`` is still cleaning up. The threads are daemon threads: one
that never returns
+ holds the run for ``timeout`` at most, and the supervisor's SIGKILL
bounds the rest.
+ """
+ deadline = time.monotonic() + timeout
+ for thread in list(self._kill_threads):
+ thread.join(max(0.0, deadline - time.monotonic()))
+
+ def stop_requested(self) -> bool:
+ """Whether :meth:`request_stop` was called; passed to the executor as
its ``stop``."""
+ return self._stop_requested
+
+ @property
+ def length(self) -> int | None:
+ """How many indexed tasks the resolved input gives, or None while it
is still being resolved."""
+ return self.resolved.length if self.resolved is not None else None
+
+
+class IndexedTaskOutcomes:
+ """
+ What the indexed tasks of one run ended with, and what the task ends with
because of it.
+
+ Entered by ``IterableOperator._run_tasks`` around the loop that drives the
indexed tasks.
+ :meth:`record` takes each outcome: a skip is noted, an outcome the
iteration cannot carry
+ raises at once, any other failure is logged and collected.
:meth:`conclude` turns what was
+ collected, and whether the run's :class:`IterationState` was asked to
stop, into the task's
+ own outcome and raises it: killed, failed on the exception the runner
judges, skipped over
+ an empty input or because every indexed task skipped. Whatever
+ exception leaves the block, those included, the failed indexed tasks'
callbacks are reported
+ on exit with the task's fate, so they say what the runner then does:
retry, or fail for good.
+ """
+
+ #: How strongly a retry policy decision speaks for the task when several
indexed tasks failed:
+ #: one the policy fails the task on outweighs one it retries on, which
outweighs the default.
+ _DECISION_WEIGHT = {RetryAction.FAIL: 2, RetryAction.RETRY: 1,
RetryAction.DEFAULT: 0}
+
+ def __init__(self, operator: IterableOperator, state: IterationState,
context: Context) -> None:
+ self._operator = operator
+ self._state = state
+ self._context = context
+ self.total = 0
+ #: Pulled before a kill and never started (see
:class:`IndexedTaskInstanceNotStarted`).
+ self.not_started = 0
+ self.exceptions: list[BaseException] = []
+ self.skipped: dict[int, AirflowSkipException] = {}
+ self._failed_runners: list[IndexedTaskRunner] = []
+ self._decision: tuple[BaseException, RetryDecision] | None = None
+
+ def __enter__(self) -> IndexedTaskOutcomes:
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback) -> None:
+ # Whatever the task ends with decides every failed indexed task's
callback, so they agree.
+ if exc_value is not None and self._failed_runners:
+ self._report(exc_value)
+
+ def note_failed(self, runner: IndexedTaskRunner) -> None:
+ """Remember a failed indexed task's runner: its callback waits for the
task's fate."""
+ self._failed_runners.append(runner)
+
+ @property
+ def log(self) -> Logger:
+ """The operator's logger: the outcomes are logged under the task they
belong to."""
+ return self._operator.log
+
+ @property
+ def failed_runners(self) -> tuple[IndexedTaskRunner, ...]:
+ """The runners of the indexed tasks that failed in this run, in the
order they failed."""
+ return tuple(self._failed_runners)
+
+ def record(self, task: IndexedTaskInstance | UnreadIndex, raised:
BaseException | None) -> None:
+ """Take one indexed task's outcome: nothing, a skip, not started, a
failure, or one ending it."""
+ if isinstance(raised, IndexedTaskInstanceNotStarted):
+ self.not_started += 1
+ return
+ self.total += 1
+ if raised is None:
+ return
+ if isinstance(raised, AirflowSkipException):
+ self.skipped[task.index] = raised
+ return
+ # An outcome the iteration cannot carry stops it at once: every later
indexed task would
+ # end the same way.
+ if (fail_fast := self._fail_fast_for(task, raised)) is not None:
+ raise fail_fast from raised
+ self.log.exception(
+ "An exception occurred for task_id %s with index %s",
task.task_id, task.index, exc_info=raised
+ )
+ self.exceptions.append(raised)
+
+ def conclude(self) -> list[int]:
+ """
+ Raise the task's outcome from what was recorded, or return the skipped
indices.
+
+ A kill comes first: with every killed indexed task returning normally
there would be no
+ failure to raise, the completion marker would be written and a result
returned over
+ XComs that do not exist, and a kill while the input resolves is not an
empty input. The
+ task then fails without a retry, as the runner treats a terminated
task, and no marker
+ is written, so a later clear resumes from the checkpoints of the
indexed tasks that did
+ finish. Then the failures, handed to the runner as the exception it
judges; then an empty
+ input, skipped as a mapped task over nothing is; then every indexed
task skipped, which
+ skips the task too.
+ """
+ if self._state.stop_requested():
+ raise AirflowTaskTerminated(self._killed_message()) from (
+ BaseExceptionGroup("Indexed task failures", self.exceptions)
if self.exceptions else None
+ )
+ if self.exceptions:
+ raise self._pick_failure_for_runner()
+ if self.total == 0:
+ raise AirflowSkipException("The input to iterate over is empty.")
+ if self.skipped and len(self.skipped) == self.total:
+ raise next(iter(self.skipped.values()))
+ return sorted(self.skipped)
+
+ def _killed_message(self) -> str:
+ """
+ Say what a kill left behind, naming only the counts that are not zero.
+
+ Before the input was resolved there is no count to give. Once it is,
the indexed tasks
+ that ran (the killed ones among them, which came back as failures),
those pulled before
+ the kill and never started, and those never pulled add up to the
length.
+ """
+ length = self._state.length
+ if length is None:
+ return "The iterated task was killed before its input was
resolved."
+ parts = [f"{self.total} of {length} indexed tasks ran"]
+ if self.not_started:
+ parts.append(f"{self.not_started} pulled but never started")
+ if never_pulled := length - self.total - self.not_started:
+ parts.append(f"{never_pulled} never pulled")
+ return f"The iterated task was killed: {', '.join(parts)}."
+
+ def _report(self, raised: BaseException) -> None:
+ """
+ Report each failed indexed task as what happens to the task: retried,
or failed for good.
+
+ Their callbacks were held back until every indexed task had run, so
the retry callback of
+ one no longer announces a retry a sibling's ``AirflowFailException``
then rules out. They
+ run one after another, on the thread that ran the iteration.
+ """
+ task_will_retry = self._task_will_retry(raised)
+ for runner in self._failed_runners:
+ runner.report_failure(task_will_retry=task_will_retry)
+
+ def _task_will_retry(self, raised: BaseException) -> bool:
+ """
+ Whether the runner retries the task for ``raised``, by the runner's
own rules.
+
+ No retry for the fail-fast exceptions nor for what is not an
``Exception`` (other than an
+ ``AirflowTaskTimeout``), none when the retry policy decides FAIL, and
otherwise a retry
+ while the parent has attempts left
(``IndexedTaskInstance.is_eligible_to_retry``). A group
+ of failures is retried when each of them is: an indexed task's own
``AirflowTaskTimeout`` next to
+ another indexed task's ``ValueError`` makes a ``BaseExceptionGroup``,
which the runner retries like
+ any failure. The policy's decision is the one
:meth:`_pick_failure_for_runner` took for
+ ``raised`` when it chose it, so the policy is not evaluated again for
the callbacks: a
+ policy that calls a model may answer differently each time, and the
callbacks must say
+ what the exception handed over says.
+ """
+ if not self._retryable(raised):
+ return False
+ if not self._failed_runners[0].task_instance.is_eligible_to_retry:
+ return False
+ if (policy := self._operator.retry_policy) is not None:
+ if self._decision is not None and self._decision[0] is raised:
+ decision = self._decision[1]
+ else:
+ ti = self._context["ti"]
+ try:
+ decision = policy.evaluate(
+ exception=raised,
+ try_number=ti.try_number,
+ max_tries=ti.max_tries,
+ context=self._context,
+ )
+ except Exception:
+ return True
+ if decision.action == RetryAction.FAIL:
+ return False
+ return True
+
+ @classmethod
+ def _retryable(cls, raised: BaseException) -> bool:
+ """Whether ``raised`` is a failure the runner may retry, judged member
by member for a group."""
+ if isinstance(raised, BaseExceptionGroup):
+ return all(cls._retryable(exc) for exc in raised.exceptions)
+ return not isinstance(raised, FAIL_WITHOUT_RETRY) and isinstance(
+ raised, (Exception, AirflowTaskTimeout)
+ )
+
+ @staticmethod
+ def _fail_fast_for(
+ task: IndexedTaskInstance | UnreadIndex, raised: BaseException
+ ) -> AirflowFailException | None:
+ """
+ Return the failure that ends the task at once for an outcome the
iteration cannot carry.
+
+ An indexed task is not a task instance of its own: it cannot defer,
reschedule, trigger a
+ DAG run or skip downstream tasks, and a ``BaseException`` that is no
``Exception``
+ (``KeyboardInterrupt``, ``SystemExit``) must never be swallowed
+ into a retry. Each of these fails the whole task without a retry, with
a message saying
+ why. Any other exception, an ``AirflowTaskTimeout`` the operator
raised itself included, is
+ the indexed task's own failure, which :meth:`record` collects.
+ """
+ indexed_task_name = f"Indexed task {task.task_id}[{task.index}]"
+ if isinstance(raised, TaskDeferred):
+ return AirflowFailException(
+ f"{indexed_task_name} attempted to defer. Deferrable operators
are not supported inside IterableOperator."
+ )
+ if isinstance(raised, (DagRunTriggerException,
DownstreamTasksSkipped)):
+ return AirflowFailException(
+ f"{indexed_task_name} raised {type(raised).__name__}.
Triggering DAG runs (TriggerDagRunOperator) and "
+ "skipping downstream tasks (ShortCircuitOperator and similar)
are not supported inside "
+ "IterableOperator: an indexed task has no downstream tasks or
DAG run of its own for "
+ "the effect to apply to."
+ )
+ if isinstance(raised, AirflowRescheduleException):
+ return AirflowFailException(
+ f"{indexed_task_name} attempted to reschedule (raised
AirflowRescheduleException). Reschedule-mode "
+ "sensors are not supported inside IterableOperator: an indexed
task has no task "
+ "instance of its own to reschedule."
+ )
+ if isinstance(raised, AirflowTaskTimeout):
+ # The parent's timeout never gets here (_run_indexed_task
re-raises it); this one the operator
+ # raised itself, off the main thread, and the runner judges it as
it does for a plain
+ # task that raises it: a failure it retries.
+ return None
+ if not isinstance(raised, Exception):
+ return AirflowFailException(
+ f"{indexed_task_name} raised a non-Exception BaseException:
{type(raised).__name__}: {raised}"
+ )
+ return None
+
+ def _pick_failure_for_runner(self) -> BaseException:
+ """
+ Pick the exception the runner decides the task's outcome on.
+
+ The runner classifies by exception type: a fail-fast exception fails
without a retry, and a
+ ``retry_policy`` matches rules against the type. A
``BaseExceptionGroup`` defeats both, so
+ an indexed task's own exception is handed over whenever one decides:
the first fail-fast
+ one, the only one, or the one whose policy decision weighs most. With
several failures the
+ others stay attached as its cause, so every traceback reaches the log.
Several failures no
+ policy decides between are raised as a group, which the task's own
retries then apply to.
+ The policy is evaluated once per failure, and the decision for the
chosen one is kept for
+ the callbacks, as is the default for a group no decision outweighs (see
+ :meth:`_task_will_retry`).
+ """
+ exceptions = self.exceptions
+ group = BaseExceptionGroup("Multiple indexed task failures",
exceptions)
+ chosen: BaseException | None = next(
+ (exc for exc in exceptions if isinstance(exc,
FAIL_WITHOUT_RETRY)), None
+ )
+ if chosen is None and len(exceptions) == 1:
+ return exceptions[0]
+ if chosen is None and (policy := self._operator.retry_policy) is not
None:
+ ti = self._context["ti"]
+ max_tries = ti.max_tries
+ weights = []
+ decisions: list[RetryDecision | None] = []
+ for exc in exceptions:
+ try:
+ decision = policy.evaluate(
+ exception=exc, try_number=ti.try_number,
max_tries=max_tries, context=self._context
+ )
+ weights.append(self._DECISION_WEIGHT.get(decision.action,
0))
+ decisions.append(decision)
+ except Exception:
+ # As the runner does: a policy that fails to evaluate
leaves the default.
+ self.log.exception("Retry policy evaluation failed for an
indexed task failure")
+ weights.append(0)
+ decisions.append(None)
+ if max(weights) > 0:
+ index = weights.index(max(weights))
+ chosen = exceptions[index]
+ if (kept := decisions[index]) is not None:
+ self._decision = (chosen, kept)
+ else:
+ # Every decision was the default, or failed to evaluate, which
the runner treats
+ # the same: the group carries the default, so the callbacks
follow it instead of
+ # evaluating the policy once more, on an exception no indexed
task raised.
+ self._decision = (group, RetryDecision.default())
+ if chosen is None:
+ return group
+ if len(exceptions) > 1:
+ chosen.__cause__ = group
+ return chosen
+
+
+class IterableOperator(BaseOperator):
+ """
+ Operator used for Iterable Tasks (IT) that runs a mapped operator over an
iterable input.
+
+ The IterableOperator wraps a :class:`MappedOperator` together with an
+ :class:`ExpandInput` and is responsible for creating and running the
+ per-index runtime task instances. The IterableOperator itself participates
+ in Airflow's native retry mechanism — its ``retries`` and ``retry_delay``
+ are inherited from the wrapped operator so that when any indexed task needs
+ a retry the whole IterableOperator is retried by Airflow. Already-succeeded
+ indexed tasks are skipped on each retry attempt because their state is
+ checkpointed in the ``task_state_store``.
+
+ The IterableOperator executes the mapped operator instances using a
+ concurrent executor with a configurable number of workers. By default
+ the worker count is taken from the mapped operator's ``partial_kwargs``
+ (``task_concurrency``) if present, otherwise falls back to
+ ``os.cpu_count()`` and finally to ``1``. ``os.cpu_count()`` counts the
CPUs of
+ the machine: a worker in a container with a CPU limit still sees every CPU
of
+ its node, so set ``task_concurrency`` there.
+
+ **Crash recovery:** When the worker crashes mid-iteration and the task is
re-run (e.g. via a
+ manual clear), already-succeeded indexed tasks are skipped and only the
pending/failed ones are
+ executed again. Every indexed task inherits its ``try_number`` from the
IterableOperator's own task
+ instance, so the attempt count reported to an indexed task matches the
attempt Airflow is currently
+ running. The checkpoint is only consulted from the second attempt onwards,
and solely to decide
+ whether an index already succeeded. Once every index has succeeded, a
completion marker is written
+ so that a *subsequent* manual clear (which does not reset ``try_number``)
re-runs every index from
+ scratch instead of replaying the previous run's stale results (see
:class:`Checkpoints`). A
+ checkpoint carries the indexed task's full result (plus its extra XComs
and outlet events) and lives in
+ the ``task_state_store`` table for the store's retention, unless a
``[workers]
+ state_store_backend`` is configured and only a reference is stored; a
custom XCom backend alone
+ does not keep large results out of the metadata database.
+
+ :param operator: The :class:`MappedOperator` to unmap and execute for
+ each element of ``expand_input``. Each indexed task runs a new
instance of it,
+ built by ``unmap()`` from that index's kwargs; the partial kwargs are
shared
+ between the instances, not copied (see the warning on shared inputs
below).
+
+ :param expand_input: Provider of the values to iterate
+ over. Its ``aresolve(context)`` method gives the indexed task count
and the
+ per-index ``mapped_kwargs`` used to unmap the operator.
+
+ :param kwargs: Additional keyword arguments forwarded to
+ :class:`BaseOperator` when instantiating the IterableOperator
+ (e.g. ``dag``, ``start_date``).
+
+ :returns: An :class:`XComIterable` if the mapped operator pushes XComs,
otherwise ``None``.
+
+ .. note::
+ ``multiple_outputs`` is ignored for iterated tasks. Each indexed
task's return value is pushed
+ whole as ``return_value_<index>`` and the task's own return value is
the ``XComIterable``
+ over them, so a ``Mapping`` return annotation on the wrapped ``@task``
does not fan its
+ keys out into separate XComs the way it does for ``.expand()``.
+
+ .. note::
+ Deferred operators (those that raise
:class:`~airflow.sdk.exceptions.TaskDeferred`) are not
+ supported yet inside IterableOperator. A ``TaskDeferred`` exception
raised by an indexed task
+ instance will propagate as an error rather than pausing and resuming
the task.
+
+ Reschedule-mode sensors (those that raise
:class:`~airflow.sdk.exceptions.AirflowRescheduleException`)
+ are also not supported. A reschedule raised by an indexed task
instance will fail the whole
+ IterableOperator immediately with a clear error rather than being
silently mishandled.
+
+ Triggering DAG runs
(:class:`~airflow.sdk.exceptions.DagRunTriggerException`, raised by
+ ``TriggerDagRunOperator``) and skipping downstream tasks are not
supported either: an indexed task
+ index has no DAG run or downstream tasks of its own for the
trigger/skip to apply to.
+ Operators that can skip downstream tasks (``ShortCircuitOperator``,
the branch operators,
+ ``@task.short_circuit``, ``@task.branch`` and any other ``SkipMixin``)
are rejected by
+ ``.iterate()`` itself, since inside an iteration they would find
nothing to skip and let
+ every downstream task run. A trigger, or a
+ :class:`~airflow.sdk.exceptions.DownstreamTasksSkipped` raised anyway,
fails the whole
+ IterableOperator immediately with a clear error rather than silently
doing nothing.
+
+ Indexed task outcomes are classified before being aggregated: if any
indexed task raises
+ :class:`~airflow.sdk.exceptions.AirflowFailException`, that exception
is re-raised directly so
+ the IterableOperator fails without retrying. An indexed task that
raises
+ :class:`~airflow.sdk.exceptions.AirflowSkipException` is skipped, as a
mapped task instance
+ would be: it pushes no XCom, does not fail the task and is not run
again on a retry. It is
+ left out of the task's :class:`~airflow.sdk.bases.xcom.XComIterable`,
so downstream tasks
+ only see the values that exist. An indexed task that returns ``None``
is not skipped: it succeeds,
+ pushes no XCom and keeps its position in the sequence, which reads
``None`` there, so a
+ downstream ``.expand()`` over the result runs over it; ``.expand()``
leaves a mapped task
+ instance that returned ``None`` out of the count instead. A direct
downstream task whose trigger rule skips it when an
+ upstream task instance is skipped (``all_success``, ``none_skipped``,
+ ``all_done_min_one_success``) is skipped, as after a mapped upstream;
one with a rule such
+ as ``none_failed`` runs over the remaining values. If *every* indexed
task is skipped, a single
+ ``AirflowSkipException`` is re-raised so the IterableOperator itself
is marked ``SKIPPED``, and
+ so is it over an empty input, as a mapped task over nothing is. All
other exceptions of the
+ indexed tasks are aggregated into a :class:`BaseExceptionGroup` and
treated as a regular
+ retryable failure.
+
+ .. 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 of them: 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. A mapped task instance gets its own copy,
because it runs in its
+ own process; an iteration does not. Treat inputs as read-only, or copy
what the task
+ changes in place.
+
+ .. note::
+ **Callbacks run per indexed task, and a failed one's wait for the
task's fate.**
+
+ ``on_success_callback`` and ``on_skipped_callback`` run once the
indexed task's checkpoint is
+ written, so they speak for work a retry will not run again, where it
ran: in its worker
+ thread for a sync operator, on the event loop for an async one. A
failed indexed task's
+ ``on_failure_callback`` or ``on_retry_callback`` runs once every
indexed task has run, on the
+ thread that ran the iteration, and says what happens to the task:
retried or failed for good
+ (see :class:`IndexedTaskOutcomes`). A failure no indexed task
+ owns, such as an error resolving the input, fires no callback, and
neither does an indexed
+ task pulled before a kill and never started: the iterated task has no
callbacks of its own.
+ Listeners fire once, for the task instance, as for any task; an
indexed task is not a task
+ instance and fires none.
+
+ .. note::
+ **Pools count the task instance, not its iterations.**
+
+ The scheduler reserves ``pool_slots`` once for the iterated task,
while up to
+ ``task_concurrency`` iterations run inside it. A pool sized to cap the
load on a shared
+ resource (database connections, the rate limit of an API) therefore
sees one reservation
+ for that many concurrent uses. Choose ``task_concurrency`` with the
pool in mind; reserving
+ slots per iteration needs support in the scheduler, which does not
exist yet.
+
+ .. note::
+ **``map_index_template`` is not applied.**
+
+ Under ``.expand()`` the template names each mapped task instance after
its item. An
+ iterated task is one task instance with no item of its own, so the
wrapped operator's
+ template is left out: rendered against the task's context, a template
that reads an item's
+ kwargs (``{{ task.op_kwargs['x'] }}``) or a value an indexed task set
in its context would fail.
+
+ .. note::
+ **Async indexed tasks should prefer async SDK calls.**
+
+ IterableOperator runs multiple async indexed tasks concurrently on the
same event loop, each
+ making async SDK calls of its own (checkpointing, XCom push). If an
async indexed task's
+ ``aexecute()`` — or a hook/callback it calls — issues a *synchronous*
SDK call instead (e.g.
+ ``Variable.get``, ``BaseHook.get_connection``/``get_hook``,
``ti.xcom_pull``, or a sync
+ ``on_success_callback``/``pre_execute``), the call works but blocks
the event loop, and so
+ every indexed task on it, for its round trip to the supervisor,
waiting first for a sibling's
+ call in flight. Use the async equivalents inside async operators:
+ :meth:`~airflow.sdk.bases.hook.BaseHook.aget_connection`/``aget_hook``,
+ ``ti.axcom_pull`` and ``Variable.aget``/``aset``. Sync indexed tasks
are not concerned: they run
+ in worker threads, their ``execute``, hooks and callbacks included,
where a synchronous SDK
+ call waits its turn without holding up the loop; so does ``on_kill``
of the indexed operators,
+ which :meth:`on_kill` runs off the loop thread.
+
+ .. warning::
+ **``execution_timeout`` caps the whole iteration, not each indexed
task.**
+
+ The IterableOperator keeps the wrapped operator's
``execution_timeout`` as a wall-clock limit
+ on the entire task instance. The runner enforces it on the main thread
exactly as for any
+ other task, so an iteration that overruns fails with
``AirflowTaskTimeout`` and
+ :meth:`on_kill` is propagated to every indexed task still in flight.
Since ``.iterate()`` runs
+ all indexed tasks in one task instance, this is the per-instance limit
of ``.expand()`` applied
+ to the whole iteration rather than to each indexed task: no limit
applies per indexed task,
+ sync or async (one of the same value, started later, could never fire
first). An
+ ``AirflowTaskTimeout`` a sync indexed task raises itself (a hook that
gave up waiting) is that
+ indexed task's own failure, as it is the mapped task instance's under
``.expand()``.
+ """
+
+ _operator: MappedOperator
+ expand_input: ExpandInput
+ partial_kwargs: dict[str, Any]
+ shallow_copy_attrs: Sequence[str] = (
+ "_operator",
+ "expand_input",
+ "partial_kwargs",
+ "_log",
+ )
+
+ @staticmethod
+ def _refuse_operators_that_skip_downstream(operator: MappedOperator) ->
None:
+ """
+ Refuse to iterate an operator that can skip downstream tasks.
+
+ An iteration has no downstream tasks of its own, so
``ShortCircuitOperator``, the branch
+ operators and any other ``SkipMixin`` would skip nothing and let every
downstream task run.
+ Checked on the class: ``MappedOperator._can_skip_downstream`` is only
derived from ``SkipMixin``
+ on the classic path, while the ``@task`` path copies a class default
that is ``False`` even for
+ ``@task.short_circuit`` and ``@task.branch``.
+ """
+ if issubclass(operator.operator_class, SkipMixin):
+ raise TypeError(
+ f"{operator.operator_name} can skip downstream tasks and
cannot be iterated: an iteration "
+ f"of {operator.task_id!r} has no downstream tasks of its own,
so it would skip nothing and "
+ "every downstream task would run. Use .expand() for it
instead."
+ )
+
+ def __init__(
+ self,
+ *,
+ operator: MappedOperator,
+ expand_input: ExpandInput,
+ **kwargs,
+ ):
+ if operator.get_closest_mapped_task_group() is not None:
+ raise NotImplementedError("operator expansion in an expanded task
group is not yet supported")
+ self._refuse_operators_that_skip_downstream(operator)
+
+ super().__init__(
+ **{
+ **kwargs,
+ "task_id": operator.task_id_without_group_prefix,
+ "owner": operator.owner,
+ "email": operator.email,
+ "email_on_retry": operator.email_on_retry,
+ "email_on_failure": operator.email_on_failure,
+ "retries": operator.retries,
+ "retry_delay": operator.retry_delay,
+ "retry_exponential_backoff":
operator.retry_exponential_backoff,
+ "max_retry_delay": operator.max_retry_delay,
+ "retry_policy": operator.retry_policy,
+ "start_date": operator.start_date,
+ "end_date": operator.end_date,
+ "depends_on_past": operator.depends_on_past,
+ "ignore_first_depends_on_past":
operator.ignore_first_depends_on_past,
+ "wait_for_past_depends_before_skipping":
operator.wait_for_past_depends_before_skipping,
+ "wait_for_downstream": operator.wait_for_downstream,
+ "dag": operator.dag,
+ "params": operator.params,
+ "priority_weight": operator.priority_weight,
+ "weight_rule": operator.weight_rule,
+ "queue": operator.queue,
+ "pool": operator.pool,
+ "pool_slots": operator.pool_slots,
+ # Kept as the wall-clock cap on the whole iteration, enforced
by the runner (see the
+ # class docstring); no limit applies per indexed task.
+ "execution_timeout": operator.execution_timeout,
+ "trigger_rule": operator.trigger_rule,
+ "resources": operator.resources,
+ "run_as_user": operator.run_as_user,
+ # Not forwarded: the template names a mapped task instance
after its item, and the
+ # iterated task instance has no item of its own. Rendered for
it against the
+ # parent's context, a template that reads an item's kwargs or
a value one indexed task set
+ # would fail the task after every indexed task succeeded, and
every retry with it.
+ "map_index_template": None,
+ "max_active_tis_per_dag": operator.max_active_tis_per_dag,
+ "max_active_tis_per_dagrun":
operator.max_active_tis_per_dagrun,
+ "executor": operator.executor,
+ "executor_config": operator.executor_config,
+ "do_xcom_push": operator.partial_kwargs.get("do_xcom_push",
True),
+ # Ignored for iterated tasks (also when passed explicitly):
the return value pushed by the
+ # runner is the XComIterable aggregate, not a dict, and every
indexed task result is pushed
+ # whole under return_value_<index>. The wrapped @task may
still infer True from a
+ # Mapping return annotation, which would make the runner
reject the aggregate.
+ "multiple_outputs": False,
+ "inlets": operator.inlets,
+ "outlets": operator.outlets,
+ "task_group": operator.task_group,
+ "doc": operator.doc,
+ "doc_md": operator.doc_md,
+ "doc_json": operator.doc_json,
+ "doc_yaml": operator.doc_yaml,
+ "doc_rst": operator.doc_rst,
+ "task_display_name": operator.task_display_name,
+ "allow_nested_operators": operator.allow_nested_operators,
+ # The iterated task has no callbacks and no execute hooks of
its own: they run per
+ # indexed task, from the wrapped operator's partial kwargs.
Passed explicitly, since
+ # _apply_defaults would otherwise fill them from the DAG's
default_args and the
+ # runner would run them for the task on top of the indexed
tasks' own.
+ "on_execute_callback": None,
+ "on_success_callback": None,
+ "on_failure_callback": None,
+ "on_retry_callback": None,
+ "on_skipped_callback": None,
+ "pre_execute": None,
+ "post_execute": None,
+ }
+ )
+ self._operator = operator
+ self.expand_input = expand_input
+ self.partial_kwargs = dict(operator.partial_kwargs) if
operator.partial_kwargs else {}
+ task_concurrency = self.partial_kwargs.pop("task_concurrency", None)
+ if task_concurrency is not None and task_concurrency < 1:
+ raise ValueError(f"task_concurrency must be at least 1, got
{task_concurrency}")
+ # pool_slots is reserved once for the task instance, not per
iteration: see the class docstring.
+ self.max_workers = task_concurrency if task_concurrency is not None
else (os.cpu_count() or 1)
+ # unmap() would normally apply these three flags to each generated
indexed operator, and
+ # __attrs_post_init__ would apply them (plus the upstream-relationship
wiring below) to the
+ # MappedOperator itself; since IterableOperator skips
__attrs_post_init__ entirely (it isn't a
+ # MappedOperator), it must reproduce that part of the contract for its
own single DAG node.
+ self.is_setup = bool(self.partial_kwargs.get("is_setup", False))
+ self.is_teardown = bool(self.partial_kwargs.get("is_teardown", False))
+ on_failure_fail_dagrun =
self.partial_kwargs.get("on_failure_fail_dagrun", False)
+ if on_failure_fail_dagrun:
+ self.on_failure_fail_dagrun = on_failure_fail_dagrun
+ XComArg.apply_upstream_relationship(self, self.expand_input.value)
+ # Mirrors MappedOperator.__attrs_post_init__: partial kwargs
corresponding to the wrapped
+ # operator's own template fields may themselves be XComArgs (e.g.
`.partial(some_field=xcom)`),
+ # and those upstream edges must be recorded too, not just the ones
from expand_input.
+ for key, value in self.partial_kwargs.items():
+ if key in self._operator.template_fields:
+ XComArg.apply_upstream_relationship(self, value)
+ # What one run remembers while it is going (see IterationState); fresh
for every run and
+ # for every copy of the operator.
+ self._state = IterationState()
+
+ def __copy__(self) -> IterableOperator:
+ # prepare_for_execution copies the operator with copy.copy, which
would share the state of
+ # the run with the Dag's operator: a kill in one attempt would then
stop the next one under
+ # dag.test(). The copy is another run and starts with a state of its
own (see IterationState).
+ other = type(self).__new__(type(self))
+ other.__setstate__({**self.__getstate__(), "_state": IterationState()})
+ return other
+
+ # How long the run waits for its threads to end once it is over: the
executor's worker and
+ # coroutine shutdown, and the thread on_kill() kills the indexed operators
from.
+ _SHUTDOWN_TIMEOUT: float = 10.0
+
+ def on_kill(self) -> None:
+ # The default BaseOperator.on_kill() is a no-op, which would otherwise
leave every
+ # currently in-flight indexed task unaware that the IterableOperator
itself was killed
+ # (SIGTERM) or hit its execution_timeout: propagate to each active
indexed operator instead.
+ # First stop the iteration from starting anything else: the killed
indexed tasks come back
+ # as failures and free their slots, which would otherwise be filled
with the next ones.
+ state = self._state
+ state.request_stop()
+ # Always in a thread of its own, which takes the indexed operators in
flight itself. The
+ # runner's SIGTERM handler calls this on the main thread, between two
bytecodes of the loop
+ # thread: a lock taken here that register() or unregister() holds at
that moment would
+ # hang the handler, so nothing here takes one (see IterationState). On
that thread a
+ # synchronous SDK call in an indexed operator's on_kill would also
hold up the event loop, and
+ # every indexed task on it, for its round trip; in its own thread the
call waits its turn and
+ # the loop goes on serving the indexed tasks. _run_tasks kills what is
in
+ # flight through _kill directly, from a thread the loop drives, and
waits for this thread
+ # before the run concludes (see IterationState.await_kill).
+ state.start_kill(lambda: self._kill(state.take_in_flight()))
+
+ def _kill(self, operators: list[BaseOperator]) -> None:
+ # One indexed operator's on_kill must not keep the kill from the
others: whatever it raises, a
+ # BaseException included, is caught here as a plain error is.
+ for operator in operators:
+ try:
+ operator.on_kill()
+ except BaseException:
+ self.log.exception("Error calling on_kill() for indexed task
operator %s", operator.task_id)
+
+ def deserialization_allowed_values(self) -> list[Any]:
+ """Return the values the operator this task iterates declares: every
indexed task pushes what it does."""
+ return self._operator.deserialization_allowed_values()
+
+ @property
+ def returns_dag_result(self) -> bool:
+ return self._operator.returns_dag_result
+
+ @returns_dag_result.setter
+ def returns_dag_result(self, value: bool) -> None:
+ self._operator.returns_dag_result = value
+
+ @property
+ def operator_name(self) -> str:
+ # Shown as the wrapped operator (its class name, or a @task callable's
custom_operator_name).
+ # task_type is not forwarded: it names the class that runs, and what
resolves a class from it
+ # (the task's class reference, OpenLineage's extractors) would
otherwise get the wrapped
+ # operator's, whose attributes this one does not have.
+ return self._operator.operator_name
+
+ @property
+ def task_retries(self) -> int:
+ return self._operator.retries or 0
+
+ def _do_render_template_fields(
+ self,
+ parent: Any,
+ template_fields: Iterable[str],
+ context: Context,
+ jinja_env: jinja2.Environment,
+ seen_oids: set[int],
+ ) -> None:
+ # IterableOperator doesn't need to render template fields as the
actual operator's template fields
+ # will be rendered in the IndexedTaskRunner when running each mapped
task instance.
+ pass
+
+ def _get_specified_expand_input(self) -> ExpandInput:
+ return self.expand_input
+
+ def _render_unmapped_operator(
+ self, context: Context, unmapped_task: BaseOperator, jinja_env:
jinja2.Environment
+ ) -> None:
+ context_update_for_unmapped(context, unmapped_task)
+
+ unmapped_task._do_render_template_fields(
+ parent=unmapped_task,
+ template_fields=self._operator.template_fields,
+ context=context,
+ jinja_env=jinja_env,
+ seen_oids=set(),
+ )
+
+ async def axcom_push(self, task: IndexedTaskInstance, value: Any) -> None:
+ await task.axcom_push(key=BaseXCom.XCOM_RETURN_KEY, value=value)
+
+ def _run_tasks(
+ self,
+ context: Context,
+ indices: AsyncIterable[int],
+ jinja_env: jinja2.Environment,
+ ) -> tuple[bool, list[int]]:
+ """Run every index to completion; return whether its tasks pushed
results, and the skipped indices."""
+ self.log.info("Running tasks with %d workers", self.max_workers)
+ with (
+ Checkpoints(context) as checkpoints,
+ IndexedTaskOutcomes(self, self._state, context) as outcomes,
+ ):
+ try:
+ with (
+ event_loop() as loop,
+ AsyncAwareExecutor(
+ loop=loop, max_workers=self.max_workers,
shutdown_timeout=self._SHUTDOWN_TIMEOUT
+ ) as executor,
+ ):
+ try:
+ for task, _result, raised in executor.imap_unordered(
+ partial(
+ self._run_index,
+ executor,
+ context,
+ jinja_env=jinja_env,
+
trust_checkpoints=checkpoints.trust_checkpoints,
+ since=checkpoints.since,
+ outcomes=outcomes,
+ ),
+ indices,
+ stop=self._state.stop_requested,
+ ):
+ outcomes.record(task, raised)
+ except BaseException:
+ # Whatever ends the loop early (the parent's
execution_timeout, a failure that
+ # stops the task) is followed by the executor
cancelling the coroutines, which
+ # would leave nothing registered for on_kill(); kill
what is in flight first,
+ # off the loop thread, so that an indexed operator's
synchronous SDK call in on_kill
+ # waits its turn while the loop goes on serving the
indexed tasks' calls in flight.
+ # Awaited here, unlike on_kill()'s own thread: the
loop keeps running meanwhile.
+ self._state.request_stop()
+ in_flight = self._state.take_in_flight()
+ # On a thread of its own, as on_kill()'s is, and
waited for as long as the
+ # SIGTERM path waits: an on_kill that never returns (a
cancel call with no
+ # network timeout) holds the run for
_SHUTDOWN_TIMEOUT, not forever.
+ self._state.start_kill(partial(self._kill, in_flight))
+ try:
+ loop.run_until_complete(
+ asyncio.to_thread(self._state.await_kill,
self._SHUTDOWN_TIMEOUT)
+ )
+ except RuntimeError:
+ self._state.await_kill(self._SHUTDOWN_TIMEOUT)
+ raise
+ finally:
+ # The kill on_kill() started runs on its own thread and may
still be cleaning an
+ # indexed operator up once the killed indexed tasks came back
and the loop closed. Waited
+ # for here, before the failure callbacks fire and the run
concludes, so the task does
+ # not end with that cleanup half done. The loop is closed, so
a synchronous SDK call
+ # in that on_kill no longer has an asend to wait for.
+ self._state.await_kill(self._SHUTDOWN_TIMEOUT)
+ skipped = outcomes.conclude()
+ return self._state.do_xcom_push, skipped
+
+ def _skip_downstream_of_a_partial_skip(self, context: Context, result:
XComIterable | None) -> None:
+ """
+ Skip the downstream tasks a skipped mapped task instance would have
skipped.
+
+ Some iterations skipped and the others succeeded, so the task
succeeds. Downstream tasks whose
+ trigger rule is satisfied by that are left to run over the values that
exist; those whose
+ rule skips them when any upstream task instance skipped are skipped
here, as they would be
+ after a mapped upstream. ``DownstreamTasksSkipped`` ends ``execute``
before the runner
+ pushes the return value, so it is pushed here first, for the
downstream tasks that do run.
+ """
+ to_skip = [
+ task.task_id
+ for task in self.downstream_list
+ if task.trigger_rule in SKIPPED_WITH_A_SKIPPED_UPSTREAM
+ ]
+ if not to_skip:
+ return
+ ti = context["ti"]
+ if TYPE_CHECKING:
+ assert isinstance(ti, RuntimeTaskInstance)
+ _push_xcom_if_needed(result, ti, self.log)
+ raise DownstreamTasksSkipped(tasks=to_skip)
+
+ async def _run_index(
+ self,
+ executor: AsyncAwareExecutor,
+ context: Context,
+ index: int,
+ *,
+ jinja_env: jinja2.Environment,
+ trust_checkpoints: bool,
+ since: int = 0,
+ outcomes: IndexedTaskOutcomes,
+ ) -> tuple[IndexedTaskInstance | UnreadIndex, Any | None, BaseException |
None]:
+ """
+ Read the input of ``index`` and run its indexed task.
+
+ Reading the input is part of the index's run, as resolving its input
is part of a mapped
+ task instance's run under ``.expand()``: a ``.map()`` callable that
raises, or a template
+ that fails to render, ends this index and no other (see
:class:`UnreadIndex`).
+ Only an ``Exception`` is taken as the index's outcome; a cancellation
or the parent's
+ timeout stops the run as a whole.
+ """
+ if self._state.stop_requested():
+ return (
+ UnreadIndex(self.task_id, index),
+ None,
+ IndexedTaskInstanceNotStarted(
+ f"Indexed task {self.task_id}[{index}] was pulled before
the kill and never started"
+ ),
+ )
+ resolved = self._state.resolved
+ if TYPE_CHECKING:
+ assert resolved is not None
+ try:
+ mapped_kwargs = await resolved.aget(index)
+ # Rendering may call the supervisor synchronously (an XComArg in a
partial kwarg,
+ # ``{{ var.value.x }}``, ``{{ conn.x }}``): on the loop thread
that call would hold up
+ # every indexed task for its round trip, from a worker thread it
waits its turn while the
+ # loop runs, as in XComArg.aresolve.
+ task = await asyncio.to_thread(
+ self._create_indexed_task,
+ context=context,
+ index=index,
+ mapped_kwargs=mapped_kwargs,
+ jinja_env=jinja_env,
+ )
+ except Exception as raised:
+ return UnreadIndex(self.task_id, index), None, raised
+ return await self._run_indexed_task(
+ executor, context, task, trust_checkpoints, since,
outcomes=outcomes
+ )
+
+ async def _run_indexed_task(
+ self,
+ executor: AsyncAwareExecutor,
+ context: Context,
+ task: IndexedTaskInstance,
+ trust_checkpoints: bool,
+ since: int = 0,
+ *,
+ outcomes: IndexedTaskOutcomes,
+ ) -> tuple[IndexedTaskInstance, Any | None, BaseException | None]:
+ # See _run_tasks: checkpoints are consulted only on a retry that is
not a rerun after a clear.
+ if (
+ trust_checkpoints
+ and (replayed := await self._replay_checkpoint(task, context,
since)) is not None
+ ):
+ return replayed
+
+ # Pulled before the kill and reached after it. on_kill() reaches the
operators that started
+ # and the stop flag keeps the executor from pulling more; an indexed
task is neither until its runner
+ # starts, so a kill that lands while its checkpoint is read above
would otherwise start it
+ # after the kill. A sync indexed task is still handed to the pool
below, whose threads are as many
+ # as the calls in flight, so only that pickup remains.
+ if self._state.stop_requested():
+ return (
+ task,
+ None,
+ IndexedTaskInstanceNotStarted(
+ f"Indexed task {task.task_id}[{task.index}] was pulled
before the kill and never started"
+ ),
+ )
+
+ # The indexed task runs against its own view of the context (see
IndexedTaskRunner.indexed_context),
+ # with its own outlet events: indexed tasks run concurrently and each
needs its events
+ # attributed correctly so they can be checkpointed and merged
individually (see
+ # IndexedTaskState.record_outlet_events and
IndexedTaskRunner.merge_outlet_events_into).
+ indexed_task_runner = IndexedTaskRunner(
+ task_instance=task,
+ register=self._state,
+ )
+ try:
+ if task.is_async:
+ with indexed_task_runner:
+ result = await indexed_task_runner.arun(context)
+ else:
+ # Entered and exited in the worker thread with execute, so
what the exit does runs
+ # there too, as the reports below do (see
_report_indexed_task).
+
+ def run_indexed_task():
+ with indexed_task_runner:
+ return indexed_task_runner.run(context)
+
+ result = await executor.run_sync(run_indexed_task)
+
+ indexed_task_state = IndexedTaskState(
+ status=TaskInstanceState.SUCCESS,
+ fingerprint=task.input_fingerprint,
+ try_number=task.try_number,
+ )
+ # The result is checkpointed as well as pushed to XCom: the runner
deletes every XCom
+ # key listed by the server before each attempt
(xcom_keys_to_clear), so XCom alone
+ # cannot survive a retry, while the state store does. Both writes
happen here, inside
+ # the coroutine, so they overlap with the indexed tasks still
running instead of adding a
+ # synchronous pass over every index once the executor has drained.
With a
+ # state_store_backend configured the checkpoint holds only a
reference to the payload.
+ if result is not None and task.do_xcom_push:
+ indexed_task_state.result = result
+
indexed_task_state.record_outlet_events(indexed_task_runner.outlet_events)
+ # Written with the one checkpoint, not per push: a retry that
skips this indexed task pushes
+ # them again, as the runner has deleted them by then.
+ if task.pushed_xcoms:
+ indexed_task_state.xcoms = dict(task.pushed_xcoms)
+ await task.aset_state(indexed_task_state)
+ except (asyncio.CancelledError, AirflowTaskTimeout) as stopped:
+ if isinstance(stopped, AirflowTaskTimeout) and
indexed_task_runner.timed_out_on_its_own:
+ # Raised by the operator in its worker thread, which the
parent's signal never
+ # reaches: this indexed task's own failure, as it is the
mapped task instance's under
+ # .expand(); the siblings go on.
+ return await self._record_failure(task, indexed_task_runner,
outcomes, stopped)
+ # Not this indexed task's outcome: it is being stopped from
outside, by the executor
+ # cancelling it or by the parent's execution_timeout, whose signal
handler raises on the
+ # main thread in whichever indexed task happens to run there. Both
go on unchanged, so a
+ # cancellation stays one and the timeout reaches the runner, which
retries the task. The
+ # indexed task the timeout struck is reported with the others (a
cancelled one has nothing).
+ if isinstance(stopped, asyncio.CancelledError):
+ # A sync indexed task's thread may still be running: it gets
no checkpoint from here on,
+ # so its exit must report nothing either (see
IndexedTaskRunner.cancel).
+ indexed_task_runner.cancel()
+ if indexed_task_runner.failure is not None:
+ outcomes.note_failed(indexed_task_runner)
+ raise
+ except AirflowSkipException as skipped:
+ return await self._record_skip(task, executor,
indexed_task_runner, skipped)
+ except BaseException as e:
+ return await self._record_failure(task, indexed_task_runner,
outcomes, e)
+
+ return await self._publish(task, executor, context,
indexed_task_runner, indexed_task_state, result)
+
+ @keeps_failure_to_its_indexed_task
+ async def _replay_checkpoint(
+ self, task: IndexedTaskInstance, context: Context, since: int
+ ) -> tuple[IndexedTaskInstance, None, BaseException | None] | None:
+ """
+ Settle an indexed task from its checkpoint on a retry, or return None
for it to run.
+
+ A checkpoint from before a clear (an older ``try_number`` than
``since``) or for another input
+ is not trusted. A SUCCESS checkpoint is replayed: its result, other
XComs and outlet events
+ are pushed again, as the runner deleted the XComs before this attempt.
A SKIPPED one is
+ skipped again. A read or a replay that fails is the indexed task's
outcome and leaves the
+ checkpoint as it is: the indexed task cannot tell whether it already
succeeded, so it does
+ not run, and the next retry reads or replays the checkpoint again. A
checkpoint serde
+ refuses to restore is the exception: it would be refused on every
retry, so the indexed
+ task runs again as if there were none, with a warning.
+ """
+ try:
+ indexed_task_state = await task.aget_state()
+ except CheckpointDeserializationError as e:
+ # Refused on this attempt, it is refused on every one: kept as the
outcome, the indexed
+ # task could never finish again, so it runs as if there were no
checkpoint.
+ self.log.warning(
+ "Running task instance %s for %s again: its checkpoint cannot
be restored: %s",
+ task.index,
+ task.task_id,
+ e,
+ )
+ return None
+ if indexed_task_state is None:
+ return None
+ if indexed_task_state.try_number < since or
indexed_task_state.fingerprint != task.input_fingerprint:
+ self.log.info(
+ "Running task instance %s for %s again: its checkpoint is from
before the task was "
+ "cleared or for another input",
+ task.index,
+ task.task_id,
+ )
+ return None
+ if indexed_task_state.status == TaskInstanceState.SUCCESS:
+ self.log.info(
+ "Skipping task instance %s for %s which already succeeded on a
previous attempt",
+ task.index,
+ task.task_id,
+ )
+ if indexed_task_state.result is not None:
+ await self.axcom_push(task, indexed_task_state.result)
+ for key, value in (indexed_task_state.xcoms or {}).items():
+ await task.axcom_push(key=key, value=value)
+ if indexed_task_state.outlet_events:
+
indexed_task_state.replay_outlet_events(context["outlet_events"])
+ return task, None, None
+ if indexed_task_state.status == TaskInstanceState.SKIPPED:
+ return (
+ task,
+ None,
+ AirflowSkipException(
+ f"Indexed task {task.task_id}[{task.index}] was skipped on
a previous attempt"
+ ),
+ )
+ return None
+
+ @keeps_failure_to_its_indexed_task
+ async def _publish(
+ self,
+ task: IndexedTaskInstance,
+ executor: AsyncAwareExecutor,
+ context: Context,
+ runner: IndexedTaskRunner,
+ indexed_task_state: IndexedTaskState,
+ result: Any | None,
+ ) -> tuple[IndexedTaskInstance, Any | None, None]:
+ """
+ Report a succeeded indexed task and publish its result, once its
SUCCESS checkpoint is written.
+
+ Reported then, not when execute returned: the callback speaks for work
a retry will not run
+ again, and a checkpoint write that fails fires nothing, as a plain
task whose result could
+ not be pushed fires no success callback. A report or a push that
fails, a callback raising a
+ BaseException included, leaves the checkpoint as it is, so the retry
replays the result
+ (see :meth:`_replay_checkpoint`) instead of running the operator
again, and reports nothing
+ again.
+ """
+ await self._report_indexed_task(executor, task, runner.report_success,
task.task.on_success_callback)
+ # The result is only checkpointed when the indexed task pushes XComs
and returned something,
+ # so the same condition decides whether there is a
return_value_<index> to push at all.
+ if indexed_task_state.result is not None:
+ await self.axcom_push(task, indexed_task_state.result)
+ runner.merge_outlet_events_into(context["outlet_events"])
+ return task, result, None
+
+ @keeps_failure_to_its_indexed_task
+ async def _record_skip(
+ self,
+ task: IndexedTaskInstance,
+ executor: AsyncAwareExecutor,
+ runner: IndexedTaskRunner,
+ skipped: AirflowSkipException,
+ ) -> tuple[IndexedTaskInstance, None, AirflowSkipException]:
+ """
+ Checkpoint a skipped indexed task as SKIPPED, report it and return the
skip as the outcome.
+
+ A checkpoint write that fails, or a BaseException from
``on_skipped_callback``, is the indexed
+ task's outcome, as for a success. A failed write leaves no checkpoint
and the retry runs the
+ indexed task again; a failed report leaves the SKIPPED checkpoint,
which the retry replays.
+ """
+ await task.aset_state(
+ IndexedTaskState(
+ status=TaskInstanceState.SKIPPED,
+ fingerprint=task.input_fingerprint,
+ try_number=task.try_number,
+ )
+ )
+ await self._report_indexed_task(executor, task, runner.report_skip,
task.task.on_skipped_callback)
+ return task, None, skipped
+
+ async def _record_failure(
+ self,
+ task: IndexedTaskInstance,
+ runner: IndexedTaskRunner,
+ outcomes: IndexedTaskOutcomes,
+ raised: BaseException,
+ ) -> tuple[IndexedTaskInstance, None, BaseException]:
+ """Note a failed indexed task for its callbacks, checkpoint it as
UP_FOR_RETRY and return it as the outcome."""
+ if runner.failure is not None:
+ outcomes.note_failed(runner)
+ # Written with the input and the attempt, like the other outcomes, so
the next attempt
+ # tells a plain retry apart from a clear or a changed input and logs
only the latter.
+ # A write that fails leaves the outcome as it is: the indexed task
failed either way, and
+ # without the checkpoint the retry runs it again, as it would after
UP_FOR_RETRY.
+ try:
+ await task.aset_state(
+ IndexedTaskState(
+ status=TaskInstanceState.UP_FOR_RETRY,
+ fingerprint=task.input_fingerprint,
+ try_number=task.try_number,
+ )
+ )
+ except (asyncio.CancelledError, AirflowTaskTimeout):
+ raise
+ except Exception:
+ self.log.exception(
+ "Could not checkpoint task instance %s for %s as up for
retry", task.index, task.task_id
+ )
+ return task, None, raised
+
+ @staticmethod
+ async def _report_indexed_task(
+ executor: AsyncAwareExecutor,
+ task: IndexedTaskInstance,
+ report: Callable[[], None],
+ callbacks: Sequence[Callable],
+ ) -> None:
+ """
+ Run an indexed task's success or skip report where its code ran.
+
+ On the loop for an async operator, in a worker thread for a sync one,
where a synchronous
+ SDK call made from the callback waits its turn for the channel instead
of holding up the loop.
+ Without ``callbacks`` the report only sets the state and ``end_date``,
which makes no SDK
+ call, so it runs on the loop and spares the indexed task a thread
hand-off.
+ """
+ if task.is_async or not callbacks:
+ report()
+ else:
+ await executor.run_sync(report)
+
+ @staticmethod
+ def _compute_fingerprint(mapped_kwargs: Mapping[str, Any]) -> str | None:
+ """
+ Digest one indexed task's input, stored on its checkpoint to tell
whether it still applies.
+
+ A retry may run on another input than the attempt that wrote the
checkpoints: the upstream was
+ cleared together with this task and produced other values. An index
then no longer means the
+ same work, and replaying its result would hand downstream a value
computed from the old value.
+ An input serde cannot serialize has no digest, and its checkpoint is
honoured by index alone.
+ """
+ try:
+ serialized = json.dumps(serialize(mapped_kwargs), sort_keys=True)
+ except (TypeError, ValueError, AttributeError, RecursionError):
+ return None
+ return hashlib.sha256(serialized.encode()).hexdigest()
+
+ def _collect_partial_inputs_from_upstream(self, unmapped_task:
BaseOperator) -> dict[str, Any]:
+ """
+ Collect the rendered values of the partial kwargs an upstream task
provides.
+
+ They belong in the fingerprint next to the iterated kwargs: clearing
the upstream together with
+ this task can change them while the input stays the same, and a
checkpoint written with the old
+ value must not be replayed. Only XComArg values count, read back from
the unmapped operator
+ once rendered, wherever they sit in a partial kwarg: at the top level,
or at any depth of the
+ mappings, lists and tuples it holds, such as a ``@task``'s
``op_kwargs`` and ``op_args``. Other
+ templated values are left out on purpose: one like ``{{ ti.try_number
}}`` changes with every
+ attempt and would make every checkpoint look stale.
+ """
+ inputs: dict[str, Any] = {}
+ for key, value in self._partial_kwargs_from_upstream.items():
+ self._collect_upstream_values(key, value, getattr(unmapped_task,
key, None), inputs)
+ return inputs
+
+ @cached_property
+ def _partial_kwargs_from_upstream(self) -> dict[str, Any]:
+ """The partial kwargs holding an XComArg at any depth: the same for
every index, so found once."""
+ return {
+ key: value
+ for key, value in self.partial_kwargs.items()
+ if any(XComArg.iter_xcom_references(value))
+ }
+
+ @classmethod
+ def _collect_upstream_values(
+ cls, path: str, declared: Any, rendered: Any, inputs: dict[str, Any]
+ ) -> None:
+ """
+ Record under its path the rendered value of every XComArg in
``declared``, at any depth.
+
+ ``declared`` is a partial kwarg as written in the DAG and ``rendered``
the same value once
+ rendered, walked side by side: a mapping by key (``op_kwargs.cfg.x``),
a list or tuple by
+ position (``op_args[0]``). A part rendered into another shape than it
was declared in is
+ recorded as None.
+ """
+ if isinstance(declared, XComArg):
+ inputs[path] = rendered
Review Comment:
When the upstream behind a partial kwarg is itself an iterated task,
`rendered` here is an `XComIterable`, and its `serialize()` keeps only
`task_id`, `run_id`, `length` and `skipped`, none of the values. So after a
failed attempt, clearing that upstream together with this task and getting
different values of the same length leaves the fingerprint unchanged, and the
succeeded indexes replay results computed from the old input. Could the
fingerprint digest the values of a lazy partial input (`XComIterable` or
`LazyXComSequence`) once per run? That would also stop every index from
fetching a whole mapped upstream through the serde serializer just to
fingerprint it.
--
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]
