diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst index b0ea5b9757583..7a6e04e0c5988 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 09ecefbf8c734..a082a718511ca 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 e7bc1eccb2c75..847d6e8e04ed4 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 271ed1d159a1a..640faa0942aa0 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 1fe5224a0d3d8..a6a33e7553351 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 12cf433d84e90..2eaa3296aad16 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 cda3ecd9a14ed..09ce1193e9018 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 7bf2a8762982b..26d0aea88781e 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 @@ 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 8b178c4f6a2eb..5ee8984a0a783 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 process(value, *, loop, ti): 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 fc116ab20bf40..f28cd7690f44a 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 @@ def _create_initial_tasks( 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 b2100d9573436..968b92248292f 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 f569fc4f36f80..4690e8adf4667 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 76133b1f85e4d..bb6c50f3507f9 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 7d646a03e8e37..661826be84c48 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 d2e58f17ca167..d620a72ace75f 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 = ; /** - * 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 7d5b4daf2ec59..3ea4123a68d3a 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 c46934774949e..4da7331d327b0 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 0fec3f0fe1566..08ed43a5308ff 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 @@ -2356,7 +2356,7 @@ def body(): 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) @@ -2368,7 +2368,7 @@ def body(): 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 f6b0c1bcebbaf..4a9865b0df514 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 @@ def body(): 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 @@ def body(): ) 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 @@ def body(): 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 @@ def body(): 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 39f73e0c77c0f..dbf7a1b67ac25 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 finish_run(): 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 finish_run(): 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 0bb45fb62a83f..91fb81b48ede4 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 body(): 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 251adadc7aeaf..67d9a6751a332 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 @@ def body(): 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 @@ def body(): 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 @@ def body(): 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 @@ def body(): ) 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 @@ def body(): 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 @@ def body(): 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 @@ def body(): 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 @@ def body(): 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 @@ def body(): 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 @@ def body(): ) ) 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 @@ def build_payload(): 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 @@ def build_payload(): 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 @@ def body(): 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 @@ def test_clear_discards_task_state_of_the_selected_loop_iteration_only( ): 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 @@ def test_clear_discards_task_state_of_the_selected_loop_iteration_only( 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 @@ def test_clear_discards_task_state_of_the_selected_loop_iteration_only( 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 @@ def test_exact_loop_mark_does_not_change_other_iterations( 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 @@ def body(): 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 @@ def body(): 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 @@ def body(): (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 @@ def test_loop_read_without_index_requires_a_region_and_index( [ 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 @@ def test_loop_mark_rejects_map_index_conflicting_with_region_index( 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 @@ def test_exact_loop_note_changes_only_selected_execution( 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 @@ def test_exact_loop_note_changes_only_selected_execution( 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 @@ def test_manual_gate_success_after_rewind_stops_and_keeps_the_reason( 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 @@ def test_manual_terminal_settles_superseded_running_execution( 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 @@ def test_manual_terminal_settles_superseded_running_execution( @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 @@ def test_bulk_exact_passes_have_distinct_execution_results( 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 @@ def test_bulk_exact_passes_have_distinct_execution_results( "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 @@ def test_bulk_exact_passes_have_distinct_execution_results( 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 @@ def test_delete_of_a_loop_pass_removes_its_archived_tries( 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 @@ def test_delete_of_a_loop_pass_removes_its_archived_tries( 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 @@ def test_delete_of_a_loop_pass_removes_its_archived_tries( "task_id": "body.first", "map_index": -1, "region_id": str(root.id), - "region_index": 2, + "region_index": 3, } ], } @@ -6249,7 +6249,7 @@ def test_delete_of_a_loop_pass_removes_its_archived_tries( 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 @@ def test_delete_of_a_loop_pass_removes_its_archived_tries( 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 @@ def test_bulk_regional_entities_do_not_lock_their_run_again(self, test_client, l "region_index": index, "new_state": "success", } - for index in (0, 2) + for index in (1, 3) ] response = test_client.patch( @@ -6289,16 +6289,16 @@ def test_group_coordinates_select_one_loop_pass(self, test_client, session, loop 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 @@ def body(): 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 @@ def test_regional_success_preserves_other_state_scope_and_notifies_listener( 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 @@ def test_regional_success_preserves_other_state_scope_and_notifies_listener( "new_state": "success", "note": "completed manually", "region_id": str(root.id), - "region_index": 1, + "region_index": 2, }, ) @@ -6404,7 +6404,7 @@ def test_regional_success_preserves_other_state_scope_and_notifies_listener( 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 @@ def test_regional_downstream_mark_uses_same_scope_as_preview( "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 @@ def changed(): 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 b03710137facf..78ee974b5f599 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 @@ def body(): 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 87be7d87815ba..65e46cd11dec3 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 @@ def test_fixed_loop_summary_reads_live_iterations(self, test_client, executed_fi 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 @@ def test_fixed_loop_summary_reads_live_iterations(self, test_client, executed_fi 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 @@ def test_summary_reports_one_invocation_across_a_fork_after_gate_clear( 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 @@ def test_summary_reports_one_invocation_across_a_fork_after_gate_clear( 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 @@ def test_summary_does_not_guess_condition_failure( 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 @@ def test_grid_counts_every_pass_but_links_to_the_newest(self, test_client, execu 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 @@ def test_grid_counts_every_pass_but_links_to_the_newest(self, test_client, execu 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 @@ def body(): 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 @@ def body(): 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 @@ def finish(): 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 58b77d6375756..6a81588ffc555 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 @@ def body(): 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 @@ def started_loop_successor(self, client, session, running_loop_gate): 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 @@ def test_loop_rewind_admits_the_gate_without_waiting_for_worker_acknowledgement( 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 @@ def clear_again(): @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 @@ def body(): ) == 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 @@ def fail_after_insert(*args, **kwargs): 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 @@ def complete(): 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 @@ def body(): 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 @@ def body(): 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 @@ def test_rerun_gate_preserves_existing_later_pass_and_reason( 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 @@ def test_rerun_gate_preserves_existing_later_pass_and_reason( 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 @@ def test_rerun_gate_preserves_existing_later_pass_and_reason( 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 @@ def body(): 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 7768fe753368b..0685fc67aec6f 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 body(): 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 body(): 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 body(): 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 body(): 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 body(): ) ) 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 body(): ) 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 record(api_session, *_): 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 record(api_session, *_): 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 @@ def test_trigger_ti_id_header_rejected(deferred_loop_pass, api_triggerer_id, hea "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 body(): 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 body(): 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 record(_, instance): 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 bc18469dcac4c..a9397e8460a88 100644 --- a/airflow-core/tests/unit/models/conftest.py +++ b/airflow-core/tests/unit/models/conftest.py @@ -46,7 +46,8 @@ def body(): 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 complete_pass(index, decision): ) 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 95d52e56301c6..5ed43c424ec0c 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 body(): 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 body(): 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 body(): 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 @@ def test_archives_every_later_pass_and_regenerates_from_the_first( 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 @@ def test_rerun_can_shorten_the_loop(self, completed_loop, session, clear_run): 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 5c191505666ca..ee63b2db37985 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.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.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 body(): 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 body(): 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 body(add_member=False): 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 body(): 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 converged(*, loop): 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 body(): 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 body(): 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 54c17bf9b2eda..573046e4e9b6e 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 gate_at(index): 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 d2bec4dc3939b..0d92f4e321030 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -5225,16 +5225,17 @@ def body(): 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 body(): 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 27fb20d1b785b..b49a74c1f6625 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 body(): 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 3eb99bce1153d..bd08fa8cf2abf 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 41a6ab4c6d394..56a6967f17e77 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 40d91d9b40286..d78d5c655976f 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 @@ def body(): 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 @@ def test_exact_scope_preserves_other_loop_pass( 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 @@ def body(): 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 @@ def body(): 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 @@ def body(): 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 e76dbaa90c462..0e677c93e7a94 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 755d5d1036245..f8aac86119154 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 @@ def max_iterations(self) -> int: @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 41a6ab4c6d394..56a6967f17e77 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 b226a8bf77b35..77e19c10df5b4 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 @@ @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 body(): @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 41a6ab4c6d394..56a6967f17e77 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" },