Expose task instance loop pass and expansion size to listeners

airflow-core/docs/authoring-and-scheduling/loops.rst

@@ -152,9 +152,21 @@ definition.
      - The previous iteration's result, or ``None`` in iteration 1.
    * - ``loop.result``
      - The terminal body task's result (the XCom pushed under the key ``return_value``) for the current iteration, used by the gate.
+   * - ``loop.region_id``
+     - The id of the loop execution (trajectory) that holds this iteration. It is the same for every task
+       instance of the iteration, mapped ones included.
+   * - ``loop.trajectory``
+     - The number of that trajectory: 1 for the loop's first execution, and one more each time a rerun
+       replaces later iterations, see :ref:`loops-trajectories`.
+   * - ``loop.forked_from_region_id``
+     - The ``loop.region_id`` of the previous trajectory, or ``None`` in trajectory 1.
+   * - ``loop.resumes_from_index``
+     - The first iteration this trajectory holds, or ``None`` in trajectory 1, which holds them from 1.
 
 Use ``loop.index`` for the loop iteration and ``ti.map_index`` for a mapped
 task instance's position. Each task instance can have several tries, each with its own ``ti.try_number``.
+The task instance also exposes the loop context as ``ti.loop``, and for a mapped task
+``ti.expanded_ti_count`` is the number of instances in its expansion, counted within the iteration.
 
 Pass data between iterations
 ============================
@@ -347,6 +359,12 @@ execution. Downstream tasks then run as they would after any successful gate: th
 gate that was marked successful by hand from one that stopped because ``until`` returned ``True``.
 Any later iterations retained after a selective clear remain unchanged.
 
+The gate's task instance exposes its decision as ``ti.loop_decision``, a ``LoopDecision``, once the gate
+has evaluated it. With ``until``, ``stop`` is ``True`` when ``until`` ended the loop, even at its last
+iteration. When ``stop`` is ``False`` at the last iteration (``loop.index == loop.max_iterations - 1``),
+the loop reached ``max_iterations`` without meeting ``until`` and the gate fails. A gate that fails before
+deciding, for example because ``until`` raised, or that was marked successful by hand, has no decision.
+
 Iterations and execution history
 ================================
 
@@ -414,3 +432,17 @@ which leave them in place. To keep later iterations, clear from the UI or the
 REST endpoint.
 
 Tasks you leave uncleared are not recomputed using the results of the rerun.
+
+.. _loops-trajectories:
+
+Reruns and trajectories
+-----------------------
+
+Clearing later iterations starts a new execution of the loop that replaces the previous ones from the
+iteration after the rerun gate. Each execution is a trajectory, identified by ``loop.region_id`` and
+numbered by ``loop.trajectory``: 1 for the first one, growing by one with each rerun. The rerun gate, and
+whatever you clear in its iteration, run again in the trajectory that iteration belongs to, with their
+next try. The replacement iterations run in the new trajectory, from ``loop.resumes_from_index``, and their
+tries count from 1 again for each task. In the example above, iteration 2 stays in trajectory 1, while
+replacement iterations 3 and 4 run in trajectory 2, with ``loop.resumes_from_index`` 3. Task log paths show
+the trajectory after the iteration, as in ``pass=3.2``.

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py

@@ -32,8 +32,9 @@
     field_validator,
     model_validator,
 )
+from pydantic.json_schema import SkipJsonSchema
 
-from airflow.api_fastapi.common.region import OmitsMissingRegion, RegionId, RegionIndex
+from airflow.api_fastapi.common.region import OmitsMissingRegion, RegionId, RegionIndex, _drop_null_default
 from airflow.api_fastapi.common.types import UtcDateTime
 from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
 from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile
@@ -457,6 +458,38 @@ class LoopContext(BaseModel):
     terminal_task_id: str
     terminal_is_mapped: bool
 
+    region_id: uuid.UUID
+    """
+    The loop region holding this iteration, which identifies the trajectory it runs in.
+
+    A mapped task's own ``region_id`` is its expansion's, so this is how its instances are tied to the
+    unmapped tasks of the same iteration.
+    """
+
+    trajectory: int = Field(ge=1)
+    """
+    The number of ``region_id`` in the loop's chain of regions: 1 for the original, n + 1 for its n-th fork.
+
+    Clearing a loop gate with the later iterations forks the newest loop region, and the iterations from
+    ``resumes_from_index`` on run again in the fork.
+    """
+
+    forked_from_region_id: Annotated[
+        SkipJsonSchema[None] | uuid.UUID, Field(json_schema_extra=_drop_null_default)
+    ] = None
+    """
+    The region ``region_id`` was forked from, the previous one in the chain. Absent for the original region.
+
+    The fork replaces the iterations from ``resumes_from_index`` on, which can be held by earlier regions of
+    the chain too.
+    """
+
+    resumes_from_index: Annotated[
+        SkipJsonSchema[None] | int,
+        Field(json_schema_extra=_drop_null_default),
+    ] = None
+    """The first iteration ``region_id`` holds. Absent for the original region, which holds them from 1."""
+
 
 class TIRunContext(BaseModel):
     """Response schema for TaskInstance run context."""
@@ -466,6 +499,16 @@ class TIRunContext(BaseModel):
 
     loop: LoopContext | None = None
 
+    expanded_ti_count: Annotated[
+        SkipJsonSchema[None] | int,
+        Field(json_schema_extra=_drop_null_default),
+    ] = None
+    """
+    Number of live, non-removed instances in the mapped task's expansion this instance belongs to.
+
+    ``None`` for a task that isn't mapped. Inside a loop, it counts the instances of this iteration only.
+    """
+
     task_reschedule_count: int = 0
     """How many times the task has been rescheduled."""
 

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -59,7 +59,6 @@
 from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
     DagRunNoteUpdatePayload,
     InactiveAssetsResponse,
-    LoopContext,
     PreviousTIResponse,
     PrevSuccessfulDagRunResponse,
     TaskBreadcrumbsResponse,
@@ -371,19 +370,13 @@ def ti_run(
             should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries),
             multi_team=conf.getboolean("core", "multi_team"),
         )
-        resolver = TaskCoordinateResolver(dag_bag, session)
-        if loop_context := resolver.loop_context(ti):
-            group, index = loop_context
-            terminal = resolver.get_task(
-                ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id
-            )
-            context.loop = LoopContext(
-                node_id=group.node_id,
-                index=index,
-                max_iterations=group.max_iterations,
-                terminal_task_id=group.terminal_task_id,
-                terminal_is_mapped=terminal.get_needs_expansion(),
-            )
+        if ti.region_id != SENTINEL_REGION_ID:
+            resolver = TaskCoordinateResolver(dag_bag, session)
+            task = resolver.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+            if (expanded_ti_count := resolver.get_expanded_ti_count(ti, task=task)) is not None:
+                context.expanded_ti_count = expanded_ti_count
+            if (loop := resolver.get_loop_context(ti, task=task)) is not None:
+                context.loop = loop
         if first_task_reschedule_start_date is not None:
             context.first_task_reschedule_start_date = first_task_reschedule_start_date
 

airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py

@@ -56,6 +56,7 @@
     AddArgBindingsToTIRunContext,
     AddCallbackRunEndpoint,
     AddDagRunNoteUpdateEndpoint,
+    AddExpandedTICountToTIRunContext,
     AddFirstTaskRescheduleStartDateField,
     AddLoopContext,
     AddLoopDecisionField,
@@ -78,6 +79,7 @@
         AddArgBindingsToTIRunContext,
         AddCallbackRunEndpoint,
         AddDagRunNoteUpdateEndpoint,
+        AddExpandedTICountToTIRunContext,
         AddLoopContext,
         AddLoopDecisionField,
         AddTerminalStateRetryReasonField,

airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py

@@ -147,6 +147,20 @@ def remove_multi_team_field(response: ResponseInfo) -> None:
         response.body.pop("multi_team", None)
 
 
+class AddExpandedTICountToTIRunContext(VersionChange):
+    """Add ``expanded_ti_count``, the size of the mapped expansion a task instance belongs to."""
+
+    description = __doc__
+
+    instructions_to_migrate_to_previous_version = (
+        schema(TIRunContext).field("expanded_ti_count").didnt_exist,
+    )
+
+    @convert_response_to_previous_version_for(TIRunContext)  # type: ignore[arg-type]
+    def remove_expanded_ti_count(response: ResponseInfo) -> None:  # type: ignore[misc]
+        response.body.pop("expanded_ti_count", None)
+
+
 class AddFirstTaskRescheduleStartDateField(VersionChange):
     """Add the ``first_task_reschedule_start_date`` field to TIRunContext."""
 

airflow-core/src/airflow/models/dynamic_region.py

@@ -18,7 +18,7 @@
 
 import hashlib
 import struct
-from collections.abc import Collection, Iterable, Iterator
+from collections.abc import Collection, Iterable, Iterator, Mapping
 from datetime import datetime
 from functools import partial
 from typing import TYPE_CHECKING, TypeVar
@@ -570,6 +570,23 @@ def loop_position(
     return None
 
 
+def get_trajectory(region: DynamicRegion, regions: Mapping[UUID, DynamicRegion]) -> int:
+    """
+    Return the 1-based position of ``region`` in its family: 1 for the original, ``n + 1`` for its n-th fork.
+
+    ``forked_from_region_id`` is unique, so a family is one chain and the position orders its regions.
+    """
+    trajectory = 1
+    lineage: set[UUID] = set()
+    while region.forked_from_region_id is not None:
+        if region.id in lineage:
+            raise ValueError("Cyclic region lineage")
+        lineage.add(region.id)
+        region = regions[region.forked_from_region_id]
+        trajectory += 1
+    return trajectory
+
+
 def build_loop_family_legs(
     query: _SelectT, family_member: DynamicRegion, *, iteration: int | None = None
 ) -> tuple[_SelectT, _SelectT]:

airflow-core/src/airflow/models/task_coordinates.py

@@ -21,8 +21,9 @@
 from uuid import UUID
 
 import attrs
-from sqlalchemy import FromClause, case, or_, select, tuple_
+from sqlalchemy import FromClause, case, func, or_, select, tuple_
 
+from airflow.api_fastapi.execution_api.datamodels.taskinstance import LoopContext
 from airflow.exceptions import TaskNotFound
 from airflow.models.dagbag import DBDagBag
 from airflow.models.dagrun import DagRun
@@ -31,6 +32,7 @@
     AmbiguousProducerError,
     DynamicRegion,
     ProducerContext,
+    get_trajectory,
     load_region_ancestry,
     loop_position,
     resolve_current_producers,
@@ -40,6 +42,7 @@
 from airflow.models.taskinstance import TaskInstance
 from airflow.serialization.definitions.mappedoperator import is_mapped
 from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup, SerializedMappedTaskGroup
+from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
     from collections.abc import Collection, Iterable, Sequence
@@ -238,15 +241,85 @@ def loop_context(
         group = enclosing_loop(task)
         if group is None:
             return None
+        position = loop_position(self._load_ancestry(ti), ti.region_id, ti.region_index, group.node_id)
+        if position is None:
+            raise ValueError("Task coordinates do not belong to the pinned loop")
+        return group, position[1]
+
+    def get_loop_context(
+        self, ti: TaskCoordinate, *, task: SerializedOperator | None = None
+    ) -> LoopContext | None:
+        """
+        Return the loop context a worker gets in ``TIRunContext.loop``, or None outside any loop.
+
+        Server-side consumers such as listeners get the same from a resolver built with :meth:`for_dag` for
+        the Dag they hold.
+        """
+        if ti.region_id == SENTINEL_REGION_ID:
+            return None
+        if task is None:
+            task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+        group = enclosing_loop(task)
+        if group is None:
+            return None
+        regions = self._load_ancestry(ti)
+        coordinate = DynamicRegion.find_physical_coordinate(
+            ti.region_id, ti.region_index, group.node_id, regions
+        )
+        if coordinate is None:
+            raise ValueError("Task coordinates do not belong to the pinned loop")
+        region_id, index = coordinate
+        region = regions[region_id]
+        terminal = self.get_task(
+            ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id
+        )
+        context = LoopContext(
+            node_id=group.node_id,
+            index=index,
+            max_iterations=group.max_iterations,
+            terminal_task_id=group.terminal_task_id,
+            terminal_is_mapped=terminal.get_needs_expansion(),
+            region_id=region_id,
+            trajectory=get_trajectory(region, regions),
+        )
+        # Set only for a fork, as the Execution API leaves out unset fields
+        if region.forked_from_region_id is not None:
+            context.forked_from_region_id = region.forked_from_region_id
+            context.resumes_from_index = region.resumes_from_index
+        return context
+
+    def _load_ancestry(self, ti: TaskCoordinate) -> dict[UUID, DynamicRegion]:
         regions = self._regions.setdefault((ti.dag_id, ti.run_id), {})
         if ti.region_id not in regions:
             regions.update(
                 load_region_ancestry([ti.region_id], dag_id=ti.dag_id, run_id=ti.run_id, session=self.session)
             )
-        position = loop_position(regions, ti.region_id, ti.region_index, group.node_id)
-        if position is None:
-            raise ValueError("Task coordinates do not belong to the pinned loop")
-        return group, position[1]
+        return regions
+
+    def get_expanded_ti_count(
+        self, ti: TaskCoordinate, *, task: SerializedOperator | None = None
+    ) -> int | None:
+        """
+        Count the live instances of a mapped task in the expansion ``ti`` belongs to, None if not mapped.
+
+        This is what a worker gets as ``TIRunContext.expanded_ti_count``.
+        """
+        if ti.region_id == SENTINEL_REGION_ID:
+            return None
+        if task is None:
+            task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+        if not task.get_needs_expansion():
+            return None
+        return self.session.scalar(
+            select(func.count()).where(
+                TaskInstance.working_set.is_(True),
+                TaskInstance.dag_id == ti.dag_id,
+                TaskInstance.run_id == ti.run_id,
+                TaskInstance.task_id == ti.task_id,
+                TaskInstance.region_id == ti.region_id,
+                or_(TaskInstance.state.is_(None), TaskInstance.state != TaskInstanceState.REMOVED),
+            )
+        )
 
     def get_task(
         self, dag_id: str, run_id: str, task_id: str, *, dag_version_id: UUID | None = None

airflow-core/src/airflow/utils/log/task_log_address.py

@@ -31,7 +31,7 @@
 
 from airflow.models.dagbag import DBDagBag
 from airflow.models.dagrun import DagRun
-from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
+from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion, get_trajectory
 from airflow.models.tasklog import LogTemplate
 from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup
 from airflow.utils.helpers import render_template
@@ -76,12 +76,8 @@ def region_log_position(
         kind = node_kinds.get(region.node_id) or ("map" if region.node_id == task_id else "loop")
         if kind not in {"loop", "map"}:
             raise ValueError(f"Unsupported log-address region kind: {kind}")
-        position = 1
-        predecessor = region
-        while predecessor.forked_from_region_id is not None:
-            predecessor = regions[predecessor.forked_from_region_id]
-            position += 1
-        suffix = f".{position}" if position > 1 else ""
+        trajectory = get_trajectory(region, regions)
+        suffix = f".{trajectory}" if trajectory > 1 else ""
         label = "pass" if kind == "loop" else "map"
         inside_loop = inside_loop or kind == "loop"
         frames.append(f"{label}={region_index}{suffix}")