You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by bb...@apache.org on 2026/10/20 23:00:51 UTC
(airflow) branch main updated: Number task loop iterations from one instead of zero (#74773)
This is an automated email from the ASF dual-hosted git repository.
bbovenzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new a8c55b78f0e Number task loop iterations from one instead of zero (#74773)
a8c55b78f0e is described below
commit a8c55b78f0e67875906cd5a87b0b158688843cce
Author: Brent Bovenzi <br...@astronomer.io>
AuthorDate: Tue Oct 20 19:00:41 2026 -0400
Number task loop iterations from one instead of zero (#74773)
A loop iteration is a number people read: it shows up in the loop filter,
on the Gantt tooltip, in the task instance list and in the `loop.index` a
Dag author branches on. Starting it at zero meant each of those readers had
to translate between the number Airflow stored and the number a person
would say, and that translation was easy to apply in one place and forget
in another.
Numbering from one in storage removes the translation instead of hiding it
behind a display-time adjustment, so the stored index, the API field and
the UI agree without anyone adding to it. Mapped expansion indexes share
the column and stay zero-based: a map index is a position, not a count, so
code reading the column has to know whether the region executes a loop or
an expansion.
Task loops have not shipped in a release, so no stored run and no client
depends on the old numbering, and the execution API version that
introduces the loop context is still unreleased.
Co-authored-by: Claude <no...@anthropic.com>
---
.../docs/authoring-and-scheduling/loops.rst | 21 ++--
.../core_api/datamodels/task_instances.py | 2 +-
.../api_fastapi/core_api/openapi/_private_ui.yaml | 2 +
.../core_api/openapi/v2-rest-api-generated.yaml | 10 +-
.../core_api/routes/public/task_instances.py | 5 +-
.../api_fastapi/core_api/services/ui/grid.py | 4 +-
.../execution_api/datamodels/taskinstance.py | 2 +-
.../execution_api/versions/v2026_10_30.py | 2 +-
.../src/airflow/example_dags/example_task_loops.py | 10 +-
airflow-core/src/airflow/models/dagrun.py | 5 +-
airflow-core/src/airflow/models/dynamic_region.py | 3 +-
airflow-core/src/airflow/models/taskinstance.py | 5 +
.../ui/openapi-gen/queries/ensureQueryData.ts | 2 +-
.../src/airflow/ui/openapi-gen/queries/prefetch.ts | 2 +-
.../src/airflow/ui/openapi-gen/queries/queries.ts | 2 +-
.../src/airflow/ui/openapi-gen/queries/suspense.ts | 2 +-
.../airflow/ui/openapi-gen/requests/schemas.gen.ts | 4 +-
.../ui/openapi-gen/requests/services.gen.ts | 2 +-
.../airflow/ui/openapi-gen/requests/types.gen.ts | 5 +-
.../layouts/Details/Gantt/GanttTimeline.test.tsx | 2 +-
.../ui/src/queries/useLoopFilterOptions.test.ts | 10 +-
.../core_api/routes/public/test_dag_run.py | 4 +-
.../core_api/routes/public/test_event_logs.py | 8 +-
.../core_api/routes/public/test_execution.py | 4 +-
.../core_api/routes/public/test_extra_links.py | 2 +-
.../core_api/routes/public/test_task_instances.py | 140 ++++++++++-----------
.../api_fastapi/core_api/routes/ui/test_gantt.py | 2 +-
.../api_fastapi/core_api/routes/ui/test_grid.py | 58 ++++++---
.../versions/head/test_task_instances.py | 50 ++++----
.../execution_api/versions/head/test_xcoms.py | 111 ++++++++--------
airflow-core/tests/unit/models/conftest.py | 7 +-
airflow-core/tests/unit/models/test_cleartasks.py | 42 +++----
airflow-core/tests/unit/models/test_dagrun.py | 107 ++++++++++------
.../tests/unit/models/test_dynamic_region.py | 52 ++++----
.../tests/unit/models/test_taskinstance.py | 21 ++--
.../unit/ti_deps/deps/test_trigger_rule_dep.py | 2 +-
.../src/airflowctl/api/datamodels/generated.py | 4 +-
go-sdk/schema/supervisor-schema.json | 2 +-
.../unit/common/ai/plugins/test_hitl_review.py | 16 +--
.../src/airflow/sdk/api/datamodels/_generated.py | 2 +-
task-sdk/src/airflow/sdk/execution_time/loop.py | 7 +-
.../airflow/sdk/execution_time/schema/schema.json | 2 +-
.../tests/task_sdk/execution_time/test_loop.py | 46 ++++---
ts-sdk/schema/supervisor-schema.json | 2 +-
44 files changed, 455 insertions(+), 338 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst
index b0ea5b97575..7a6e04e0c59 100644
--- a/airflow-core/docs/authoring-and-scheduling/loops.rst
+++ b/airflow-core/docs/authoring-and-scheduling/loops.rst
@@ -145,11 +145,11 @@ definition.
* - Attribute
- Meaning
* - ``loop.index``
- - The current iteration number, starting at zero.
+ - The current iteration number, starting at one.
* - ``loop.max_iterations``
- The configured limit on the number of iterations.
* - ``loop.previous``
- - The previous iteration's result, or ``None`` in iteration 0.
+ - 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.
@@ -159,7 +159,7 @@ task instance's position. Each task instance can have several tries, each with i
Pass data between iterations
============================
-Between iterations, use ``loop.previous``. Check ``loop.index == 0`` to handle
+Between iterations, use ``loop.previous``. Check ``loop.index == 1`` to handle
the first iteration. Do not test ``loop.previous is None``: it is also ``None``
when the previous terminal task returned nothing, and a previous result of
``0``, ``False``, or an empty collection can be valid data.
@@ -240,8 +240,8 @@ The task receives a sequence (``LazyXComSequence``) with one element for each
iteration the loop ran, in iteration order, so ``results[-1]`` is the last
iteration's result and ``len(results)`` is the number of iterations. A mapped
loop task gives one element per mapped instance for each iteration, ordered by
-iteration first and then by map index: all instances of iteration 0, then all
-instances of iteration 1, and so on. ``ti.xcom_pull(task_ids=...)`` and
+iteration first and then by map index: all instances of iteration 1, then all
+instances of iteration 2, and so on. ``ti.xcom_pull(task_ids=...)`` and
``{{ ti.xcom_pull(task_ids=...) }}`` in a template return the same sequence.
Only the current executions are included. Iterations that a rerun replaced do
@@ -290,9 +290,10 @@ first iteration, then each previous result plus one.
:start-after: [START mapped_loop]
:end-before: [END mapped_loop]
-The loop iteration and mapped position are separate "coordinates". In this
-example, each iteration has mapped instances 0 and 1, so mapped instance 1 in
-iteration 0 is distinct from mapped instance 1 in iteration 1.
+The loop iteration and mapped position are separate "coordinates", and they do
+not count from the same number: iterations start at 1 and map indexes start at
+0. In this example iteration 1 has mapped instances 0 and 1, so mapped instance
+1 in iteration 1 is distinct from mapped instance 1 in iteration 2.
Mapped instances appear within their loop iteration, so you can inspect their
states, tries, and logs separately.
@@ -385,14 +386,14 @@ Suppose each iteration contains this sequence:
prepare → process → consume → gate
-The loop has already run iterations 0 through 4. You clear ``process`` in
+The loop has already run iterations 1 through 4. You clear ``process`` in
iteration 2 with **Clear downstream tasks** and **Clear later loop iterations** both selected:
* ``prepare`` in iteration 2 remains completed.
* ``process``, ``consume``, and the gate in iteration 2 run again.
* Later iterations 3 and 4 are cleared. Whether replacement iterations run
depends on the new gate decisions.
-* Iterations 0 and 1 remain untouched.
+* Iteration 1 remains untouched.
If the gate now stops at iteration 2, replacement iterations 3 and 4 are not
created.
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
index 09ecefbf8c7..a082a718511 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
@@ -54,7 +54,7 @@ class LoopIterationResponse(BaseModel):
"""A task execution's position in an enclosing loop."""
loop_id: str
- iteration: int
+ iteration: int = Field(ge=1, description="Position in the enclosing loop, counting from one.")
class TaskInstanceResponse(OmitsMissingRegion, BaseModel):
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
index e7bc1eccb2c..847d6e8e04e 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
@@ -5042,7 +5042,9 @@ components:
title: Loop Id
iteration:
type: integer
+ minimum: 1.0
title: Iteration
+ description: Position in the enclosing loop, counting from one.
type: object
required:
- loop_id
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
index 271ed1d159a..640faa0942a 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -9223,11 +9223,13 @@ paths:
schema:
anyOf:
- type: integer
- minimum: 0
+ minimum: 1
- type: 'null'
- description: Narrow ``loop_id`` to a single pass of that loop.
+ description: Narrow ``loop_id`` to a single pass of that loop. Iterations
+ count from one.
title: Iteration
- description: Narrow ``loop_id`` to a single pass of that loop.
+ description: Narrow ``loop_id`` to a single pass of that loop. Iterations
+ count from one.
- name: cursor
in: query
required: false
@@ -16990,7 +16992,9 @@ components:
title: Loop Id
iteration:
type: integer
+ minimum: 1.0
title: Iteration
+ description: Position in the enclosing loop, counting from one.
type: object
required:
- loop_id
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
index 1fe5224a0d3..a6a33e75533 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
@@ -642,7 +642,10 @@ def get_task_instances(
] = None,
iteration: Annotated[
int | None,
- Query(ge=0, description="Narrow ``loop_id`` to a single pass of that loop."),
+ Query(
+ ge=1,
+ description="Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.",
+ ),
] = None,
cursor: str | None = Query(
None,
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
index 12cf433d84e..2eaa3296aad 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
@@ -443,7 +443,7 @@ def summarize_loop_run(
end_date=max(ends) if ends and all(s in State.finished for s in states) else None,
)
)
- latest = max(iterations, default=-1)
+ latest = max(iterations, default=0)
last_gate = next((ti for ti in iterations.get(latest, []) if ti.task_id == group.gate_task_id), None)
active = any(ti.state not in State.finished for tasks in iterations.values() for ti in tasks)
status: Literal["running", "stopped_early", "ran_to_cap", "failed", "skipped", "removed"] = "running"
@@ -459,7 +459,7 @@ def summarize_loop_run(
elif last_gate.state == TaskInstanceState.UPSTREAM_FAILED:
status = "failed"
elif last_gate.state == TaskInstanceState.SUCCESS:
- if latest + 1 == group.max_iterations:
+ if latest == group.max_iterations:
status = "ran_to_cap"
if not group.has_until:
reason = "cap_reached"
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 cda3ecd9a14..09ce1193e90 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
@@ -452,7 +452,7 @@ class LoopContext(BaseModel):
"""Pinned loop definition and enclosing iteration for a task execution."""
node_id: str
- index: int = Field(ge=0)
+ index: int = Field(ge=1)
max_iterations: int = Field(gt=0)
terminal_task_id: str
terminal_is_mapped: bool
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 7bf2a876298..26d0aea8878 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
@@ -38,7 +38,7 @@ from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
class AddLoopContext(VersionChange):
- """Provide the pinned loop definition and enclosing iteration at task startup."""
+ """Provide the pinned loop definition and enclosing iteration, numbered from one, at task startup."""
description = __doc__
instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("loop").didnt_exist,)
diff --git a/airflow-core/src/airflow/example_dags/example_task_loops.py b/airflow-core/src/airflow/example_dags/example_task_loops.py
index 8b178c4f6a2..5ee8984a0a7 100644
--- a/airflow-core/src/airflow/example_dags/example_task_loops.py
+++ b/airflow-core/src/airflow/example_dags/example_task_loops.py
@@ -28,7 +28,7 @@ def refine_estimate():
def refine():
@task
def improve(*, loop):
- estimate = 1.0 if loop.index == 0 else loop.previous["estimate"]
+ estimate = 1.0 if loop.index == 1 else loop.previous["estimate"]
return (estimate + 2.0 / estimate) / 2.0
@task
@@ -61,7 +61,7 @@ def fixed_task_loop():
def accumulate():
@task
def increment(*, loop):
- return (0 if loop.index == 0 else loop.previous) + 1
+ return (0 if loop.index == 1 else loop.previous) + 1
increment()
@@ -79,7 +79,7 @@ def mapped_task_loop():
def process_batch():
@task
def choose_items(*, loop):
- return [1, 2] if loop.index == 0 else [value + 1 for value in loop.previous]
+ return [1, 2] if loop.index == 1 else [value + 1 for value in loop.previous]
@task
def process(value, *, loop, ti):
@@ -89,7 +89,7 @@ def mapped_task_loop():
process.expand(value=choose_items())
def batch_ready(*, loop):
- return min(loop.result) >= 2
+ return min(loop.result) >= 3
process_batch.loop(max_iterations=3, until=batch_ready)
@@ -105,7 +105,7 @@ def partial_override_loop():
def accumulate(increment_by):
@task
def increment(increment_by, *, loop):
- return (0 if loop.index == 0 else loop.previous) + increment_by
+ return (0 if loop.index == 1 else loop.previous) + increment_by
increment(increment_by)
diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py
index fc116ab20bf..f28cd7690f4 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -154,7 +154,7 @@ class InvalidLoopDecision(ValueError):
def validate_loop_decision(stop: bool | None, region_index: int, group: SerializedLoopTaskGroup) -> None:
"""Raise :class:`InvalidLoopDecision` unless ``stop`` is valid for the gate of pass ``region_index``."""
- at_limit = region_index + 1 >= group.max_iterations
+ at_limit = region_index >= group.max_iterations
if stop is None:
raise InvalidLoopDecision("Successful loop gate requires a decision")
if not stop and at_limit:
@@ -2433,7 +2433,8 @@ class DagRun(Base, LoggingMixin):
region = DynamicRegion.get_or_create(
dag_id=self.dag_id, run_id=self.run_id, node_id=loops[group_id].node_id, session=session
)
- coordinates = [(region.id, 0)]
+ # Iterations count from one; a mapped expansion's slots still count from zero.
+ coordinates = [(region.id, 1)]
for region_id, region_index in coordinates:
if DynamicRegion.is_coordinate_superseded(region_id, region_index, regions):
continue
diff --git a/airflow-core/src/airflow/models/dynamic_region.py b/airflow-core/src/airflow/models/dynamic_region.py
index b2100d95734..968b9224829 100644
--- a/airflow-core/src/airflow/models/dynamic_region.py
+++ b/airflow-core/src/airflow/models/dynamic_region.py
@@ -742,7 +742,8 @@ def _build_candidate_query(
raise ValueError("Caller is not inside the requested loop")
if context.previous_iteration:
position = position[0], position[1] - 1
- if position[1] < 0:
+ # Iterations count from one, so the first pass has no predecessor to look for.
+ if position[1] < 1:
return None
query = _filter_producers(
diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py
index f569fc4f36f..4690e8adf46 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -1107,6 +1107,11 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
task_id: Mapped[str] = mapped_column(StringID(), nullable=False)
dag_id: Mapped[str] = mapped_column(StringID(), nullable=False)
run_id: Mapped[str] = mapped_column(StringID(), nullable=False)
+ # The slot this instance fills in its region, so its meaning follows the region's node: a loop
+ # region numbers its iterations from one and a mapped expansion numbers its slots from zero.
+ # -1 covers both a plain task, which belongs to no region at all and is the common case, and a
+ # mapped task that has not been expanded yet. Code that reads the column without knowing which
+ # node the region executes cannot tell an iteration from a map index.
region_index: Mapped[int] = mapped_column(Integer, nullable=False, server_default="-1")
map_index = synonym("region_index")
region_id: Mapped[UUID] = mapped_column(
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
index 76133b1f85e..bb6c50f3507 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -1380,7 +1380,7 @@ export const ensureUseTaskInstanceServiceGetMappedTaskInstanceData = (queryClien
* @param data.regionId
* @param data.regionIndex
* @param data.loopId Task group id of a loop; matches every instance it produced.
-* @param data.iteration Narrow ``loop_id`` to a single pass of that loop.
+* @param data.iteration Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
* @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored.
* @param data.taskId
* @param data.runAfterGte
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
index 7d646a03e8e..661826be84c 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -1380,7 +1380,7 @@ export const prefetchUseTaskInstanceServiceGetMappedTaskInstance = (queryClient:
* @param data.regionId
* @param data.regionIndex
* @param data.loopId Task group id of a loop; matches every instance it produced.
-* @param data.iteration Narrow ``loop_id`` to a single pass of that loop.
+* @param data.iteration Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
* @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored.
* @param data.taskId
* @param data.runAfterGte
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
index d2e58f17ca1..d620a72ace7 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -1380,7 +1380,7 @@ export const useTaskInstanceServiceGetMappedTaskInstance = <TData = Common.TaskI
* @param data.regionId
* @param data.regionIndex
* @param data.loopId Task group id of a loop; matches every instance it produced.
-* @param data.iteration Narrow ``loop_id`` to a single pass of that loop.
+* @param data.iteration Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
* @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored.
* @param data.taskId
* @param data.runAfterGte
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
index 7c83461a6ab..5e7a0cef977 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -1380,7 +1380,7 @@ export const useTaskInstanceServiceGetMappedTaskInstanceSuspense = <TData = Comm
* @param data.regionId
* @param data.regionIndex
* @param data.loopId Task group id of a loop; matches every instance it produced.
-* @param data.iteration Narrow ``loop_id`` to a single pass of that loop.
+* @param data.iteration Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
* @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored.
* @param data.taskId
* @param data.runAfterGte
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
index b1fb2e97730..2eb5f596b63 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -6708,7 +6708,9 @@ export const $LoopIterationResponse = {
},
iteration: {
type: 'integer',
- title: 'Iteration'
+ minimum: 1,
+ title: 'Iteration',
+ description: 'Position in the enclosing loop, counting from one.'
}
},
type: 'object',
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index 5e5fb1e981d..369e8e386f0 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -3136,7 +3136,7 @@ export class TaskInstanceService {
* @param data.regionId
* @param data.regionIndex
* @param data.loopId Task group id of a loop; matches every instance it produced.
- * @param data.iteration Narrow ``loop_id`` to a single pass of that loop.
+ * @param data.iteration Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
* @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored.
* @param data.taskId
* @param data.runAfterGte
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index 5321f56062b..59af66a1f7a 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -1833,6 +1833,9 @@ export type LastAssetEventResponse = {
*/
export type LoopIterationResponse = {
loop_id: string;
+ /**
+ * Position in the enclosing loop, counting from one.
+ */
iteration: number;
};
@@ -4657,7 +4660,7 @@ export type GetTaskInstancesData = {
endDateLte?: string | null;
executor?: Array<(string)>;
/**
- * Narrow ``loop_id`` to a single pass of that loop.
+ * Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
*/
iteration?: number | null;
limit?: number;
diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/Gantt/GanttTimeline.test.tsx b/airflow-core/src/airflow/ui/src/layouts/Details/Gantt/GanttTimeline.test.tsx
index 7d5b4daf2ec..3ea4123a68d 100644
--- a/airflow-core/src/airflow/ui/src/layouts/Details/Gantt/GanttTimeline.test.tsx
+++ b/airflow-core/src/airflow/ui/src/layouts/Details/Gantt/GanttTimeline.test.tsx
@@ -123,7 +123,7 @@ describe("GanttTimeline segment bars", () => {
vi.unstubAllGlobals();
});
- it.each([0, 3])("shows loop iteration %s in the hovered bar's tooltip", async (iteration) => {
+ it.each([1, 3])("shows loop iteration %s in the hovered bar's tooltip", async (iteration) => {
const segment: GanttDataItem = {
loopIteration: { iteration, loop_id: "loop" },
regionId: "11111111-1111-1111-1111-111111111111",
diff --git a/airflow-core/src/airflow/ui/src/queries/useLoopFilterOptions.test.ts b/airflow-core/src/airflow/ui/src/queries/useLoopFilterOptions.test.ts
index c4693477494..4da7331d327 100644
--- a/airflow-core/src/airflow/ui/src/queries/useLoopFilterOptions.test.ts
+++ b/airflow-core/src/airflow/ui/src/queries/useLoopFilterOptions.test.ts
@@ -65,7 +65,7 @@ describe("loopsForNode", () => {
const pass = (iteration: number | undefined): LoopOption => ({ iteration, loopId: "batch" }) as LoopOption;
// [whole loop, then the passes newest first]
-const OPTIONS = [pass(undefined), pass(2), pass(1), pass(0)];
+const OPTIONS = [pass(undefined), pass(3), pass(2), pass(1)];
describe("selectActivePass", () => {
it("takes the pass the server resolved for the row", () => {
@@ -79,17 +79,17 @@ describe("selectActivePass", () => {
});
it("falls back to the newest pass when nothing pins one", () => {
- expect(selectActivePass(OPTIONS, new URLSearchParams())?.iteration).toBe(2);
+ expect(selectActivePass(OPTIONS, new URLSearchParams())?.iteration).toBe(3);
});
it("reads a pass pinned as a filter", () => {
- expect(selectActivePass(OPTIONS, new URLSearchParams("iteration=0"))?.iteration).toBe(0);
+ expect(selectActivePass(OPTIONS, new URLSearchParams("iteration=1"))?.iteration).toBe(1);
});
it("ignores a pass belonging to a different loop", () => {
- const params = new URLSearchParams("loop_id=other&iteration=0");
+ const params = new URLSearchParams("loop_id=other&iteration=1");
- expect(selectActivePass(OPTIONS, params)?.iteration).toBe(2);
+ expect(selectActivePass(OPTIONS, params)?.iteration).toBe(3);
});
});
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
index c0e3bf9ede8..6c6aa57990d 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
@@ -2373,7 +2373,7 @@ class TestClearDagRun:
run_id=run.run_id,
dag_version_id=run.created_dag_version_id,
region_id=region.id,
- region_index=1,
+ region_index=2,
)
later.state = State.SUCCESS
session.add(later)
@@ -2385,7 +2385,7 @@ class TestClearDagRun:
assert response.status_code == 200, response.text
work = [row for row in response.json()["task_instances"] if row["task_id"] == "body.work"]
- assert {row["region_index"] for row in work} == {0, 1}
+ assert {row["region_index"] for row in work} == {1, 2}
assert {row["map_index"] for row in work} == {-1}
assert all(
row["loop_iteration"] == {"loop_id": loop.group_id, "iteration": row["region_index"]}
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_event_logs.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_event_logs.py
index f6b0c1bcebb..4a9865b0df5 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_event_logs.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_event_logs.py
@@ -409,7 +409,7 @@ class TestGetEventLogs(TestEventLogsEndpoint):
dag_id=run.dag_id, run_id=run.run_id, node_id=loop.group_id, session=session
)
passes = {}
- for index in (1, 2):
+ for index in (2, 3):
passes[index] = TaskInstance(
task=dag.get_task("body.work"),
run_id=run.run_id,
@@ -419,7 +419,7 @@ class TestGetEventLogs(TestEventLogsEndpoint):
)
session.add(passes[index])
session.flush()
- archived = passes[2]
+ archived = passes[3]
retried = archived.prepare_db_for_next_try(session=session)
events = {
name: Log(
@@ -428,7 +428,7 @@ class TestGetEventLogs(TestEventLogsEndpoint):
for name, ti_id in [
("first try", archived.id),
("second try", retried.id),
- ("other pass", passes[1].id),
+ ("other pass", passes[2].id),
("unattributed", None),
]
}
@@ -436,7 +436,7 @@ class TestGetEventLogs(TestEventLogsEndpoint):
session.commit()
response = test_client.get(
- "/eventLogs", params={"region_id": str(region.id), "region_index": 2, "task_id": "body.work"}
+ "/eventLogs", params={"region_id": str(region.id), "region_index": 3, "task_id": "body.work"}
)
assert response.status_code == 200, response.text
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_execution.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_execution.py
index 39f73e0c77c..dbf7a1b67ac 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_execution.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_execution.py
@@ -128,7 +128,7 @@ def test_loop_repeated_selective_clears_archive_only_replaced_task_instances(
initial = live_executions()
prepared = next(ti for ti in initial if ti["task_id"] == "prepare")
archived_ids = set()
- for index, keep_future in [(1, retain_later), (0, False)]:
+ for index, keep_future in [(2, retain_later), (1, False)]:
before = live_executions()
selected = next(
ti for ti in before if ti["task_id"] == "body.work" and ti.get("region_index") == index
@@ -163,7 +163,7 @@ def test_loop_repeated_selective_clears_archive_only_replaced_task_instances(
params={"region_id": ti["region_id"], "region_index": ti["region_index"]},
)
assert value.status_code == 200, value.text
- assert value.json()["value"] == ti["region_index"] + 1
+ assert value.json()["value"] == ti["region_index"]
archived = session.scalars(
select(TaskInstance)
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py
index 0bb45fb62a8..91fb81b48ed 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py
@@ -142,7 +142,7 @@ def test_serialized_links_read_xcom_from_selected_region(test_client, dag_maker,
with dag_maker(serialized=True) as dag:
create_loop(body, max_iterations=4)
run = dag_maker.create_dagrun()
- sibling = next(ti for ti in run.task_instances if ti.task_id == "body.work" and ti.region_index == 0)
+ sibling = next(ti for ti in run.task_instances if ti.task_id == "body.work" and ti.region_index == 1)
source_region = session.get(DynamicRegion, sibling.region_id)
selected_region = DynamicRegion(
dag_id=run.dag_id,
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index 251adadc7ae..67d9a6751a3 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -1518,7 +1518,7 @@ class TestGetTaskInstances(TestTaskInstanceEndpoint):
assert response.status_code == 200, response.text
assert response.json()["total_entries"] == 2
assert {ti["region_index"] for ti in response.json()["task_instances"]} == {2, 7}
- response = test_client.get(url, params={**params, "iteration": 0, "limit": 100})
+ response = test_client.get(url, params={**params, "iteration": 1, "limit": 100})
assert response.json()["total_entries"] == 2
# region_id names one region exactly; the logical scope of a pass is what loop_id and
# iteration are for, so this does not reach the expansion nested inside it.
@@ -3613,12 +3613,12 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
run_id=dr.run_id,
dag_version_id=dr.created_dag_version_id,
region_id=root.id,
- region_index=1,
+ region_index=2,
state=State.SUCCESS,
)
)
session.commit()
- suffix_ids = {ti.id for ti in dr.get_task_instances(session=session) if ti.region_index == 1}
+ suffix_ids = {ti.id for ti in dr.get_task_instances(session=session) if ti.region_index == 2}
dag_id, run_id, gate_id = dr.dag_id, dr.run_id, loop.gate_task_id
with dag_maker(dag_id=dag_id, serialized=True, session=session):
MockOperator(task_id="replacement")
@@ -3838,7 +3838,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
DynamicRegion.dag_id == dr.dag_id, DynamicRegion.node_id == loop.group_id
)
)
- for index in (1, 2):
+ for index in (2, 3):
for loop_task in loop.iter_tasks():
session.add(
TaskInstance(
@@ -3852,7 +3852,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
)
session.commit()
tis = list(dr.get_task_instances(session=session))
- seed = next(ti for ti in tis if ti.task_id == "body.member" and ti.region_index == 1)
+ seed = next(ti for ti in tis if ti.task_id == "body.member" and ti.region_index == 2)
if only_failed:
seed.state = State.FAILED
session.commit()
@@ -3879,8 +3879,8 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
ti.id
for ti in tis
if ti.task_id == "outside"
- or (ti.region_index == 1 and ti.task_id == loop.gate_task_id)
- or (later and ti.region_id == root.id and ti.region_index > 1)
+ or (ti.region_index == 2 and ti.task_id == loop.gate_task_id)
+ or (later and ti.region_id == root.id and ti.region_index > 2)
)
assert {
(row["task_id"], row.get("region_id"), row.get("region_index"))
@@ -3899,7 +3899,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
elif state in (State.SUCCESS, State.FAILED):
historical = session.get(TaskInstance, identity)
assert historical.working_set is None
- archived = coordinates[identity][2] > 1 and downstream and later
+ archived = coordinates[identity][2] > 2 and downstream and later
assert historical.archived_reason == ("superseded" if archived else "retry")
def test_loop_clear_reports_conflict_when_a_worker_archived_the_execution_first(
@@ -3924,7 +3924,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
run_id=dr.run_id,
dag_version_id=dr.created_dag_version_id,
region_id=root.id,
- region_index=1,
+ region_index=2,
state=State.SUCCESS,
)
)
@@ -3932,7 +3932,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
gate = next(
ti
for ti in dr.get_task_instances(session=session)
- if ti.task_id == loop.gate_task_id and ti.region_index == 0
+ if ti.task_id == loop.gate_task_id and ti.region_index == 1
)
mocker.patch.object(
TaskInstance,
@@ -3976,7 +3976,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
DynamicRegion.dag_id == dr.dag_id, DynamicRegion.node_id == loop.group_id
)
)
- for index in (1, 2, 3):
+ for index in (2, 3, 4):
for loop_task in loop.iter_tasks():
session.add(
TaskInstance(
@@ -3989,7 +3989,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
)
)
for ti in dr.get_task_instances(session=session):
- ti.state = State.FAILED if only_failed and ti.region_index in (1, 2) else State.SUCCESS
+ ti.state = State.FAILED if only_failed and ti.region_index in (2, 3) else State.SUCCESS
session.commit()
task_id = loop.gate_task_id if seed_task == "gate" else seed_task
@@ -3998,7 +3998,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
select(TaskInstance).where(
TaskInstance.run_id == dr.run_id,
TaskInstance.task_id == task_id,
- TaskInstance.region_index == 1,
+ TaskInstance.region_index == 2,
TaskInstance.working_set.is_(True),
)
)
@@ -4012,7 +4012,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
test_client.post(f"/dags/{dr.dag_id}/clearTaskInstances", json={**build_payload(), "dry_run": False})
for ti in session.scalars(select(TaskInstance).where(TaskInstance.working_set.is_(True))):
- ti.state = State.FAILED if only_failed and ti.region_index in (1, 2) else State.SUCCESS
+ ti.state = State.FAILED if only_failed and ti.region_index in (2, 3) else State.SUCCESS
session.commit()
payload = build_payload()
archived_before = set(
@@ -5870,7 +5870,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
loop = create_loop(body, max_iterations=3, until=lambda loop: True)
dr = dag_maker.create_dagrun(run_id="manual-run")
root = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id == loop.group_id))
- for index in (1, 2):
+ for index in (2, 3):
for operator in loop.iter_tasks():
session.add(
TaskInstance(
@@ -5894,7 +5894,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
):
dr, loop, root, tis = loop_instances
backend = MetastoreBackend()
- for index in (1, 2):
+ for index in (2, 3):
backend.set(
TaskScope(
dag_id=dr.dag_id,
@@ -5914,7 +5914,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
json={
"dry_run": False,
"dag_run_id": dr.run_id,
- "task_instance_ids": [str(tis["body.first", 1].id)],
+ "task_instance_ids": [str(tis["body.first", 2].id)],
"keep_task_state": keep_task_state,
},
)
@@ -5928,7 +5928,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
TaskStateStoreModel.task_id == "body.first",
)
)
- ) == ({1, 2} if keep_task_state else {2})
+ ) == ({2, 3} if keep_task_state else {3})
@pytest.mark.parametrize("dry_run", [False, True])
def test_exact_loop_mark_does_not_change_other_iterations(
@@ -5939,26 +5939,26 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
response = test_client.patch(
url + ("/dry_run" if dry_run else ""),
- json={"new_state": "success", "region_id": str(root.id), "region_index": 1},
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 2},
)
assert response.status_code == 200, response.text
assert [
(item["task_id"], item["map_index"], item["region_index"])
for item in response.json()["task_instances"]
- ] == [("body.first", -1, 1)]
+ ] == [("body.first", -1, 2)]
session.expire_all()
- for index in (0, 2):
+ for index in (1, 3):
assert tis["body.first", index].state == State.FAILED
assert tis["body.last", index].state == State.FAILED
- assert tis["body.first", 1].state == (State.FAILED if dry_run else State.SUCCESS)
+ assert tis["body.first", 2].state == (State.FAILED if dry_run else State.SUCCESS)
current_last = session.scalar(
select(TaskInstance).where(
TaskInstance.dag_id == dr.dag_id,
TaskInstance.run_id == dr.run_id,
TaskInstance.task_id == "body.last",
TaskInstance.region_id == root.id,
- TaskInstance.region_index == 1,
+ TaskInstance.region_index == 2,
TaskInstance.working_set.is_(True),
)
)
@@ -5981,7 +5981,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
start >> loop >> MockOperator(task_id="finish")
dr = dag_maker.create_dagrun()
root = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id == loop.group_id))
- for index in (1, 2):
+ for index in (2, 3):
for loop_task in loop.iter_tasks():
session.add(
TaskInstance(
@@ -6002,7 +6002,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
json={
"new_state": "failed",
"region_id": str(root.id),
- "region_index": 1,
+ "region_index": 2,
"include_upstream": True,
"include_downstream": True,
},
@@ -6013,10 +6013,10 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
(item["task_id"], item.get("region_index")) for item in response.json()["task_instances"]
} == {
("start", None),
- ("body.prepare", 1),
- ("body.process", 1),
- ("body.combine", 1),
- (loop.gate_task_id, 1),
+ ("body.prepare", 2),
+ ("body.process", 2),
+ ("body.combine", 2),
+ (loop.gate_task_id, 2),
("finish", None),
}
@@ -6050,7 +6050,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
[
pytest.param(7, 400, id="conflicting"),
pytest.param(-1, 200, id="default"),
- pytest.param(1, 200, id="matching"),
+ pytest.param(2, 200, id="matching"),
],
)
def test_loop_mark_rejects_map_index_conflicting_with_region_index(
@@ -6060,7 +6060,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
response = test_client.patch(
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.first/{map_index}",
- json={"new_state": "success", "region_id": str(root.id), "region_index": 1},
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 2},
)
assert response.status_code == expected_status, response.text
@@ -6072,7 +6072,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
self, test_client, session, loop_instances, coordinates_in_query
):
dr, loop, root, tis = loop_instances
- coordinates = {"region_id": str(root.id), "region_index": 1}
+ coordinates = {"region_id": str(root.id), "region_index": 2}
response = test_client.patch(
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.first",
@@ -6082,20 +6082,20 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
assert response.status_code == 200, response.text
session.expire_all()
- assert tis["body.first", 1].note == "selected iteration"
- assert tis["body.first", 0].note is None
- assert tis["body.first", 2].note is None
+ assert tis["body.first", 2].note == "selected iteration"
+ assert tis["body.first", 1].note is None
+ assert tis["body.first", 3].note is None
def test_manual_gate_success_after_rewind_stops_and_keeps_the_reason(
self, test_client, session, loop_instances
):
dr, loop, root, tis = loop_instances
- clear_loop_task_instances([tis[loop.gate_task_id, 0]], downstream=False, session=session)
+ clear_loop_task_instances([tis[loop.gate_task_id, 1]], downstream=False, session=session)
gate = session.scalars(
select(TaskInstance).where(
TaskInstance.dag_id == dr.dag_id,
TaskInstance.task_id == loop.gate_task_id,
- TaskInstance.region_index == 0,
+ TaskInstance.region_index == 1,
TaskInstance.working_set.is_(True),
)
).one()
@@ -6105,13 +6105,13 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
response = test_client.patch(
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/{gate.task_id}",
- json={"new_state": "success", "region_id": str(root.id), "region_index": 0},
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 1},
)
assert response.status_code == 200, response.text
session.expire_all()
assert session.get(TaskInstance, gate_id).state == State.SUCCESS
- assert all(ti.region_index == 0 for ti in dr.get_task_instances(session=session))
+ assert all(ti.region_index == 1 for ti in dr.get_task_instances(session=session))
assert session.scalars(
select(XComModelV2.value).where(XComModelV2.task_instance_id == gate_id)
).all() == ["stale"]
@@ -6121,17 +6121,17 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
self, test_client, session, loop_instances, outcome
):
dr, loop, root, tis = loop_instances
- archiving = tis["body.first", 2]
+ archiving = tis["body.first", 3]
archiving.state = State.RUNNING
session.flush()
archiving_id = archiving.id
- clear_loop_task_instances([tis[loop.gate_task_id, 0]], downstream=False, session=session)
+ clear_loop_task_instances([tis[loop.gate_task_id, 1]], downstream=False, session=session)
assert archiving.state == State.RESTARTING
session.commit()
response = test_client.patch(
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/{archiving.task_id}",
- json={"new_state": outcome, "region_id": str(root.id), "region_index": 2},
+ json={"new_state": outcome, "region_id": str(root.id), "region_index": 3},
)
assert response.status_code == 200, response.text
@@ -6146,15 +6146,15 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
@pytest.mark.parametrize(
("body", "params"),
[
- ({"region_index": 1}, {}),
- ({}, {"region_index": 1}),
+ ({"region_index": 2}, {}),
+ ({}, {"region_index": 2}),
(
- {"region_index": 1, "region_id": "00000000-0000-0000-0000-000000000000"},
{"region_index": 2, "region_id": "00000000-0000-0000-0000-000000000000"},
+ {"region_index": 3, "region_id": "00000000-0000-0000-0000-000000000000"},
),
(
{
- "region_index": 1,
+ "region_index": 2,
"region_id": "00000000-0000-0000-0000-000000000000",
"include_past": True,
},
@@ -6178,7 +6178,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
self, test_client, session, loop_instances, action
):
dr, loop, root, tis = loop_instances
- ids = {str(tis["body.first", index].id) for index in (0, 2)}
+ ids = {str(tis["body.first", index].id) for index in (1, 3)}
entities = [
{
"task_id": "body.first",
@@ -6187,7 +6187,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
"region_index": index,
**({"new_state": "success"} if action == "update" else {}),
}
- for index in (0, 2)
+ for index in (1, 3)
]
response = test_client.patch(
@@ -6199,7 +6199,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
assert not response.json()[action]["errors"]
assert set(response.json()[action]["success"]) == ids
session.expire_all()
- assert tis["body.first", 1].state == State.FAILED
+ assert tis["body.first", 2].state == State.FAILED
live = {str(ti.id): ti.state for ti in dr.get_task_instances(session=session)}
if action == "delete":
assert not ids & live.keys()
@@ -6211,7 +6211,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
self, test_client, session, loop_instances, via
):
dr, loop, root, tis = loop_instances
- archived = tis["body.first", 2]
+ archived = tis["body.first", 3]
retried = archived.prepare_db_for_next_try(session)
session.commit()
pass_rows = (
@@ -6221,7 +6221,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
TaskInstance.run_id == dr.run_id,
TaskInstance.task_id == "body.first",
TaskInstance.region_id == root.id,
- TaskInstance.region_index == 2,
+ TaskInstance.region_index == 3,
)
.execution_options(include_all_attempts=True)
)
@@ -6239,7 +6239,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
"task_id": "body.first",
"map_index": -1,
"region_id": str(root.id),
- "region_index": 2,
+ "region_index": 3,
}
],
}
@@ -6249,7 +6249,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
else:
response = test_client.delete(
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.first",
- params={"region_id": str(root.id), "region_index": 2},
+ params={"region_id": str(root.id), "region_index": 3},
)
assert response.status_code == 200, response.text
@@ -6258,7 +6258,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
session.expire_all()
assert set(session.scalars(pass_rows)) == set()
remaining = {(ti.task_id, ti.region_index) for ti in dr.get_task_instances(session=session)}
- assert remaining == set(tis) - {("body.first", 2)}
+ assert remaining == set(tis) - {("body.first", 3)}
def test_bulk_regional_entities_do_not_lock_their_run_again(self, test_client, loop_instances, mocker):
dr, loop, root, tis = loop_instances
@@ -6271,7 +6271,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
"region_index": index,
"new_state": "success",
}
- for index in (0, 2)
+ for index in (1, 3)
]
response = test_client.patch(
@@ -6289,16 +6289,16 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
response = test_client.patch(
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskGroupInstances/body"
+ ("/dry_run" if dry_run else ""),
- json={"new_state": "success", "region_id": str(root.id), "region_index": 1},
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 2},
)
assert response.status_code == 200, response.text
assert {item["id"] for item in response.json()["task_instances"]} == {
- str(ti.id) for (task_id, index), ti in tis.items() if index == 1
+ str(ti.id) for (task_id, index), ti in tis.items() if index == 2
}
session.expire_all()
for (_task_id, index), ti in tis.items():
- assert ti.state == (State.SUCCESS if index == 1 and not dry_run else State.FAILED)
+ assert ti.state == (State.SUCCESS if index == 2 and not dry_run else State.FAILED)
@pytest.mark.parametrize("action", ["update", "delete"])
@pytest.mark.parametrize("map_index", [None, -1])
@@ -6345,7 +6345,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
response = test_client.patch(
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskGroupInstances/body",
- json={"new_state": "success", "region_id": str(root.id), "region_index": 0},
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 1},
)
assert response.status_code == 200, response.text
@@ -6376,7 +6376,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
region_id=root.id,
region_index=index,
)
- for index in (0, 1)
+ for index in (1, 2)
]
for scope in scopes:
backend.set(scope, "job_id", "external-job", session=session)
@@ -6390,7 +6390,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
"new_state": "success",
"note": "completed manually",
"region_id": str(root.id),
- "region_index": 1,
+ "region_index": 2,
},
)
@@ -6404,7 +6404,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
TaskStateStoreModel.task_id == "body.first",
)
)
- ) == {0}
+ ) == {1}
@pytest.mark.parametrize("dry_run", [False, True])
def test_regional_downstream_mark_uses_same_scope_as_preview(
@@ -6419,17 +6419,17 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
"new_state": "success",
"include_downstream": True,
"region_id": str(root.id),
- "region_index": 1,
+ "region_index": 2,
},
)
assert response.status_code == 200, response.text
assert {item["id"] for item in response.json()["task_instances"]} == {
- str(ti.id) for (_task_id, index), ti in tis.items() if index == 1
+ str(ti.id) for (_task_id, index), ti in tis.items() if index == 2
}
session.expire_all()
assert all(
- ti.state == (State.SUCCESS if index == 1 and not dry_run else State.FAILED)
+ ti.state == (State.SUCCESS if index == 2 and not dry_run else State.FAILED)
for (_task_id, index), ti in tis.items()
)
@@ -6449,29 +6449,29 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
response = test_client.patch(
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskGroupInstances/body",
- json={"new_state": "success", "region_id": str(root.id), "region_index": 1},
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 2},
)
assert response.status_code == 200, response.text
assert {item["id"] for item in response.json()["task_instances"]} == {
- str(ti.id) for (_task_id, index), ti in tis.items() if index == 1
+ str(ti.id) for (_task_id, index), ti in tis.items() if index == 2
}
def test_group_terminal_response_retains_executions_archived_by_manual_settlement(
self, test_client, session, loop_instances
):
dr, loop, root, tis = loop_instances
- draining = [tis["body.first", 2], tis[loop.gate_task_id, 2]]
+ draining = [tis["body.first", 3], tis[loop.gate_task_id, 3]]
for ti in draining:
ti.state = State.RUNNING
session.flush()
ids = {str(ti.id) for ti in draining}
- clear_loop_task_instances([tis[loop.gate_task_id, 0]], downstream=False, session=session)
+ clear_loop_task_instances([tis[loop.gate_task_id, 1]], downstream=False, session=session)
session.commit()
response = test_client.patch(
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskGroupInstances/body",
- json={"new_state": "skipped", "region_id": str(root.id), "region_index": 2},
+ json={"new_state": "skipped", "region_id": str(root.id), "region_index": 3},
)
assert response.status_code == 200, response.text
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_gantt.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_gantt.py
index b03710137fa..78ee974b5f5 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_gantt.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_gantt.py
@@ -400,7 +400,7 @@ class TestGetGanttDataEndpoint:
assert response.status_code == 200, response.text
rows = [row for row in response.json()["task_instances"] if row["task_id"] == "body.work"]
assert {row["id"] for row in rows} == {str(first.id), str(second.id)}
- assert {row["region_index"] for row in rows} == {0, 2}
+ assert {row["region_index"] for row in rows} == {1, 2}
assert {row["region_id"] for row in rows} == {str(first.region_id)}
assert {row["map_index"] for row in rows} == {-1}
assert all(
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
index 87be7d87815..65e46cd11de 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
@@ -1333,11 +1333,11 @@ class TestGetLoopSummaryEndpoint:
("until", "expected"),
[
(
- lambda loop: LoopDecision(stop=loop.index >= 1, reason=f"iteration {loop.index} decided"),
- ["iteration 0 decided", "iteration 1 decided"],
+ lambda loop: LoopDecision(stop=loop.index >= 2, reason=f"iteration {loop.index} decided"),
+ ["iteration 1 decided", "iteration 2 decided"],
),
- (lambda loop: loop.index >= 1, [None, None]),
- (lambda loop: LoopDecision(stop=loop.index >= 1), [None, None]),
+ (lambda loop: loop.index >= 2, [None, None]),
+ (lambda loop: LoopDecision(stop=loop.index >= 2), [None, None]),
],
ids=["with_reason", "plain_bool", "decision_without_reason"],
)
@@ -1402,7 +1402,7 @@ class TestGetLoopSummaryEndpoint:
assert summary["exit_criteria_name"] is None
assert "exit_task_id" not in summary
assert summary["iterations_ran"] == 3
- assert [iteration["index"] for iteration in summary["iterations"]] == [0, 1, 2]
+ assert [iteration["index"] for iteration in summary["iterations"]] == [1, 2, 3]
family = summary["loop_region_id"]
assert summary["loop_regions"] == [{"region_id": family}]
assert test_client.get(path, params={"loop_region_id": str(uuid4())}).status_code == 404
@@ -1412,13 +1412,13 @@ class TestGetLoopSummaryEndpoint:
TaskInstance.dag_id == run.dag_id,
TaskInstance.run_id == run.run_id,
TaskInstance.task_id == self.gate_task_id,
- TaskInstance.region_index == 1,
+ TaskInstance.region_index == 2,
)
)
clear_loop_task_instances([gate], session=session, later_loop_iterations=True)
session.commit()
current = test_client.get(path).json()
- assert [iteration["index"] for iteration in current["iterations"]] == [0, 1]
+ assert [iteration["index"] for iteration in current["iterations"]] == [1, 2]
assert current["loop_region_id"] == family
assert current["status"] == "running"
@@ -1434,13 +1434,13 @@ class TestGetLoopSummaryEndpoint:
TaskInstance.dag_id == run.dag_id,
TaskInstance.run_id == run.run_id,
TaskInstance.task_id == self.gate_task_id,
- TaskInstance.region_index == 0,
+ TaskInstance.region_index == 1,
)
)
clear_loop_task_instances([gate], session=session, later_loop_iterations=True)
fork = session.scalar(select(DynamicRegion).where(DynamicRegion.forked_from_region_id.is_not(None)))
dag = DBDagBag().get_dag(run.created_dag_version_id, session=session)
- for index in (1, 2):
+ for index in (2, 3):
for task_id in ("body.work", self.gate_task_id):
session.add(
TaskInstance(
@@ -1458,7 +1458,7 @@ class TestGetLoopSummaryEndpoint:
assert summary["loop_region_id"] == family
assert summary["loop_regions"] == [{"region_id": family}]
- assert [iteration["index"] for iteration in summary["iterations"]] == [0, 1, 2]
+ assert [iteration["index"] for iteration in summary["iterations"]] == [1, 2, 3]
assert test_client.get(path, params={"loop_region_id": str(fork.id)}).status_code == 404
@pytest.mark.parametrize(("gate_state", "expected"), [("skipped", "skipped"), ("failed", "failed")])
@@ -1472,7 +1472,7 @@ class TestGetLoopSummaryEndpoint:
TaskInstance.dag_id == run.dag_id,
TaskInstance.run_id == run.run_id,
TaskInstance.task_id == self.gate_task_id,
- TaskInstance.region_index == 2,
+ TaskInstance.region_index == 3,
)
)
gate.state = TaskInstanceState(gate_state)
@@ -1576,7 +1576,7 @@ class TestGetLoopSummaryEndpoint:
TaskInstance.dag_id == run.dag_id,
TaskInstance.run_id == run.run_id,
TaskInstance.task_id == "body.work",
- TaskInstance.region_index == 0,
+ TaskInstance.region_index == 1,
)
):
ti.state = TaskInstanceState.RUNNING
@@ -1593,7 +1593,7 @@ class TestGetLoopSummaryEndpoint:
assert by_task_id["body"]["child_states"]["success"] >= 2
# The coordinate is recorded per row, and a link still has to address one, so it stays
# on the newest pass even though an earlier one is the thing now running.
- assert by_task_id["body.work"]["latest_region_index"] == 2
+ assert by_task_id["body.work"]["latest_region_index"] == 3
def test_loop_history_reads_runs_in_a_fixed_number_of_queries(self, test_client, dag_maker, session):
@task
@@ -1617,13 +1617,13 @@ class TestGetLoopSummaryEndpoint:
session.add(region)
session.flush()
for ti in run.task_instances:
- ti.region_id, ti.region_index, ti.state = region.id, 0, TaskInstanceState.SUCCESS
+ ti.region_id, ti.region_index, ti.state = region.id, 1, TaskInstanceState.SUCCESS
later = TaskInstance(
task=dag.get_task(ti.task_id),
run_id=run.run_id,
dag_version_id=ti.dag_version_id,
region_id=region.id,
- region_index=1,
+ region_index=2,
state=TaskInstanceState.FAILED
if index == 1 and ti.task_id == "body.work"
else TaskInstanceState.SUCCESS,
@@ -1743,7 +1743,7 @@ class TestGetLoopSummaryEndpoint:
response = test_client.get(f"/grid/loop/{run.dag_id}/{run.run_id}/body")
assert response.status_code == 200, response.text
assert response.json()["iterations_ran"] == 2
- assert [iteration["index"] for iteration in response.json()["iterations"]] == [0, 1]
+ assert [iteration["index"] for iteration in response.json()["iterations"]] == [1, 2]
response = test_client.get(f"/grid/ti_summaries/{run.dag_id}", params={"run_ids": [run.run_id]})
assert response.status_code == 200, response.text
tasks = json.loads(response.text)["task_instances"]
@@ -1781,6 +1781,32 @@ class TestGetLoopSummaryEndpoint:
assert response.json()["max_iterations"] == 3
assert [iteration["state"] for iteration in response.json()["iterations"]] == ["removed"] * 3
+ def test_summary_separates_stopping_below_the_cap_from_reaching_it(self, test_client, dag_maker):
+ """A loop that converges one pass below its cap stopped early; it did not run to the cap."""
+
+ @task
+ def work():
+ return 1
+
+ @task_group
+ def body():
+ work()
+
+ def converged(*, loop):
+ return loop.index == 2
+
+ with dag_maker(serialized=False) as dag:
+ body.loop(max_iterations=3, until=converged)
+ run = dag.test()
+
+ summary = test_client.get(f"/grid/loop/{run.dag_id}/{run.run_id}/body").json()
+
+ # A zero-based cap check reads the second of three passes as the cap and says ran_to_cap.
+ assert summary["status"] == "stopped_early"
+ assert summary["stopped_at_iteration"] == 2
+ assert summary["reason"] is None
+ assert summary["iterations_ran"] == 2
+
def test_condition_docstring_reaches_loop_summary_and_structure(self, test_client, dag_maker):
@task
def work():
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 58b77d63757..6a81588ffc5 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
@@ -2329,7 +2329,7 @@ class TestTIUpdateState:
session.expire_all()
assert session.get(TaskInstance, right_id).state == State.RUNNING
ready = dr.task_instance_scheduling_decisions(session=session).schedulable_tis
- assert {(ti.task_id, ti.region_index) for ti in ready} == {("body.left", 1), ("body.right", 1)}
+ assert {(ti.task_id, ti.region_index) for ti in ready} == {("body.left", 2), ("body.right", 2)}
@pytest.fixture
def started_loop_successor(self, client, session, running_loop_gate):
@@ -2344,11 +2344,11 @@ class TestTIUpdateState:
running = next(
ti
for ti in dr.get_task_instances(session=session)
- if ti.task_id != gate.task_id and ti.region_index == 1
+ if ti.task_id != gate.task_id and ti.region_index == 2
)
running.state = State.QUEUED
for ti in dr.get_task_instances(session=session):
- if ti.task_id != gate.task_id and ti.region_index == 0:
+ if ti.task_id != gate.task_id and ti.region_index == 1:
ti.state = State.SUCCESS
session.commit()
running_id = running.id
@@ -2409,7 +2409,7 @@ class TestTIUpdateState:
dr.dag = dag_maker.serialized_dag
decision = dr.task_instance_scheduling_decisions(session=session)
assert gate.id in {ti.id for ti in decision.schedulable_tis}
- assert not any(ti.region_index == 1 for ti in dr.get_task_instances(session=session))
+ assert not any(ti.region_index == 2 for ti in dr.get_task_instances(session=session))
def test_restart_ack_locks_run_before_task(self, client, session, create_task_instance, mocker):
if session.bind.dialect.name == "sqlite":
@@ -2499,17 +2499,17 @@ class TestTIUpdateState:
@pytest.mark.parametrize(
("decision", "state", "max_iterations", "conditional", "rejected", "passes"),
[
- ("continue", State.SUCCESS, 3, False, False, [0, 1]),
- ("stop", State.SUCCESS, 3, False, True, [0]),
- (None, State.SUCCESS, 3, False, True, [0]),
- ("continue", State.FAILED, 3, False, False, [0]),
- ("continue", State.SKIPPED, 3, False, False, [0]),
- ("stop", State.SUCCESS, 1, False, False, [0]),
- ("continue", State.SUCCESS, 1, False, True, [0]),
- ("stop", State.SUCCESS, 3, True, False, [0]),
- ("stop", State.SUCCESS, 1, True, False, [0]),
- ("continue", State.SUCCESS, 1, True, True, [0]),
- (None, State.FAILED, 1, True, False, [0]),
+ ("continue", State.SUCCESS, 3, False, False, [1, 2]),
+ ("stop", State.SUCCESS, 3, False, True, [1]),
+ (None, State.SUCCESS, 3, False, True, [1]),
+ ("continue", State.FAILED, 3, False, False, [1]),
+ ("continue", State.SKIPPED, 3, False, False, [1]),
+ ("stop", State.SUCCESS, 1, False, False, [1]),
+ ("continue", State.SUCCESS, 1, False, True, [1]),
+ ("stop", State.SUCCESS, 3, True, False, [1]),
+ ("stop", State.SUCCESS, 1, True, False, [1]),
+ ("continue", State.SUCCESS, 1, True, True, [1]),
+ (None, State.FAILED, 1, True, False, [1]),
],
)
def test_loop_gate_completion_consumes_decision_atomically(
@@ -2581,8 +2581,8 @@ class TestTIUpdateState:
)
== passes
)
- if mapped and passes == [0, 1]:
- successor_region = next(region for region in children if region.parent_region_index == 1)
+ if mapped and passes == [1, 2]:
+ successor_region = next(region for region in children if region.parent_region_index == 2)
assert [ti.region_index for ti in current if ti.region_id == successor_region.id] == [-1]
dr.dag = dag_maker.serialized_dag
@@ -2637,7 +2637,7 @@ class TestTIUpdateState:
assert response.status_code == 500
session.expire_all()
- assert observed == [State.RUNNING, [0]]
+ assert observed == [State.RUNNING, [1]]
assert gate.state == State.RUNNING
assert len(dr.get_task_instances(session=session)) == 2
@@ -2696,7 +2696,7 @@ class TestTIUpdateState:
assert gate.state == State.SUCCESS
assert sorted(
ti.region_index for ti in dr.get_task_instances(session=session) if ti.task_id == gate.task_id
- ) == [0, 1]
+ ) == [1, 2]
def test_loop_gate_completion_uses_its_pinned_limit(self, client, session, running_loop_gate, dag_maker):
gate = running_loop_gate
@@ -2716,7 +2716,7 @@ class TestTIUpdateState:
assert response.status_code == 204, response.text
session.expire_all()
- successors = [ti for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1]
+ successors = [ti for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 2]
assert len(successors) == 2
assert all(ti.dag_version_id == original_version for ti in successors)
@@ -2746,7 +2746,7 @@ class TestTIUpdateState:
assert response.status_code == 204, response.text
session.expire_all()
assert gate.state == State.FAILED
- assert [ti.region_index for ti in gate.dag_run.get_task_instances(session=session)] == [0, 0]
+ assert [ti.region_index for ti in gate.dag_run.get_task_instances(session=session)] == [1, 1]
def test_loop_gate_state_update_error_rolls_back_next_pass(
self, client, session, running_loop_gate, mocker
@@ -2822,7 +2822,7 @@ class TestTIUpdateState:
assert response.status_code == 204
session.refresh(gate)
successor_ids = {
- ti.id for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1
+ ti.id for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 2
}
archived_id = gate.id
dag_run = gate.dag_run
@@ -2833,7 +2833,7 @@ class TestTIUpdateState:
TaskInstance.dag_id == dag_run.dag_id,
TaskInstance.run_id == dag_run.run_id,
TaskInstance.task_id == gate.task_id,
- TaskInstance.region_index == 0,
+ TaskInstance.region_index == 1,
TaskInstance.working_set.is_(True),
)
)
@@ -2854,7 +2854,7 @@ class TestTIUpdateState:
session.expire_all()
assert gate.state == state
assert {
- ti.id for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1
+ ti.id for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 2
} == successor_ids
assert _stored_loop_decision_reasons(session, archived_id) == ["first"]
assert _stored_loop_decision_reasons(session, gate.id) == []
@@ -4810,7 +4810,7 @@ class TestTISkipDownstream:
other_pass = TaskInstance(
task=dag.get_task(target.task_id), run_id=dr.run_id, dag_version_id=target.dag_version_id
)
- other_pass.region_id, other_pass.region_index = region.id, 1
+ other_pass.region_id, other_pass.region_index = region.id, 2
session.add(other_pass)
session.commit()
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
index 7768fe75336..0685fc67aec 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
@@ -151,6 +151,19 @@ def test_previous_iteration_xcom_uses_retained_predecessor(client, loop_xcoms, s
assert client.head(url, params=params).headers["Content-Range"] == "map_indexes 1"
+def test_previous_iteration_of_the_first_pass_finds_nothing(client, loop_xcoms, authenticate_as):
+ """Asking the first pass for its predecessor asks for a pass below the first, which has none."""
+ dr, _, _, previous = loop_xcoms
+ authenticate_as(previous)
+
+ response = client.get(
+ f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{previous.task_id}/key",
+ params={"previous_iteration": True},
+ )
+
+ assert response.status_code == 404
+
+
def test_previous_iteration_rejects_prior_dates(client, loop_xcoms):
dr, producer, _, _ = loop_xcoms
@@ -182,7 +195,7 @@ def test_previous_mapped_iteration_reads_retained_slots_before_count_and_slice(
run_id=dr.run_id,
node_id="body.mapped",
parent_region_id=parent.id,
- parent_region_index=0,
+ parent_region_index=1,
session=session,
)
current = DynamicRegion.get_or_create(
@@ -190,7 +203,7 @@ def test_previous_mapped_iteration_reads_retained_slots_before_count_and_slice(
run_id=dr.run_id,
node_id="body.mapped",
parent_region_id=parent.id,
- parent_region_index=1,
+ parent_region_index=2,
session=session,
)
session.add_all([previous, current])
@@ -200,13 +213,13 @@ def test_previous_mapped_iteration_reads_retained_slots_before_count_and_slice(
run_id=dr.run_id,
node_id="body.mapped",
parent_region_id=parent.id,
- parent_region_index=0,
+ parent_region_index=1,
forked_from_region_id=previous.id,
)
session.add(replacement)
session.flush()
consumer = next(ti for ti in dr.task_instances if ti.task_id == "body.consumer")
- consumer.region_id, consumer.region_index = parent.id, 1
+ consumer.region_id, consumer.region_index = parent.id, 2
for ti in dr.task_instances:
if ti.task_id == "body.mapped":
ti.region_id = current.id
@@ -242,8 +255,8 @@ def test_previous_mapped_iteration_reads_retained_slots_before_count_and_slice(
LOOP_SEQUENCES = {
- "unmapped": ("body.step", ["step-0", "step-1", "step-2"]),
- "mapped": ("body.fanout", ["0-0", "0-1", "1-0", "1-1", "2-0", "2-1"]),
+ "unmapped": ("body.step", ["step-1", "step-2", "step-3"]),
+ "mapped": ("body.fanout", ["1-0", "1-1", "2-0", "2-1", "3-0", "3-1"]),
}
@@ -271,9 +284,9 @@ def outside_loop_run(dag_maker, session, authenticate_as):
)
)
tis = {(ti.task_id, ti.region_id, ti.region_index): ti for ti in dr.task_instances}
- steps = {0: tis["body.step", root.id, 0]}
- fanouts = {(0, ti.region_index): ti for ti in dr.task_instances if ti.task_id == "body.fanout"}
- for iteration in (1, 2):
+ steps = {1: tis["body.step", root.id, 1]}
+ fanouts = {(1, ti.region_index): ti for ti in dr.task_instances if ti.task_id == "body.fanout"}
+ for iteration in (2, 3):
steps[iteration] = TaskInstance(
dag.get_task("body.step"),
dr.created_dag_version_id,
@@ -301,7 +314,7 @@ def outside_loop_run(dag_maker, session, authenticate_as):
)
session.add(fanouts[iteration, slot])
session.flush()
- for iteration in (2, 1, 0):
+ for iteration in (3, 2, 1):
XComModel.set_for_attempt(
task_instance_id=steps[iteration].id,
key="key",
@@ -359,13 +372,13 @@ def test_parsing_attempt_id_does_not_select_a_loop_iteration(
dr, _, _, steps, _ = outside_loop_run
token = ExecutionToken.model_validate(
{
- "id": steps[0].id,
+ "id": steps[1].id,
"claims": {
"scope": "dag_parse",
"exp": 2000000000,
"dag_bundles": [session.get(DagModel, dr.dag_id).bundle_name],
"job_id": 1,
- "session_id": steps[0].id,
+ "session_id": steps[1].id,
"relative_fileloc": "dag.py",
},
}
@@ -393,8 +406,8 @@ def test_outside_task_that_uses_loop_outputs_waits_for_the_gate_and_reads_the_lo
assert consumer.upstream_task_ids == {"body.__loop_gate"}
url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/body.step/{XCOM_RETURN_KEY}"
assert client.head(url).headers["Content-Range"] == "map_indexes 3"
- assert client.get(f"{url}/slice").json() == ["returned-0", "returned-1", "returned-2"]
- assert client.get(f"{url}/item/-1").json() == "returned-2"
+ assert client.get(f"{url}/slice").json() == ["returned-1", "returned-2", "returned-3"]
+ assert client.get(f"{url}/item/-1").json() == "returned-3"
@pytest.mark.parametrize("shape", LOOP_SEQUENCES)
@@ -430,7 +443,7 @@ def test_outside_caller_reads_a_forked_loop_from_the_region_that_owns_each_pass(
run_id=dr.run_id,
node_id="body",
forked_from_region_id=root.id,
- resumes_from_index=2,
+ resumes_from_index=3,
)
nested = DynamicRegion(
dag_id=dr.dag_id,
@@ -441,7 +454,7 @@ def test_outside_caller_reads_a_forked_loop_from_the_region_that_owns_each_pass(
)
session.add(fork)
session.flush()
- nested.parent_region_id, nested.parent_region_index = fork.id, 2
+ nested.parent_region_id, nested.parent_region_index = fork.id, 3
session.add(nested)
session.flush()
replacements = [
@@ -450,7 +463,7 @@ def test_outside_caller_reads_a_forked_loop_from_the_region_that_owns_each_pass(
dr.created_dag_version_id,
run_id=dr.run_id,
region_id=fork.id,
- region_index=2,
+ region_index=3,
),
*(
TaskInstance(
@@ -465,15 +478,15 @@ def test_outside_caller_reads_a_forked_loop_from_the_region_that_owns_each_pass(
]
session.add_all(replacements)
session.flush()
- for ti, value in zip(replacements, ["step-2-again", "2-0-again", "2-1-again"]):
+ for ti, value in zip(replacements, ["step-3-again", "3-0-again", "3-1-again"]):
XComModel.set_for_attempt(
task_instance_id=ti.id, key="key", value=value, serialize=False, session=session
)
session.commit()
task_id = LOOP_SEQUENCES[shape][0]
expected = {
- "unmapped": ["step-0", "step-1", "step-2-again"],
- "mapped": ["0-0", "0-1", "1-0", "1-1", "2-0-again", "2-1-again"],
+ "unmapped": ["step-1", "step-2", "step-3-again"],
+ "mapped": ["1-0", "1-1", "2-0", "2-1", "3-0-again", "3-1-again"],
}[shape]
url = get_loop_url(dr, task_id)
@@ -504,7 +517,7 @@ def test_outside_caller_cannot_read_a_loop_with_several_executions(client, sessi
run_id=dr.run_id,
node_id="body",
parent_region_id=root.id,
- parent_region_index=0,
+ parent_region_index=1,
)
)
session.commit()
@@ -517,15 +530,15 @@ def test_outside_caller_cannot_read_a_loop_with_several_executions(client, sessi
def test_loop_task_reads_from_inside_the_loop_stay_on_one_pass(client, outside_loop_run):
dr, _, _, steps, authenticate_as = outside_loop_run
- authenticate_as(steps[1])
+ authenticate_as(steps[2])
step_url = get_loop_url(dr, "body.step")
fanout_url = get_loop_url(dr, "body.fanout")
assert client.head(step_url).headers["Content-Range"] == "map_indexes 1"
- assert client.get(f"{step_url}/slice").json() == ["step-1"]
- assert client.get(f"{step_url}/slice", params={"previous_iteration": True}).json() == ["step-0"]
- assert client.get(f"{fanout_url}/slice").json() == ["1-0", "1-1"]
- assert client.get(f"{fanout_url}/item/-1").json() == "1-1"
+ assert client.get(f"{step_url}/slice").json() == ["step-2"]
+ assert client.get(f"{step_url}/slice", params={"previous_iteration": True}).json() == ["step-1"]
+ assert client.get(f"{fanout_url}/slice").json() == ["2-0", "2-1"]
+ assert client.get(f"{fanout_url}/item/-1").json() == "2-1"
@pytest.fixture
@@ -535,8 +548,8 @@ def deferred_loop_pass(session, outside_loop_run):
trigger.triggerer_id = TRIGGERER_ID
session.add(trigger)
session.flush()
- steps[1].state = TaskInstanceState.DEFERRED
- steps[1].trigger_id = trigger.id
+ steps[2].state = TaskInstanceState.DEFERRED
+ steps[2].trigger_id = trigger.id
session.commit()
yield dr, steps, trigger
clear_db_triggers()
@@ -544,7 +557,7 @@ def deferred_loop_pass(session, outside_loop_run):
def test_trigger_in_a_loop_reads_and_writes_as_its_deferred_pass(session, deferred_loop_pass):
dr, steps, _ = deferred_loop_pass
- headers = {TRIGGER_TI_ID_HEADER: str(steps[1].id)}
+ headers = {TRIGGER_TI_ID_HEADER: str(steps[2].id)}
step_url = f"/xcoms/{dr.dag_id}/{dr.run_id}/body.step/key"
sessions_begun = []
@@ -556,10 +569,10 @@ def test_trigger_in_a_loop_reads_and_writes_as_its_deferred_pass(session, deferr
event.listen(Session, "after_begin", record)
try:
with capture_orm_selects("task_instance") as task_instance_selects:
- assert api.get(step_url, headers=headers).json() == {"key": "key", "value": "step-1"}
+ assert api.get(step_url, headers=headers).json() == {"key": "key", "value": "step-2"}
assert api.head(step_url, headers=headers).headers["Content-Range"] == "map_indexes 1"
- assert api.get(f"{step_url}/slice", headers=headers).json() == ["step-1"]
- assert api.get(f"{step_url}/item/0", headers=headers).json() == "step-1"
+ assert api.get(f"{step_url}/slice", headers=headers).json() == ["step-2"]
+ assert api.get(f"{step_url}/item/0", headers=headers).json() == "step-2"
assert api.post(f"{step_url}-written", json="from-trigger", headers=headers).status_code == 201
assert api.delete(step_url, headers=headers).status_code == 200
finally:
@@ -573,9 +586,9 @@ def test_trigger_in_a_loop_reads_and_writes_as_its_deferred_pass(session, deferr
assert len(liveness_checks) == 1
assert len(set(map(id, sessions_begun))) == 6
session.expire_all()
- assert XComModelV2.get_for_attempt(steps[1].id, "key-written", session=session).value == "from-trigger"
- assert XComModelV2.get_for_attempt(steps[1].id, "key", session=session) is None
- assert XComModelV2.get_for_attempt(steps[0].id, "key", session=session).value == "step-0"
+ assert XComModelV2.get_for_attempt(steps[2].id, "key-written", session=session).value == "from-trigger"
+ assert XComModelV2.get_for_attempt(steps[2].id, "key", session=session) is None
+ assert XComModelV2.get_for_attempt(steps[1].id, "key", session=session).value == "step-1"
@pytest.mark.parametrize(
@@ -591,7 +604,7 @@ def test_trigger_ti_id_header_rejected(deferred_loop_pass, api_triggerer_id, hea
with TestClient(InProcessExecutionAPI(triggerer_id=api_triggerer_id).app) as api:
response = api.get(
f"/xcoms/{dr.dag_id}/{dr.run_id}/body.step/key",
- headers={TRIGGER_TI_ID_HEADER: header or str(steps[1].id)},
+ headers={TRIGGER_TI_ID_HEADER: header or str(steps[2].id)},
)
assert response.status_code == expected_status
@@ -632,7 +645,7 @@ NOT_A_LIVE_ATTEMPT = "Task instance is not a live attempt"
"get",
{"state": TaskInstanceState.FAILED, "trigger_removed": True},
200,
- "step-1",
+ "step-2",
None,
id="read-live-task-no-longer-deferred",
),
@@ -652,20 +665,20 @@ def test_trigger_ti_id_claim_check(
):
dr, steps, trigger = deferred_loop_pass
trigger.triggerer_id = change.get("triggerer_id", TRIGGERER_ID)
- steps[1].state = change.get("state", TaskInstanceState.DEFERRED)
+ steps[2].state = change.get("state", TaskInstanceState.DEFERRED)
if change.get("trigger_removed"):
- steps[1].trigger_id = None
+ steps[2].trigger_id = None
session.delete(trigger)
session.commit()
if change.get("archived"):
- steps[1].archive(reason="cleared", session=session)
+ steps[2].archive(reason="cleared", session=session)
session.commit()
with TestClient(InProcessExecutionAPI(triggerer_id=TRIGGERER_ID).app) as api:
response = api.request(
method,
f"/xcoms/{dr.dag_id}/{dr.run_id}/body.step/key",
- headers={TRIGGER_TI_ID_HEADER: str(steps[1].id)},
+ headers={TRIGGER_TI_ID_HEADER: str(steps[2].id)},
json="from-trigger" if method == "post" else None,
)
@@ -697,15 +710,15 @@ def test_loop_result_item_read_loads_no_producer_rows(dag_maker, session, client
create_loop(body, max_iterations=3)
dr = dag_maker.create_dagrun()
tis = {(ti.task_id, ti.map_index): ti for ti in dr.task_instances}
- consumer = tis["body.consumer", 0]
+ consumer = tis["body.consumer", 1]
loop_region = DynamicRegion.get_or_create(
dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session
)
session.add(loop_region)
session.flush()
- consumer.region_id, consumer.region_index = loop_region.id, 1
+ consumer.region_id, consumer.region_index = loop_region.id, 2
mapped_task = dag.get_task("body.mapped")
- for iteration in range(3):
+ for iteration in range(1, 4):
region = DynamicRegion.get_or_create(
dag_id=dr.dag_id,
run_id=dr.run_id,
@@ -717,7 +730,7 @@ def test_loop_result_item_read_loads_no_producer_rows(dag_maker, session, client
session.add(region)
session.flush()
for map_index in range(3):
- ti = tis.get(("body.mapped", map_index)) if iteration == 0 else None
+ ti = tis.get(("body.mapped", map_index)) if iteration == 1 else None
if ti is None:
ti = TaskInstance(mapped_task, consumer.dag_version_id, run_id=dr.run_id, map_index=map_index)
session.add(ti)
@@ -745,7 +758,7 @@ def test_loop_result_item_read_loads_no_producer_rows(dag_maker, session, client
finally:
event.remove(Session, "loaded_as_persistent", record)
- assert response.json() == "1-0"
+ assert response.json() == "2-0"
assert loaded == []
@@ -813,8 +826,8 @@ def test_xcom_change_to_a_loop_task_from_outside_the_loop_is_rejected(client, si
@pytest.mark.parametrize(
("caller", "target", "params"),
[
- pytest.param(("body.step", 0), "body.step", {}, id="own-xcom-inside-loop"),
- pytest.param(("body.step", 0), "outside", {}, id="from-inside-loop-to-outside"),
+ pytest.param(("body.step", 1), "body.step", {}, id="own-xcom-inside-loop"),
+ pytest.param(("body.step", 1), "outside", {}, id="from-inside-loop-to-outside"),
pytest.param(("outside", -1), "mapped", {"map_index": 1}, id="mapped-task-outside-loops"),
],
)
@@ -854,7 +867,7 @@ def test_write_and_delete_conflict_when_several_task_instances_match(
def test_archived_attempt_in_a_loop_keeps_writing_its_own_xcom(client, single_pass_loop, session):
dr, tis, authenticate_as = single_pass_loop
- step = tis["body.step", 0]
+ step = tis["body.step", 1]
authenticate_as(step)
step.archive(reason="test", session=session)
session.commit()
diff --git a/airflow-core/tests/unit/models/conftest.py b/airflow-core/tests/unit/models/conftest.py
index bc18469dcac..a9397e8460a 100644
--- a/airflow-core/tests/unit/models/conftest.py
+++ b/airflow-core/tests/unit/models/conftest.py
@@ -46,7 +46,8 @@ def loop_run(dag_maker, session):
loop >> EmptyOperator(task_id="outside")
dr = dag_maker.create_dagrun()
root = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id == loop.group_id))
- for index in range(1, 5):
+ # The run already holds the first iteration; add the rest up to the cap.
+ for index in range(2, 6):
child = DynamicRegion(
dag_id=dr.dag_id,
run_id=dr.run_id,
@@ -108,6 +109,6 @@ def completed_loop(dag_maker, session):
)
session.flush()
- for index in range(5):
- complete_pass(index, "stop" if index == 4 else "continue")
+ for index in range(1, 6):
+ complete_pass(index, "stop" if index == 5 else "continue")
return dr, loop, complete_pass
diff --git a/airflow-core/tests/unit/models/test_cleartasks.py b/airflow-core/tests/unit/models/test_cleartasks.py
index 95d52e56301..5ed43c424ec 100644
--- a/airflow-core/tests/unit/models/test_cleartasks.py
+++ b/airflow-core/tests/unit/models/test_cleartasks.py
@@ -251,7 +251,7 @@ def test_partition_clear_keeps_loop_runs_whole_across_batches(loop_run, session)
assert result == (1, len(original))
for ti_id, task_id, pass_index in original:
archived = session.get(TaskInstance, ti_id)
- later_pass = task_id != "outside" and pass_index >= 1
+ later_pass = task_id != "outside" and pass_index >= 2
assert (archived.working_set, archived.archived_reason) == (
None,
"superseded" if later_pass else "retry",
@@ -298,7 +298,7 @@ def test_dry_run_clear_lists_later_loop_passes_that_clearing_a_gate_archives(
assert {ti.id for ti in listed} == {
ti.id
for ti in tis
- if ti.task_id == loop.gate_task_id or (ti.task_id != "outside" and (iteration(ti) or 0) >= 1)
+ if ti.task_id == loop.gate_task_id or (ti.task_id != "outside" and (iteration(ti) or 0) >= 2)
}
assert len(listed) == len({ti.id for ti in listed})
@@ -2179,7 +2179,7 @@ def test_loop_upstream_and_downstream_scope_excludes_other_descendants_of_upstre
start >> loop >> EmptyOperator(task_id="finish")
dr = dag_maker.create_dagrun()
root = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id == loop.group_id))
- for index in (1, 2):
+ for index in (2, 3):
for loop_task in loop.iter_tasks():
session.add(
TaskInstance(
@@ -2193,18 +2193,18 @@ def test_loop_upstream_and_downstream_scope_excludes_other_descendants_of_upstre
session.flush()
tis = {(ti.task_id, ti.region_index): ti for ti in dr.get_task_instances(session=session)}
- scope = select_loop_clear_scope([tis["body.process", 1]], upstream=True, downstream=True, session=session)
+ scope = select_loop_clear_scope([tis["body.process", 2]], upstream=True, downstream=True, session=session)
assert {key for key, ti in tis.items() if ti.id in scope.retry_ids} == {
("start", -1),
- ("body.prepare", 1),
- ("body.process", 1),
- ("body.combine", 1),
- (loop.gate_task_id, 1),
+ ("body.prepare", 2),
+ ("body.process", 2),
+ ("body.combine", 2),
+ (loop.gate_task_id, 2),
("finish", -1),
}
assert {key for key, ti in tis.items() if ti.id in scope.archive_ids} == {
- (loop_task.task_id, 2) for loop_task in loop.iter_tasks()
+ (loop_task.task_id, 3) for loop_task in loop.iter_tasks()
}
@@ -2214,8 +2214,8 @@ def test_outside_seed_upstream_selects_all_current_loop_occurrences(loop_run, se
scope = select_loop_clear_scope([outside], upstream=True, downstream=False, session=session)
- assert scope.retry_ids == {ti.id for ti in tis if ti.task_id == "outside" or iteration(ti) == 0}
- assert scope.archive_ids == {ti.id for ti in tis if ti.task_id != "outside" and iteration(ti) > 0}
+ assert scope.retry_ids == {ti.id for ti in tis if ti.task_id == "outside" or iteration(ti) == 1}
+ assert scope.archive_ids == {ti.id for ti in tis if ti.task_id != "outside" and iteration(ti) > 1}
@pytest.fixture
@@ -2235,13 +2235,13 @@ def mapped_loop_after_predecessor(dag_maker, session):
run_id=dr.run_id,
node_id="body.process",
parent_region_id=root.id,
- parent_region_index=1,
+ parent_region_index=2,
)
session.add(child)
session.flush()
for loop_task in loop.iter_tasks():
mapped = loop_task.task_id == "body.process"
- for index in range(2) if mapped else [1]:
+ for index in range(2) if mapped else [2]:
session.add(
TaskInstance(
task=loop_task,
@@ -2288,7 +2288,7 @@ def test_loop_upstream_and_downstream_scope_excludes_other_map_indexes(
assert scope.retry_ids == {seed.id} | {
ti.id
for ti in tis
- if ti.task_id == "before" or (ti.task_id == loop.gate_task_id and ti.region_index == 1)
+ if ti.task_id == "before" or (ti.task_id == loop.gate_task_id and ti.region_index == 2)
}
assert not scope.archive_ids
@@ -2442,19 +2442,19 @@ class TestWholeRunClear:
clear_run(dr, session)
live = dr.get_task_instances(session=session)
- assert {ti.region_index for ti in live} == {0}
+ assert {ti.region_index for ti in live} == {1}
for (task_id, index), ti_id in original.items():
archived = session.get(TaskInstance, ti_id)
- if index == 0:
+ if index == 1:
assert (archived.working_set, archived.archived_reason) == (None, "retry"), task_id
else:
assert (archived.working_set, archived.archived_reason) == (None, "superseded"), task_id
fork = session.scalar(select(DynamicRegion).where(DynamicRegion.forked_from_region_id.is_not(None)))
- assert fork.resumes_from_index == 1
+ assert fork.resumes_from_index == 2
- complete_pass(0, "continue")
+ complete_pass(1, "continue")
- regenerated = [ti for ti in dr.get_task_instances(session=session) if ti.region_index == 1]
+ regenerated = [ti for ti in dr.get_task_instances(session=session) if ti.region_index == 2]
assert len(regenerated) == 4
assert {ti.region_id for ti in regenerated} == {fork.id}
@@ -2462,9 +2462,9 @@ class TestWholeRunClear:
dr, loop, complete_pass = completed_loop
clear_run(dr, session)
- complete_pass(0, "stop")
+ complete_pass(1, "stop")
- assert {ti.region_index for ti in dr.get_task_instances(session=session)} == {0}
+ assert {ti.region_index for ti in dr.get_task_instances(session=session)} == {1}
def test_selected_state_patch_locks_dag_run_before_task_instances(completed_loop, session):
diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py
index 5c191505666..ee63b2db379 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -62,9 +62,11 @@ from airflow.models.dagbag import DBDagBag
from airflow.models.dagrun import (
DagRun,
DagRunNote,
+ InvalidLoopDecision,
clear_partition_runs,
get_or_create_dagrun,
lock_dag_runs,
+ validate_loop_decision,
)
from airflow.models.deadline import Deadline
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
@@ -99,6 +101,7 @@ from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.definitions.callback import AsyncCallback
from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference, VariableInterval
from airflow.serialization.definitions.deadline import SerializedReferenceModels
+from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup
from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG
from airflow.settings import get_policy_plugin_manager
from airflow.task.trigger_rule import TriggerRule
@@ -156,7 +159,7 @@ def test_loop_first_pass_uses_one_region_and_mapped_child(dag_maker, session, va
assert {region.node_id for region in regions} == {loop.group_id, "body.mapped"}
parent = next(region for region in regions if region.node_id == loop.group_id)
mapped_region = next(region for region in regions if region.node_id == "body.mapped")
- assert (mapped_region.parent_region_id, mapped_region.parent_region_index) == (parent.id, 0)
+ assert (mapped_region.parent_region_id, mapped_region.parent_region_index) == (parent.id, 1)
original_ids = {ti.id for ti in dr.task_instances}
for ti in dr.task_instances:
if ti.task_id == "outside":
@@ -164,7 +167,7 @@ def test_loop_first_pass_uses_one_region_and_mapped_child(dag_maker, session, va
elif ti.task_id == "body.mapped":
assert ti.region_id == mapped_region.id
else:
- assert (ti.region_id, ti.region_index) == (parent.id, 0)
+ assert (ti.region_id, ti.region_index) == (parent.id, 1)
dr.verify_integrity(session=session, dag_version_id=dr.created_dag_version_id)
@@ -237,7 +240,7 @@ def test_loop_new_member_joins_live_pass_after_reserialization(dag_maker, sessio
current = dr.get_task_instances(session=session)
assert {ti.id for ti in current if ti.task_id != "body.new"} == original_ids
added = next(ti for ti in current if ti.task_id == "body.new")
- assert (added.region_id, added.region_index, added.dag_version_id) == (gate.region_id, 0, version_id)
+ assert (added.region_id, added.region_index, added.dag_version_id) == (gate.region_id, 1, version_id)
assert session.get(TI, gate.id).state == State.SUCCESS
@@ -329,7 +332,7 @@ def test_loop_runs_through_scheduler_sdk_and_api(dag_maker, session, mode):
terminal(1)
def converged(*, loop):
- return mode == "until_true" and loop.index == 1
+ return mode == "until_true" and loop.index == 2
with dag_maker(serialized=False) as dag:
loop = body.loop(max_iterations=2, until=converged if mode.startswith("until") else None)
@@ -343,17 +346,51 @@ def test_loop_runs_through_scheduler_sdk_and_api(dag_maker, session, mode):
if ti.task_id == loop.gate_task_id
)
if mode == "empty":
- assert gates == [(0, State.SKIPPED)]
+ assert gates == [(1, State.SKIPPED)]
elif mode == "until_false":
- assert gates == [(0, State.SUCCESS), (1, State.FAILED)]
+ assert gates == [(1, State.SUCCESS), (2, State.FAILED)]
elif mode == "unhandled_failure":
- assert gates == [(0, State.UPSTREAM_FAILED)]
+ assert gates == [(1, State.UPSTREAM_FAILED)]
else:
- assert gates == [(0, State.SUCCESS), (1, State.SUCCESS)]
+ assert gates == [(1, State.SUCCESS), (2, State.SUCCESS)]
outside = dr.get_task_instance(task_id="outside", session=session)
assert outside.state == State.SUCCESS
+@pytest.mark.parametrize("has_until", [False, True])
+@pytest.mark.parametrize(
+ ("region_index", "stop", "accepted"),
+ [
+ # One pass below the cap there is still an iteration left, so continuing is the only
+ # valid answer for a fixed-count loop. A zero-based cap check treats this as the limit.
+ pytest.param(2, False, True, id="continue below the cap"),
+ pytest.param(2, True, False, id="stop below the cap"),
+ # At the cap the loop cannot produce another pass.
+ pytest.param(3, False, False, id="continue at the cap"),
+ pytest.param(3, True, True, id="stop at the cap"),
+ ],
+)
+def test_validate_loop_decision_counts_the_cap_from_one(region_index, stop, accepted, has_until):
+ group = mock.create_autospec(SerializedLoopTaskGroup, instance=True)
+ group.max_iterations, group.has_until = 3, has_until
+ # A conditional loop may stop early, so only the fixed-count loop rejects an early stop.
+ expected = accepted or (has_until and stop)
+
+ if expected:
+ validate_loop_decision(stop, region_index, group)
+ else:
+ with pytest.raises(InvalidLoopDecision):
+ validate_loop_decision(stop, region_index, group)
+
+
+def test_validate_loop_decision_requires_a_decision():
+ group = mock.create_autospec(SerializedLoopTaskGroup, instance=True)
+ group.max_iterations, group.has_until = 3, True
+
+ with pytest.raises(InvalidLoopDecision, match="requires a decision"):
+ validate_loop_decision(None, 1, group)
+
+
@pytest.fixture
def clear_loop_example_runs():
yield
@@ -385,7 +422,7 @@ def test_documented_task_loop_examples_complete(session, dag_id, expected_iterat
assert dr.state == DagRunState.SUCCESS
gates = [ti for ti in dr.get_task_instances(session=session) if ti.operator == "LoopGateOperator"]
assert sorted((ti.region_index, ti.state) for ti in gates) == [
- (index, State.SUCCESS) for index in range(expected_iterations)
+ (index, State.SUCCESS) for index in range(1, expected_iterations + 1)
]
@@ -6353,9 +6390,9 @@ def test_integrity_does_not_create_members_at_draining_gate_coordinates(complete
for ti in dr.get_task_instances(session=session)
if ti.task_id == loop.gate_task_id
}
- gates[4].state = State.RUNNING
+ gates[5].state = State.RUNNING
session.flush()
- clear_loop_task_instances([gates[2]], downstream=False, session=session)
+ clear_loop_task_instances([gates[3]], downstream=False, session=session)
@task_group
def body():
@@ -6373,7 +6410,7 @@ def test_integrity_does_not_create_members_at_draining_gate_coordinates(complete
dr.verify_integrity(session=session, dag_version_id=version)
- assert gates[4].state == State.RESTARTING
+ assert gates[5].state == State.RESTARTING
assert set(
session.scalars(
select(TaskInstance.region_index).where(
@@ -6382,7 +6419,7 @@ def test_integrity_does_not_create_members_at_draining_gate_coordinates(complete
TaskInstance.task_id == "body.new",
)
)
- ) == {0, 1, 2}
+ ) == {1, 2, 3}
def test_run_stays_running_after_clearing_a_gate_while_the_last_pass_is_running(completed_loop, session):
@@ -6392,15 +6429,15 @@ def test_run_stays_running_after_clearing_a_gate_while_the_last_pass_is_running(
for ti in dr.get_task_instances(session=session)
if ti.task_id == loop.gate_task_id
}
- assert sorted(gates) == [0, 1, 2, 3, 4], "Checking pre-conditions: every pass has a gate"
+ assert sorted(gates) == [1, 2, 3, 4, 5], "Checking pre-conditions: every pass has a gate"
# Pretend the last pass's gate is still running, so that clearing an earlier gate supersedes it and
# leaves it RESTARTING until its worker acknowledges the termination.
- gates[4].state = State.RUNNING
+ gates[5].state = State.RUNNING
session.flush()
- clear_loop_task_instances([gates[2]], downstream=False, session=session)
+ clear_loop_task_instances([gates[3]], downstream=False, session=session)
dr.dag = DBDagBag().get_dag_for_run(dr, session=session)
- assert gates[4].state == State.RESTARTING
+ assert gates[5].state == State.RESTARTING
dr.update_state(session=session)
@@ -6411,26 +6448,26 @@ def test_run_stays_running_when_clearing_an_early_gate_restarts_a_running_body_t
dr, loop, _ = completed_loop
members = {(ti.task_id, ti.region_index): ti for ti in dr.get_task_instances(session=session)}
# `completed_loop` creates and completes every pass up front, but a real loop creates pass N + 1 only
- # when gate N completes. Remove pass 4 so that pass 3 is the last pass and is still in progress.
- for ti in [ti for (_, index), ti in members.items() if index == 4]:
+ # when gate N completes. Remove pass 5 so that pass 4 is the last pass and is still in progress.
+ for ti in [ti for (_, index), ti in members.items() if index == 5]:
session.delete(ti)
- # Pass 3 is mid-flight: its `process` task is running and `consume` and the gate have not started.
+ # Pass 4 is mid-flight: its `process` task is running and `consume` and the gate have not started.
for task_id in ("body.consume", loop.gate_task_id):
- members[task_id, 3].state = None
- members["body.process", 3].state = State.RUNNING
+ members[task_id, 4].state = None
+ members["body.process", 4].state = State.RUNNING
session.flush()
- assert {ti.region_index for ti in dr.get_task_instances(session=session)} == {0, 1, 2, 3}, (
- "Checking pre-conditions: pass 3 is the last pass"
+ assert {ti.region_index for ti in dr.get_task_instances(session=session)} == {1, 2, 3, 4}, (
+ "Checking pre-conditions: pass 4 is the last pass"
)
- # Clearing gate 1 supersedes passes 2 and 3: `prepare[3]` is archived and the running `process[3]`
+ # Clearing gate 2 supersedes passes 3 and 4: `prepare[4]` is archived and the running `process[4]`
# goes RESTARTING, so it is left with no live upstream while it waits for its worker to acknowledge.
clear_loop_task_instances(
- [members[loop.gate_task_id, 1]], downstream=False, dag_run_state=False, session=session
+ [members[loop.gate_task_id, 2]], downstream=False, dag_run_state=False, session=session
)
dr.dag = DBDagBag().get_dag_for_run(dr, session=session)
- assert members["body.process", 3].state == State.RESTARTING
- assert members["body.prepare", 3].working_set is None
+ assert members["body.process", 4].state == State.RESTARTING
+ assert members["body.prepare", 4].working_set is None
dr.update_state(session=session)
@@ -6444,30 +6481,30 @@ def test_rerun_gate_continues_while_superseded_later_passes_terminate(completed_
for ti in dr.get_task_instances(session=session)
if ti.task_id == loop.gate_task_id
}
- gates[4].state = State.RUNNING
+ gates[5].state = State.RUNNING
session.flush()
- clear_loop_task_instances([gates[2]], downstream=False, session=session)
- rerun = session.get(TaskInstance, gates[2].id)
+ clear_loop_task_instances([gates[3]], downstream=False, session=session)
+ rerun = session.get(TaskInstance, gates[3].id)
rerun.state = State.SUCCESS
session.flush()
group, _ = TaskCoordinateResolver(DBDagBag(), session).loop_context(rerun)
dr.complete_loop_gate(rerun, group, State.SUCCESS, stop=False, reason=None, session=session)
- assert gates[4].state == State.RESTARTING
+ assert gates[5].state == State.RESTARTING
assert session.scalars(
select(TaskInstance.state).where(
TaskInstance.dag_id == dr.dag_id,
TaskInstance.run_id == dr.run_id,
TaskInstance.task_id == "body.prepare",
- TaskInstance.region_index == 3,
+ TaskInstance.region_index == 4,
)
).all() == [None]
@pytest.mark.parametrize(
("index", "state", "stored"),
- [(4, State.SUCCESS, ["why"]), (4, State.FAILED, ["why"]), (2, State.SUCCESS, []), (2, State.FAILED, [])],
+ [(5, State.SUCCESS, ["why"]), (5, State.FAILED, ["why"]), (3, State.SUCCESS, []), (3, State.FAILED, [])],
ids=["last-success", "last-failed", "earlier-success", "earlier-failed"],
)
def test_complete_loop_gate_records_the_reason_only_without_a_later_pass(
@@ -6483,7 +6520,7 @@ def test_complete_loop_gate_records_the_reason_only_without_a_later_pass(
group, _ = resolver.loop_context(gate)
before = len(dr.get_task_instances(session=session))
- dr.complete_loop_gate(gate, group, state, stop=index == 4, reason="why", session=session)
+ dr.complete_loop_gate(gate, group, state, stop=index == 5, reason="why", session=session)
assert len(dr.get_task_instances(session=session)) == before
assert (
diff --git a/airflow-core/tests/unit/models/test_dynamic_region.py b/airflow-core/tests/unit/models/test_dynamic_region.py
index 54c17bf9b2e..573046e4e9b 100644
--- a/airflow-core/tests/unit/models/test_dynamic_region.py
+++ b/airflow-core/tests/unit/models/test_dynamic_region.py
@@ -910,31 +910,31 @@ def test_gate_rerun_after_chained_forks_keeps_later_passes_in_later_regions(comp
if ti.task_id == loop.gate_task_id and ti.region_index == index
)
+ clear_loop_task_instances([gate_at(3)], downstream=False, session=session)
+ complete_pass(3, "continue")
+ complete_pass(4, "continue")
clear_loop_task_instances([gate_at(2)], downstream=False, session=session)
complete_pass(2, "continue")
complete_pass(3, "continue")
- clear_loop_task_instances([gate_at(1)], downstream=False, session=session)
- complete_pass(1, "continue")
- complete_pass(2, "continue")
before = {
(ti.task_id, ti.region_index): (ti.id, ti.region_id) for ti in dr.get_task_instances(session=session)
}
- clear_loop_task_instances([gate_at(1)], downstream=False, later_loop_iterations=False, session=session)
+ clear_loop_task_instances([gate_at(2)], downstream=False, later_loop_iterations=False, session=session)
- complete_pass(1, "continue")
+ complete_pass(2, "continue")
current = {
(ti.task_id, ti.region_index): (ti.id, ti.region_id) for ti in dr.get_task_instances(session=session)
}
- assert {key: value for key, value in current.items() if key[1] >= 2} == {
- key: value for key, value in before.items() if key[1] >= 2
+ assert {key: value for key, value in current.items() if key[1] >= 3} == {
+ key: value for key, value in before.items() if key[1] >= 3
}
def test_repeated_rewind_appends_empty_forks_and_generates_in_latest_region(completed_loop, session):
dr, loop, complete_pass = completed_loop
original_region = next(ti.region_id for ti in dr.get_task_instances(session=session))
- for index in (2, 1):
+ for index in (3, 2):
gate = next(
ti
for ti in dr.get_task_instances(session=session)
@@ -945,21 +945,21 @@ def test_repeated_rewind_appends_empty_forks_and_generates_in_latest_region(comp
regions = list(session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)))
first = next(region for region in regions if region.forked_from_region_id == original_region)
second = next(region for region in regions if region.forked_from_region_id == first.id)
- assert (first.resumes_from_index, second.resumes_from_index) == (3, 2)
+ assert (first.resumes_from_index, second.resumes_from_index) == (4, 3)
assert not session.scalar(
select(TaskInstance.id).where(TaskInstance.region_id.in_([first.id, second.id]))
)
gate = next(
ti
for ti in dr.get_task_instances(session=session)
- if ti.task_id == loop.gate_task_id and ti.region_index == 1
+ if ti.task_id == loop.gate_task_id and ti.region_index == 2
)
clear_loop_task_instances([gate], downstream=False, session=session)
- complete_pass(1, "continue")
+ complete_pass(2, "continue")
third = session.scalar(select(DynamicRegion).where(DynamicRegion.forked_from_region_id == second.id))
- generated = [ti for ti in dr.get_task_instances(session=session) if ti.region_index == 2]
+ generated = [ti for ti in dr.get_task_instances(session=session) if ti.region_index == 3]
assert len(generated) == 4
assert {ti.region_id for ti in generated} == {third.id}
@@ -982,7 +982,7 @@ def test_loop_sequence_of_unmapped_task_is_ordered_by_iteration(loop_run, sessio
sequence = read_loop_sequence(session, dr, loop, "body.prepare", is_mapped=False)
- assert [ti.region_index for ti in sequence] == [0, 1, 2, 3, 4]
+ assert [ti.region_index for ti in sequence] == [1, 2, 3, 4, 5]
assert {ti.task_id for ti in sequence} == {"body.prepare"}
@@ -992,7 +992,7 @@ def test_loop_sequence_of_mapped_task_is_iteration_major_then_map_index(loop_run
sequence = read_loop_sequence(session, dr, loop, "body.process", is_mapped=True)
assert [(iteration(ti), ti.region_index) for ti in sequence] == [
- (index, slot) for index in range(5) for slot in range(2)
+ (index, slot) for index in range(1, 6) for slot in range(2)
]
@@ -1016,7 +1016,7 @@ def test_loop_sequence_skips_archived_and_other_regions(loop_run, session):
sequence = read_loop_sequence(session, dr, loop, "body.prepare", is_mapped=False)
- assert [ti.region_index for ti in sequence] == [0, 1, 3, 4]
+ assert [ti.region_index for ti in sequence] == [1, 3, 4, 5]
def test_loop_sequence_without_a_loop_region_is_empty(dag_maker, session):
@@ -1055,13 +1055,13 @@ def test_loop_sequence_rejects_a_loop_with_several_executions(loop_run, session)
def test_loop_sequence_takes_each_pass_from_the_region_that_owns_it(loop_run, session):
dr, dag, loop, root, tis, iteration = loop_run
first = add_region(
- session, dr, node_id=loop.group_id, forked_from_region_id=root.id, resumes_from_index=3
+ session, dr, node_id=loop.group_id, forked_from_region_id=root.id, resumes_from_index=4
)
second = add_region(
- session, dr, node_id=loop.group_id, forked_from_region_id=first.id, resumes_from_index=2
+ session, dr, node_id=loop.group_id, forked_from_region_id=first.id, resumes_from_index=3
)
prepare = dag.get_task("body.prepare")
- for region, index in [(first, 3), (first, 4), (second, 2)]:
+ for region, index in [(first, 4), (first, 5), (second, 3)]:
session.add(
TaskInstance(
prepare, dr.created_dag_version_id, run_id=dr.run_id, region_id=region.id, region_index=index
@@ -1072,21 +1072,21 @@ def test_loop_sequence_takes_each_pass_from_the_region_that_owns_it(loop_run, se
sequence = read_loop_sequence(session, dr, loop, "body.prepare", is_mapped=False)
assert [(ti.region_id, ti.region_index) for ti in sequence] == [
- (root.id, 0),
(root.id, 1),
- (second.id, 2),
+ (root.id, 2),
+ (second.id, 3),
]
def test_loop_sequence_of_mapped_task_skips_passes_a_fork_replaced(loop_run, session):
dr, dag, loop, root, tis, iteration = loop_run
- fork = add_region(session, dr, node_id=loop.group_id, forked_from_region_id=root.id, resumes_from_index=3)
+ fork = add_region(session, dr, node_id=loop.group_id, forked_from_region_id=root.id, resumes_from_index=4)
replacement = add_region(
session,
dr,
node_id="body.process",
parent_region_id=fork.id,
- parent_region_index=3,
+ parent_region_index=4,
)
process = dag.get_task("body.process")
for slot in range(2):
@@ -1105,12 +1105,12 @@ def test_loop_sequence_of_mapped_task_skips_passes_a_fork_replaced(loop_run, ses
passes = {region.id: region.parent_region_index for region in session.scalars(select(DynamicRegion))}
assert [(ti.region_id == replacement.id, passes[ti.region_id], ti.region_index) for ti in sequence] == [
- (False, 0, 0),
- (False, 0, 1),
(False, 1, 0),
(False, 1, 1),
(False, 2, 0),
(False, 2, 1),
- (True, 3, 0),
- (True, 3, 1),
+ (False, 3, 0),
+ (False, 3, 1),
+ (True, 4, 0),
+ (True, 4, 1),
]
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py
index d2bec4dc393..0d92f4e3210 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -5225,16 +5225,17 @@ def display_coordinates_tis(dag_maker, session):
run_id=dr.run_id,
node_id="body.inner_mapped",
parent_region_id=loop_region.id,
- parent_region_index=1,
+ parent_region_index=2,
session=session,
)
- for task_id, region in (("body.step", loop_region), ("body.inner_mapped", inner_region)):
- tis[task_id, 1] = TaskInstance(
+ # Keyed by iteration: `step` carries it as its own index, `inner_mapped` carries a map index instead.
+ for task_id, region, index in (("body.step", loop_region, 2), ("body.inner_mapped", inner_region, 1)):
+ tis[task_id, 2] = TaskInstance(
dag.get_task(task_id),
run_id=dr.run_id,
dag_version_id=dr.created_dag_version_id,
region_id=region.id,
- region_index=1,
+ region_index=index,
)
return tis
@@ -5245,7 +5246,7 @@ def display_coordinates_tis(dag_maker, session):
pytest.param(("mapped", 1), False, {"map_index": 1}, id="mapped task outside a loop"),
pytest.param(("plain", -1), False, {}, id="unmapped task"),
pytest.param(("plain", -1), True, {"map_index": -1}, id="unmapped task, index requested"),
- pytest.param(("body.step", 1), False, {"iteration": 1}, id="unmapped loop-body task"),
+ pytest.param(("body.step", 2), False, {"iteration": 2}, id="unmapped loop-body task"),
],
)
def test_display_coordinates(display_coordinates_tis, key, include_unmapped, expected):
@@ -5254,7 +5255,7 @@ def test_display_coordinates(display_coordinates_tis, key, include_unmapped, exp
def test_display_coordinates_of_a_mapped_loop_body_task_keep_the_region(display_coordinates_tis):
- ti = display_coordinates_tis["body.inner_mapped", 1]
+ ti = display_coordinates_tis["body.inner_mapped", 2]
assert ti.get_display_coordinates() == {"region_id": str(ti.region_id), "region_index": 1}
@@ -5263,7 +5264,7 @@ def test_display_coordinates_of_a_mapped_loop_body_task_keep_the_region(display_
[
pytest.param(("mapped", 1), True, True, id="mapped task outside a loop"),
pytest.param(("mapped", 1), False, False, id="detached mapped task outside a loop"),
- pytest.param(("body.inner_mapped", 1), True, False, id="mapped loop-body task"),
+ pytest.param(("body.inner_mapped", 2), True, False, id="mapped loop-body task"),
],
)
def test_display_coordinates_without_the_task_definition(
@@ -5457,7 +5458,7 @@ def test_archive_decides_from_the_state_committed_after_the_lock(completed_loop,
stale = next(
ti
for ti in dr.get_task_instances(session=session)
- if ti.task_id == "body.consume" and ti.region_index == 2
+ if ti.task_id == "body.consume" and ti.region_index == 3
)
stale.state, stale.start_date, stale.end_date = State.QUEUED, None, None
session.flush()
@@ -5479,11 +5480,11 @@ def test_archive_decides_from_the_state_committed_after_the_lock(completed_loop,
def test_mapped_archival_ack_uses_fork_facts_when_definition_is_missing(loop_run, session):
dr, dag, loop, root, tis, iteration = loop_run
archiving = next(
- ti for ti in tis if ti.task_id == "body.process" and iteration(ti) == 3 and ti.region_index == 0
+ ti for ti in tis if ti.task_id == "body.process" and iteration(ti) == 4 and ti.region_index == 0
)
archiving.state = State.RUNNING
session.flush()
- gate = next(ti for ti in tis if ti.task_id == loop.gate_task_id and ti.region_index == 2)
+ gate = next(ti for ti in tis if ti.task_id == loop.gate_task_id and ti.region_index == 3)
clear_loop_task_instances([gate], downstream=False, session=session)
session.execute(
delete(SerializedDagModel).where(SerializedDagModel.dag_version_id == archiving.dag_version_id)
diff --git a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py
index 27fb20d1b78..b49a74c1f66 100644
--- a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py
+++ b/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py
@@ -75,7 +75,7 @@ def test_outside_dependency_observes_only_latest_live_loop_gate(dag_maker, sessi
run_id=dr.run_id,
dag_version_id=previous.dag_version_id,
region_id=previous.region_id,
- region_index=1,
+ region_index=2,
state=latest_state,
)
session.add(latest)
diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
index 3eb99bce115..bd08fa8cf2a 100644
--- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py
+++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
@@ -1098,7 +1098,9 @@ class LoopIterationResponse(BaseModel):
"""
loop_id: Annotated[str, Field(title="Loop Id")]
- iteration: Annotated[int, Field(title="Iteration")]
+ iteration: Annotated[
+ int, Field(description="Position in the enclosing loop, counting from one.", ge=1, title="Iteration")
+ ]
class MaterializeAssetBody(BaseModel):
diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json
index 41a6ab4c6d3..56a6967f17e 100644
--- a/go-sdk/schema/supervisor-schema.json
+++ b/go-sdk/schema/supervisor-schema.json
@@ -3250,7 +3250,7 @@
"type": "string"
},
"index": {
- "minimum": 0,
+ "minimum": 1,
"title": "Index",
"type": "integer"
},
diff --git a/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py b/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py
index 40d91d9b402..d78d5c65597 100644
--- a/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py
+++ b/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py
@@ -1056,7 +1056,7 @@ class TestRegionalReview:
run_id=dr.run_id,
dag_version_id=first.dag_version_id,
region_id=first.region_id,
- region_index=1,
+ region_index=2,
state="deferred",
)
session.add(second)
@@ -1111,19 +1111,19 @@ class TestRegionalReview:
assert response.status_code == 200, response.text
data = response.json()
assert data["status"] == expected_status
- assert data["current_output"] == "pass 1"
- assert data["conversation"][0]["content"] == "pass 1"
+ assert data["current_output"] == "pass 2"
+ assert data["conversation"][0]["content"] == "pass 2"
assert data["task_completed"] is False
session.expire_all()
earlier = session.scalar(
select(XComModel.value).where(
XComModel.key == XCOM_AGENT_SESSION,
XComModel.region_id == ti.region_id,
- XComModel.region_index == 0,
+ XComModel.region_index == 1,
)
)
assert earlier["status"] == "pending_review"
- assert earlier["current_output"] == "pass 0"
+ assert earlier["current_output"] == "pass 1"
@mock.patch("airflow.providers.common.ai.plugins.hitl_review.AIRFLOW_V_3_4_PLUS", False)
def test_region_selectors_are_rejected_on_hosts_without_regions(self, regional_review, test_client):
@@ -1181,7 +1181,7 @@ class TestForkedRegionReview:
run_id=dr.run_id,
dag_version_id=first.dag_version_id,
region_id=region.id,
- region_index=1,
+ region_index=2,
state="failed",
)
forked = TaskInstance(
@@ -1189,7 +1189,7 @@ class TestForkedRegionReview:
run_id=dr.run_id,
dag_version_id=first.dag_version_id,
region_id=fork.id,
- region_index=1,
+ region_index=2,
state="running",
)
session.add_all([beyond_cut, forked])
@@ -1350,7 +1350,7 @@ class TestMappedReviewInLoop:
run_id=first_region.run_id,
node_id=first_region.node_id,
parent_region_id=first_region.parent_region_id,
- parent_region_index=1,
+ parent_region_index=2,
session=session,
)
session.flush()
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index e76dbaa90c4..0e677c93e7a 100644
--- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -338,7 +338,7 @@ class LoopContext(BaseModel):
"""
node_id: Annotated[str, Field(title="Node Id")]
- index: Annotated[int, Field(ge=0, title="Index")]
+ index: Annotated[int, Field(ge=1, title="Index")]
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")]
diff --git a/task-sdk/src/airflow/sdk/execution_time/loop.py b/task-sdk/src/airflow/sdk/execution_time/loop.py
index 755d5d10362..f8aac861191 100644
--- a/task-sdk/src/airflow/sdk/execution_time/loop.py
+++ b/task-sdk/src/airflow/sdk/execution_time/loop.py
@@ -56,7 +56,8 @@ class LoopContextAccessor:
@property
def previous(self) -> Any:
- return None if self.index == 0 else self._result(previous_iteration=True)
+ # Iterations count from one, so the first pass has no previous result.
+ return None if self.index == 1 else self._result(previous_iteration=True)
@property
def result(self) -> Any:
@@ -82,7 +83,7 @@ def execute_loop_gate(gate: LoopGateOperator, context: Context) -> None:
ti = cast("RuntimeTaskInstance", context["ti"])
reason: str | None = None
if gate.until is None:
- stop = loop.index + 1 >= loop.max_iterations
+ stop = loop.index >= loop.max_iterations
else:
match result := gate.until(**determine_kwargs(gate.until, (), context)):
case bool():
@@ -117,7 +118,7 @@ def execute_loop_gate(gate: LoopGateOperator, context: Context) -> None:
raise AirflowFailException("LoopDecision reason cannot contain NUL bytes")
ti._loop_decision = LoopDecisionPayload(stop=stop, reason=reason)
suffix = "" if reason is None else f": {reason}"
- if not stop and loop.index + 1 >= loop.max_iterations:
+ if not stop and loop.index >= loop.max_iterations:
raise LoopMaxIterationsExceeded(
f"Loop condition remained false at max_iterations={loop.max_iterations}{suffix}"
)
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 41a6ab4c6d3..56a6967f17e 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
@@ -3250,7 +3250,7 @@
"type": "string"
},
"index": {
- "minimum": 0,
+ "minimum": 1,
"title": "Index",
"type": "integer"
},
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 b226a8bf77b..77e19c10df5 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_loop.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_loop.py
@@ -44,7 +44,7 @@ from airflow.sdk.execution_time.loop import LoopMaxIterationsExceeded
@pytest.fixture
def loop_ti(create_runtime_ti):
- def make(*, index=0, max_iterations=3, until=None, map_index=-1):
+ def make(*, index=1, max_iterations=3, until=None, map_index=-1):
@task_group
def body():
BaseOperator(task_id="terminal")
@@ -68,7 +68,7 @@ def test_first_iteration_previous_does_not_read_xcom(loop_ti, mock_supervisor_co
ti = loop_ti(map_index=7)
loop = ti.get_template_context()["loop"]
- assert loop.index == 0
+ assert loop.index == 1
assert loop.max_iterations == 3
assert loop.previous is None
assert ti.map_index == 7
@@ -107,7 +107,7 @@ def test_body_callable_receives_loop_context_after_unmapping(create_runtime_ti,
@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=1)
+ ti = loop_ti(index=2)
mock_supervisor_comms.send.return_value = XComResult(key="return_value", value=value)
loop = ti.get_template_context()["loop"]
@@ -123,31 +123,34 @@ def test_previous_and_current_result_preserve_falsey_values(loop_ti, mock_superv
@pytest.mark.parametrize(
("index", "until", "expected"),
[
- (0, None, LoopDecisionPayload(stop=False)),
- (2, None, LoopDecisionPayload(stop=True)),
- (0, lambda: False, LoopDecisionPayload(stop=False)),
- (0, lambda: True, LoopDecisionPayload(stop=True)),
- (2, lambda: True, LoopDecisionPayload(stop=True)),
- (0, lambda: LoopDecision(stop=True), LoopDecisionPayload(stop=True)),
+ (1, None, LoopDecisionPayload(stop=False)),
+ # max_iterations - 1: a zero-based cap check stops here instead, so this pins the base.
+ (2, None, LoopDecisionPayload(stop=False)),
+ (3, None, LoopDecisionPayload(stop=True)),
+ (1, lambda: False, LoopDecisionPayload(stop=False)),
+ (1, lambda: True, LoopDecisionPayload(stop=True)),
+ (3, lambda: True, LoopDecisionPayload(stop=True)),
+ (1, lambda: LoopDecision(stop=True), LoopDecisionPayload(stop=True)),
(
- 0,
+ 1,
lambda: LoopDecision(stop=False, reason="score 0.7, needs > 0.9"),
LoopDecisionPayload(stop=False, reason="score 0.7, needs > 0.9"),
),
(
- 0,
+ 1,
lambda: LoopDecision(stop=True, reason="é" * 256),
LoopDecisionPayload(stop=True, reason="é" * 256),
),
- (0, lambda: LoopDecision(stop=True, reason=""), LoopDecisionPayload(stop=True)),
+ (1, lambda: LoopDecision(stop=True, reason=""), LoopDecisionPayload(stop=True)),
(
- 0,
+ 1,
lambda: LoopDecision(stop=True, reason="token hunter2 expired"),
LoopDecisionPayload(stop=True, reason="token *** expired"),
),
],
ids=[
"fixed_continue",
+ "fixed_continue_below_cap",
"fixed_stop_at_cap",
"bool_continue",
"bool_stop",
@@ -221,7 +224,7 @@ def _raise_value_error():
],
)
def test_gate_fails_without_a_decision(loop_ti, mock_supervisor_comms, until, error, message):
- ti = loop_ti(index=0, until=until)
+ ti = loop_ti(index=1, until=until)
state, sent, raised = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
@@ -232,6 +235,17 @@ def test_gate_fails_without_a_decision(loop_ti, mock_supervisor_comms, until, er
assert sent.loop_decision is None
+def test_unmet_condition_before_the_cap_continues_instead_of_failing(loop_ti, mock_supervisor_comms):
+ """The pass before the cap has an iteration left, so a false condition just continues."""
+ ti = loop_ti(index=2, until=lambda: False)
+
+ state, message, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
+
+ assert (state, error) == (TaskInstanceState.SUCCESS, None)
+ assert isinstance(message, SucceedTask)
+ assert message.loop_decision == LoopDecisionPayload(stop=False)
+
+
@pytest.mark.parametrize(
("result", "message", "reason"),
[
@@ -247,7 +261,7 @@ def test_gate_fails_without_a_decision(loop_ti, mock_supervisor_comms, until, er
def test_gate_fails_without_retrying_when_the_condition_is_false_at_the_cap(
loop_ti, mock_supervisor_comms, result, message, reason
):
- ti = loop_ti(index=2, until=lambda: result)
+ ti = loop_ti(index=3, until=lambda: result)
state, sent, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
@@ -263,7 +277,7 @@ def test_gate_fails_without_retrying_when_the_condition_is_false_at_the_cap(
def test_mapped_terminal_keeps_iteration_across_lazy_reads(
loop_ti, mock_supervisor_comms, previous_iteration, values
):
- ti = loop_ti(index=1)
+ ti = loop_ti(index=2)
ti._ti_context_from_server.loop.terminal_is_mapped = True
def respond(message):
diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json
index 41a6ab4c6d3..56a6967f17e 100644
--- a/ts-sdk/schema/supervisor-schema.json
+++ b/ts-sdk/schema/supervisor-schema.json
@@ -3250,7 +3250,7 @@
"type": "string"
},
"index": {
- "minimum": 0,
+ "minimum": 1,
"title": "Index",
"type": "integer"
},