dabla commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r3991928513
########## task-sdk/src/airflow/sdk/definitions/iterableoperator.py: ########## @@ -0,0 +1,487 @@ +# +# 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 warnings +from collections.abc import Iterable, Mapping, Sequence +from itertools import repeat +from typing import TYPE_CHECKING, Any + +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.definitions._internal.expandinput import BatchedExpandInput +from airflow.sdk.definitions.context import clone_context +from airflow.sdk.definitions.mappedoperator import MappedOperator +from airflow.sdk.definitions.xcom_arg import MapXComArg, XComArg # noqa: F401 +from airflow.sdk.exceptions import ( + AirflowFailException, + TaskDeferred, +) +from airflow.sdk.execution_time.executor import AsyncAwareExecutor, TaskExecutor +from airflow.sdk.execution_time.task_runner import IndexedTaskInstance + +if TYPE_CHECKING: + import jinja2 + + from airflow.sdk.bases.xcom import XComIterable + from airflow.sdk.definitions._internal.expandinput import ExpandInput + from airflow.sdk.definitions.context import Context + +_ITERABLE_CHECKPOINT_KEY_PREFIX = "_iterable_task_" + + +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 pending/failed sub-tasks are recreated + with their accumulated ``try_number`` so that retries are not wasted. + + :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 meaningfully supported: the exception is treated like any other sub-task failure and + counts towards the IterableOperator's own ``retries``, but the requested ``reschedule_date`` is not + honored — the worker is not released and the next attempt follows the IterableOperator's own + ``retry_delay`` instead of waiting until ``reschedule_date``. + + .. 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", + ) + + def __init__( + self, + *, + operator: MappedOperator, + expand_input: ExpandInput, + **kwargs, + ): + 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, + "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, + "priority_weight": operator.priority_weight, + "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, + "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, + ) + XComArg.apply_upstream_relationship(self, self.expand_input.value) + + @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: + return self._operator.__class__.__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 _unmap_operator( + self, context: Context, mapped_kwargs: Context, jinja_env: jinja2.Environment + ) -> BaseOperator: + from airflow.sdk.execution_time.context import context_update_for_unmapped + + unmapped_task = self._operator.unmap(mapped_kwargs) + # Make sure deferred operators will always raise a DeferredTask exception when executed + unmapped_task.start_from_trigger = False + 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(), + ) + return unmapped_task + + async def _xcom_push(self, task: IndexedTaskInstance, value: Any) -> None: + if task.xcom_pushed: + self.log.debug( + "XCom already pushed for task_id %s with index %s", + task.task_id, + task.index, + ) + else: + self.log.debug( + "Pushing XCom for task_id %s with index %s", + task.task_id, + task.index, + ) + + 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] = [] + 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, + ): + 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." + ) + + # 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 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: + raise BaseExceptionGroup("Multiple sub-task failures", exceptions) + if do_xcom_push: + from airflow.sdk.bases.xcom import XComIterable + + 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 + + def _checkpoint_key(self, index: int) -> str: + return f"{_ITERABLE_CHECKPOINT_KEY_PREFIX}{index}" + + async def _run_task( + self, + executor: AsyncAwareExecutor, + context: Context, + task: IndexedTaskInstance, + ) -> tuple[IndexedTaskInstance, Any | None, BaseException | None]: + task_state_store = context["task_state_store"] + checkpoint = await task_state_store.aget(self._checkpoint_key(task.index)) + if isinstance(checkpoint, dict): + try_number = checkpoint.get("try_number", 0) + if not isinstance(try_number, int): + try_number = 0 + if checkpoint.get("status") == "succeeded": + self.log.info( + "Skipping task instance %s for %s which already finished successfully after %s attempts", + task.index, + task.task_id, + try_number + 1, + ) + return task, None, None + task.try_number = try_number + + try: + if task.is_async: + result = await self._run_async_operator(context, task) + else: + result = await executor.run_sync(self._run_operator, context, task) + + if result is not None and task.do_xcom_push: + await self._xcom_push(task, result) + + await task_state_store.aset( + self._checkpoint_key(task.index), + {"status": "succeeded", "try_number": task.try_number}, + ) + return task, result, None + except BaseException as e: + await task_state_store.aset( + self._checkpoint_key(task.index), + {"status": "pending", "try_number": task.try_number}, + ) + return task, None, e + + def _run_operator(self, context: Context, task_instance: IndexedTaskInstance): + with TaskExecutor(task_instance=task_instance) as executor: + return executor.run( + context={ + **clone_context(context), + **{ + "ti": task_instance, + "task_instance": task_instance, + }, + } + ) + + async def _run_async_operator(self, context: Context, task_instance: IndexedTaskInstance): Review Comment: This is done explicitly, that's why the DeadlockImminentError has been implemented, to avoid deadlocks when users are mixing sync calls within async context. The has even been an additional [fix](https://github.com/apache/airflow/pull/71890) regarding this which has been merged and backported to 3.3.2. For Variable, there is another [PR](https://github.com/apache/airflow/pull/72329) open which adds the async accessors. -- 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]
