This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e6a2c9f4e78 Expose task instance loop pass and expansion size to
listeners (#75132)
e6a2c9f4e78 is described below
commit e6a2c9f4e7817012cd503a8663f64aa14c2a490a
Author: Kacper Muda <[email protected]>
AuthorDate: Wed Oct 21 13:30:54 2026 +0200
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]>
---
.../docs/authoring-and-scheduling/loops.rst | 32 +++++++
.../execution_api/datamodels/taskinstance.py | 45 +++++++++-
.../execution_api/routes/task_instances.py | 21 ++---
.../api_fastapi/execution_api/versions/__init__.py | 2 +
.../execution_api/versions/v2026_10_30.py | 14 ++++
airflow-core/src/airflow/models/dynamic_region.py | 19 ++++-
.../src/airflow/models/task_coordinates.py | 83 ++++++++++++++++--
.../src/airflow/utils/log/task_log_address.py | 10 +--
.../versions/head/test_task_instances.py | 51 ++++++++++++
.../tests/unit/models/test_dynamic_region.py | 17 +++-
.../tests/unit/models/test_task_coordinates.py | 82 ++++++++++++++++++
go-sdk/pkg/execution/genmodels/models.gen.go | 16 ++++
go-sdk/schema/supervisor-schema.json | 97 ++++++++++++++--------
.../src/airflow/sdk/api/datamodels/_generated.py | 5 ++
task-sdk/src/airflow/sdk/execution_time/loop.py | 18 ++++
.../airflow/sdk/execution_time/schema/schema.json | 97 ++++++++++++++--------
.../sdk/execution_time/schema/versions/__init__.py | 2 +
.../execution_time/schema/versions/v2026_10_30.py | 14 ++++
.../src/airflow/sdk/execution_time/task_runner.py | 26 +++++-
task-sdk/src/airflow/sdk/types.py | 11 +++
.../execution_time/schema/test_migrator.py | 5 ++
.../tests/task_sdk/execution_time/test_loop.py | 54 ++++++++++++
ts-sdk/schema/supervisor-schema.json | 97 ++++++++++++++--------
ts-sdk/src/generated/supervisor.ts | 14 +++-
24 files changed, 688 insertions(+), 144 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst
b/airflow-core/docs/authoring-and-scheduling/loops.rst
index 7a6e04e0c59..9e414563e59 100644
--- a/airflow-core/docs/authoring-and-scheduling/loops.rst
+++ b/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``.
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
index 09ce1193e90..8a1deb203fd 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
@@ -32,8 +32,9 @@ from pydantic import (
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."""
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index c27e2f74f6a..226ec01b91e 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -59,7 +59,6 @@ from
airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_ar
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
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index 3df478de386..e85afc58b9b 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
@@ -56,6 +56,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30
import (
AddArgBindingsToTIRunContext,
AddCallbackRunEndpoint,
AddDagRunNoteUpdateEndpoint,
+ AddExpandedTICountToTIRunContext,
AddFirstTaskRescheduleStartDateField,
AddLoopContext,
AddLoopDecisionField,
@@ -78,6 +79,7 @@ bundle = VersionBundle(
AddArgBindingsToTIRunContext,
AddCallbackRunEndpoint,
AddDagRunNoteUpdateEndpoint,
+ AddExpandedTICountToTIRunContext,
AddLoopContext,
AddLoopDecisionField,
AddTerminalStateRetryReasonField,
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
index 26d0aea8878..1f7780ee321 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
@@ -147,6 +147,20 @@ class AddMultiTeamToTIRunContext(VersionChange):
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."""
diff --git a/airflow-core/src/airflow/models/dynamic_region.py
b/airflow-core/src/airflow/models/dynamic_region.py
index 968b9224829..63a6574b265 100644
--- a/airflow-core/src/airflow/models/dynamic_region.py
+++ b/airflow-core/src/airflow/models/dynamic_region.py
@@ -18,7 +18,7 @@ from __future__ import annotations
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]:
diff --git a/airflow-core/src/airflow/models/task_coordinates.py
b/airflow-core/src/airflow/models/task_coordinates.py
index 73c2f87ed8d..fd85cbddc60 100644
--- a/airflow-core/src/airflow/models/task_coordinates.py
+++ b/airflow-core/src/airflow/models/task_coordinates.py
@@ -21,8 +21,9 @@ from typing import TYPE_CHECKING, Protocol
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 @@ from airflow.models.dynamic_region import (
AmbiguousProducerError,
DynamicRegion,
ProducerContext,
+ get_trajectory,
load_region_ancestry,
loop_position,
resolve_current_producers,
@@ -40,6 +42,7 @@ from airflow.models.dynamic_region import (
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 @@ class TaskCoordinateResolver:
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
diff --git a/airflow-core/src/airflow/utils/log/task_log_address.py
b/airflow-core/src/airflow/utils/log/task_log_address.py
index b72eaab72ec..27be654a274 100644
--- a/airflow-core/src/airflow/utils/log/task_log_address.py
+++ b/airflow-core/src/airflow/utils/log/task_log_address.py
@@ -31,7 +31,7 @@ from sqlalchemy.orm.attributes import NO_VALUE
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}")
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
index 6a81588ffc5..d241600d281 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
@@ -657,6 +657,7 @@ class TestTIRunState:
assert response.status_code == 200
if version == "2026-06-30":
assert "loop" not in response.json()
+ assert "expanded_ti_count" not in response.json()
return
assert response.json()["loop"] == {
"node_id": loop.group_id,
@@ -664,7 +665,57 @@ class TestTIRunState:
"max_iterations": 4,
"terminal_task_id": loop.terminal_task_id,
"terminal_is_mapped": mapped,
+ "region_id": str(parent.id),
+ "trajectory": 1,
}
+ assert response.json().get("expanded_ti_count") == (1 if mapped else
None)
+
+ def
test_startup_loop_context_describes_the_fork_holding_the_iteration(self,
client, session, dag_maker):
+ @task_group
+ def body():
+ EmptyOperator(task_id="terminal")
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=4)
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ ti = next(ti for ti in dr.task_instances if ti.task_id ==
loop.terminal_task_id)
+ original = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id, run_id=dr.run_id, node_id=loop.group_id,
session=session
+ )
+ session.add(original)
+ session.flush()
+ fork = DynamicRegion.fork_family(original.id, 2, {original.id:
original}, session=session)
+ session.flush()
+ ti.region_id, ti.region_index = fork.id, 2
+ ti.state = State.QUEUED
+ session.commit()
+
+ response = client.patch(f"/execution/task-instances/{ti.id}/run",
json=self.RUN_PAYLOAD)
+
+ assert response.status_code == 200
+ loop_context = response.json()["loop"]
+ assert loop_context["region_id"] == str(fork.id)
+ assert loop_context["trajectory"] == 2
+ assert loop_context["forked_from_region_id"] == str(original.id)
+ assert loop_context["resumes_from_index"] == 2
+
+ @pytest.mark.parametrize("removed", [0, 1])
+ def test_startup_expanded_ti_count_counts_live_instances_of_the_expansion(
+ self, client, session, dag_maker, removed
+ ):
+ with dag_maker(serialized=True):
+ PythonOperator.partial(task_id="mapped",
python_callable=str).expand(op_args=[[1], [2], [3]])
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ tis = sorted(dr.task_instances, key=lambda ti: ti.map_index)
+ for ti in tis[3 - removed :]:
+ ti.state = State.REMOVED
+ tis[0].state = State.QUEUED
+ session.commit()
+
+ response = client.patch(f"/execution/task-instances/{tis[0].id}/run",
json=self.RUN_PAYLOAD)
+
+ assert response.status_code == 200
+ assert response.json()["expanded_ti_count"] == 3 - removed
@pytest.mark.parametrize("regional", [False, True])
def test_startup_xcom_cleanup_keys_belong_to_the_starting_try(
diff --git a/airflow-core/tests/unit/models/test_dynamic_region.py
b/airflow-core/tests/unit/models/test_dynamic_region.py
index 573046e4e9b..591cbbadd24 100644
--- a/airflow-core/tests/unit/models/test_dynamic_region.py
+++ b/airflow-core/tests/unit/models/test_dynamic_region.py
@@ -18,7 +18,6 @@
from __future__ import annotations
import struct
-from typing import TYPE_CHECKING
from uuid import uuid4
import pytest
@@ -36,6 +35,7 @@ from airflow.models.dynamic_region import (
ProducerContext,
build_loop_family_legs,
build_loop_sequence_order,
+ get_trajectory,
resolve_current_producers,
select_loop_producer_ids,
)
@@ -52,9 +52,6 @@ from airflow.utils.state import TaskInstanceState
from tests_common.test_utils.asserts import capture_orm_selects
from tests_common.test_utils.db import clear_db_runs
-if TYPE_CHECKING:
- from airflow.models.dagrun import DagRun
-
pytestmark = pytest.mark.db_test
@@ -1114,3 +1111,15 @@ def
test_loop_sequence_of_mapped_task_skips_passes_a_fork_replaced(loop_run, ses
(True, 4, 0),
(True, 4, 1),
]
+
+
[email protected]("forks", [0, 1, 2])
+def test_trajectory_counts_the_forks_a_region_descends_from(forks):
+ dag_run = DagRun(dag_id="dag", run_id="run")
+ region = make_region(dag_run, id=uuid4())
+ regions = {region.id: region}
+ for _ in range(forks):
+ region = make_region(dag_run, id=uuid4(),
forked_from_region_id=region.id)
+ regions[region.id] = region
+
+ assert get_trajectory(region, regions) == forks + 1
diff --git a/airflow-core/tests/unit/models/test_task_coordinates.py
b/airflow-core/tests/unit/models/test_task_coordinates.py
index b74f38460a2..57af3454f80 100644
--- a/airflow-core/tests/unit/models/test_task_coordinates.py
+++ b/airflow-core/tests/unit/models/test_task_coordinates.py
@@ -105,6 +105,9 @@ def
test_mapped_task_loop_iteration_is_separate_from_map_index(dag_maker, sessio
assert resolver.loop_iterations(ti) == [("body", 3)]
assert resolver.public_map_index(ti) == 1
+ loop = resolver.get_loop_context(ti)
+ assert loop is not None
+ assert (loop.index, loop.region_id, loop.trajectory) == (3,
loop_region.id, 1)
@pytest.mark.parametrize("run_count", [1, 3])
@@ -165,6 +168,85 @@ def
test_loop_context_reuses_ancestry_loaded_by_loop_iterations(loop_coordinates
assert (group.group_id, index) == ("body", 1)
+def
test_loop_context_describes_the_region_holding_the_iteration(forked_loop_coordinates,
session):
+ _, producer, consumer, _, outside = forked_loop_coordinates
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+ common = {
+ "node_id": "body",
+ "index": 2,
+ "max_iterations": 3,
+ "terminal_task_id": "body.consumer",
+ "terminal_is_mapped": False,
+ }
+
+ original = resolver.get_loop_context(producer)
+ fork = resolver.get_loop_context(consumer)
+
+ assert original is not None
+ assert fork is not None
+ assert original.model_dump(exclude_unset=True) == {
+ **common,
+ "region_id": producer.region_id,
+ "trajectory": 1,
+ }
+ assert fork.model_dump(exclude_unset=True) == {
+ **common,
+ "region_id": consumer.region_id,
+ "trajectory": 2,
+ "forked_from_region_id": producer.region_id,
+ "resumes_from_index": 2,
+ }
+ assert resolver.get_loop_context(outside) is None
+
+
+def
test_expanded_ti_count_counts_live_instances_of_the_expansion(loop_coordinates,
dag_maker, session):
+ with dag_maker("mapped-expansion", serialized=True):
+ PythonOperator.partial(task_id="mapped",
python_callable=str).expand(op_args=[[1], [2], [3]])
+ mapped = sorted(dag_maker.create_dagrun().task_instances, key=lambda ti:
ti.map_index)
+ mapped[2].state = TaskInstanceState.REMOVED
+ _, _, consumer, _, outside = loop_coordinates
+ session.flush()
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+ resolver.get_expanded_ti_count(mapped[0])
+
+ with assert_queries_count(1):
+ assert resolver.get_expanded_ti_count(mapped[1]) == 2
+ assert resolver.get_expanded_ti_count(consumer) is None
+ assert resolver.get_expanded_ti_count(outside) is None
+
+
+def
test_expanded_ti_count_counts_each_instance_of_one_loop_iteration_once(dag_maker,
session):
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="mapped",
python_callable=str).expand(op_args=[[1], [2]])
+
+ with dag_maker(serialized=True):
+ create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ first, second = sorted(
+ (ti for ti in dr.task_instances if ti.task_id == "body.mapped"),
key=lambda ti: ti.map_index
+ )
+ first.prepare_db_for_next_try(session)
+ next_expansion = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id=first.task_id,
+ parent_region_id=session.get(DynamicRegion,
first.region_id).parent_region_id,
+ parent_region_index=2,
+ session=session,
+ )
+ session.add(next_expansion)
+ session.flush()
+ next_iteration = TaskInstance(task=first.task, run_id=dr.run_id,
dag_version_id=first.dag_version_id)
+ next_iteration.region_id, next_iteration.region_index = next_expansion.id, 0
+ session.add(next_iteration)
+ session.flush()
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+
+ assert resolver.get_expanded_ti_count(second) == 2
+ assert resolver.get_expanded_ti_count(next_iteration) == 1
+
+
def test_removed_task_coordinates_degrade_to_stored_region_data(dag_maker,
session):
with dag_maker("removed-coordinates", serialized=True):
EmptyOperator(task_id="kept")
diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go
b/go-sdk/pkg/execution/genmodels/models.gen.go
index 0a5dba694e6..fbf36031cea 100644
--- a/go-sdk/pkg/execution/genmodels/models.gen.go
+++ b/go-sdk/pkg/execution/genmodels/models.gen.go
@@ -1320,6 +1320,10 @@ type LogicalDates []time.Time
// Pinned loop definition and enclosing iteration for a task execution.
type LoopContext struct {
+ // ForkedFromRegionID corresponds to the JSON schema field
+ // "forked_from_region_id".
+ ForkedFromRegionID *string `msgpack:"forked_from_region_id,omitempty"`
+
// Index corresponds to the JSON schema field "index".
Index int `msgpack:"index"`
@@ -1329,11 +1333,20 @@ type LoopContext struct {
// NodeID corresponds to the JSON schema field "node_id".
NodeID string `msgpack:"node_id"`
+ // RegionID corresponds to the JSON schema field "region_id".
+ RegionID string `msgpack:"region_id"`
+
+ // ResumesFromIndex corresponds to the JSON schema field
"resumes_from_index".
+ ResumesFromIndex *int `msgpack:"resumes_from_index,omitempty"`
+
// TerminalIsMapped corresponds to the JSON schema field
"terminal_is_mapped".
TerminalIsMapped bool `msgpack:"terminal_is_mapped"`
// TerminalTaskID corresponds to the JSON schema field
"terminal_task_id".
TerminalTaskID string `msgpack:"terminal_task_id"`
+
+ // Trajectory corresponds to the JSON schema field "trajectory".
+ Trajectory int `msgpack:"trajectory"`
}
// A loop gate's evaluation of its “until“ callable.
@@ -1696,6 +1709,9 @@ type TIRunContext struct {
// DagRun corresponds to the JSON schema field "dag_run".
DagRun DagRun `msgpack:"dag_run"`
+ // ExpandedTICount corresponds to the JSON schema field
"expanded_ti_count".
+ ExpandedTICount *int `msgpack:"expanded_ti_count,omitempty"`
+
// FirstTaskRescheduleStartDate corresponds to the JSON schema field
// "first_task_reschedule_start_date".
FirstTaskRescheduleStartDate interface{}
`msgpack:"first_task_reschedule_start_date,omitempty"`
diff --git a/go-sdk/schema/supervisor-schema.json
b/go-sdk/schema/supervisor-schema.json
index 56a6967f17e..e739708d12d 100644
--- a/go-sdk/schema/supervisor-schema.json
+++ b/go-sdk/schema/supervisor-schema.json
@@ -3242,42 +3242,6 @@
"title": "LazyDeserializedDAG",
"type": "object"
},
- "LoopContext": {
- "description": "Pinned loop definition and enclosing iteration for a
task execution.",
- "properties": {
- "node_id": {
- "title": "Node Id",
- "type": "string"
- },
- "index": {
- "minimum": 1,
- "title": "Index",
- "type": "integer"
- },
- "max_iterations": {
- "exclusiveMinimum": 0,
- "title": "Max Iterations",
- "type": "integer"
- },
- "terminal_task_id": {
- "title": "Terminal Task Id",
- "type": "string"
- },
- "terminal_is_mapped": {
- "title": "Terminal Is Mapped",
- "type": "boolean"
- }
- },
- "required": [
- "node_id",
- "index",
- "max_iterations",
- "terminal_task_id",
- "terminal_is_mapped"
- ],
- "title": "LoopContext",
- "type": "object"
- },
"LoopDecisionPayload": {
"additionalProperties": false,
"description": "A loop gate's evaluation of its ``until`` callable.",
@@ -5253,6 +5217,63 @@
"title": "DagRun",
"type": "object"
},
+ "LoopContext": {
+ "description": "Pinned loop definition and enclosing iteration for a
task execution.",
+ "properties": {
+ "node_id": {
+ "title": "Node Id",
+ "type": "string"
+ },
+ "index": {
+ "minimum": 1,
+ "title": "Index",
+ "type": "integer"
+ },
+ "max_iterations": {
+ "exclusiveMinimum": 0,
+ "title": "Max Iterations",
+ "type": "integer"
+ },
+ "terminal_task_id": {
+ "title": "Terminal Task Id",
+ "type": "string"
+ },
+ "terminal_is_mapped": {
+ "title": "Terminal Is Mapped",
+ "type": "boolean"
+ },
+ "region_id": {
+ "format": "uuid",
+ "title": "Region Id",
+ "type": "string"
+ },
+ "trajectory": {
+ "minimum": 1,
+ "title": "Trajectory",
+ "type": "integer"
+ },
+ "forked_from_region_id": {
+ "format": "uuid",
+ "title": "Forked From Region Id",
+ "type": "string"
+ },
+ "resumes_from_index": {
+ "title": "Resumes From Index",
+ "type": "integer"
+ }
+ },
+ "required": [
+ "node_id",
+ "index",
+ "max_iterations",
+ "terminal_task_id",
+ "terminal_is_mapped",
+ "region_id",
+ "trajectory"
+ ],
+ "title": "LoopContext",
+ "type": "object"
+ },
"TIRunContext": {
"description": "Response schema for TaskInstance run context.",
"properties": {
@@ -5270,6 +5291,10 @@
],
"default": null
},
+ "expanded_ti_count": {
+ "title": "Expanded Ti Count",
+ "type": "integer"
+ },
"task_reschedule_count": {
"default": 0,
"title": "Task Reschedule Count",
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index 0e677c93e7a..6c4b2536fc4 100644
--- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -342,6 +342,10 @@ class LoopContext(BaseModel):
max_iterations: Annotated[int, Field(gt=0, title="Max Iterations")]
terminal_task_id: Annotated[str, Field(title="Terminal Task Id")]
terminal_is_mapped: Annotated[bool, Field(title="Terminal Is Mapped")]
+ region_id: Annotated[UUID, Field(title="Region Id")]
+ trajectory: Annotated[int, Field(ge=1, title="Trajectory")]
+ forked_from_region_id: Annotated[UUID | None, Field(title="Forked From
Region Id")] = None
+ resumes_from_index: Annotated[int | None, Field(title="Resumes From
Index")] = None
class LoopDecisionPayload(BaseModel):
@@ -999,6 +1003,7 @@ class TIRunContext(BaseModel):
dag_run: DagRun
loop: LoopContext | None = None
+ expanded_ti_count: Annotated[int | None, Field(title="Expanded Ti Count")]
= None
task_reschedule_count: Annotated[int | None, Field(title="Task Reschedule
Count")] = 0
first_task_reschedule_start_date: Annotated[
AwareDatetime | None, Field(title="First Task Reschedule Start Date")
diff --git a/task-sdk/src/airflow/sdk/execution_time/loop.py
b/task-sdk/src/airflow/sdk/execution_time/loop.py
index f8aac861191..b515001fe65 100644
--- a/task-sdk/src/airflow/sdk/execution_time/loop.py
+++ b/task-sdk/src/airflow/sdk/execution_time/loop.py
@@ -31,6 +31,8 @@ from airflow.sdk.execution_time.lazy_sequence import
LazyXComSequence
from airflow.sdk.execution_time.xcom import XCom
if TYPE_CHECKING:
+ from uuid import UUID
+
from airflow.sdk.api.datamodels._generated import LoopContext
from airflow.sdk.definitions._internal.loop import LoopGateOperator
from airflow.sdk.definitions.context import Context
@@ -54,6 +56,22 @@ class LoopContextAccessor:
def max_iterations(self) -> int:
return self._descriptor.max_iterations
+ @property
+ def region_id(self) -> UUID:
+ return self._descriptor.region_id
+
+ @property
+ def trajectory(self) -> int:
+ return self._descriptor.trajectory
+
+ @property
+ def forked_from_region_id(self) -> UUID | None:
+ return self._descriptor.forked_from_region_id
+
+ @property
+ def resumes_from_index(self) -> int | None:
+ return self._descriptor.resumes_from_index
+
@property
def previous(self) -> Any:
# Iterations count from one, so the first pass has no previous result.
diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
index 56a6967f17e..e739708d12d 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
@@ -3242,42 +3242,6 @@
"title": "LazyDeserializedDAG",
"type": "object"
},
- "LoopContext": {
- "description": "Pinned loop definition and enclosing iteration for a
task execution.",
- "properties": {
- "node_id": {
- "title": "Node Id",
- "type": "string"
- },
- "index": {
- "minimum": 1,
- "title": "Index",
- "type": "integer"
- },
- "max_iterations": {
- "exclusiveMinimum": 0,
- "title": "Max Iterations",
- "type": "integer"
- },
- "terminal_task_id": {
- "title": "Terminal Task Id",
- "type": "string"
- },
- "terminal_is_mapped": {
- "title": "Terminal Is Mapped",
- "type": "boolean"
- }
- },
- "required": [
- "node_id",
- "index",
- "max_iterations",
- "terminal_task_id",
- "terminal_is_mapped"
- ],
- "title": "LoopContext",
- "type": "object"
- },
"LoopDecisionPayload": {
"additionalProperties": false,
"description": "A loop gate's evaluation of its ``until`` callable.",
@@ -5253,6 +5217,63 @@
"title": "DagRun",
"type": "object"
},
+ "LoopContext": {
+ "description": "Pinned loop definition and enclosing iteration for a
task execution.",
+ "properties": {
+ "node_id": {
+ "title": "Node Id",
+ "type": "string"
+ },
+ "index": {
+ "minimum": 1,
+ "title": "Index",
+ "type": "integer"
+ },
+ "max_iterations": {
+ "exclusiveMinimum": 0,
+ "title": "Max Iterations",
+ "type": "integer"
+ },
+ "terminal_task_id": {
+ "title": "Terminal Task Id",
+ "type": "string"
+ },
+ "terminal_is_mapped": {
+ "title": "Terminal Is Mapped",
+ "type": "boolean"
+ },
+ "region_id": {
+ "format": "uuid",
+ "title": "Region Id",
+ "type": "string"
+ },
+ "trajectory": {
+ "minimum": 1,
+ "title": "Trajectory",
+ "type": "integer"
+ },
+ "forked_from_region_id": {
+ "format": "uuid",
+ "title": "Forked From Region Id",
+ "type": "string"
+ },
+ "resumes_from_index": {
+ "title": "Resumes From Index",
+ "type": "integer"
+ }
+ },
+ "required": [
+ "node_id",
+ "index",
+ "max_iterations",
+ "terminal_task_id",
+ "terminal_is_mapped",
+ "region_id",
+ "trajectory"
+ ],
+ "title": "LoopContext",
+ "type": "object"
+ },
"TIRunContext": {
"description": "Response schema for TaskInstance run context.",
"properties": {
@@ -5270,6 +5291,10 @@
],
"default": null
},
+ "expanded_ti_count": {
+ "title": "Expanded Ti Count",
+ "type": "integer"
+ },
"task_reschedule_count": {
"default": 0,
"title": "Task Reschedule Count",
diff --git
a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
index 0f35e9f0e6f..40da65651ae 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
@@ -40,6 +40,7 @@ def get_bundle() -> VersionBundle:
from airflow.sdk.execution_time.schema.versions.v2026_10_30 import (
AddArgBindingsToSupervisorTIRunContext,
AddDagDefinitionsToDagFileParsingResult,
+ AddExpandedTICountToSupervisorTIRunContext,
AddFirstTaskRescheduleStartDateToSupervisorTIRunContext,
AddLoopRuntime,
AddRegionCoordinatesToSupervisorTaskInstance,
@@ -56,6 +57,7 @@ def get_bundle() -> VersionBundle:
AddDagDefinitionsToDagFileParsingResult,
AddRegionCoordinatesToSupervisorTaskInstance,
AddFirstTaskRescheduleStartDateToSupervisorTIRunContext,
+ AddExpandedTICountToSupervisorTIRunContext,
),
Version("2026-06-16"),
)
diff --git
a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
index c9a80625c3e..d69abf63b5f 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
@@ -119,3 +119,17 @@ class
AddFirstTaskRescheduleStartDateToSupervisorTIRunContext(VersionChange):
instructions_to_migrate_to_previous_version = (
schema(TIRunContext).field("first_task_reschedule_start_date").didnt_exist,
)
+
+
+class AddExpandedTICountToSupervisorTIRunContext(VersionChange):
+ """
+ Add ``expanded_ti_count``, the size of the mapped expansion a task
instance belongs to.
+
+ The supervisor-schema mirror of the execution API's
``AddExpandedTICountToTIRunContext``.
+ """
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version = (
+ schema(TIRunContext).field("expanded_ti_count").didnt_exist,
+ )
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 92a8b1ed0cb..ee78312e95e 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -74,6 +74,7 @@ from airflow.sdk.definitions.asset import (
AssetUriRef,
)
from airflow.sdk.definitions.context import render_template_to_string
+from airflow.sdk.definitions.loop import LoopDecision
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.sdk.definitions.param import process_params
from airflow.sdk.definitions.xcom_arg import PlainXComArg
@@ -267,6 +268,27 @@ class RuntimeTaskInstance(TaskInstance):
if self.region_index is None:
self.region_index = self.map_index
+ @property
+ def loop(self) -> LoopContextAccessor | None:
+ """The iteration of the innermost loop this task instance runs in,
``None`` outside a loop."""
+ if self._ti_context_from_server is None or
self._ti_context_from_server.loop is None:
+ return None
+ return LoopContextAccessor(self._ti_context_from_server.loop, self)
+
+ @property
+ def loop_decision(self) -> LoopDecision | None:
+ """A loop gate's decision, set once the gate has evaluated whether the
loop stops."""
+ if self._loop_decision is None:
+ return None
+ return LoopDecision(stop=self._loop_decision.stop,
reason=self._loop_decision.reason)
+
+ @property
+ def expanded_ti_count(self) -> int | None:
+ """Number of instances in this mapped task's expansion in its loop
iteration, None if not mapped."""
+ if self._ti_context_from_server is None:
+ return None
+ return self._ti_context_from_server.expanded_ti_count
+
@property
def stats_tags(self) -> dict[str, str]:
"""Metric tags for this task instance, including dag tags and
team_name when available."""
@@ -360,8 +382,8 @@ class RuntimeTaskInstance(TaskInstance):
if TYPE_CHECKING:
assert self._cached_template_context is not None
if from_server:
- if from_server.loop is not None:
- self._cached_template_context["loop"] =
LoopContextAccessor(from_server.loop, self)
+ if (loop := self.loop) is not None:
+ self._cached_template_context["loop"] = loop
dag_run = from_server.dag_run
context_from_server: Context = {
# TODO: Assess if we need to pass these through
timezone.coerce_datetime
diff --git a/task-sdk/src/airflow/sdk/types.py
b/task-sdk/src/airflow/sdk/types.py
index eccee7a26a1..83b268780d2 100644
--- a/task-sdk/src/airflow/sdk/types.py
+++ b/task-sdk/src/airflow/sdk/types.py
@@ -51,8 +51,10 @@ if TYPE_CHECKING:
BaseAssetUniqueKey,
)
from airflow.sdk.definitions.context import Context
+ from airflow.sdk.definitions.loop import LoopDecision
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.sdk.execution_time.comms import DagResult
+ from airflow.sdk.execution_time.loop import LoopContextAccessor
Operator: TypeAlias = BaseOperator | MappedOperator
@@ -157,6 +159,15 @@ class RuntimeTaskInstanceProtocol(Protocol):
@property
def stats_tags(self) -> dict[str, str]: ...
+ @property
+ def loop(self) -> LoopContextAccessor | None: ...
+
+ @property
+ def loop_decision(self) -> LoopDecision | None: ...
+
+ @property
+ def expanded_ti_count(self) -> int | None: ...
+
def xcom_pull(
self,
task_ids: str | Iterable[str] | None = None,
diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
index e8c4a16aff2..6e86037fe05 100644
--- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
+++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
@@ -526,13 +526,18 @@ class TestRealBundleArgBindingsDowngrade:
max_iterations=4,
terminal_task_id="body.terminal",
terminal_is_mapped=True,
+ region_id=UUID("0199a8c0-0000-7000-8000-000000000001"),
+ trajectory=1,
)
+ startup_details.ti_context.expanded_ti_count = 2
context = real_migrator.downgrade(startup_details,
version).model_dump()["ti_context"]
if version == "2026-06-16":
assert "loop" not in context
+ assert "expanded_ti_count" not in context
else:
assert context["loop"] ==
startup_details.ti_context.loop.model_dump()
+ assert context["expanded_ti_count"] == 2
@pytest.mark.parametrize("version", ["2026-06-16", "2026-10-30"])
def test_task_coordinates_follow_supervisor_version(self, real_migrator,
startup_details, version):
diff --git a/task-sdk/tests/task_sdk/execution_time/test_loop.py
b/task-sdk/tests/task_sdk/execution_time/test_loop.py
index 77e19c10df5..8dd474773b1 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_loop.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_loop.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import re
+from uuid import UUID
import pytest
@@ -41,6 +42,9 @@ from airflow.sdk.execution_time.comms import (
from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
from airflow.sdk.execution_time.loop import LoopMaxIterationsExceeded
+LOOP_REGION_ID = UUID("0199a8c0-0000-7000-8000-000000000001")
+FORKED_REGION_ID = UUID("0199a8c0-0000-7000-8000-000000000002")
+
@pytest.fixture
def loop_ti(create_runtime_ti):
@@ -58,6 +62,8 @@ def loop_ti(create_runtime_ti):
max_iterations=max_iterations,
terminal_task_id=group.terminal_task_id,
terminal_is_mapped=False,
+ region_id=LOOP_REGION_ID,
+ trajectory=1,
)
return ti
@@ -100,11 +106,59 @@ def
test_body_callable_receives_loop_context_after_unmapping(create_runtime_ti,
max_iterations=4,
terminal_task_id=group.terminal_task_id,
terminal_is_mapped=mapped,
+ region_id=LOOP_REGION_ID,
+ trajectory=1,
)
assert operator.execute(ti.get_template_context()) == (5, 2, 0 if mapped
else -1)
+def test_task_instance_exposes_the_loop_iteration_and_its_trajectory(loop_ti):
+ ti = loop_ti(index=2)
+ ti._ti_context_from_server.loop =
ti._ti_context_from_server.loop.model_copy(
+ update={"trajectory": 2, "forked_from_region_id": FORKED_REGION_ID,
"resumes_from_index": 2}
+ )
+
+ loop = ti.loop
+
+ assert loop is not None
+ assert (loop.index, loop.max_iterations) == (2, 3)
+ assert (loop.region_id, loop.trajectory) == (LOOP_REGION_ID, 2)
+ assert (loop.forked_from_region_id, loop.resumes_from_index) ==
(FORKED_REGION_ID, 2)
+
+
+def test_task_instance_outside_a_loop_has_no_loop(create_runtime_ti):
+ ti = create_runtime_ti(task=BaseOperator(task_id="plain"))
+
+ assert ti.loop is None
+ assert ti.loop_decision is None
+
+
[email protected]("expanded_ti_count", [None, 4])
+def test_task_instance_exposes_the_size_of_its_expansion(create_runtime_ti,
expanded_ti_count):
+ ti = create_runtime_ti(task=BaseOperator(task_id="mapped"))
+ ti._ti_context_from_server.expanded_ti_count = expanded_ti_count
+
+ assert ti.expanded_ti_count == expanded_ti_count
+
+
[email protected](
+ ("index", "until", "expected"),
+ [
+ (1, lambda: LoopDecision(stop=False, reason="not yet"),
LoopDecision(stop=False, reason="not yet")),
+ (3, lambda: True, LoopDecision(stop=True)),
+ (3, lambda: False, LoopDecision(stop=False)),
+ ],
+ ids=["continue", "until_met", "cap_reached"],
+)
+def test_gate_exposes_its_decision(loop_ti, mock_supervisor_comms, index,
until, expected):
+ ti = loop_ti(index=index, until=until)
+
+ task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
+
+ assert ti.loop_decision == expected
+
+
@pytest.mark.parametrize("value", [0, False, [], ""])
def test_previous_and_current_result_preserve_falsey_values(loop_ti,
mock_supervisor_comms, value):
ti = loop_ti(index=2)
diff --git a/ts-sdk/schema/supervisor-schema.json
b/ts-sdk/schema/supervisor-schema.json
index 56a6967f17e..e739708d12d 100644
--- a/ts-sdk/schema/supervisor-schema.json
+++ b/ts-sdk/schema/supervisor-schema.json
@@ -3242,42 +3242,6 @@
"title": "LazyDeserializedDAG",
"type": "object"
},
- "LoopContext": {
- "description": "Pinned loop definition and enclosing iteration for a
task execution.",
- "properties": {
- "node_id": {
- "title": "Node Id",
- "type": "string"
- },
- "index": {
- "minimum": 1,
- "title": "Index",
- "type": "integer"
- },
- "max_iterations": {
- "exclusiveMinimum": 0,
- "title": "Max Iterations",
- "type": "integer"
- },
- "terminal_task_id": {
- "title": "Terminal Task Id",
- "type": "string"
- },
- "terminal_is_mapped": {
- "title": "Terminal Is Mapped",
- "type": "boolean"
- }
- },
- "required": [
- "node_id",
- "index",
- "max_iterations",
- "terminal_task_id",
- "terminal_is_mapped"
- ],
- "title": "LoopContext",
- "type": "object"
- },
"LoopDecisionPayload": {
"additionalProperties": false,
"description": "A loop gate's evaluation of its ``until`` callable.",
@@ -5253,6 +5217,63 @@
"title": "DagRun",
"type": "object"
},
+ "LoopContext": {
+ "description": "Pinned loop definition and enclosing iteration for a
task execution.",
+ "properties": {
+ "node_id": {
+ "title": "Node Id",
+ "type": "string"
+ },
+ "index": {
+ "minimum": 1,
+ "title": "Index",
+ "type": "integer"
+ },
+ "max_iterations": {
+ "exclusiveMinimum": 0,
+ "title": "Max Iterations",
+ "type": "integer"
+ },
+ "terminal_task_id": {
+ "title": "Terminal Task Id",
+ "type": "string"
+ },
+ "terminal_is_mapped": {
+ "title": "Terminal Is Mapped",
+ "type": "boolean"
+ },
+ "region_id": {
+ "format": "uuid",
+ "title": "Region Id",
+ "type": "string"
+ },
+ "trajectory": {
+ "minimum": 1,
+ "title": "Trajectory",
+ "type": "integer"
+ },
+ "forked_from_region_id": {
+ "format": "uuid",
+ "title": "Forked From Region Id",
+ "type": "string"
+ },
+ "resumes_from_index": {
+ "title": "Resumes From Index",
+ "type": "integer"
+ }
+ },
+ "required": [
+ "node_id",
+ "index",
+ "max_iterations",
+ "terminal_task_id",
+ "terminal_is_mapped",
+ "region_id",
+ "trajectory"
+ ],
+ "title": "LoopContext",
+ "type": "object"
+ },
"TIRunContext": {
"description": "Response schema for TaskInstance run context.",
"properties": {
@@ -5270,6 +5291,10 @@
],
"default": null
},
+ "expanded_ti_count": {
+ "title": "Expanded Ti Count",
+ "type": "integer"
+ },
"task_reschedule_count": {
"default": 0,
"title": "Task Reschedule Count",
diff --git a/ts-sdk/src/generated/supervisor.ts
b/ts-sdk/src/generated/supervisor.ts
index 8b018e48514..75a05dde12c 100644
--- a/ts-sdk/src/generated/supervisor.ts
+++ b/ts-sdk/src/generated/supervisor.ts
@@ -236,6 +236,11 @@ export type Index = number;
export type MaxIterations = number;
export type TerminalTaskId = string;
export type TerminalIsMapped = boolean;
+export type RegionId1 = string;
+export type Trajectory = number;
+export type ForkedFromRegionId = string;
+export type ResumesFromIndex = number;
+export type ExpandedTiCount = number;
export type TaskRescheduleCount = number;
export type FirstTaskRescheduleStartDate = string | null;
export type MaxTries = number;
@@ -547,7 +552,7 @@ export type EndDate4 = string | null;
export type State4 = string | null;
export type TryNumber2 = number;
export type MapIndex6 = number | null;
-export type RegionId1 = string | null;
+export type RegionId2 = string | null;
export type RegionIndex1 = number | null;
export type Duration = number | null;
export type Type60 = "PreviousTIResult";
@@ -1051,6 +1056,7 @@ export interface TaskCallbackRequest {
export interface TIRunContext {
dag_run: DagRun;
loop?: LoopContext | null;
+ expanded_ti_count?: ExpandedTiCount;
task_reschedule_count?: TaskRescheduleCount;
first_task_reschedule_start_date?: FirstTaskRescheduleStartDate;
max_tries: MaxTries;
@@ -1076,6 +1082,10 @@ export interface LoopContext {
max_iterations: MaxIterations;
terminal_task_id: TerminalTaskId;
terminal_is_mapped: TerminalIsMapped;
+ region_id: RegionId1;
+ trajectory: Trajectory;
+ forked_from_region_id?: ForkedFromRegionId;
+ resumes_from_index?: ResumesFromIndex;
}
/**
* Variable schema for responses with fields that are needed for Runtime.
@@ -1686,7 +1696,7 @@ export interface PreviousTIResponse {
state?: State4;
try_number: TryNumber2;
map_index?: MapIndex6;
- region_id?: RegionId1;
+ region_id?: RegionId2;
region_index?: RegionIndex1;
duration?: Duration;
}
