diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst index 7a6e04e0c5988..9e414563e5977 100644 --- a/airflow-core/docs/authoring-and-scheduling/loops.rst +++ b/airflow-core/docs/authoring-and-scheduling/loops.rst @@ -152,9 +152,21 @@ definition. - The previous iteration's result, or ``None`` in iteration 1. * - ``loop.result`` - The terminal body task's result (the XCom pushed under the key ``return_value``) for the current iteration, used by the gate. + * - ``loop.region_id`` + - The id of the loop execution (trajectory) that holds this iteration. It is the same for every task + instance of the iteration, mapped ones included. + * - ``loop.trajectory`` + - The number of that trajectory: 1 for the loop's first execution, and one more each time a rerun + replaces later iterations, see :ref:`loops-trajectories`. + * - ``loop.forked_from_region_id`` + - The ``loop.region_id`` of the previous trajectory, or ``None`` in trajectory 1. + * - ``loop.resumes_from_index`` + - The first iteration this trajectory holds, or ``None`` in trajectory 1, which holds them from 1. Use ``loop.index`` for the loop iteration and ``ti.map_index`` for a mapped task instance's position. Each task instance can have several tries, each with its own ``ti.try_number``. +The task instance also exposes the loop context as ``ti.loop``, and for a mapped task +``ti.expanded_ti_count`` is the number of instances in its expansion, counted within the iteration. Pass data between iterations ============================ @@ -347,6 +359,12 @@ execution. Downstream tasks then run as they would after any successful gate: th gate that was marked successful by hand from one that stopped because ``until`` returned ``True``. Any later iterations retained after a selective clear remain unchanged. +The gate's task instance exposes its decision as ``ti.loop_decision``, a ``LoopDecision``, once the gate +has evaluated it. With ``until``, ``stop`` is ``True`` when ``until`` ended the loop, even at its last +iteration. When ``stop`` is ``False`` at the last iteration (``loop.index == loop.max_iterations - 1``), +the loop reached ``max_iterations`` without meeting ``until`` and the gate fails. A gate that fails before +deciding, for example because ``until`` raised, or that was marked successful by hand, has no decision. + Iterations and execution history ================================ @@ -414,3 +432,17 @@ which leave them in place. To keep later iterations, clear from the UI or the REST endpoint. Tasks you leave uncleared are not recomputed using the results of the rerun. + +.. _loops-trajectories: + +Reruns and trajectories +----------------------- + +Clearing later iterations starts a new execution of the loop that replaces the previous ones from the +iteration after the rerun gate. Each execution is a trajectory, identified by ``loop.region_id`` and +numbered by ``loop.trajectory``: 1 for the first one, growing by one with each rerun. The rerun gate, and +whatever you clear in its iteration, run again in the trajectory that iteration belongs to, with their +next try. The replacement iterations run in the new trajectory, from ``loop.resumes_from_index``, and their +tries count from 1 again for each task. In the example above, iteration 2 stays in trajectory 1, while +replacement iterations 3 and 4 run in trajectory 2, with ``loop.resumes_from_index`` 3. Task log paths show +the trajectory after the iteration, as in ``pass=3.2``. diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index 09ce1193e9018..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 @@ -32,8 +32,9 @@ field_validator, model_validator, ) +from pydantic.json_schema import SkipJsonSchema -from airflow.api_fastapi.common.region import OmitsMissingRegion, RegionId, RegionIndex +from airflow.api_fastapi.common.region import OmitsMissingRegion, RegionId, RegionIndex, _drop_null_default from airflow.api_fastapi.common.types import UtcDateTime from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile @@ -457,6 +458,38 @@ class LoopContext(BaseModel): terminal_task_id: str terminal_is_mapped: bool + region_id: uuid.UUID + """ + The loop region holding this iteration, which identifies the trajectory it runs in. + + A mapped task's own ``region_id`` is its expansion's, so this is how its instances are tied to the + unmapped tasks of the same iteration. + """ + + trajectory: int = Field(ge=1) + """ + The number of ``region_id`` in the loop's chain of regions: 1 for the original, n + 1 for its n-th fork. + + Clearing a loop gate with the later iterations forks the newest loop region, and the iterations from + ``resumes_from_index`` on run again in the fork. + """ + + forked_from_region_id: Annotated[ + SkipJsonSchema[None] | uuid.UUID, Field(json_schema_extra=_drop_null_default) + ] = None + """ + The region ``region_id`` was forked from, the previous one in the chain. Absent for the original region. + + The fork replaces the iterations from ``resumes_from_index`` on, which can be held by earlier regions of + the chain too. + """ + + resumes_from_index: Annotated[ + SkipJsonSchema[None] | int, + Field(json_schema_extra=_drop_null_default), + ] = None + """The first iteration ``region_id`` holds. Absent for the original region, which holds them from 1.""" + class TIRunContext(BaseModel): """Response schema for TaskInstance run context.""" @@ -466,6 +499,16 @@ class TIRunContext(BaseModel): loop: LoopContext | None = None + expanded_ti_count: Annotated[ + SkipJsonSchema[None] | int, + Field(json_schema_extra=_drop_null_default), + ] = None + """ + Number of live, non-removed instances in the mapped task's expansion this instance belongs to. + + ``None`` for a task that isn't mapped. Inside a loop, it counts the instances of this iteration only. + """ + task_reschedule_count: int = 0 """How many times the task has been rescheduled.""" 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..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, @@ -371,19 +370,13 @@ def ti_run( should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), multi_team=conf.getboolean("core", "multi_team"), ) - resolver = TaskCoordinateResolver(dag_bag, session) - if loop_context := resolver.loop_context(ti): - group, index = loop_context - terminal = resolver.get_task( - ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id - ) - context.loop = LoopContext( - node_id=group.node_id, - index=index, - max_iterations=group.max_iterations, - terminal_task_id=group.terminal_task_id, - terminal_is_mapped=terminal.get_needs_expansion(), - ) + if ti.region_id != SENTINEL_REGION_ID: + resolver = TaskCoordinateResolver(dag_bag, session) + task = resolver.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id) + if (expanded_ti_count := resolver.get_expanded_ti_count(ti, task=task)) is not None: + context.expanded_ti_count = expanded_ti_count + if (loop := resolver.get_loop_context(ti, task=task)) is not None: + context.loop = loop if first_task_reschedule_start_date is not None: context.first_task_reschedule_start_date = first_task_reschedule_start_date diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index 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..fd85cbddc60d0 100644 --- a/airflow-core/src/airflow/models/task_coordinates.py +++ b/airflow-core/src/airflow/models/task_coordinates.py @@ -21,8 +21,9 @@ from uuid import UUID import attrs -from sqlalchemy import FromClause, case, or_, select, tuple_ +from sqlalchemy import FromClause, case, func, or_, select, tuple_ +from airflow.api_fastapi.execution_api.datamodels.taskinstance import LoopContext from airflow.exceptions import TaskNotFound from airflow.models.dagbag import DBDagBag from airflow.models.dagrun import DagRun @@ -31,6 +32,7 @@ AmbiguousProducerError, DynamicRegion, ProducerContext, + get_trajectory, load_region_ancestry, loop_position, resolve_current_producers, @@ -40,6 +42,7 @@ from airflow.models.taskinstance import TaskInstance from airflow.serialization.definitions.mappedoperator import is_mapped from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup, SerializedMappedTaskGroup +from airflow.utils.state import TaskInstanceState if TYPE_CHECKING: from collections.abc import Collection, Iterable, Sequence @@ -238,15 +241,85 @@ def loop_context( group = enclosing_loop(task) if group is None: return None + position = loop_position(self._load_ancestry(ti), ti.region_id, ti.region_index, group.node_id) + if position is None: + raise ValueError("Task coordinates do not belong to the pinned loop") + return group, position[1] + + def get_loop_context( + self, ti: TaskCoordinate, *, task: SerializedOperator | None = None + ) -> LoopContext | None: + """ + Return the loop context a worker gets in ``TIRunContext.loop``, or None outside any loop. + + Server-side consumers such as listeners get the same from a resolver built with :meth:`for_dag` for + the Dag they hold. + """ + if ti.region_id == SENTINEL_REGION_ID: + return None + if task is None: + task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id) + group = enclosing_loop(task) + if group is None: + return None + regions = self._load_ancestry(ti) + coordinate = DynamicRegion.find_physical_coordinate( + ti.region_id, ti.region_index, group.node_id, regions + ) + if coordinate is None: + raise ValueError("Task coordinates do not belong to the pinned loop") + region_id, index = coordinate + region = regions[region_id] + terminal = self.get_task( + ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id + ) + context = LoopContext( + node_id=group.node_id, + index=index, + max_iterations=group.max_iterations, + terminal_task_id=group.terminal_task_id, + terminal_is_mapped=terminal.get_needs_expansion(), + region_id=region_id, + trajectory=get_trajectory(region, regions), + ) + # Set only for a fork, as the Execution API leaves out unset fields + if region.forked_from_region_id is not None: + context.forked_from_region_id = region.forked_from_region_id + context.resumes_from_index = region.resumes_from_index + return context + + def _load_ancestry(self, ti: TaskCoordinate) -> dict[UUID, DynamicRegion]: regions = self._regions.setdefault((ti.dag_id, ti.run_id), {}) if ti.region_id not in regions: regions.update( load_region_ancestry([ti.region_id], dag_id=ti.dag_id, run_id=ti.run_id, session=self.session) ) - position = loop_position(regions, ti.region_id, ti.region_index, group.node_id) - if position is None: - raise ValueError("Task coordinates do not belong to the pinned loop") - return group, position[1] + return regions + + def get_expanded_ti_count( + self, ti: TaskCoordinate, *, task: SerializedOperator | None = None + ) -> int | None: + """ + Count the live instances of a mapped task in the expansion ``ti`` belongs to, None if not mapped. + + This is what a worker gets as ``TIRunContext.expanded_ti_count``. + """ + if ti.region_id == SENTINEL_REGION_ID: + return None + if task is None: + task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id) + if not task.get_needs_expansion(): + return None + return self.session.scalar( + select(func.count()).where( + TaskInstance.working_set.is_(True), + TaskInstance.dag_id == ti.dag_id, + TaskInstance.run_id == ti.run_id, + TaskInstance.task_id == ti.task_id, + TaskInstance.region_id == ti.region_id, + or_(TaskInstance.state.is_(None), TaskInstance.state != TaskInstanceState.REMOVED), + ) + ) def get_task( self, dag_id: str, run_id: str, task_id: str, *, dag_version_id: UUID | None = None 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..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 @@ -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,57 @@ def body(): "max_iterations": 4, "terminal_task_id": loop.terminal_task_id, "terminal_is_mapped": mapped, + "region_id": str(parent.id), + "trajectory": 1, } + assert response.json().get("expanded_ti_count") == (1 if mapped else None) + + def test_startup_loop_context_describes_the_fork_holding_the_iteration(self, client, session, dag_maker): + @task_group + def body(): + EmptyOperator(task_id="terminal") + + with dag_maker(serialized=True): + loop = create_loop(body, max_iterations=4) + dr = dag_maker.create_dagrun(state=DagRunState.RUNNING) + ti = next(ti for ti in dr.task_instances if ti.task_id == loop.terminal_task_id) + original = DynamicRegion.get_or_create( + dag_id=dr.dag_id, run_id=dr.run_id, node_id=loop.group_id, session=session + ) + session.add(original) + session.flush() + fork = DynamicRegion.fork_family(original.id, 2, {original.id: original}, session=session) + session.flush() + ti.region_id, ti.region_index = fork.id, 2 + ti.state = State.QUEUED + session.commit() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + + assert response.status_code == 200 + loop_context = response.json()["loop"] + assert loop_context["region_id"] == str(fork.id) + assert loop_context["trajectory"] == 2 + assert loop_context["forked_from_region_id"] == str(original.id) + assert loop_context["resumes_from_index"] == 2 + + @pytest.mark.parametrize("removed", [0, 1]) + def test_startup_expanded_ti_count_counts_live_instances_of_the_expansion( + self, client, session, dag_maker, removed + ): + with dag_maker(serialized=True): + PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2], [3]]) + dr = dag_maker.create_dagrun(state=DagRunState.RUNNING) + tis = sorted(dr.task_instances, key=lambda ti: ti.map_index) + for ti in tis[3 - removed :]: + ti.state = State.REMOVED + tis[0].state = State.QUEUED + session.commit() + + response = client.patch(f"/execution/task-instances/{tis[0].id}/run", json=self.RUN_PAYLOAD) + + assert response.status_code == 200 + assert response.json()["expanded_ti_count"] == 3 - removed @pytest.mark.parametrize("regional", [False, True]) def test_startup_xcom_cleanup_keys_belong_to_the_starting_try( diff --git a/airflow-core/tests/unit/models/test_dynamic_region.py b/airflow-core/tests/unit/models/test_dynamic_region.py index 573046e4e9b6e..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 @@ -36,6 +35,7 @@ ProducerContext, build_loop_family_legs, build_loop_sequence_order, + get_trajectory, resolve_current_producers, select_loop_producer_ids, ) @@ -52,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 @@ -1114,3 +1111,15 @@ def test_loop_sequence_of_mapped_task_skips_passes_a_fork_replaced(loop_run, ses (True, 4, 0), (True, 4, 1), ] + + +@pytest.mark.parametrize("forks", [0, 1, 2]) +def test_trajectory_counts_the_forks_a_region_descends_from(forks): + dag_run = DagRun(dag_id="dag", run_id="run") + region = make_region(dag_run, id=uuid4()) + regions = {region.id: region} + for _ in range(forks): + region = make_region(dag_run, id=uuid4(), forked_from_region_id=region.id) + regions[region.id] = region + + assert get_trajectory(region, regions) == forks + 1 diff --git a/airflow-core/tests/unit/models/test_task_coordinates.py b/airflow-core/tests/unit/models/test_task_coordinates.py index b74f38460a24d..57af3454f802a 100644 --- a/airflow-core/tests/unit/models/test_task_coordinates.py +++ b/airflow-core/tests/unit/models/test_task_coordinates.py @@ -105,6 +105,9 @@ def body(): assert resolver.loop_iterations(ti) == [("body", 3)] assert resolver.public_map_index(ti) == 1 + loop = resolver.get_loop_context(ti) + assert loop is not None + assert (loop.index, loop.region_id, loop.trajectory) == (3, loop_region.id, 1) @pytest.mark.parametrize("run_count", [1, 3]) @@ -165,6 +168,85 @@ def test_loop_context_reuses_ancestry_loaded_by_loop_iterations(loop_coordinates assert (group.group_id, index) == ("body", 1) +def test_loop_context_describes_the_region_holding_the_iteration(forked_loop_coordinates, session): + _, producer, consumer, _, outside = forked_loop_coordinates + resolver = TaskCoordinateResolver(DBDagBag(), session) + common = { + "node_id": "body", + "index": 2, + "max_iterations": 3, + "terminal_task_id": "body.consumer", + "terminal_is_mapped": False, + } + + original = resolver.get_loop_context(producer) + fork = resolver.get_loop_context(consumer) + + assert original is not None + assert fork is not None + assert original.model_dump(exclude_unset=True) == { + **common, + "region_id": producer.region_id, + "trajectory": 1, + } + assert fork.model_dump(exclude_unset=True) == { + **common, + "region_id": consumer.region_id, + "trajectory": 2, + "forked_from_region_id": producer.region_id, + "resumes_from_index": 2, + } + assert resolver.get_loop_context(outside) is None + + +def test_expanded_ti_count_counts_live_instances_of_the_expansion(loop_coordinates, dag_maker, session): + with dag_maker("mapped-expansion", serialized=True): + PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2], [3]]) + mapped = sorted(dag_maker.create_dagrun().task_instances, key=lambda ti: ti.map_index) + mapped[2].state = TaskInstanceState.REMOVED + _, _, consumer, _, outside = loop_coordinates + session.flush() + resolver = TaskCoordinateResolver(DBDagBag(), session) + resolver.get_expanded_ti_count(mapped[0]) + + with assert_queries_count(1): + assert resolver.get_expanded_ti_count(mapped[1]) == 2 + assert resolver.get_expanded_ti_count(consumer) is None + assert resolver.get_expanded_ti_count(outside) is None + + +def test_expanded_ti_count_counts_each_instance_of_one_loop_iteration_once(dag_maker, session): + @task_group + def body(): + PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2]]) + + with dag_maker(serialized=True): + create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + first, second = sorted( + (ti for ti in dr.task_instances if ti.task_id == "body.mapped"), key=lambda ti: ti.map_index + ) + first.prepare_db_for_next_try(session) + next_expansion = DynamicRegion.get_or_create( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id=first.task_id, + parent_region_id=session.get(DynamicRegion, first.region_id).parent_region_id, + parent_region_index=2, + session=session, + ) + session.add(next_expansion) + session.flush() + next_iteration = TaskInstance(task=first.task, run_id=dr.run_id, dag_version_id=first.dag_version_id) + next_iteration.region_id, next_iteration.region_index = next_expansion.id, 0 + session.add(next_iteration) + session.flush() + resolver = TaskCoordinateResolver(DBDagBag(), session) + + assert resolver.get_expanded_ti_count(second) == 2 + assert resolver.get_expanded_ti_count(next_iteration) == 1 + + def test_removed_task_coordinates_degrade_to_stored_region_data(dag_maker, session): with dag_maker("removed-coordinates", serialized=True): EmptyOperator(task_id="kept") diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 0a5dba694e675..fbf36031cea99 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1320,6 +1320,10 @@ type LogicalDates []time.Time // Pinned loop definition and enclosing iteration for a task execution. type LoopContext struct { + // ForkedFromRegionID corresponds to the JSON schema field + // "forked_from_region_id". + ForkedFromRegionID *string `msgpack:"forked_from_region_id,omitempty"` + // Index corresponds to the JSON schema field "index". Index int `msgpack:"index"` @@ -1329,11 +1333,20 @@ type LoopContext struct { // NodeID corresponds to the JSON schema field "node_id". NodeID string `msgpack:"node_id"` + // RegionID corresponds to the JSON schema field "region_id". + RegionID string `msgpack:"region_id"` + + // ResumesFromIndex corresponds to the JSON schema field "resumes_from_index". + ResumesFromIndex *int `msgpack:"resumes_from_index,omitempty"` + // TerminalIsMapped corresponds to the JSON schema field "terminal_is_mapped". TerminalIsMapped bool `msgpack:"terminal_is_mapped"` // TerminalTaskID corresponds to the JSON schema field "terminal_task_id". TerminalTaskID string `msgpack:"terminal_task_id"` + + // Trajectory corresponds to the JSON schema field "trajectory". + Trajectory int `msgpack:"trajectory"` } // A loop gate's evaluation of its “until“ callable. @@ -1696,6 +1709,9 @@ type TIRunContext struct { // DagRun corresponds to the JSON schema field "dag_run". DagRun DagRun `msgpack:"dag_run"` + // ExpandedTICount corresponds to the JSON schema field "expanded_ti_count". + ExpandedTICount *int `msgpack:"expanded_ti_count,omitempty"` + // FirstTaskRescheduleStartDate corresponds to the JSON schema field // "first_task_reschedule_start_date". FirstTaskRescheduleStartDate interface{} `msgpack:"first_task_reschedule_start_date,omitempty"` diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index 56a6967f17e77..e739708d12dfc 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -3242,42 +3242,6 @@ "title": "LazyDeserializedDAG", "type": "object" }, - "LoopContext": { - "description": "Pinned loop definition and enclosing iteration for a task execution.", - "properties": { - "node_id": { - "title": "Node Id", - "type": "string" - }, - "index": { - "minimum": 1, - "title": "Index", - "type": "integer" - }, - "max_iterations": { - "exclusiveMinimum": 0, - "title": "Max Iterations", - "type": "integer" - }, - "terminal_task_id": { - "title": "Terminal Task Id", - "type": "string" - }, - "terminal_is_mapped": { - "title": "Terminal Is Mapped", - "type": "boolean" - } - }, - "required": [ - "node_id", - "index", - "max_iterations", - "terminal_task_id", - "terminal_is_mapped" - ], - "title": "LoopContext", - "type": "object" - }, "LoopDecisionPayload": { "additionalProperties": false, "description": "A loop gate's evaluation of its ``until`` callable.", @@ -5253,6 +5217,63 @@ "title": "DagRun", "type": "object" }, + "LoopContext": { + "description": "Pinned loop definition and enclosing iteration for a task execution.", + "properties": { + "node_id": { + "title": "Node Id", + "type": "string" + }, + "index": { + "minimum": 1, + "title": "Index", + "type": "integer" + }, + "max_iterations": { + "exclusiveMinimum": 0, + "title": "Max Iterations", + "type": "integer" + }, + "terminal_task_id": { + "title": "Terminal Task Id", + "type": "string" + }, + "terminal_is_mapped": { + "title": "Terminal Is Mapped", + "type": "boolean" + }, + "region_id": { + "format": "uuid", + "title": "Region Id", + "type": "string" + }, + "trajectory": { + "minimum": 1, + "title": "Trajectory", + "type": "integer" + }, + "forked_from_region_id": { + "format": "uuid", + "title": "Forked From Region Id", + "type": "string" + }, + "resumes_from_index": { + "title": "Resumes From Index", + "type": "integer" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped", + "region_id", + "trajectory" + ], + "title": "LoopContext", + "type": "object" + }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { @@ -5270,6 +5291,10 @@ ], "default": null }, + "expanded_ti_count": { + "title": "Expanded Ti Count", + "type": "integer" + }, "task_reschedule_count": { "default": 0, "title": "Task Reschedule Count", diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 0e677c93e7a94..6c4b2536fc44c 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -342,6 +342,10 @@ class LoopContext(BaseModel): max_iterations: Annotated[int, Field(gt=0, title="Max Iterations")] terminal_task_id: Annotated[str, Field(title="Terminal Task Id")] terminal_is_mapped: Annotated[bool, Field(title="Terminal Is Mapped")] + region_id: Annotated[UUID, Field(title="Region Id")] + trajectory: Annotated[int, Field(ge=1, title="Trajectory")] + forked_from_region_id: Annotated[UUID | None, Field(title="Forked From Region Id")] = None + resumes_from_index: Annotated[int | None, Field(title="Resumes From Index")] = None class LoopDecisionPayload(BaseModel): @@ -999,6 +1003,7 @@ class TIRunContext(BaseModel): dag_run: DagRun loop: LoopContext | None = None + expanded_ti_count: Annotated[int | None, Field(title="Expanded Ti Count")] = None task_reschedule_count: Annotated[int | None, Field(title="Task Reschedule Count")] = 0 first_task_reschedule_start_date: Annotated[ AwareDatetime | None, Field(title="First Task Reschedule Start Date") diff --git a/task-sdk/src/airflow/sdk/execution_time/loop.py b/task-sdk/src/airflow/sdk/execution_time/loop.py index f8aac86119154..b515001fe655d 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,22 @@ def index(self) -> int: def max_iterations(self) -> int: return self._descriptor.max_iterations + @property + def region_id(self) -> UUID: + return self._descriptor.region_id + + @property + def trajectory(self) -> int: + return self._descriptor.trajectory + + @property + def forked_from_region_id(self) -> UUID | None: + return self._descriptor.forked_from_region_id + + @property + def resumes_from_index(self) -> int | None: + return self._descriptor.resumes_from_index + @property def previous(self) -> Any: # Iterations count from one, so the first pass has no previous result. diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 56a6967f17e77..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,42 +3242,6 @@ "title": "LazyDeserializedDAG", "type": "object" }, - "LoopContext": { - "description": "Pinned loop definition and enclosing iteration for a task execution.", - "properties": { - "node_id": { - "title": "Node Id", - "type": "string" - }, - "index": { - "minimum": 1, - "title": "Index", - "type": "integer" - }, - "max_iterations": { - "exclusiveMinimum": 0, - "title": "Max Iterations", - "type": "integer" - }, - "terminal_task_id": { - "title": "Terminal Task Id", - "type": "string" - }, - "terminal_is_mapped": { - "title": "Terminal Is Mapped", - "type": "boolean" - } - }, - "required": [ - "node_id", - "index", - "max_iterations", - "terminal_task_id", - "terminal_is_mapped" - ], - "title": "LoopContext", - "type": "object" - }, "LoopDecisionPayload": { "additionalProperties": false, "description": "A loop gate's evaluation of its ``until`` callable.", @@ -5253,6 +5217,63 @@ "title": "DagRun", "type": "object" }, + "LoopContext": { + "description": "Pinned loop definition and enclosing iteration for a task execution.", + "properties": { + "node_id": { + "title": "Node Id", + "type": "string" + }, + "index": { + "minimum": 1, + "title": "Index", + "type": "integer" + }, + "max_iterations": { + "exclusiveMinimum": 0, + "title": "Max Iterations", + "type": "integer" + }, + "terminal_task_id": { + "title": "Terminal Task Id", + "type": "string" + }, + "terminal_is_mapped": { + "title": "Terminal Is Mapped", + "type": "boolean" + }, + "region_id": { + "format": "uuid", + "title": "Region Id", + "type": "string" + }, + "trajectory": { + "minimum": 1, + "title": "Trajectory", + "type": "integer" + }, + "forked_from_region_id": { + "format": "uuid", + "title": "Forked From Region Id", + "type": "string" + }, + "resumes_from_index": { + "title": "Resumes From Index", + "type": "integer" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped", + "region_id", + "trajectory" + ], + "title": "LoopContext", + "type": "object" + }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { @@ -5270,6 +5291,10 @@ ], "default": null }, + "expanded_ti_count": { + "title": "Expanded Ti Count", + "type": "integer" + }, "task_reschedule_count": { "default": 0, "title": "Task Reschedule Count", diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 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..ee78312e95ed2 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 in its loop iteration, None if not mapped.""" + if self._ti_context_from_server is None: + return None + return self._ti_context_from_server.expanded_ti_count + @property def stats_tags(self) -> dict[str, str]: """Metric tags for this task instance, including dag tags and team_name when available.""" @@ -360,8 +382,8 @@ 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/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..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 @@ -526,13 +526,18 @@ 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"), + trajectory=1, ) + startup_details.ti_context.expanded_ti_count = 2 context = real_migrator.downgrade(startup_details, version).model_dump()["ti_context"] if version == "2026-06-16": assert "loop" not in context + assert "expanded_ti_count" not in context else: assert context["loop"] == startup_details.ti_context.loop.model_dump() + assert context["expanded_ti_count"] == 2 @pytest.mark.parametrize("version", ["2026-06-16", "2026-10-30"]) def test_task_coordinates_follow_supervisor_version(self, real_migrator, startup_details, version): diff --git a/task-sdk/tests/task_sdk/execution_time/test_loop.py b/task-sdk/tests/task_sdk/execution_time/test_loop.py index 77e19c10df5b4..8dd474773b138 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,8 @@ def body(): max_iterations=max_iterations, terminal_task_id=group.terminal_task_id, terminal_is_mapped=False, + region_id=LOOP_REGION_ID, + trajectory=1, ) return ti @@ -100,11 +106,59 @@ def body(): max_iterations=4, terminal_task_id=group.terminal_task_id, terminal_is_mapped=mapped, + region_id=LOOP_REGION_ID, + trajectory=1, ) assert operator.execute(ti.get_template_context()) == (5, 2, 0 if mapped else -1) +def test_task_instance_exposes_the_loop_iteration_and_its_trajectory(loop_ti): + ti = loop_ti(index=2) + ti._ti_context_from_server.loop = ti._ti_context_from_server.loop.model_copy( + update={"trajectory": 2, "forked_from_region_id": FORKED_REGION_ID, "resumes_from_index": 2} + ) + + loop = ti.loop + + assert loop is not None + assert (loop.index, loop.max_iterations) == (2, 3) + assert (loop.region_id, loop.trajectory) == (LOOP_REGION_ID, 2) + assert (loop.forked_from_region_id, loop.resumes_from_index) == (FORKED_REGION_ID, 2) + + +def test_task_instance_outside_a_loop_has_no_loop(create_runtime_ti): + ti = create_runtime_ti(task=BaseOperator(task_id="plain")) + + assert ti.loop is None + assert ti.loop_decision is None + + +@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"), + [ + (1, lambda: LoopDecision(stop=False, reason="not yet"), LoopDecision(stop=False, reason="not yet")), + (3, lambda: True, LoopDecision(stop=True)), + (3, lambda: False, LoopDecision(stop=False)), + ], + ids=["continue", "until_met", "cap_reached"], +) +def test_gate_exposes_its_decision(loop_ti, mock_supervisor_comms, index, until, expected): + ti = loop_ti(index=index, until=until) + + task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log) + + assert ti.loop_decision == expected + + @pytest.mark.parametrize("value", [0, False, [], ""]) def test_previous_and_current_result_preserve_falsey_values(loop_ti, mock_supervisor_comms, value): ti = loop_ti(index=2) diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index 56a6967f17e77..e739708d12dfc 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -3242,42 +3242,6 @@ "title": "LazyDeserializedDAG", "type": "object" }, - "LoopContext": { - "description": "Pinned loop definition and enclosing iteration for a task execution.", - "properties": { - "node_id": { - "title": "Node Id", - "type": "string" - }, - "index": { - "minimum": 1, - "title": "Index", - "type": "integer" - }, - "max_iterations": { - "exclusiveMinimum": 0, - "title": "Max Iterations", - "type": "integer" - }, - "terminal_task_id": { - "title": "Terminal Task Id", - "type": "string" - }, - "terminal_is_mapped": { - "title": "Terminal Is Mapped", - "type": "boolean" - } - }, - "required": [ - "node_id", - "index", - "max_iterations", - "terminal_task_id", - "terminal_is_mapped" - ], - "title": "LoopContext", - "type": "object" - }, "LoopDecisionPayload": { "additionalProperties": false, "description": "A loop gate's evaluation of its ``until`` callable.", @@ -5253,6 +5217,63 @@ "title": "DagRun", "type": "object" }, + "LoopContext": { + "description": "Pinned loop definition and enclosing iteration for a task execution.", + "properties": { + "node_id": { + "title": "Node Id", + "type": "string" + }, + "index": { + "minimum": 1, + "title": "Index", + "type": "integer" + }, + "max_iterations": { + "exclusiveMinimum": 0, + "title": "Max Iterations", + "type": "integer" + }, + "terminal_task_id": { + "title": "Terminal Task Id", + "type": "string" + }, + "terminal_is_mapped": { + "title": "Terminal Is Mapped", + "type": "boolean" + }, + "region_id": { + "format": "uuid", + "title": "Region Id", + "type": "string" + }, + "trajectory": { + "minimum": 1, + "title": "Trajectory", + "type": "integer" + }, + "forked_from_region_id": { + "format": "uuid", + "title": "Forked From Region Id", + "type": "string" + }, + "resumes_from_index": { + "title": "Resumes From Index", + "type": "integer" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped", + "region_id", + "trajectory" + ], + "title": "LoopContext", + "type": "object" + }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { @@ -5270,6 +5291,10 @@ ], "default": null }, + "expanded_ti_count": { + "title": "Expanded Ti Count", + "type": "integer" + }, "task_reschedule_count": { "default": 0, "title": "Task Reschedule Count", diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 8b018e485149d..75a05dde12cd6 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -236,6 +236,11 @@ export type Index = number; export type MaxIterations = number; export type TerminalTaskId = string; export type TerminalIsMapped = boolean; +export type RegionId1 = string; +export type Trajectory = number; +export type ForkedFromRegionId = string; +export type ResumesFromIndex = number; +export type ExpandedTiCount = number; export type TaskRescheduleCount = number; export type FirstTaskRescheduleStartDate = string | null; export type MaxTries = number; @@ -547,7 +552,7 @@ export type EndDate4 = string | null; export type State4 = string | null; export type TryNumber2 = number; export type MapIndex6 = number | null; -export type RegionId1 = string | null; +export type RegionId2 = string | null; export type RegionIndex1 = number | null; export type Duration = number | null; export type Type60 = "PreviousTIResult"; @@ -1051,6 +1056,7 @@ export interface TaskCallbackRequest { export interface TIRunContext { dag_run: DagRun; loop?: LoopContext | null; + expanded_ti_count?: ExpandedTiCount; task_reschedule_count?: TaskRescheduleCount; first_task_reschedule_start_date?: FirstTaskRescheduleStartDate; max_tries: MaxTries; @@ -1076,6 +1082,10 @@ export interface LoopContext { max_iterations: MaxIterations; terminal_task_id: TerminalTaskId; terminal_is_mapped: TerminalIsMapped; + region_id: RegionId1; + trajectory: Trajectory; + forked_from_region_id?: ForkedFromRegionId; + resumes_from_index?: ResumesFromIndex; } /** * Variable schema for responses with fields that are needed for Runtime. @@ -1686,7 +1696,7 @@ export interface PreviousTIResponse { state?: State4; try_number: TryNumber2; map_index?: MapIndex6; - region_id?: RegionId1; + region_id?: RegionId2; region_index?: RegionIndex1; duration?: Duration; }