Skip to content

Commit e6a2c9f

Browse files
kacpermudaclaude
andauthored
Expose task instance loop pass and expansion size to listeners (#75132)
A single event about a task instance should be enough to tell which execution it describes: which loop pass and rerun of the loop it ran in, how many instances its mapped expansion has, and why a loop gate stopped. Workers only had the loop iteration, a mapped task inside a loop could not be tied to the pass it ran in, reruns could only be ordered by comparing region ids from two events, and server side consumers such as listeners had to read region rows to place a stored task instance. Co-authored-by: Claude <[email protected]>
1 parent f3cf550 commit e6a2c9f

24 files changed

Lines changed: 688 additions & 144 deletions

File tree

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,9 +152,21 @@ definition.
152152
- The previous iteration's result, or ``None`` in iteration 1.
153153
* - ``loop.result``
154154
- The terminal body task's result (the XCom pushed under the key ``return_value``) for the current iteration, used by the gate.
155+
* - ``loop.region_id``
156+
- The id of the loop execution (trajectory) that holds this iteration. It is the same for every task
157+
instance of the iteration, mapped ones included.
158+
* - ``loop.trajectory``
159+
- The number of that trajectory: 1 for the loop's first execution, and one more each time a rerun
160+
replaces later iterations, see :ref:`loops-trajectories`.
161+
* - ``loop.forked_from_region_id``
162+
- The ``loop.region_id`` of the previous trajectory, or ``None`` in trajectory 1.
163+
* - ``loop.resumes_from_index``
164+
- The first iteration this trajectory holds, or ``None`` in trajectory 1, which holds them from 1.
155165

156166
Use ``loop.index`` for the loop iteration and ``ti.map_index`` for a mapped
157167
task instance's position. Each task instance can have several tries, each with its own ``ti.try_number``.
168+
The task instance also exposes the loop context as ``ti.loop``, and for a mapped task
169+
``ti.expanded_ti_count`` is the number of instances in its expansion, counted within the iteration.
158170

159171
Pass data between iterations
160172
============================
@@ -347,6 +359,12 @@ execution. Downstream tasks then run as they would after any successful gate: th
347359
gate that was marked successful by hand from one that stopped because ``until`` returned ``True``.
348360
Any later iterations retained after a selective clear remain unchanged.
349361

362+
The gate's task instance exposes its decision as ``ti.loop_decision``, a ``LoopDecision``, once the gate
363+
has evaluated it. With ``until``, ``stop`` is ``True`` when ``until`` ended the loop, even at its last
364+
iteration. When ``stop`` is ``False`` at the last iteration (``loop.index == loop.max_iterations - 1``),
365+
the loop reached ``max_iterations`` without meeting ``until`` and the gate fails. A gate that fails before
366+
deciding, for example because ``until`` raised, or that was marked successful by hand, has no decision.
367+
350368
Iterations and execution history
351369
================================
352370

@@ -414,3 +432,17 @@ which leave them in place. To keep later iterations, clear from the UI or the
414432
REST endpoint.
415433

416434
Tasks you leave uncleared are not recomputed using the results of the rerun.
435+
436+
.. _loops-trajectories:
437+
438+
Reruns and trajectories
439+
-----------------------
440+
441+
Clearing later iterations starts a new execution of the loop that replaces the previous ones from the
442+
iteration after the rerun gate. Each execution is a trajectory, identified by ``loop.region_id`` and
443+
numbered by ``loop.trajectory``: 1 for the first one, growing by one with each rerun. The rerun gate, and
444+
whatever you clear in its iteration, run again in the trajectory that iteration belongs to, with their
445+
next try. The replacement iterations run in the new trajectory, from ``loop.resumes_from_index``, and their
446+
tries count from 1 again for each task. In the example above, iteration 2 stays in trajectory 1, while
447+
replacement iterations 3 and 4 run in trajectory 2, with ``loop.resumes_from_index`` 3. Task log paths show
448+
the trajectory after the iteration, as in ``pass=3.2``.

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

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,9 @@
3232
field_validator,
3333
model_validator,
3434
)
35+
from pydantic.json_schema import SkipJsonSchema
3536

36-
from airflow.api_fastapi.common.region import OmitsMissingRegion, RegionId, RegionIndex
37+
from airflow.api_fastapi.common.region import OmitsMissingRegion, RegionId, RegionIndex, _drop_null_default
3738
from airflow.api_fastapi.common.types import UtcDateTime
3839
from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
3940
from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile
@@ -457,6 +458,38 @@ class LoopContext(BaseModel):
457458
terminal_task_id: str
458459
terminal_is_mapped: bool
459460

461+
region_id: uuid.UUID
462+
"""
463+
The loop region holding this iteration, which identifies the trajectory it runs in.
464+
465+
A mapped task's own ``region_id`` is its expansion's, so this is how its instances are tied to the
466+
unmapped tasks of the same iteration.
467+
"""
468+
469+
trajectory: int = Field(ge=1)
470+
"""
471+
The number of ``region_id`` in the loop's chain of regions: 1 for the original, n + 1 for its n-th fork.
472+
473+
Clearing a loop gate with the later iterations forks the newest loop region, and the iterations from
474+
``resumes_from_index`` on run again in the fork.
475+
"""
476+
477+
forked_from_region_id: Annotated[
478+
SkipJsonSchema[None] | uuid.UUID, Field(json_schema_extra=_drop_null_default)
479+
] = None
480+
"""
481+
The region ``region_id`` was forked from, the previous one in the chain. Absent for the original region.
482+
483+
The fork replaces the iterations from ``resumes_from_index`` on, which can be held by earlier regions of
484+
the chain too.
485+
"""
486+
487+
resumes_from_index: Annotated[
488+
SkipJsonSchema[None] | int,
489+
Field(json_schema_extra=_drop_null_default),
490+
] = None
491+
"""The first iteration ``region_id`` holds. Absent for the original region, which holds them from 1."""
492+
460493

461494
class TIRunContext(BaseModel):
462495
"""Response schema for TaskInstance run context."""
@@ -466,6 +499,16 @@ class TIRunContext(BaseModel):
466499

467500
loop: LoopContext | None = None
468501

502+
expanded_ti_count: Annotated[
503+
SkipJsonSchema[None] | int,
504+
Field(json_schema_extra=_drop_null_default),
505+
] = None
506+
"""
507+
Number of live, non-removed instances in the mapped task's expansion this instance belongs to.
508+
509+
``None`` for a task that isn't mapped. Inside a loop, it counts the instances of this iteration only.
510+
"""
511+
469512
task_reschedule_count: int = 0
470513
"""How many times the task has been rescheduled."""
471514

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

Lines changed: 7 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,6 @@
5959
from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
6060
DagRunNoteUpdatePayload,
6161
InactiveAssetsResponse,
62-
LoopContext,
6362
PreviousTIResponse,
6463
PrevSuccessfulDagRunResponse,
6564
TaskBreadcrumbsResponse,
@@ -371,19 +370,13 @@ def ti_run(
371370
should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries),
372371
multi_team=conf.getboolean("core", "multi_team"),
373372
)
374-
resolver = TaskCoordinateResolver(dag_bag, session)
375-
if loop_context := resolver.loop_context(ti):
376-
group, index = loop_context
377-
terminal = resolver.get_task(
378-
ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id
379-
)
380-
context.loop = LoopContext(
381-
node_id=group.node_id,
382-
index=index,
383-
max_iterations=group.max_iterations,
384-
terminal_task_id=group.terminal_task_id,
385-
terminal_is_mapped=terminal.get_needs_expansion(),
386-
)
373+
if ti.region_id != SENTINEL_REGION_ID:
374+
resolver = TaskCoordinateResolver(dag_bag, session)
375+
task = resolver.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id)
376+
if (expanded_ti_count := resolver.get_expanded_ti_count(ti, task=task)) is not None:
377+
context.expanded_ti_count = expanded_ti_count
378+
if (loop := resolver.get_loop_context(ti, task=task)) is not None:
379+
context.loop = loop
387380
if first_task_reschedule_start_date is not None:
388381
context.first_task_reschedule_start_date = first_task_reschedule_start_date
389382

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
AddArgBindingsToTIRunContext,
5757
AddCallbackRunEndpoint,
5858
AddDagRunNoteUpdateEndpoint,
59+
AddExpandedTICountToTIRunContext,
5960
AddFirstTaskRescheduleStartDateField,
6061
AddLoopContext,
6162
AddLoopDecisionField,
@@ -78,6 +79,7 @@
7879
AddArgBindingsToTIRunContext,
7980
AddCallbackRunEndpoint,
8081
AddDagRunNoteUpdateEndpoint,
82+
AddExpandedTICountToTIRunContext,
8183
AddLoopContext,
8284
AddLoopDecisionField,
8385
AddTerminalStateRetryReasonField,

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,20 @@ def remove_multi_team_field(response: ResponseInfo) -> None:
147147
response.body.pop("multi_team", None)
148148

149149

150+
class AddExpandedTICountToTIRunContext(VersionChange):
151+
"""Add ``expanded_ti_count``, the size of the mapped expansion a task instance belongs to."""
152+
153+
description = __doc__
154+
155+
instructions_to_migrate_to_previous_version = (
156+
schema(TIRunContext).field("expanded_ti_count").didnt_exist,
157+
)
158+
159+
@convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type]
160+
def remove_expanded_ti_count(response: ResponseInfo) -> None: # type: ignore[misc]
161+
response.body.pop("expanded_ti_count", None)
162+
163+
150164
class AddFirstTaskRescheduleStartDateField(VersionChange):
151165
"""Add the ``first_task_reschedule_start_date`` field to TIRunContext."""
152166

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
import hashlib
2020
import struct
21-
from collections.abc import Collection, Iterable, Iterator
21+
from collections.abc import Collection, Iterable, Iterator, Mapping
2222
from datetime import datetime
2323
from functools import partial
2424
from typing import TYPE_CHECKING, TypeVar
@@ -570,6 +570,23 @@ def loop_position(
570570
return None
571571

572572

573+
def get_trajectory(region: DynamicRegion, regions: Mapping[UUID, DynamicRegion]) -> int:
574+
"""
575+
Return the 1-based position of ``region`` in its family: 1 for the original, ``n + 1`` for its n-th fork.
576+
577+
``forked_from_region_id`` is unique, so a family is one chain and the position orders its regions.
578+
"""
579+
trajectory = 1
580+
lineage: set[UUID] = set()
581+
while region.forked_from_region_id is not None:
582+
if region.id in lineage:
583+
raise ValueError("Cyclic region lineage")
584+
lineage.add(region.id)
585+
region = regions[region.forked_from_region_id]
586+
trajectory += 1
587+
return trajectory
588+
589+
573590
def build_loop_family_legs(
574591
query: _SelectT, family_member: DynamicRegion, *, iteration: int | None = None
575592
) -> tuple[_SelectT, _SelectT]:

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

Lines changed: 78 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,9 @@
2121
from uuid import UUID
2222

2323
import attrs
24-
from sqlalchemy import FromClause, case, or_, select, tuple_
24+
from sqlalchemy import FromClause, case, func, or_, select, tuple_
2525

26+
from airflow.api_fastapi.execution_api.datamodels.taskinstance import LoopContext
2627
from airflow.exceptions import TaskNotFound
2728
from airflow.models.dagbag import DBDagBag
2829
from airflow.models.dagrun import DagRun
@@ -31,6 +32,7 @@
3132
AmbiguousProducerError,
3233
DynamicRegion,
3334
ProducerContext,
35+
get_trajectory,
3436
load_region_ancestry,
3537
loop_position,
3638
resolve_current_producers,
@@ -40,6 +42,7 @@
4042
from airflow.models.taskinstance import TaskInstance
4143
from airflow.serialization.definitions.mappedoperator import is_mapped
4244
from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup, SerializedMappedTaskGroup
45+
from airflow.utils.state import TaskInstanceState
4346

4447
if TYPE_CHECKING:
4548
from collections.abc import Collection, Iterable, Sequence
@@ -238,15 +241,85 @@ def loop_context(
238241
group = enclosing_loop(task)
239242
if group is None:
240243
return None
244+
position = loop_position(self._load_ancestry(ti), ti.region_id, ti.region_index, group.node_id)
245+
if position is None:
246+
raise ValueError("Task coordinates do not belong to the pinned loop")
247+
return group, position[1]
248+
249+
def get_loop_context(
250+
self, ti: TaskCoordinate, *, task: SerializedOperator | None = None
251+
) -> LoopContext | None:
252+
"""
253+
Return the loop context a worker gets in ``TIRunContext.loop``, or None outside any loop.
254+
255+
Server-side consumers such as listeners get the same from a resolver built with :meth:`for_dag` for
256+
the Dag they hold.
257+
"""
258+
if ti.region_id == SENTINEL_REGION_ID:
259+
return None
260+
if task is None:
261+
task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id)
262+
group = enclosing_loop(task)
263+
if group is None:
264+
return None
265+
regions = self._load_ancestry(ti)
266+
coordinate = DynamicRegion.find_physical_coordinate(
267+
ti.region_id, ti.region_index, group.node_id, regions
268+
)
269+
if coordinate is None:
270+
raise ValueError("Task coordinates do not belong to the pinned loop")
271+
region_id, index = coordinate
272+
region = regions[region_id]
273+
terminal = self.get_task(
274+
ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id
275+
)
276+
context = LoopContext(
277+
node_id=group.node_id,
278+
index=index,
279+
max_iterations=group.max_iterations,
280+
terminal_task_id=group.terminal_task_id,
281+
terminal_is_mapped=terminal.get_needs_expansion(),
282+
region_id=region_id,
283+
trajectory=get_trajectory(region, regions),
284+
)
285+
# Set only for a fork, as the Execution API leaves out unset fields
286+
if region.forked_from_region_id is not None:
287+
context.forked_from_region_id = region.forked_from_region_id
288+
context.resumes_from_index = region.resumes_from_index
289+
return context
290+
291+
def _load_ancestry(self, ti: TaskCoordinate) -> dict[UUID, DynamicRegion]:
241292
regions = self._regions.setdefault((ti.dag_id, ti.run_id), {})
242293
if ti.region_id not in regions:
243294
regions.update(
244295
load_region_ancestry([ti.region_id], dag_id=ti.dag_id, run_id=ti.run_id, session=self.session)
245296
)
246-
position = loop_position(regions, ti.region_id, ti.region_index, group.node_id)
247-
if position is None:
248-
raise ValueError("Task coordinates do not belong to the pinned loop")
249-
return group, position[1]
297+
return regions
298+
299+
def get_expanded_ti_count(
300+
self, ti: TaskCoordinate, *, task: SerializedOperator | None = None
301+
) -> int | None:
302+
"""
303+
Count the live instances of a mapped task in the expansion ``ti`` belongs to, None if not mapped.
304+
305+
This is what a worker gets as ``TIRunContext.expanded_ti_count``.
306+
"""
307+
if ti.region_id == SENTINEL_REGION_ID:
308+
return None
309+
if task is None:
310+
task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id)
311+
if not task.get_needs_expansion():
312+
return None
313+
return self.session.scalar(
314+
select(func.count()).where(
315+
TaskInstance.working_set.is_(True),
316+
TaskInstance.dag_id == ti.dag_id,
317+
TaskInstance.run_id == ti.run_id,
318+
TaskInstance.task_id == ti.task_id,
319+
TaskInstance.region_id == ti.region_id,
320+
or_(TaskInstance.state.is_(None), TaskInstance.state != TaskInstanceState.REMOVED),
321+
)
322+
)
250323

251324
def get_task(
252325
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‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131

3232
from airflow.models.dagbag import DBDagBag
3333
from airflow.models.dagrun import DagRun
34-
from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
34+
from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion, get_trajectory
3535
from airflow.models.tasklog import LogTemplate
3636
from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup
3737
from airflow.utils.helpers import render_template
@@ -76,12 +76,8 @@ def region_log_position(
7676
kind = node_kinds.get(region.node_id) or ("map" if region.node_id == task_id else "loop")
7777
if kind not in {"loop", "map"}:
7878
raise ValueError(f"Unsupported log-address region kind: {kind}")
79-
position = 1
80-
predecessor = region
81-
while predecessor.forked_from_region_id is not None:
82-
predecessor = regions[predecessor.forked_from_region_id]
83-
position += 1
84-
suffix = f".{position}" if position > 1 else ""
79+
trajectory = get_trajectory(region, regions)
80+
suffix = f".{trajectory}" if trajectory > 1 else ""
8581
label = "pass" if kind == "loop" else "map"
8682
inside_loop = inside_loop or kind == "loop"
8783
frames.append(f"{label}={region_index}{suffix}")

0 commit comments

Comments
 (0)