kaxil commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4338223595
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -955,6 +990,633 @@ def mark_success_url(self) -> str:
return self.log_url
+class CheckpointDeserializationError(Exception):
+ """An indexed task's checkpoint holds a value serde refuses to restore,
such as a class off the allow-list."""
+
+
+@dataclass
+class IndexedTaskState:
+ status: TaskInstanceState
+ result: Any | None = None
+ # Outlet asset events the indexed 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 indexed tasks
emitted. 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; persisting a
snapshot here lets
+ # IndexedTaskLifecycle 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 indexed 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 indexed
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 indexed task pushed, by their key before the index
suffix. The runner deletes
+ # every XCom before an attempt, so an indexed 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 an indexed
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:
+ """
+ Restore a checkpoint, or None when there is none.
+
+ A result or XCom value serde fails to restore raises
:class:`CheckpointDeserializationError`:
+ a class off the allow-list (``ImportError``), one it has no
deserializer for or whose
+ constructor refuses the data (``TypeError``), a pydantic model that
changed since the
+ checkpoint was written (``ValidationError``, a ``ValueError``), or
whatever else a class
+ raises. Refused once, it is refused on every retry. ``raw`` is already
read, so nothing
+ transient can fail here, and every failure is taken as a refusal.
+ """
+ if not isinstance(raw, Mapping):
+ return None
+ try:
+ result = serde_deserialize(raw.get("result"))
+ xcoms = (
+ {key: serde_deserialize(value) for key, value in
raw["xcoms"].items()}
+ if raw.get("xcoms")
+ else None
+ )
+ except Exception as e:
+ raise CheckpointDeserializationError(str(e)) from e
+ return cls(
+ status=TaskInstanceState(raw["status"]),
+ result=result,
+ outlet_events=raw.get("outlet_events"),
+ fingerprint=raw.get("fingerprint"),
+ try_number=raw.get("try_number", 0),
+ xcoms=xcoms,
+ )
+
+
+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 indexed 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)
+ # The params processed for this index's operator when they differ from the
iterated task's:
+ # the index iterates ``params``, or the operator set its own. None when
the parent's apply.
+ own_params: dict[str, Any] | None = None
+
+ @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:
Review Comment:
If an operator pushes its result with `ti.xcom_push("return_value", v)` and
returns None, nothing of it reaches the checkpoint: this skips the key for
`pushed_xcoms`, and `result` is None. That's fine until a sibling fails. On the
retry the runner has already deleted the XComs, the replay pushes nothing for
this index, and downstream reads None. Could `_record_push` keep it when the
operator returns None?
##########
providers/openlineage/tests/unit/openlineage/extractors/test_manager.py:
##########
@@ -402,6 +402,34 @@ def execute(self, context: Context) -> Any:
assert metadata.outputs == outlets
[email protected](not AIRFLOW_V_3_4_PLUS, reason=".iterate() is available
from Airflow 3.4")
Review Comment:
The PR is on the 3.5.0 milestone now and `v3-4-test` has branched, so once
3.4.x is in the provider compatibility matrix this test runs on an Airflow
without `.iterate()`. Gating on `hasattr(OperatorPartial, "iterate")` would
hold whichever release it lands in. The `versionadded:: 3.4.0` on the docs page
needs the same bump.
##########
task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py:
##########
@@ -88,6 +113,36 @@ def __hash__(self):
def __iter__(self) -> Iterator[T]:
return LazyXComIterator(seq=self)
+ def __aiter__(self) -> AsyncIterator[T]:
+ return AsyncLazyXComIterator(seq=self)
+
+ async def aget(self, index: int) -> T:
+ """Async counterpart of ``self[index]``; the same
``GetXComSequenceItem`` request via ``asend``."""
+ from airflow.sdk.execution_time.comms import (
+ ErrorResponse,
+ GetXComSequenceItem,
+ XComSequenceIndexResult,
+ )
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+ from airflow.sdk.execution_time.xcom import XCom
+
+ source = (xcom_arg := self._xcom_arg).operator
+ msg = await SUPERVISOR_COMMS.asend(
+ GetXComSequenceItem(
+ previous_iteration=self._previous_iteration,
+ key=xcom_arg.key,
+ dag_id=source.dag_id,
+ task_id=source.task_id,
+ run_id=self._ti.run_id,
+ offset=index,
+ ),
+ )
+ if isinstance(msg, ErrorResponse):
+ raise IndexError(index)
+ if not isinstance(msg, XComSequenceIndexResult):
+ raise TypeError(f"Got unexpected response to GetXComSequenceItem:
{msg!r}")
+ return XCom._load(_XComWrapper(msg.root),
previous_iteration=self._previous_iteration)
Review Comment:
`XCom._load` runs the backend's `deserialize_value` here on the event loop
thread, so with an object-storage XCom backend every indexed task in flight
waits for each item's download. `aget_one` already moves the same call off the
loop with `asyncio.to_thread` (xcom.py:600), and the same change fits here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
