dabla commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r4027410570
########## task-sdk/src/airflow/sdk/definitions/iterableoperator.py: ########## @@ -0,0 +1,725 @@ +# +# 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 copy +import os +import threading +import warnings +from collections.abc import Iterable, Mapping, Sequence +from itertools import repeat +from typing import TYPE_CHECKING, Any +from uuid import UUID + +try: + # Python 3.11+ + BaseExceptionGroup +except NameError: + from exceptiongroup import BaseExceptionGroup + +from airflow.sdk import BaseXCom, TaskInstanceState +from airflow.sdk.bases.operator import BaseAsyncOperator, BaseOperator, event_loop +from airflow.sdk.bases.xcom import XComIterable +from airflow.sdk.definitions._internal.expandinput import BatchedExpandInput +from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAliasEvent, AssetUniqueKey +from airflow.sdk.definitions.context import clone_context +from airflow.sdk.definitions.mappedoperator import MappedOperator +from airflow.sdk.definitions.xcom_arg import XComArg +from airflow.sdk.exceptions import ( + AirflowFailException, + AirflowRescheduleException, + AirflowSkipException, + DagRunTriggerException, + DownstreamTasksSkipped, + TaskDeferred, +) +from airflow.sdk.execution_time.comms import DeadlockImminentError +from airflow.sdk.execution_time.context import OutletEventAccessors, context_update_for_unmapped +from airflow.sdk.execution_time.executor import AsyncAwareExecutor, TaskExecutor +from airflow.sdk.execution_time.task_runner import IndexedTaskInstance, IndexedTaskState + +if TYPE_CHECKING: + import jinja2 + + from airflow.sdk.definitions._internal.expandinput import ExpandInput + from airflow.sdk.definitions.context import Context + from airflow.sdk.types import OutletEventAccessorsProtocol + + +def _serialize_outlet_events(accessors: OutletEventAccessors) -> list[dict[str, Any]]: + """ + Snapshot the outlet asset events one sub-task recorded into a JSON-safe list. + + Persisted on the sub-task's checkpoint so a later attempt can replay them via + ``_replay_outlet_events`` when the sub-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, + } + ) + return events + + +def _merge_outlet_events(target: OutletEventAccessorsProtocol, source: OutletEventAccessors) -> None: + """ + Merge every outlet asset event recorded in ``source`` into ``target``. + + Used both to fold a sub-task's isolated accessor into the IterableOperator's shared + ``context["outlet_events"]`` right after it succeeds, and to replay a checkpointed + snapshot (via ``_replay_outlet_events``) for a sub-task skipped on retry. + """ + for asset_or_alias, accessor in source.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 _replay_outlet_events(target: OutletEventAccessorsProtocol, events: list[dict[str, Any]]) -> None: + """ + Re-populate ``target`` with events a sub-task recorded on a previous attempt. + + 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. + """ + replayed = OutletEventAccessors() + for event in events: + if event["kind"] == "asset": + accessor = replayed[Asset(name=event["name"], uri=event["uri"])] + accessor.extra.update(event["extra"]) + if event["partition_keys"]: + accessor.add_partitions(event["partition_keys"]) + else: + accessor = replayed[AssetAlias(name=event["source_alias_name"])] + accessor.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"], + ) + ) + _merge_outlet_events(target, replayed) + + +class IterableOperator(BaseOperator): + """ + Operator used for Task Iteration (TI) 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``. + + **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, all checkpoints are dropped + 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. + + :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 (or batches) to iterate + over. Its ``iter_values(context)`` method is used to produce 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:: + 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 + (:class:`~airflow.sdk.exceptions.DownstreamTasksSkipped`, raised e.g. by + ``ShortCircuitOperator``) 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. Either exception raised by a + sub-task 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. If *every* sub-task raises + :class:`~airflow.sdk.exceptions.AirflowSkipException`, a single ``AirflowSkipException`` is + re-raised so the IterableOperator is marked ``SKIPPED`` instead of ``UP_FOR_RETRY``. Any other + mix of sub-task exceptions (including a partial skip alongside other failures) is aggregated + into a :class:`BaseExceptionGroup` and treated as a regular retryable failure. + + .. warning:: + **Async sub-tasks must only make 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``), it can collide with another sub-task's async SDK + call that is concurrently holding the communication lock, which is detected and raised + eagerly as a non-retryable failure rather than silently deadlocking. Use the async-safe + equivalents inside async operators: :meth:`~airflow.sdk.bases.hook.BaseHook.aget_connection`/ + ``aget_hook``, ``ti.axcom_pull``. ``Variable`` has no async equivalent yet. + + .. warning:: + **``execution_timeout`` is only enforced for async sub-tasks.** + + Async sub-tasks (instances of :class:`~airflow.sdk.bases.operator.BaseAsyncOperator`) respect + ``execution_timeout`` via ``asyncio.wait_for``. Sync sub-tasks run in worker threads and rely on + :class:`~airflow.sdk.execution_time.timeout.TimeoutPosix`, which requires ``signal.SIGALRM`` and + only works in the main thread. Because sync sub-tasks execute in a thread pool, ``SIGALRM`` cannot + be delivered to them, so their ``execution_timeout`` is silently ignored. Use + :class:`~airflow.sdk.bases.operator.BaseAsyncOperator` if per-sub-task time limits are required. + """ + + _operator: MappedOperator + expand_input: ExpandInput + partial_kwargs: dict[str, Any] + shallow_copy_attrs: Sequence[str] = ( + "_operator", + "expand_input", + "partial_kwargs", + "_log", + "_active_sub_operators", + "_active_sub_operators_lock", + ) + + 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") + + super().__init__( + **{ + **kwargs, + "task_id": operator.task_id, + "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, + "execution_timeout": None, + "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), + "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, + } + ) + 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}") + # Known v1 limitation: pool_slots is reserved once by the scheduler for this IterableOperator TI, + # but up to max_workers sub-tasks run concurrently inside it. Operators that set pool_slots > 1 to + # protect a shared resource (e.g. a DB connection pool) will be under-accounted — the pool sees one + # reservation while max_workers connections can be active simultaneously. A proper fix requires the + # scheduler to reserve pool_slots * max_workers slots, which needs scheduler-side changes. + self.max_workers = task_concurrency if task_concurrency is not None else (os.cpu_count() or 1) + if operator.execution_timeout and not issubclass(operator.operator_class, BaseAsyncOperator): + warnings.warn( + f"Operator {operator.task_id!r} has execution_timeout set, but sync operators run in " + "worker threads where TimeoutPosix (SIGALRM) cannot be delivered. " + "The execution_timeout will not be enforced for sync sub-tasks inside IterableOperator. " + "Use BaseAsyncOperator if per-sub-task time limits are required.", + UserWarning, + stacklevel=2, + ) + # 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) + # Populated with each sub-task's unmapped operator while it is actively executing, so + # on_kill() (see below) can propagate a kill/timeout signal to whichever sub-tasks happen + # to be in flight; guarded by a lock since sub-tasks execute concurrently. + self._active_sub_operators: set[BaseOperator] = set() + self._active_sub_operators_lock = threading.Lock() + + 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. + with self._active_sub_operators_lock: + active_operators = list(self._active_sub_operators) + for operator in active_operators: + try: + operator.on_kill() + except Exception: + self.log.exception("Error calling on_kill() for sub-task operator %s", operator.task_id) + + @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 task_type(self) -> str: + # self._operator is the MappedOperator/DecoratedMappedOperator wrapper used to unmap + # each sub-task; its own task_type field already holds the wrapped operator's class + # name (set from operator_class.__name__ when the wrapper is built), which is what + # should be reported here (e.g. as TaskInstance.operator), not the wrapper's class name. + return self._operator.task_type + + @property + def operator_name(self) -> str: + # Same forwarding as task_type above: self._operator already resolves to the wrapped + # operator's display name (e.g. a @task-decorated callable's custom_operator_name), + # so report that instead of falling back to this wrapper's own task_type. + 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 TaskExecutor 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(), + ) + + @classmethod + async def axcom_push(cls, task: IndexedTaskInstance, value: Any) -> None: + await task.axcom_push(key=BaseXCom.XCOM_RETURN_KEY, value=value) + + def _run_tasks( + self, + context: Context, + tasks: Iterable[IndexedTaskInstance], + ) -> XComIterable | None: + exceptions: list[Exception] = [] + total = 0 + do_xcom_push = True + + self.log.info("Running tasks with %d workers", self.max_workers) + + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=self.max_workers) as executor: + for task, _result, raised in executor.map( + self._run_task, + repeat(executor), + repeat(context), + tasks, + ): + total += 1 + do_xcom_push = task.do_xcom_push + + if raised is None: + continue + + if isinstance(raised, TaskDeferred): + raise AirflowFailException( + f"Sub-task {task.task_id}[{task.index}] attempted to defer. " + "Deferrable operators are not supported inside IterableOperator." + ) + + if isinstance(raised, (DagRunTriggerException, DownstreamTasksSkipped)): + raise AirflowFailException( + f"Sub-task {task.task_id}[{task.index}] raised " + f"{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." + ) from raised + + if isinstance(raised, AirflowRescheduleException): + raise AirflowFailException( + f"Sub-task {task.task_id}[{task.index}] 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." + ) from raised + + # Non-Exception BaseExceptions (e.g. DeadlockImminentError, + # KeyboardInterrupt, SystemExit) must never be swallowed: they + # signal conditions where continuing iteration is meaningless + # because every subsequent task would fail for the same reason. + # Re-raise immediately to stop all task iteration. + if isinstance(raised, DeadlockImminentError): + raise AirflowFailException( + f"Sub-task {task.task_id}[{task.index}] made a synchronous SDK call " + "(e.g. Variable.get, BaseHook.get_connection/get_hook, ti.xcom_pull, or a " + "sync callback) from an async sub-task. Synchronous SDK calls are not safe " + "inside an async operator's aexecute(): they can collide with another " + "concurrently running sub-task's async SDK call and deadlock the event " + "loop, so this is detected and raised eagerly instead. Use the async-safe " + "equivalents (e.g. Variable.aget/aset, Hook.aget_connection/aget_hook, ti.axcom_pull) inside " + "async operators." + ) from raised + if not isinstance(raised, Exception): + raise AirflowFailException( + f"Sub-task {task.task_id}[{task.index}] raised a non-Exception BaseException: " + f"{type(raised).__name__}: {raised}" + ) from raised + + self.log.exception( + "An exception occurred for task_id %s with index %s", + task.task_id, + task.index, + exc_info=raised, + ) + exceptions.append(raised) + + if exceptions: + # An AirflowFailException means the task must not be retried — propagate the first one + # directly rather than burying it in a BaseExceptionGroup, which _run_task_and_map_outcome + # would otherwise dispatch to the generic BaseException branch (i.e. eligible for retry). + for exc in exceptions: + if isinstance(exc, AirflowFailException): + raise exc + # If every sub-task was skipped, propagate a single AirflowSkipException so the runner + # marks the whole IterableOperator SKIPPED instead of UP_FOR_RETRY. + if len(exceptions) == total and all(isinstance(exc, AirflowSkipException) for exc in exceptions): + raise exceptions[0] + raise BaseExceptionGroup("Multiple sub-task failures", exceptions) + # Every index succeeded: drop the checkpoints. A manual clear does not reset try_number + # (models.taskinstance.clear_task_instances only raises max_tries), so leaving these behind + # would make the next attempt's try_number > 1, causing _run_task to treat every index as + # already-succeeded and replay stale results instead of re-running anything. + context["task_state_store"].clear() + if do_xcom_push: + return XComIterable( + task_id=self.task_id, + dag_id=self.dag_id, + run_id=context["run_id"], + length=len(self.expand_input), + map_index=context["ti"].map_index, + ) + return None + + async def _run_task( + self, + executor: AsyncAwareExecutor, + context: Context, + task: IndexedTaskInstance, + ) -> tuple[IndexedTaskInstance, Any | None, BaseException | None]: + indexed_task_state = await task.aget_state() + # We only rely on task state if it's not the first attempt + if ( + indexed_task_state is not None + and task.try_number > 1 + and indexed_task_state.status == TaskInstanceState.SUCCESS + ): + self.log.info( + "Skipping task instance %s for %s which already finished successfully after %s attempts", + task.index, + task.task_id, + indexed_task_state.try_number, + ) + await self.axcom_push(task, indexed_task_state.result) + if indexed_task_state.outlet_events: + _replay_outlet_events(context["outlet_events"], indexed_task_state.outlet_events) Review Comment: It is the second case: the failed attempt never commits. Outlet events are serialised only in `_handle_current_task_success` and travel only on `TISuccessStatePayload`; the retry and terminal payloads have no such field. So on the first failed attempt item 5's events never reached the server, and the replay on the retry that skips item 5 is the only thing that preserves them; there is no double emit. Written into the `IndexedTaskState.outlet_events` comment and the `_replay_outlet_events` docstring in baa4934825. --- Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
