dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4142570349
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -897,6 +923,337 @@ 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 items, 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
+
+ @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
+ 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),
+ )
+
+
+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`. 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
+
+ @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
+ item, 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(
Review Comment:
Reproduced: the indexed TI had `logical_date=None` and no `dag_run`/`ds` in
its template context. 1000b9175a passes
`_ti_context_from_server=parent._ti_context_from_server` to `model_construct`.
`next_method` from it stays `None`, since the IterableOperator never defers.
The test builds the parent with the runtime-TI fixture and checks the logical
date and the `dag_run`/`ds` of the indexed TI's template context.
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]
