kaxil commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4338223566
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -805,6 +821,193 @@ def delete(
cls.purge(xcom_result)
+class XComIterable(Sequence):
+ """
+ An iterable that lazily fetches XCom values one by one instead of loading
all at once.
+
+ This is a read-only :class:`collections.abc.Sequence` over the
``return_value_<index>`` XComs an
+ iterated task pushed, one per index: the values are written by the
producing task's runner as
+ each indexed task finishes (see ``IndexedTaskLifecycle._push_result``),
and the iterable only ever reads
+ them. Nothing on this class mutates the underlying XComs.
+
+ Indexing follows the usual sequence rules, negative indices included:
``result[-1]`` is the last
+ value. Iterations that were skipped pushed nothing and are left out, as
the XComs of skipped
+ mapped task instances are: ``length`` counts every input item, ``skipped``
lists the indices
+ that were skipped, and positions in the sequence run over the others only.
An iteration that
+ returned ``None`` pushed nothing either but keeps its position: reading it
gives ``None``,
+ where ``.expand()`` would not count such a mapped task instance at all.
+
+ Every element is a remote fetch, so random access costs one XCom read per
element, and so do
+ iterating and slicing: N elements are N requests. Reading several in one
request needs an
+ Execution API endpoint that takes a list of keys, which does not exist yet.
+
+ An iterable read from the previous iteration (``loop.previous`` inside a
loop) is built with
+ ``previous_iteration`` set (see :meth:`deserialize`), so its elements are
read from that
+ iteration too and not from the current one. "Previous" is relative to
whoever reads, so
+ such an iterable cannot be pushed on: :meth:`serialize` refuses it (see
there).
+ """
+
+ def __init__(
+ self,
+ task_id: str,
+ dag_id: str,
+ run_id: str,
+ map_index: int | None = None,
+ length: int | None = None,
+ skipped: Sequence[int] = (),
+ previous_iteration: bool = False,
+ ):
+ self.task_id = task_id
+ self.dag_id = dag_id
+ self.run_id = run_id
+ self.map_index = map_index
+ self.length = length or 0
+ self.skipped: list[int] = sorted(skipped)
+ self.previous_iteration = previous_iteration
+
+ def _index_of(self, position: int) -> int:
+ """Map ``position`` in the sequence to its input index, stepping over
the skipped indices."""
+ # range() maps a negative position from the end and raises IndexError
past either end.
+ index = range(len(self))[position]
+ for skipped_index in self.skipped:
+ if skipped_index > index:
+ break
+ index += 1
+ return index
+
+ def __iter__(self) -> Iterator[Any]:
+ return _XComIterator(self)
+
+ def __len__(self) -> int:
+ return self.length - len(self.skipped)
+
+ async def alen(self) -> int:
+ """Async twin of ``len(self)``, for readers on the event loop that
take a length before each read."""
+ return len(self)
+
+ def __getitem__(self, key: int | slice) -> Any | Sequence[Any]:
+ """Allow direct indexing so this works like a sequence."""
+ from airflow.sdk.execution_time.xcom import XCom
+
+ if isinstance(key, slice):
+ # TODO: This issues one XCom.get_one call per element — N
round-trips for a full slice.
+ # XComIterable stores results under distinct keys (return_value_0,
return_value_1, …)
+ # with the same map_index, so the existing GetXComSequenceSlice
endpoint (which ranges
+ # over map_index for a single key) cannot be reused. A new POST
endpoint that accepts
+ # a list of keys and returns values in a single query is needed;
once that lands, replace
+ # this loop with a single batched fetch.
+ start, stop, step = key.indices(len(self))
+ return [self[i] for i in range(start, stop, step)]
+
+ return XCom.get_one(**self._item_query(key))
+
+ async def aget(self, index: int) -> Any:
+ """
+ Async counterpart of ``self[index]``: fetch one value through
``XCom.aget_one``.
+
+ Use it, or ``async for``, from code running on an event loop that has
other SDK calls in
+ flight (an iterated task consuming this iterable as its input): a
synchronous read there
+ would block the loop, and every call in flight on it, for its round
trip to the supervisor.
+ """
+ from airflow.sdk.execution_time.xcom import XCom
+
+ return await XCom.aget_one(**self._item_query(index))
+
+ def __aiter__(self) -> AsyncIterator[Any]:
+ return _AsyncXComIterator(self)
+
+ def _item_query(self, position: int) -> dict[str, Any]:
Review Comment:
An `XComIterable` doesn't record which loop pass it came from, and its item
reads go out as a plain `GetXCom` with no scope. A plain task outside a loop
that reads an iterated task inside it gets one of these per pass, and
`_build_producer_request` in `models/task_coordinates.py` rejects a scalar read
of a loop producer from outside its loop, so every item read comes back as a
400. The `previous_iteration` contextvar has the reverse gap: an `XComIterable`
for a task outside the loop, loaded through `loop.previous`, gets the flag too
and hits "Previous-iteration lookup requires a shared loop scope". I traced
this in the code and didn't run it. Could the pointer carry the producer's
region and pass, and could one test read an iterated loop task from a plain
task outside the loop end to end?
##########
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),
Review Comment:
The indexed task instance doesn't get the parent's `bundle_instance`, and of
the `TaskInstance` fields `dag_version_id`, `hostname`, `queue` and
`context_carrier` are missing too, while this is the `context["ti"]` the
operator sees. The virtualenv-based Python operators read `ti.bundle_instance`
to put the Dag bundle on the subprocess `PYTHONPATH` (standard
`operators/python.py:568`), so under `.iterate()` a callable that imports from
the bundle can't find it, and the OpenLineage bundle facets come out as None.
Copying those from `parent` here would cover it.
--
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]
