dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4251950477
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -945,6 +972,581 @@ def mark_success_url(self) -> str:
return self.log_url
+@dataclass
+class IndexedTaskState:
+ status: TaskInstanceState
+ result: Any | None = None
+ # Outlet asset events the sub-task recorded on a previous successful
attempt. Outlet events
+ # only reach the server on the parent's success payload
(_handle_current_task_success), so an
+ # attempt that fails never registers what its succeeded sub-tasks emitted.
A sub-task skipped
+ # on retry (because it already succeeded) never re-executes, so it never
re-emits into the
+ # fresh OutletEventAccessors created for the new attempt; persisting a
snapshot here lets
+ # IterableOperator._run_task replay it, which is the only way those events
survive, and it
+ # cannot double-emit because the failed attempt sent nothing.
+ outlet_events: list[dict[str, Any]] | None = None
+ # Digest of the input the sub-task ran with. An index only means the same
work while its input
+ # is the same, so a checkpoint is honoured on a later attempt only when
this still matches:
+ # after the upstream was cleared and produced other values, the sub-task
runs again.
+ fingerprint: str | None = None
+ # The attempt that wrote the checkpoint. After a manual clear only the
checkpoints written since
+ # are resumed from (see Checkpoints), which this tells apart from those
left by the run before.
+ try_number: int = 0
+ # The other XComs the sub-task pushed, by their key before the index
suffix. The runner deletes
+ # every XCom before an attempt, so a sub-task skipped on retry has them
pushed again from here,
+ # next to its result, as a mapped task instance that succeeded keeps its
own.
+ xcoms: dict[str, Any] | None = None
+
+ def record_outlet_events(self, accessors: OutletEventAccessors) -> None:
+ """
+ Keep a JSON-safe snapshot of the outlet asset events the indexed task
recorded.
+
+ Persisted with the checkpoint so a later attempt can
:meth:`replay_outlet_events` when
+ the indexed task is skipped because it already succeeded.
+ """
+ events: list[dict[str, Any]] = []
+ for _asset_or_alias, accessor in accessors.items():
+ if isinstance(accessor.key, AssetUniqueKey):
+ events.append(
+ {
+ "kind": "asset",
+ "name": accessor.key.name,
+ "uri": accessor.key.uri,
+ "extra": accessor.extra,
+ "partition_keys": sorted(accessor.partition_keys),
+ }
+ )
+ for alias_event in accessor.asset_alias_events:
+ events.append(
+ {
+ "kind": "asset_alias",
+ "source_alias_name": alias_event.source_alias_name,
+ "dest_asset_key": {
+ "name": alias_event.dest_asset_key.name,
+ "uri": alias_event.dest_asset_key.uri,
+ },
+ "dest_asset_extra": alias_event.dest_asset_extra,
+ "extra": alias_event.extra,
+ }
+ )
+ if events:
+ self.outlet_events = events
+
+ def replay_outlet_events(self, target: OutletEventAccessorsProtocol) ->
None:
+ """
+ Re-emit into ``target`` the events the indexed task recorded on a
previous attempt.
+
+ An indexed task skipped on retry (because it already succeeded) never
re-executes, so it
+ never re-emits into the fresh ``OutletEventAccessors`` created for the
new attempt. The
+ failed attempt sent nothing to the server either (outlet events travel
only on the success
+ payload), so replaying cannot emit an event twice; without it the
events would be lost.
+ They land in the parent's accessors as
:meth:`IndexedTaskRunner.merge_outlet_events_into`
+ lands the live ones: one event per asset per task instance.
+ """
+ for event in self.outlet_events or ():
+ if event["kind"] == "asset":
+ accessor = target[Asset(name=event["name"], uri=event["uri"])]
+ accessor.extra.update(event["extra"])
+ if event["partition_keys"]:
+ accessor.add_partitions(event["partition_keys"])
+ else:
+
target[AssetAlias(name=event["source_alias_name"])].asset_alias_events.append(
+ AssetAliasEvent(
+ source_alias_name=event["source_alias_name"],
+
dest_asset_key=AssetUniqueKey(**event["dest_asset_key"]),
+ dest_asset_extra=event["dest_asset_extra"],
+ extra=event["extra"],
+ )
+ )
+
+ @staticmethod
+ def build_key(index: int) -> str:
+ # The task state store is already scoped to the parent task instance
(dag, run, task and
+ # map index), so the key carries no identity, only a namespace that
keeps the operator's own
+ # entries apart from anything user code stores from inside a sub-task.
+ return f"_iterable_{index}"
+
+ def serialize(self) -> dict[str, Any]:
+ # The checkpoint travels to the supervisor as a JsonValue, which
rejects anything that is not
+ # plain JSON (tuples, datetimes, models, ...). Serde turns those into
JSON-compatible
+ # structures and restores them on read, exactly as XCom does with the
same result.
+ data: dict[str, Any] = {"status": self.status.value}
+ if self.result is not None:
+ data["result"] = serde_serialize(self.result)
+ if self.outlet_events:
+ data["outlet_events"] = self.outlet_events
+ if self.fingerprint:
+ data["fingerprint"] = self.fingerprint
+ if self.try_number:
+ data["try_number"] = self.try_number
+ if self.xcoms:
+ data["xcoms"] = {key: serde_serialize(value) for key, value in
self.xcoms.items()}
+ return data
+
+ @classmethod
+ def deserialize(cls, raw: Any) -> IndexedTaskState | None:
+ if not isinstance(raw, Mapping):
+ return None
+ return cls(
+ status=TaskInstanceState(raw["status"]),
+ result=serde_deserialize(raw.get("result")),
+ outlet_events=raw.get("outlet_events"),
+ fingerprint=raw.get("fingerprint"),
+ try_number=raw.get("try_number", 0),
+ xcoms={key: serde_deserialize(value) for key, value in
raw["xcoms"].items()}
+ if raw.get("xcoms")
+ else None,
+ )
+
+
+class IndexedTaskInstance(RuntimeTaskInstance):
+ """
+ Indexed task instance to run a mapped operator.
+
+ It shares the parent task instance's identity, so what an iteration pushes
or stores lands in
+ the parent's scope, suffixed with the index so that iterations never
overwrite each other:
+ XComs through :meth:`xcom_push`, task state through
:attr:`task_state_store`. Reading back
+ follows the same rule: :meth:`xcom_pull` of the iteration's own XComs adds
the index, a pull
+ from another task does not, as the store's accessor adds it to every key
of its own. The
+ operator's own checkpoints go to the parent's store unsuffixed, under their
+ ``_iterable_<index>`` keys.
+ """
+
+ index: int
+ parent_task_state_store: TaskStateStoreAccessor
+ input_fingerprint: str | None = None
+ # The XComs the sub-task pushed other than its return value, by their key
before the index
+ # suffix: kept in memory while it runs and written once, with its SUCCESS
checkpoint.
+ pushed_xcoms: dict[str, Any] = Field(default_factory=dict)
+
+ @classmethod
+ def create_indexed_task(
+ cls, *, context: Context, index: int, operator: BaseOperator,
input_fingerprint: str | None = None
+ ) -> IndexedTaskInstance:
+ """
+ Create the runtime instance for one index of an iterated task from the
parent's context.
+
+ The instance shares the parent task instance's identity (id, run, map
index, try number and
+ retry budget), so XComs and task state land in the parent's scope, and
carries the unmapped
+ operator for that index. The budget is the parent's ``max_tries``
rather than the operator's
+ ``retries``: a manual clear raises it, and it is the parent Airflow
retries. The parent's
+ state store comes from the context, the same accessor the operator's
+ checkpoints use. ``model_construct`` skips Pydantic validation on
purpose: one instance is built per
+ indexed task, and the parent was validated already, so only the index
needs checking here.
+ """
+ if index < 0:
+ raise ValueError(f"IndexedTaskInstance requires index >= 0, got
{index}")
+ parent = context["ti"]
+ return cls.model_construct(
+ id=parent.id,
+ parent_task_state_store=context["task_state_store"],
+ # The parent's: an iteration's XComs and state belong to the task
instance that runs it.
+ task_id=parent.task_id,
+ dag_id=operator.dag_id,
+ run_id=parent.run_id,
+ map_index=parent.map_index,
+ region_id=getattr(parent, "region_id", None),
+ region_index=getattr(parent, "region_index", None),
+ index=index,
+ input_fingerprint=input_fingerprint,
+ max_tries=parent.max_tries,
+ start_date=parent.start_date,
+ state=TaskInstanceState.SCHEDULED.value,
+ is_mapped=True,
+ task=operator,
+ try_number=parent.try_number,
+ # The dag run, logical date and the rest the server sent for the
task instance that
+ # runs this iteration: its template context, get_previous_ti() and
the like read them.
+ _ti_context_from_server=getattr(parent, "_ti_context_from_server",
None),
+ )
+
+ def xcom_push(
+ self,
+ key: str,
+ value: Any,
+ ):
+ super().xcom_push(key=f"{key}_{self.index}", value=value)
+ self._record_push(key, value)
+
+ async def axcom_push(
+ self,
+ key: str,
+ value: Any,
+ ):
+ await super().axcom_push(key=f"{key}_{self.index}", value=value)
+ self._record_push(key, value)
+
+ def _record_push(self, key: str, value: Any) -> None:
+ # The return value has its own slot on the checkpoint and is published
by the operator.
+ if key != BaseXCom.XCOM_RETURN_KEY:
+ self.pushed_xcoms[key] = value
+
+ def _own_key(self, task_ids: str | Iterable[str] | None, dag_id: str |
None, key: str) -> str:
+ """
+ Suffix ``key`` with the index for a pull of this iteration's own XComs.
+
+ What :meth:`xcom_push` wrote under ``<key>_<index>`` is read back
under the same name when
+ the pull names no task or this task instance's own (in its own DAG); a
pull from another
+ task, or from several, keeps its key, since those XComs carry no index.
+ """
+ own_task = task_ids is None or task_ids == self.task_id
+ own_dag = dag_id is None or dag_id == self.dag_id
+ return f"{key}_{self.index}" if own_task and own_dag else key
+
+ def xcom_pull(
+ self,
+ task_ids: str | Iterable[str] | None = None,
+ dag_id: str | None = None,
+ key: str = BaseXCom.XCOM_RETURN_KEY,
+ include_prior_dates: bool = False,
+ *,
+ map_indexes: int | Iterable[int] | None | ArgNotSet = NOTSET,
+ default: Any = None,
+ run_id: str | None = None,
+ ) -> Any:
+ return super().xcom_pull(
+ task_ids=task_ids,
+ dag_id=dag_id,
+ key=self._own_key(task_ids, dag_id, key),
+ include_prior_dates=include_prior_dates,
+ map_indexes=map_indexes,
+ default=default,
+ run_id=run_id,
+ )
+
+ async def axcom_pull(
+ self,
+ task_ids: str | Iterable[str] | None = None,
+ dag_id: str | None = None,
+ key: str = BaseXCom.XCOM_RETURN_KEY,
+ include_prior_dates: bool = False,
+ *,
+ map_indexes: int | Iterable[int] | None | ArgNotSet = NOTSET,
+ default: Any = None,
+ run_id: str | None = None,
+ ) -> Any:
+ return await super().axcom_pull(
+ task_ids=task_ids,
+ dag_id=dag_id,
+ key=self._own_key(task_ids, dag_id, key),
+ include_prior_dates=include_prior_dates,
+ map_indexes=map_indexes,
+ default=default,
+ run_id=run_id,
+ )
+
+ @cached_property
+ def task_state_store(self) -> TaskStateStoreAccessor: # type:
ignore[override]
+ """The parent's store seen from this iteration: keys are suffixed with
the index."""
+ return IndexedTaskStateStoreAccessor(self.parent_task_state_store,
self.index)
+
+ def context_for(self, context: Context, *, outlet_events:
OutletEventAccessors | None = None) -> Context:
+ """
+ Return the parent's context as this indexed task sees it.
+
+ A clone of ``context`` with this task instance under ``ti`` and
``task_instance``, its
+ unmapped operator under ``task`` and its indexed view of the store
under
+ ``task_state_store``: the keys ``context_update_for_unmapped`` sets
for a mapped task
+ instance, and the store next to them, so that a template and
``execute`` read the same
+ store. The one place these keys are listed:
:meth:`IndexedTaskRunner.indexed_context`
+ builds the context the task runs in from it, and
``IterableOperator._create_task`` the
+ one its templates are rendered against. ``outlet_events`` is swapped
when given, for the
+ run; before it, nothing is emitted and the parent's accessor stays.
+ """
+ indexed: Context = {
+ **clone_context(context),
+ "ti": self,
+ "task_instance": self,
+ "task": self.task,
+ "task_state_store": self.task_state_store,
+ }
+ if outlet_events is not None:
+ indexed["outlet_events"] = outlet_events
+ return indexed
+
+ async def aget_state(self) -> IndexedTaskState | None:
+ return IndexedTaskState.deserialize(await
self.parent_task_state_store.aget(self.state_key))
+
+ async def aset_state(self, state: IndexedTaskState) -> None:
+ await self.parent_task_state_store.aset(self.state_key,
state.serialize())
+
+ @property
+ def is_async(self) -> bool:
+ return self.task.is_async
+
+ @property
+ def is_eligible_to_retry(self) -> bool:
+ """
+ Whether Airflow runs the parent again after this attempt fails.
+
+ The same rule the API server applies to the parent
(``_is_eligible_to_retry``), so the
+ callbacks of an iteration agree with what happens to the task instance.
+ """
+ return self.max_tries != 0 and self.try_number <= self.max_tries
+
+ @property
+ def state_key(self) -> str:
+ return IndexedTaskState.build_key(self.index)
+
+ @property
+ def do_xcom_push(self) -> bool:
+ return self.task.do_xcom_push
+
+
+class SubOperatorRegister(Protocol):
+ """Where an indexed task's operator is noted while its code runs, so that
a kill reaches it."""
+
+ def register(self, operator: BaseOperator) -> None: ...
+
+ def unregister(self, operator: BaseOperator) -> None: ...
+
+
+class IndexedTaskRunner(LoggingMixin):
+ """
+ Run one indexed task of an iterated task: its operator, against its own
view of the context.
+
+ Named apart from Airflow's executors, which schedule task instances, and
from the task runner
+ process, which runs the parent: this runs one index inside that process,
sync or async.
+ """
+
+ def __init__(
+ self,
+ task_instance: IndexedTaskInstance,
+ register: SubOperatorRegister | None = None,
+ outlet_events: OutletEventAccessors | None = None,
+ ):
+ """
+ Run an operator or trigger for one sub-task instance.
+
+ :param outlet_events: The accessor the sub-task's asset events are
collected in, its own
+ so they can be checkpointed and merged apart from its siblings'.
Created here when not
+ given; the caller reads it back through :attr:`outlet_events`
after the run.
+ :param register: Optional register the operator is entered in while
its code runs (see
+ :meth:`in_flight`), so that IterableOperator.on_kill() can reach
whichever sub-tasks
+ are in flight; the iterated task's ``IterationState``.
+ """
+ super().__init__()
+ self.task_instance = task_instance
+ self.outlet_events = outlet_events if outlet_events is not None else
OutletEventAccessors()
+ self._result: Any | None = None
+ self._start_time: float | None = None
+ self._context: Context | None = None
+ self._register = register
+ #: The exception this indexed task failed with, noted by __exit__ and
reported by
+ #: :meth:`report_failure` once the whole task's fate is known.
+ self.failure: BaseException | None = None
+ #: Whether the ``AirflowTaskTimeout`` this indexed task ended with was
raised by its operator
+ #: off the main thread, where the parent's limit never strikes; see
:meth:`in_flight`.
+ self.timed_out_on_its_own = False
+ self._cancelled = False
+
+ def merge_outlet_events_into(self, target: OutletEventAccessorsProtocol)
-> None:
+ """
+ Fold the outlet asset events this indexed task recorded into the
parent's ``target``.
+
+ Called right after the indexed task succeeds, on the
IterableOperator's shared
+ ``context["outlet_events"]``. A task instance sends one event per
asset, so indexed tasks
+ that emit to the same asset end up in that one event: ``extra`` keeps
what the last one to
+ finish wrote, while partition keys and alias events accumulate.
``.expand()`` sends one
+ event per mapped task instance instead; the docs page says so in its
comparison table.
+ """
+ for asset_or_alias, accessor in self.outlet_events.items():
+ target_accessor = target[asset_or_alias]
+ target_accessor.extra.update(accessor.extra)
+
target_accessor.asset_alias_events.extend(accessor.asset_alias_events)
+ target_accessor.partition_keys.update(accessor.partition_keys)
+
+ def cancel(self) -> None:
+ """
+ Note that the coroutine waiting for this sync indexed task was
cancelled.
+
+ Its thread may go on, but the task gets no checkpoint from here on,
and no report either:
+ the next attempt runs it again and reports it then. Best effort: a
thread already past the
+ check in its exit still notes a failure.
+ """
+ self._cancelled = True
+
+ @property
+ def cancelled(self) -> bool:
+ """Whether :meth:`cancel` was called: the indexed task's exit then
notes nothing."""
+ return self._cancelled
+
+ @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 operator(self) -> BaseOperator:
+ return self.task_instance.task
+
+ @property
+ def is_async(self) -> bool:
+ return self.task_instance.is_async
+
+ @contextmanager
+ def indexed_context(self, context: Context) -> Generator[Context, None,
None]:
+ """
+ Enter the parent's context as this indexed task sees it.
+
+ Yields the task instance's :meth:`IndexedTaskInstance.context_for`
view of the parent's
+ context with this runner's own outlet events, remembered on the runner
and made the
+ current context for the duration of the block, so user code reads the
indexed task's
+ unmapped operator under ``context["task"]``, not the IterableOperator.
The parent's
+ context is left untouched: ``context_update_for_unmapped`` sets
``ti.task`` on whatever
+ ``ti`` it finds, which must be this task's, not the parent's.
+ """
+ indexed_context = self.task_instance.context_for(context,
outlet_events=self.outlet_events)
+ self._context = indexed_context
+ with set_indexed_context(indexed_context):
+ yield indexed_context
+
+ @contextmanager
+ def in_flight(self) -> Generator[None, None, None]:
+ """
+ Register the operator as running for exactly as long as its code runs.
+
+ Entered by :meth:`run` and :meth:`arun`, so in the worker thread or
coroutine that executes
+ the operator: a sync operator stays registered while its thread is
still inside ``execute``,
+ even after the coroutine waiting for it was cancelled, and
``IterableOperator.on_kill`` can
+ reach it. The operator the parent's execution timeout strikes stays
registered as well:
+ the timeout lands on the main thread, where the loop runs async
operators, and its
+ ``on_kill`` must not run there, where a synchronous SDK call raises, so
+ ``IterableOperator._run_tasks`` kills it off the loop thread with the
others once the
+ timeout has unwound. ``TimeoutPosix`` raises on the main thread only,
so an
+ ``AirflowTaskTimeout`` on a worker thread was raised by the operator
itself (a hook that
+ gave up waiting): that is this indexed task's own failure, noted as
+ :attr:`timed_out_on_its_own`, and the operator is unregistered like
any other.
+ """
+ if self._register is not None:
+ self._register.register(self.operator)
+ struck_by_the_parent = False
+ try:
+ yield
+ except AirflowTaskTimeout:
+ self.timed_out_on_its_own = threading.current_thread() is not
threading.main_thread()
Review Comment:
Agreed, documented as a limitation for now (060fe8c84b). The iterable tasks
page has a warning in its async section, next to the one on synchronous calls:
an async task raising `AirflowTaskTimeout` itself is taken for the task's own
`execution_timeout`, so the iteration stops, the running iterations are killed
and the task is retried as on its own timeout; a sync task is not concerned,
and for a time limit of the operator's own it should raise another exception
such as `TimeoutError`. The `in_flight` docstring names it too. Telling the two
apart with a flag set by the timeout handler would be the real fix; I would
rather keep that as a follow-up than change the timeout path in this PR.
Drafted-by: Claude Opus 5.5; reviewed by @dabla before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
