dabla commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r4252254414
########## task-sdk/src/airflow/sdk/definitions/iterableoperator.py: ########## @@ -0,0 +1,1444 @@ +# +# 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 partial, wraps +from typing import TYPE_CHECKING, Any, TypeVar + +try: + # Python 3.11+ + BaseExceptionGroup +except NameError: + from exceptiongroup import BaseExceptionGroup + +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 ( + 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_item. +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) + + +# 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. +def keeps_failure_to_its_item( + 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 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._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_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 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 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 + # sub-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 a sub-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 sub-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 + + 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 sub-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 sub-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 a sub-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("Sub-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} items 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 item's own ``AirflowTaskTimeout`` next to + another item'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. + """ + sub_task = f"Sub-task {task.task_id}[{task.index}]" + if isinstance(raised, TaskDeferred): + return AirflowFailException( + f"{sub_task} attempted to defer. Deferrable operators are not supported inside IterableOperator." + ) + if isinstance(raised, (DagRunTriggerException, DownstreamTasksSkipped)): + return AirflowFailException( + f"{sub_task} raised {type(raised).__name__}. Triggering DAG runs (TriggerDagRunOperator) and " + "skipping downstream tasks (ShortCircuitOperator and similar) are not supported inside " + "IterableOperator: the sub-task's index has no downstream tasks or DAG run of its own for " + "the effect to apply to." + ) + if isinstance(raised, AirflowRescheduleException): + return AirflowFailException( + f"{sub_task} attempted to reschedule (raised AirflowRescheduleException). Reschedule-mode " + "sensors are not supported inside IterableOperator: the sub-task's index has no task " + "instance of its own to reschedule." + ) + if isinstance(raised, AirflowTaskTimeout): + # The parent's timeout never gets here (_run_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"{sub_task} 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 sub-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 a sub-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 item 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 sub-task needs + a retry the whole IterableOperator is retried by Airflow. Already-succeeded + sub-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 sub-tasks are skipped and only the pending/failed ones are + executed again. Every sub-task inherits its ``try_number`` from the IterableOperator's own task + instance, so the attempt count reported to a sub-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 runtime receives a + deep copy/unmapped instance of this operator. + + :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 sub-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: a sub-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. + + Sub-task outcomes are classified before being aggregated: if any sub-task raises + :class:`~airflow.sdk.exceptions.AirflowFailException`, that exception is re-raised directly so + the IterableOperator fails without retrying. A sub-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. A sub-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* sub-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 sub-task exceptions 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:: + **Async sub-tasks should prefer async SDK calls.** + + IterableOperator runs multiple async sub-tasks concurrently on the same event loop, each + making async SDK calls of its own (checkpointing, XCom push). If an async sub-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 sub-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 sub-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 sub-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 sub-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 sub-task raises itself (a hook that gave up waiting) is that + sub-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, + "map_index_template": operator.map_index_template, + "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 sub-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 sub-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 sub-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 sub-task unaware that the IterableOperator itself was killed + # (SIGTERM) or hit its execution_timeout: propagate to each active sub-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 sub-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 a sub-operator's on_kill would also hold up the event loop, and + # every sub-task on it, for its round trip; in its own thread the call waits its turn and + # the loop goes on serving the sub-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 sub-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 sub-task operator %s", operator.task_id) + + def deserialization_allowed_values(self) -> list[Any]: + """Return the values the operator this task iterates declares: every item 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.""" + do_xcom_push = True + + 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_item, + executor, + context, + jinja_env=jinja_env, + trust_checkpoints=checkpoints.trust_checkpoints, + since=checkpoints.since, + outcomes=outcomes, + ), + indices, + stop=self._state.stop_requested, + ): + if isinstance(task, IndexedTaskInstance): + do_xcom_push = task.do_xcom_push + 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 a sub-operator's synchronous SDK call in on_kill + # waits for the sub-tasks' calls in flight instead of raising. Awaited here, + # unlike on_kill()'s own thread: the loop keeps running meanwhile. + self._state.request_stop() + in_flight = self._state.take_in_flight() + try: + loop.run_until_complete(asyncio.to_thread(self._kill, in_flight)) + except RuntimeError: + self._kill(in_flight) + raise + finally: + # The kill on_kill() started runs on its own thread and may still be cleaning a + # sub-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 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_item( + 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"Sub-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 sub-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_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_task(executor, context, task, trust_checkpoints, since, outcomes=outcomes) + + async def _run_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 item 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 item 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"Sub-task {task.task_id}[{task.index}] was pulled before the kill and never started" + ), + ) + + # The sub-task runs against its own view of the context (see IndexedTaskRunner.indexed_context), + # with its own outlet events: sub-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_item). + + 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 sub-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 sub-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 sub-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 sub-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 sub-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 + # sub-task the timeout struck is reported with the others (a cancelled one has nothing). + if isinstance(stopped, asyncio.CancelledError): + # A sync sub-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_item + 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. + """ + indexed_task_state = await task.aget_state() + 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"Sub-task {task.task_id}[{task.index}] was skipped on a previous attempt" + ), + ) + return None + + @keeps_failure_to_its_item + 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_item(executor, task, runner.report_success) + # The result is only checkpointed when the sub-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_item + 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_item(executor, task, runner.report_skip) + 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_item( + executor: AsyncAwareExecutor, task: IndexedTaskInstance, report: Callable[[], None] + ) -> 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 for the comms lock instead of raising. Review Comment: Right, fixed in 843d6b3a3a: the `_report_item` docstring, the kill comment in `_run_tasks` and the `in_flight` docstring now give the reason that holds since the call waits, that on the loop thread it would hold up every item for its round trip. The leftover copy of the `_DECISION_WEIGHT` comment is gone. I also went through the whole change for other passages saying a sync call on the loop raises; the remaining `DeadlockImminentError` wording is about the re-entry case only, which still raises. Drafted-by: Claude Opus 5.5; reviewed by @dabla before posting ########## task-sdk/src/airflow/sdk/execution_time/task_runner.py: ########## @@ -945,6 +972,583 @@ def mark_success_url(self) -> str: return self.log_url +@dataclass +class IndexedTaskState: + status: TaskInstanceState + result: Any | None = None + # Outlet asset events the sub-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 sub-tasks emitted. A sub-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_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 sub-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 sub-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 sub-task pushed, by their key before the index suffix. The runner deletes + # every XCom before an attempt, so a sub-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 a sub-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: + if not isinstance(raw, Mapping): + return None + return cls( + status=TaskInstanceState(raw["status"]), + result=serde_deserialize(raw.get("result")), + outlet_events=raw.get("outlet_events"), + fingerprint=raw.get("fingerprint"), + try_number=raw.get("try_number", 0), + xcoms={key: serde_deserialize(value) for key, value in raw["xcoms"].items()} + if raw.get("xcoms") + else None, + ) + + +class IndexedTaskInstance(RuntimeTaskInstance): + """ + Indexed task instance to run a mapped operator. + + It shares the parent task instance's identity, so what an iteration pushes or stores lands in + the parent's scope, suffixed with the index so that iterations never overwrite each other: + XComs through :meth:`xcom_push`, task state through :attr:`task_state_store`. Reading back + follows the same rule: :meth:`xcom_pull` of the iteration's own XComs adds the index, a pull + from another task does not, as the store's accessor adds it to every key of its own. The + operator's own checkpoints go to the parent's store unsuffixed, under their + ``_iterable_<index>`` keys. + """ + + index: int + parent_task_state_store: TaskStateStoreAccessor + input_fingerprint: str | None = None + # The XComs the sub-task pushed other than its return value, by their key before the index + # suffix: kept in memory while it runs and written once, with its SUCCESS checkpoint. + pushed_xcoms: dict[str, Any] = Field(default_factory=dict) + + @classmethod + def create_indexed_task( + cls, *, context: Context, index: int, operator: BaseOperator, input_fingerprint: str | None = None + ) -> IndexedTaskInstance: + """ + Create the runtime instance for one index of an iterated task from the parent's context. + + The instance shares the parent task instance's identity (id, run, map index, try number and + retry budget), so XComs and task state land in the parent's scope, and carries the unmapped + operator for that index. The budget is the parent's ``max_tries`` rather than the operator's + ``retries``: a manual clear raises it, and it is the parent Airflow retries. The parent's + state store comes from the context, the same accessor the operator's + checkpoints use. ``model_construct`` skips Pydantic validation on purpose: one instance is built per + indexed task, and the parent was validated already, so only the index needs checking here. + """ + if index < 0: + raise ValueError(f"IndexedTaskInstance requires index >= 0, got {index}") + parent = context["ti"] + return cls.model_construct( + id=parent.id, + parent_task_state_store=context["task_state_store"], + # The parent's: an iteration's XComs and state belong to the task instance that runs it. + task_id=parent.task_id, + dag_id=operator.dag_id, + run_id=parent.run_id, + map_index=parent.map_index, + region_id=getattr(parent, "region_id", None), + region_index=getattr(parent, "region_index", None), + index=index, + input_fingerprint=input_fingerprint, + max_tries=parent.max_tries, + start_date=parent.start_date, + state=TaskInstanceState.SCHEDULED.value, + is_mapped=True, + task=operator, + try_number=parent.try_number, + # The dag run, logical date and the rest the server sent for the task instance that + # runs this iteration: its template context, get_previous_ti() and the like read them. + _ti_context_from_server=getattr(parent, "_ti_context_from_server", None), + ) + + def xcom_push( + self, + key: str, + value: Any, + ): + super().xcom_push(key=f"{key}_{self.index}", value=value) + self._record_push(key, value) + + async def axcom_push( + self, + key: str, + value: Any, + ): + await super().axcom_push(key=f"{key}_{self.index}", value=value) + self._record_push(key, value) + + def _record_push(self, key: str, value: Any) -> None: + # The return value has its own slot on the checkpoint and is published by the operator. + if key != BaseXCom.XCOM_RETURN_KEY: + self.pushed_xcoms[key] = value + + def _own_key(self, task_ids: str | Iterable[str] | None, dag_id: str | None, key: str) -> str: + """ + Suffix ``key`` with the index for a pull of this iteration's own XComs. + + What :meth:`xcom_push` wrote under ``<key>_<index>`` is read back under the same name when + the pull names no task or this task instance's own (in its own DAG); a pull from another + task, or from several, keeps its key, since those XComs carry no index. + """ + own_task = task_ids is None or task_ids == self.task_id + own_dag = dag_id is None or dag_id == self.dag_id + return f"{key}_{self.index}" if own_task and own_dag else key + + def xcom_pull( + self, + task_ids: str | Iterable[str] | None = None, + dag_id: str | None = None, + key: str = BaseXCom.XCOM_RETURN_KEY, + include_prior_dates: bool = False, + *, + map_indexes: int | Iterable[int] | None | ArgNotSet = NOTSET, + default: Any = None, + run_id: str | None = None, + ) -> Any: + return super().xcom_pull( + task_ids=task_ids, + dag_id=dag_id, + key=self._own_key(task_ids, dag_id, key), + include_prior_dates=include_prior_dates, + map_indexes=map_indexes, + default=default, + run_id=run_id, + ) + + async def axcom_pull( + self, + task_ids: str | Iterable[str] | None = None, + dag_id: str | None = None, + key: str = BaseXCom.XCOM_RETURN_KEY, + include_prior_dates: bool = False, + *, + map_indexes: int | Iterable[int] | None | ArgNotSet = NOTSET, + default: Any = None, + run_id: str | None = None, + ) -> Any: + return await super().axcom_pull( + task_ids=task_ids, + dag_id=dag_id, + key=self._own_key(task_ids, dag_id, key), + include_prior_dates=include_prior_dates, + map_indexes=map_indexes, + default=default, + run_id=run_id, + ) + + @cached_property + def task_state_store(self) -> TaskStateStoreAccessor: # type: ignore[override] + """The parent's store seen from this iteration: keys are suffixed with the index.""" + return IndexedTaskStateStoreAccessor(self.parent_task_state_store, self.index) + + def context_for(self, context: Context, *, outlet_events: OutletEventAccessors | None = None) -> Context: + """ + Return the parent's context as this indexed task sees it. + + A clone of ``context`` with this task instance under ``ti`` and ``task_instance``, its + unmapped operator under ``task`` and its indexed view of the store under + ``task_state_store``: the keys ``context_update_for_unmapped`` sets for a mapped task + instance, and the store next to them, so that a template and ``execute`` read the same + store. The one place these keys are listed: :meth:`IndexedTaskRunner.indexed_context` + builds the context the task runs in from it, and ``IterableOperator._create_task`` the + one its templates are rendered against. ``outlet_events`` is swapped when given, for the + run; before it, nothing is emitted and the parent's accessor stays. + """ + indexed: Context = { + **clone_context(context), + "ti": self, + "task_instance": self, + "task": self.task, + "task_state_store": self.task_state_store, + } + if outlet_events is not None: + indexed["outlet_events"] = outlet_events + return indexed + + async def aget_state(self) -> IndexedTaskState | None: + return IndexedTaskState.deserialize(await self.parent_task_state_store.aget(self.state_key)) + + async def aset_state(self, state: IndexedTaskState) -> None: + await self.parent_task_state_store.aset(self.state_key, state.serialize()) + + @property + def is_async(self) -> bool: + return self.task.is_async + + @property + def is_eligible_to_retry(self) -> bool: + """ + Whether Airflow runs the parent again after this attempt fails. + + The same rule the API server applies to the parent (``_is_eligible_to_retry``), so the + callbacks of an iteration agree with what happens to the task instance. + """ + return self.max_tries != 0 and self.try_number <= self.max_tries + + @property + def state_key(self) -> str: + return IndexedTaskState.build_key(self.index) + + @property + def do_xcom_push(self) -> bool: + return self.task.do_xcom_push + + +class SubOperatorRegister(Protocol): + """Where an indexed task's operator is noted while its code runs, so that a kill reaches it.""" + + def register(self, operator: BaseOperator) -> None: ... + + def unregister(self, operator: BaseOperator) -> None: ... + + +class IndexedTaskRunner(LoggingMixin): + """ + Run one indexed task of an iterated task: its operator, against its own view of the context. + + Named apart from Airflow's executors, which schedule task instances, and from the task runner + process, which runs the parent: this runs one index inside that process, sync or async. + """ + + def __init__( + self, + task_instance: IndexedTaskInstance, + register: SubOperatorRegister | None = None, + outlet_events: OutletEventAccessors | None = None, + ): + """ + Run an operator or trigger for one sub-task instance. + + :param outlet_events: The accessor the sub-task's asset events are collected in, its own + so they can be checkpointed and merged apart from its siblings'. Created here when not + given; the caller reads it back through :attr:`outlet_events` after the run. + :param register: Optional register the operator is entered in while its code runs (see + :meth:`in_flight`), so that IterableOperator.on_kill() can reach whichever sub-tasks + are in flight; the iterated task's ``IterationState``. + """ + super().__init__() + self.task_instance = task_instance + self.outlet_events = outlet_events if outlet_events is not None else OutletEventAccessors() + self._result: Any | None = None + self._start_time: float | None = None + self._context: Context | None = None + self._register = register + #: The exception this indexed task failed with, noted by __exit__ and reported by + #: :meth:`report_failure` once the whole task's fate is known. + self.failure: BaseException | None = None + #: Whether the ``AirflowTaskTimeout`` this indexed task ended with was raised by its operator + #: off the main thread, where the parent's limit never strikes; see :meth:`in_flight`. + self.timed_out_on_its_own = False + self._cancelled = False + + def merge_outlet_events_into(self, target: OutletEventAccessorsProtocol) -> None: + """ + Fold the outlet asset events this indexed task recorded into the parent's ``target``. + + Called right after the indexed task succeeds, on the IterableOperator's shared + ``context["outlet_events"]``. A task instance sends one event per asset, so indexed tasks + that emit to the same asset end up in that one event: ``extra`` keeps what the last one to + finish wrote, while partition keys and alias events accumulate. ``.expand()`` sends one + event per mapped task instance instead; the docs page says so in its comparison table. + """ + for asset_or_alias, accessor in self.outlet_events.items(): + target_accessor = target[asset_or_alias] + target_accessor.extra.update(accessor.extra) + target_accessor.asset_alias_events.extend(accessor.asset_alias_events) + target_accessor.partition_keys.update(accessor.partition_keys) + + def cancel(self) -> None: + """ + Note that the coroutine waiting for this sync indexed task was cancelled. + + Its thread may go on, but the task gets no checkpoint from here on, and no report either: + the next attempt runs it again and reports it then. Best effort: a thread already past the + check in its exit still notes a failure. + """ + self._cancelled = True + + @property + def cancelled(self) -> bool: + """Whether :meth:`cancel` was called: the indexed task's exit then notes nothing.""" + return self._cancelled + + @property + def dag_id(self) -> str: + return self.task_instance.dag_id + + @property + def task_id(self) -> str: + return self.task_instance.task_id + + @property + def task_index(self) -> int: + return self.task_instance.index + + @property + def operator(self) -> BaseOperator: + return self.task_instance.task + + @property + def is_async(self) -> bool: + return self.task_instance.is_async + + @contextmanager + def indexed_context(self, context: Context) -> Generator[Context, None, None]: + """ + Enter the parent's context as this indexed task sees it. + + Yields the task instance's :meth:`IndexedTaskInstance.context_for` view of the parent's + context with this runner's own outlet events, remembered on the runner and made the + current context for the duration of the block, so user code reads the indexed task's + unmapped operator under ``context["task"]``, not the IterableOperator. The parent's + context is left untouched: ``context_update_for_unmapped`` sets ``ti.task`` on whatever + ``ti`` it finds, which must be this task's, not the parent's. + """ + indexed_context = self.task_instance.context_for(context, outlet_events=self.outlet_events) + self._context = indexed_context + with set_indexed_context(indexed_context): + yield indexed_context + + @contextmanager + def in_flight(self) -> Generator[None, None, None]: + """ + Register the operator as running for exactly as long as its code runs. + + Entered by :meth:`run` and :meth:`arun`, so in the worker thread or coroutine that executes + the operator: a sync operator stays registered while its thread is still inside ``execute``, + even after the coroutine waiting for it was cancelled, and ``IterableOperator.on_kill`` can + reach it. The operator the parent's execution timeout strikes stays registered as well: + the timeout lands on the main thread, where the loop runs async operators, and its + ``on_kill`` must not run there, where a synchronous SDK call raises, so + ``IterableOperator._run_tasks`` kills it off the loop thread with the others once the + timeout has unwound. ``TimeoutPosix`` raises on the main thread only, so an + ``AirflowTaskTimeout`` on a worker thread was raised by the operator itself (a hook that + gave up waiting): that is this indexed task's own failure, noted as + :attr:`timed_out_on_its_own`, and the operator is unregistered like any other. An async + operator runs on the main thread, so one that raises ``AirflowTaskTimeout`` itself cannot be + told apart from the parent's timeout and is taken for it, a limitation the docs name. + """ + if self._register is not None: + self._register.register(self.operator) + struck_by_the_parent = False + try: + yield + except AirflowTaskTimeout: + self.timed_out_on_its_own = threading.current_thread() is not threading.main_thread() + struck_by_the_parent = not self.timed_out_on_its_own + raise + finally: + if self._register is not None and not struck_by_the_parent: + self._register.unregister(self.operator) + + def run(self, context: Context): + """Run the operator synchronously against this indexed task's own view of ``context``.""" + with self.in_flight(), self.indexed_context(context) as indexed_context: + return _execute_task(indexed_context, self.task_instance, self.log) + + async def arun(self, context: Context): + """Run the async operator against this indexed task's own view of ``context``.""" + with self.in_flight(), self.indexed_context(context) as indexed_context: + return await _execute_async_task(indexed_context, self.task_instance, self.log) + + def __enter__(self): + self._start_time = time.monotonic() + + if self.log.isEnabledFor(logging.INFO): + self.log.info( + "Running attempt %s of %s for %s with index %s in %s mode.", + self.task_instance.try_number, + self.task_instance.max_tries + 1, + self.task_instance.task_id, + self.task_index, + "async" if self.is_async else "sync", + ) + return self + + def __exit__(self, exc_type, exc_value, traceback): + elapsed = time.monotonic() - self._start_time if self._start_time else 0.0 + + if self._cancelled: + return None + if exc_value: + # Cancelled because the task is stopping, for a reason another iteration raised: this + # iteration neither failed nor will be retried on its own account, so it gets no state + # and no callback. The iteration that stopped the task reports its own outcome. + if isinstance(exc_value, CancelledError): + raise exc_value + # A skip is reported by report_skip() once its checkpoint is written, like a success. + if isinstance(exc_value, AirflowSkipException): + raise exc_value + # A failure is only noted here. Whether it is retried is the whole task's fate, which + # the other iterations decide too (a sibling's AirflowFailException fails it without a + # retry), so IterableOperator reports it through report_failure() once all have run. + self.log.error( + "Task instance %s for %s failed on attempt %s in %.2f seconds due to: %s", + self.task_index, + self.task_instance.task_id, + self.task_instance.try_number, + elapsed, + exc_value, + ) + self.failure = exc_value + raise exc_value + + # The state and the success callback follow once the checkpoint is written: report_success(). + if self.log.isEnabledFor(logging.INFO): + self.log.info( + "Task instance %s for %s finished successfully on attempt %s in %.2f seconds", + self.task_index, + self.task_instance.task_id, + self.task_instance.try_number, + elapsed, + ) + + def report_success(self) -> None: + """ + Report this indexed task as succeeded: its state, ``end_date`` and ``on_success_callback``. + + Called by ``IterableOperator._run_task`` once the SUCCESS checkpoint is written, where the + indexed task's code ran. The callback then speaks for work a retry will not run again: a + checkpoint write that fails fires nothing, and the attempt that runs the indexed task again + reports it then; a result push that fails after the checkpoint is replayed from it, so the + callback fires once. + """ + self.task_instance.end_date = datetime.now(tz=UTC) + self.task_instance.state = TaskInstanceState.SUCCESS + if self._context is not None: + _run_task_state_change_callbacks( + self.task_instance.task, "on_success_callback", self._context, self.log + ) + + def report_skip(self) -> None: + """Report this indexed task as skipped: its state, ``end_date`` and ``on_skipped_callback``, once its SKIPPED checkpoint is written.""" + self.task_instance.end_date = datetime.now(tz=UTC) + self.task_instance.state = TaskInstanceState.SKIPPED + if self._context is not None: + _run_task_state_change_callbacks( + self.task_instance.task, "on_skipped_callback", self._context, self.log + ) + + def report_failure(self, task_will_retry: bool) -> None: + """ + Report this indexed task's failure as what happens to the whole task. + + Called once every iteration has run, so the state and the callback agree with the task: + ``UP_FOR_RETRY`` and ``on_retry_callback`` when it is retried, ``FAILED`` and + ``on_failure_callback`` otherwise. + """ + self.task_instance.end_date = datetime.now(tz=UTC) + if task_will_retry: + self.task_instance.state = TaskInstanceState.UP_FOR_RETRY + else: + self.task_instance.state = TaskInstanceState.FAILED + if self._context is not None: Review Comment: Right, fixed in dec578bccf as you suggested: `report_failure` sets `self._context["exception"] = self.failure` before running the callbacks. A new test runs two failing items through `.partial(...).iterate(...)` and checks that each `on_retry_callback` (retries left) and each `on_failure_callback` (last attempt) gets its own item's error; both fail on the previous code. Drafted-by: Claude Opus 5.5; reviewed by @dabla before posting ########## task-sdk/src/airflow/sdk/execution_time/executor.py: ########## @@ -0,0 +1,262 @@ +# +# 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 inspect +import logging +import os +import time +from asyncio import ( + FIRST_COMPLETED, + AbstractEventLoop, + CancelledError, + Future, + Semaphore, + Task, + TimeoutError as AsyncTimeoutError, + gather, + wait, + wait_for, + wrap_future, +) +from collections.abc import AsyncIterable, Callable, Iterator +from concurrent.futures import Executor, ThreadPoolExecutor +from contextlib import suppress +from typing import Any + +_log = logging.getLogger(__name__) + + +class AsyncAwareExecutor(Executor): + """ + Executes both sync and async functions concurrently. + + Sync functions run in a ThreadPoolExecutor. + Async coroutines run on an asyncio event loop with a semaphore limit. + + :param loop: Event loop used to schedule async tasks and coordinate mixed execution. + :param max_workers: Maximum concurrent workers used by both thread pool and async semaphore. + :param shutdown_timeout: Maximum time to wait, in seconds, for in-flight async tasks and + thread-pool workers to finish during ``shutdown(wait=True)``. Python threads cannot be + forcibly stopped, so a worker stuck in blocking user code (slow HTTP call, blocked C + extension, a deadlocked DB driver, ...) would otherwise hang ``shutdown()`` forever. + """ + + def __init__( + self, loop: AbstractEventLoop, max_workers: int | None = None, shutdown_timeout: float = 10.0 + ): + if max_workers is None: + max_workers = os.cpu_count() or 1 + if max_workers <= 0: + raise ValueError("max_workers must be greater than 0") + + self._loop = loop + self._max_workers = max_workers + self._shutdown_timeout = shutdown_timeout + self._semaphore = Semaphore(max_workers) + self._thread_pool = ThreadPoolExecutor(max_workers=max_workers) + self._async_tasks: set[Task[Any]] = set() + self._shutdown = False + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_type is not None: + # On error path, cancel futures but still wait briefly for them to + # process CancelledError and release what they hold. A task + # cancelled in the middle of a request is CommsDecoder's concern: + # asend runs the request in a thread, which drains the response and + # releases the channel's lock after the await was abandoned. + self.shutdown(wait=True, cancel_futures=True) + else: + self.shutdown(wait=True) + + def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None: + if self._shutdown: + return + + self._shutdown = True + + if cancel_futures: + for task in list(self._async_tasks): + task.cancel() + + if wait and self._async_tasks: + with suppress(TimeoutError, AsyncTimeoutError): + self._loop.run_until_complete( + wait_for( + gather(*self._async_tasks, return_exceptions=True), + timeout=self._shutdown_timeout, + ) + ) + + # ThreadPoolExecutor.shutdown(wait=True) blocks until every worker thread + # finishes its current work item, with no way to bound that wait or forcibly + # stop a thread stuck in blocking user code (slow HTTP call, blocked C + # extension, a deadlocked DB driver, ...). Ask the pool to stop accepting new + # work (and cancel anything not yet started) up front, then bound how long we + # personally wait on the worker threads instead of blocking indefinitely. + self._thread_pool.shutdown(wait=False, cancel_futures=cancel_futures) + + if wait: + threads = list(getattr(self._thread_pool, "_threads", ())) + deadline = time.monotonic() + self._shutdown_timeout + for thread in threads: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + thread.join(timeout=remaining) + + stuck = [thread.name for thread in threads if thread.is_alive()] + if stuck: + _log.error( + "%d worker thread(s) still running %.1fs after shutdown was requested; " + "giving up waiting to avoid blocking indefinitely. This may leak resources. " + "Affected threads: %s", + len(stuck), + self._shutdown_timeout, + stuck, + ) + + def submit(self, func: Callable[..., Any] | Any, *args, **kwargs) -> Future[Any]: # type: ignore[override] + """ + Submit a callable for execution. + + Always returns an asyncio.Future for consistency, whether the callable + is sync (run in thread pool) or async (run on the event loop). + """ + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + + if inspect.iscoroutine(func): + coro = func + elif inspect.iscoroutinefunction(func): + coro = func(*args, **kwargs) + else: + # Wrap thread pool future as asyncio.Future for consistent return type + return wrap_future(self._thread_pool.submit(func, *args, **kwargs), loop=self._loop) + + async def guarded(): + try: + async with self._semaphore: + return await coro + except CancelledError: + # If cancellation occurs while waiting for the semaphore, + # the inner coroutine was never awaited. Close it to prevent + # "coroutine was never awaited" RuntimeWarning. + coro.close() + raise + + task = self._loop.create_task(guarded()) + self._async_tasks.add(task) + task.add_done_callback(self._async_tasks.discard) + return task + + async def run_sync(self, func: Callable[..., Any], *args, **kwargs) -> Any: + """Run a sync callable in this executor's thread pool and await its result.""" + future = self._thread_pool.submit(func, *args, **kwargs) Review Comment: Yes, done in 4a2e548a51: `run_sync` and the sync path of `submit`, which had the same gap, now hand the pool `contextvars.copy_context().run`, a fresh copy per call, as `asyncio.to_thread` does. A new executor test with `max_workers=1` checks that each call sees what its caller set and not what the previous call on that thread set, for both paths; it fails on the previous code. The copy has a small cost, measured with CodSpeed's simulation mode on a benchmark of the executor (instructions): | Benchmark | Change | |---|---| | `run_sync` round trip | +1.8% | | `submit` round trip | +1.7% | | `imap_unordered` over 50 sync items | +1.4% | That is a few thousand instructions per call, against an item's own work, and what `asyncio.to_thread` pays too. Drafted-by: Claude Opus 5.5; 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]
