dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4263966461
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -950,6 +985,623 @@ 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
+ # IterableOperator._run_indexed_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 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 refuses to restore, a class off the
allow-list
+ (``ImportError``) or one it has no deserializer for (``TypeError``),
raises
+ :class:`CheckpointDeserializationError`: refused once, it is refused
on every retry.
+ """
+ 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 (ImportError, TypeError) as e:
Review Comment:
Fixed in 1d06e0522d: `IndexedTaskState.deserialize` now takes any exception
from the two `serde_deserialize` calls as a refusal, so the indexed task runs
again instead of failing every retry. The test is parametrized over a class off
the allow-list, a pydantic model that gained a required field, and a class
whose constructor raises `TypeError`; without the fix only the pydantic case
fails, the other two pin what was already caught.
Drafted-by: Claude Fable 5.1; 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]
