uranusjr commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r3818998449
########## task-sdk/src/airflow/sdk/definitions/iterableoperator.py: ########## @@ -0,0 +1,540 @@ +# +# 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 +from collections import deque +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, timezone +from airflow.sdk.bases.operator import 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, + AirflowRescheduleException, + AirflowRescheduleTaskInstanceException, + 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.providers.standard.triggers.temporal import DateTimeTrigger + from airflow.sdk.definitions._internal.expandinput import ExpandInput + from airflow.sdk.definitions.context import Context + from airflow.sdk.execution_time.lazy_sequence import XComIterable + + +ExternalDateTimeTrigger: type[DateTimeTrigger] | None + +try: + from airflow.providers.standard.triggers.temporal import DateTimeTrigger as ExternalDateTimeTrigger +except ModuleNotFoundError: + # If the providers package with DateTimeTrigger is not available (e.g. in + # minimal installs or tests), set the symbol to None so callers can + # explicitly check for availability. Using hasattr(self, DateTimeTrigger) + # is incorrect because hasattr expects a string attribute name. + ExternalDateTimeTrigger = None + + +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 is a + lightweight, non-retrying wrapper — retries, timeouts and deferred + execution are handled by the individual indexed task instances that the + IterableOperator creates for each element produced by the + ``expand_input``. + + 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``. + + :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``). Note that the IterableOperator + overrides retry-related parameters because retries are managed by + the per-index tasks. + + :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. + """ + + _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": 0, # We should not retry the IterableOperator, only the indexed runtime ti's should be retried + # Known v1 limitation: there is no durable per-index checkpoint. On a worker crash mid-iteration, + # clearing or re-running the parent TI reconstructs every IndexedTaskInstance with xcom_pushed=False + # and re-executes from index 0, duplicating any external effects that already completed. Until + # durable checkpointing is implemented, iteration should be constrained to idempotent work. + # AIP-103 Task State Management could be a solution to this known v1 limitation. + "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) + 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[BaseException] = [] + reschedule_date = timezone.utcnow() + failed_tasks: deque[IndexedTaskInstance] = deque() + do_xcom_push = True + + self.log.info("Running tasks with %d workers", self.max_workers) + + while True: + 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." + ) + + if isinstance(raised, AirflowRescheduleException): + if not isinstance(raised, AirflowRescheduleTaskInstanceException): + raise AirflowFailException( + f"Sub-task {task.task_id}[{task.index}] attempted to reschedule. " + "Reschedule-mode sensors are not supported inside IterableOperator." + ) + reschedule_date = max(reschedule_date, raised.reschedule_date) + self.log.exception( + "An exception occurred for task_id %s with index %s, it has been rescheduled at %s", + task.task_id, + task.index, + reschedule_date, + ) + failed_tasks.append(raised.task) + continue + + # 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 not failed_tasks: + if exceptions: + # If this IterableOperator is backed by a batched expand input + # (created from a MappedIterableOperator), the parent mapped + # task should never be retried; retries are handled by the + # individual indexed runtime tasks. In that case raise + # AirflowFailException to mark failure without retrying the + # parent TaskInstance. For regular (non-batched) IterableOperator + # behavior, preserve the previous behavior and raise the + # BaseExceptionGroup so callers/tests that expect it keep working. + if isinstance(self.expand_input, BatchedExpandInput): + raise AirflowFailException(f"Multiple sub-task failures: {exceptions}") + raise BaseExceptionGroup("Multiple sub-task failures", exceptions) Review Comment: This should probably be folded into ExpandInput subclasses (something like `raise self.expand_input.wrap_exceptions(exceptions)`?) instead of doing isinstance checks here. -- 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]
