kaxil commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r3471412951
##########
airflow-core/src/airflow/serialization/definitions/mappedoperator.py:
##########
@@ -515,6 +515,11 @@ 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
+ batch_size = task.partial_kwargs.get("batch_size")
Review Comment:
This `batch_size` special-case is only in the runtime `get_mapped_ti_count`.
`get_parse_time_mapped_ti_count` (line 458) wasn't given the same treatment and
still returns the raw expand-input length.
So `op.partial(...).batch(size=2).iterate(x=[0..9])` creates 10 parent
mapped TIs at dag-run creation instead of 2; `BatchedExpandInput.iter_values()`
then only fills indexes 0 and 1, and indexes 2..9 run with empty work and
succeed silently. Override the parse-time count for the batched case too, and
define behavior for `batch_size > input_length`.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -2091,9 +2223,12 @@ def _execute_task(context: Context, ti:
RuntimeTaskInstance, log: Logger):
# Populate the context var so ExecutorSafeguard doesn't complain
ctx.run(ExecutorSafeguard.tracker.set, task)
- # Export context in os.environ to make it available for operators to use.
+ # Export context vars thread-safely to avoid race conditions in concurrent
execution
+ # (e.g., IterableOperator with AsyncAwareExecutor). Use
get_airflow_context_var() to read these.
+ # We intentionally do NOT update os.environ here as that would cause races
between tasks.
+ # We set these in the copied context so they're available inside
ctx.run(execute, ...).
airflow_context_vars = context_to_airflow_vars(context,
in_env_var_format=True)
- os.environ.update(airflow_context_vars)
+ ctx.run(_airflow_context_vars.set, airflow_context_vars)
Review Comment:
This drops the `AIRFLOW_CTX_*` env-var contract for every task, not just
iterated ones. `_execute_task` is the common execution path, and removing
`os.environ.update(airflow_context_vars)` means provider code that reads these
from the environment now gets `None`: trino's `impersonate_as_owner` reads
`os.getenv("AIRFLOW_CTX_DAG_OWNER")` (trino.py:205), and amazon's
user-agent/team metadata reads
`os.environ.get("AIRFLOW_CTX_DAG_ID"/"AIRFLOW_CTX_TEAM_NAME")`
(base_aws.py:596/611/617). The replacement `_airflow_context_vars` contextvar
has no production readers -- `get_airflow_context_var()` is only called by the
new test.
The race you're avoiding is real for the concurrent in-process case, but the
fix shouldn't remove the global env contract. Keep exporting to `os.environ`
for normal task execution and isolate the iterate path another way (or run
indexed sub-tasks in processes if env isolation is required). Note
`test_task_runner.py:1775` currently asserts `os.environ` is *not* updated, so
the test locks in the regression rather than catching it.
##########
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 asyncio
+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,
DecoratedDeferredAsyncOperator, 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,
+ AirflowTaskTimeout,
+ 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 Dynamic Task Iteration (DTI) 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``.
+ """
+
+ _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)
+ self._number_of_tasks: int = 0
+ XComArg.apply_upstream_relationship(self, self.expand_input.value)
+
+ @property
+ def task_type(self) -> str:
+ """@property: type of the task."""
+ return self._operator.__class__.__name__
+
+ @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
+
+ self._number_of_tasks += 1
+ 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()
+ deferred_tasks: deque[IndexedTaskInstance] = deque()
+ 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:
+ self.log.debug("result: %s", result)
+ continue
+
+ if isinstance(raised, TaskDeferred):
+ operator = DecoratedDeferredAsyncOperator(
+ operator=task.task, task_deferred=raised
+ )
+ deferred_tasks.append(
+ self._create_mapped_task(
+ run_id=task.run_id,
+ index=task.index,
+ map_index=task.map_index, # type:
ignore[arg-type]
+ try_number=task.try_number,
+ operator=operator,
+ )
+ )
+ continue
+
+ if isinstance(raised, asyncio.TimeoutError):
+ self.log.warning("A timeout occurred for task_id
%s", task.task_id)
+ if task.next_try_number > (self.retries or 0):
Review Comment:
`self.retries` is hardcoded to 0 on `IterableOperator` (line 131), so
`task.next_try_number > (self.retries or 0)` is always true and a timed-out
sub-task goes straight to `exceptions` instead of being rescheduled -- the
per-index retry budget is never honored. Compare against the sub-task's
`max_tries` instead (the way the executor does at its equivalent check).
Separately, the sync timeout path won't fire at all: `_execute_task` uses
`signal`-based timeout, which disables itself off the main thread, so sync
sub-tasks running in the thread pool aren't bounded, and
`AsyncAwareExecutor.map(..., timeout=...)` surfaces the timeout from the
iterator rather than as a per-task `(task, result, raised)` tuple, so this
branch never sees it for sync tasks.
##########
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 asyncio
+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,
DecoratedDeferredAsyncOperator, 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,
+ AirflowTaskTimeout,
+ 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 Dynamic Task Iteration (DTI) 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``.
+ """
+
+ _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)
+ self._number_of_tasks: int = 0
+ XComArg.apply_upstream_relationship(self, self.expand_input.value)
+
+ @property
+ def task_type(self) -> str:
+ """@property: type of the task."""
+ return self._operator.__class__.__name__
+
+ @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
+
+ self._number_of_tasks += 1
+ 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()
+ deferred_tasks: deque[IndexedTaskInstance] = deque()
+ 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:
+ self.log.debug("result: %s", result)
+ continue
+
+ if isinstance(raised, TaskDeferred):
+ operator = DecoratedDeferredAsyncOperator(
+ operator=task.task, task_deferred=raised
+ )
+ deferred_tasks.append(
+ self._create_mapped_task(
+ run_id=task.run_id,
+ index=task.index,
+ map_index=task.map_index, # type:
ignore[arg-type]
+ try_number=task.try_number,
+ operator=operator,
+ )
+ )
+ continue
+
+ if isinstance(raised, asyncio.TimeoutError):
+ self.log.warning("A timeout occurred for task_id
%s", task.task_id)
+ if task.next_try_number > (self.retries or 0):
+ exceptions.append(AirflowTaskTimeout(raised))
+ else:
+ reschedule_date = max(reschedule_date,
task.next_retry_datetime())
+ failed_tasks.append(task)
+ continue
+
+ 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
+
+ self.log.exception(
+ "An exception occurred for task_id %s with index
%s",
+ task.task_id,
+ task.index,
+ exc_info=raised,
+ )
+ exceptions.append(raised)
+
+ # Deferred tasks are re-fed as a new pass once the current
batch completes,
+ # because the event loop is restarted at the top of the
outer while loop.
+ if deferred_tasks:
+ tasks = list(deferred_tasks)
+ deferred_tasks.clear()
+ continue
+
+ 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)
+ if do_xcom_push:
+ from airflow.sdk.execution_time.lazy_sequence import
XComIterable
+
+ return XComIterable(
+ task_id=self.task_id,
+ dag_id=self.dag_id,
+ run_id=context["run_id"],
+ length=self._number_of_tasks,
+ map_index=context["ti"].map_index,
+ )
+ return None
+
+ # If the retry time is still in the future we defer the operator
so the worker
+ # slot is released. If the retry time has already passed we
immediately re-run
+ # the failed tasks without deferring.
+ if reschedule_date > timezone.utcnow():
+ if ExternalDateTimeTrigger is not None:
+ self.defer(
+ trigger=ExternalDateTimeTrigger(reschedule_date),
+ method_name=self.execute_failed_tasks.__name__,
+ kwargs={
+ "failed_tasks": {
+ failed_task.index: failed_task.try_number for
failed_task in failed_tasks
+ },
+ },
+ )
+ else:
+ self.log.warning(
+ "DateTimeTrigger is not available; failed tasks cannot
be rescheduled at %s and will be retried immediately",
+ reschedule_date,
+ )
+
+ tasks = list(failed_tasks)
+ failed_tasks.clear()
+ exceptions.clear()
+ reschedule_date = timezone.utcnow()
+
+ async def _run_task(
+ self,
+ executor: AsyncAwareExecutor,
+ context: Context,
+ task: IndexedTaskInstance,
+ ) -> tuple[IndexedTaskInstance, Any | None, BaseException | None]:
+ try:
+ if task.is_async:
+ result = await self._run_async_operator(context, task)
+ else:
+ result = await executor.run_sync(self._run_operator, context,
task)
+
+ # Push XCom asynchronously (non-blocking)
+ if result is not None and task.do_xcom_push:
+ await self._xcom_push(task, result)
+
+ return task, result, None
+ except BaseException as e:
+ 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),
Review Comment:
Sub-tasks run here via `TaskExecutor` without `set_current_context()`. The
normal execution path wraps the operator in `with set_current_context(context)`
(task_runner.py:1699), but this path skips it, and `_CURRENT_CONTEXT` is a
process-global list (contextmanager.py:35), not a contextvar/thread-local stack.
So an iterated TaskFlow callable or operator that calls
`airflow.sdk.get_current_context()` either observes the parent
`IterableOperator`'s context (the last value pushed) instead of its own indexed
context, or races with sibling threads. Wrap each indexed execution in
`set_current_context(per_task_context)` and back the current-context stack with
a `ContextVar` so concurrent sub-tasks stay isolated.
##########
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 asyncio
+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,
DecoratedDeferredAsyncOperator, 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,
+ AirflowTaskTimeout,
+ 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 Dynamic Task Iteration (DTI) 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``.
+ """
+
+ _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)
+ self._number_of_tasks: int = 0
+ XComArg.apply_upstream_relationship(self, self.expand_input.value)
+
+ @property
+ def task_type(self) -> str:
+ """@property: type of the task."""
+ return self._operator.__class__.__name__
+
+ @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
+
+ self._number_of_tasks += 1
+ 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()
+ deferred_tasks: deque[IndexedTaskInstance] = deque()
+ 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:
+ self.log.debug("result: %s", result)
+ continue
+
+ if isinstance(raised, TaskDeferred):
Review Comment:
A sub-task that raises `TaskDeferred` gets wrapped in
`DecoratedDeferredAsyncOperator` and re-run on the same in-process executor, so
`run_trigger()` is awaited inside the worker event loop instead of persisting
deferred state and releasing the slot to the triggerer. An iterated deferrable
sensor that defers for an hour keeps the parent worker occupied the whole time,
and if the worker dies there's no per-index TI state to resume from. For v1 it
may be cleaner to reject deferrable operators under `IterableOperator`, or to
persist parent-level deferred state with per-index entries.
Related, over in `_internal/expandinput.py:255`: `iter_values` does
`tuple(v)` on each resolved value before `itertools.product`, so a lazy
`XComIterable` upstream gets drained in full before any work starts, which is a
memory spike plus a burst of XCom fetches ahead of backpressure. `product`
needs replayable inputs, so if laziness matters here it'll need explicit
streaming or a documented size cap.
##########
task-sdk/src/airflow/sdk/definitions/batchedoperator.py:
##########
@@ -0,0 +1,510 @@
+#
+# 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 inspect
+from abc import ABCMeta, abstractmethod
+from collections.abc import Callable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+import attrs
+
+from airflow.sdk import TriggerRule, timezone
+from airflow.sdk.bases.decorator import (
+ DecoratedMappedOperator,
+ FParams,
+ FReturn,
+ OperatorSubclass,
+ _TaskDecorator,
+ get_unique_task_id,
+)
+from airflow.sdk.bases.operator import (
+ BaseOperator,
+ coerce_resources,
+ coerce_timedelta,
+ get_merged_defaults,
+ parse_retries,
+)
+from airflow.sdk.definitions._internal.contextmanager import (
+ DagContext,
+ TaskGroupContext,
+)
+from airflow.sdk.definitions._internal.expandinput import (
+ EXPAND_INPUT_EMPTY,
+ DecoratedExpandInput,
+ DictOfListsExpandInput,
+ ExpandInput,
+ ListOfDictsExpandInput,
+ OperatorExpandArgument,
+ OperatorExpandKwargsArgument,
+)
+from airflow.sdk.definitions._internal.types import NOTSET
+from airflow.sdk.definitions.mappedoperator import (
+ MappedOperator,
+ OperatorPartial,
+ ensure_xcomarg_return_value,
+ prevent_duplicates,
+ validate_mapping_kwargs,
+)
+from airflow.sdk.definitions.xcom_arg import XComArg
+
+if TYPE_CHECKING:
+ from airflow.sdk.definitions.iterableoperator import IterableOperator,
MappedIterableOperator
+ from airflow.sdk.definitions.mappedoperator import ValidationSource
+ from airflow.sdk.definitions.param import ParamsDict
+
+T = TypeVar("T", bound=OperatorPartial | _TaskDecorator)
+
+
[email protected](kw_only=True, repr=False)
+class BatchableOperator(Generic[T], metaclass=ABCMeta):
+ """
+ Intermediate abstraction for batched mapping.
+
+ This class decorates an OperatorPartial and stores configuration for
batched mapping.
+ It is used to facilitate batched expansion of operators, allowing tasks to
be mapped over batches
+ of data and then iterate over the batched data.
+
+ :param operator_partial: The partial operator to be batched.
+ :param size: The number of batches to create.
+ """
+
+ operator_partial: T
+ size: int
+
+ @property
+ def operator_class(self) -> type[BaseOperator]:
+ return self.operator_partial.operator_class
+
+ @property
+ def kwargs(self) -> dict[str, Any]:
+ return self.operator_partial.kwargs
+
+ @abstractmethod
+ def iterate(self, **mapped_kwargs: OperatorExpandArgument) -> Any:
+ """
+ Iterate the operator over the provided mapped keyword arguments.
+
+ :param mapped_kwargs: Keyword arguments to expand against.
+ :return: An expanded operator or XComArg, depending on the subclass
implementation.
+ """
+
+ @abstractmethod
+ def iterate_kwargs(self, kwargs: OperatorExpandKwargsArgument, *, strict:
bool = True) -> Any:
+ """
+ Iterate the operator over a list of dictionaries or XComArg.
+
+ :param kwargs: List of dicts or XComArg to expand against.
+ :param strict: Whether to enforce strict argument checking.
+ :return: An expanded operator or XComArg, depending on the subclass
implementation.
+ """
+
+ @abstractmethod
+ def _iterate(
+ self,
+ expand_input: ExpandInput,
+ *,
+ strict: bool,
+ ) -> IterableOperator | MappedIterableOperator:
+ """
+ Create an iterable operator for the given expansion input.
+
+ This method calls the _expand method first to get a MappedOperator
based on expansion input,
+ then wraps it in either an IterableOperator or MappedIterableOperator
depending on the batch size.
+
+ :param expand_input: The input to iterate against.
+ :param strict: Whether to enforce strict argument checking.
+ :return: An IterableOperator or MappedIterableOperator.
+ """
+
+ @abstractmethod
+ def _expand(
+ self,
+ expand_input: ExpandInput,
+ *,
+ strict: bool,
+ register_with_dag: bool = True,
+ ) -> MappedOperator:
+ """
+ Create a mapped operator for the given expansion input.
+
+ :param expand_input: The input to expand against.
+ :param strict: Whether to enforce strict argument checking.
+ :param register_with_dag: Whether to apply upstream relationships.
+ :return: A MappedOperator instance.
+ """
+
+
[email protected](kw_only=True, repr=False)
+class BatchedOperator(BatchableOperator[OperatorPartial]):
+ """
+ Concrete implementation of BatchableOperator for classic (non-decorated)
operators.
+
+ This class wraps an OperatorPartial and provides batched expansion and
iteration logic
+ for classic Airflow operators. It enables mapping tasks over batches of
data, supporting
+ both direct expansion via keyword arguments and expansion via a list of
dictionaries or XComArg.
+
+ :param operator_partial: The OperatorPartial instance to be batched and
expanded.
+ :param size: The number of batches to create for mapping.
+ """
+
+ @property
+ def params(self) -> ParamsDict | dict:
+ return self.operator_partial.params
+
+ @property
+ def _expand_called(self) -> bool:
+ return self.operator_partial._expand_called
+
+ @_expand_called.setter
+ def _expand_called(self, value: bool) -> None:
+ self.operator_partial._expand_called = value
+
+ def iterate(self, **mapped_kwargs: OperatorExpandArgument) ->
IterableOperator | MappedIterableOperator:
+ if not mapped_kwargs:
+ raise TypeError("no arguments to iterate against")
+
+ validate_mapping_kwargs(self.operator_class, "iterate", mapped_kwargs)
+ prevent_duplicates(
+ self.kwargs,
+ mapped_kwargs,
+ fail_reason="unmappable or already specified",
+ )
+ # Since the input is already checked at parse time, we can set strict
+ # to False to skip the checks on execution.
+ expand_input = DictOfListsExpandInput(mapped_kwargs)
+ return self._iterate(expand_input, strict=False)
+
+ def iterate_kwargs(
+ self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = True
+ ) -> IterableOperator | MappedIterableOperator:
+ if isinstance(kwargs, Sequence):
+ for item in kwargs:
+ if not isinstance(item, (XComArg, Mapping)):
+ raise TypeError(f"expected XComArg or list[dict], not
{type(kwargs).__name__}")
+ elif not isinstance(kwargs, XComArg):
+ raise TypeError(f"expected XComArg or list[dict], not
{type(kwargs).__name__}")
+
+ expand_input = ListOfDictsExpandInput(kwargs)
+ return self._iterate(expand_input, strict=strict)
+
+ def _iterate(
+ self,
+ expand_input: ExpandInput,
+ *,
+ strict: bool,
+ ) -> IterableOperator | MappedIterableOperator:
+ from airflow.sdk.definitions.iterableoperator import IterableOperator,
MappedIterableOperator
+
+ operator = self._expand(expand_input, strict=strict,
register_with_dag=False)
+
+ if self.size > 0:
+ return MappedIterableOperator(
+ mapped_operator=operator,
+ expand_input=expand_input,
+ batch_size=self.size,
+ )
+ return IterableOperator(
+ operator=operator,
+ expand_input=expand_input,
+ )
+
+ def _expand(
+ self,
+ expand_input: ExpandInput,
+ *,
+ strict: bool,
+ register_with_dag: bool = True,
+ ) -> MappedOperator:
+ from airflow.providers.standard.operators.empty import EmptyOperator
+ from airflow.providers.standard.utils.skipmixin import SkipMixin
Review Comment:
Importing `SkipMixin` from `airflow.providers.standard.utils.skipmixin`
here, then doing `issubclass(self.operator_class, SkipMixin)` at line 276,
checks against the wrong class. On Airflow 3, `ShortCircuitOperator` (and other
standard skip operators) get `SkipMixin` via
`airflow.providers.common.compat.sdk`, which resolves to
`airflow.sdk.bases.skipmixin.SkipMixin` first (compat `sdk.py:158`,
first-importable path wins). That's a different class object from the
`standard.utils` one.
Since `.expand()` now routes through this `_expand` path
(mappedoperator.py:225), `ShortCircuitOperator.partial(...).expand(...)`
serializes with `can_skip_downstream=False`, and downstream skip propagation
stops working. Use the SDK `SkipMixin`, or better, the operator's own
`_can_skip_downstream` / `inherits_from_skipmixin` metadata, rather than
importing a provider class into task-sdk.
##########
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 asyncio
+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,
DecoratedDeferredAsyncOperator, 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,
+ AirflowTaskTimeout,
+ 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 Dynamic Task Iteration (DTI) 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``.
+ """
+
+ _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,
Review Comment:
This kwargs block copies a long list of attributes from the wrapped operator
but not `returns_dag_result`. The decorated path passes it to the inner
`DecoratedMappedOperator` (batchedoperator.py:508), but the `IterableOperator`
that actually runs and pushes the final `XComIterable` never gets it, and the
parent XCom is pushed with `dag_result=ti.task.returns_dag_result`
(task_runner.py:1022).
So a `@result @task` used with `.iterate()` / `.batch().iterate()` won't
have its final `return_value` marked as a dag result, and the dag-run result
API filtering on `dag_result=True` misses it. Propagate `returns_dag_result`
onto `IterableOperator` (and through `MappedIterableOperator.unmap()`).
--
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]
