dabla commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r4243514256
########## task-sdk/src/airflow/sdk/definitions/iterableoperator.py: ########## @@ -0,0 +1,1292 @@ +# +# 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, Callable, Iterable, Mapping, Sequence +from functools import partial +from typing import TYPE_CHECKING, Any + +try: + # Python 3.11+ + BaseExceptionGroup +except NameError: + from exceptiongroup import BaseExceptionGroup + +from airflow.sdk import BaseXCom, TaskInstanceState, TriggerRule +from airflow.sdk.bases.operator import BaseOperator, event_loop +from airflow.sdk.bases.skipmixin import SkipMixin +from airflow.sdk.bases.xcom import XComIterable +from airflow.sdk.definitions.retry_policy import RetryAction, RetryDecision +from airflow.sdk.definitions.xcom_arg import XComArg +from airflow.sdk.exceptions import ( + AirflowFailException, + AirflowRescheduleException, + AirflowSensorTimeout, + AirflowSkipException, + AirflowTaskTerminated, + AirflowTaskTimeout, + DagRunTriggerException, + DownstreamTasksSkipped, + TaskDeferred, +) +from airflow.sdk.execution_time.comms import DeadlockImminentError +from airflow.sdk.execution_time.context import context_update_for_unmapped +from airflow.sdk.execution_time.executor import AsyncAwareExecutor +from airflow.sdk.execution_time.task_runner import ( + IndexedTaskInstance, + IndexedTaskRunner, + IndexedTaskState, + _push_xcom_if_needed, +) +from airflow.sdk.serde import serialize + +if TYPE_CHECKING: + import jinja2 + + from airflow.sdk.definitions._internal.expandinput import ExpandInput, Resolved + from airflow.sdk.definitions.context import Context + from airflow.sdk.definitions.mappedoperator import MappedOperator + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + from airflow.sdk.types import Logger + + +# The trigger rules under which one skipped upstream task instance skips a task, whatever the other +# upstream task instances did (see TriggerRuleDep). A skipped iteration has the same effect on a +# downstream task with one of these rules as a skipped mapped task instance would. +SKIPPED_WITH_A_SKIPPED_UPSTREAM = frozenset( + {TriggerRule.ALL_SUCCESS, TriggerRule.NONE_SKIPPED, TriggerRule.ALL_DONE_MIN_ONE_SUCCESS} +) + + +# Raised by an indexed task, these fail the task without a retry, as the runner does for a task that raises +# them itself (see _run_task_and_map_outcome and _handle_handler_failure). +FAIL_WITHOUT_RETRY = (AirflowFailException, AirflowSensorTimeout, AirflowTaskTerminated) + + +# How strongly a retry policy decision speaks for the task when several indexed tasks failed: one the +# policy says must not be retried fails the task, one it says to retry makes it retry on its terms. +class Checkpoints: + """ + Decide whether one attempt of an IterableOperator may resume from its per-index checkpoints. + + Checkpoints are only consulted from the second attempt onwards. A completion marker left by a + previous fully successful run means this attempt follows a manual clear (which raises + ``max_tries`` but does not reset ``try_number``), so every index must run again: the stale + ``SUCCESS`` checkpoints are ignored and overwritten. On entry the marker is replaced by the + attempt the rerun starts at, so that a crash during the rerun resumes from the checkpoints + written since, and only from those: an index the crashed rerun did not reach still holds its + checkpoint from before the clear, which must not be replayed. The marker is written again when + the block exits without an exception, and when every iteration skipped: the task is then + ``SKIPPED``, a final state, and a clear of it must run every iteration again rather than + replay the skips, as a cleared mapped task instance would. An attempt that fails, even with + some iterations skipped, writes no marker, so a retry or a clear after it resumes: iterations + that succeeded or skipped keep that outcome and only the others run again. + + The checkpoints themselves are never deleted: one marker write costs the same whatever the indexed task + count, and the store is scoped to the parent task instance, so state a sub-task stored for itself + is never touched. They expire with the store's default retention (``[state_store] + default_retention_days``, 30 days unless configured, 0 disables expiry), which also bounds how + long a task that exhausted its retries keeps them. Keeping them until then is intended: a manual + clear of such a task resumes from the checkpoints instead of re-running every index, and the + marker is what tells a clear-after-success apart from that. + """ + + # Same namespace as IndexedTaskState.build_key, for the same reason. + COMPLETION_KEY = "_iterable_completed" + + def __init__(self, context: Context) -> None: + self._store = context["task_state_store"] + self._try_number = context["ti"].try_number + self.trust_checkpoints = False + # The attempt from which checkpoints may be resumed; older ones predate a manual clear. + self.since = 0 + + def __enter__(self) -> Checkpoints: + if self._try_number > 1: + marker = self._store.get(self.COMPLETION_KEY) + if isinstance(marker, Mapping) and marker.get("completed"): + self._store.set(self.COMPLETION_KEY, {"completed": False, "since": self._try_number}) + else: + self.trust_checkpoints = True + since = marker.get("since") if isinstance(marker, Mapping) else None + if isinstance(since, int): + self.since = since + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + if exc_type is None or issubclass(exc_type, AirflowSkipException): + self._store.set(self.COMPLETION_KEY, {"completed": True, "try_number": self._try_number}) + + +class IndexedTaskInstanceNotStarted(Exception): + """ + The outcome of an indexed task instance pulled before a kill and reached after it. + + Never raised: ``IterableOperator._run_task`` returns it as the instance's outcome, so that + :meth:`IndexedTaskOutcomes.record` counts it apart from the failures. It ran no code, wrote + no checkpoint and fires no callback; the next attempt runs it. + """ + + +class IterationState: + """ + The state of one run of an iterated task, kept apart from the operator's configuration. + + It holds what the run needs to remember while it is going: the sub-operators in flight and + those already killed, so that :meth:`IterableOperator.on_kill` reaches each one once; the + stop flag a kill sets, which the executor consults before starting the next indexed task; + and the resolved input. What the indexed tasks ended with is :class:`IndexedTaskOutcomes`. + + A 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, 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._failure_for_the_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 the + parent's timeout), none when the retry policy decides FAIL, and otherwise a retry while the + parent has attempts left (``IndexedTaskInstance.is_eligible_to_retry``). The policy's + decision is the one :meth:`_failure_for_the_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 isinstance(raised, FAIL_WITHOUT_RETRY) or not isinstance(raised, (Exception, AirflowTaskTimeout)): Review Comment: Right, reproduced: with one item timing out on its own and one raising `ValueError`, both items fired `on_failure_callback` while the runner retried the group. `_task_will_retry` now judges a group member by member, so it counts as retryable when each failure is, and the test with exactly that pair sees a retry callback for both items (25f3154699). 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]
