kaxil commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r3625190606
##########
airflow-core/src/airflow/serialization/definitions/mappedoperator.py:
##########
@@ -519,6 +528,9 @@ def _(task: SerializedBaseOperator | TaskSDKBaseOperator,
run_id: str, *, sessio
def _(task: SerializedMappedOperator | TaskSDKMappedOperator, run_id: str, *,
session: Session) -> int:
from airflow.serialization.serialized_objects import BaseSerialization,
_ExpandInputRef
+ if task.batch_size > 0:
Review Comment:
This returns `task.batch_size` unconditionally, but
`get_parse_time_mapped_ti_count` (line 467) clamps to `min(batch_size,
mapped_ti_count)`. When `batch_size > item_count` the two disagree:
`batch(size=5).iterate(x=[a, b, c])` makes the scheduler create 5 TIs from this
runtime count, then `_check_for_removed_or_restored_tasks` uses the parse-time
count of 3 and marks map_index 3 and 4 REMOVED on every run, while downstream
`expanded_ti_count` still expects 5. The round-robin partition keeps the
removed batches empty so there's no item-level data loss, but it produces
spurious REMOVED TIs and a downstream count mismatch on a plausible authoring
input. Clamp here too, e.g. `return min(task.batch_size, mapped_ti_count)`
after the real length is computed.
##########
task-sdk/src/airflow/sdk/definitions/iterableoperator.py:
##########
@@ -0,0 +1,525 @@
+#
+# 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 functools import cached_property
+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,
+ 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.
+ """
+
+ _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
+ "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}")
+ 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
+
+ @cached_property
+ def timeout(self) -> float | None:
+ if self._operator.execution_timeout:
+ return self._operator.execution_timeout.total_seconds()
+ return None
+
+ 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,
+ timeout=self.timeout,
+ ):
+ 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,
AirflowRescheduleTaskInstanceException):
Review Comment:
This only catches `AirflowRescheduleTaskInstanceException`, but a
reschedule-mode sensor raises the base `AirflowRescheduleException` (the TI
variant subclasses it, exceptions.py:180). So a base reschedule never matches
here and falls through to `exceptions.append(raised)` below, failing the whole
IterableOperator.
It's worse with the default `retries=0`: in `TaskExecutor.__exit__`
(executor.py:296-306) a non-deferred exception is gated on the retry budget, so
`next_try_number(1) > max_tries(0)` re-raises the original base
`AirflowRescheduleException` as a hard failure on the very first not-ready
poke. Even with `retries>0`, the sensor's own `reschedule_date` (its poke
interval) gets replaced by `next_retry_datetime()` (retry backoff), and it
fails after `retries` pokes regardless of readiness.
Either reject reschedule-mode sensors explicitly the way deferrable
operators are rejected at line 279, or handle the base
`AirflowRescheduleException` here, preserving `raised.reschedule_date` and not
consuming the retry budget (the normal runner does this in task_runner.py).
##########
task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py:
##########
@@ -166,6 +169,80 @@ def __getitem__(self, key: int | slice) -> T | Sequence[T]:
return XCom.deserialize_value(_XComWrapper(msg.root))
+class XComIterable(Sequence):
+ """An iterable that lazily fetches XCom values one by one instead of
loading all at once."""
+
+ def __init__(self, task_id: str, dag_id: str, run_id: str, length: int,
map_index: int | None = None):
+ self.task_id = task_id
+ self.dag_id = dag_id
+ self.run_id = run_id
+ self.length = length
+ self.map_index = map_index
+
+ def __iter__(self) -> Iterator:
+ return _XComIterator(self)
+
+ def __len__(self) -> int:
+ return self.length
+
+ @overload
+ def __getitem__(self, key: int) -> Any: ...
+
+ @overload
+ def __getitem__(self, key: slice) -> Sequence[Any]: ...
+
+ def __getitem__(self, key: int | slice) -> Any | Sequence[Any]:
+ """Allow direct indexing so this works like a sequence."""
+ if isinstance(key, slice):
+ start, stop, step = key.indices(len(self))
+ return [self[i] for i in range(start, stop, step)]
Review Comment:
`XComIterable` fetches one XCom per index -- `__getitem__` calls
`XCom.get_one` per element (line 203), and this slice branch is `[self[i] for i
in range(...)]`, so iterating or slicing a 1000-item iterate result is 1000 API
round-trips. `LazyXComSequence` in this same file already batches this with a
single `GetXComSequenceSlice` comms message. Give `XComIterable` the same
batched slice-fetch, or share a base with `LazyXComSequence`, so full iteration
is one or a few round-trips instead of N.
##########
task-sdk/src/airflow/sdk/definitions/iterableoperator.py:
##########
@@ -0,0 +1,525 @@
+#
+# 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 functools import cached_property
+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,
+ 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.
+ """
+
+ _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
Review Comment:
The parent is forced to `retries=0` while per-index attempt and success
state live only as in-memory `IndexedTaskInstance` objects, and `execute()`
doesn't consult already-pushed results as completion checkpoints. So if item 0
completes a non-idempotent write and the worker dies during item 1, there's no
scheduler-visible retry, and clearing or rerunning the parent reconstructs
every index with `xcom_pushed=False` and repeats item 0, duplicating the
external effect. What's the intended crash-recovery story here? If durable
per-index checkpointing is out of scope for v1, it would help to say so and
constrain iteration to idempotent work explicitly, rather than silently
disabling parent retries.
##########
task-sdk/src/airflow/sdk/execution_time/executor.py:
##########
@@ -0,0 +1,367 @@
+#
+# 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 contextvars
+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 Callable, Iterable, Iterator
+from concurrent.futures import Executor, ThreadPoolExecutor
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any, cast
+
+from airflow.sdk import BaseAsyncOperator, BaseOperator, TaskInstanceState,
timezone
+from airflow.sdk.bases.operator import ExecutorSafeguard
+from airflow.sdk.definitions._internal.logging_mixin import LoggingMixin
+from airflow.sdk.exceptions import AirflowRescheduleTaskInstanceException,
TaskDeferred
+from airflow.sdk.execution_time.callback_runner import create_executable_runner
+from airflow.sdk.execution_time.context import context_get_outlet_events,
set_current_context
+from airflow.sdk.execution_time.task_runner import (
+ RuntimeTaskInstance,
+ _execute_task,
+ _run_task_state_change_callbacks,
+)
+
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk import Context
+ from airflow.sdk.execution_time.task_runner import IndexedTaskInstance
+
+
+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.
+ """
+
+ def __init__(self, loop: AbstractEventLoop, max_workers: int | None =
None):
+ 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._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 resources (e.g., threading
+ # locks). Without waiting, cancelled tasks that hold _thread_lock
+ # never execute their finally blocks, permanently leaking the lock
+ # and causing subsequent comms.send() calls to deadlock.
+ 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=10.0,
+ )
+ )
+
+ self._thread_pool.shutdown(wait=wait, cancel_futures=cancel_futures)
+
+ 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)
+ return await wrap_future(future, loop=self._loop)
+
+ def map(
+ self,
+ fn: Callable[..., Any],
+ *iterables: Iterable[Any],
+ timeout: float | None = None,
+ chunksize: int = 1,
+ ) -> Iterator[Any]:
+ """Apply fn to iterables and stream results in completion order."""
+ if chunksize < 1:
+ raise ValueError("chunksize must be >= 1")
+
+ if self._shutdown:
+ raise RuntimeError("cannot schedule new futures after shutdown")
+
+ start = time.monotonic()
+ iterator = zip(*iterables)
+ pending: dict[Future[Any], Future[Any]] = {}
+ exhausted = False
+
+ def _remaining_timeout() -> float | None:
+ if timeout is None:
+ return None
+ remaining = timeout - (time.monotonic() - start)
+ if remaining <= 0:
+ raise TimeoutError()
+ return remaining
+
+ def _submit_next() -> bool:
+ nonlocal exhausted
+
+ if exhausted:
+ return False
+
+ try:
+ args = next(iterator)
+ except StopIteration:
+ exhausted = True
+ return False
+
+ future = self.submit(fn, *args)
+ pending[future] = future
+ return True
+
+ def _fill_pending() -> None:
+ """Submit tasks until pending reaches max_workers or iterator is
exhausted."""
+ while len(pending) < self._max_workers and _submit_next():
+ pass
+
+ # Submit up to max_workers tasks initially
+ _fill_pending()
+
+ while pending:
+ wait_timeout = _remaining_timeout()
+ done, _ = self._loop.run_until_complete(
+ wait(set(pending), timeout=wait_timeout,
return_when=FIRST_COMPLETED)
+ )
+
+ if not done:
+ raise TimeoutError()
+
+ for completed in done:
+ pending.pop(completed)
+ _fill_pending()
+ yield completed.result()
+
+
+class TaskExecutor(LoggingMixin):
+ """Base class to run an operator or trigger with given task context and
task instance."""
+
+ def __init__(
+ self,
+ task_instance: IndexedTaskInstance,
+ ):
+ super().__init__()
+ self.task_instance = task_instance
+ self._result: Any | None = None
+ self._start_time: float | None = None
+
+ @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 xcom_key(self):
+ return self.task_instance.xcom_key
+
+ @property
+ def operator(self) -> BaseOperator:
+ return self.task_instance.task
+
+ @property
+ def is_async(self) -> bool:
+ return self.task_instance.is_async
+
+ def run(self, context: Context):
+ with set_current_context(context):
+ return _execute_task(context, self.task_instance, self.log)
+
+ async def arun(self, context: Context):
+ with set_current_context(context):
+ return await _execute_async_task(context, self.task_instance,
self.log)
+
+ def __enter__(self):
+ self._start_time = time.monotonic()
+
+ if self.log.isEnabledFor(logging.INFO):
+ self.log.info(
+ "Attempting running task %s of %s for %s with index %s in %s
mode.",
+ self.task_instance.try_number,
+ self.operator.retries,
+ 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 exc_value:
+ # Non-Exception BaseExceptions (e.g. DeadlockImminentError,
+ # KeyboardInterrupt, SystemExit) must never be retried: they
+ # signal conditions where continuing is meaningless.
+ # Re-raise immediately without retry.
+ if not isinstance(exc_value, Exception):
+ self.task_instance.state = TaskInstanceState.FAILED
+ raise exc_value
+ if not isinstance(exc_value, TaskDeferred):
+ if self.task_instance.next_try_number >
self.task_instance.max_tries:
+ self.log.error(
+ "Task instance %s for %s failed after %s attempts in
%.2f seconds due to: %s",
+ self.task_index,
+ self.task_instance.task_id,
+ self.task_instance.max_tries,
+ elapsed,
+ exc_value,
+ )
+ self.task_instance.state = TaskInstanceState.FAILED
+ raise exc_value
+ self.task_instance.try_number += 1
+ self.task_instance.end_date = timezone.utcnow()
+ self.task_instance.state = TaskInstanceState.UP_FOR_RESCHEDULE
+ raise AirflowRescheduleTaskInstanceException(
+ task=self.task_instance,
+ reschedule_date=self.task_instance.next_retry_datetime(),
+ )
+ raise exc_value
+
+ self.task_instance.state = TaskInstanceState.SUCCESS
+ if self.log.isEnabledFor(logging.INFO):
+ self.log.info(
+ "Task instance %s for %s finished successfully in %s attempts
in %.2f seconds",
+ self.task_index,
+ self.task_instance.task_id,
+ self.task_instance.next_try_number,
+ elapsed,
+ )
+
+
+async def _execute_async_task(context: Context, ti: RuntimeTaskInstance, log:
Logger):
+ """Execute Task (optionally with a Timeout) and push Xcom results."""
+ # set-up
+ task = cast("BaseAsyncOperator", ti.task)
+ execute = task.aexecute # here we must use aexecute instead of execute
+
+ # async tasks can't originate from deferred operator, so no need to check
next_method
+
+ ctx = contextvars.copy_context()
+ # Populate the context var so ExecutorSafeguard doesn't complain
+ ctx.run(ExecutorSafeguard.tracker.set, task)
+
+ outlet_events = context_get_outlet_events(context)
+
+ if (pre_execute_hook := task._pre_execute_hook) is not None:
+ create_executable_runner(pre_execute_hook, outlet_events,
logger=log).run(context)
+ if getattr(pre_execute_hook := task.pre_execute, "__func__", None) is not
BaseOperator.pre_execute:
+ create_executable_runner(pre_execute_hook, outlet_events,
logger=log).run(context)
+
+ _run_task_state_change_callbacks(task, "on_execute_callback", context, log)
Review Comment:
Only `on_execute_callback` is dispatched for indexed sub-tasks.
`on_success_callback`, `on_failure_callback`, and `on_retry_callback` on the
wrapped operator never fire per index, and the outer IterableOperator's
finalization doesn't inherit them from the mapped operator either. So cleanup,
notifications, credential refresh, and lineage hooks attached to the original
operator silently stop running once it's used with `.iterate()`. Dispatch the
matching callbacks for each indexed lifecycle transition, or reject operators
that carry them the way deferrable operators are rejected.
##########
task-sdk/src/airflow/sdk/definitions/context.py:
##########
@@ -95,6 +95,73 @@ class Context(TypedDict, total=False):
KNOWN_CONTEXT_KEYS: set[str] = set(Context.__annotations__.keys())
+def clone_context(context: Context) -> Context:
+ """
+ Create a safe, per-task copy of an execution ``Context`` for concurrent
execution.
+
+ The execution context is a mutable mapping that contains many nested
+ structures (``params``, ``templates_dict``, ``outlet_events``, ``dag_run``,
+ etc.). When running the same logical task concurrently (for example when
+ the ``IterableOperator`` spawns multiple indexed task instances that run in
+ parallel using threads, processes or asyncio tasks), those mutable objects
+ could be mutated by one indexed runtime and unintentionally observed by
+ another. That leads to subtle race conditions, corrupted state, and
+ flakiness in task execution.
+
+ ``clone_context`` returns a new :class:`Context` mapping where the
top-level
+ mapping is copied and specific mutable sub-objects that are commonly
+ mutated during execution are deep-copied or shallow-copied as appropriate:
+
+ - ``params`` and ``templates_dict`` are deep-copied because they are
+ dictionaries that users and operators commonly mutate.
+ - ``inlets`` and ``outlets`` are converted to new lists (shallow copy)
+ because the sequence identity must be isolated but the elements are
+ typically read-only accessor objects.
+ - ``dag_run`` is deep-copied because it carries nested state that must
+ not be shared between concurrent executions.
+ - ``outlet_events`` is intentionally **not** copied so that events emitted
+ by sub-tasks are captured in the parent's accessor and serialized when
+ the parent task completes. Concurrent writes to the same asset may race,
+ but that is expected behavior for parallel sub-tasks.
+
+ Use cases
+ - Multithreading: when using thread-based executors (``concurrent.futures``
+ ThreadPoolExecutor) multiple threads share memory; cloning prevents
+ concurrent mutation of shared structures.
+ - Async concurrency: when running coroutine-based tasks concurrently in
+ the same event loop, tasks may still mutate shared mappings; cloning
+ avoids interference.
+ - Multiprocessing: while processes do not share memory, cloning keeps the
+ semantics consistent and avoids accidentally capturing references that
+ would be pickled.
+
+ Performance
+ - The implementation intentionally copies only a small set of commonly
+ mutated fields rather than performing a blanket deep copy of the entire
+ context to keep the operation cheap. If future code stores additional
+ mutable state in the context that needs isolation, this function should
+ be extended appropriately.
+
+ :param context: The original execution context to clone.
+ :returns: A new :class:`Context` safe to hand to a concurrently running
task.
+
+ :meta private:
+ """
+ cloned_context = Context()
+ cloned_context.update(context)
+ cloned_context["params"] = copy.deepcopy(context.get("params", {}))
+ cloned_context["inlets"] = list(context.get("inlets", []))
+ cloned_context["outlets"] = list(context.get("outlets", []))
+ templates_dict = cloned_context.get("templates_dict")
+ if templates_dict is not None:
+ cloned_context["templates_dict"] = copy.deepcopy(templates_dict)
+ cloned_context["inlet_events"] = copy.deepcopy(context["inlet_events"])
+ # outlet_events is intentionally NOT copied - sub-tasks must emit into the
Review Comment:
Sharing the parent's `outlet_events` accessor across concurrently running
sub-tasks can drop a whole accessor's events, not just race on a single value.
`OutletEventAccessors.__getitem__` (execution_time/context.py) is
check-then-act: `if key not in self._dict: self._dict[key] =
OutletEventAccessor(...)`. With sync sub-tasks on worker threads and async
sub-tasks on the loop running in parallel, thread A can create an accessor and
write `extra["k1"]`, then thread B's `__getitem__` for the same asset
overwrites the slot with a fresh accessor, so A's accumulated events are gone
when the parent serializes `_dict`. The docstring frames same-asset writes as
expected, but this loses more than a racing value. Guard the insert (a lock or
`setdefault`), or give each sub-task its own accessor and merge at the end.
##########
task-sdk/src/airflow/sdk/definitions/iterableoperator.py:
##########
@@ -0,0 +1,525 @@
+#
+# 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 functools import cached_property
+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,
+ 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.
+ """
+
+ _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
+ "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}")
+ self.max_workers = task_concurrency if task_concurrency is not None
else (os.cpu_count() or 1)
Review Comment:
Worth confirming how the fan-out interacts with pool accounting. The
scheduler reserves the wrapped operator's `pool_slots` once for the single
IterableOperator TI, but this launches up to `task_concurrency` (or
`cpu_count()`) concurrent in-process executions, and batched mapping multiplies
that across batch TIs. So a database operator with `pool_slots=1` can issue
dozens of concurrent operations while Airflow accounts for one slot, which
defeats the pool as a protection for a shared resource. If AIP-104 already
scoped this as a known v1 limitation, fair enough; if not, the local
concurrency probably needs to be reflected in the pool reservation (or capped
when it can't be reserved) rather than defaulting pooled operators to CPU-count
fan-out.
##########
task-sdk/src/airflow/sdk/definitions/iterableoperator.py:
##########
@@ -0,0 +1,525 @@
+#
+# 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 functools import cached_property
+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,
+ 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.
+ """
+
+ _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
+ "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}")
+ 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
+
+ @cached_property
+ def timeout(self) -> float | None:
+ if self._operator.execution_timeout:
+ return self._operator.execution_timeout.total_seconds()
+ return None
+
+ 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,
+ timeout=self.timeout,
Review Comment:
`self.timeout` (line 197) reads the *wrapped* operator's `execution_timeout`
and passes it as `AsyncAwareExecutor.map`'s `timeout`, but that argument is a
cumulative wall-clock deadline for the whole stream, not a per-item limit. The
same `execution_timeout` is also applied per task downstream
(`_execute_async_task` / `_run_execute_callable`), so it ends up enforced twice
with two different meanings.
Concrete case: `execution_timeout=timedelta(seconds=30)`, `.iterate()` over
100 items at ~5s each with `max_workers=4` -- cumulative ~125s exceeds 30s, so
`map` raises a bare builtin `TimeoutError` mid-stream even though no single
item exceeded 30s. Since the parent has `retries=0`, the task fails permanently
and the final `XComIterable` is never returned (though orphan `return_value_i`
XComs were already pushed). Don't pass `execution_timeout` as the `map`
deadline -- either pass `timeout=None` and rely on the per-task enforcement, or
add a distinct, separately configured batch timeout.
##########
task-sdk/src/airflow/sdk/definitions/iterableoperator.py:
##########
@@ -0,0 +1,525 @@
+#
+# 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 functools import cached_property
+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,
+ 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.
+ """
+
+ _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
+ "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}")
+ 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
+
+ @cached_property
+ def timeout(self) -> float | None:
+ if self._operator.execution_timeout:
+ return self._operator.execution_timeout.total_seconds()
+ return None
+
+ 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,
+ timeout=self.timeout,
+ ):
+ 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,
AirflowRescheduleTaskInstanceException):
+ 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 from raised
Review Comment:
`raise AirflowFailException from raised` raises the class rather than an
instance, so this surfaces an empty-message failure. The other two sites (lines
279 and 322) instantiate it. This most likely meant `raise
AirflowFailException(...) from raised` with a message naming the non-Exception
BaseException that triggered it.
##########
task-sdk/src/airflow/sdk/definitions/context.py:
##########
@@ -95,6 +95,73 @@ class Context(TypedDict, total=False):
KNOWN_CONTEXT_KEYS: set[str] = set(Context.__annotations__.keys())
+def clone_context(context: Context) -> Context:
+ """
+ Create a safe, per-task copy of an execution ``Context`` for concurrent
execution.
+
+ The execution context is a mutable mapping that contains many nested
+ structures (``params``, ``templates_dict``, ``outlet_events``, ``dag_run``,
+ etc.). When running the same logical task concurrently (for example when
+ the ``IterableOperator`` spawns multiple indexed task instances that run in
+ parallel using threads, processes or asyncio tasks), those mutable objects
+ could be mutated by one indexed runtime and unintentionally observed by
+ another. That leads to subtle race conditions, corrupted state, and
+ flakiness in task execution.
+
+ ``clone_context`` returns a new :class:`Context` mapping where the
top-level
+ mapping is copied and specific mutable sub-objects that are commonly
+ mutated during execution are deep-copied or shallow-copied as appropriate:
+
+ - ``params`` and ``templates_dict`` are deep-copied because they are
+ dictionaries that users and operators commonly mutate.
+ - ``inlets`` and ``outlets`` are converted to new lists (shallow copy)
+ because the sequence identity must be isolated but the elements are
+ typically read-only accessor objects.
+ - ``dag_run`` is deep-copied because it carries nested state that must
+ not be shared between concurrent executions.
+ - ``outlet_events`` is intentionally **not** copied so that events emitted
+ by sub-tasks are captured in the parent's accessor and serialized when
+ the parent task completes. Concurrent writes to the same asset may race,
+ but that is expected behavior for parallel sub-tasks.
+
+ Use cases
+ - Multithreading: when using thread-based executors (``concurrent.futures``
+ ThreadPoolExecutor) multiple threads share memory; cloning prevents
+ concurrent mutation of shared structures.
+ - Async concurrency: when running coroutine-based tasks concurrently in
+ the same event loop, tasks may still mutate shared mappings; cloning
+ avoids interference.
+ - Multiprocessing: while processes do not share memory, cloning keeps the
+ semantics consistent and avoids accidentally capturing references that
+ would be pickled.
+
+ Performance
+ - The implementation intentionally copies only a small set of commonly
+ mutated fields rather than performing a blanket deep copy of the entire
+ context to keep the operation cheap. If future code stores additional
+ mutable state in the context that needs isolation, this function should
+ be extended appropriately.
+
+ :param context: The original execution context to clone.
+ :returns: A new :class:`Context` safe to hand to a concurrently running
task.
+
+ :meta private:
+ """
+ cloned_context = Context()
+ cloned_context.update(context)
+ cloned_context["params"] = copy.deepcopy(context.get("params", {}))
+ cloned_context["inlets"] = list(context.get("inlets", []))
+ cloned_context["outlets"] = list(context.get("outlets", []))
+ templates_dict = cloned_context.get("templates_dict")
+ if templates_dict is not None:
+ cloned_context["templates_dict"] = copy.deepcopy(templates_dict)
+ cloned_context["inlet_events"] = copy.deepcopy(context["inlet_events"])
Review Comment:
`clone_context` runs once per sub-task (via `_run_operator` /
`_run_async_operator`), and it deep-copies `inlet_events` here and `dag_run` on
line 161 every time. For a 1000-item iterate that's ~2000 extra `copy.deepcopy`
graph-walks on objects sub-tasks don't mutate: `InletEventsAccessors` builds
its state in `__init__` and its `__getitem__` returns a fresh accessor per call
(read-only during execute), and `dag_run` is a read-only data carrier.
Isolating `params` / `templates_dict` (which rendering can mutate) is enough.
Either drop these two deepcopies, or take one clone before the `map` loop and
reuse it. Note
`test_context.py::test_clone_context_deep_and_shallow_copy_semantics` asserts
`inlet_events` / `dag_run` are copied, so it would change with the fix.
--
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]
