Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 11 additions & 10 deletions airflow-core/docs/authoring-and-scheduling/loops.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Comment thread
bbovenzi marked this conversation as resolved.
Resolved
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.

Expand Down Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,)
Expand Down
10 changes: 5 additions & 5 deletions airflow-core/src/airflow/example_dags/example_task_loops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand All @@ -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):
Expand All @@ -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)

Expand All @@ -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)

Expand Down
5 changes: 3 additions & 2 deletions airflow-core/src/airflow/models/dagrun.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion airflow-core/src/airflow/models/dynamic_region.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
5 changes: 5 additions & 0 deletions airflow-core/src/airflow/models/taskinstance.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1380,7 +1380,7 @@ export const useTaskInstanceServiceGetMappedTaskInstance = <TData = Common.TaskI
* @param data.regionId
* @param data.regionIndex
* @param data.loopId Task group id of a loop; matches every instance it produced.
* @param data.iteration Narrow ``loop_id`` to a single pass of that loop.
* @param data.iteration Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
* @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored.
* @param data.taskId
* @param data.runAfterGte
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1380,7 +1380,7 @@ export const useTaskInstanceServiceGetMappedTaskInstanceSuspense = <TData = Comm
* @param data.regionId
* @param data.regionIndex
* @param data.loopId Task group id of a loop; matches every instance it produced.
* @param data.iteration Narrow ``loop_id`` to a single pass of that loop.
* @param data.iteration Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
* @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored.
* @param data.taskId
* @param data.runAfterGte
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6708,7 +6708,9 @@ export const $LoopIterationResponse = {
},
iteration: {
type: 'integer',
title: 'Iteration'
minimum: 1,
title: 'Iteration',
description: 'Position in the enclosing loop, counting from one.'
}
},
type: 'object',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3136,7 +3136,7 @@ export class TaskInstanceService {
* @param data.regionId
* @param data.regionIndex
* @param data.loopId Task group id of a loop; matches every instance it produced.
* @param data.iteration Narrow ``loop_id`` to a single pass of that loop.
* @param data.iteration Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
* @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored.
* @param data.taskId
* @param data.runAfterGte
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1833,6 +1833,9 @@ export type LastAssetEventResponse = {
*/
export type LoopIterationResponse = {
loop_id: string;
/**
* Position in the enclosing loop, counting from one.
*/
iteration: number;
};

Expand Down Expand Up @@ -4657,7 +4660,7 @@ export type GetTaskInstancesData = {
endDateLte?: string | null;
executor?: Array<(string)>;
/**
* Narrow ``loop_id`` to a single pass of that loop.
* Narrow ``loop_id`` to a single pass of that loop. Iterations count from one.
*/
iteration?: number | null;
limit?: number;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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", () => {
Expand All @@ -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);
});
});

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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"]}
Expand Down
Loading
Loading