From 82544a1a3080e915386a752ec1c79f7f8673abeb Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Mon, 19 Oct 2026 19:25:53 +0200 Subject: [PATCH 1/5] Expose task instance loop and expansion position to workers and listeners 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 --- .../docs/authoring-and-scheduling/loops.rst | 35 +++++ .../execution_api/datamodels/taskinstance.py | 29 ++++ .../execution_api/routes/task_instances.py | 20 ++- .../execution_api/versions/__init__.py | 2 + .../execution_api/versions/v2026_10_30.py | 14 ++ .../src/airflow/models/dynamic_region.py | 19 ++- .../src/airflow/models/task_coordinates.py | 108 +++++++++++++- .../src/airflow/utils/log/task_log_address.py | 10 +- .../versions/head/test_task_instances.py | 53 +++++++ .../tests/unit/models/test_dynamic_region.py | 15 ++ .../unit/models/test_task_coordinates.py | 51 +++++++ go-sdk/pkg/execution/genmodels/models.gen.go | 138 ++++++++++-------- go-sdk/schema/supervisor-schema.json | 115 ++++++++++----- .../airflow/sdk/api/datamodels/_generated.py | 5 + .../src/airflow/sdk/execution_time/loop.py | 22 +++ .../sdk/execution_time/schema/schema.json | 115 ++++++++++----- .../schema/versions/__init__.py | 2 + .../schema/versions/v2026_10_30.py | 14 ++ .../airflow/sdk/execution_time/task_runner.py | 22 +++ task-sdk/src/airflow/sdk/types.py | 11 ++ .../execution_time/schema/test_migrator.py | 4 + .../task_sdk/execution_time/test_loop.py | 52 +++++++ ts-sdk/schema/supervisor-schema.json | 115 ++++++++++----- ts-sdk/src/generated/supervisor.ts | 14 +- 24 files changed, 797 insertions(+), 188 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst index 7a6e04e0c5988..f1bf5a0d205a9 100644 --- a/airflow-core/docs/authoring-and-scheduling/loops.rst +++ b/airflow-core/docs/authoring-and-scheduling/loops.rst @@ -152,9 +152,23 @@ 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.group_id`` + - The group id of the loop. + * - ``loop.region_id`` + - The id of the loop execution that holds this iteration. It is the same for every task instance of + the iteration, mapped ones included. + * - ``loop.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 execution this one replaced, or ``None`` in trajectory 1. + * - ``loop.resumes_from_index`` + - The first iteration this execution holds, 0 in trajectory 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 +361,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 +434,18 @@ 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 old one from the +iteration after the rerun gate. Each execution is a trajectory: ``loop.trajectory`` is 1 for the +first one and grows 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 the start again. 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 09ce1193e9018..c12a83e291a86 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 @@ -457,6 +457,28 @@ class LoopContext(BaseModel): terminal_task_id: str terminal_is_mapped: bool + region_id: uuid.UUID + """ + The loop region holding this iteration. + + 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 pass. + """ + + trajectory: int = Field(default=1, ge=1) + """ + 1 for the loop's original region, and one more for each fork of it. + + Clearing a loop with its later iterations forks the loop region, and the iterations from + ``resumes_from_index`` on run again in the fork. + """ + + forked_from_region_id: uuid.UUID | None = None + """The loop region this fork replaces from ``resumes_from_index`` on, ``None`` for an original region.""" + + resumes_from_index: int = Field(default=0, ge=0) + """The first iteration ``region_id`` holds, 0 for an original region.""" + class TIRunContext(BaseModel): """Response schema for TaskInstance run context.""" @@ -466,6 +488,13 @@ class TIRunContext(BaseModel): loop: LoopContext | None = None + expanded_ti_count: int | None = 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 c27e2f74f6a74..80f199ef66c2a 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 @@ -372,17 +372,23 @@ def ti_run( multi_team=conf.getboolean("core", "multi_team"), ) resolver = TaskCoordinateResolver(dag_bag, session) - if loop_context := resolver.loop_context(ti): - group, index = loop_context + position = resolver.get_execution_position(ti) + if position.expanded_ti_count is not None: + context.expanded_ti_count = position.expanded_ti_count + if (loop := position.loop) is not None: terminal = resolver.get_task( - ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id + ti.dag_id, ti.run_id, loop.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, + node_id=loop.group.node_id, + index=loop.index, + max_iterations=loop.group.max_iterations, + terminal_task_id=loop.group.terminal_task_id, terminal_is_mapped=terminal.get_needs_expansion(), + region_id=loop.region_id, + trajectory=loop.trajectory, + forked_from_region_id=loop.forked_from_region_id, + resumes_from_index=loop.resumes_from_index, ) 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 3df478de386d1..e85afc58b9bdd 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 @@ AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, AddDagRunNoteUpdateEndpoint, + AddExpandedTICountToTIRunContext, AddFirstTaskRescheduleStartDateField, AddLoopContext, AddLoopDecisionField, @@ -78,6 +79,7 @@ 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 26d0aea88781e..1f7780ee32171 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 @@ 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.""" diff --git a/airflow-core/src/airflow/models/dynamic_region.py b/airflow-core/src/airflow/models/dynamic_region.py index 968b92248292f..63a6574b26549 100644 --- a/airflow-core/src/airflow/models/dynamic_region.py +++ b/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]: diff --git a/airflow-core/src/airflow/models/task_coordinates.py b/airflow-core/src/airflow/models/task_coordinates.py index 73c2f87ed8df5..53b72e740dd89 100644 --- a/airflow-core/src/airflow/models/task_coordinates.py +++ b/airflow-core/src/airflow/models/task_coordinates.py @@ -21,7 +21,7 @@ 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.exceptions import TaskNotFound from airflow.models.dagbag import DBDagBag @@ -31,6 +31,7 @@ AmbiguousProducerError, DynamicRegion, ProducerContext, + get_trajectory, load_region_ancestry, loop_position, resolve_current_producers, @@ -40,6 +41,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 @@ -69,6 +71,40 @@ class _ProducerRequest: loop_node_id: str | None = None +@attrs.frozen(kw_only=True) +class LoopPosition: + """ + The pass of the innermost loop a task instance runs in. + + ``region_id`` is the loop region that holds the pass, also for a mapped task, whose own ``region_id`` is + its expansion's. A clear that reruns later passes forks the loop region: ``trajectory`` is 1 for the + original region and grows by one with each fork, ``forked_from_region_id`` is the region the fork + replaced and ``resumes_from_index`` the first pass the fork owns (0 for an original region). + """ + + group: SerializedLoopTaskGroup + index: int + region_id: UUID + trajectory: int + forked_from_region_id: UUID | None + resumes_from_index: int + + +@attrs.frozen(kw_only=True) +class TaskInstanceExecutionPosition: + """ + Where a task instance runs, as the worker gets it when the task starts. + + ``map_index`` is the map index users see, -1 for a task that isn't mapped. ``expanded_ti_count`` is the + number of live instances of a mapped task in its own expansion, and ``None`` for a task that isn't mapped. + ``loop`` is ``None`` outside a loop. + """ + + map_index: int + expanded_ti_count: int | None + loop: LoopPosition | None + + class TaskCoordinate(Protocol): """Stored task identity shared by live and retained task data.""" @@ -231,6 +267,14 @@ def get_loop_iteration(self, ti: TaskCoordinate) -> tuple[str, int] | None: def loop_context( self, ti: TaskCoordinate, *, task: SerializedOperator | None = None ) -> tuple[SerializedLoopTaskGroup, int] | None: + if (position := self.get_loop_position(ti, task=task)) is None: + return None + return position.group, position.index + + def get_loop_position( + self, ti: TaskCoordinate, *, task: SerializedOperator | None = None + ) -> LoopPosition | None: + """Return the pass of the innermost loop ``ti`` runs in, or None outside any loop.""" if ti.region_id == SENTINEL_REGION_ID: return None if task is None: @@ -244,9 +288,52 @@ def loop_context( 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: + coordinate = DynamicRegion.find_physical_coordinate( + ti.region_id, ti.region_index, group.node_id, regions + ) + if position is None or coordinate is None: raise ValueError("Task coordinates do not belong to the pinned loop") - return group, position[1] + region = regions[coordinate[0]] + return LoopPosition( + group=group, + index=position[1], + region_id=region.id, + trajectory=get_trajectory(region, regions), + forked_from_region_id=region.forked_from_region_id, + resumes_from_index=region.resumes_from_index, + ) + + 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.""" + 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_execution_position(self, ti: TaskCoordinate) -> TaskInstanceExecutionPosition: + """Place ``ti`` in its mapped expansion and loop pass from its pinned task definition.""" + if ti.region_id == SENTINEL_REGION_ID: + return TaskInstanceExecutionPosition(map_index=ti.region_index, expanded_ti_count=None, loop=None) + task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id) + return TaskInstanceExecutionPosition( + map_index=ti.region_index if task.get_needs_expansion() else -1, + expanded_ti_count=self.get_expanded_ti_count(ti, task=task), + loop=self.get_loop_position(ti, task=task), + ) def get_task( self, dag_id: str, run_id: str, task_id: str, *, dag_version_id: UUID | None = None @@ -839,3 +926,18 @@ def select_producer_ids( region_index=request.region_index, session=self.session, ) + + +def get_task_instance_execution_position( + task_instance: TaskInstance, *, session: Session +) -> TaskInstanceExecutionPosition: + """ + Place a stored task instance in its mapped expansion and loop pass. + + This is what the Execution API sends a worker when the task starts (``TIRunContext.loop`` and + ``TIRunContext.expanded_ti_count``), so server-side consumers such as listeners reporting a manual state + change describe the task instance as the worker would, without reading region rows themselves. + """ + return TaskCoordinateResolver(DBDagBag(load_op_links=False), session).get_execution_position( + task_instance + ) 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 b72eaab72ecaf..27be654a2741c 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 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 6a81588ffc555..f9825d9044bed 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 @@ def body(): 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,59 @@ def body(): "max_iterations": 4, "terminal_task_id": loop.terminal_task_id, "terminal_is_mapped": mapped, + "region_id": str(parent.id), + "trajectory": 1, + "forked_from_region_id": None, + "resumes_from_index": 0, } + 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 573046e4e9b6e..e34ba8ee2b048 100644 --- a/airflow-core/tests/unit/models/test_dynamic_region.py +++ b/airflow-core/tests/unit/models/test_dynamic_region.py @@ -36,6 +36,7 @@ ProducerContext, build_loop_family_legs, build_loop_sequence_order, + get_trajectory, resolve_current_producers, select_loop_producer_ids, ) @@ -1114,3 +1115,17 @@ def test_loop_sequence_of_mapped_task_skips_passes_a_fork_replaced(loop_run, ses (True, 4, 0), (True, 4, 1), ] + + +@pytest.mark.parametrize("forks", [0, 1, 2]) +def test_trajectory_counts_the_forks_a_region_descends_from(forks): + regions = {} + region = DynamicRegion(id=uuid4(), dag_id="dag", run_id="run", node_id="loop") + regions[region.id] = region + for _ in range(forks): + region = DynamicRegion( + id=uuid4(), dag_id="dag", run_id="run", node_id="loop", 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 b74f38460a24d..1fa67bbd785ea 100644 --- a/airflow-core/tests/unit/models/test_task_coordinates.py +++ b/airflow-core/tests/unit/models/test_task_coordinates.py @@ -33,6 +33,7 @@ TaskCoordinateResolver, build_coordinate_filters, get_public_map_index, + get_task_instance_execution_position, is_plain_expansion, ) from airflow.models.taskinstance import TaskInstance @@ -105,6 +106,9 @@ def body(): assert resolver.loop_iterations(ti) == [("body", 3)] assert resolver.public_map_index(ti) == 1 + position = resolver.get_loop_position(ti) + assert position is not None + assert (position.index, position.region_id, position.trajectory) == (3, loop_region.id, 1) @pytest.mark.parametrize("run_count", [1, 3]) @@ -165,6 +169,53 @@ def test_loop_context_reuses_ancestry_loaded_by_loop_iterations(loop_coordinates assert (group.group_id, index) == ("body", 1) +def test_loop_position_describes_the_region_holding_the_iteration(forked_loop_coordinates, session): + _, producer, consumer, _, outside = forked_loop_coordinates + resolver = TaskCoordinateResolver(DBDagBag(), session) + + original = resolver.get_loop_position(producer) + fork = resolver.get_loop_position(consumer) + + assert original is not None + assert fork is not None + assert (original.region_id, original.trajectory, original.forked_from_region_id) == ( + producer.region_id, + 1, + None, + ) + assert original.resumes_from_index == 0 + assert (fork.region_id, fork.trajectory, fork.forked_from_region_id) == ( + consumer.region_id, + 2, + producer.region_id, + ) + assert fork.resumes_from_index == 2 + assert resolver.get_loop_position(outside) is None + + +def test_execution_position_of_a_stored_task_instance(loop_coordinates, dag_maker, session): + with dag_maker("mapped-execution-position", 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() + + in_loop = get_task_instance_execution_position(consumer, session=session) + plain = get_task_instance_execution_position(outside, session=session) + expanded = get_task_instance_execution_position(mapped[1], session=session) + + assert (in_loop.map_index, in_loop.expanded_ti_count) == (-1, None) + assert in_loop.loop is not None + assert (in_loop.loop.group.group_id, in_loop.loop.index, in_loop.loop.region_id) == ( + "body", + 2, + consumer.region_id, + ) + assert (plain.map_index, plain.expanded_ti_count, plain.loop) == (-1, None, None) + assert (expanded.map_index, expanded.expanded_ti_count, expanded.loop) == (1, 2, None) + + 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 0a5dba694e675..6eebf1bdb1e41 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 interface{} `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,omitempty"` } // 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 interface{} `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"` @@ -1734,21 +1750,53 @@ type TIRunContext struct { type TaskArgBinding interface{} -type TriggerKwargs map[string]interface{} +type TaskBreadcrumbsResult struct { + // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". + Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` -type Warnings []interface{} + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` +} -type TaskIds []string +type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` +// Task callback status information. +// +// A Class with information about the success/failure TI callback to be executed. +// Currently, only failure +// callbacks when tasks are externally killed or experience heartbeat timeouts are +// run via DagFileProcessorProcess. +type TaskCallbackRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` + // BundleVersion corresponds to the JSON schema field "bundle_version". + BundleVersion interface{} `msgpack:"bundle_version"` + + // ContextFromServer corresponds to the JSON schema field "context_from_server". + ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` + + // Filepath corresponds to the JSON schema field "filepath". + Filepath string `msgpack:"filepath"` + + // Msg corresponds to the JSON schema field "msg". + Msg interface{} `msgpack:"msg,omitempty"` + + // TaskCallbackType corresponds to the JSON schema field "task_callback_type". + TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` + + // TI corresponds to the JSON schema field "ti". + TI TaskInstance `msgpack:"ti"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` + + // VersionData corresponds to the JSON schema field "version_data". + VersionData *VersionData `msgpack:"version_data,omitempty"` } +type TaskIds []string + // Schema for TaskInstance model with minimal required fields needed for Runtime. type TaskInstance struct { // ContextCarrier corresponds to the JSON schema field "context_carrier". @@ -1788,69 +1836,41 @@ type TaskInstance struct { TryNumber int `msgpack:"try_number"` } +type TaskInstanceState string + const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" const TaskInstanceStateDeferred TaskInstanceState = "deferred" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" const TaskInstanceStateFailed TaskInstanceState = "failed" -const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateRunning TaskInstanceState = "running" const TaskInstanceStateQueued TaskInstanceState = "queued" - -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" const TaskInstanceStateRemoved TaskInstanceState = "removed" -type TaskInstanceState string - -type TaskOutlets []AssetProfile - -type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} - -type TaskBreadcrumbsResult struct { - // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". - Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` - - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` -} - -// Task callback status information. -// -// A Class with information about the success/failure TI callback to be executed. -// Currently, only failure -// callbacks when tasks are externally killed or experience heartbeat timeouts are -// run via DagFileProcessorProcess. -type TaskCallbackRequest struct { - // BundleName corresponds to the JSON schema field "bundle_name". - BundleName string `msgpack:"bundle_name"` +type VersionData map[string]interface{} - // BundleVersion corresponds to the JSON schema field "bundle_version". - BundleVersion interface{} `msgpack:"bundle_version"` +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" +const TaskInstanceStateRunning TaskInstanceState = "running" +const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" - // ContextFromServer corresponds to the JSON schema field "context_from_server". - ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" - // Filepath corresponds to the JSON schema field "filepath". - Filepath string `msgpack:"filepath"` +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" - // Msg corresponds to the JSON schema field "msg". - Msg interface{} `msgpack:"msg,omitempty"` +type TaskOutlets []AssetProfile - // TaskCallbackType corresponds to the JSON schema field "task_callback_type". - TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` - // TI corresponds to the JSON schema field "ti". - TI TaskInstance `msgpack:"ti"` + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` +type TriggerKwargs map[string]interface{} - // VersionData corresponds to the JSON schema field "version_data". - VersionData *VersionData `msgpack:"version_data,omitempty"` -} +type Warnings []interface{} // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { @@ -1953,8 +1973,6 @@ type UpdateDagRunNote struct { Type string `msgpack:"type,omitempty"` } -type VersionData map[string]interface{} - // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { // ChosenOptions corresponds to the JSON schema field "chosen_options". diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index 56a6967f17e77..8fb379f1b6589 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,73 @@ "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": { + "default": 1, + "minimum": 1, + "title": "Trajectory", + "type": "integer" + }, + "forked_from_region_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Forked From Region Id" + }, + "resumes_from_index": { + "default": 0, + "minimum": 0, + "title": "Resumes From Index", + "type": "integer" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped", + "region_id" + ], + "title": "LoopContext", + "type": "object" + }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { @@ -5270,6 +5301,18 @@ ], "default": null }, + "expanded_ti_count": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Expanded Ti Count" + }, "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 0e677c93e7a94..011ff4e139bb3 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 | None, Field(ge=1, title="Trajectory")] = 1 + forked_from_region_id: Annotated[UUID | None, Field(title="Forked From Region Id")] = None + resumes_from_index: Annotated[int | None, Field(ge=0, title="Resumes From Index")] = 0 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 f8aac86119154..f7adb18360395 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.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,26 @@ def index(self) -> int: def max_iterations(self) -> int: return self._descriptor.max_iterations + @property + def group_id(self) -> str: + return self._descriptor.node_id + + @property + def region_id(self) -> UUID: + return self._descriptor.region_id + + @property + def trajectory(self) -> int: + return cast("int", 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: + return cast("int", 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 56a6967f17e77..8fb379f1b6589 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,73 @@ "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": { + "default": 1, + "minimum": 1, + "title": "Trajectory", + "type": "integer" + }, + "forked_from_region_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Forked From Region Id" + }, + "resumes_from_index": { + "default": 0, + "minimum": 0, + "title": "Resumes From Index", + "type": "integer" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped", + "region_id" + ], + "title": "LoopContext", + "type": "object" + }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { @@ -5270,6 +5301,18 @@ ], "default": null }, + "expanded_ti_count": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Expanded Ti Count" + }, "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 0f35e9f0e6f2b..40da65651ae83 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 c9a80625c3ea0..d69abf63b5f34 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 92a8b1ed0cb83..c4d5310bf6e5e 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 @@ 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 @@ def model_post_init(self, context: Any) -> None: 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, ``None`` for a task that isn't 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.""" diff --git a/task-sdk/src/airflow/sdk/types.py b/task-sdk/src/airflow/sdk/types.py index eccee7a26a12b..83b268780d2c5 100644 --- a/task-sdk/src/airflow/sdk/types.py +++ b/task-sdk/src/airflow/sdk/types.py @@ -51,8 +51,10 @@ 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 @@ def mark_success_url(self) -> str: ... @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 e8c4a16aff248..9ab934e2cb9c9 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,17 @@ def test_loop_context_follows_supervisor_version(self, real_migrator, startup_de max_iterations=4, terminal_task_id="body.terminal", terminal_is_mapped=True, + region_id=UUID("0199a8c0-0000-7000-8000-000000000001"), ) + 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 77e19c10df5b4..cfcf76797de30 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.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,7 @@ def body(): max_iterations=max_iterations, terminal_task_id=group.terminal_task_id, terminal_is_mapped=False, + region_id=LOOP_REGION_ID, ) return ti @@ -100,11 +105,58 @@ def body(): max_iterations=4, terminal_task_id=group.terminal_task_id, terminal_is_mapped=mapped, + region_id=LOOP_REGION_ID, ) 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.group_id, loop.index, loop.max_iterations) == ("body", 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 + + +@pytest.mark.parametrize("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 + + +@pytest.mark.parametrize( + ("index", "until", "expected"), + [ + (0, lambda: LoopDecision(stop=False, reason="not yet"), LoopDecision(stop=False, reason="not yet")), + (2, lambda: True, LoopDecision(stop=True)), + (2, 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 56a6967f17e77..8fb379f1b6589 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,73 @@ "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": { + "default": 1, + "minimum": 1, + "title": "Trajectory", + "type": "integer" + }, + "forked_from_region_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Forked From Region Id" + }, + "resumes_from_index": { + "default": 0, + "minimum": 0, + "title": "Resumes From Index", + "type": "integer" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped", + "region_id" + ], + "title": "LoopContext", + "type": "object" + }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { @@ -5270,6 +5301,18 @@ ], "default": null }, + "expanded_ti_count": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Expanded Ti Count" + }, "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 8b018e485149d..1e31f770f664f 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 | null; +export type ResumesFromIndex = number; +export type ExpandedTiCount = number | null; 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; } From d9981434364900d2b1d969258e20a658a36fe775 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Tue, 20 Oct 2026 09:52:13 +0200 Subject: [PATCH 2/5] Address review of task instance loop and expansion position Co-Authored-By: Claude --- .../docs/authoring-and-scheduling/loops.rst | 20 ++- .../execution_api/datamodels/taskinstance.py | 21 ++- .../execution_api/routes/task_instances.py | 41 ++--- .../src/airflow/models/task_coordinates.py | 80 ++++------ .../versions/head/test_task_instances.py | 2 - .../tests/unit/models/test_dynamic_region.py | 14 +- .../unit/models/test_task_coordinates.py | 26 ++-- go-sdk/pkg/execution/genmodels/models.gen.go | 30 ++-- go-sdk/schema/supervisor-schema.json | 140 +++++++++--------- .../airflow/sdk/api/datamodels/_generated.py | 4 +- .../src/airflow/sdk/execution_time/loop.py | 10 +- .../sdk/execution_time/schema/schema.json | 140 +++++++++--------- .../airflow/sdk/execution_time/task_runner.py | 6 +- .../execution_time/schema/test_migrator.py | 1 + .../task_sdk/execution_time/test_loop.py | 4 +- ts-sdk/schema/supervisor-schema.json | 140 +++++++++--------- ts-sdk/src/generated/supervisor.ts | 4 +- 17 files changed, 330 insertions(+), 353 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst index f1bf5a0d205a9..9893175c5a9b2 100644 --- a/airflow-core/docs/authoring-and-scheduling/loops.rst +++ b/airflow-core/docs/authoring-and-scheduling/loops.rst @@ -152,18 +152,16 @@ 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.group_id`` - - The group id of the loop. * - ``loop.region_id`` - - The id of the loop execution that holds this iteration. It is the same for every task instance of - the iteration, mapped ones included. + - 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`` - - 1 for the loop's first execution, and one more each time a rerun replaces later iterations, see - :ref:`loops-trajectories`. + - 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 execution this one replaced, or ``None`` in trajectory 1. + - The ``loop.region_id`` of the previous trajectory, or ``None`` in trajectory 1. * - ``loop.resumes_from_index`` - - The first iteration this execution holds, 0 in trajectory 1. + - The first iteration this trajectory holds, or ``None`` in trajectory 1, which holds them from 0. 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``. @@ -440,9 +438,9 @@ Tasks you leave uncleared are not recomputed using the results of the rerun. Reruns and trajectories ----------------------- -Clearing later iterations starts a new execution of the loop that replaces the old one from the -iteration after the rerun gate. Each execution is a trajectory: ``loop.trajectory`` is 1 for the -first one and grows by one with each rerun. The rerun gate, and whatever you clear in its +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 the start again. In the 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 c12a83e291a86..dd1fc0c90b5c9 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 @@ -459,25 +459,30 @@ class LoopContext(BaseModel): region_id: uuid.UUID """ - The loop region holding this iteration. + 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 pass. + unmapped tasks of the same iteration. """ - trajectory: int = Field(default=1, ge=1) + trajectory: int = Field(ge=1) """ - 1 for the loop's original region, and one more for each fork of it. + 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 with its later iterations forks the loop region, and the iterations from + 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: uuid.UUID | None = None - """The loop region this fork replaces from ``resumes_from_index`` on, ``None`` for an original region.""" + """ + 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: int = Field(default=0, ge=0) - """The first iteration ``region_id`` holds, 0 for an original region.""" + resumes_from_index: int | None = None + """The first iteration ``region_id`` holds. Absent for the original region, which holds them from 0.""" class TIRunContext(BaseModel): 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 80f199ef66c2a..e7472245d6af0 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 @@ -371,25 +371,28 @@ 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) - position = resolver.get_execution_position(ti) - if position.expanded_ti_count is not None: - context.expanded_ti_count = position.expanded_ti_count - if (loop := position.loop) is not None: - terminal = resolver.get_task( - ti.dag_id, ti.run_id, loop.group.terminal_task_id, dag_version_id=ti.dag_version_id - ) - context.loop = LoopContext( - node_id=loop.group.node_id, - index=loop.index, - max_iterations=loop.group.max_iterations, - terminal_task_id=loop.group.terminal_task_id, - terminal_is_mapped=terminal.get_needs_expansion(), - region_id=loop.region_id, - trajectory=loop.trajectory, - forked_from_region_id=loop.forked_from_region_id, - resumes_from_index=loop.resumes_from_index, - ) + 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_position(ti, task=task)) is not None: + terminal = resolver.get_task( + ti.dag_id, ti.run_id, loop.group.terminal_task_id, dag_version_id=ti.dag_version_id + ) + context.loop = LoopContext( + node_id=loop.group.node_id, + index=loop.index, + max_iterations=loop.group.max_iterations, + terminal_task_id=loop.group.terminal_task_id, + terminal_is_mapped=terminal.get_needs_expansion(), + region_id=loop.region_id, + trajectory=loop.trajectory, + ) + # Set only for a fork, as the route leaves out unset fields + if loop.forked_from_region_id is not None: + context.loop.forked_from_region_id = loop.forked_from_region_id + context.loop.resumes_from_index = loop.resumes_from_index 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/models/task_coordinates.py b/airflow-core/src/airflow/models/task_coordinates.py index 53b72e740dd89..333025dba304f 100644 --- a/airflow-core/src/airflow/models/task_coordinates.py +++ b/airflow-core/src/airflow/models/task_coordinates.py @@ -74,12 +74,13 @@ class _ProducerRequest: @attrs.frozen(kw_only=True) class LoopPosition: """ - The pass of the innermost loop a task instance runs in. + The iteration of the innermost loop a task instance runs in. - ``region_id`` is the loop region that holds the pass, also for a mapped task, whose own ``region_id`` is - its expansion's. A clear that reruns later passes forks the loop region: ``trajectory`` is 1 for the - original region and grows by one with each fork, ``forked_from_region_id`` is the region the fork - replaced and ``resumes_from_index`` the first pass the fork owns (0 for an original region). + ``region_id`` is the loop region that holds the iteration, also for a mapped task, whose own ``region_id`` + is its expansion's, and identifies the trajectory. A clear that reruns later iterations forks the newest + loop region: ``trajectory`` numbers the regions of that chain from 1, ``forked_from_region_id`` is the + previous region of the chain and ``resumes_from_index`` the first iteration the fork holds. Both are + ``None`` for the original region. """ group: SerializedLoopTaskGroup @@ -87,22 +88,7 @@ class LoopPosition: region_id: UUID trajectory: int forked_from_region_id: UUID | None - resumes_from_index: int - - -@attrs.frozen(kw_only=True) -class TaskInstanceExecutionPosition: - """ - Where a task instance runs, as the worker gets it when the task starts. - - ``map_index`` is the map index users see, -1 for a task that isn't mapped. ``expanded_ti_count`` is the - number of live instances of a mapped task in its own expansion, and ``None`` for a task that isn't mapped. - ``loop`` is ``None`` outside a loop. - """ - - map_index: int - expanded_ti_count: int | None - loop: LoopPosition | None + resumes_from_index: int | None class TaskCoordinate(Protocol): @@ -274,7 +260,12 @@ def loop_context( def get_loop_position( self, ti: TaskCoordinate, *, task: SerializedOperator | None = None ) -> LoopPosition | None: - """Return the pass of the innermost loop ``ti`` runs in, or None outside any loop.""" + """ + Return the iteration of the innermost loop ``ti`` runs in, or None outside any loop. + + This is what a worker gets as ``TIRunContext.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: @@ -287,26 +278,31 @@ def get_loop_position( 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) coordinate = DynamicRegion.find_physical_coordinate( ti.region_id, ti.region_index, group.node_id, regions ) - if position is None or coordinate is None: + if coordinate is None: raise ValueError("Task coordinates do not belong to the pinned loop") - region = regions[coordinate[0]] + region_id, index = coordinate + region = regions[region_id] + is_fork = region.forked_from_region_id is not None return LoopPosition( group=group, - index=position[1], - region_id=region.id, + index=index, + region_id=region_id, trajectory=get_trajectory(region, regions), forked_from_region_id=region.forked_from_region_id, - resumes_from_index=region.resumes_from_index, + resumes_from_index=region.resumes_from_index if is_fork else None, ) 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.""" + """ + 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: @@ -324,17 +320,6 @@ def get_expanded_ti_count( ) ) - def get_execution_position(self, ti: TaskCoordinate) -> TaskInstanceExecutionPosition: - """Place ``ti`` in its mapped expansion and loop pass from its pinned task definition.""" - if ti.region_id == SENTINEL_REGION_ID: - return TaskInstanceExecutionPosition(map_index=ti.region_index, expanded_ti_count=None, loop=None) - task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id) - return TaskInstanceExecutionPosition( - map_index=ti.region_index if task.get_needs_expansion() else -1, - expanded_ti_count=self.get_expanded_ti_count(ti, task=task), - loop=self.get_loop_position(ti, task=task), - ) - def get_task( self, dag_id: str, run_id: str, task_id: str, *, dag_version_id: UUID | None = None ) -> SerializedOperator: @@ -926,18 +911,3 @@ def select_producer_ids( region_index=request.region_index, session=self.session, ) - - -def get_task_instance_execution_position( - task_instance: TaskInstance, *, session: Session -) -> TaskInstanceExecutionPosition: - """ - Place a stored task instance in its mapped expansion and loop pass. - - This is what the Execution API sends a worker when the task starts (``TIRunContext.loop`` and - ``TIRunContext.expanded_ti_count``), so server-side consumers such as listeners reporting a manual state - change describe the task instance as the worker would, without reading region rows themselves. - """ - return TaskCoordinateResolver(DBDagBag(load_op_links=False), session).get_execution_position( - task_instance - ) 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 f9825d9044bed..d241600d281da 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 @@ -667,8 +667,6 @@ def body(): "terminal_is_mapped": mapped, "region_id": str(parent.id), "trajectory": 1, - "forked_from_region_id": None, - "resumes_from_index": 0, } assert response.json().get("expanded_ti_count") == (1 if mapped else None) diff --git a/airflow-core/tests/unit/models/test_dynamic_region.py b/airflow-core/tests/unit/models/test_dynamic_region.py index e34ba8ee2b048..591cbbadd2442 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 @@ -53,9 +52,6 @@ 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 @@ -1119,13 +1115,11 @@ def test_loop_sequence_of_mapped_task_skips_passes_a_fork_replaced(loop_run, ses @pytest.mark.parametrize("forks", [0, 1, 2]) def test_trajectory_counts_the_forks_a_region_descends_from(forks): - regions = {} - region = DynamicRegion(id=uuid4(), dag_id="dag", run_id="run", node_id="loop") - regions[region.id] = region + 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 = DynamicRegion( - id=uuid4(), dag_id="dag", run_id="run", node_id="loop", forked_from_region_id=region.id - ) + 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 1fa67bbd785ea..de1213e3d9ed0 100644 --- a/airflow-core/tests/unit/models/test_task_coordinates.py +++ b/airflow-core/tests/unit/models/test_task_coordinates.py @@ -33,7 +33,6 @@ TaskCoordinateResolver, build_coordinate_filters, get_public_map_index, - get_task_instance_execution_position, is_plain_expansion, ) from airflow.models.taskinstance import TaskInstance @@ -183,7 +182,7 @@ def test_loop_position_describes_the_region_holding_the_iteration(forked_loop_co 1, None, ) - assert original.resumes_from_index == 0 + assert original.resumes_from_index is None assert (fork.region_id, fork.trajectory, fork.forked_from_region_id) == ( consumer.region_id, 2, @@ -193,27 +192,20 @@ def test_loop_position_describes_the_region_holding_the_iteration(forked_loop_co assert resolver.get_loop_position(outside) is None -def test_execution_position_of_a_stored_task_instance(loop_coordinates, dag_maker, session): - with dag_maker("mapped-execution-position", serialized=True): +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]) - in_loop = get_task_instance_execution_position(consumer, session=session) - plain = get_task_instance_execution_position(outside, session=session) - expanded = get_task_instance_execution_position(mapped[1], session=session) - - assert (in_loop.map_index, in_loop.expanded_ti_count) == (-1, None) - assert in_loop.loop is not None - assert (in_loop.loop.group.group_id, in_loop.loop.index, in_loop.loop.region_id) == ( - "body", - 2, - consumer.region_id, - ) - assert (plain.map_index, plain.expanded_ti_count, plain.loop) == (-1, None, None) - assert (expanded.map_index, expanded.expanded_ti_count, expanded.loop) == (1, 2, None) + 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_removed_task_coordinates_degrade_to_stored_region_data(dag_maker, session): diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 6eebf1bdb1e41..004ccb0d518eb 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1337,7 +1337,7 @@ type LoopContext struct { RegionID string `msgpack:"region_id"` // ResumesFromIndex corresponds to the JSON schema field "resumes_from_index". - ResumesFromIndex int `msgpack:"resumes_from_index,omitempty"` + ResumesFromIndex interface{} `msgpack:"resumes_from_index,omitempty"` // TerminalIsMapped corresponds to the JSON schema field "terminal_is_mapped". TerminalIsMapped bool `msgpack:"terminal_is_mapped"` @@ -1346,7 +1346,7 @@ type LoopContext struct { TerminalTaskID string `msgpack:"terminal_task_id"` // Trajectory corresponds to the JSON schema field "trajectory". - Trajectory *int `msgpack:"trajectory,omitempty"` + Trajectory int `msgpack:"trajectory"` } // A loop gate's evaluation of its “until“ callable. @@ -1844,29 +1844,29 @@ const TaskInstanceStateFailed TaskInstanceState = "failed" const TaskInstanceStateQueued TaskInstanceState = "queued" const TaskInstanceStateRemoved TaskInstanceState = "removed" -type VersionData map[string]interface{} +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} const TaskInstanceStateScheduled TaskInstanceState = "scheduled" const TaskInstanceStateRunning TaskInstanceState = "running" const TaskInstanceStateSuccess TaskInstanceState = "success" const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +type VersionData map[string]interface{} -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" type TaskOutlets []AssetProfile -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` - - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` -} +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" type TriggerKwargs map[string]interface{} diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index 8fb379f1b6589..b10b3e7633470 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -3242,6 +3242,79 @@ "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" + }, + "region_id": { + "format": "uuid", + "title": "Region Id", + "type": "string" + }, + "trajectory": { + "minimum": 1, + "title": "Trajectory", + "type": "integer" + }, + "forked_from_region_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Forked From Region Id" + }, + "resumes_from_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Resumes From Index" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped", + "region_id", + "trajectory" + ], + "title": "LoopContext", + "type": "object" + }, "LoopDecisionPayload": { "additionalProperties": false, "description": "A loop gate's evaluation of its ``until`` callable.", @@ -5217,73 +5290,6 @@ "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": { - "default": 1, - "minimum": 1, - "title": "Trajectory", - "type": "integer" - }, - "forked_from_region_id": { - "anyOf": [ - { - "format": "uuid", - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Forked From Region Id" - }, - "resumes_from_index": { - "default": 0, - "minimum": 0, - "title": "Resumes From Index", - "type": "integer" - } - }, - "required": [ - "node_id", - "index", - "max_iterations", - "terminal_task_id", - "terminal_is_mapped", - "region_id" - ], - "title": "LoopContext", - "type": "object" - }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 011ff4e139bb3..6c4b2536fc44c 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -343,9 +343,9 @@ class LoopContext(BaseModel): 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 | None, Field(ge=1, title="Trajectory")] = 1 + 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(ge=0, title="Resumes From Index")] = 0 + resumes_from_index: Annotated[int | None, Field(title="Resumes From Index")] = None class LoopDecisionPayload(BaseModel): diff --git a/task-sdk/src/airflow/sdk/execution_time/loop.py b/task-sdk/src/airflow/sdk/execution_time/loop.py index f7adb18360395..b515001fe655d 100644 --- a/task-sdk/src/airflow/sdk/execution_time/loop.py +++ b/task-sdk/src/airflow/sdk/execution_time/loop.py @@ -56,25 +56,21 @@ def index(self) -> int: def max_iterations(self) -> int: return self._descriptor.max_iterations - @property - def group_id(self) -> str: - return self._descriptor.node_id - @property def region_id(self) -> UUID: return self._descriptor.region_id @property def trajectory(self) -> int: - return cast("int", self._descriptor.trajectory) + 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: - return cast("int", self._descriptor.resumes_from_index) + def resumes_from_index(self) -> int | None: + return self._descriptor.resumes_from_index @property def previous(self) -> Any: 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 8fb379f1b6589..b10b3e7633470 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -3242,6 +3242,79 @@ "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" + }, + "region_id": { + "format": "uuid", + "title": "Region Id", + "type": "string" + }, + "trajectory": { + "minimum": 1, + "title": "Trajectory", + "type": "integer" + }, + "forked_from_region_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Forked From Region Id" + }, + "resumes_from_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Resumes From Index" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped", + "region_id", + "trajectory" + ], + "title": "LoopContext", + "type": "object" + }, "LoopDecisionPayload": { "additionalProperties": false, "description": "A loop gate's evaluation of its ``until`` callable.", @@ -5217,73 +5290,6 @@ "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": { - "default": 1, - "minimum": 1, - "title": "Trajectory", - "type": "integer" - }, - "forked_from_region_id": { - "anyOf": [ - { - "format": "uuid", - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Forked From Region Id" - }, - "resumes_from_index": { - "default": 0, - "minimum": 0, - "title": "Resumes From Index", - "type": "integer" - } - }, - "required": [ - "node_id", - "index", - "max_iterations", - "terminal_task_id", - "terminal_is_mapped", - "region_id" - ], - "title": "LoopContext", - "type": "object" - }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { 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 c4d5310bf6e5e..ee78312e95ed2 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -284,7 +284,7 @@ def loop_decision(self) -> LoopDecision | None: @property def expanded_ti_count(self) -> int | None: - """Number of instances in this mapped task's expansion, ``None`` for a task that isn't mapped.""" + """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 @@ -382,8 +382,8 @@ def get_template_context(self) -> Context: 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/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 9ab934e2cb9c9..6e86037fe0550 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 @@ -527,6 +527,7 @@ def test_loop_context_follows_supervisor_version(self, real_migrator, startup_de 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"] 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 cfcf76797de30..a1fe5a1c0e232 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_loop.py +++ b/task-sdk/tests/task_sdk/execution_time/test_loop.py @@ -63,6 +63,7 @@ def body(): terminal_task_id=group.terminal_task_id, terminal_is_mapped=False, region_id=LOOP_REGION_ID, + trajectory=1, ) return ti @@ -106,6 +107,7 @@ def body(): 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) @@ -120,7 +122,7 @@ def test_task_instance_exposes_the_loop_iteration_and_its_trajectory(loop_ti): loop = ti.loop assert loop is not None - assert (loop.group_id, loop.index, loop.max_iterations) == ("body", 2, 3) + 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) diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index 8fb379f1b6589..b10b3e7633470 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -3242,6 +3242,79 @@ "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" + }, + "region_id": { + "format": "uuid", + "title": "Region Id", + "type": "string" + }, + "trajectory": { + "minimum": 1, + "title": "Trajectory", + "type": "integer" + }, + "forked_from_region_id": { + "anyOf": [ + { + "format": "uuid", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Forked From Region Id" + }, + "resumes_from_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Resumes From Index" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped", + "region_id", + "trajectory" + ], + "title": "LoopContext", + "type": "object" + }, "LoopDecisionPayload": { "additionalProperties": false, "description": "A loop gate's evaluation of its ``until`` callable.", @@ -5217,73 +5290,6 @@ "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": { - "default": 1, - "minimum": 1, - "title": "Trajectory", - "type": "integer" - }, - "forked_from_region_id": { - "anyOf": [ - { - "format": "uuid", - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Forked From Region Id" - }, - "resumes_from_index": { - "default": 0, - "minimum": 0, - "title": "Resumes From Index", - "type": "integer" - } - }, - "required": [ - "node_id", - "index", - "max_iterations", - "terminal_task_id", - "terminal_is_mapped", - "region_id" - ], - "title": "LoopContext", - "type": "object" - }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 1e31f770f664f..2e9b159844b22 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -239,7 +239,7 @@ export type TerminalIsMapped = boolean; export type RegionId1 = string; export type Trajectory = number; export type ForkedFromRegionId = string | null; -export type ResumesFromIndex = number; +export type ResumesFromIndex = number | null; export type ExpandedTiCount = number | null; export type TaskRescheduleCount = number; export type FirstTaskRescheduleStartDate = string | null; @@ -1083,7 +1083,7 @@ export interface LoopContext { terminal_task_id: TerminalTaskId; terminal_is_mapped: TerminalIsMapped; region_id: RegionId1; - trajectory?: Trajectory; + trajectory: Trajectory; forked_from_region_id?: ForkedFromRegionId; resumes_from_index?: ResumesFromIndex; } From f823c8f28b45f7caead3ad517407b45083afb0f0 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Tue, 20 Oct 2026 15:28:09 +0200 Subject: [PATCH 3/5] Test that expanded_ti_count counts one loop iteration and skips retried attempts Co-Authored-By: Claude --- .../unit/models/test_task_coordinates.py | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/airflow-core/tests/unit/models/test_task_coordinates.py b/airflow-core/tests/unit/models/test_task_coordinates.py index de1213e3d9ed0..46e1665bea44b 100644 --- a/airflow-core/tests/unit/models/test_task_coordinates.py +++ b/airflow-core/tests/unit/models/test_task_coordinates.py @@ -208,6 +208,38 @@ def test_expanded_ti_count_counts_live_instances_of_the_expansion(loop_coordinat 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=1, + 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") From abcbe8c9b0174e6567859b733723ecf10655e49b Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 21 Oct 2026 11:26:24 +0200 Subject: [PATCH 4/5] Address PR review Co-Authored-By: Claude --- .../docs/authoring-and-scheduling/loops.rst | 13 +- .../execution_api/datamodels/taskinstance.py | 17 ++- .../execution_api/routes/task_instances.py | 20 +-- .../src/airflow/models/task_coordinates.py | 75 ++++----- .../unit/models/test_task_coordinates.py | 45 +++--- go-sdk/pkg/execution/genmodels/models.gen.go | 126 ++++++++-------- go-sdk/schema/supervisor-schema.json | 142 ++++++++---------- .../sdk/execution_time/schema/schema.json | 142 ++++++++---------- ts-sdk/schema/supervisor-schema.json | 142 ++++++++---------- ts-sdk/src/generated/supervisor.ts | 6 +- 10 files changed, 327 insertions(+), 401 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst index 9893175c5a9b2..c85ca06c34d75 100644 --- a/airflow-core/docs/authoring-and-scheduling/loops.rst +++ b/airflow-core/docs/authoring-and-scheduling/loops.rst @@ -440,10 +440,9 @@ 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 the start again. 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``. +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 dd1fc0c90b5c9..f3ee8c0aa60f5 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 @@ 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 @@ -473,7 +474,9 @@ class LoopContext(BaseModel): ``resumes_from_index`` on run again in the fork. """ - forked_from_region_id: uuid.UUID | None = None + 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. @@ -481,7 +484,10 @@ class LoopContext(BaseModel): the chain too. """ - resumes_from_index: int | None = None + 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 0.""" @@ -493,7 +499,10 @@ class TIRunContext(BaseModel): loop: LoopContext | None = None - expanded_ti_count: int | 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. 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 e7472245d6af0..226ec01b91edf 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.taskinstance import ( DagRunNoteUpdatePayload, InactiveAssetsResponse, - LoopContext, PreviousTIResponse, PrevSuccessfulDagRunResponse, TaskBreadcrumbsResponse, @@ -376,23 +375,8 @@ def ti_run( 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_position(ti, task=task)) is not None: - terminal = resolver.get_task( - ti.dag_id, ti.run_id, loop.group.terminal_task_id, dag_version_id=ti.dag_version_id - ) - context.loop = LoopContext( - node_id=loop.group.node_id, - index=loop.index, - max_iterations=loop.group.max_iterations, - terminal_task_id=loop.group.terminal_task_id, - terminal_is_mapped=terminal.get_needs_expansion(), - region_id=loop.region_id, - trajectory=loop.trajectory, - ) - # Set only for a fork, as the route leaves out unset fields - if loop.forked_from_region_id is not None: - context.loop.forked_from_region_id = loop.forked_from_region_id - context.loop.resumes_from_index = loop.resumes_from_index + 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/models/task_coordinates.py b/airflow-core/src/airflow/models/task_coordinates.py index 333025dba304f..fd85cbddc60d0 100644 --- a/airflow-core/src/airflow/models/task_coordinates.py +++ b/airflow-core/src/airflow/models/task_coordinates.py @@ -23,6 +23,7 @@ import attrs 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 @@ -71,26 +72,6 @@ class _ProducerRequest: loop_node_id: str | None = None -@attrs.frozen(kw_only=True) -class LoopPosition: - """ - The iteration of the innermost loop a task instance runs in. - - ``region_id`` is the loop region that holds the iteration, also for a mapped task, whose own ``region_id`` - is its expansion's, and identifies the trajectory. A clear that reruns later iterations forks the newest - loop region: ``trajectory`` numbers the regions of that chain from 1, ``forked_from_region_id`` is the - previous region of the chain and ``resumes_from_index`` the first iteration the fork holds. Both are - ``None`` for the original region. - """ - - group: SerializedLoopTaskGroup - index: int - region_id: UUID - trajectory: int - forked_from_region_id: UUID | None - resumes_from_index: int | None - - class TaskCoordinate(Protocol): """Stored task identity shared by live and retained task data.""" @@ -253,18 +234,26 @@ def get_loop_iteration(self, ti: TaskCoordinate) -> tuple[str, int] | None: def loop_context( self, ti: TaskCoordinate, *, task: SerializedOperator | None = None ) -> tuple[SerializedLoopTaskGroup, int] | None: - if (position := self.get_loop_position(ti, task=task)) is None: + 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 - return position.group, position.index + 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_position( + def get_loop_context( self, ti: TaskCoordinate, *, task: SerializedOperator | None = None - ) -> LoopPosition | None: + ) -> LoopContext | None: """ - Return the iteration of the innermost loop ``ti`` runs in, or None outside any loop. + Return the loop context a worker gets in ``TIRunContext.loop``, or None outside any loop. - This is what a worker gets as ``TIRunContext.loop``. Server-side consumers such as listeners get the - same from a resolver built with :meth:`for_dag` for the Dag they hold. + 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 @@ -273,11 +262,7 @@ def get_loop_position( group = enclosing_loop(task) if group is None: return None - 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) - ) + regions = self._load_ancestry(ti) coordinate = DynamicRegion.find_physical_coordinate( ti.region_id, ti.region_index, group.node_id, regions ) @@ -285,15 +270,31 @@ def get_loop_position( raise ValueError("Task coordinates do not belong to the pinned loop") region_id, index = coordinate region = regions[region_id] - is_fork = region.forked_from_region_id is not None - return LoopPosition( - group=group, + 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), - forked_from_region_id=region.forked_from_region_id, - resumes_from_index=region.resumes_from_index if is_fork else None, ) + # 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) + ) + return regions def get_expanded_ti_count( self, ti: TaskCoordinate, *, task: SerializedOperator | None = None diff --git a/airflow-core/tests/unit/models/test_task_coordinates.py b/airflow-core/tests/unit/models/test_task_coordinates.py index 46e1665bea44b..e43ec928d3292 100644 --- a/airflow-core/tests/unit/models/test_task_coordinates.py +++ b/airflow-core/tests/unit/models/test_task_coordinates.py @@ -105,9 +105,9 @@ def body(): assert resolver.loop_iterations(ti) == [("body", 3)] assert resolver.public_map_index(ti) == 1 - position = resolver.get_loop_position(ti) - assert position is not None - assert (position.index, position.region_id, position.trajectory) == (3, loop_region.id, 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]) @@ -168,28 +168,35 @@ def test_loop_context_reuses_ancestry_loaded_by_loop_iterations(loop_coordinates assert (group.group_id, index) == ("body", 1) -def test_loop_position_describes_the_region_holding_the_iteration(forked_loop_coordinates, session): +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_position(producer) - fork = resolver.get_loop_position(consumer) + original = resolver.get_loop_context(producer) + fork = resolver.get_loop_context(consumer) assert original is not None assert fork is not None - assert (original.region_id, original.trajectory, original.forked_from_region_id) == ( - producer.region_id, - 1, - None, - ) - assert original.resumes_from_index is None - assert (fork.region_id, fork.trajectory, fork.forked_from_region_id) == ( - consumer.region_id, - 2, - producer.region_id, - ) - assert fork.resumes_from_index == 2 - assert resolver.get_loop_position(outside) is 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): diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 004ccb0d518eb..fbf36031cea99 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1322,7 +1322,7 @@ type LogicalDates []time.Time type LoopContext struct { // ForkedFromRegionID corresponds to the JSON schema field // "forked_from_region_id". - ForkedFromRegionID interface{} `msgpack:"forked_from_region_id,omitempty"` + ForkedFromRegionID *string `msgpack:"forked_from_region_id,omitempty"` // Index corresponds to the JSON schema field "index". Index int `msgpack:"index"` @@ -1337,7 +1337,7 @@ type LoopContext struct { RegionID string `msgpack:"region_id"` // ResumesFromIndex corresponds to the JSON schema field "resumes_from_index". - ResumesFromIndex interface{} `msgpack:"resumes_from_index,omitempty"` + ResumesFromIndex *int `msgpack:"resumes_from_index,omitempty"` // TerminalIsMapped corresponds to the JSON schema field "terminal_is_mapped". TerminalIsMapped bool `msgpack:"terminal_is_mapped"` @@ -1710,7 +1710,7 @@ type TIRunContext struct { DagRun DagRun `msgpack:"dag_run"` // ExpandedTICount corresponds to the JSON schema field "expanded_ti_count". - ExpandedTICount interface{} `msgpack:"expanded_ti_count,omitempty"` + ExpandedTICount *int `msgpack:"expanded_ti_count,omitempty"` // FirstTaskRescheduleStartDate corresponds to the JSON schema field // "first_task_reschedule_start_date". @@ -1750,53 +1750,21 @@ type TIRunContext struct { type TaskArgBinding interface{} -type TaskBreadcrumbsResult struct { - // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". - Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` - - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` -} - -type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} - -// Task callback status information. -// -// A Class with information about the success/failure TI callback to be executed. -// Currently, only failure -// callbacks when tasks are externally killed or experience heartbeat timeouts are -// run via DagFileProcessorProcess. -type TaskCallbackRequest struct { - // BundleName corresponds to the JSON schema field "bundle_name". - BundleName string `msgpack:"bundle_name"` - - // BundleVersion corresponds to the JSON schema field "bundle_version". - BundleVersion interface{} `msgpack:"bundle_version"` - - // ContextFromServer corresponds to the JSON schema field "context_from_server". - ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` - - // Filepath corresponds to the JSON schema field "filepath". - Filepath string `msgpack:"filepath"` - - // Msg corresponds to the JSON schema field "msg". - Msg interface{} `msgpack:"msg,omitempty"` +type TriggerKwargs map[string]interface{} - // TaskCallbackType corresponds to the JSON schema field "task_callback_type". - TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` +type Warnings []interface{} - // TI corresponds to the JSON schema field "ti". - TI TaskInstance `msgpack:"ti"` +type TaskIds []string - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` - // VersionData corresponds to the JSON schema field "version_data". - VersionData *VersionData `msgpack:"version_data,omitempty"` + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` } -type TaskIds []string - // Schema for TaskInstance model with minimal required fields needed for Runtime. type TaskInstance struct { // ContextCarrier corresponds to the JSON schema field "context_carrier". @@ -1836,41 +1804,69 @@ type TaskInstance struct { TryNumber int `msgpack:"try_number"` } -type TaskInstanceState string - const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" const TaskInstanceStateDeferred TaskInstanceState = "deferred" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" const TaskInstanceStateFailed TaskInstanceState = "failed" +const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateRunning TaskInstanceState = "running" const TaskInstanceStateQueued TaskInstanceState = "queued" + +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" const TaskInstanceStateRemoved TaskInstanceState = "removed" -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` +type TaskInstanceState string - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` +type TaskOutlets []AssetProfile + +type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} + +type TaskBreadcrumbsResult struct { + // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". + Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` } -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" -const TaskInstanceStateRunning TaskInstanceState = "running" -const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateRestarting TaskInstanceState = "restarting" +// Task callback status information. +// +// A Class with information about the success/failure TI callback to be executed. +// Currently, only failure +// callbacks when tasks are externally killed or experience heartbeat timeouts are +// run via DagFileProcessorProcess. +type TaskCallbackRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` -type VersionData map[string]interface{} + // BundleVersion corresponds to the JSON schema field "bundle_version". + BundleVersion interface{} `msgpack:"bundle_version"` -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" + // ContextFromServer corresponds to the JSON schema field "context_from_server". + ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` -type TaskOutlets []AssetProfile + // Filepath corresponds to the JSON schema field "filepath". + Filepath string `msgpack:"filepath"` -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" + // Msg corresponds to the JSON schema field "msg". + Msg interface{} `msgpack:"msg,omitempty"` -type TriggerKwargs map[string]interface{} + // TaskCallbackType corresponds to the JSON schema field "task_callback_type". + TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` -type Warnings []interface{} + // TI corresponds to the JSON schema field "ti". + TI TaskInstance `msgpack:"ti"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` + + // VersionData corresponds to the JSON schema field "version_data". + VersionData *VersionData `msgpack:"version_data,omitempty"` +} // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { @@ -1973,6 +1969,8 @@ type UpdateDagRunNote struct { Type string `msgpack:"type,omitempty"` } +type VersionData map[string]interface{} + // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { // ChosenOptions corresponds to the JSON schema field "chosen_options". diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index b10b3e7633470..e739708d12dfc 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -3242,79 +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" - }, - "region_id": { - "format": "uuid", - "title": "Region Id", - "type": "string" - }, - "trajectory": { - "minimum": 1, - "title": "Trajectory", - "type": "integer" - }, - "forked_from_region_id": { - "anyOf": [ - { - "format": "uuid", - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Forked From Region Id" - }, - "resumes_from_index": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Resumes From Index" - } - }, - "required": [ - "node_id", - "index", - "max_iterations", - "terminal_task_id", - "terminal_is_mapped", - "region_id", - "trajectory" - ], - "title": "LoopContext", - "type": "object" - }, "LoopDecisionPayload": { "additionalProperties": false, "description": "A loop gate's evaluation of its ``until`` callable.", @@ -5290,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": { @@ -5308,16 +5292,8 @@ "default": null }, "expanded_ti_count": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Expanded Ti Count" + "title": "Expanded Ti Count", + "type": "integer" }, "task_reschedule_count": { "default": 0, 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 b10b3e7633470..e739708d12dfc 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -3242,79 +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" - }, - "region_id": { - "format": "uuid", - "title": "Region Id", - "type": "string" - }, - "trajectory": { - "minimum": 1, - "title": "Trajectory", - "type": "integer" - }, - "forked_from_region_id": { - "anyOf": [ - { - "format": "uuid", - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Forked From Region Id" - }, - "resumes_from_index": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Resumes From Index" - } - }, - "required": [ - "node_id", - "index", - "max_iterations", - "terminal_task_id", - "terminal_is_mapped", - "region_id", - "trajectory" - ], - "title": "LoopContext", - "type": "object" - }, "LoopDecisionPayload": { "additionalProperties": false, "description": "A loop gate's evaluation of its ``until`` callable.", @@ -5290,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": { @@ -5308,16 +5292,8 @@ "default": null }, "expanded_ti_count": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Expanded Ti Count" + "title": "Expanded Ti Count", + "type": "integer" }, "task_reschedule_count": { "default": 0, diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index b10b3e7633470..e739708d12dfc 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -3242,79 +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" - }, - "region_id": { - "format": "uuid", - "title": "Region Id", - "type": "string" - }, - "trajectory": { - "minimum": 1, - "title": "Trajectory", - "type": "integer" - }, - "forked_from_region_id": { - "anyOf": [ - { - "format": "uuid", - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Forked From Region Id" - }, - "resumes_from_index": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Resumes From Index" - } - }, - "required": [ - "node_id", - "index", - "max_iterations", - "terminal_task_id", - "terminal_is_mapped", - "region_id", - "trajectory" - ], - "title": "LoopContext", - "type": "object" - }, "LoopDecisionPayload": { "additionalProperties": false, "description": "A loop gate's evaluation of its ``until`` callable.", @@ -5290,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": { @@ -5308,16 +5292,8 @@ "default": null }, "expanded_ti_count": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Expanded Ti Count" + "title": "Expanded Ti Count", + "type": "integer" }, "task_reschedule_count": { "default": 0, diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 2e9b159844b22..75a05dde12cd6 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -238,9 +238,9 @@ export type TerminalTaskId = string; export type TerminalIsMapped = boolean; export type RegionId1 = string; export type Trajectory = number; -export type ForkedFromRegionId = string | null; -export type ResumesFromIndex = number | null; -export type ExpandedTiCount = number | null; +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; From f3d974d6c63c4285b50930b6a288501c13ab4223 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 21 Oct 2026 11:42:14 +0200 Subject: [PATCH 5/5] Follow one-based task loop iterations Co-Authored-By: Claude --- airflow-core/docs/authoring-and-scheduling/loops.rst | 2 +- .../api_fastapi/execution_api/datamodels/taskinstance.py | 2 +- airflow-core/tests/unit/models/test_task_coordinates.py | 2 +- task-sdk/tests/task_sdk/execution_time/test_loop.py | 6 +++--- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst index c85ca06c34d75..9e414563e5977 100644 --- a/airflow-core/docs/authoring-and-scheduling/loops.rst +++ b/airflow-core/docs/authoring-and-scheduling/loops.rst @@ -161,7 +161,7 @@ definition. * - ``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 0. + - 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``. 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 f3ee8c0aa60f5..8a1deb203fd89 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 @@ -488,7 +488,7 @@ class LoopContext(BaseModel): 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 0.""" + """The first iteration ``region_id`` holds. Absent for the original region, which holds them from 1.""" class TIRunContext(BaseModel): diff --git a/airflow-core/tests/unit/models/test_task_coordinates.py b/airflow-core/tests/unit/models/test_task_coordinates.py index e43ec928d3292..57af3454f802a 100644 --- a/airflow-core/tests/unit/models/test_task_coordinates.py +++ b/airflow-core/tests/unit/models/test_task_coordinates.py @@ -232,7 +232,7 @@ def body(): 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=1, + parent_region_index=2, session=session, ) session.add(next_expansion) 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 a1fe5a1c0e232..8dd474773b138 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_loop.py +++ b/task-sdk/tests/task_sdk/execution_time/test_loop.py @@ -145,9 +145,9 @@ def test_task_instance_exposes_the_size_of_its_expansion(create_runtime_ti, expa @pytest.mark.parametrize( ("index", "until", "expected"), [ - (0, lambda: LoopDecision(stop=False, reason="not yet"), LoopDecision(stop=False, reason="not yet")), - (2, lambda: True, LoopDecision(stop=True)), - (2, lambda: False, LoopDecision(stop=False)), + (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"], )