From 0b3e636e89abe0a1a2aad086a752fd39e2e7ec2a Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Fri, 16 Oct 2026 16:32:07 +0200 Subject: [PATCH 01/10] Improve task loop information and tests for OpenLineage provider --- .../openlineage_tests/prepare_dags.py | 1 + .../openlineage/facets/AirflowRunFacet.json | 32 + .../providers/openlineage/plugins/listener.py | 87 +- .../providers/openlineage/utils/utils.py | 190 ++- .../example_openlineage_retry_dag.py | 68 + ...ample_openlineage_sensor_reschedule_dag.py | 67 + .../example_openlineage_task_loop_dag.py | 122 ++ .../openlineage_manual_lineage_dag.json | 64 + .../openlineage_retry_dag.json | 76 + .../openlineage_sensor_reschedule_dag.json | 34 + .../openlineage_task_loop_dag.json | 1486 +++++++++++++++++ .../tests/system/openlineage/operator.py | 178 +- .../unit/openlineage/plugins/test_listener.py | 82 + .../test_system_tests_run_id_integrity.py | 250 +++ .../unit/openlineage/utils/test_utils.py | 292 +++- 15 files changed, 2956 insertions(+), 73 deletions(-) create mode 100644 providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py create mode 100644 providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py create mode 100644 providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py create mode 100644 providers/openlineage/tests/system/openlineage/expected_events/openlineage_retry_dag.json create mode 100644 providers/openlineage/tests/system/openlineage/expected_events/openlineage_sensor_reschedule_dag.json create mode 100644 providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json create mode 100644 providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/openlineage_tests/prepare_dags.py b/airflow-e2e-tests/tests/airflow_e2e_tests/openlineage_tests/prepare_dags.py index 7d99e83b09ef4..bff41c6ff6e10 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/openlineage_tests/prepare_dags.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/openlineage_tests/prepare_dags.py @@ -50,6 +50,7 @@ # otherwise raise at parse time and take the whole dag-processor down, not just skip that one DAG. MIN_AIRFLOW_VERSION_FOR_DAG: dict[str, tuple[int, int]] = { "example_openlineage_hitl_dag.py": (3, 1), + "example_openlineage_task_loop_dag.py": (3, 4), } diff --git a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json index 621eb6a761453..8312c2b37d9a0 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json +++ b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json @@ -193,13 +193,45 @@ "duration": { "type": "number" }, + "id": { + "description": "UUID of this attempt (Airflow 3+). Every try, and every task instance a clear creates again, gets a new one.", + "type": "string", + "format": "uuid" + }, "map_index": { + "description": "Position in a mapped task's expansion. Absent for a task that isn't mapped.", + "type": "integer" + }, + "loop_index": { + "description": "Zero-based iteration of the innermost loop the task runs in (Airflow 3.4+). Absent outside a loop.", + "type": "integer" + }, + "loop_max_iterations": { + "description": "Maximum number of iterations of the innermost loop the task runs in (Airflow 3.4+). Absent outside a loop.", + "type": "integer" + }, + "loop_group_id": { + "description": "Group id of the innermost loop the task runs in (Airflow 3.4+), which loop_index counts the iterations of. Absent outside a loop.", + "type": "string" + }, + "region_id": { + "description": "Id of the region (one execution of a loop or of a mapped task's expansion) the task instance belongs to (Airflow 3.4+). Shared by the task instances of that region. Clearing a loop forks its region, so a re-run of an iteration gets a new region_id. Region ids are UUIDv7, so a later trajectory has a greater one, which is an implementation detail. Absent for a task that is neither mapped nor in a loop.", + "type": "string", + "format": "uuid" + }, + "region_index": { + "description": "Position inside the region (Airflow 3.4+): the iteration for a task in a loop that isn't mapped, the map index for a mapped task. Together with dag_id, the Dag run's run_id, task_id, region_id and try_number it identifies one attempt. Absent with region_id.", "type": "integer" }, "pool": { "type": "string" }, "try_number": { + "description": "Try of the task instance, counted per region_id and region_index.", + "type": "integer" + }, + "max_tries": { + "description": "Last try_number that is retried on failure, so a failed try with try_number > max_tries is final. Equals the task's retries until the task instance is cleared, which sets it to the last try_number plus retries.", "type": "integer" }, "queued_dttm": { diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py index 6f04c7881f310..b1fc356f577f3 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py @@ -53,8 +53,8 @@ get_dag_run_dag_and_task_from_ti, get_job_name, get_orm_task_instance_positions, - get_regional_task_instance_run_id, get_task_documentation, + get_task_instance_position, get_task_parent_run_facet, get_user_provided_run_facets, is_dag_run_asset_triggered, @@ -263,15 +263,17 @@ def on_running(): clear_number=clear_number, ) - task_uuid = get_regional_task_instance_run_id( - task_instance - ) or self.adapter.build_task_instance_run_id( - dag_id=task_instance.dag_id, - task_id=task_instance.task_id, - try_number=task_instance.try_number, - logical_date=date, - map_index=task_instance.map_index, - ) + position = get_task_instance_position(task_instance) + if position.in_loop: + task_uuid = str(task_instance.id) + else: + task_uuid = self.adapter.build_task_instance_run_id( + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + try_number=task_instance.try_number, + logical_date=date, + map_index=position.map_index, + ) event_type = RunState.RUNNING.value.lower() operator_name = task.task_type.lower() @@ -341,6 +343,7 @@ def on_running(): task, task_uuid, include_full_task_info=controls.include_full_task_info, + ti_position=position, ), **debug_facet, }, @@ -427,15 +430,17 @@ def on_success(): clear_number=dagrun.clear_number, ) - task_uuid = get_regional_task_instance_run_id( - task_instance - ) or self.adapter.build_task_instance_run_id( - dag_id=task_instance.dag_id, - task_id=task_instance.task_id, - try_number=task_instance.try_number, - logical_date=date, - map_index=task_instance.map_index, - ) + position = get_task_instance_position(task_instance) + if position.in_loop: + task_uuid = str(task_instance.id) + else: + task_uuid = self.adapter.build_task_instance_run_id( + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + try_number=task_instance.try_number, + logical_date=date, + map_index=position.map_index, + ) event_type = RunState.COMPLETE.value.lower() operator_name = task.task_type.lower() @@ -504,6 +509,7 @@ def on_success(): task, task_uuid, include_full_task_info=controls.include_full_task_info, + ti_position=position, ), **get_airflow_debug_facet(), }, @@ -605,15 +611,17 @@ def on_failure(): clear_number=dagrun.clear_number, ) - task_uuid = get_regional_task_instance_run_id( - task_instance - ) or self.adapter.build_task_instance_run_id( - dag_id=task_instance.dag_id, - task_id=task_instance.task_id, - try_number=task_instance.try_number, - logical_date=date, - map_index=task_instance.map_index, - ) + position = get_task_instance_position(task_instance) + if position.in_loop: + task_uuid = str(task_instance.id) + else: + task_uuid = self.adapter.build_task_instance_run_id( + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + try_number=task_instance.try_number, + logical_date=date, + map_index=position.map_index, + ) event_type = RunState.FAIL.value.lower() operator_name = task.task_type.lower() @@ -682,6 +690,7 @@ def on_failure(): task, task_uuid, include_full_task_info=controls.include_full_task_info, + ti_position=position, ), **get_airflow_debug_facet(), }, @@ -759,15 +768,17 @@ def on_skipped(): clear_number=dagrun.clear_number, ) - task_uuid = get_regional_task_instance_run_id( - task_instance - ) or self.adapter.build_task_instance_run_id( - dag_id=task_instance.dag_id, - task_id=task_instance.task_id, - try_number=task_instance.try_number, - logical_date=date, - map_index=task_instance.map_index, - ) + position = get_task_instance_position(task_instance) + if position.in_loop: + task_uuid = str(task_instance.id) + else: + task_uuid = self.adapter.build_task_instance_run_id( + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + try_number=task_instance.try_number, + logical_date=date, + map_index=position.map_index, + ) event_type = RunState.COMPLETE.value.lower() operator_name = task.task_type.lower() @@ -835,6 +846,7 @@ def on_skipped(): task, task_uuid, include_full_task_info=controls.include_full_task_info, + ti_position=position, ), **get_airflow_debug_facet(), }, @@ -959,6 +971,7 @@ def _on_task_instance_manual_state_change( task, task_uuid, include_full_task_info=include_full_task_info, + ti_position=position, ) adapter_kwargs: dict = { diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index ca43441870d47..437de4fb2c9b9 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1123,9 +1123,51 @@ def team_name(cls, dagrun: DagRun) -> str | None: class TaskInstanceInfo(InfoJsonEncodable): - """Defines encoding TaskInstance object to JSON.""" + """ + Encode a task instance as ``run.facets.airflow.taskInstance``, which tells which execution sent an event. + + The fields that identify the execution, each left out when it doesn't apply: + + * ``id`` (Airflow 3+): the UUID of the attempt. Every try, and every task instance a clear creates again, + gets a new one, so it identifies one execution, and it is the ``runId`` of a task inside a loop. Airflow's + API can look it up, but it says nothing by itself. + * ``try_number``: the try, counted per ``region_id`` and ``region_index`` (see below). + * ``max_tries``: the last ``try_number`` that still gets retried on failure, so a failed try with + ``try_number > max_tries`` is final. It equals the task's ``retries`` (in ``task``) only until a clear: + clearing sets it to the last try number plus ``retries``, giving the task its retries again. + * ``map_index``: the position in a mapped task's expansion that users see. Absent for a task that isn't + mapped. + * ``loop_index``, ``loop_max_iterations`` and ``loop_group_id`` (Airflow 3.4+): the zero-based iteration + of the innermost loop around the task, that loop's limit and its group id, which can differ from + ``task.task_group`` when the task sits in a task group inside the loop. The iterations of loops further + out aren't known. ``(loop_group_id, loop_index)`` names an iteration, but not which trajectory it ran in, + nor in a nested loop which iteration of the outer loop. + * ``region_id`` and ``region_index`` (Airflow 3.4+): the coordinates Airflow stores the task instance under. + A region is one execution of a loop, or of a mapped task's expansion, and is shared by all the task + instances it creates. ``region_index`` is the position inside it: the iteration for a task in a loop that + isn't mapped, the map index for a mapped task. Neither is unique: ``(dag_id, dagRun.run_id, task_id, + region_id, region_index, try_number)`` is, as Airflow enforces it (``task_instance_try_key``). Both are + absent for a task in no region, i.e. one that is neither mapped nor in a loop. + + How a cleared loop reads: clearing a loop from iteration ``k`` retries what was cleared in iteration ``k`` in + place, with the same region and the next ``try_number``. The iterations after ``k`` are created again in a + new region forked from the loop's region, and try numbers count per region, so the new execution of + iteration ``k + 1`` can repeat the task, ``loop_index``, ``map_index`` and ``try_number`` of the one it + replaces. Only ``region_id`` (and ``id``) tell the two apart. Each such re-run of a loop is a trajectory. + ``dagRun.clear_number`` counts clears of the whole Dag run only. + + Ordering: region ids are UUIDv7, so of two executions of one task with the same ``loop_index`` and + ``map_index``, the later trajectory has the greater ``region_id``. That is how Airflow generates them, not a + documented guarantee, and it only orders the trajectories seen: an event doesn't say whether it comes from + the second or the third trajectory, nor from which iteration its trajectory was forked. + + What events can't tell: a mapped task's ``region_id`` is the one of its own expansion, which sits inside the + loop's region without naming it. So a mapped task instance in a loop can't be matched to the trajectory of + the loop, or to the task instances of the same pass that aren't mapped (whose ``region_id`` is the loop's), + other than by ``loop_index`` and time. + """ - includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"] + includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "try_number"] casts = { "log_url": lambda ti: getattr(ti, "log_url", None), "note": lambda ti: safe_getattr(ti, "note", None), # From manual state changes only @@ -1135,11 +1177,39 @@ class TaskInstanceInfo(InfoJsonEncodable): "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None, } + def __init__(self, obj, position: TaskInstancePosition | None = None): + # Pass the position when the caller has already placed the task instance, e.g. from its stored region + self._position = position + super().__init__(obj) + def _extend_fields(self) -> None: - map_index = get_task_instance_map_index(self.obj) + # A stored task instance whose task isn't attached can't be placed without its region, and a facet must + # not fail the whole event, so then it describes the task instance without its position. + position = self._position or _get_task_instance_position(self.obj) + if position is not None: + position = _add_runtime_loop_iteration(self.obj, position) + map_index = position.map_index if position else -1 self.map_index = map_index if map_index != -1 else None self.rendered_map_index = getattr(self.obj, "rendered_map_index", None) if map_index != -1 else None self._fields.extend(("map_index", "rendered_map_index")) + if position and position.loop_index is not None: + self.loop_index = position.loop_index + self._fields.append("loop_index") + if position and position.loop_max_iterations is not None: + self.loop_max_iterations = position.loop_max_iterations + self._fields.append("loop_max_iterations") + if position and position.loop_group_id is not None: + self.loop_group_id = position.loop_group_id + self._fields.append("loop_group_id") + if AIRFLOW_V_3_0_PLUS and (ti_id := getattr(self.obj, "id", None)) is not None: + self.id = str(ti_id) + self._fields.append("id") + # A runtime task instance in no region carries the sentinel region id instead of None + region_id = getattr(self.obj, "region_id", None) + if isinstance(region_id, UUID) and region_id.int != 0: + self.region_id = str(region_id) + self.region_index = getattr(self.obj, "region_index", None) + self._fields.extend(("region_id", "region_index")) class AssetInfo(InfoJsonEncodable): @@ -1354,6 +1424,10 @@ class TaskGroupInfo(InfoJsonEncodable): "tooltip", "upstream_group_ids", "upstream_task_ids", + "gate_task_id", # in a loop group + "has_until", # in a loop group + "max_iterations", # in a loop group + "terminal_task_id", # in a loop group ] @@ -1409,13 +1483,14 @@ def get_airflow_run_facet( task: BaseOperator, task_uuid: str, include_full_task_info: bool = False, + ti_position: TaskInstancePosition | None = None, ) -> dict[str, AirflowRunFacet]: runtime_assets = get_runtime_outlet_assets(task_instance) return { "airflow": AirflowRunFacet( dag=DagInfo(dag), dagRun=DagRunInfo(dag_run), - taskInstance=TaskInstanceInfo(task_instance), + taskInstance=TaskInstanceInfo(task_instance, position=ti_position), task=( TaskInfoComplete(task, runtime_assets=runtime_assets) if include_full_task_info @@ -1461,6 +1536,8 @@ def get_task_duration(ti): return { "airflowState": AirflowStateRunFacet( dagRunState=dag_run_state, + # TODO: keyed by task_id only, so the mapped instances and loop iterations of one task collapse + # into a single entry (the last one by region wins). Needs the map index and loop iteration here. tasksState={ti.task_id: ti.state for ti in tis}, tasksDuration={ti.task_id: get_task_duration(ti) for ti in tis}, ) @@ -1489,18 +1566,28 @@ def is_dag_run_asset_triggered( @attrs.frozen class TaskInstancePosition: """ - Where a task instance runs: inside a loop body or not, and the map index users see. + Where a task instance runs: inside a loop body or not, the map index users see and the loop iteration. A plain mapped task keeps the region of its own expansion, so its dag, task, try, logical date and map index stay unique; inside a loop they repeat on every pass, so only the execution UUID identifies - the run. + the run. ``loop_index`` is the zero-based iteration and ``loop_max_iterations`` the limit of the loop, the + counterparts of ``try_number`` and the maximum number of tries, and ``loop_group_id`` the loop's group id. + All three describe the innermost loop and are ``None`` outside a loop. A stored task whose task isn't + attached has no ``loop_max_iterations``, nor ``loop_group_id`` when it is mapped, and a stored mapped task + in a loop has no ``loop_index`` unless ``get_orm_task_instance_positions`` could read its region. A runtime + task instance has no ``loop_index`` nor ``loop_max_iterations`` unless placed with + ``get_task_instance_position``, as only its template context holds them. + ``TaskInstanceInfo`` explains how these read together with the try and the region in events. """ in_loop: bool map_index: int + loop_index: int | None = None + loop_max_iterations: int | None = None + loop_group_id: str | None = None -def _is_task_in_loop(task: Any) -> bool: +def _get_enclosing_loop(task: Any) -> Any | None: from airflow.sdk.definitions._internal.loop import LoopTaskGroup from airflow.sdk.definitions.taskgroup import TaskGroup from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup, SerializedTaskGroup @@ -1508,9 +1595,9 @@ def _is_task_in_loop(task: Any) -> bool: group = getattr(task, "task_group", None) while isinstance(group, (TaskGroup, SerializedTaskGroup)): if isinstance(group, (LoopTaskGroup, SerializedLoopTaskGroup)): - return True + return group group = group.parent_group - return group is not None + return None def _get_task_instance_position( @@ -1526,9 +1613,26 @@ def _get_task_instance_position( return TaskInstancePosition(in_loop=False, map_index=map_index) if (task := getattr(task_instance, "task", None)) is None: return None + enclosing_loop = _get_enclosing_loop(task) + in_loop = enclosing_loop is not None + loop_index: int | None = None + loop_max_iterations: int | None = None + loop_group_id = enclosing_loop.group_id if in_loop else None + if in_loop and is_stored: + loop_max_iterations = getattr(enclosing_loop, "max_iterations", None) if is_stored and not task.get_needs_expansion(): + # The stored region_index means a map index only for a mapped task. A task that isn't mapped has none, + # and inside a loop its region_index is the iteration instead. + if in_loop: + loop_index = map_index map_index = -1 - return TaskInstancePosition(in_loop=_is_task_in_loop(task), map_index=map_index) + return TaskInstancePosition( + in_loop=in_loop, + map_index=map_index, + loop_index=loop_index, + loop_max_iterations=loop_max_iterations, + loop_group_id=loop_group_id, + ) def _require_task_instance_position( @@ -1542,25 +1646,46 @@ def _require_task_instance_position( return position +def _add_runtime_loop_iteration( + task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition +) -> TaskInstancePosition: + # Run ids, map indexes and the lineage macros don't need the iteration, so only this reads the template + # context: it re-validates the params on every call, which can raise from inside a user's template. + if not position.in_loop or position.loop_index is not None or isinstance(task_instance, TaskInstance): + return position + if (loop := task_instance.get_template_context().get("loop")) is None: + return position + return attrs.evolve(position, loop_index=loop.index, loop_max_iterations=loop.max_iterations) + + +def get_task_instance_position(task_instance: TaskInstance | RuntimeTaskInstance) -> TaskInstancePosition: + """Place a task instance with its loop iteration, once per event, to pass to ``get_airflow_run_facet``.""" + return _add_runtime_loop_iteration(task_instance, _require_task_instance_position(task_instance)) + + def get_orm_task_instance_positions( task_instances: Collection[TaskInstance], *, session: Session | None ) -> dict[UUID, TaskInstancePosition]: """ - Place stored task instances, reading the regions of those without their task in one query. + Place stored task instances, reading in one query the regions of those without their task and of mapped tasks in a loop. - Server side only: workers never hold a session, and an instance that needs its region read without one - raises instead of opening a connection. + Server side only: workers never hold a session. An instance without its task raises when there is no session to + read its region, while a mapped task in a loop is then placed without its iteration. """ positions: dict[UUID, TaskInstancePosition] = {} unplaced: list[TaskInstance] = [] + without_iteration: list[TaskInstance] = [] for task_instance in task_instances: if (position := _get_task_instance_position(task_instance)) is None: unplaced.append(task_instance) else: positions[task_instance.id] = position - if not unplaced: + # A stored mapped task in a loop keeps its iteration on its region only. Without a session it stays unknown. + if position.in_loop and position.loop_index is None and session is not None: + without_iteration.append(task_instance) + if not unplaced and not without_iteration: return positions - if session is None: + if unplaced and session is None: raise ValueError("Placing a stored task instance without its task needs the session it was loaded in") from sqlalchemy import select @@ -1570,8 +1695,15 @@ def get_orm_task_instance_positions( regions = { row.id: row for row in session.execute( - select(DynamicRegion.id, DynamicRegion.node_id, DynamicRegion.parent_region_id).where( - DynamicRegion.id.in_({task_instance.region_id for task_instance in unplaced}) + select( + DynamicRegion.id, + DynamicRegion.node_id, + DynamicRegion.parent_region_id, + DynamicRegion.parent_region_index, + ).where( + DynamicRegion.id.in_( + {task_instance.region_id for task_instance in (*unplaced, *without_iteration)} + ) ) ) } @@ -1579,10 +1711,25 @@ def get_orm_task_instance_positions( region = regions.get(task_instance.region_id) if region is not None and region.node_id == task_instance.task_id: positions[task_instance.id] = TaskInstancePosition( - in_loop=region.parent_region_id is not None, map_index=task_instance.region_index + in_loop=region.parent_region_id is not None, + map_index=task_instance.region_index, + loop_index=region.parent_region_index, ) else: - positions[task_instance.id] = TaskInstancePosition(in_loop=True, map_index=-1) + # Only a found region tells that the task isn't mapped, so that its region is the loop's and its + # region_index the iteration + positions[task_instance.id] = TaskInstancePosition( + in_loop=True, + map_index=-1, + loop_index=task_instance.region_index if region is not None else None, + loop_group_id=region.node_id if region is not None else None, + ) + for task_instance in without_iteration: + region = regions.get(task_instance.region_id) + if region is not None and region.parent_region_index is not None: + positions[task_instance.id] = attrs.evolve( + positions[task_instance.id], loop_index=region.parent_region_index + ) return positions @@ -2049,6 +2196,11 @@ def _get_task_groups_details(dag: DAG | SerializedDAG, edge_map: dict[str, tuple "ui_label": tg.label, "downstream_task_edges": task_edges, "downstream_group_edges": group_edges, + **{ + field: getattr(tg, field) # Only a loop group has these + for field in ("gate_task_id", "has_until", "max_iterations", "terminal_task_id") + if hasattr(tg, field) + }, } return result diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py new file mode 100644 index 0000000000000..39a63053cc2b4 --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py @@ -0,0 +1,68 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +DAG with a task that fails on its first try and succeeds on the retry. + +It checks: + - the first try emits START and FAIL, the retry emits START and COMPLETE + - taskInstance's `try_number` is 1 for the first try and 2 for the retry, and `max_tries` is 1 for both + - each try has its own runId (the runId check of OpenLineageTestOperator is enabled by default) +""" + +from __future__ import annotations + +from datetime import datetime, timedelta + +from airflow import DAG + +try: + from airflow.sdk import task +except ImportError: + # Airflow 2 path + from airflow.decorators import task # type: ignore[attr-defined,no-redef] + +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT +from system.openlineage.expected_events import get_expected_event_file_path +from system.openlineage.operator import OpenLineageTestOperator + +DAG_ID = "openlineage_retry_dag" + +with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, + dag_id=DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + catchup=False, + default_args={"retries": 1, "retry_delay": timedelta(seconds=1)}, +) as dag: + + @task + def flaky_task(**context): + if context["ti"].try_number == 1: + raise ValueError("Intentional failure of the first try") + + check_events = OpenLineageTestOperator( + task_id="check_events", file_path=get_expected_event_file_path(DAG_ID) + ) + + flaky_task() >> check_events + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py new file mode 100644 index 0000000000000..a75ea3bd78f39 --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py @@ -0,0 +1,67 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +DAG with a sensor in reschedule mode. + +It checks: + - the sensor emits one START (the START of later pokes is skipped) and one COMPLETE + - taskInstance's `try_number` stays at 1 for both, as rescheduling does not start a new try + - START and COMPLETE share the runId (the runId check of OpenLineageTestOperator is enabled by default) +""" + +from __future__ import annotations + +import os +from datetime import datetime, timedelta + +from airflow import DAG +from airflow.providers.standard.sensors.time_delta import TimeDeltaSensor + +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT +from system.openlineage.expected_events import get_expected_event_file_path +from system.openlineage.operator import OpenLineageTestOperator + +DAG_ID = "openlineage_sensor_reschedule_dag" + +with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, + dag_id=DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + catchup=False, + default_args={"retries": 0}, +) as dag: + # Timedelta is compared to the DAGRun start timestamp, which can occur long before a worker picks up the + # task. We need to ensure the sensor is rescheduled at least once, so setting 180s. + wait = TimeDeltaSensor( + task_id="wait", + delta=timedelta(seconds=int(os.getenv("SYSTEM_TESTS_OL_TIMEDELTA", "180"))), + poke_interval=10, + mode="reschedule", + ) + + check_events = OpenLineageTestOperator( + task_id="check_events", file_path=get_expected_event_file_path(DAG_ID) + ) + + wait >> check_events + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py new file mode 100644 index 0000000000000..2942ef16d166a --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py @@ -0,0 +1,122 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +DAG with a plain task and task loops, some of their tasks retrying (requires Airflow 3.4+). + +It checks, for a plain task, tasks in a loop, a mapped task in a loop and the loop gates: + - taskInstance's `try_number`, `map_index`, `loop_index` and `loop_max_iterations` of every event: the + iteration is the zero-based `loop_index`, `map_index` the mapped position, and a plain task has no + `loop_index` + - taskInstance's `loop_group_id` of every event: the loop group of a task in a loop, absent for a plain task + - taskInstance's `max_tries` of the tasks other than the gates, the `retries` of each task + - taskInstance's `id` of every event, and `region_id` and `region_index` of the tasks in a loop: the + `region_index` is the iteration for a task that isn't mapped and the map index for a mapped one, and a + plain task has neither + - the loop groups' `max_iterations`, `has_until`, `gate_task_id` and `terminal_task_id`, in the task's + `task_group` and in the `taskGroups` of the Dag's job facet + - a retried task emits START and FAIL for each failed try, then START and COMPLETE for the last one + - every execution (iteration, map index, try) has its own runId, shared by its START and COMPLETE/FAIL + (the runId check of OpenLineageTestOperator is enabled by default) +""" + +from __future__ import annotations + +from datetime import datetime, timedelta + +from airflow import DAG +from airflow.sdk import task, task_group + +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT +from system.openlineage.expected_events import get_expected_event_file_path +from system.openlineage.operator import OpenLineageTestOperator + +DAG_ID = "openlineage_task_loop_dag" + + +def sort_by_execution(event: dict) -> tuple[int, int, int]: + # Dag run events have no airflow facet + task_instance = event["run"]["facets"].get("airflow", {}).get("taskInstance", {}) + return ( + task_instance.get("loop_index", -1), + task_instance.get("map_index") if task_instance.get("map_index") is not None else -1, + task_instance.get("try_number", -1), + ) + + +with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, + dag_id=DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + catchup=False, + default_args={"retries": 0}, +) as dag: + + @task + def plain_task(): + return None + + @task_group + def accumulate(): + @task(retries=2, retry_delay=timedelta(seconds=1)) + def increment(*, loop, ti): + # The second iteration fails on its first try + if loop.index == 1 and ti.try_number < 2: + raise ValueError(f"Intentional failure, iteration {loop.index}, try {ti.try_number}") + return (0 if loop.index == 0 else loop.previous) + 1 + + @task + def report(value): + return value + + report(increment()) + + @task_group + def process_batch(): + @task + def choose_items(): + return [1, 2] + + @task( + max_active_tis_per_dagrun=1, # Execute sequentially to not overwrite each other OL events + retries=3, + retry_delay=timedelta(seconds=1), + ) + def process(value, *, loop, ti): + # The mapped instance of the second item fails on its first two tries, in every iteration + if value == 2 and ti.try_number < 3: + raise ValueError(f"Intentional failure, iteration {loop.index}, try {ti.try_number}") + return value + loop.index + + process.expand(value=choose_items()) + + accumulate_loop = accumulate.loop(max_iterations=3) + process_batch_loop = process_batch.loop(max_iterations=2) + + check_events = OpenLineageTestOperator( + task_id="check_events", + file_path=get_expected_event_file_path(DAG_ID), + event_sort_fn=sort_by_execution, + ) + + [plain_task(), accumulate_loop, process_batch_loop] >> check_events + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_manual_lineage_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_manual_lineage_dag.json index 452f3551746fc..2d3187f82d89c 100644 --- a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_manual_lineage_dag.json +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_manual_lineage_dag.json @@ -1,4 +1,20 @@ [ + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": "{{ result is number }}" + } + } + } + }, + "job": { + "name": "openlineage_manual_lineage_dag.datasets_minimal" + } + }, { "eventType": "RUNNING", "eventTime": "{{ is_datetime(result) }}", @@ -133,6 +149,38 @@ } } }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": "{{ result is number }}" + } + } + } + }, + "job": { + "name": "openlineage_manual_lineage_dag.datasets_minimal" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": "{{ result is number }}" + } + } + } + }, + "job": { + "name": "openlineage_manual_lineage_dag.datasets_maximal" + } + }, { "eventType": "RUNNING", "eventTime": "{{ is_datetime(result) }}", @@ -321,6 +369,22 @@ } } }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": "{{ result is number }}" + } + } + } + }, + "job": { + "name": "openlineage_manual_lineage_dag.datasets_maximal" + } + }, { "eventType": "START", "eventTime": "{{ is_datetime(result) }}", diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_retry_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_retry_dag.json new file mode 100644 index 0000000000000..678816d96241a --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_retry_dag.json @@ -0,0 +1,76 @@ +[ + { + "eventType": "START", + "job": { + "name": "openlineage_retry_dag" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 1 + } + } + } + }, + "job": { + "name": "openlineage_retry_dag.flaky_task" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 1 + } + } + } + }, + "job": { + "name": "openlineage_retry_dag.flaky_task" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "max_tries": 1 + } + } + } + }, + "job": { + "name": "openlineage_retry_dag.flaky_task" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "max_tries": 1 + } + } + } + }, + "job": { + "name": "openlineage_retry_dag.flaky_task" + } + } +] diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_sensor_reschedule_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_sensor_reschedule_dag.json new file mode 100644 index 0000000000000..8ab51fa095221 --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_sensor_reschedule_dag.json @@ -0,0 +1,34 @@ +[ + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1 + } + } + } + }, + "job": { + "name": "openlineage_sensor_reschedule_dag.wait" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1 + } + } + } + }, + "job": { + "name": "openlineage_sensor_reschedule_dag.wait" + } + } +] diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json new file mode 100644 index 0000000000000..4ba3373812418 --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json @@ -0,0 +1,1486 @@ +[ + { + "eventType": "START", + "job": { + "name": "openlineage_task_loop_dag", + "facets": { + "airflow": { + "taskGroups": { + "accumulate": { + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + }, + "process_batch": { + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": null, + "loop_max_iterations": null, + "id": "{{ is_uuid(result) }}", + "region_id": null, + "region_index": null, + "max_tries": 0, + "loop_group_id": null + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.plain_task" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": null, + "loop_max_iterations": null, + "id": "{{ is_uuid(result) }}", + "region_id": null, + "region_index": null, + "max_tries": 0, + "loop_group_id": null + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.plain_task" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.increment" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "max_tries": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.report" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 2, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 2, + "loop_group_id": "accumulate" + }, + "task": { + "task_group": { + "group_id": "accumulate", + "gate_task_id": "accumulate.__loop_gate", + "has_until": false, + "max_iterations": 3, + "terminal_task_id": "accumulate.report" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.accumulate.__loop_gate" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.choose_items" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.choose_items" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.choose_items" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.choose_items" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 0, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 3, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 0, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 3, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 0, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 3, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 0, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 3, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 0, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "FAIL", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 2, + "loop_index": 1, + "loop_max_iterations": 2, + "map_index": 1, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "max_tries": 3, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.process" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.__loop_gate" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.__loop_gate" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "process_batch" + }, + "task": { + "task_group": { + "group_id": "process_batch", + "gate_task_id": "process_batch.__loop_gate", + "has_until": false, + "max_iterations": 2, + "terminal_task_id": "process_batch.process" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.__loop_gate" + } + } +] diff --git a/providers/openlineage/tests/system/openlineage/operator.py b/providers/openlineage/tests/system/openlineage/operator.py index 802ac24320d60..1e4922f4bf360 100644 --- a/providers/openlineage/tests/system/openlineage/operator.py +++ b/providers/openlineage/tests/system/openlineage/operator.py @@ -23,8 +23,9 @@ import re import time import uuid -from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from collections import defaultdict +from collections.abc import Callable, Collection +from typing import TYPE_CHECKING, Any, NamedTuple from urllib.parse import urlparse from dateutil.parser import parse @@ -251,6 +252,128 @@ def match(expected, result, env: Environment, path: list | None = None) -> bool: return True +class Attempt(NamedTuple): + """ + Which execution of a job an event belongs to. + + ``try_number``, ``map_index``, ``loop_index`` and ``region_id`` are None for a Dag event and when the task has + none. ``clear_number`` plays for a Dag run the role ``try_number`` plays for a task instance, and ``region_id`` + tells apart two executions of one loop iteration after the loop was cleared, which can share every other field. + """ + + dag_run_id: Any + clear_number: Any + try_number: Any + map_index: Any + loop_index: Any + region_id: Any + + +def _get_attempt(event: dict) -> Attempt | None: + """Return the attempt of an event, or None for an event without the airflow facet.""" + facets = (event.get("run") or {}).get("facets") or {} + # Task events carry the airflow facet, Dag run events carry the same content in airflowDagRun. + airflow_facet = facets.get("airflow") or facets.get("airflowDagRun") + if not isinstance(airflow_facet, dict): + return None + dag_run = airflow_facet.get("dagRun") or {} + task_instance = airflow_facet.get("taskInstance") or {} + return Attempt( + dag_run_id=dag_run.get("run_id"), + clear_number=dag_run.get("clear_number"), + try_number=task_instance.get("try_number"), + map_index=task_instance.get("map_index"), + loop_index=task_instance.get("loop_index"), + region_id=task_instance.get("region_id"), + ) + + +def _get_job_type(event: dict) -> str | None: + return (((event.get("job") or {}).get("facets") or {}).get("jobType") or {}).get("jobType") + + +def check_run_id_integrity( + events_by_key: dict[str, list[dict]], repeated_start_keys: Collection[str] = () +) -> list[str]: + """ + Check that ``run.runId`` identifies exactly one attempt of exactly one job, and return the problems found. + + Events are grouped by ``runId``. The attempt of an event is its Dag run id and clear number, plus for a task + event ``(try_number, map_index, loop_index, region_id)``, all from ``run.facets.airflow`` (``airflowDagRun`` for Dag + events). Query events, which are child runs of a task, have no attempt. For every group: + + 1. all events share one ``job.name``; + 2. all events share one attempt, so another Dag run, clear, retry, map index or loop iteration can't reuse + a runId; + 3. the lifecycle is complete, so START and COMPLETE/FAIL of an attempt can't carry different runIds: + + * task and query: exactly one START, or several when the key of the START events is in + ``repeated_start_keys`` (deferral re-emits START on every resume), and exactly one COMPLETE or FAIL; + * Dag run: one START and at most one COMPLETE or FAIL. The Dag run that contains the checking task is + still running when the check happens, so its COMPLETE/FAIL is missing, while Dag runs triggered by it + can already have finished. + + Across groups: + + 4. one attempt of a job has one runId. + + Only the events passed in are checked, i.e. those of the keys listed in the expected events. + """ + problems: list[str] = [] + groups: dict[str, list[tuple[str, dict]]] = defaultdict(list) + for key, events in events_by_key.items(): + for event in events: + run_id = (event.get("run") or {}).get("runId") + if run_id: + groups[run_id].append((key, event)) + else: + problems.append(f"[{key}] {event.get('eventType')} event has no run.runId") + + run_ids_by_attempt: dict[tuple, set[str]] = defaultdict(set) + for run_id, members in sorted(groups.items()): + job_names = {event["job"]["name"] for _, event in members} + if len(job_names) > 1: + problems.append(f"runId {run_id} is shared by different jobs: {sorted(job_names)}") + continue + job_name = next(iter(job_names)) + + attempts = {attempt for _, event in members if (attempt := _get_attempt(event)) is not None} + if len(attempts) > 1: + problems.append( + f"runId {run_id} of `{job_name}` is shared by different attempts: {sorted(attempts, key=str)}" + ) + for attempt in attempts: + run_ids_by_attempt[(job_name, attempt)].add(run_id) + + job_type = _get_job_type(members[0][1]) + if job_type not in ("TASK", "QUERY", "DAG"): + continue + event_types = [event.get("eventType") for _, event in members] + start_keys = {key for key, event in members if event.get("eventType") == "START"} + terminals = [event_type for event_type in event_types if event_type in ("COMPLETE", "FAIL")] + starts = event_types.count("START") + repeated_starts_allowed = job_type != "DAG" and bool(start_keys & set(repeated_start_keys)) + too_many_starts = starts > 1 and not repeated_starts_allowed + too_many_terminals = len(terminals) > 1 or (job_type != "DAG" and not terminals) + if starts < 1 or too_many_starts or too_many_terminals: + expected = ( + "one START (several only if the Dag expects repeated STARTs, see event_count_assertions) " + "and one COMPLETE/FAIL" + if job_type != "DAG" + else "one START and at most one COMPLETE/FAIL" + ) + problems.append( + f"runId {run_id} of {job_type} `{job_name}` should have {expected}, " + f"got: {sorted(str(event_type) for event_type in event_types)}" + ) + + for attempt, run_ids in sorted(run_ids_by_attempt.items(), key=str): + if len(run_ids) > 1: + job_name, attempt_fields = attempt + problems.append(f"`{job_name}` {attempt_fields} has several runIds: {sorted(run_ids)}") + return problems + + class OpenLineageTestOperator(BaseOperator): """ This operator is added for system testing purposes. @@ -289,6 +412,10 @@ class OpenLineageTestOperator(BaseOperator): ascending before comparison. Useful for mapped tasks whose events arrive in non-deterministic order. When the template for a key is a list, sorting is required for pairwise comparison to be meaningful. + :param check_run_ids: if True (default), additionally check the collected events with + ``check_run_id_integrity``: START and COMPLETE/FAIL of one attempt share a ``runId``, while another + try, map index or loop iteration gets its own. Set to False only for a Dag whose expected events + intentionally list a task's lifecycle partially. Special sentinel values in expected dicts: @@ -316,6 +443,7 @@ def __init__( clear_variables: bool = True, fail_fast: bool = False, event_sort_fn: Callable[[dict], Any] | None = None, + check_run_ids: bool = True, **kwargs, ): super().__init__(**kwargs) @@ -325,6 +453,7 @@ def __init__( self.event_count_assertions = event_count_assertions self.clear_variables = clear_variables self.fail_fast = fail_fast + self.check_run_ids = check_run_ids self.event_sort_fn = event_sort_fn if event_templates is None and file_path is None: raise ValueError("Either event_templates or file_path must be provided") @@ -351,15 +480,25 @@ def execute(self, context: Context) -> None: else: self.event_templates[key] = [existing, event] errors: dict[str, str] = {} + events_by_key: dict[str, list[dict]] = {} try: for key, template in self.event_templates.items(): # type: ignore[union-attr] log.info("Checking key: `%s`", key) try: - self._check_key(key, template) + events_by_key[key] = self._check_key(key, template) except ValueError as e: if self.fail_fast: raise errors[key] = str(e) + if self.check_run_ids: + if problems := check_run_id_integrity( + events_by_key, + repeated_start_keys={key for key in events_by_key if self._expects_repeated_events(key)}, + ): + message = "\n".join(f" {problem}" for problem in problems) + if self.fail_fast: + raise ValueError(f"runId integrity check failed:\n{message}") + errors[""] = f"\n{message}" finally: if self.clear_variables: for key in self.event_templates: # type: ignore[union-attr] @@ -367,9 +506,10 @@ def execute(self, context: Context) -> None: Variable.delete(key=key) if errors: failures = "\n".join(f" [{key}] {error}" for key, error in errors.items()) + log.error("Event check failures:\n%s", failures) raise ValueError(f"Event check failures:\n{failures}") - def _check_key(self, key: str, template: Any) -> None: + def _check_key(self, key: str, template: Any) -> list[dict]: if isinstance(template, dict) and "$absent" in template: reason = template["$absent"] try: @@ -377,7 +517,7 @@ def _check_key(self, key: str, template: Any) -> None: raise ValueError(f"Expected no events for key {key} ({reason}), got {len(actual)}") except NoVariableError: log.info("Key `%s` absent as expected (%s).", key, reason) - return + return [] try: actual_events = Variable.get(key=key, deserialize_json=True) @@ -403,15 +543,17 @@ def _check_key(self, key: str, template: Any) -> None: default_count = str(len(template)) if isinstance(template, list) else "1" self._assert_event_count(key, len(actual_events), default=default_count) + parsed_events = [json.loads(evt_str) for evt_str in actual_events] if isinstance(template, list): # Multiple expected events: compare each one after sorting. - for i, (tmpl, evt_str) in enumerate(zip(template, actual_events)): - if not match(tmpl, json.loads(evt_str), self.env, [i]): + for i, (tmpl, event) in enumerate(zip(template, parsed_events)): + if not match(tmpl, event, self.env, [i]): raise ValueError(f"Event at index {i} does not match template for key `{key}`") else: # Last event is checked against the template - if not match(template, json.loads(actual_events[-1]), self.env): + if not match(template, parsed_events[-1], self.env): raise ValueError(f"Event received does not match one specified in test for key `{key}`") + return parsed_events def _assert_event_count(self, key: str, actual: int, default: str = "1") -> None: """ @@ -422,13 +564,25 @@ def _assert_event_count(self, key: str, actual: int, default: str = "1") -> None templates, str(len(template)) for list templates — supplied by the caller). Delegates the numeric check to _check_count_expr. """ - expr = default + self._check_count_expr(key, actual, self._get_count_expr(key, default)) + + def _get_count_expr(self, key: str, default: str) -> str: if self.event_count_assertions: for pattern, count_expr in self.event_count_assertions.items(): if re.fullmatch(pattern, key): - expr = str(count_expr) - break - self._check_count_expr(key, actual, expr) + return str(count_expr) + return default + + def _expects_repeated_events(self, key: str) -> bool: + """ + Tell whether the Dag declares that a START key holds several START events of one run. + + An open-ended ``">=N"`` (N >= 2) assertion on a START key does that, e.g. for a deferrable task that + emits START again after every resume. Exact counts don't: they also cover mapped tasks, where each + instance emits its own START under its own runId. + """ + expr = self._get_count_expr(key, default="1").strip() + return key.endswith(".event.start") and expr.startswith(">=") and int(expr[2:]) >= 2 @staticmethod def _check_count_expr(key: str, actual: int, expr: str) -> None: diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index d377e55d74abf..cf9e96c27367b 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -49,6 +49,7 @@ from airflow.providers.openlineage.plugins.macros import lineage_run_id from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy from airflow.providers.openlineage.utils.selective_enable import disable_lineage, enable_lineage +from airflow.providers.openlineage.utils.utils import TaskInstancePosition from airflow.utils import types from airflow.utils.state import DagRunState, State @@ -1209,6 +1210,87 @@ def body(): assert lineage_run_id(ti) == str(ti.id) listener.adapter.build_task_instance_run_id.assert_not_called() + @pytest.mark.parametrize("state", ["running", "success", "failed", "skipped"]) + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_worker_event_of_task_in_loop_has_its_iteration(self, state): + from airflow.sdk import DAG, task_group + from airflow.sdk.api.datamodels._generated import LoopContext + from airflow.sdk.definitions._internal.loop import create_loop + + listener, ti = self._create_listener_and_task_instance() + listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id + with DAG("loop_dag") as loop_dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + ti.task = loop_dag.get_task("body.work") + ti.region_id = uuid.uuid4() + ti.region_index = 1 + ti._ti_context_from_server.loop = LoopContext( + node_id="body", index=1, max_iterations=3, terminal_task_id="body.work", terminal_is_mapped=False + ) + kwargs = {"error": ValueError("test")} if state == "failed" else {} + + getattr(listener, f"on_task_instance_{state}")(None, ti, **kwargs) + + method = {"running": "start_task", "failed": "fail_task"}.get(state, "complete_task") + run_facets = getattr(listener.adapter, method).call_args.kwargs["run_facets"] + task_instance_facet = run_facets["airflow"].taskInstance + assert ( + task_instance_facet["loop_index"], + task_instance_facet["loop_max_iterations"], + task_instance_facet["loop_group_id"], + ) == (1, 3, "body") + + @pytest.mark.parametrize("state", ["running", "success", "failed", "skipped"]) + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_worker_event_run_id_of_mapped_task_has_its_map_index(self, state): + listener, ti = self._create_listener_and_task_instance() + ti.map_index = 2 + kwargs = {"error": ValueError("test")} if state == "failed" else {} + + getattr(listener, f"on_task_instance_{state}")(None, ti, **kwargs) + + assert listener.adapter.build_task_instance_run_id.call_args.kwargs["map_index"] == 2 + + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") + def test_manual_event_passes_its_stored_position_to_the_airflow_run_facet(self, mocker): + from airflow.sdk import DAG, task_group + from airflow.sdk.definitions._internal.loop import create_loop + + listener, ti = self._create_listener_and_task_instance(runtime_ti=False) + listener._executor = mocker.Mock(spec=ProcessPoolExecutor) + with DAG("loop_dag") as loop_dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + ti.task = loop_dag.get_task("body.work") + ti.region_id = uuid.uuid4() + ti.region_index = 2 + get_airflow_run_facet = mocker.patch( + "airflow.providers.openlineage.plugins.listener.get_airflow_run_facet", + autospec=True, + return_value={}, + ) + mocker.patch.object(listener, "submit_callable", autospec=True) + + listener._on_task_instance_manual_state_change(ti, ti.dag_run, "success") + + assert get_airflow_run_facet.call_args.kwargs["ti_position"] == TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=2, loop_max_iterations=3, loop_group_id="body" + ) + @pytest.mark.skip("Rendering fields is not migrated yet in Airflow 3") @patch("airflow.models.BaseOperator.render_template") def test_listener_does_not_change_task_instance(self, render_mock, mock_supervisor_comms, spy_agency): diff --git a/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py new file mode 100644 index 0000000000000..9403ee2e5a045 --- /dev/null +++ b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py @@ -0,0 +1,250 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Tests for ``check_run_id_integrity``, the runId check of the OpenLineage system tests' checking operator.""" + +from __future__ import annotations + +import pytest + +from system.openlineage.operator import OpenLineageTestOperator, check_run_id_integrity + + +def _event( + run_id: str | None, + event_type: str = "START", + *, + job: str = "dag.task", + job_type: str = "TASK", + dag_run_id: str = "run_1", + clear_number: int = 0, + try_number: int = 1, + map_index: int | None = None, + loop_index: int | None = None, + region_id: str | None = None, +) -> dict: + run: dict = {"facets": {}} + if run_id is not None: + run["runId"] = run_id + dag_run = {"run_id": dag_run_id, "clear_number": clear_number} + if job_type == "TASK": + task_instance = {"try_number": try_number} + if map_index is not None: + task_instance["map_index"] = map_index + if loop_index is not None: + task_instance["loop_index"] = loop_index + if region_id is not None: + task_instance["region_id"] = region_id + run["facets"]["airflow"] = {"dagRun": dag_run, "taskInstance": task_instance} + elif job_type == "DAG": + run["facets"]["airflowDagRun"] = {"dagRun": dag_run} + return { + "eventType": event_type, + "job": {"name": job, "facets": {"jobType": {"jobType": job_type}}}, + "run": run, + } + + +def _dag_event(run_id: str, event_type: str = "START", **kwargs) -> dict: + return _event(run_id, event_type, job="dag", job_type="DAG", **kwargs) + + +def _query_event(run_id: str, event_type: str = "START", **kwargs) -> dict: + return _event(run_id, event_type, job="dag.task.query.1", job_type="QUERY", **kwargs) + + +VALID_SCENARIOS = { + "start_and_complete": [_event("a"), _event("a", "COMPLETE")], + "start_and_fail": [_event("a"), _event("a", "FAIL")], + "retry": [ + _event("a"), + _event("a", "FAIL"), + _event("b", try_number=2), + _event("b", "COMPLETE", try_number=2), + ], + "mapped": [ + _event("a", map_index=0), + _event("a", "COMPLETE", map_index=0), + _event("b", map_index=1), + _event("b", "COMPLETE", map_index=1), + ], + "loop_iterations": [ + event + for index, run_id in enumerate("abc") + for event in (_event(run_id, loop_index=index), _event(run_id, "COMPLETE", loop_index=index)) + ], + "mapped_in_loop": [ + event + for run_id, map_index, loop_index in (("a", 0, 0), ("b", 1, 0), ("c", 0, 1), ("d", 1, 1)) + for event in ( + _event(run_id, map_index=map_index, loop_index=loop_index), + _event(run_id, "COMPLETE", map_index=map_index, loop_index=loop_index), + ) + ], + "loop_iteration_run_again_after_a_clear": [ + _event("a", loop_index=1, region_id="loop_region"), + _event("a", "COMPLETE", loop_index=1, region_id="loop_region"), + _event("b", loop_index=1, region_id="forked_loop_region"), + _event("b", "COMPLETE", loop_index=1, region_id="forked_loop_region"), + ], + "same_task_in_two_dag_runs": [ + _event("a", dag_run_id="run_1"), + _event("a", "COMPLETE", dag_run_id="run_1"), + _event("b", dag_run_id="run_2"), + _event("b", "COMPLETE", dag_run_id="run_2"), + ], + "query_start_and_complete": [_query_event("a"), _query_event("a", "COMPLETE")], + "query_start_and_fail": [_query_event("a"), _query_event("a", "FAIL")], + "dag_run_start_only": [_dag_event("a")], + "dag_run_start_and_complete": [_dag_event("a"), _dag_event("a", "COMPLETE")], + "dag_run_cleared": [ + _dag_event("a"), + _dag_event("a", "FAIL"), + _dag_event("b", clear_number=1), + _dag_event("b", "COMPLETE", clear_number=1), + ], +} + +INVALID_SCENARIOS = { + "start_and_complete_with_different_run_ids": ( + [_event("a"), _event("b", "COMPLETE")], + "has several runIds", + ), + "complete_without_start": ([_event("a", "COMPLETE")], "one START"), + "start_without_terminal": ([_event("a")], "one COMPLETE/FAIL"), + "two_terminal_events": ([_event("a"), _event("a", "FAIL"), _event("a", "COMPLETE")], "one COMPLETE/FAIL"), + "repeated_start_not_declared": ( + [_event("a"), _event("a"), _event("a", "COMPLETE")], + "one START (several only if", + ), + "retry_reusing_run_id": ( + [_event("a"), _event("a", "FAIL"), _event("a", try_number=2), _event("a", "COMPLETE", try_number=2)], + "shared by different attempts", + ), + "map_indexes_sharing_run_id": ( + [ + _event("a", map_index=0), + _event("a", map_index=1), + _event("a", "COMPLETE", map_index=0), + ], + "shared by different attempts", + ), + "loop_iterations_sharing_run_id": ( + [_event("a", loop_index=0), _event("a", "COMPLETE", loop_index=0), _event("a", loop_index=1)], + "shared by different attempts", + ), + "loop_iteration_runs_of_two_trajectories_sharing_run_id": ( + [ + _event("a", loop_index=1, region_id="loop_region"), + _event("a", "COMPLETE", loop_index=1, region_id="loop_region"), + _event("a", loop_index=1, region_id="forked_loop_region"), + ], + "shared by different attempts", + ), + "dag_runs_sharing_run_id": ( + [_event("a", dag_run_id="run_1"), _event("a", "COMPLETE", dag_run_id="run_2")], + "shared by different attempts", + ), + "clear_numbers_sharing_run_id": ( + [_dag_event("a"), _dag_event("a", "COMPLETE", clear_number=1)], + "shared by different attempts", + ), + "dag_run_with_different_run_ids_per_event": ( + [_dag_event("a"), _dag_event("b", "COMPLETE")], + "has several runIds", + ), + "query_start_without_terminal": ([_query_event("a")], "one COMPLETE/FAIL"), + "query_complete_without_start": ([_query_event("a", "COMPLETE")], "one START"), + "query_two_terminal_events": ( + [_query_event("a"), _query_event("a", "COMPLETE"), _query_event("a", "FAIL")], + "one COMPLETE/FAIL", + ), + "dag_run_complete_without_start": ([_dag_event("a", "COMPLETE")], "one START and at most one"), + "dag_run_two_terminal_events": ( + [_dag_event("a"), _dag_event("a", "COMPLETE"), _dag_event("a", "FAIL")], + "one START and at most one", + ), + "jobs_sharing_run_id": ( + [_event("a", job="dag.one"), _event("a", "COMPLETE", job="dag.two")], + "shared by different jobs", + ), + "missing_run_id": ([_event(None)], "has no run.runId"), +} + + +@pytest.mark.parametrize("events", VALID_SCENARIOS.values(), ids=VALID_SCENARIOS.keys()) +def test_valid_events_have_no_problems(events): + assert check_run_id_integrity({"dag.task.event.start": events}) == [] + + +@pytest.mark.parametrize( + ("events", "expected_problem"), INVALID_SCENARIOS.values(), ids=INVALID_SCENARIOS.keys() +) +def test_invalid_events_are_reported(events, expected_problem): + problems = check_run_id_integrity({"dag.task.event.start": events}) + + assert any(expected_problem in problem for problem in problems), problems + + +def test_events_of_one_attempt_are_joined_across_keys(): + events_by_key = { + "dag.task.event.start": [_event("a")], + "dag.task.event.complete": [_event("a", "COMPLETE")], + } + + assert check_run_id_integrity(events_by_key) == [] + + +@pytest.mark.parametrize( + ("repeated_start_keys", "has_problems"), + [((), True), (("dag.task.event.start",), False), (("other.event.start",), True)], +) +def test_repeated_starts_are_allowed_only_for_declared_keys(repeated_start_keys, has_problems): + events_by_key = { + "dag.task.event.start": [_event("a"), _event("a"), _event("a")], + "dag.task.event.complete": [_event("a", "COMPLETE")], + } + + problems = check_run_id_integrity(events_by_key, repeated_start_keys=repeated_start_keys) + + assert bool(problems) is has_problems + + +def test_repeated_starts_are_never_allowed_for_a_dag_run(): + events_by_key = {"dag.event.start": [_dag_event("a"), _dag_event("a")]} + + assert check_run_id_integrity(events_by_key, repeated_start_keys=("dag.event.start",)) + + +@pytest.mark.parametrize( + ("event_count_assertions", "expected"), + [ + ({"dag.task.event.start": ">=2"}, True), + ({"dag.task.event.start": ">=3"}, True), + ({"dag.task.event.start": ">=1"}, False), + ({"dag.task.event.start": "==2"}, False), + ({"dag.task.event.start": 2}, False), + ({".*\\.event\\.start": ">=2"}, True), + ({"dag.task.event.complete": ">=2"}, False), + (None, False), + ], +) +def test_operator_derives_repeated_starts_from_event_count_assertions(event_count_assertions, expected): + operator = OpenLineageTestOperator( + task_id="check_events", event_templates={}, event_count_assertions=event_count_assertions + ) + + assert operator._expects_repeated_events("dag.task.event.start") is expected diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index 1b66b80ab5ad0..947f3293bfcea 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -29,6 +29,7 @@ import pendulum import pytest from openlineage.client.facet_v2 import parent_run +from sqlalchemy.orm import Session from sqlalchemy.orm.exc import DetachedInstanceError from uuid6 import uuid7 @@ -81,6 +82,8 @@ get_root_information_from_dagrun_conf, get_runtime_outlet_assets, get_task_documentation, + get_task_instance_map_index, + get_task_instance_position, get_task_parent_run_facet, get_user_provided_run_facets, is_dag_run_asset_triggered, @@ -2011,6 +2014,33 @@ def test_get_task_groups_details(): assert result == expected +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +def test_get_task_groups_details_has_loop_attributes_only_for_a_loop(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("loop_dag", schedule=None, start_date=datetime.datetime(2024, 6, 1)) as dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + with TaskGroup("plain"): + EmptyOperator(task_id="work") + + result = _get_task_groups_details(dag, edge_map=_build_labeled_edge_map(dag)) + + loop_attributes = {"max_iterations", "has_until", "terminal_task_id", "gate_task_id"} + assert {k: result["body"][k] for k in loop_attributes} == { + "max_iterations": 3, + "has_until": False, + "terminal_task_id": "body.work", + "gate_task_id": "body.__loop_gate", + } + assert not loop_attributes & set(result["plain"]) + + def test_get_task_groups_details_nested(): with DAG("test_dag", schedule=None, start_date=datetime.datetime(2024, 6, 1)) as dag: with TaskGroup("tg1", prefix_group_id=True) as tg: @@ -3266,10 +3296,13 @@ def test_taskinstance_info_af3(): bundle_instance = MagicMock(version="bundle_version") bundle_instance.name = "bundle_name" runtime_ti.bundle_instance = bundle_instance + runtime_ti.max_tries = 3 assert dict(TaskInstanceInfo(runtime_ti)) == { + "id": str(runtime_ti.id), "log_url": runtime_ti.log_url, "map_index": 2, + "max_tries": 3, "note": None, "rendered_map_index": None, "try_number": 1, @@ -3305,6 +3338,7 @@ def test_taskinstance_info_af2(): "rendered_map_index": None, "pool": "default_pool", "try_number": 0, + "max_tries": 0, "queued_dttm": "2024-06-01T00:00:00+00:00", "log_url": "some_log_url", "dag_bundle_name": None, @@ -3924,15 +3958,251 @@ def body(): create_loop(body, max_iterations=2) else: body() - ti = SimpleNamespace(id=uuid4(), region_id=uuid4(), task=dag.get_task("body.work")) + ti = SimpleNamespace( + id=uuid4(), + region_id=uuid4(), + map_index=-1, + task=dag.get_task("body.work"), + get_template_context=MagicMock(), + ) assert get_regional_task_instance_run_id(ti) == (str(ti.id) if in_loop else None) + assert get_task_instance_map_index(ti) == -1 + ti.get_template_context.assert_not_called() + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +@pytest.mark.parametrize("has_loop_context", [True, False]) +def test_runtime_task_instance_position_has_the_iteration_of_its_loop_context(has_loop_context): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("position_dag") as dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + loop = SimpleNamespace(index=2, max_iterations=3) if has_loop_context else None + ti = SimpleNamespace( + id=uuid4(), + region_id=uuid4(), + map_index=-1, + task=dag.get_task("body.work"), + get_template_context=MagicMock(return_value={"loop": loop}), + ) + + assert get_task_instance_position(ti) == TaskInstancePosition( + in_loop=True, + map_index=-1, + loop_index=2 if has_loop_context else None, + loop_max_iterations=3 if has_loop_context else None, + loop_group_id="body", + ) + ti.get_template_context.assert_called_once_with() -@pytest.mark.db_test @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +def test_stored_task_instance_without_task_or_region_has_no_loop_index(): + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + ti.dag_id = "dag_id" + ti.task_id = "task_id" + ti.region_id = uuid4() + ti.region_index = 1 + ti.task = None + session = MagicMock(spec=Session) + session.execute.return_value = [] + + assert get_orm_task_instance_positions([ti], session=session) == { + ti.id: TaskInstancePosition(in_loop=True, map_index=-1, loop_index=None) + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") @pytest.mark.parametrize("mapped", [False, True]) -def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag_maker, session, mapped): +@pytest.mark.parametrize("in_loop", [False, True]) +def test_runtime_task_instance_info_has_loop_index_only_inside_a_loop(in_loop, mapped): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("loop_index_dag") as dag: + + @task_group + def body(): + if mapped: + PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) + else: + EmptyOperator(task_id="work") + + if in_loop: + create_loop(body, max_iterations=3) + else: + body() + ti = SimpleNamespace( + id=uuid4(), + region_id=uuid4(), + map_index=1 if mapped else -1, + task=dag.get_task("body.work"), + ) + if in_loop: + ti.get_template_context = lambda: {"loop": SimpleNamespace(index=2, max_iterations=3)} + + info = dict(TaskInstanceInfo(ti)) + + assert info["map_index"] == (1 if mapped else None) + if in_loop: + assert (info["loop_index"], info["loop_max_iterations"], info["loop_group_id"]) == (2, 3, "body") + else: + assert "loop_index" not in info + assert "loop_max_iterations" not in info + assert "loop_group_id" not in info + + +@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Attempt UUIDs require Airflow 3") +@pytest.mark.parametrize("stored", [False, True]) +@pytest.mark.parametrize("in_region", [False, True]) +def test_task_instance_info_has_attempt_id_and_region_only_in_a_region(stored, in_region): + region_id = uuid4() if in_region else UUID(int=0) + if stored: + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + else: + ti = SimpleNamespace(id=uuid7(), map_index=-1) + ti.region_id = region_id + ti.region_index = 2 if in_region else -1 + + info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1))) + + assert info["id"] == str(ti.id) + if in_region: + assert (info["region_id"], info["region_index"]) == (str(region_id), 2) + else: + assert "region_id" not in info + assert "region_index" not in info + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +def test_stored_task_instance_in_loop_has_loop_index_and_max_iterations(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("stored_loop_dag") as dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 2 + + assert get_orm_task_instance_positions([ti], session=None) == { + ti.id: TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=2, loop_max_iterations=3, loop_group_id="body" + ) + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +@pytest.mark.parametrize("with_session", [True, False]) +def test_stored_mapped_task_in_loop_gets_its_iteration_from_its_region(with_session): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("stored_mapped_loop_dag") as dag: + + @task_group + def body(): + PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) + + create_loop(body, max_iterations=3) + ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 1 + session = MagicMock(spec=Session) + session.execute.return_value = [ + SimpleNamespace(id=ti.region_id, node_id="body.work", parent_region_id=uuid4(), parent_region_index=2) + ] + + positions = get_orm_task_instance_positions([ti], session=session if with_session else None) + + assert positions == { + ti.id: TaskInstancePosition( + in_loop=True, + map_index=1, + loop_index=2 if with_session else None, + loop_max_iterations=3, + loop_group_id="body", + ) + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +def test_task_instance_info_of_stored_instance_without_task_does_not_raise(): + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + ti.dag_id = "dag_id" + ti.task_id = "task_id" + ti.region_id = uuid4() + ti.region_index = 2 + ti.task = None + + info = dict(TaskInstanceInfo(ti)) + + assert info["map_index"] is None + assert "loop_index" not in info + + +def test_task_instance_info_uses_the_position_it_is_given(): + ti = SimpleNamespace(id=uuid4(), region_id=uuid4(), map_index=-1) + position = TaskInstancePosition( + in_loop=True, map_index=1, loop_index=2, loop_max_iterations=3, loop_group_id="body" + ) + + info = dict(TaskInstanceInfo(ti, position=position)) + + assert (info["map_index"], info["loop_index"], info["loop_max_iterations"], info["loop_group_id"]) == ( + 1, + 2, + 3, + "body", + ) + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +def test_task_group_info_has_loop_attributes_only_for_a_loop(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("loop_group_dag"): + + @task_group + def body(): + EmptyOperator(task_id="work") + + loop = create_loop(body, max_iterations=3) + + @task_group + def plain(): + EmptyOperator(task_id="work") + + plain_group = plain() + + loop_info = dict(TaskGroupInfo(loop)) + assert loop_info["max_iterations"] == 3 + assert loop_info["has_until"] is False + assert loop_info["terminal_task_id"] == "body.work" + assert loop_info["gate_task_id"] == "body.__loop_gate" + assert not {"max_iterations", "has_until", "terminal_task_id", "gate_task_id"} & set( + TaskGroupInfo(plain_group) + ) + + +@pytest.mark.db_test +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +@pytest.mark.parametrize("kind", ["plain_loop", "mapped", "mapped_in_loop"]) +def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag_maker, session, kind): + mapped = kind != "plain_loop" with dag_maker(serialized=True): if mapped: PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) @@ -3940,12 +4210,19 @@ def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag EmptyOperator(task_id="work") run = dag_maker.create_dagrun() ti = max(run.task_instances, key=lambda ti: ti.region_index) - if not mapped: + if kind == "plain_loop": region = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id, node_id="loop") session.add(region) session.flush() ti.region_id = region.id ti.region_index = 2 + elif kind == "mapped_in_loop": + loop_region = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id, node_id="loop") + session.add(loop_region) + session.flush() + mapped_region = session.get(DynamicRegion, ti.region_id) + mapped_region.parent_region_id = loop_region.id + mapped_region.parent_region_index = 2 session.flush() ti.task = None @@ -3955,7 +4232,12 @@ def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag positions = get_orm_task_instance_positions([ti], session=session) assert positions == { - ti.id: TaskInstancePosition(in_loop=not mapped, map_index=ti.region_index if mapped else -1) + ti.id: TaskInstancePosition( + in_loop=kind != "mapped", + map_index=ti.region_index if mapped else -1, + loop_index=2 if kind != "mapped" else None, + loop_group_id="loop" if kind == "plain_loop" else None, + ) } From d6b586cbb85713cb4e71e7077b13dfeeaa30d22e Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Fri, 16 Oct 2026 17:49:58 +0200 Subject: [PATCH 02/10] Fix mypy and docs spelling errors in OpenLineage task loop changes Provider mypy and the docs spellcheck run in CI through breeze rather than the default prek hooks, so they were not caught locally before the previous push. Co-Authored-By: Claude --- .../src/airflow/providers/openlineage/utils/utils.py | 5 +++-- .../openlineage/example_openlineage_retry_dag.py | 2 +- .../example_openlineage_sensor_reschedule_dag.py | 2 +- .../openlineage/example_openlineage_task_loop_dag.py | 4 ++-- .../openlineage/tests/system/openlineage/operator.py | 11 +++++------ .../openlineage/test_system_tests_run_id_integrity.py | 2 +- 6 files changed, 13 insertions(+), 13 deletions(-) diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index 437de4fb2c9b9..31a509c5efb27 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1617,7 +1617,7 @@ def _get_task_instance_position( in_loop = enclosing_loop is not None loop_index: int | None = None loop_max_iterations: int | None = None - loop_group_id = enclosing_loop.group_id if in_loop else None + loop_group_id = enclosing_loop.group_id if enclosing_loop is not None else None if in_loop and is_stored: loop_max_iterations = getattr(enclosing_loop, "max_iterations", None) if is_stored and not task.get_needs_expansion(): @@ -1685,7 +1685,8 @@ def get_orm_task_instance_positions( without_iteration.append(task_instance) if not unplaced and not without_iteration: return positions - if unplaced and session is None: + if session is None: + # Only unplaced instances get here without a session, as no mapped task is added to without_iteration then raise ValueError("Placing a stored task instance without its task needs the session it was loaded in") from sqlalchemy import select diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py index 39a63053cc2b4..67da53b777c4d 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py @@ -20,7 +20,7 @@ It checks: - the first try emits START and FAIL, the retry emits START and COMPLETE - taskInstance's `try_number` is 1 for the first try and 2 for the retry, and `max_tries` is 1 for both - - each try has its own runId (the runId check of OpenLineageTestOperator is enabled by default) + - each try has its own `runId` (the `runId` check of OpenLineageTestOperator is enabled by default) """ from __future__ import annotations diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py index a75ea3bd78f39..802974519c687 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_sensor_reschedule_dag.py @@ -20,7 +20,7 @@ It checks: - the sensor emits one START (the START of later pokes is skipped) and one COMPLETE - taskInstance's `try_number` stays at 1 for both, as rescheduling does not start a new try - - START and COMPLETE share the runId (the runId check of OpenLineageTestOperator is enabled by default) + - START and COMPLETE share the `runId` (the `runId` check of OpenLineageTestOperator is enabled by default) """ from __future__ import annotations diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py index 2942ef16d166a..a3747b063b6e4 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py @@ -29,8 +29,8 @@ - the loop groups' `max_iterations`, `has_until`, `gate_task_id` and `terminal_task_id`, in the task's `task_group` and in the `taskGroups` of the Dag's job facet - a retried task emits START and FAIL for each failed try, then START and COMPLETE for the last one - - every execution (iteration, map index, try) has its own runId, shared by its START and COMPLETE/FAIL - (the runId check of OpenLineageTestOperator is enabled by default) + - every execution (iteration, map index, try) has its own `runId`, shared by its START and COMPLETE/FAIL + (the `runId` check of OpenLineageTestOperator is enabled by default) """ from __future__ import annotations diff --git a/providers/openlineage/tests/system/openlineage/operator.py b/providers/openlineage/tests/system/openlineage/operator.py index 1e4922f4bf360..b187590d3632d 100644 --- a/providers/openlineage/tests/system/openlineage/operator.py +++ b/providers/openlineage/tests/system/openlineage/operator.py @@ -304,8 +304,8 @@ def check_run_id_integrity( 1. all events share one ``job.name``; 2. all events share one attempt, so another Dag run, clear, retry, map index or loop iteration can't reuse - a runId; - 3. the lifecycle is complete, so START and COMPLETE/FAIL of an attempt can't carry different runIds: + a ``runId``; + 3. the lifecycle is complete, so START and COMPLETE/FAIL of an attempt can't carry different ``runId`` values: * task and query: exactly one START, or several when the key of the START events is in ``repeated_start_keys`` (deferral re-emits START on every resume), and exactly one COMPLETE or FAIL; @@ -315,7 +315,7 @@ def check_run_id_integrity( Across groups: - 4. one attempt of a job has one runId. + 4. one attempt of a job has one ``runId``. Only the events passed in are checked, i.e. those of the keys listed in the expected events. """ @@ -367,9 +367,8 @@ def check_run_id_integrity( f"got: {sorted(str(event_type) for event_type in event_types)}" ) - for attempt, run_ids in sorted(run_ids_by_attempt.items(), key=str): + for (job_name, attempt_fields), run_ids in sorted(run_ids_by_attempt.items(), key=str): if len(run_ids) > 1: - job_name, attempt_fields = attempt problems.append(f"`{job_name}` {attempt_fields} has several runIds: {sorted(run_ids)}") return problems @@ -579,7 +578,7 @@ def _expects_repeated_events(self, key: str) -> bool: An open-ended ``">=N"`` (N >= 2) assertion on a START key does that, e.g. for a deferrable task that emits START again after every resume. Exact counts don't: they also cover mapped tasks, where each - instance emits its own START under its own runId. + instance emits its own START under its own ``runId``. """ expr = self._get_count_expr(key, default="1").strip() return key.endswith(".event.start") and expr.startswith(">=") and int(expr[2:]) >= 2 diff --git a/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py index 9403ee2e5a045..92ba3afedf7b4 100644 --- a/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py +++ b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py @@ -41,7 +41,7 @@ def _event( run["runId"] = run_id dag_run = {"run_id": dag_run_id, "clear_number": clear_number} if job_type == "TASK": - task_instance = {"try_number": try_number} + task_instance: dict[str, int | str] = {"try_number": try_number} if map_index is not None: task_instance["map_index"] = map_index if loop_index is not None: From 5e3414aa90c5cc2ca2e94195f3e171426fe38a39 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Fri, 16 Oct 2026 19:41:37 +0200 Subject: [PATCH 03/10] Expose loop gate decisions in OpenLineage task instance facet Alerting on why a loop stopped (for example, its until condition still being false at max_iterations) needs the gate's decision in the event itself, without reading the Airflow metadata database. Co-Authored-By: Claude --- .../openlineage/facets/AirflowRunFacet.json | 14 + .../providers/openlineage/utils/utils.py | 21 ++ .../example_openlineage_task_loop_dag.py | 24 +- .../openlineage_task_loop_dag.json | 314 +++++++++++++++++- .../unit/openlineage/plugins/test_listener.py | 25 ++ .../unit/openlineage/utils/test_utils.py | 43 +++ 6 files changed, 426 insertions(+), 15 deletions(-) diff --git a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json index 8312c2b37d9a0..a6c2e16be28a0 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json +++ b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json @@ -214,6 +214,20 @@ "description": "Group id of the innermost loop the task runs in (Airflow 3.4+), which loop_index counts the iterations of. Absent outside a loop.", "type": "string" }, + "loop_decision": { + "description": "Decision of a loop gate (Airflow 3.4+), on its COMPLETE or FAIL event only. A FAIL with stop false at the last iteration means the loop's until was still false at max_iterations.", + "type": "object", + "properties": { + "stop": { + "description": "Whether the gate stopped the loop.", + "type": "boolean" + }, + "reason": { + "description": "Reason given by until, if any.", + "type": "string" + } + } + }, "region_id": { "description": "Id of the region (one execution of a loop or of a mapped task's expansion) the task instance belongs to (Airflow 3.4+). Shared by the task instances of that region. Clearing a loop forks its region, so a re-run of an iteration gets a new region_id. Region ids are UUIDv7, so a later trajectory has a greater one, which is an implementation detail. Absent for a task that is neither mapped nor in a loop.", "type": "string", diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index 31a509c5efb27..0a76400f87b86 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1149,6 +1149,22 @@ class TaskInstanceInfo(InfoJsonEncodable): region_id, region_index, try_number)`` is, as Airflow enforces it (``task_instance_try_key``). Both are absent for a task in no region, i.e. one that is neither mapped nor in a loop. + ``loop_decision`` (Airflow 3.4+): ``stop`` and the optional ``reason`` returned by ``until``, as evaluated by + a loop's gate, the task that ends each iteration (``gate_task_id`` of the loop group, in ``task.task_group`` + and in ``taskGroups`` of ``AirflowJobFacet``). Only the gate's COMPLETE and FAIL events carry it, as the gate + decides while it runs. Read with ``loop_index``, ``loop_max_iterations`` and the group's ``has_until``: + + * COMPLETE with ``stop`` false: the loop goes on to iteration ``loop_index + 1``. + * COMPLETE with ``stop`` true and ``has_until``: ``until`` ended the loop at iteration ``loop_index``, which + can be the last one. + * COMPLETE with ``stop`` true and no ``has_until``: the loop ran all its iterations, as one without ``until`` + always stops at ``loop_index == loop_max_iterations - 1``. + * FAIL with ``stop`` false: ``until`` was still false at the last iteration (``loop_index == + loop_max_iterations - 1``), so the loop ran out of iterations and fails. That's the only failure after the + decision is made. + * FAIL without ``loop_decision``: the gate failed before deciding, e.g. ``until`` raised or returned something + other than a bool or a ``LoopDecision``; the error facet tells which. + How a cleared loop reads: clearing a loop from iteration ``k`` retries what was cleared in iteration ``k`` in place, with the same region and the next ``try_number``. The iterations after ``k`` are created again in a new region forked from the loop's region, and try numbers count per region, so the new execution of @@ -1204,6 +1220,11 @@ def _extend_fields(self) -> None: if AIRFLOW_V_3_0_PLUS and (ti_id := getattr(self.obj, "id", None)) is not None: self.id = str(ti_id) self._fields.append("id") + # Only a loop gate's runtime task instance has one, once it has evaluated ``until``. Read from a private + # attribute until the task SDK exposes it. + if (decision := getattr(self.obj, "_loop_decision", None)) is not None: + self.loop_decision = {"stop": decision.stop, "reason": decision.reason} + self._fields.append("loop_decision") # A runtime task instance in no region carries the sentinel region id instead of None region_id = getattr(self.obj, "region_id", None) if isinstance(region_id, UUID) and region_id.int != 0: diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py index a3747b063b6e4..ccd5e90a60fab 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. """ -DAG with a plain task and task loops, some of their tasks retrying (requires Airflow 3.4+). +DAG with a plain task and task loops, some of their tasks retrying and one stopped early by `until` +(requires Airflow 3.4+). It checks, for a plain task, tasks in a loop, a mapped task in a loop and the loop gates: - taskInstance's `try_number`, `map_index`, `loop_index` and `loop_max_iterations` of every event: the @@ -26,6 +27,10 @@ - taskInstance's `id` of every event, and `region_id` and `region_index` of the tasks in a loop: the `region_index` is the iteration for a task that isn't mapped and the map index for a mapped one, and a plain task has neither + - the gates' `loop_decision`: absent on START; for the loops without `until`, `stop` true only at the last + iteration and no `reason`; for the loop with `until` (whose gate is named after it), the `reason` it + returns and `stop` true at the iteration it ends the loop, before `max_iterations`; absent for the plain + task - the loop groups' `max_iterations`, `has_until`, `gate_task_id` and `terminal_task_id`, in the task's `task_group` and in the `taskGroups` of the Dag's job facet - a retried task emits START and FAIL for each failed try, then START and COMPLETE for the last one @@ -38,7 +43,7 @@ from datetime import datetime, timedelta from airflow import DAG -from airflow.sdk import task, task_group +from airflow.sdk import LoopDecision, task, task_group from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path @@ -104,8 +109,21 @@ def process(value, *, loop, ti): process.expand(value=choose_items()) + @task_group + def poll(): + @task + def check(*, loop): + return loop.index + + check() + + def ready(*, loop): + # Stops at the second of three iterations, giving a reason both when continuing and when stopping + return LoopDecision(stop=loop.result >= 1, reason=f"check returned {loop.result}") + accumulate_loop = accumulate.loop(max_iterations=3) process_batch_loop = process_batch.loop(max_iterations=2) + poll_loop = poll.loop(max_iterations=3, until=ready) check_events = OpenLineageTestOperator( task_id="check_events", @@ -113,7 +131,7 @@ def process(value, *, loop, ti): event_sort_fn=sort_by_execution, ) - [plain_task(), accumulate_loop, process_batch_loop] >> check_events + [plain_task(), accumulate_loop, process_batch_loop, poll_loop] >> check_events from tests_common.test_utils.system_tests import get_test_run # noqa: E402 diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json index 4ba3373812418..b42dec0cb2216 100644 --- a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json @@ -17,6 +17,12 @@ "has_until": false, "max_iterations": 2, "terminal_task_id": "process_batch.process" + }, + "poll": { + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" } } } @@ -37,7 +43,8 @@ "region_id": null, "region_index": null, "max_tries": 0, - "loop_group_id": null + "loop_group_id": null, + "loop_decision": null } } } @@ -60,7 +67,8 @@ "region_id": null, "region_index": null, "max_tries": 0, - "loop_group_id": null + "loop_group_id": null, + "loop_decision": null } } } @@ -530,7 +538,8 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 0, - "loop_group_id": "accumulate" + "loop_group_id": "accumulate", + "loop_decision": null }, "task": { "task_group": { @@ -561,7 +570,8 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, - "loop_group_id": "accumulate" + "loop_group_id": "accumulate", + "loop_decision": null }, "task": { "task_group": { @@ -592,7 +602,8 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 2, - "loop_group_id": "accumulate" + "loop_group_id": "accumulate", + "loop_decision": null }, "task": { "task_group": { @@ -623,7 +634,10 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 0, - "loop_group_id": "accumulate" + "loop_group_id": "accumulate", + "loop_decision": { + "stop": false + } }, "task": { "task_group": { @@ -654,7 +668,10 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, - "loop_group_id": "accumulate" + "loop_group_id": "accumulate", + "loop_decision": { + "stop": false + } }, "task": { "task_group": { @@ -685,7 +702,10 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 2, - "loop_group_id": "accumulate" + "loop_group_id": "accumulate", + "loop_decision": { + "stop": true + } }, "task": { "task_group": { @@ -1372,7 +1392,8 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 0, - "loop_group_id": "process_batch" + "loop_group_id": "process_batch", + "loop_decision": null }, "task": { "task_group": { @@ -1403,7 +1424,8 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, - "loop_group_id": "process_batch" + "loop_group_id": "process_batch", + "loop_decision": null }, "task": { "task_group": { @@ -1434,7 +1456,10 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 0, - "loop_group_id": "process_batch" + "loop_group_id": "process_batch", + "loop_decision": { + "stop": false + } }, "task": { "task_group": { @@ -1465,7 +1490,10 @@ "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, - "loop_group_id": "process_batch" + "loop_group_id": "process_batch", + "loop_decision": { + "stop": true + } }, "task": { "task_group": { @@ -1482,5 +1510,267 @@ "job": { "name": "openlineage_task_loop_dag.process_batch.__loop_gate" } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "poll", + "max_tries": 0 + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.check" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "poll", + "max_tries": 0 + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.check" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "poll", + "max_tries": 0 + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.check" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "poll", + "max_tries": 0 + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.check" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "poll", + "loop_decision": null + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.ready" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "poll", + "loop_decision": null + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.ready" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 0, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0, + "loop_group_id": "poll", + "loop_decision": { + "stop": false, + "reason": "check returned 0" + } + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.ready" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "loop_index": 1, + "loop_max_iterations": 3, + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1, + "loop_group_id": "poll", + "loop_decision": { + "stop": true, + "reason": "check returned 1" + } + }, + "task": { + "task_group": { + "group_id": "poll", + "gate_task_id": "poll.ready", + "has_until": true, + "max_iterations": 3, + "terminal_task_id": "poll.check" + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.poll.ready" + } } ] diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index cf9e96c27367b..69aadf932bed5 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -1248,6 +1248,31 @@ def body(): task_instance_facet["loop_group_id"], ) == (1, 3, "body") + @pytest.mark.parametrize("state", ["running", "success", "failed"]) + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") + @mock.patch( + "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call + ) + def test_worker_event_of_loop_gate_has_its_decision_once_evaluated(self, state): + from airflow.sdk.api.datamodels._generated import LoopDecisionPayload + + listener, ti = self._create_listener_and_task_instance() + listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id + if state != "running": + ti._loop_decision = LoopDecisionPayload(stop=state == "success", reason="checked") + kwargs = {"error": ValueError("test")} if state == "failed" else {} + + getattr(listener, f"on_task_instance_{state}")(None, ti, **kwargs) + + method = {"running": "start_task", "failed": "fail_task"}.get(state, "complete_task") + task_instance_facet = ( + getattr(listener.adapter, method).call_args.kwargs["run_facets"]["airflow"].taskInstance + ) + if state == "running": + assert "loop_decision" not in task_instance_facet + else: + assert task_instance_facet["loop_decision"] == {"stop": state == "success", "reason": "checked"} + @pytest.mark.parametrize("state", ["running", "success", "failed", "skipped"]) @mock.patch( "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index 947f3293bfcea..e6fbcebb19b86 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -4059,6 +4059,49 @@ def body(): assert "loop_group_id" not in info +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +@pytest.mark.parametrize( + ("until", "loop_index", "expected", "runs_out"), + [ + pytest.param(None, 1, {"stop": False, "reason": None}, False, id="no_until_continues"), + pytest.param(None, 2, {"stop": True, "reason": None}, False, id="no_until_stops_at_last"), + pytest.param("done", 0, {"stop": True, "reason": "done"}, False, id="until_stops_with_reason"), + pytest.param(False, 2, {"stop": False, "reason": None}, True, id="until_false_at_last"), + ], +) +def test_task_instance_info_has_the_decision_of_a_loop_gate(until, loop_index, expected, runs_out): + from airflow.sdk import LoopDecision + from airflow.sdk.definitions._internal.loop import LoopGateOperator + from airflow.sdk.execution_time.loop import LoopMaxIterationsExceeded, execute_loop_gate + + if until is False: + until_callable = lambda: False + elif until is not None: + until_callable = lambda: LoopDecision(stop=True, reason=until) + else: + until_callable = None + gate = LoopGateOperator(task_id="__loop_gate", until=until_callable) + ti = SimpleNamespace(id=uuid7(), map_index=-1) + context = {"loop": SimpleNamespace(index=loop_index, max_iterations=3), "ti": ti} + + if runs_out: + with pytest.raises(LoopMaxIterationsExceeded): + execute_loop_gate(gate, context) + else: + execute_loop_gate(gate, context) + info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1))) + + assert info["loop_decision"] == expected + + +def test_task_instance_info_has_no_loop_decision_for_other_tasks(): + ti = SimpleNamespace(id=uuid7(), map_index=-1) + + info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1))) + + assert "loop_decision" not in info + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Attempt UUIDs require Airflow 3") @pytest.mark.parametrize("stored", [False, True]) @pytest.mark.parametrize("in_region", [False, True]) From c3daca8c68e4ac16f0fb3f719cadfcac28fff0cf Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Mon, 19 Oct 2026 13:18:37 +0200 Subject: [PATCH 04/10] Keep OpenLineage events when the loop iteration can't be read The run id must not depend on the optional loop iteration, and the lineage macro should place a task instance the same way the listener does. Co-Authored-By: Claude --- .../providers/openlineage/plugins/macros.py | 6 ++-- .../providers/openlineage/utils/utils.py | 16 ++++++--- .../unit/openlineage/utils/test_utils.py | 33 ++++++++++++++++--- 3 files changed, 44 insertions(+), 11 deletions(-) diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py index 5ce197f0ae632..bce984aeac15a 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py @@ -23,9 +23,9 @@ from airflow.providers.openlineage.utils.utils import ( get_job_name, get_parent_information_from_dagrun_conf, - get_regional_task_instance_run_id, get_root_information_from_dagrun_conf, get_task_instance_map_index, + is_task_instance_in_loop, ) from airflow.providers.openlineage.version_compat import AIRFLOW_V_3_0_PLUS @@ -65,7 +65,9 @@ def lineage_run_id(task_instance: TaskInstance): For more information take a look at the guide: :ref:`howto/macros:openlineage` """ - return get_regional_task_instance_run_id(task_instance) or OpenLineageAdapter.build_task_instance_run_id( + if is_task_instance_in_loop(task_instance): + return str(task_instance.id) + return OpenLineageAdapter.build_task_instance_run_id( dag_id=task_instance.dag_id, task_id=task_instance.task_id, try_number=task_instance.try_number, diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index 0a76400f87b86..ce23ed1efd693 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1674,7 +1674,14 @@ def _add_runtime_loop_iteration( # context: it re-validates the params on every call, which can raise from inside a user's template. if not position.in_loop or position.loop_index is not None or isinstance(task_instance, TaskInstance): return position - if (loop := task_instance.get_template_context().get("loop")) is None: + try: + loop = task_instance.get_template_context().get("loop") + except Exception: + # The iteration only describes the run, so failing to read it must not cost the event its run id + log.warning("OpenLineage failed to read the loop iteration of task instance %s", task_instance.id) + log.debug("Exception details:", exc_info=True) + return position + if loop is None: return position return attrs.evolve(position, loop_index=loop.index, loop_max_iterations=loop.max_iterations) @@ -1755,10 +1762,9 @@ def get_orm_task_instance_positions( return positions -def get_regional_task_instance_run_id(task_instance: TaskInstance | RuntimeTaskInstance) -> str | None: - if _require_task_instance_position(task_instance).in_loop: - return str(task_instance.id) - return None +def is_task_instance_in_loop(task_instance: TaskInstance | RuntimeTaskInstance) -> bool: + """Tell whether a task instance runs in a loop, where only its attempt UUID identifies the run.""" + return _require_task_instance_position(task_instance).in_loop def build_task_instance_ol_run_id( diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index e6fbcebb19b86..be2d76b56c08a 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -78,7 +78,6 @@ get_operator_provider_version, get_orm_task_instance_positions, get_parent_information_from_dagrun_conf, - get_regional_task_instance_run_id, get_root_information_from_dagrun_conf, get_runtime_outlet_assets, get_task_documentation, @@ -87,6 +86,7 @@ get_task_parent_run_facet, get_user_provided_run_facets, is_dag_run_asset_triggered, + is_task_instance_in_loop, is_valid_uuid, translate_airflow_asset, ) @@ -3944,7 +3944,7 @@ def test_orm_lineage_facets_project_public_index(dag_maker, session, mapped): @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") @pytest.mark.parametrize("in_loop", [False, True]) -def test_regional_run_id_is_the_execution_uuid_only_inside_a_loop(in_loop): +def test_task_instance_is_in_loop_only_inside_a_loop(in_loop): from airflow.sdk import task_group from airflow.sdk.definitions._internal.loop import create_loop @@ -3966,7 +3966,7 @@ def body(): get_template_context=MagicMock(), ) - assert get_regional_task_instance_run_id(ti) == (str(ti.id) if in_loop else None) + assert is_task_instance_in_loop(ti) is in_loop assert get_task_instance_map_index(ti) == -1 ti.get_template_context.assert_not_called() @@ -4003,6 +4003,31 @@ def body(): ti.get_template_context.assert_called_once_with() +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +def test_runtime_task_instance_position_without_its_iteration_when_the_context_fails(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("position_dag") as dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + ti = SimpleNamespace( + id=uuid4(), + region_id=uuid4(), + map_index=-1, + task=dag.get_task("body.work"), + get_template_context=MagicMock(side_effect=ValueError("invalid params")), + ) + + assert get_task_instance_position(ti) == TaskInstancePosition( + in_loop=True, map_index=-1, loop_group_id="body" + ) + + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") def test_stored_task_instance_without_task_or_region_has_no_loop_index(): ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) @@ -4270,7 +4295,7 @@ def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag ti.task = None with pytest.raises(ValueError, match="has no task attached"): - get_regional_task_instance_run_id(ti) + is_task_instance_in_loop(ti) with assert_queries_count(1): positions = get_orm_task_instance_positions([ti], session=session) From e6454fb1115e9835a98d1f50cc40ffcbfb48fe92 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Mon, 19 Oct 2026 16:26:07 +0200 Subject: [PATCH 05/10] Address review of OpenLineage task loop changes A region of an unexpected shape, or a failed lookup of a mapped task's iteration, must not merge the runs of loop iterations or cost an event. Tests now pin the region shapes Airflow creates, so a change to them fails the provider's tests instead of silently changing events. Co-Authored-By: Claude --- .../openlineage/facets/AirflowJobFacet.json | 25 +- .../openlineage/facets/AirflowRunFacet.json | 16 + .../providers/openlineage/plugins/listener.py | 56 +--- .../providers/openlineage/plugins/macros.py | 13 +- .../providers/openlineage/utils/utils.py | 189 +++++++---- .../example_openlineage_task_loop_dag.py | 16 +- .../openlineage_task_loop_dag.json | 271 ++++++++++++++++ .../tests/system/openlineage/operator.py | 7 +- .../unit/openlineage/plugins/test_listener.py | 87 ++--- .../test_system_tests_run_id_integrity.py | 17 + .../unit/openlineage/utils/test_utils.py | 303 +++++++++++++++++- 11 files changed, 827 insertions(+), 173 deletions(-) diff --git a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowJobFacet.json b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowJobFacet.json index 51a9954de376a..b6095f00677b7 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowJobFacet.json +++ b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowJobFacet.json @@ -15,9 +15,30 @@ "additionalProperties": true }, "taskGroups": { - "description": "Information about all task groups within the DAG.", + "description": "Information about all task groups within the DAG, keyed by group id.", "type": "object", - "additionalProperties": true + "additionalProperties": { + "type": "object", + "properties": { + "max_iterations": { + "description": "Maximum number of iterations of the loop (Airflow 3.4+). Present for a loop group only.", + "type": "integer" + }, + "has_until": { + "description": "Whether the loop has an until condition that can stop it before max_iterations (Airflow 3.4+). Present for a loop group only.", + "type": "boolean" + }, + "gate_task_id": { + "description": "Id of the loop's gate, the task that ends each iteration and decides whether the loop goes on (Airflow 3.4+). Present for a loop group only.", + "type": "string" + }, + "terminal_task_id": { + "description": "Id of the last task of the loop body, whose result the gate and the next iteration read (Airflow 3.4+). Present for a loop group only.", + "type": "string" + } + }, + "additionalProperties": true + } }, "tasks": { "description": "Details of all individual tasks within the DAG.", diff --git a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json index a6c2e16be28a0..bfba75b5415fe 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json +++ b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json @@ -142,6 +142,22 @@ }, "upstream_task_ids": { "type": "string" + }, + "max_iterations": { + "description": "Maximum number of iterations of the loop (Airflow 3.4+). Present for a loop group only.", + "type": "integer" + }, + "has_until": { + "description": "Whether the loop has an until condition that can stop it before max_iterations (Airflow 3.4+). Present for a loop group only.", + "type": "boolean" + }, + "gate_task_id": { + "description": "Id of the loop's gate, the task that ends each iteration and decides whether the loop goes on (Airflow 3.4+). Present for a loop group only.", + "type": "string" + }, + "terminal_task_id": { + "description": "Id of the last task of the loop body, whose result the gate and the next iteration read (Airflow 3.4+). Present for a loop group only.", + "type": "string" } }, "additionalProperties": true, diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py index b1fc356f577f3..dda70e50fecd8 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py @@ -55,6 +55,7 @@ get_orm_task_instance_positions, get_task_documentation, get_task_instance_position, + get_task_instance_run_id, get_task_parent_run_facet, get_user_provided_run_facets, is_dag_run_asset_triggered, @@ -264,16 +265,7 @@ def on_running(): ) position = get_task_instance_position(task_instance) - if position.in_loop: - task_uuid = str(task_instance.id) - else: - task_uuid = self.adapter.build_task_instance_run_id( - dag_id=task_instance.dag_id, - task_id=task_instance.task_id, - try_number=task_instance.try_number, - logical_date=date, - map_index=position.map_index, - ) + task_uuid = get_task_instance_run_id(task_instance, logical_date=date, position=position) event_type = RunState.RUNNING.value.lower() operator_name = task.task_type.lower() @@ -431,16 +423,7 @@ def on_success(): ) position = get_task_instance_position(task_instance) - if position.in_loop: - task_uuid = str(task_instance.id) - else: - task_uuid = self.adapter.build_task_instance_run_id( - dag_id=task_instance.dag_id, - task_id=task_instance.task_id, - try_number=task_instance.try_number, - logical_date=date, - map_index=position.map_index, - ) + task_uuid = get_task_instance_run_id(task_instance, logical_date=date, position=position) event_type = RunState.COMPLETE.value.lower() operator_name = task.task_type.lower() @@ -612,16 +595,7 @@ def on_failure(): ) position = get_task_instance_position(task_instance) - if position.in_loop: - task_uuid = str(task_instance.id) - else: - task_uuid = self.adapter.build_task_instance_run_id( - dag_id=task_instance.dag_id, - task_id=task_instance.task_id, - try_number=task_instance.try_number, - logical_date=date, - map_index=position.map_index, - ) + task_uuid = get_task_instance_run_id(task_instance, logical_date=date, position=position) event_type = RunState.FAIL.value.lower() operator_name = task.task_type.lower() @@ -769,16 +743,7 @@ def on_skipped(): ) position = get_task_instance_position(task_instance) - if position.in_loop: - task_uuid = str(task_instance.id) - else: - task_uuid = self.adapter.build_task_instance_run_id( - dag_id=task_instance.dag_id, - task_id=task_instance.task_id, - try_number=task_instance.try_number, - logical_date=date, - map_index=position.map_index, - ) + task_uuid = get_task_instance_run_id(task_instance, logical_date=date, position=position) event_type = RunState.COMPLETE.value.lower() operator_name = task.task_type.lower() @@ -925,16 +890,7 @@ def _on_task_instance_manual_state_change( # the equivalent note in `on_dag_run_running` (listener.py ~868). date = dagrun.logical_date or dagrun.run_after position = get_orm_task_instance_positions([ti], session=object_session(ti))[ti.id] - if position.in_loop: - task_uuid = str(ti.id) - else: - task_uuid = self.adapter.build_task_instance_run_id( - dag_id=ti.dag_id, - task_id=ti.task_id, - try_number=ti.try_number, - logical_date=date, - map_index=position.map_index, - ) + task_uuid = get_task_instance_run_id(ti, logical_date=date, position=position) parent_run_id = self.adapter.build_dag_run_id( dag_id=ti.dag_id, logical_date=date, diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py index bce984aeac15a..5039e235b0b6d 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/macros.py @@ -24,8 +24,7 @@ get_job_name, get_parent_information_from_dagrun_conf, get_root_information_from_dagrun_conf, - get_task_instance_map_index, - is_task_instance_in_loop, + get_task_instance_run_id, ) from airflow.providers.openlineage.version_compat import AIRFLOW_V_3_0_PLUS @@ -65,15 +64,7 @@ def lineage_run_id(task_instance: TaskInstance): For more information take a look at the guide: :ref:`howto/macros:openlineage` """ - if is_task_instance_in_loop(task_instance): - return str(task_instance.id) - return OpenLineageAdapter.build_task_instance_run_id( - dag_id=task_instance.dag_id, - task_id=task_instance.task_id, - try_number=task_instance.try_number, - logical_date=_get_logical_date(task_instance), - map_index=get_task_instance_map_index(task_instance), - ) + return get_task_instance_run_id(task_instance, logical_date=_get_logical_date(task_instance)) def lineage_parent_id(task_instance: TaskInstance): diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index ce23ed1efd693..1ce42cf421aca 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1194,16 +1194,19 @@ class TaskInstanceInfo(InfoJsonEncodable): } def __init__(self, obj, position: TaskInstancePosition | None = None): - # Pass the position when the caller has already placed the task instance, e.g. from its stored region + # Pass the position when the caller has already placed the task instance, e.g. from its stored region or + # with get_task_instance_position; it is then used as is. self._position = position super().__init__(obj) def _extend_fields(self) -> None: - # A stored task instance whose task isn't attached can't be placed without its region, and a facet must - # not fail the whole event, so then it describes the task instance without its position. - position = self._position or _get_task_instance_position(self.obj) - if position is not None: - position = _add_runtime_loop_iteration(self.obj, position) + position = self._position + if position is None: + # A stored task instance whose task isn't attached can't be placed without its region, and a facet + # must not fail the whole event, so then it describes the task instance without its position. + position = _get_task_instance_position(self.obj) + if position is not None: + position = _add_runtime_loop_iteration(self.obj, position) map_index = position.map_index if position else -1 self.map_index = map_index if map_index != -1 else None self.rendered_map_index = getattr(self.obj, "rendered_map_index", None) if map_index != -1 else None @@ -1223,7 +1226,9 @@ def _extend_fields(self) -> None: # Only a loop gate's runtime task instance has one, once it has evaluated ``until``. Read from a private # attribute until the task SDK exposes it. if (decision := getattr(self.obj, "_loop_decision", None)) is not None: - self.loop_decision = {"stop": decision.stop, "reason": decision.reason} + self.loop_decision = {"stop": decision.stop} + if decision.reason is not None: + self.loop_decision["reason"] = decision.reason self._fields.append("loop_decision") # A runtime task instance in no region carries the sentinel region id instead of None region_id = getattr(self.obj, "region_id", None) @@ -1587,18 +1592,13 @@ def is_dag_run_asset_triggered( @attrs.frozen class TaskInstancePosition: """ - Where a task instance runs: inside a loop body or not, the map index users see and the loop iteration. - - A plain mapped task keeps the region of its own expansion, so its dag, task, try, logical date and - map index stay unique; inside a loop they repeat on every pass, so only the execution UUID identifies - the run. ``loop_index`` is the zero-based iteration and ``loop_max_iterations`` the limit of the loop, the - counterparts of ``try_number`` and the maximum number of tries, and ``loop_group_id`` the loop's group id. - All three describe the innermost loop and are ``None`` outside a loop. A stored task whose task isn't - attached has no ``loop_max_iterations``, nor ``loop_group_id`` when it is mapped, and a stored mapped task - in a loop has no ``loop_index`` unless ``get_orm_task_instance_positions`` could read its region. A runtime - task instance has no ``loop_index`` nor ``loop_max_iterations`` unless placed with - ``get_task_instance_position``, as only its template context holds them. - ``TaskInstanceInfo`` explains how these read together with the try and the region in events. + Where a task instance runs: in a loop or not, the map index users see, and its loop iteration. + + The loop fields are ``None`` outside a loop, and also when they can't be known: a stored task whose task + isn't attached has no ``loop_max_iterations``, nor ``loop_group_id`` when it is mapped; a stored mapped task + has no ``loop_index`` unless ``get_orm_task_instance_positions`` could read its region; and a runtime task + instance has no ``loop_index`` nor ``loop_max_iterations`` unless placed with ``get_task_instance_position``, + as only its template context holds them. ``TaskInstanceInfo`` explains what these fields mean in events. """ in_loop: bool @@ -1608,7 +1608,8 @@ class TaskInstancePosition: loop_group_id: str | None = None -def _get_enclosing_loop(task: Any) -> Any | None: +def _find_enclosing_loop(task: Any) -> tuple[bool, Any | None]: + """Tell whether a task runs in a loop, and return that loop's group when it is known.""" from airflow.sdk.definitions._internal.loop import LoopTaskGroup from airflow.sdk.definitions.taskgroup import TaskGroup from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup, SerializedTaskGroup @@ -1616,9 +1617,19 @@ def _get_enclosing_loop(task: Any) -> Any | None: group = getattr(task, "task_group", None) while isinstance(group, (TaskGroup, SerializedTaskGroup)): if isinstance(group, (LoopTaskGroup, SerializedLoopTaskGroup)): - return group + return True, group group = group.parent_group - return None + if group is None: + return False, None + # A group of an unknown type could be a loop. Assuming one only gives each iteration its own runId, while + # missing one would merge the runs of all its iterations, which share every coordinate but the UUID. + log.warning( + "OpenLineage can't tell whether task `%s` runs in a loop, as its task group has an unknown type `%s`; " + "assuming it does", + getattr(task, "task_id", None), + type(group).__name__, + ) + return True, None def _get_task_instance_position( @@ -1634,12 +1645,11 @@ def _get_task_instance_position( return TaskInstancePosition(in_loop=False, map_index=map_index) if (task := getattr(task_instance, "task", None)) is None: return None - enclosing_loop = _get_enclosing_loop(task) - in_loop = enclosing_loop is not None + in_loop, enclosing_loop = _find_enclosing_loop(task) loop_index: int | None = None loop_max_iterations: int | None = None loop_group_id = enclosing_loop.group_id if enclosing_loop is not None else None - if in_loop and is_stored: + if enclosing_loop is not None and is_stored: loop_max_iterations = getattr(enclosing_loop, "max_iterations", None) if is_stored and not task.get_needs_expansion(): # The stored region_index means a map index only for a mapped task. A task that isn't mapped has none, @@ -1718,55 +1728,95 @@ def get_orm_task_instance_positions( raise ValueError("Placing a stored task instance without its task needs the session it was loaded in") from sqlalchemy import select + from sqlalchemy.orm import aliased from airflow.models.dynamic_region import DynamicRegion - regions = { - row.id: row - for row in session.execute( - select( - DynamicRegion.id, - DynamicRegion.node_id, - DynamicRegion.parent_region_id, - DynamicRegion.parent_region_index, - ).where( - DynamicRegion.id.in_( - {task_instance.region_id for task_instance in (*unplaced, *without_iteration)} - ) + parent = aliased(DynamicRegion) + query = ( + select( + DynamicRegion.id, + DynamicRegion.node_id, + DynamicRegion.parent_region_id, + DynamicRegion.parent_region_index, + parent.node_id.label("parent_node_id"), + parent.parent_region_id.label("grandparent_region_id"), + ) + .outerjoin(parent, parent.id == DynamicRegion.parent_region_id) + .where( + DynamicRegion.id.in_( + {task_instance.region_id for task_instance in (*unplaced, *without_iteration)} ) ) - } + ) + if unplaced: + regions = {row.id: row for row in session.execute(query)} + else: + # Only mapped tasks' iterations are missing, and they only describe the run, so failing to read them + # must not cost the event. The savepoint keeps the caller's transaction usable after a failure. + try: + with session.begin_nested(): + regions = {row.id: row for row in session.execute(query)} + except Exception: + log.warning( + "OpenLineage failed to read the loop iteration of mapped task instances from their regions" + ) + log.debug("Exception details:", exc_info=True) + return positions + + # Airflow creates a region for each run of a loop and for the expansion of each mapped task, a task in a + # mapped task group included. As loops can't be nested nor sit in a mapped task group, a loop's region is + # top-level and a mapped task's region sits directly in its loop's region, if any. The provider's + # test_region_shapes_* tests check this against Airflow. Any other shape is placed as a loop without its + # iteration, which keeps the runId of each execution unique, instead of guessing. for task_instance in unplaced: region = regions.get(task_instance.region_id) + position = None + map_index = -1 if region is not None and region.node_id == task_instance.task_id: - positions[task_instance.id] = TaskInstancePosition( - in_loop=region.parent_region_id is not None, - map_index=task_instance.region_index, - loop_index=region.parent_region_index, - ) - else: - # Only a found region tells that the task isn't mapped, so that its region is the loop's and its - # region_index the iteration - positions[task_instance.id] = TaskInstancePosition( + map_index = task_instance.region_index + if region.parent_region_id is None: + position = TaskInstancePosition(in_loop=False, map_index=map_index) + elif region.grandparent_region_id is None: + position = TaskInstancePosition( + in_loop=True, + map_index=map_index, + loop_index=region.parent_region_index, + loop_group_id=region.parent_node_id, + ) + elif region is not None and region.parent_region_id is None: + position = TaskInstancePosition( in_loop=True, map_index=-1, - loop_index=task_instance.region_index if region is not None else None, - loop_group_id=region.node_id if region is not None else None, + loop_index=task_instance.region_index, + loop_group_id=region.node_id, + ) + if position is None: + log.warning( + "OpenLineage can't place task instance %s from its region, which has an unexpected shape; " + "assuming it runs in a loop", + task_instance.id, ) + position = TaskInstancePosition(in_loop=True, map_index=map_index) + positions[task_instance.id] = position for task_instance in without_iteration: + position = positions[task_instance.id] region = regions.get(task_instance.region_id) - if region is not None and region.parent_region_index is not None: - positions[task_instance.id] = attrs.evolve( - positions[task_instance.id], loop_index=region.parent_region_index + if ( + region is not None + and region.parent_node_id == position.loop_group_id + and region.grandparent_region_id is None + ): + positions[task_instance.id] = attrs.evolve(position, loop_index=region.parent_region_index) + else: + log.warning( + "OpenLineage can't read the loop iteration of task instance %s, as its region isn't directly " + "in its loop's region", + task_instance.id, ) return positions -def is_task_instance_in_loop(task_instance: TaskInstance | RuntimeTaskInstance) -> bool: - """Tell whether a task instance runs in a loop, where only its attempt UUID identifies the run.""" - return _require_task_instance_position(task_instance).in_loop - - def build_task_instance_ol_run_id( dag_id: str, task_id: str, @@ -1795,6 +1845,31 @@ def build_task_instance_ol_run_id( ) +def get_task_instance_run_id( + task_instance: TaskInstance | RuntimeTaskInstance, + *, + logical_date: datetime.datetime, + position: TaskInstancePosition | None = None, +) -> str: + """ + Return the OpenLineage runId of a task instance's attempt, which all its events must share. + + In a loop the iterations share every coordinate but the attempt UUID, so that is the runId. Otherwise it is + built from the coordinates. + """ + if position is None: + position = _require_task_instance_position(task_instance) + if position.in_loop: + return str(task_instance.id) + return build_task_instance_ol_run_id( + dag_id=task_instance.dag_id, + task_id=task_instance.task_id, + try_number=task_instance.try_number, + logical_date=logical_date, + map_index=position.map_index, + ) + + def is_valid_uuid(uuid_string: str | None) -> bool: """Validate that a string is a valid UUID format.""" if uuid_string is None: diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py index ccd5e90a60fab..cbb4cd7916e51 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py @@ -18,7 +18,8 @@ DAG with a plain task and task loops, some of their tasks retrying and one stopped early by `until` (requires Airflow 3.4+). -It checks, for a plain task, tasks in a loop, a mapped task in a loop and the loop gates: +It checks, for a plain task, tasks in a loop, a mapped task and a mapped task group in a loop, and the loop +gates: - taskInstance's `try_number`, `map_index`, `loop_index` and `loop_max_iterations` of every event: the iteration is the zero-based `loop_index`, `map_index` the mapped position, and a plain task has no `loop_index` @@ -31,6 +32,8 @@ iteration and no `reason`; for the loop with `until` (whose gate is named after it), the `reason` it returns and `stop` true at the iteration it ends the loop, before `max_iterations`; absent for the plain task + - a mapped task group inside a loop: its task's `map_index` and `region_index` are the mapped position, + its `loop_index` the iteration, and its `loop_group_id` the loop's, unlike its own `task_group` - the loop groups' `max_iterations`, `has_until`, `gate_task_id` and `terminal_task_id`, in the task's `task_group` and in the `taskGroups` of the Dag's job facet - a retried task emits START and FAIL for each failed try, then START and COMPLETE for the last one @@ -107,7 +110,16 @@ def process(value, *, loop, ti): raise ValueError(f"Intentional failure, iteration {loop.index}, try {ti.try_number}") return value + loop.index - process.expand(value=choose_items()) + @task_group + def fetch(source): + @task(max_active_tis_per_dagrun=1) # Execute sequentially to not overwrite each other OL events + def download(source, *, loop): + return f"{source}-{loop.index}" + + download(source) + + items = choose_items() + fetch.expand(source=["a", "b"]) >> process.expand(value=items) @task_group def poll(): diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json index b42dec0cb2216..6fc30d3d73258 100644 --- a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json @@ -23,6 +23,13 @@ "has_until": true, "max_iterations": 3, "terminal_task_id": "poll.check" + }, + "process_batch.fetch": { + "parent_group": "process_batch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null } } } @@ -1772,5 +1779,269 @@ "job": { "name": "openlineage_task_loop_dag.poll.ready" } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 0, + "loop_index": 0, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 0, + "loop_index": 1, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "START", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 0, + "loop_index": 0, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 1, + "loop_index": 0, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 0, + "loop_index": 1, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 0 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } + }, + { + "eventType": "COMPLETE", + "run": { + "runId": "{{ is_uuid(result) }}", + "facets": { + "airflow": { + "taskInstance": { + "try_number": 1, + "max_tries": 0, + "map_index": 1, + "loop_index": 1, + "loop_max_iterations": 2, + "loop_group_id": "process_batch", + "id": "{{ is_uuid(result) }}", + "region_id": "{{ is_uuid(result) }}", + "region_index": 1 + }, + "task": { + "task_group": { + "group_id": "fetch", + "max_iterations": null, + "has_until": null, + "gate_task_id": null, + "terminal_task_id": null + } + } + } + } + }, + "job": { + "name": "openlineage_task_loop_dag.process_batch.fetch.download" + } } ] diff --git a/providers/openlineage/tests/system/openlineage/operator.py b/providers/openlineage/tests/system/openlineage/operator.py index b187590d3632d..104ee964d404b 100644 --- a/providers/openlineage/tests/system/openlineage/operator.py +++ b/providers/openlineage/tests/system/openlineage/operator.py @@ -414,7 +414,7 @@ class OpenLineageTestOperator(BaseOperator): :param check_run_ids: if True (default), additionally check the collected events with ``check_run_id_integrity``: START and COMPLETE/FAIL of one attempt share a ``runId``, while another try, map index or loop iteration gets its own. Set to False only for a Dag whose expected events - intentionally list a task's lifecycle partially. + intentionally list a task's lifecycle partially. It runs only once every key matches its template. Special sentinel values in expected dicts: @@ -489,7 +489,10 @@ def execute(self, context: Context) -> None: if self.fail_fast: raise errors[key] = str(e) - if self.check_run_ids: + if self.check_run_ids and errors: + # A key that failed its template is missing from events_by_key, so its runIds would look incomplete + log.info("Skipping the runId integrity check until the event templates match") + elif self.check_run_ids: if problems := check_run_id_integrity( events_by_key, repeated_start_keys={key for key in events_by_key if self._expects_repeated_events(key)}, diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index 69aadf932bed5..cff6981807b26 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -292,8 +292,27 @@ def test_executor_initializer_with_real_stats(self, mock_configure_orm): mock_configure_orm.assert_called_once() +def _build_fake_task_run_id(dag_id, task_id, try_number, logical_date, map_index): + return f"{logical_date.isoformat()}.{dag_id}.{task_id}.{try_number}.{map_index}" + + +@pytest.fixture +def task_run_id_builder(): + """Patch how a task's runId is built outside a loop, with a readable fake that keeps its inputs.""" + with mock.patch( + "airflow.providers.openlineage.utils.utils.build_task_instance_ol_run_id", + autospec=True, + side_effect=_build_fake_task_run_id, + ) as builder: + yield builder + + @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Airflow 2 tests") class TestOpenLineageListenerAirflow2: + @pytest.fixture(autouse=True) + def _use_task_run_id_builder(self, task_run_id_builder): + self.build_task_run_id = task_run_id_builder + @patch("airflow.models.TaskInstance.xcom_push") @patch("airflow.models.BaseOperator.render_template") def test_listener_does_not_change_task_instance(self, render_mock, xcom_push_mock): @@ -421,9 +440,6 @@ def _create_listener_and_task_instance(self) -> tuple[OpenLineageListener, TaskI def mock_dag_id(dag_id, logical_date, clear_number): return f"{logical_date.isoformat()}.{dag_id}.{clear_number}" - def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): - return f"{logical_date.isoformat()}.{dag_id}.{task_id}.{try_number}.{map_index}" - listener = OpenLineageListener() listener.extractor_manager = mock.Mock() @@ -433,7 +449,6 @@ def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): adapter = mock.Mock() adapter.build_dag_run_id.side_effect = mock_dag_id - adapter.build_task_instance_run_id.side_effect = mock_task_id adapter.start_task = mock.Mock() adapter.fail_task = mock.Mock() adapter.complete_task = mock.Mock() @@ -875,12 +890,12 @@ def test_on_task_instance_running_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is in the running state. This test ensures that when an Airflow task instance transitions to the running state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_running(None, task_instance, None) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -895,7 +910,7 @@ def test_on_task_instance_failed_correctly_calls_openlineage_adapter_run_id_meth """Tests the OpenLineageListener's response when a task instance is in the failed state. This test ensures that when an Airflow task instance transitions to the failed state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() @@ -904,7 +919,7 @@ def test_on_task_instance_failed_correctly_calls_openlineage_adapter_run_id_meth listener.on_task_instance_failed( previous_state=None, task_instance=task_instance, **on_task_failed_kwargs, session=None ) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -919,12 +934,12 @@ def test_on_task_instance_success_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is in the success state. This test ensures that when an Airflow task instance transitions to the success state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_success(None, task_instance, None) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -1011,7 +1026,7 @@ def test_listener_on_task_instance_running_do_not_call_adapter_when_disabled_ope listener.on_task_instance_running(None, task_instance, None) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.start_task.assert_not_called() @@ -1037,7 +1052,7 @@ def test_listener_on_task_instance_failed_do_not_call_adapter_when_disabled_oper ) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.fail_task.assert_not_called() @@ -1059,7 +1074,7 @@ def test_listener_on_task_instance_success_do_not_call_adapter_when_disabled_ope listener.on_task_instance_success(None, task_instance, None) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.complete_task.assert_not_called() @@ -1141,6 +1156,10 @@ def set_result(*args, **kwargs): @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Airflow 3 tests") class TestOpenLineageListenerAirflow3: + @pytest.fixture(autouse=True) + def _use_task_run_id_builder(self, task_run_id_builder): + self.build_task_run_id = task_run_id_builder + @pytest.mark.parametrize("state", ["success", "failed", "skipped"]) @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") def test_manual_regional_event_uses_same_execution_uuid(self, mocker, state): @@ -1150,9 +1169,7 @@ def test_manual_regional_event_uses_same_execution_uuid(self, mocker, state): listener, ti = self._create_listener_and_task_instance(runtime_ti=False) listener._executor = mocker.Mock(spec=ProcessPoolExecutor) listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id - listener.adapter.build_task_instance_run_id.side_effect = ( - OpenLineageAdapter.build_task_instance_run_id - ) + self.build_task_run_id.side_effect = OpenLineageAdapter.build_task_instance_run_id with DAG("loop_dag") as loop_dag: @task_group @@ -1173,7 +1190,7 @@ def body(): listener._on_task_instance_manual_state_change(ti, ti.dag_run, state) assert submit.call_args.kwargs["run_id"] == str(ti.id) - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() @pytest.mark.parametrize("state", ["running", "success", "failed", "skipped"]) @pytest.mark.parametrize("map_index", [-1, 2]) @@ -1195,9 +1212,7 @@ def body(): create_loop(body, max_iterations=2) ti.task = loop_dag.get_task("body.work") listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id - listener.adapter.build_task_instance_run_id.side_effect = ( - OpenLineageAdapter.build_task_instance_run_id - ) + self.build_task_run_id.side_effect = OpenLineageAdapter.build_task_instance_run_id ti.region_id = uuid.uuid4() ti.region_index = 2 ti.map_index = map_index @@ -1208,7 +1223,7 @@ def body(): method = {"running": "start_task", "failed": "fail_task"}.get(state, "complete_task") assert getattr(listener.adapter, method).call_args.kwargs["run_id"] == str(ti.id) assert lineage_run_id(ti) == str(ti.id) - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() @pytest.mark.parametrize("state", ["running", "success", "failed", "skipped"]) @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") @@ -1284,7 +1299,7 @@ def test_worker_event_run_id_of_mapped_task_has_its_map_index(self, state): getattr(listener, f"on_task_instance_{state}")(None, ti, **kwargs) - assert listener.adapter.build_task_instance_run_id.call_args.kwargs["map_index"] == 2 + assert self.build_task_run_id.call_args.kwargs["map_index"] == 2 @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") def test_manual_event_passes_its_stored_position_to_the_airflow_run_facet(self, mocker): @@ -1572,9 +1587,6 @@ def _create_listener_and_task_instance( def mock_dag_id(dag_id, logical_date, clear_number): return f"{logical_date.isoformat()}.{dag_id}.{clear_number}" - def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): - return f"{logical_date.isoformat()}.{dag_id}.{task_id}.{try_number}.{map_index}" - listener = OpenLineageListener() listener.extractor_manager = mock.Mock() @@ -1584,7 +1596,6 @@ def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): adapter = mock.Mock() adapter.build_dag_run_id.side_effect = mock_dag_id - adapter.build_task_instance_run_id.side_effect = mock_task_id adapter.start_task = mock.Mock() adapter.fail_task = mock.Mock() adapter.complete_task = mock.Mock() @@ -2479,12 +2490,12 @@ def test_on_task_instance_running_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is in the running state. This test ensures that when an Airflow task instance transitions to the running state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_running(None, task_instance) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -2557,7 +2568,7 @@ def test_on_task_instance_failed_correctly_calls_openlineage_adapter_run_id_meth """Tests the OpenLineageListener's response when a task instance is in the failed state. This test ensures that when an Airflow task instance transitions to the failed state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() @@ -2566,7 +2577,7 @@ def test_on_task_instance_failed_correctly_calls_openlineage_adapter_run_id_meth listener.on_task_instance_failed( previous_state=None, task_instance=task_instance, **on_task_failed_kwargs ) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -2581,12 +2592,12 @@ def test_on_task_instance_success_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is in the success state. This test ensures that when an Airflow task instance transitions to the success state, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_success(None, task_instance) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -2614,7 +2625,7 @@ def test_listener_on_task_instance_running_do_not_call_adapter_when_disabled_ope listener.on_task_instance_running(None, task_instance) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.start_task.assert_not_called() @@ -2640,7 +2651,7 @@ def test_listener_on_task_instance_failed_do_not_call_adapter_when_disabled_oper ) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.fail_task.assert_not_called() @@ -2662,7 +2673,7 @@ def test_listener_on_task_instance_success_do_not_call_adapter_when_disabled_ope listener.on_task_instance_success(None, task_instance) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.complete_task.assert_not_called() @@ -2673,12 +2684,12 @@ def test_on_task_instance_skipped_correctly_calls_openlineage_adapter_run_id_met """Tests the OpenLineageListener's response when a task instance is skipped. This test ensures that when an Airflow task instance is skipped via AirflowSkipException, - the OpenLineageAdapter's `build_task_instance_run_id` method is called exactly once with the correct + the task's runId is built exactly once from the correct parameters derived from the task instance. """ listener, task_instance = self._create_listener_and_task_instance() listener.on_task_instance_skipped(previous_state=None, task_instance=task_instance) - listener.adapter.build_task_instance_run_id.assert_called_once_with( + self.build_task_run_id.assert_called_once_with( dag_id="dag_id", task_id="task_id", logical_date=timezone.datetime(2020, 1, 1, 1, 1, 1), @@ -2704,7 +2715,7 @@ def test_listener_on_task_instance_skipped_do_not_call_adapter_when_disabled_ope listener.on_task_instance_skipped(previous_state=None, task_instance=task_instance) mock_disabled.assert_called_once_with(operator=task_instance.task, dag_id="dag_id", task_id="task_id") listener.adapter.build_dag_run_id.assert_not_called() - listener.adapter.build_task_instance_run_id.assert_not_called() + self.build_task_run_id.assert_not_called() listener.extractor_manager.extract_metadata.assert_not_called() listener.adapter.complete_task.assert_not_called() diff --git a/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py index 92ba3afedf7b4..fb4ccc0f30225 100644 --- a/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py +++ b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py @@ -18,6 +18,8 @@ from __future__ import annotations +from unittest import mock + import pytest from system.openlineage.operator import OpenLineageTestOperator, check_run_id_integrity @@ -248,3 +250,18 @@ def test_operator_derives_repeated_starts_from_event_count_assertions(event_coun ) assert operator._expects_repeated_events("dag.task.event.start") is expected + + +@mock.patch("system.openlineage.operator.time.sleep") +@mock.patch("system.openlineage.operator.check_run_id_integrity", autospec=True) +def test_operator_skips_the_run_id_check_when_a_template_fails(check_integrity, _sleep): + operator = OpenLineageTestOperator( + task_id="check_events", event_templates={"dag.task.event.start": {}}, clear_variables=False + ) + + with mock.patch.object(operator, "_check_key", side_effect=ValueError("template mismatch")): + with pytest.raises(ValueError, match="template mismatch") as err: + operator.execute(context={}) + + check_integrity.assert_not_called() + assert "runId integrity" not in str(err.value) diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index be2d76b56c08a..03f817745f999 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -83,10 +83,10 @@ get_task_documentation, get_task_instance_map_index, get_task_instance_position, + get_task_instance_run_id, get_task_parent_run_facet, get_user_provided_run_facets, is_dag_run_asset_triggered, - is_task_instance_in_loop, is_valid_uuid, translate_airflow_asset, ) @@ -3944,7 +3944,7 @@ def test_orm_lineage_facets_project_public_index(dag_maker, session, mapped): @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") @pytest.mark.parametrize("in_loop", [False, True]) -def test_task_instance_is_in_loop_only_inside_a_loop(in_loop): +def test_task_instance_run_id_is_the_attempt_uuid_only_inside_a_loop(in_loop): from airflow.sdk import task_group from airflow.sdk.definitions._internal.loop import create_loop @@ -3960,17 +3960,42 @@ def body(): body() ti = SimpleNamespace( id=uuid4(), + dag_id="run_id_dag", + task_id="body.work", + try_number=1, region_id=uuid4(), map_index=-1, task=dag.get_task("body.work"), get_template_context=MagicMock(), ) + logical_date = timezone.datetime(2024, 1, 1) + + run_id = get_task_instance_run_id(ti, logical_date=logical_date) - assert is_task_instance_in_loop(ti) is in_loop + assert run_id == ( + str(ti.id) + if in_loop + else build_task_instance_ol_run_id( + dag_id="run_id_dag", task_id="body.work", try_number=1, logical_date=logical_date, map_index=-1 + ) + ) assert get_task_instance_map_index(ti) == -1 ti.get_template_context.assert_not_called() +def test_task_instance_run_id_uses_the_map_index_of_the_given_position(): + ti = SimpleNamespace(id=uuid4(), dag_id="dag", task_id="task", try_number=2) + logical_date = timezone.datetime(2024, 1, 1) + + run_id = get_task_instance_run_id( + ti, logical_date=logical_date, position=TaskInstancePosition(in_loop=False, map_index=3) + ) + + assert run_id == build_task_instance_ol_run_id( + dag_id="dag", task_id="task", try_number=2, logical_date=logical_date, map_index=3 + ) + + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") @pytest.mark.parametrize("has_loop_context", [True, False]) def test_runtime_task_instance_position_has_the_iteration_of_its_loop_context(has_loop_context): @@ -4028,6 +4053,62 @@ def body(): ) +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +def test_stored_mapped_task_in_loop_is_placed_without_iteration_when_its_region_query_fails(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("stored_mapped_loop_dag") as dag: + + @task_group + def body(): + PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) + + create_loop(body, max_iterations=3) + ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 1 + session = MagicMock(spec=Session) + session.execute.side_effect = OSError("connection lost") + + positions = get_orm_task_instance_positions([ti], session=session) + + assert positions == { + ti.id: TaskInstancePosition(in_loop=True, map_index=1, loop_max_iterations=3, loop_group_id="body") + } + session.begin_nested.assert_called_once_with() + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +def test_stored_task_instance_without_task_still_raises_when_its_region_query_fails(): + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 1 + ti.task = None + session = MagicMock(spec=Session) + session.execute.side_effect = OSError("connection lost") + + with pytest.raises(OSError, match="connection lost"): + get_orm_task_instance_positions([ti], session=session) + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +def test_task_in_a_task_group_of_unknown_type_is_assumed_in_a_loop(): + task = SimpleNamespace(task_id="work", task_group=SimpleNamespace(group_id="unknown")) + ti = SimpleNamespace(id=uuid4(), region_id=uuid4(), map_index=-1, task=task) + + assert get_task_instance_position(ti) == TaskInstancePosition(in_loop=True, map_index=-1) + + +def test_task_instance_info_uses_a_given_position_as_is(): + ti = SimpleNamespace(id=uuid4(), map_index=-1, get_template_context=MagicMock()) + + info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=True, map_index=-1))) + + assert "loop_index" not in info + ti.get_template_context.assert_not_called() + + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") def test_stored_task_instance_without_task_or_region_has_no_loop_index(): ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) @@ -4088,10 +4169,10 @@ def body(): @pytest.mark.parametrize( ("until", "loop_index", "expected", "runs_out"), [ - pytest.param(None, 1, {"stop": False, "reason": None}, False, id="no_until_continues"), - pytest.param(None, 2, {"stop": True, "reason": None}, False, id="no_until_stops_at_last"), + pytest.param(None, 1, {"stop": False}, False, id="no_until_continues"), + pytest.param(None, 2, {"stop": True}, False, id="no_until_stops_at_last"), pytest.param("done", 0, {"stop": True, "reason": "done"}, False, id="until_stops_with_reason"), - pytest.param(False, 2, {"stop": False, "reason": None}, True, id="until_false_at_last"), + pytest.param(False, 2, {"stop": False}, True, id="until_false_at_last"), ], ) def test_task_instance_info_has_the_decision_of_a_loop_gate(until, loop_index, expected, runs_out): @@ -4189,9 +4270,7 @@ def body(): ti.region_id = uuid4() ti.region_index = 1 session = MagicMock(spec=Session) - session.execute.return_value = [ - SimpleNamespace(id=ti.region_id, node_id="body.work", parent_region_id=uuid4(), parent_region_index=2) - ] + session.execute.return_value = [_region_row(ti.region_id, "body.work", parent=(uuid4(), 2, "body"))] positions = get_orm_task_instance_positions([ti], session=session if with_session else None) @@ -4266,6 +4345,208 @@ def plain(): ) +def _region_row(region_id, node_id, *, parent=None, grandparent_region_id=None): + """Build a row of the region query: ``parent`` is ``(parent_region_id, parent_region_index, parent_node_id)``.""" + parent_region_id, parent_region_index, parent_node_id = parent or (None, None, None) + return SimpleNamespace( + id=region_id, + node_id=node_id, + parent_region_id=parent_region_id, + parent_region_index=parent_region_index, + parent_node_id=parent_node_id, + grandparent_region_id=grandparent_region_id, + ) + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") +@pytest.mark.parametrize( + ("own_region", "parent", "grandparent_region_id", "map_index"), + [ + pytest.param(True, (uuid4(), 2, "outer"), uuid4(), 1, id="mapped_task_two_levels_deep"), + pytest.param(False, (uuid4(), 2, "outer"), None, -1, id="enclosing_region_not_top_level"), + pytest.param(False, None, None, -1, id="region_missing"), + ], +) +def test_stored_task_instance_in_an_unexpected_region_shape_is_placed_without_iteration( + own_region, parent, grandparent_region_id, map_index +): + ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) + ti.task_id = "work" + ti.region_id = uuid4() + ti.region_index = 1 + ti.task = None + session = MagicMock(spec=Session) + rows = [] + if parent is not None or own_region: + rows = [ + _region_row( + ti.region_id, + "work" if own_region else "group", + parent=parent, + grandparent_region_id=grandparent_region_id, + ) + ] + session.execute.return_value = rows + + assert get_orm_task_instance_positions([ti], session=session) == { + ti.id: TaskInstancePosition(in_loop=True, map_index=map_index) + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +@pytest.mark.parametrize( + ("parent", "grandparent_region_id"), + [ + pytest.param((uuid4(), 2, "another_loop"), None, id="parent_is_not_its_loop"), + pytest.param((uuid4(), 2, "body"), uuid4(), id="loop_region_not_top_level"), + ], +) +def test_stored_mapped_task_in_loop_keeps_no_iteration_from_an_unexpected_region_shape( + parent, grandparent_region_id +): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("stored_mapped_loop_dag") as dag: + + @task_group + def body(): + PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) + + create_loop(body, max_iterations=3) + ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 1 + session = MagicMock(spec=Session) + session.execute.return_value = [ + _region_row(ti.region_id, "body.work", parent=parent, grandparent_region_id=grandparent_region_id) + ] + + assert get_orm_task_instance_positions([ti], session=session) == { + ti.id: TaskInstancePosition(in_loop=True, map_index=1, loop_max_iterations=3, loop_group_id="body") + } + + +@pytest.mark.db_test +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +@pytest.mark.parametrize("task_attached", [False, True]) +def test_region_shapes_of_airflow_are_the_ones_openlineage_places(dag_maker, session, task_attached): + """ + Pin the regions Airflow creates, which ``get_orm_task_instance_positions`` relies on to place task instances. + + If this fails, Airflow changed which regions it creates or how they nest: update the region parsing in + ``get_orm_task_instance_positions`` and its comment before changing these expectations. + """ + from airflow.sdk import task, task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with dag_maker(serialized=True) as dag: + + @task_group + def body(): + plain = EmptyOperator(task_id="plain") + mapped = PythonOperator.partial(task_id="mapped", python_callable=list).expand(op_kwargs=[{}, {}]) + + @task_group + def mapped_group(x): + @task + def member(x): + return x + + member(x) + + last = EmptyOperator(task_id="last") + plain >> mapped >> mapped_group.expand(x=[1, 2]) >> last + + create_loop(body, max_iterations=2) + + @task_group + def top_group(x): + @task + def top_member(x): + return x + + top_member(x) + + top_group.expand(x=[1, 2]) + run = dag_maker.create_dagrun() + tis = [ti for ti in run.task_instances if ti.task_id != "body.__loop_gate"] + for ti in tis: + ti.task = dag.get_task(ti.task_id) if task_attached else None + in_loop_extra = {"loop_max_iterations": 2} if task_attached else {} + + positions = get_orm_task_instance_positions(tis, session=session) + + assert {(ti.task_id, ti.region_index): positions[ti.id] for ti in tis} == { + ("body.plain", 0): TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=0, loop_group_id="body", **in_loop_extra + ), + ("body.last", 0): TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=0, loop_group_id="body", **in_loop_extra + ), + **{ + (task_id, map_index): TaskInstancePosition( + in_loop=True, map_index=map_index, loop_index=0, loop_group_id="body", **in_loop_extra + ) + for task_id in ("body.mapped", "body.mapped_group.member") + for map_index in (0, 1) + }, + **{ + ("top_group.top_member", map_index): TaskInstancePosition(in_loop=False, map_index=map_index) + for map_index in (0, 1) + }, + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") +@pytest.mark.parametrize("shape", ["nested_loop", "loop_in_mapped_group", "expand_in_mapped_group"]) +def test_region_shapes_stay_limited_by_the_task_sdk(shape): + """ + Pin the Dag shapes the task SDK rejects, which ``get_orm_task_instance_positions`` relies on. + + If one of them stops raising, a task can sit two region levels deep, which OpenLineage places as a loop + without its iteration: support the new shape in ``get_orm_task_instance_positions``, then update this test. + """ + from airflow.sdk import task, task_group + from airflow.sdk.definitions._internal.loop import create_loop + + def define_dag(): + with DAG(f"{shape}_dag"): + + @task_group + def inner(): + EmptyOperator(task_id="work") + + if shape == "nested_loop": + + @task_group + def outer(): + create_loop(inner, max_iterations=2) + + create_loop(outer, max_iterations=2) + elif shape == "loop_in_mapped_group": + + @task_group + def mapped_outer(x): + create_loop(inner, max_iterations=2) + + mapped_outer.expand(x=[1, 2]) + else: + + @task_group + def mapped_outer(x): + @task + def member(y): + return y + + member.expand(y=[x]) + + mapped_outer.expand(x=[1, 2]) + + with pytest.raises(NotImplementedError): + define_dag() + + @pytest.mark.db_test @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") @pytest.mark.parametrize("kind", ["plain_loop", "mapped", "mapped_in_loop"]) @@ -4295,7 +4576,7 @@ def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag ti.task = None with pytest.raises(ValueError, match="has no task attached"): - is_task_instance_in_loop(ti) + get_task_instance_run_id(ti, logical_date=timezone.datetime(2024, 1, 1)) with assert_queries_count(1): positions = get_orm_task_instance_positions([ti], session=session) @@ -4304,7 +4585,7 @@ def test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag in_loop=kind != "mapped", map_index=ti.region_index if mapped else -1, loop_index=2 if kind != "mapped" else None, - loop_group_id="loop" if kind == "plain_loop" else None, + loop_group_id="loop" if kind != "mapped" else None, ) } From ee7132c9f39af5fe94a6bc1dcf8d9f896ab5096b Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Mon, 19 Oct 2026 18:25:45 +0200 Subject: [PATCH 06/10] Stub lineage macros instead of OpenLineage internals in provider tests The dbt Cloud and Dataproc tests patched how OpenLineage builds a task's run id, so they broke when that changed. Stubbing the macros they actually use keeps them independent of it. Co-Authored-By: Claude --- .../unit/dbt/cloud/utils/test_openlineage.py | 4 +-- .../google/cloud/operators/test_dataproc.py | 36 +++++++++---------- .../providers/openlineage/utils/utils.py | 6 ++-- .../example_openlineage_task_loop_dag.py | 3 +- .../unit/openlineage/plugins/test_listener.py | 2 -- 5 files changed, 24 insertions(+), 27 deletions(-) diff --git a/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py b/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py index 2f38284efbafb..eae4bde3a2f71 100644 --- a/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py +++ b/providers/dbt/cloud/tests/unit/dbt/cloud/utils/test_openlineage.py @@ -162,8 +162,8 @@ def test_get_parent_run_metadata(): class TestGenerateOpenLineageEventsFromDbtCloudRun: @patch("importlib.metadata.version", return_value="3.0.0") @patch("airflow.providers.openlineage.api.emit") - @patch("airflow.providers.openlineage.plugins.adapter.OpenLineageAdapter.build_task_instance_run_id") - @patch("airflow.providers.openlineage.plugins.adapter.OpenLineageAdapter.build_dag_run_id") + @patch("airflow.providers.openlineage.plugins.macros.lineage_run_id") + @patch("airflow.providers.openlineage.plugins.macros.lineage_root_run_id") @patch.object(DbtCloudHook, "get_job_run") @patch.object(DbtCloudHook, "get_project") @patch.object(DbtCloudHook, "get_job_run_artifact") diff --git a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py index c8b61bfb12f8d..29facc7c96788 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py +++ b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py @@ -1898,8 +1898,8 @@ def test_dataproc_operator_execute_async_done_before_defer(self, mock_submit_job op.execute(context=self.mock_context) assert not mock_defer.called - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id") - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id") @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible") @mock.patch(DATAPROC_PATH.format("DataprocHook")) def test_execute_openlineage_parent_job_info_injection( @@ -1969,8 +1969,8 @@ def test_execute_openlineage_parent_job_info_injection( metadata=METADATA, ) - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id") - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id") @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible") @mock.patch(DATAPROC_PATH.format("DataprocHook")) @@ -2028,8 +2028,8 @@ def test_execute_openlineage_http_transport_info_injection( metadata=METADATA, ) - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id") - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id") @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible") @mock.patch(DATAPROC_PATH.format("DataprocHook")) @@ -3120,8 +3120,8 @@ def test_wait_for_operation_on_execute(self, mock_hook): ) mock_op.return_value.result.assert_not_called() - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id") - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id") @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible") @mock.patch(DATAPROC_PATH.format("DataprocHook")) def test_execute_openlineage_parent_job_info_injection( @@ -3315,8 +3315,8 @@ def test_execute_openlineage_parent_job_info_injection_skipped_when_ol_not_acces metadata=METADATA, ) - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id") - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id") @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible") @mock.patch(DATAPROC_PATH.format("DataprocHook")) @@ -3425,8 +3425,8 @@ def test_execute_openlineage_transport_info_injection( metadata=METADATA, ) - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id") - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id") @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible") @mock.patch(DATAPROC_PATH.format("DataprocHook")) @@ -4124,8 +4124,8 @@ def test_execute_batch_already_exists_cancelled(self, mock_hook, mock_log): mock_log.info.assert_any_call("Batch with given id already exists.") @mock.patch.object(DataprocCreateBatchOperator, "log", new_callable=mock.MagicMock) - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id") - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id") @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible") @mock.patch(DATAPROC_PATH.format("Batch.to_dict")) @mock.patch(DATAPROC_PATH.format("DataprocHook")) @@ -4182,8 +4182,8 @@ def test_execute_openlineage_parent_job_info_injection( mock_log.info.assert_any_call("Batch job %s completed.\nDriver logs: %s", BATCH_ID, logs_link) @mock.patch.object(DataprocCreateBatchOperator, "log", new_callable=mock.MagicMock) - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id") - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id") @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible") @mock.patch(DATAPROC_PATH.format("Batch.to_dict")) @@ -4250,8 +4250,8 @@ def test_execute_openlineage_transport_info_injection( logs_link, ) - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_dag_run_ol_run_id") - @mock.patch("airflow.providers.openlineage.plugins.adapter.build_task_instance_ol_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_root_run_id") + @mock.patch("airflow.providers.openlineage.utils.spark.lineage_run_id") @mock.patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") @mock.patch("airflow.providers.google.cloud.openlineage.utils._is_openlineage_provider_accessible") @mock.patch(DATAPROC_PATH.format("Batch.to_dict")) diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index 1ce42cf421aca..a8a1a76d33217 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1852,10 +1852,10 @@ def get_task_instance_run_id( position: TaskInstancePosition | None = None, ) -> str: """ - Return the OpenLineage runId of a task instance's attempt, which all its events must share. + Return the OpenLineage ``runId`` of a task instance's attempt, which all its events must share. - In a loop the iterations share every coordinate but the attempt UUID, so that is the runId. Otherwise it is - built from the coordinates. + In a loop the iterations share every coordinate but the attempt UUID, so that is the ``runId``. Otherwise + it is built from the coordinates. """ if position is None: position = _require_task_instance_position(task_instance) diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py index cbb4cd7916e51..e4cf123c12fc7 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py @@ -18,8 +18,7 @@ DAG with a plain task and task loops, some of their tasks retrying and one stopped early by `until` (requires Airflow 3.4+). -It checks, for a plain task, tasks in a loop, a mapped task and a mapped task group in a loop, and the loop -gates: +It checks, for a plain task, tasks, a mapped task and a mapped task group in a loop, and loop gates: - taskInstance's `try_number`, `map_index`, `loop_index` and `loop_max_iterations` of every event: the iteration is the zero-based `loop_index`, `map_index` the mapped position, and a plain task has no `loop_index` diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index cff6981807b26..32440f2402697 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -1169,7 +1169,6 @@ def test_manual_regional_event_uses_same_execution_uuid(self, mocker, state): listener, ti = self._create_listener_and_task_instance(runtime_ti=False) listener._executor = mocker.Mock(spec=ProcessPoolExecutor) listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id - self.build_task_run_id.side_effect = OpenLineageAdapter.build_task_instance_run_id with DAG("loop_dag") as loop_dag: @task_group @@ -1212,7 +1211,6 @@ def body(): create_loop(body, max_iterations=2) ti.task = loop_dag.get_task("body.work") listener.adapter.build_dag_run_id.side_effect = OpenLineageAdapter.build_dag_run_id - self.build_task_run_id.side_effect = OpenLineageAdapter.build_task_instance_run_id ti.region_id = uuid.uuid4() ti.region_index = 2 ti.map_index = map_index From 49e797a07357db9bf2a46906a1269935f12bfbc3 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 21 Oct 2026 13:36:38 +0200 Subject: [PATCH 07/10] Describe loop reruns and expansion size in OpenLineage task events Uses what apache/airflow#75132 exposes, so a single event tells which rerun of a loop it comes from, whether that rerun superseded an earlier one, and how many mapped instances to expect, and events for manual state changes get the same details as worker events. Loop details are grouped in one object so it can describe enclosing loops if Airflow ever nests them, and iterations are counted from one as in Airflow. Co-Authored-By: Claude --- .../docs/guides/task_execution_identity.rst | 227 +++++ providers/openlineage/docs/index.rst | 1 + .../openlineage/facets/AirflowRunFacet.json | 82 +- .../providers/openlineage/plugins/listener.py | 4 +- .../providers/openlineage/utils/utils.py | 219 ++-- .../example_openlineage_task_loop_dag.py | 23 +- .../openlineage_task_loop_dag.json | 948 +++++++++++++----- .../tests/system/openlineage/operator.py | 4 +- .../unit/openlineage/plugins/test_listener.py | 33 +- .../test_system_tests_run_id_integrity.py | 4 +- .../unit/openlineage/utils/test_utils.py | 198 +++- 11 files changed, 1342 insertions(+), 401 deletions(-) create mode 100644 providers/openlineage/docs/guides/task_execution_identity.rst diff --git a/providers/openlineage/docs/guides/task_execution_identity.rst b/providers/openlineage/docs/guides/task_execution_identity.rst new file mode 100644 index 0000000000000..8f2339b262288 --- /dev/null +++ b/providers/openlineage/docs/guides/task_execution_identity.rst @@ -0,0 +1,227 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _guides/task_execution_identity:openlineage: + +Identifying task executions in events +------------------------------------- + +Every OpenLineage event the provider emits for a task tells which execution of the task it describes: +which Dag run, which try, which mapped instance and, inside a :doc:`task loop +`, which iteration and which rerun of the loop. This page +explains how to read that from a single event, without access to the Airflow metadata database, for example to +group events per execution, deduplicate alerts or tell whether a failure is final. + +All the fields below are in the ``airflow`` run facet of task events (``run.facets.airflow``), unless said +otherwise. Dag run events carry the ``dag`` and ``dagRun`` parts in the ``airflowDagRun`` run facet. A field +that doesn't apply to an execution is left out of the event. + +Run ids +^^^^^^^ + +The ``runId`` of an event identifies one attempt, and all the events of that attempt (START, then COMPLETE or +FAIL) share it. + +* **Dag run:** a UUIDv7 derived from the OpenLineage namespace, the ``dag_id`` and ``dagRun.clear_number``, + with the Dag run's logical date (or ``run_after`` when it has none) as its timestamp. Clearing the whole Dag + run gives it a new ``runId``. +* **Task outside a loop:** a UUIDv7 derived from the namespace, ``dag_id``, ``task_id``, ``try_number`` and + ``map_index``, with the same timestamp. The same attempt always gets the same ``runId``, so it can be rebuilt + from those fields. +* **Task inside a loop:** ``taskInstance.id``, the UUID Airflow gives the attempt (Airflow 3.4+). The iterations + of a loop share every other coordinate, so a derived id would not be unique. + +The ``ParentRunFacet`` of a task event points at its Dag run's ``runId``. + +Fields +^^^^^^ + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Field + - Meaning + * - ``dag.dag_id``, ``task.task_id`` + - The Dag and the task. ``job.name`` is ``.`` for a task and ```` for a Dag run. + * - ``dagRun.run_id`` + - The Dag run. It stays the same when the Dag run is cleared. + * - ``dagRun.clear_number`` + - How many times the whole Dag run was cleared. Clearing single tasks doesn't change it. + * - ``taskInstance.id`` + - The UUID of the attempt (Airflow 3+). Every try, and every task instance a clear creates again, gets a + new one. + * - ``taskInstance.try_number`` + - The try. It counts per ``region_id`` and ``region_index``, so a task instance that a loop rerun creates + again starts from the first try. + * - ``taskInstance.max_tries`` + - The last ``try_number`` that is retried on failure. A FAIL with ``try_number > max_tries`` is final; + otherwise the task will be retried. Clearing a task instance raises it to give the task its retries again. + * - ``taskInstance.map_index`` + - The position in a mapped task's expansion, as shown in the UI. Absent for a task that isn't mapped. + * - ``taskInstance.expanded_ti_count`` + - The number of live instances in the mapped task's expansion (Airflow 3.4+), counted within the loop + iteration inside a loop. ``map_index`` is one of ``range(expanded_ti_count)``, so it tells how many events + to expect and is the denominator of "3 of 20 instances failed". Absent for a task that isn't mapped. + * - ``taskInstance.region_id``, ``taskInstance.region_index`` + - The coordinates Airflow stores the task instance under (Airflow 3.4+). A region is one execution of a + loop or of a mapped task's expansion. ``region_index`` is the iteration for a task in a loop that isn't + mapped, and the map index for a mapped task. Absent for a task that is neither mapped nor in a loop. + * - ``taskInstance.loop`` + - The iteration of the loop the task runs in (Airflow 3.4+), see below. Absent outside a loop. + * - ``task.task_group`` + - The task's task group. For a loop group it also has ``max_iterations``, ``has_until``, ``gate_task_id`` + and ``terminal_task_id``, also listed in the ``taskGroups`` of the Dag's ``AirflowJobFacet``. + +The ``taskInstance.loop`` object: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Field + - Meaning + * - ``group_id`` + - The loop's group id. It can differ from ``task.task_group`` when the task sits in a task group inside + the loop. + * - ``index`` + - The iteration, counted from 1. + * - ``max_iterations`` + - The loop's limit on iterations. + * - ``region_id`` + - The loop region that holds the iteration. All the task instances of one pass of the loop share it, + mapped ones included, whose own ``taskInstance.region_id`` is their expansion's. + * - ``trajectory`` + - 1 for the loop's first execution, and one more for each rerun of it, see :ref:`loop reruns + `. Of two executions of the same iteration, the higher + trajectory is the current one. + * - ``forked_from_region_id`` + - The ``loop.region_id`` of the previous trajectory. Absent in trajectory 1. + * - ``resumes_from_index`` + - The first iteration this trajectory holds. Absent in trajectory 1, which holds them from 1. + * - ``decision`` + - Only on the COMPLETE and FAIL events of a loop's gate: ``stop`` and the optional ``reason`` the gate + decided, see :ref:`loop gates `. + +Loops can't be nested, so ``loop`` always describes the only loop around the task. + +Grouping events +^^^^^^^^^^^^^^^ + +Each of these keys is built from a single event: + +.. list-table:: + :header-rows: 1 + :widths: 35 65 + + * - Group + - Key + * - A Dag run, across its clears + - ``(dag_id, dagRun.run_id)`` + * - One execution of a Dag run + - ``(dag_id, dagRun.run_id, dagRun.clear_number)`` + * - A task in a Dag run, across retries and clears + - ``(dag_id, dagRun.run_id, task_id, loop.group_id, loop.index, map_index)``, the loop and map parts left + empty when absent + * - One attempt + - ``runId``, or ``taskInstance.id`` on Airflow 3+ + * - One mapped expansion + - ``(dag_id, dagRun.run_id, task_id, taskInstance.region_id)``, with ``expanded_ti_count`` instances + * - One iteration of a loop, in one trajectory + - ``(dag_id, dagRun.run_id, loop.region_id, loop.index)`` + * - One iteration of a loop, across reruns + - ``(dag_id, dagRun.run_id, loop.group_id, loop.index)``, the current execution being the one with the + highest ``loop.trajectory`` + +How each kind of task emits events +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Plain task.** START when it starts, then COMPLETE or FAIL. A skipped task emits COMPLETE. + +**Retries.** Each try is an attempt with its own ``runId``, ``try_number`` and ``taskInstance.id``. A FAIL with +``try_number <= max_tries`` is followed by another try. + +**Sensors in reschedule mode and deferred tasks.** Rescheduling or deferring a task doesn't start a new try, so +its events keep one ``try_number`` and ``runId``. A rescheduled sensor emits a single START, while a deferred +task emits START again when it resumes. + +**Mapped task.** Each mapped instance has its own ``map_index`` and ``runId``, and every instance of one +expansion carries the same ``taskInstance.region_id`` and ``expanded_ti_count``. Tasks of a mapped task group +are mapped tasks in this sense. A task whose expansion turns out empty doesn't run, so no instance of it emits +events. + +**Task in a loop.** Each iteration emits its own events, with ``loop.index`` the iteration. For a task that isn't +mapped, ``taskInstance.region_id`` is the same as ``loop.region_id``. + +**Mapped task in a loop.** Each iteration expands the task again: ``map_index`` and ``expanded_ti_count`` are +within the iteration, ``taskInstance.region_id`` is the iteration's expansion and ``loop.region_id`` ties the +instances to the other tasks of the same pass. + +.. _guides/task_execution_identity:loop-gates: + +**Loop gate.** The gate ends each iteration and decides whether the loop goes on. Its task id is the loop +group's ``gate_task_id``. Read its ``loop.decision`` with ``loop.index``, ``loop.max_iterations`` and the loop +group's ``has_until``: + +.. list-table:: + :header-rows: 1 + :widths: 40 60 + + * - Event + - Meaning + * - COMPLETE, ``stop`` false + - The loop goes on to iteration ``index + 1``. + * - COMPLETE, ``stop`` true, ``has_until`` + - ``until`` ended the loop at iteration ``index``, which can be the last one. + * - COMPLETE, ``stop`` true, no ``has_until`` + - The loop ran all its iterations: without ``until`` it always stops at ``index == max_iterations``. + * - FAIL, ``stop`` false + - ``until`` was still false at the last iteration, so the loop reached ``max_iterations`` and fails. + * - FAIL without ``decision`` + - The gate failed before deciding, for example because ``until`` raised; the error facet tells why. + +.. _guides/task_execution_identity:loop-reruns: + +**Loop rerun.** Clearing a loop's gate together with its later iterations reruns the loop from there. Say a +loop ran iterations 1 to 4 and the gate of iteration 2 is cleared: + +* iteration 2's cleared tasks run again in place: same ``loop.region_id`` and ``loop.trajectory``, next + ``try_number``; +* iterations 3 and 4 are replaced by a new execution of the loop, if the gate decides to go on. Its events + have a new ``loop.region_id``, ``loop.trajectory`` one more, ``loop.forked_from_region_id`` the previous + trajectory's region and ``loop.resumes_from_index`` 3. Their tries count from the start again, so ``loop.index``, + ``map_index`` and ``try_number`` can repeat those of the replaced executions: ``loop.region_id`` and + ``loop.trajectory`` tell them apart. + +An event from a trajectory greater than 1 tells that the iterations from ``resumes_from_index`` on of every +earlier trajectory of that loop in the Dag run were superseded. They can span several earlier trajectories, +not only ``forked_from_region_id``, when an earlier rerun resumed later than this one. + +**Dag run cleared.** ``dagRun.clear_number`` grows and the Dag run gets a new ``runId``. The tasks it reruns +get a new try. + +**State set by hand.** Marking a task instance successful or failed in the UI or the API emits COMPLETE or FAIL +from the API server, with the same fields as the worker sends, ``loop`` and ``expanded_ti_count`` included +from Airflow 3.4.0 on. A task that never started has no START. A gate marked by hand has no ``loop.decision``. + +Limitations +^^^^^^^^^^^ + +* Loops need Airflow 3.4. Pre-releases of Airflow 3.4.0 give no ``loop`` details beyond ``group_id``, + ``index`` and ``max_iterations``, and no ``expanded_ti_count``. +* The ``airflowState`` run facet of Dag run events keys task states by ``task_id`` only, so mapped instances + and loop iterations of one task share one entry. diff --git a/providers/openlineage/docs/index.rst b/providers/openlineage/docs/index.rst index 5286019bee69d..cc0fee8e1b95e 100644 --- a/providers/openlineage/docs/index.rst +++ b/providers/openlineage/docs/index.rst @@ -39,6 +39,7 @@ Supported classes Custom Operators Job Hierarchy & Macros + Identifying task executions Spark Integration Troubleshooting diff --git a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json index bfba75b5415fe..ba86c9b856b7c 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json +++ b/providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json @@ -218,34 +218,15 @@ "description": "Position in a mapped task's expansion. Absent for a task that isn't mapped.", "type": "integer" }, - "loop_index": { - "description": "Zero-based iteration of the innermost loop the task runs in (Airflow 3.4+). Absent outside a loop.", + "expanded_ti_count": { + "description": "Number of live instances in the mapped task's expansion (Airflow 3.4+), counted within the loop iteration inside a loop. Absent for a task that isn't mapped.", "type": "integer" }, - "loop_max_iterations": { - "description": "Maximum number of iterations of the innermost loop the task runs in (Airflow 3.4+). Absent outside a loop.", - "type": "integer" - }, - "loop_group_id": { - "description": "Group id of the innermost loop the task runs in (Airflow 3.4+), which loop_index counts the iterations of. Absent outside a loop.", - "type": "string" - }, - "loop_decision": { - "description": "Decision of a loop gate (Airflow 3.4+), on its COMPLETE or FAIL event only. A FAIL with stop false at the last iteration means the loop's until was still false at max_iterations.", - "type": "object", - "properties": { - "stop": { - "description": "Whether the gate stopped the loop.", - "type": "boolean" - }, - "reason": { - "description": "Reason given by until, if any.", - "type": "string" - } - } + "loop": { + "$ref": "#/$defs/TaskInstanceLoop" }, "region_id": { - "description": "Id of the region (one execution of a loop or of a mapped task's expansion) the task instance belongs to (Airflow 3.4+). Shared by the task instances of that region. Clearing a loop forks its region, so a re-run of an iteration gets a new region_id. Region ids are UUIDv7, so a later trajectory has a greater one, which is an implementation detail. Absent for a task that is neither mapped nor in a loop.", + "description": "Id of the region (one execution of a loop or of a mapped task's expansion) the task instance belongs to (Airflow 3.4+). Shared by the task instances of that region. For a mapped task it is its expansion's, see loop.region_id for the loop's. Absent for a task that is neither mapped nor in a loop.", "type": "string", "format": "uuid" }, @@ -279,6 +260,59 @@ "try_number" ] }, + "TaskInstanceLoop": { + "description": "Iteration of the innermost loop the task runs in (Airflow 3.4+). Absent outside a loop. Fields that can't be known are left out.", + "type": "object", + "properties": { + "group_id": { + "description": "Group id of the loop.", + "type": "string" + }, + "index": { + "description": "Iteration, counted from 1.", + "type": "integer" + }, + "max_iterations": { + "description": "Maximum number of iterations of the loop.", + "type": "integer" + }, + "region_id": { + "description": "Id of the loop region holding the iteration, shared by every task instance of that pass of the loop, mapped ones included.", + "type": "string", + "format": "uuid" + }, + "trajectory": { + "description": "1 for the loop's original region, and one more for each fork of it. Clearing a loop gate with the later iterations forks the loop region, and the higher trajectory is the current one.", + "type": "integer", + "minimum": 1 + }, + "forked_from_region_id": { + "description": "Id of the previous loop region of the chain this fork extends. Absent for the original region. The fork replaces the iterations from resumes_from_index on, which earlier regions of the chain can hold too.", + "type": "string", + "format": "uuid" + }, + "resumes_from_index": { + "description": "First iteration the fork holds. Absent for the original region, which holds them from 1.", + "type": "integer", + "minimum": 0 + }, + "decision": { + "description": "Decision of a loop gate, on its COMPLETE or FAIL event only. A FAIL with stop false at the last iteration means the loop's until was still false at max_iterations.", + "type": "object", + "properties": { + "stop": { + "description": "Whether the gate stopped the loop.", + "type": "boolean" + }, + "reason": { + "description": "Reason given by until, if any.", + "type": "string" + } + } + } + }, + "additionalProperties": true + }, "DagRun": { "type": "object", "properties": { diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py index dda70e50fecd8..7a2a2e0cc7033 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py @@ -52,7 +52,7 @@ get_dag_parent_run_facet, get_dag_run_dag_and_task_from_ti, get_job_name, - get_orm_task_instance_positions, + get_stored_task_instance_position, get_task_documentation, get_task_instance_position, get_task_instance_run_id, @@ -889,7 +889,7 @@ def _on_task_instance_manual_state_change( # pickler loses TaskGroup attributes and crashes event emission -- see # the equivalent note in `on_dag_run_running` (listener.py ~868). date = dagrun.logical_date or dagrun.run_after - position = get_orm_task_instance_positions([ti], session=object_session(ti))[ti.id] + position = get_stored_task_instance_position(ti, session=object_session(ti)) task_uuid = get_task_instance_run_id(ti, logical_date=date, position=position) parent_run_id = self.adapter.build_dag_run_id( dag_id=ti.dag_id, diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index a8a1a76d33217..ea470e7c5896d 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1137,50 +1137,52 @@ class TaskInstanceInfo(InfoJsonEncodable): clearing sets it to the last try number plus ``retries``, giving the task its retries again. * ``map_index``: the position in a mapped task's expansion that users see. Absent for a task that isn't mapped. - * ``loop_index``, ``loop_max_iterations`` and ``loop_group_id`` (Airflow 3.4+): the zero-based iteration - of the innermost loop around the task, that loop's limit and its group id, which can differ from - ``task.task_group`` when the task sits in a task group inside the loop. The iterations of loops further - out aren't known. ``(loop_group_id, loop_index)`` names an iteration, but not which trajectory it ran in, - nor in a nested loop which iteration of the outer loop. + * ``expanded_ti_count`` (Airflow 3.4+): the number of live instances in the mapped task's expansion, counted + within the loop iteration inside a loop, so ``map_index`` is one of ``range(expanded_ti_count)``. Absent + for a task that isn't mapped. * ``region_id`` and ``region_index`` (Airflow 3.4+): the coordinates Airflow stores the task instance under. A region is one execution of a loop, or of a mapped task's expansion, and is shared by all the task instances it creates. ``region_index`` is the position inside it: the iteration for a task in a loop that isn't mapped, the map index for a mapped task. Neither is unique: ``(dag_id, dagRun.run_id, task_id, region_id, region_index, try_number)`` is, as Airflow enforces it (``task_instance_try_key``). Both are absent for a task in no region, i.e. one that is neither mapped nor in a loop. - - ``loop_decision`` (Airflow 3.4+): ``stop`` and the optional ``reason`` returned by ``until``, as evaluated by - a loop's gate, the task that ends each iteration (``gate_task_id`` of the loop group, in ``task.task_group`` - and in ``taskGroups`` of ``AirflowJobFacet``). Only the gate's COMPLETE and FAIL events carry it, as the gate - decides while it runs. Read with ``loop_index``, ``loop_max_iterations`` and the group's ``has_until``: - - * COMPLETE with ``stop`` false: the loop goes on to iteration ``loop_index + 1``. - * COMPLETE with ``stop`` true and ``has_until``: ``until`` ended the loop at iteration ``loop_index``, which - can be the last one. + * ``loop`` (Airflow 3.4+): the iteration of the innermost loop around the task, absent outside a loop. + Loops can't be nested. Its fields, each left out when unknown: + + * ``group_id``: the loop's group id, which can differ from ``task.task_group`` when the task sits in a + task group inside the loop. + * ``index`` and ``max_iterations``: the iteration, counted from 1, and the loop's limit. + * ``region_id``: the loop region holding the iteration, shared by every task instance of that pass of + the loop, mapped ones included, whose own ``region_id`` is their expansion's. + * ``trajectory``, ``forked_from_region_id`` and ``resumes_from_index``: clearing a loop gate with the + later iterations forks the loop region, see below. ``trajectory`` is 1 for the original region and one + more for each fork, ``forked_from_region_id`` the previous region of the chain and + ``resumes_from_index`` the first iteration the fork holds. Both are absent for the original region. + * ``decision``: ``stop`` and the optional ``reason`` returned by ``until``, as evaluated by a loop's gate, + the task that ends each iteration (``gate_task_id`` of the loop group, in ``task.task_group`` and in + ``taskGroups`` of ``AirflowJobFacet``). Only the gate's COMPLETE and FAIL events carry it, as the gate + decides while it runs. + + How a gate's ``loop.decision`` reads, with ``loop.index``, ``loop.max_iterations`` and the group's + ``has_until``: + + * COMPLETE with ``stop`` false: the loop goes on to iteration ``index + 1``. + * COMPLETE with ``stop`` true and ``has_until``: ``until`` ended the loop at iteration ``index``, which can be + the last one. * COMPLETE with ``stop`` true and no ``has_until``: the loop ran all its iterations, as one without ``until`` - always stops at ``loop_index == loop_max_iterations - 1``. - * FAIL with ``stop`` false: ``until`` was still false at the last iteration (``loop_index == - loop_max_iterations - 1``), so the loop ran out of iterations and fails. That's the only failure after the - decision is made. - * FAIL without ``loop_decision``: the gate failed before deciding, e.g. ``until`` raised or returned something + always stops at ``index == max_iterations``. + * FAIL with ``stop`` false: ``until`` was still false at the last iteration (``index == max_iterations``), + so the loop ran out of iterations and fails. That's the only failure after the decision is made. + * FAIL without ``decision``: the gate failed before deciding, e.g. ``until`` raised or returned something other than a bool or a ``LoopDecision``; the error facet tells which. How a cleared loop reads: clearing a loop from iteration ``k`` retries what was cleared in iteration ``k`` in place, with the same region and the next ``try_number``. The iterations after ``k`` are created again in a - new region forked from the loop's region, and try numbers count per region, so the new execution of - iteration ``k + 1`` can repeat the task, ``loop_index``, ``map_index`` and ``try_number`` of the one it - replaces. Only ``region_id`` (and ``id``) tell the two apart. Each such re-run of a loop is a trajectory. + new region forked from the loop's newest region, with ``loop.trajectory`` one more and ``loop.resumes_from_index`` + ``k + 1``, and try numbers count per region, so the new execution of iteration ``k + 1`` can repeat the task, + ``loop.index``, ``map_index`` and ``try_number`` of the one it replaces. ``loop.region_id`` and + ``loop.trajectory`` tell the two apart, and the higher trajectory is the current one. ``dagRun.clear_number`` counts clears of the whole Dag run only. - - Ordering: region ids are UUIDv7, so of two executions of one task with the same ``loop_index`` and - ``map_index``, the later trajectory has the greater ``region_id``. That is how Airflow generates them, not a - documented guarantee, and it only orders the trajectories seen: an event doesn't say whether it comes from - the second or the third trajectory, nor from which iteration its trajectory was forked. - - What events can't tell: a mapped task's ``region_id`` is the one of its own expansion, which sits inside the - loop's region without naming it. So a mapped task instance in a loop can't be matched to the trajectory of - the loop, or to the task instances of the same pass that aren't mapped (whose ``region_id`` is the loop's), - other than by ``loop_index`` and time. """ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "try_number"] @@ -1211,25 +1213,23 @@ def _extend_fields(self) -> None: self.map_index = map_index if map_index != -1 else None self.rendered_map_index = getattr(self.obj, "rendered_map_index", None) if map_index != -1 else None self._fields.extend(("map_index", "rendered_map_index")) - if position and position.loop_index is not None: - self.loop_index = position.loop_index - self._fields.append("loop_index") - if position and position.loop_max_iterations is not None: - self.loop_max_iterations = position.loop_max_iterations - self._fields.append("loop_max_iterations") - if position and position.loop_group_id is not None: - self.loop_group_id = position.loop_group_id - self._fields.append("loop_group_id") + if position and position.expanded_ti_count is not None: + self.expanded_ti_count = position.expanded_ti_count + self._fields.append("expanded_ti_count") if AIRFLOW_V_3_0_PLUS and (ti_id := getattr(self.obj, "id", None)) is not None: self.id = str(ti_id) self._fields.append("id") - # Only a loop gate's runtime task instance has one, once it has evaluated ``until``. Read from a private - # attribute until the task SDK exposes it. - if (decision := getattr(self.obj, "_loop_decision", None)) is not None: - self.loop_decision = {"stop": decision.stop} + loop = _build_loop_info(position) if position and position.in_loop else {} + # Only a loop gate's runtime task instance has one, once it has evaluated ``until``. Airflow 3.4.0 exposes + # it, a pre-release only as a private attribute. + decision = getattr(self.obj, "loop_decision", None) or getattr(self.obj, "_loop_decision", None) + if decision is not None: + loop["decision"] = {"stop": decision.stop} if decision.reason is not None: - self.loop_decision["reason"] = decision.reason - self._fields.append("loop_decision") + loop["decision"]["reason"] = decision.reason + if loop: + self.loop = loop + self._fields.append("loop") # A runtime task instance in no region carries the sentinel region id instead of None region_id = getattr(self.obj, "region_id", None) if isinstance(region_id, UUID) and region_id.int != 0: @@ -1238,6 +1238,21 @@ def _extend_fields(self) -> None: self._fields.extend(("region_id", "region_index")) +def _build_loop_info(position: TaskInstancePosition) -> dict[str, Any]: + fields = { + "group_id": position.loop_group_id, + "index": position.loop_index, + "max_iterations": position.loop_max_iterations, + "region_id": str(position.loop_region_id) if position.loop_region_id else None, + "trajectory": position.loop_trajectory, + "forked_from_region_id": ( + str(position.loop_forked_from_region_id) if position.loop_forked_from_region_id else None + ), + "resumes_from_index": position.loop_resumes_from_index, + } + return {name: value for name, value in fields.items() if value is not None} + + class AssetInfo(InfoJsonEncodable): """ Defines encoding Airflow Asset object to JSON. @@ -1594,11 +1609,12 @@ class TaskInstancePosition: """ Where a task instance runs: in a loop or not, the map index users see, and its loop iteration. - The loop fields are ``None`` outside a loop, and also when they can't be known: a stored task whose task - isn't attached has no ``loop_max_iterations``, nor ``loop_group_id`` when it is mapped; a stored mapped task - has no ``loop_index`` unless ``get_orm_task_instance_positions`` could read its region; and a runtime task - instance has no ``loop_index`` nor ``loop_max_iterations`` unless placed with ``get_task_instance_position``, - as only its template context holds them. ``TaskInstanceInfo`` explains what these fields mean in events. + The loop fields are ``None`` outside a loop, and also when they can't be known. Airflow 3.4+ tells all of + them, and ``expanded_ti_count`` for a mapped task: a runtime task instance through its ``loop`` and + ``expanded_ti_count``, read by ``get_task_instance_position``, and a stored one through Airflow's + ``TaskCoordinateResolver``, called by ``get_stored_task_instance_position``. Before that, a + runtime task instance only has the iteration and limit from its template context, and a stored one only what + its task and region rows tell. ``TaskInstanceInfo`` explains what these fields mean in events. """ in_loop: bool @@ -1606,6 +1622,11 @@ class TaskInstancePosition: loop_index: int | None = None loop_max_iterations: int | None = None loop_group_id: str | None = None + loop_region_id: UUID | None = None + loop_trajectory: int | None = None + loop_forked_from_region_id: UUID | None = None + loop_resumes_from_index: int | None = None + expanded_ti_count: int | None = None def _find_enclosing_loop(task: Any) -> tuple[bool, Any | None]: @@ -1680,20 +1701,33 @@ def _require_task_instance_position( def _add_runtime_loop_iteration( task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition ) -> TaskInstancePosition: - # Run ids, map indexes and the lineage macros don't need the iteration, so only this reads the template - # context: it re-validates the params on every call, which can raise from inside a user's template. - if not position.in_loop or position.loop_index is not None or isinstance(task_instance, TaskInstance): + if isinstance(task_instance, TaskInstance): return position - try: - loop = task_instance.get_template_context().get("loop") - except Exception: - # The iteration only describes the run, so failing to read it must not cost the event its run id - log.warning("OpenLineage failed to read the loop iteration of task instance %s", task_instance.id) - log.debug("Exception details:", exc_info=True) + # Only runtime task instances of Airflow 3.4+ expose these, the expansion size from 3.4.0 on + position = attrs.evolve(position, expanded_ti_count=getattr(task_instance, "expanded_ti_count", None)) + if not position.in_loop or position.loop_index is not None: return position + if (loop := getattr(task_instance, "loop", None)) is None: + # Run ids, map indexes and the lineage macros don't need the iteration, so only this reads the template + # context: it re-validates the params on every call, which can raise from inside a user's template. + try: + loop = task_instance.get_template_context().get("loop") + except Exception: + # The iteration only describes the run, so failing to read it must not cost the event its run id + log.warning("OpenLineage failed to read the loop iteration of task instance %s", task_instance.id) + log.debug("Exception details:", exc_info=True) + return position if loop is None: return position - return attrs.evolve(position, loop_index=loop.index, loop_max_iterations=loop.max_iterations) + return attrs.evolve( + position, + loop_index=loop.index, + loop_max_iterations=loop.max_iterations, + loop_region_id=getattr(loop, "region_id", None), + loop_trajectory=getattr(loop, "trajectory", None), + loop_forked_from_region_id=getattr(loop, "forked_from_region_id", None), + loop_resumes_from_index=getattr(loop, "resumes_from_index", None), + ) def get_task_instance_position(task_instance: TaskInstance | RuntimeTaskInstance) -> TaskInstancePosition: @@ -1701,6 +1735,67 @@ def get_task_instance_position(task_instance: TaskInstance | RuntimeTaskInstance return _add_runtime_loop_iteration(task_instance, _require_task_instance_position(task_instance)) +def _get_core_task_instance_position( + task_instance: TaskInstance, *, session: Session | None +) -> TaskInstancePosition | None: + """Place a stored task instance in a region as Airflow places it for the worker, from Airflow 3.4.0 on.""" + region_id = task_instance.region_id + if session is None or not AIRFLOW_V_3_4_PLUS or not isinstance(region_id, UUID) or region_id.int == 0: + return None + # Reusing the Dag the task instance holds spares deserializing it again, unless it pins another version + if (dag := getattr(getattr(task_instance, "task", None), "dag", None)) is None: + return None + from airflow.models.task_coordinates import TaskCoordinateResolver + + if not hasattr(TaskCoordinateResolver, "get_loop_context"): # An Airflow 3.4 pre-release + return None + try: + # The savepoint keeps the caller's transaction usable after a failure + with session.begin_nested(): + resolver = TaskCoordinateResolver.for_dag(dag, session) + task = resolver.get_task( + task_instance.dag_id, + task_instance.run_id, + task_instance.task_id, + dag_version_id=task_instance.dag_version_id, + ) + loop = resolver.get_loop_context(task_instance, task=task) + expanded_ti_count = resolver.get_expanded_ti_count(task_instance, task=task) + except Exception: + log.warning("OpenLineage failed to place task instance %s with Airflow", task_instance.id) + log.debug("Exception details:", exc_info=True) + return None + map_index = task_instance.region_index if task.get_needs_expansion() else -1 + if loop is None: + return TaskInstancePosition(in_loop=False, map_index=map_index, expanded_ti_count=expanded_ti_count) + return TaskInstancePosition( + in_loop=True, + map_index=map_index, + expanded_ti_count=expanded_ti_count, + loop_index=loop.index, + loop_max_iterations=loop.max_iterations, + loop_group_id=loop.node_id, + loop_region_id=loop.region_id, + loop_trajectory=loop.trajectory, + loop_forked_from_region_id=loop.forked_from_region_id, + loop_resumes_from_index=loop.resumes_from_index, + ) + + +def get_stored_task_instance_position( + task_instance: TaskInstance, *, session: Session | None +) -> TaskInstancePosition: + """ + Place a stored task instance for an event about it, as Airflow places it for the worker when it can. + + That needs Airflow 3.4.0+ and a session. Otherwise, or when Airflow fails to, it is placed with + ``get_orm_task_instance_positions``, which only knows what the task and region rows tell. + """ + if (position := _get_core_task_instance_position(task_instance, session=session)) is not None: + return position + return get_orm_task_instance_positions([task_instance], session=session)[task_instance.id] + + def get_orm_task_instance_positions( task_instances: Collection[TaskInstance], *, session: Session | None ) -> dict[UUID, TaskInstancePosition]: diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py index e4cf123c12fc7..dac1d9e2e0390 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py @@ -19,20 +19,23 @@ (requires Airflow 3.4+). It checks, for a plain task, tasks, a mapped task and a mapped task group in a loop, and loop gates: - - taskInstance's `try_number`, `map_index`, `loop_index` and `loop_max_iterations` of every event: the - iteration is the zero-based `loop_index`, `map_index` the mapped position, and a plain task has no - `loop_index` - - taskInstance's `loop_group_id` of every event: the loop group of a task in a loop, absent for a plain task + - taskInstance's `try_number`, `map_index` and `loop` of every event: `loop.index` is the iteration, + counted from 1, `loop.max_iterations` the loop's limit and `loop.group_id` its group, `map_index` the + mapped position, and a plain task has no `loop` + - `loop.region_id`, the loop region of the iteration, shared by the tasks of one iteration, mapped ones + included, `loop.trajectory` 1, and no `loop.forked_from_region_id` nor `loop.resumes_from_index`, as + no loop is cleared + - taskInstance's `expanded_ti_count` of the mapped tasks, the size of their expansion in the iteration - taskInstance's `max_tries` of the tasks other than the gates, the `retries` of each task - taskInstance's `id` of every event, and `region_id` and `region_index` of the tasks in a loop: the `region_index` is the iteration for a task that isn't mapped and the map index for a mapped one, and a plain task has neither - - the gates' `loop_decision`: absent on START; for the loops without `until`, `stop` true only at the last + - the gates' `loop.decision`: absent on START; for the loops without `until`, `stop` true only at the last iteration and no `reason`; for the loop with `until` (whose gate is named after it), the `reason` it returns and `stop` true at the iteration it ends the loop, before `max_iterations`; absent for the plain task - a mapped task group inside a loop: its task's `map_index` and `region_index` are the mapped position, - its `loop_index` the iteration, and its `loop_group_id` the loop's, unlike its own `task_group` + its `loop.index` the iteration, and its `loop.group_id` the loop's, unlike its own `task_group` - the loop groups' `max_iterations`, `has_until`, `gate_task_id` and `terminal_task_id`, in the task's `task_group` and in the `taskGroups` of the Dag's job facet - a retried task emits START and FAIL for each failed try, then START and COMPLETE for the last one @@ -58,7 +61,7 @@ def sort_by_execution(event: dict) -> tuple[int, int, int]: # Dag run events have no airflow facet task_instance = event["run"]["facets"].get("airflow", {}).get("taskInstance", {}) return ( - task_instance.get("loop_index", -1), + task_instance.get("loop", {}).get("index", -1), task_instance.get("map_index") if task_instance.get("map_index") is not None else -1, task_instance.get("try_number", -1), ) @@ -82,9 +85,9 @@ def accumulate(): @task(retries=2, retry_delay=timedelta(seconds=1)) def increment(*, loop, ti): # The second iteration fails on its first try - if loop.index == 1 and ti.try_number < 2: + if loop.index == 2 and ti.try_number < 2: raise ValueError(f"Intentional failure, iteration {loop.index}, try {ti.try_number}") - return (0 if loop.index == 0 else loop.previous) + 1 + return (0 if loop.index == 1 else loop.previous) + 1 @task def report(value): @@ -130,7 +133,7 @@ def check(*, loop): def ready(*, loop): # Stops at the second of three iterations, giving a reason both when continuing and when stopping - return LoopDecision(stop=loop.result >= 1, reason=f"check returned {loop.result}") + return LoopDecision(stop=loop.result >= 2, reason=f"check returned {loop.result}") accumulate_loop = accumulate.loop(max_iterations=3) process_batch_loop = process_batch.loop(max_iterations=2) diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json index 6fc30d3d73258..c72e563ec84c3 100644 --- a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_task_loop_dag.json @@ -44,14 +44,12 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": null, - "loop_max_iterations": null, "id": "{{ is_uuid(result) }}", "region_id": null, "region_index": null, "max_tries": 0, - "loop_group_id": null, - "loop_decision": null + "loop": null, + "expanded_ti_count": null } } } @@ -68,14 +66,12 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": null, - "loop_max_iterations": null, "id": "{{ is_uuid(result) }}", "region_id": null, "region_index": null, "max_tries": 0, - "loop_group_id": null, - "loop_decision": null + "loop": null, + "expanded_ti_count": null } } } @@ -92,13 +88,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, + "region_index": 1, "max_tries": 2, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -124,13 +127,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, + "region_index": 2, "max_tries": 2, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -156,13 +166,20 @@ "airflow": { "taskInstance": { "try_number": 2, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, + "region_index": 2, "max_tries": 2, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -188,13 +205,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 2, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 2, + "region_index": 3, "max_tries": 2, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 3, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -220,13 +244,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, + "region_index": 1, "max_tries": 2, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -252,13 +283,20 @@ "airflow": { "taskInstance": { "try_number": 2, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, + "region_index": 2, "max_tries": 2, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -284,13 +322,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 2, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 2, + "region_index": 3, "max_tries": 2, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 3, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -316,13 +361,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, + "region_index": 2, "max_tries": 2, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -348,13 +400,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, + "region_index": 1, "max_tries": 0, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -380,13 +439,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, + "region_index": 2, "max_tries": 0, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -412,13 +478,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 2, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 2, + "region_index": 3, "max_tries": 0, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 3, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -444,13 +517,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, + "region_index": 1, "max_tries": 0, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -476,13 +556,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, + "region_index": 2, "max_tries": 0, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -508,13 +595,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 2, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 2, + "region_index": 3, "max_tries": 0, - "loop_group_id": "accumulate" + "loop": { + "group_id": "accumulate", + "index": 3, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -540,13 +634,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, - "loop_group_id": "accumulate", - "loop_decision": null + "region_index": 1, + "loop": { + "group_id": "accumulate", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -572,13 +673,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, - "loop_group_id": "accumulate", - "loop_decision": null + "region_index": 2, + "loop": { + "group_id": "accumulate", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -604,13 +712,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 2, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 2, - "loop_group_id": "accumulate", - "loop_decision": null + "region_index": 3, + "loop": { + "group_id": "accumulate", + "index": 3, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -636,15 +751,22 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, - "loop_group_id": "accumulate", - "loop_decision": { - "stop": false - } + "region_index": 1, + "loop": { + "group_id": "accumulate", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": { + "stop": false + } + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -670,15 +792,22 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, - "loop_group_id": "accumulate", - "loop_decision": { - "stop": false - } + "region_index": 2, + "loop": { + "group_id": "accumulate", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": { + "stop": false + } + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -704,15 +833,22 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 2, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 2, - "loop_group_id": "accumulate", - "loop_decision": { - "stop": true - } + "region_index": 3, + "loop": { + "group_id": "accumulate", + "index": 3, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": { + "stop": true + } + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -738,13 +874,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 2, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, + "region_index": 1, "max_tries": 0, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -770,13 +913,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 2, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, + "region_index": 2, "max_tries": 0, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -802,13 +952,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 2, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, + "region_index": 1, "max_tries": 0, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -834,13 +991,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 2, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, + "region_index": 2, "max_tries": 0, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -866,14 +1030,21 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 2, "map_index": 0, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 0, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -899,14 +1070,21 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -932,14 +1110,21 @@ "airflow": { "taskInstance": { "try_number": 2, - "loop_index": 0, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -965,14 +1150,21 @@ "airflow": { "taskInstance": { "try_number": 3, - "loop_index": 0, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -998,14 +1190,21 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 2, "map_index": 0, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 0, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1031,14 +1230,21 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1064,14 +1270,21 @@ "airflow": { "taskInstance": { "try_number": 2, - "loop_index": 1, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1097,14 +1310,21 @@ "airflow": { "taskInstance": { "try_number": 3, - "loop_index": 1, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1130,14 +1350,21 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 2, "map_index": 0, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 0, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1163,14 +1390,21 @@ "airflow": { "taskInstance": { "try_number": 3, - "loop_index": 0, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1196,14 +1430,21 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 2, "map_index": 0, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 0, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1229,14 +1470,21 @@ "airflow": { "taskInstance": { "try_number": 3, - "loop_index": 1, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1262,14 +1510,21 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1295,14 +1550,21 @@ "airflow": { "taskInstance": { "try_number": 2, - "loop_index": 0, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1328,14 +1590,21 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1361,14 +1630,21 @@ "airflow": { "taskInstance": { "try_number": 2, - "loop_index": 1, - "loop_max_iterations": 2, "map_index": 1, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", "region_index": 1, "max_tries": 3, - "loop_group_id": "process_batch" + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1394,13 +1670,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 2, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, - "loop_group_id": "process_batch", - "loop_decision": null + "region_index": 1, + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1426,13 +1709,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 2, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, - "loop_group_id": "process_batch", - "loop_decision": null + "region_index": 2, + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1458,15 +1748,22 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 2, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, - "loop_group_id": "process_batch", - "loop_decision": { - "stop": false - } + "region_index": 1, + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": { + "stop": false + } + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1492,15 +1789,22 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 2, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, - "loop_group_id": "process_batch", - "loop_decision": { - "stop": true - } + "region_index": 2, + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": { + "stop": true + } + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1526,13 +1830,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, - "loop_group_id": "poll", - "max_tries": 0 + "region_index": 1, + "max_tries": 0, + "loop": { + "group_id": "poll", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1558,13 +1869,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, - "loop_group_id": "poll", - "max_tries": 0 + "region_index": 2, + "max_tries": 0, + "loop": { + "group_id": "poll", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1590,13 +1908,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, - "loop_group_id": "poll", - "max_tries": 0 + "region_index": 1, + "max_tries": 0, + "loop": { + "group_id": "poll", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1622,13 +1947,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, - "loop_group_id": "poll", - "max_tries": 0 + "region_index": 2, + "max_tries": 0, + "loop": { + "group_id": "poll", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1654,13 +1986,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, - "loop_group_id": "poll", - "loop_decision": null + "region_index": 1, + "loop": { + "group_id": "poll", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1686,13 +2025,20 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, - "loop_group_id": "poll", - "loop_decision": null + "region_index": 2, + "loop": { + "group_id": "poll", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": null + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1718,16 +2064,23 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 0, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0, - "loop_group_id": "poll", - "loop_decision": { - "stop": false, - "reason": "check returned 0" - } + "region_index": 1, + "loop": { + "group_id": "poll", + "index": 1, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": { + "stop": false, + "reason": "check returned 1" + } + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1753,16 +2106,23 @@ "airflow": { "taskInstance": { "try_number": 1, - "loop_index": 1, - "loop_max_iterations": 3, "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1, - "loop_group_id": "poll", - "loop_decision": { - "stop": true, - "reason": "check returned 1" - } + "region_index": 2, + "loop": { + "group_id": "poll", + "index": 2, + "max_iterations": 3, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null, + "decision": { + "stop": true, + "reason": "check returned 2" + } + }, + "expanded_ti_count": null }, "task": { "task_group": { @@ -1790,12 +2150,19 @@ "try_number": 1, "max_tries": 0, "map_index": 0, - "loop_index": 0, - "loop_max_iterations": 2, - "loop_group_id": "process_batch", "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0 + "region_index": 0, + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1823,12 +2190,19 @@ "try_number": 1, "max_tries": 0, "map_index": 1, - "loop_index": 0, - "loop_max_iterations": 2, - "loop_group_id": "process_batch", "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1 + "region_index": 1, + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1856,12 +2230,19 @@ "try_number": 1, "max_tries": 0, "map_index": 0, - "loop_index": 1, - "loop_max_iterations": 2, - "loop_group_id": "process_batch", "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0 + "region_index": 0, + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1889,12 +2270,19 @@ "try_number": 1, "max_tries": 0, "map_index": 1, - "loop_index": 1, - "loop_max_iterations": 2, - "loop_group_id": "process_batch", "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1 + "region_index": 1, + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1922,12 +2310,19 @@ "try_number": 1, "max_tries": 0, "map_index": 0, - "loop_index": 0, - "loop_max_iterations": 2, - "loop_group_id": "process_batch", "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0 + "region_index": 0, + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1955,12 +2350,19 @@ "try_number": 1, "max_tries": 0, "map_index": 1, - "loop_index": 0, - "loop_max_iterations": 2, - "loop_group_id": "process_batch", "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1 + "region_index": 1, + "loop": { + "group_id": "process_batch", + "index": 1, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -1988,12 +2390,19 @@ "try_number": 1, "max_tries": 0, "map_index": 0, - "loop_index": 1, - "loop_max_iterations": 2, - "loop_group_id": "process_batch", "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 0 + "region_index": 0, + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { @@ -2021,12 +2430,19 @@ "try_number": 1, "max_tries": 0, "map_index": 1, - "loop_index": 1, - "loop_max_iterations": 2, - "loop_group_id": "process_batch", "id": "{{ is_uuid(result) }}", "region_id": "{{ is_uuid(result) }}", - "region_index": 1 + "region_index": 1, + "loop": { + "group_id": "process_batch", + "index": 2, + "max_iterations": 2, + "region_id": "{{ is_uuid(result) }}", + "trajectory": 1, + "forked_from_region_id": null, + "resumes_from_index": null + }, + "expanded_ti_count": 2 }, "task": { "task_group": { diff --git a/providers/openlineage/tests/system/openlineage/operator.py b/providers/openlineage/tests/system/openlineage/operator.py index 104ee964d404b..76546c269441e 100644 --- a/providers/openlineage/tests/system/openlineage/operator.py +++ b/providers/openlineage/tests/system/openlineage/operator.py @@ -283,7 +283,7 @@ def _get_attempt(event: dict) -> Attempt | None: clear_number=dag_run.get("clear_number"), try_number=task_instance.get("try_number"), map_index=task_instance.get("map_index"), - loop_index=task_instance.get("loop_index"), + loop_index=(task_instance.get("loop") or {}).get("index"), region_id=task_instance.get("region_id"), ) @@ -299,7 +299,7 @@ def check_run_id_integrity( Check that ``run.runId`` identifies exactly one attempt of exactly one job, and return the problems found. Events are grouped by ``runId``. The attempt of an event is its Dag run id and clear number, plus for a task - event ``(try_number, map_index, loop_index, region_id)``, all from ``run.facets.airflow`` (``airflowDagRun`` for Dag + event ``(try_number, map_index, loop.index, region_id)``, all from ``run.facets.airflow`` (``airflowDagRun`` for Dag events). Query events, which are child runs of a task, have no attempt. For every group: 1. all events share one ``job.name``; diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index 32440f2402697..b513adc90b437 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -1245,8 +1245,17 @@ def body(): ti.task = loop_dag.get_task("body.work") ti.region_id = uuid.uuid4() ti.region_index = 1 + forked_from_region_id = uuid.uuid4() ti._ti_context_from_server.loop = LoopContext( - node_id="body", index=1, max_iterations=3, terminal_task_id="body.work", terminal_is_mapped=False + node_id="body", + index=1, + max_iterations=3, + terminal_task_id="body.work", + terminal_is_mapped=False, + region_id=ti.region_id, + trajectory=2, + forked_from_region_id=forked_from_region_id, + resumes_from_index=1, ) kwargs = {"error": ValueError("test")} if state == "failed" else {} @@ -1254,12 +1263,15 @@ def body(): method = {"running": "start_task", "failed": "fail_task"}.get(state, "complete_task") run_facets = getattr(listener.adapter, method).call_args.kwargs["run_facets"] - task_instance_facet = run_facets["airflow"].taskInstance - assert ( - task_instance_facet["loop_index"], - task_instance_facet["loop_max_iterations"], - task_instance_facet["loop_group_id"], - ) == (1, 3, "body") + assert run_facets["airflow"].taskInstance["loop"] == { + "group_id": "body", + "index": 1, + "max_iterations": 3, + "region_id": str(ti.region_id), + "trajectory": 2, + "forked_from_region_id": str(forked_from_region_id), + "resumes_from_index": 1, + } @pytest.mark.parametrize("state", ["running", "success", "failed"]) @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") @@ -1282,9 +1294,12 @@ def test_worker_event_of_loop_gate_has_its_decision_once_evaluated(self, state): getattr(listener.adapter, method).call_args.kwargs["run_facets"]["airflow"].taskInstance ) if state == "running": - assert "loop_decision" not in task_instance_facet + assert "loop" not in task_instance_facet else: - assert task_instance_facet["loop_decision"] == {"stop": state == "success", "reason": "checked"} + assert task_instance_facet["loop"]["decision"] == { + "stop": state == "success", + "reason": "checked", + } @pytest.mark.parametrize("state", ["running", "success", "failed", "skipped"]) @mock.patch( diff --git a/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py index fb4ccc0f30225..e347be0169e11 100644 --- a/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py +++ b/providers/openlineage/tests/unit/openlineage/test_system_tests_run_id_integrity.py @@ -43,11 +43,11 @@ def _event( run["runId"] = run_id dag_run = {"run_id": dag_run_id, "clear_number": clear_number} if job_type == "TASK": - task_instance: dict[str, int | str] = {"try_number": try_number} + task_instance: dict[str, int | str | dict[str, int]] = {"try_number": try_number} if map_index is not None: task_instance["map_index"] = map_index if loop_index is not None: - task_instance["loop_index"] = loop_index + task_instance["loop"] = {"index": loop_index} if region_id is not None: task_instance["region_id"] = region_id run["facets"]["airflow"] = {"dagRun": dag_run, "taskInstance": task_instance} diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index 03f817745f999..180ad26005418 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -80,6 +80,7 @@ get_parent_information_from_dagrun_conf, get_root_information_from_dagrun_conf, get_runtime_outlet_assets, + get_stored_task_instance_position, get_task_documentation, get_task_instance_map_index, get_task_instance_position, @@ -4053,6 +4054,143 @@ def body(): ) +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +def test_runtime_task_instance_position_reads_its_loop_and_expansion_from_the_task_instance(): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("position_dag") as dag: + + @task_group + def body(): + PythonOperator.partial(task_id="work", python_callable=list).expand(op_kwargs=[{}, {}]) + + create_loop(body, max_iterations=3) + loop_region_id, forked_from_region_id = uuid4(), uuid4() + ti = SimpleNamespace( + id=uuid4(), + region_id=uuid4(), + map_index=1, + task=dag.get_task("body.work"), + expanded_ti_count=2, + loop=SimpleNamespace( + index=2, + max_iterations=3, + region_id=loop_region_id, + trajectory=2, + forked_from_region_id=forked_from_region_id, + resumes_from_index=1, + ), + get_template_context=MagicMock(), + ) + + assert get_task_instance_position(ti) == TaskInstancePosition( + in_loop=True, + map_index=1, + expanded_ti_count=2, + loop_index=2, + loop_max_iterations=3, + loop_group_id="body", + loop_region_id=loop_region_id, + loop_trajectory=2, + loop_forked_from_region_id=forked_from_region_id, + loop_resumes_from_index=1, + ) + ti.get_template_context.assert_not_called() + + +@pytest.mark.db_test +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +def test_stored_task_instance_is_placed_as_airflow_places_it_for_the_worker(dag_maker, session): + from airflow.models.dynamic_region import DynamicRegion + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with dag_maker(serialized=True): + + @task_group + def body(): + EmptyOperator(task_id="plain") >> PythonOperator.partial( + task_id="mapped", python_callable=list + ).expand(op_kwargs=[{}, {}]) + + create_loop(body, max_iterations=3) + PythonOperator.partial(task_id="outside", python_callable=list).expand(op_kwargs=[{}, {}, {}]) + run = dag_maker.create_dagrun() + tis = {(ti.task_id, ti.region_index): ti for ti in run.task_instances} + loop_region = session.get(DynamicRegion, tis["body.plain", 1].region_id) + fork = DynamicRegion.fork_family(loop_region.id, 2, {loop_region.id: loop_region}, session=session) + session.flush() + rerun = TaskInstance( + task=tis["body.plain", 1].task, run_id=run.run_id, dag_version_id=tis["body.plain", 1].dag_version_id + ) + rerun.region_id, rerun.region_index = fork.id, 2 + session.add(rerun) + session.flush() + in_loop = {"in_loop": True, "loop_max_iterations": 3, "loop_group_id": "body", "loop_trajectory": 1} + + # Placing reuses the Dag the task instance holds: its region rows, then the count of a mapped expansion + with assert_queries_count(4): + plain = get_stored_task_instance_position(tis["body.plain", 1], session=session) + with assert_queries_count(5): + mapped = get_stored_task_instance_position(tis["body.mapped", 1], session=session) + outside = get_stored_task_instance_position(tis["outside", 2], session=session) + rerun_position = get_stored_task_instance_position(rerun, session=session) + + assert plain == TaskInstancePosition(map_index=-1, loop_index=1, loop_region_id=loop_region.id, **in_loop) + assert mapped == TaskInstancePosition( + map_index=1, expanded_ti_count=2, loop_index=1, loop_region_id=loop_region.id, **in_loop + ) + assert outside == TaskInstancePosition(in_loop=False, map_index=2, expanded_ti_count=3) + assert rerun_position == TaskInstancePosition( + in_loop=True, + map_index=-1, + loop_index=2, + loop_max_iterations=3, + loop_group_id="body", + loop_region_id=fork.id, + loop_trajectory=2, + loop_forked_from_region_id=loop_region.id, + loop_resumes_from_index=2, + ) + assert dict(TaskInstanceInfo(rerun, position=rerun_position))["loop"] == { + "group_id": "body", + "index": 2, + "max_iterations": 3, + "region_id": str(fork.id), + "trajectory": 2, + "forked_from_region_id": str(loop_region.id), + "resumes_from_index": 2, + } + + +@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +@mock.patch( + "airflow.models.task_coordinates.TaskCoordinateResolver.for_dag", side_effect=ValueError("Dag not found") +) +def test_stored_task_instance_is_placed_from_its_rows_when_airflow_fails_to(mock_for_dag): + from airflow.sdk import task_group + from airflow.sdk.definitions._internal.loop import create_loop + + with DAG("fallback_dag") as dag: + + @task_group + def body(): + EmptyOperator(task_id="work") + + create_loop(body, max_iterations=3) + ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) + ti.region_id = uuid4() + ti.region_index = 1 + session = MagicMock(spec=Session) + + assert get_stored_task_instance_position(ti, session=session) == TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=1, loop_max_iterations=3, loop_group_id="body" + ) + mock_for_dag.assert_called_once_with(dag, session) + session.execute.assert_not_called() + + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") def test_stored_mapped_task_in_loop_is_placed_without_iteration_when_its_region_query_fails(): from airflow.sdk import task_group @@ -4105,7 +4243,7 @@ def test_task_instance_info_uses_a_given_position_as_is(): info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=True, map_index=-1))) - assert "loop_index" not in info + assert "loop" not in info ti.get_template_context.assert_not_called() @@ -4128,7 +4266,7 @@ def test_stored_task_instance_without_task_or_region_has_no_loop_index(): @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") @pytest.mark.parametrize("mapped", [False, True]) @pytest.mark.parametrize("in_loop", [False, True]) -def test_runtime_task_instance_info_has_loop_index_only_inside_a_loop(in_loop, mapped): +def test_runtime_task_instance_info_has_loop_only_inside_a_loop(in_loop, mapped): from airflow.sdk import task_group from airflow.sdk.definitions._internal.loop import create_loop @@ -4158,11 +4296,9 @@ def body(): assert info["map_index"] == (1 if mapped else None) if in_loop: - assert (info["loop_index"], info["loop_max_iterations"], info["loop_group_id"]) == (2, 3, "body") + assert info["loop"] == {"group_id": "body", "index": 2, "max_iterations": 3} else: - assert "loop_index" not in info - assert "loop_max_iterations" not in info - assert "loop_group_id" not in info + assert "loop" not in info @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") @@ -4170,9 +4306,9 @@ def body(): ("until", "loop_index", "expected", "runs_out"), [ pytest.param(None, 1, {"stop": False}, False, id="no_until_continues"), - pytest.param(None, 2, {"stop": True}, False, id="no_until_stops_at_last"), - pytest.param("done", 0, {"stop": True, "reason": "done"}, False, id="until_stops_with_reason"), - pytest.param(False, 2, {"stop": False}, True, id="until_false_at_last"), + pytest.param(None, 3, {"stop": True}, False, id="no_until_stops_at_last"), + pytest.param("done", 1, {"stop": True, "reason": "done"}, False, id="until_stops_with_reason"), + pytest.param(False, 3, {"stop": False}, True, id="until_false_at_last"), ], ) def test_task_instance_info_has_the_decision_of_a_loop_gate(until, loop_index, expected, runs_out): @@ -4197,7 +4333,7 @@ def test_task_instance_info_has_the_decision_of_a_loop_gate(until, loop_index, e execute_loop_gate(gate, context) info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1))) - assert info["loop_decision"] == expected + assert info["loop"]["decision"] == expected def test_task_instance_info_has_no_loop_decision_for_other_tasks(): @@ -4205,7 +4341,7 @@ def test_task_instance_info_has_no_loop_decision_for_other_tasks(): info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1))) - assert "loop_decision" not in info + assert "loop" not in info @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Attempt UUIDs require Airflow 3") @@ -4297,23 +4433,37 @@ def test_task_instance_info_of_stored_instance_without_task_does_not_raise(): info = dict(TaskInstanceInfo(ti)) assert info["map_index"] is None - assert "loop_index" not in info + assert "loop" not in info def test_task_instance_info_uses_the_position_it_is_given(): ti = SimpleNamespace(id=uuid4(), region_id=uuid4(), map_index=-1) + loop_region_id, forked_from_region_id = uuid4(), uuid4() position = TaskInstancePosition( - in_loop=True, map_index=1, loop_index=2, loop_max_iterations=3, loop_group_id="body" + in_loop=True, + map_index=1, + expanded_ti_count=4, + loop_index=2, + loop_max_iterations=3, + loop_group_id="body", + loop_region_id=loop_region_id, + loop_trajectory=2, + loop_forked_from_region_id=forked_from_region_id, + loop_resumes_from_index=2, ) info = dict(TaskInstanceInfo(ti, position=position)) - assert (info["map_index"], info["loop_index"], info["loop_max_iterations"], info["loop_group_id"]) == ( - 1, - 2, - 3, - "body", - ) + assert (info["map_index"], info["expanded_ti_count"]) == (1, 4) + assert info["loop"] == { + "group_id": "body", + "index": 2, + "max_iterations": 3, + "region_id": str(loop_region_id), + "trajectory": 2, + "forked_from_region_id": str(forked_from_region_id), + "resumes_from_index": 2, + } @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4") @@ -4478,15 +4628,15 @@ def top_member(x): positions = get_orm_task_instance_positions(tis, session=session) assert {(ti.task_id, ti.region_index): positions[ti.id] for ti in tis} == { - ("body.plain", 0): TaskInstancePosition( - in_loop=True, map_index=-1, loop_index=0, loop_group_id="body", **in_loop_extra + ("body.plain", 1): TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=1, loop_group_id="body", **in_loop_extra ), - ("body.last", 0): TaskInstancePosition( - in_loop=True, map_index=-1, loop_index=0, loop_group_id="body", **in_loop_extra + ("body.last", 1): TaskInstancePosition( + in_loop=True, map_index=-1, loop_index=1, loop_group_id="body", **in_loop_extra ), **{ (task_id, map_index): TaskInstancePosition( - in_loop=True, map_index=map_index, loop_index=0, loop_group_id="body", **in_loop_extra + in_loop=True, map_index=map_index, loop_index=1, loop_group_id="body", **in_loop_extra ) for task_id in ("body.mapped", "body.mapped_group.member") for map_index in (0, 1) From 72264e05daf9c5882a9d661c03d9264cae719051 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 21 Oct 2026 15:12:56 +0200 Subject: [PATCH 08/10] Fix OpenLineage manual state change events on Airflow before 3.4 Stored task instances have no region_id before Airflow 3.4, so reading it before checking the version failed the event. Co-Authored-By: Claude --- .../src/airflow/providers/openlineage/utils/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index ea470e7c5896d..9170010609058 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1739,8 +1739,10 @@ def _get_core_task_instance_position( task_instance: TaskInstance, *, session: Session | None ) -> TaskInstancePosition | None: """Place a stored task instance in a region as Airflow places it for the worker, from Airflow 3.4.0 on.""" + if session is None or not AIRFLOW_V_3_4_PLUS: + return None region_id = task_instance.region_id - if session is None or not AIRFLOW_V_3_4_PLUS or not isinstance(region_id, UUID) or region_id.int == 0: + if not isinstance(region_id, UUID) or region_id.int == 0: return None # Reusing the Dag the task instance holds spares deserializing it again, unless it pins another version if (dag := getattr(getattr(task_instance, "task", None), "dag", None)) is None: From cbd52152363df00416f3fbea9de47b4ac34ee6db Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 21 Oct 2026 17:06:34 +0200 Subject: [PATCH 09/10] Clarify OpenLineage events for UI state changes and drop pre-release code The API server holds task instances without their task when a user sets a state, so those events never carried the airflow run facet the guide promised. The fallbacks for loop details only served Airflow builds from before apache/airflow#75132, which no release contains. Timing the stored position lookup makes its per task instance queries in the scheduler visible. Co-Authored-By: Claude --- .../docs/guides/task_execution_identity.rst | 25 ++++- .../providers/openlineage/plugins/listener.py | 4 +- .../providers/openlineage/utils/utils.py | 103 +++--------------- .../unit/openlineage/plugins/test_listener.py | 10 ++ .../unit/openlineage/utils/test_utils.py | 77 ++++--------- .../metrics/metrics_template.yaml | 6 + 6 files changed, 79 insertions(+), 146 deletions(-) diff --git a/providers/openlineage/docs/guides/task_execution_identity.rst b/providers/openlineage/docs/guides/task_execution_identity.rst index 8f2339b262288..b5175b459fb48 100644 --- a/providers/openlineage/docs/guides/task_execution_identity.rst +++ b/providers/openlineage/docs/guides/task_execution_identity.rst @@ -28,7 +28,9 @@ group events per execution, deduplicate alerts or tell whether a failure is fina All the fields below are in the ``airflow`` run facet of task events (``run.facets.airflow``), unless said otherwise. Dag run events carry the ``dag`` and ``dagRun`` parts in the ``airflowDagRun`` run facet. A field -that doesn't apply to an execution is left out of the event. +that doesn't apply to an execution is left out of the event. Events for a task state set in the UI or the API +have no ``airflow`` run facet, see :ref:`state set outside the worker +`. Run ids ^^^^^^^ @@ -214,14 +216,25 @@ not only ``forked_from_region_id``, when an earlier rerun resumed later than thi **Dag run cleared.** ``dagRun.clear_number`` grows and the Dag run gets a new ``runId``. The tasks it reruns get a new try. -**State set by hand.** Marking a task instance successful or failed in the UI or the API emits COMPLETE or FAIL -from the API server, with the same fields as the worker sends, ``loop`` and ``expanded_ti_count`` included -from Airflow 3.4.0 on. A task that never started has no START. A gate marked by hand has no ``loop.decision``. +.. _guides/task_execution_identity:outside-worker: + +**State set outside the worker.** Two kinds of state change emit COMPLETE or FAIL without the worker, so a task +that never started has no START for them, and a gate's event has no ``loop.decision``: + +* The scheduler failing a task instance, e.g. when the executor reports a failure the worker couldn't send. + The FAIL event has the same ``airflow`` run facet as a worker's, ``loop`` and ``expanded_ti_count`` + included on Airflow 3.4+. +* Marking a task instance successful, failed or skipped in the UI or the API. The API server holds the task + instance without its task, so the event has its ``runId`` (the attempt's UUID inside a loop, as a worker's + event of that attempt), ``job.name`` and ``ParentRunFacet``, but no ``airflow`` run facet: none of + ``taskInstance``, ``loop`` or ``expanded_ti_count``. To place such an event in a loop iteration or a mapped + expansion, match its ``runId`` with the worker's START of the same attempt, which has them. Limitations ^^^^^^^^^^^ -* Loops need Airflow 3.4. Pre-releases of Airflow 3.4.0 give no ``loop`` details beyond ``group_id``, - ``index`` and ``max_iterations``, and no ``expanded_ti_count``. +* Loops, ``region_id``, ``region_index`` and ``expanded_ti_count`` need Airflow 3.4. +* Events for task states set in the UI or the API have no ``airflow`` run facet, see + :ref:`state set outside the worker `. * The ``airflowState`` run facet of Dag run events keys task states by ``task_id`` only, so mapped instances and loop iterations of one task share one entry. diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py index 7a2a2e0cc7033..eca259e426a32 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py @@ -889,7 +889,9 @@ def _on_task_instance_manual_state_change( # pickler loses TaskGroup attributes and crashes event emission -- see # the equivalent note in `on_dag_run_running` (listener.py ~868). date = dagrun.logical_date or dagrun.run_after - position = get_stored_task_instance_position(ti, session=object_session(ti)) + # Runs in the scheduler or API server process and can query the region tables for each task instance + with Stats.timer("ol.manual_state_change.place_task_instance"): + position = get_stored_task_instance_position(ti, session=object_session(ti)) task_uuid = get_task_instance_run_id(ti, logical_date=date, position=position) parent_run_id = self.adapter.build_dag_run_id( dag_id=ti.dag_id, diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index 9170010609058..2234eea60745d 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1126,63 +1126,10 @@ class TaskInstanceInfo(InfoJsonEncodable): """ Encode a task instance as ``run.facets.airflow.taskInstance``, which tells which execution sent an event. - The fields that identify the execution, each left out when it doesn't apply: - - * ``id`` (Airflow 3+): the UUID of the attempt. Every try, and every task instance a clear creates again, - gets a new one, so it identifies one execution, and it is the ``runId`` of a task inside a loop. Airflow's - API can look it up, but it says nothing by itself. - * ``try_number``: the try, counted per ``region_id`` and ``region_index`` (see below). - * ``max_tries``: the last ``try_number`` that still gets retried on failure, so a failed try with - ``try_number > max_tries`` is final. It equals the task's ``retries`` (in ``task``) only until a clear: - clearing sets it to the last try number plus ``retries``, giving the task its retries again. - * ``map_index``: the position in a mapped task's expansion that users see. Absent for a task that isn't - mapped. - * ``expanded_ti_count`` (Airflow 3.4+): the number of live instances in the mapped task's expansion, counted - within the loop iteration inside a loop, so ``map_index`` is one of ``range(expanded_ti_count)``. Absent - for a task that isn't mapped. - * ``region_id`` and ``region_index`` (Airflow 3.4+): the coordinates Airflow stores the task instance under. - A region is one execution of a loop, or of a mapped task's expansion, and is shared by all the task - instances it creates. ``region_index`` is the position inside it: the iteration for a task in a loop that - isn't mapped, the map index for a mapped task. Neither is unique: ``(dag_id, dagRun.run_id, task_id, - region_id, region_index, try_number)`` is, as Airflow enforces it (``task_instance_try_key``). Both are - absent for a task in no region, i.e. one that is neither mapped nor in a loop. - * ``loop`` (Airflow 3.4+): the iteration of the innermost loop around the task, absent outside a loop. - Loops can't be nested. Its fields, each left out when unknown: - - * ``group_id``: the loop's group id, which can differ from ``task.task_group`` when the task sits in a - task group inside the loop. - * ``index`` and ``max_iterations``: the iteration, counted from 1, and the loop's limit. - * ``region_id``: the loop region holding the iteration, shared by every task instance of that pass of - the loop, mapped ones included, whose own ``region_id`` is their expansion's. - * ``trajectory``, ``forked_from_region_id`` and ``resumes_from_index``: clearing a loop gate with the - later iterations forks the loop region, see below. ``trajectory`` is 1 for the original region and one - more for each fork, ``forked_from_region_id`` the previous region of the chain and - ``resumes_from_index`` the first iteration the fork holds. Both are absent for the original region. - * ``decision``: ``stop`` and the optional ``reason`` returned by ``until``, as evaluated by a loop's gate, - the task that ends each iteration (``gate_task_id`` of the loop group, in ``task.task_group`` and in - ``taskGroups`` of ``AirflowJobFacet``). Only the gate's COMPLETE and FAIL events carry it, as the gate - decides while it runs. - - How a gate's ``loop.decision`` reads, with ``loop.index``, ``loop.max_iterations`` and the group's - ``has_until``: - - * COMPLETE with ``stop`` false: the loop goes on to iteration ``index + 1``. - * COMPLETE with ``stop`` true and ``has_until``: ``until`` ended the loop at iteration ``index``, which can be - the last one. - * COMPLETE with ``stop`` true and no ``has_until``: the loop ran all its iterations, as one without ``until`` - always stops at ``index == max_iterations``. - * FAIL with ``stop`` false: ``until`` was still false at the last iteration (``index == max_iterations``), - so the loop ran out of iterations and fails. That's the only failure after the decision is made. - * FAIL without ``decision``: the gate failed before deciding, e.g. ``until`` raised or returned something - other than a bool or a ``LoopDecision``; the error facet tells which. - - How a cleared loop reads: clearing a loop from iteration ``k`` retries what was cleared in iteration ``k`` in - place, with the same region and the next ``try_number``. The iterations after ``k`` are created again in a - new region forked from the loop's newest region, with ``loop.trajectory`` one more and ``loop.resumes_from_index`` - ``k + 1``, and try numbers count per region, so the new execution of iteration ``k + 1`` can repeat the task, - ``loop.index``, ``map_index`` and ``try_number`` of the one it replaces. ``loop.region_id`` and - ``loop.trajectory`` tell the two apart, and the higher trajectory is the current one. - ``dagRun.clear_number`` counts clears of the whole Dag run only. + Adds to the included fields, each left out when it doesn't apply: ``id``, ``map_index``, + ``rendered_map_index``, ``expanded_ti_count``, ``region_id``, ``region_index`` and ``loop``. The provider's + "Identifying task executions in events" guide (``guides/task_execution_identity.rst``) explains how + consumers read them, and ``AirflowRunFacet.json`` describes each one. """ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "try_number"] @@ -1220,9 +1167,8 @@ def _extend_fields(self) -> None: self.id = str(ti_id) self._fields.append("id") loop = _build_loop_info(position) if position and position.in_loop else {} - # Only a loop gate's runtime task instance has one, once it has evaluated ``until``. Airflow 3.4.0 exposes - # it, a pre-release only as a private attribute. - decision = getattr(self.obj, "loop_decision", None) or getattr(self.obj, "_loop_decision", None) + # Only a loop gate's runtime task instance has one, once it has evaluated ``until`` + decision = getattr(self.obj, "loop_decision", None) if decision is not None: loop["decision"] = {"stop": decision.stop} if decision.reason is not None: @@ -1609,12 +1555,11 @@ class TaskInstancePosition: """ Where a task instance runs: in a loop or not, the map index users see, and its loop iteration. - The loop fields are ``None`` outside a loop, and also when they can't be known. Airflow 3.4+ tells all of - them, and ``expanded_ti_count`` for a mapped task: a runtime task instance through its ``loop`` and - ``expanded_ti_count``, read by ``get_task_instance_position``, and a stored one through Airflow's - ``TaskCoordinateResolver``, called by ``get_stored_task_instance_position``. Before that, a - runtime task instance only has the iteration and limit from its template context, and a stored one only what - its task and region rows tell. ``TaskInstanceInfo`` explains what these fields mean in events. + The loop fields are ``None`` outside a loop, and also when they can't be known. A runtime task instance + tells all of them, and ``expanded_ti_count`` for a mapped task, through its ``loop`` and + ``expanded_ti_count``, read by ``get_task_instance_position``. A stored one tells them through Airflow's + ``TaskCoordinateResolver``, called by ``get_stored_task_instance_position``, and otherwise only what its + task and region rows tell. ``AirflowRunFacet.json`` describes what these fields mean in events. """ in_loop: bool @@ -1703,30 +1648,20 @@ def _add_runtime_loop_iteration( ) -> TaskInstancePosition: if isinstance(task_instance, TaskInstance): return position - # Only runtime task instances of Airflow 3.4+ expose these, the expansion size from 3.4.0 on + # Only runtime task instances of Airflow 3.4+ expose these position = attrs.evolve(position, expanded_ti_count=getattr(task_instance, "expanded_ti_count", None)) if not position.in_loop or position.loop_index is not None: return position if (loop := getattr(task_instance, "loop", None)) is None: - # Run ids, map indexes and the lineage macros don't need the iteration, so only this reads the template - # context: it re-validates the params on every call, which can raise from inside a user's template. - try: - loop = task_instance.get_template_context().get("loop") - except Exception: - # The iteration only describes the run, so failing to read it must not cost the event its run id - log.warning("OpenLineage failed to read the loop iteration of task instance %s", task_instance.id) - log.debug("Exception details:", exc_info=True) - return position - if loop is None: return position return attrs.evolve( position, loop_index=loop.index, loop_max_iterations=loop.max_iterations, - loop_region_id=getattr(loop, "region_id", None), - loop_trajectory=getattr(loop, "trajectory", None), - loop_forked_from_region_id=getattr(loop, "forked_from_region_id", None), - loop_resumes_from_index=getattr(loop, "resumes_from_index", None), + loop_region_id=loop.region_id, + loop_trajectory=loop.trajectory, + loop_forked_from_region_id=loop.forked_from_region_id, + loop_resumes_from_index=loop.resumes_from_index, ) @@ -1738,7 +1673,7 @@ def get_task_instance_position(task_instance: TaskInstance | RuntimeTaskInstance def _get_core_task_instance_position( task_instance: TaskInstance, *, session: Session | None ) -> TaskInstancePosition | None: - """Place a stored task instance in a region as Airflow places it for the worker, from Airflow 3.4.0 on.""" + """Place a stored task instance in a region as Airflow places it for the worker, from Airflow 3.4 on.""" if session is None or not AIRFLOW_V_3_4_PLUS: return None region_id = task_instance.region_id @@ -1749,8 +1684,6 @@ def _get_core_task_instance_position( return None from airflow.models.task_coordinates import TaskCoordinateResolver - if not hasattr(TaskCoordinateResolver, "get_loop_context"): # An Airflow 3.4 pre-release - return None try: # The savepoint keeps the caller's transaction usable after a failure with session.begin_nested(): @@ -1790,7 +1723,7 @@ def get_stored_task_instance_position( """ Place a stored task instance for an event about it, as Airflow places it for the worker when it can. - That needs Airflow 3.4.0+ and a session. Otherwise, or when Airflow fails to, it is placed with + That needs Airflow 3.4+ and a session. Otherwise, or when Airflow fails to, it is placed with ``get_orm_task_instance_positions``, which only knows what the task and region rows tell. """ if (position := _get_core_task_instance_position(task_instance, session=session)) is not None: diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index b513adc90b437..c16161d784912 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -1344,6 +1344,16 @@ def body(): in_loop=True, map_index=-1, loop_index=2, loop_max_iterations=3, loop_group_id="body" ) + def test_manual_event_times_placing_its_task_instance(self, mocker): + listener, ti = self._create_listener_and_task_instance(runtime_ti=False) + listener._executor = mocker.Mock(spec=ProcessPoolExecutor) + mocker.patch.object(listener, "submit_callable", autospec=True) + stats = mocker.patch("airflow.providers.openlineage.plugins.listener.Stats", autospec=True) + + listener._on_task_instance_manual_state_change(ti, ti.dag_run, "success") + + stats.timer.assert_called_once_with("ol.manual_state_change.place_task_instance") + @pytest.mark.skip("Rendering fields is not migrated yet in Airflow 3") @patch("airflow.models.BaseOperator.render_template") def test_listener_does_not_change_task_instance(self, render_mock, mock_supervisor_comms, spy_agency): diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index 180ad26005418..b1ddb9972bf37 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -3967,7 +3967,6 @@ def body(): region_id=uuid4(), map_index=-1, task=dag.get_task("body.work"), - get_template_context=MagicMock(), ) logical_date = timezone.datetime(2024, 1, 1) @@ -3981,7 +3980,6 @@ def body(): ) ) assert get_task_instance_map_index(ti) == -1 - ti.get_template_context.assert_not_called() def test_task_instance_run_id_uses_the_map_index_of_the_given_position(): @@ -3998,39 +3996,7 @@ def test_task_instance_run_id_uses_the_map_index_of_the_given_position(): @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") -@pytest.mark.parametrize("has_loop_context", [True, False]) -def test_runtime_task_instance_position_has_the_iteration_of_its_loop_context(has_loop_context): - from airflow.sdk import task_group - from airflow.sdk.definitions._internal.loop import create_loop - - with DAG("position_dag") as dag: - - @task_group - def body(): - EmptyOperator(task_id="work") - - create_loop(body, max_iterations=3) - loop = SimpleNamespace(index=2, max_iterations=3) if has_loop_context else None - ti = SimpleNamespace( - id=uuid4(), - region_id=uuid4(), - map_index=-1, - task=dag.get_task("body.work"), - get_template_context=MagicMock(return_value={"loop": loop}), - ) - - assert get_task_instance_position(ti) == TaskInstancePosition( - in_loop=True, - map_index=-1, - loop_index=2 if has_loop_context else None, - loop_max_iterations=3 if has_loop_context else None, - loop_group_id="body", - ) - ti.get_template_context.assert_called_once_with() - - -@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") -def test_runtime_task_instance_position_without_its_iteration_when_the_context_fails(): +def test_runtime_task_instance_position_without_its_loop_context_has_no_iteration(): from airflow.sdk import task_group from airflow.sdk.definitions._internal.loop import create_loop @@ -4042,11 +4008,7 @@ def body(): create_loop(body, max_iterations=3) ti = SimpleNamespace( - id=uuid4(), - region_id=uuid4(), - map_index=-1, - task=dag.get_task("body.work"), - get_template_context=MagicMock(side_effect=ValueError("invalid params")), + id=uuid4(), region_id=uuid4(), map_index=-1, task=dag.get_task("body.work"), loop=None ) assert get_task_instance_position(ti) == TaskInstancePosition( @@ -4081,7 +4043,6 @@ def body(): forked_from_region_id=forked_from_region_id, resumes_from_index=1, ), - get_template_context=MagicMock(), ) assert get_task_instance_position(ti) == TaskInstancePosition( @@ -4096,7 +4057,6 @@ def body(): loop_forked_from_region_id=forked_from_region_id, loop_resumes_from_index=1, ) - ti.get_template_context.assert_not_called() @pytest.mark.db_test @@ -4238,15 +4198,6 @@ def test_task_in_a_task_group_of_unknown_type_is_assumed_in_a_loop(): assert get_task_instance_position(ti) == TaskInstancePosition(in_loop=True, map_index=-1) -def test_task_instance_info_uses_a_given_position_as_is(): - ti = SimpleNamespace(id=uuid4(), map_index=-1, get_template_context=MagicMock()) - - info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=True, map_index=-1))) - - assert "loop" not in info - ti.get_template_context.assert_not_called() - - @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires Airflow 3.4") def test_stored_task_instance_without_task_or_region_has_no_loop_index(): ti = TaskInstance(task=MagicMock(), run_id="run_id", dag_version_id=uuid7()) @@ -4290,7 +4241,14 @@ def body(): task=dag.get_task("body.work"), ) if in_loop: - ti.get_template_context = lambda: {"loop": SimpleNamespace(index=2, max_iterations=3)} + ti.loop = SimpleNamespace( + index=2, + max_iterations=3, + region_id=None, + trajectory=None, + forked_from_region_id=None, + resumes_from_index=None, + ) info = dict(TaskInstanceInfo(ti)) @@ -4315,6 +4273,10 @@ def test_task_instance_info_has_the_decision_of_a_loop_gate(until, loop_index, e from airflow.sdk import LoopDecision from airflow.sdk.definitions._internal.loop import LoopGateOperator from airflow.sdk.execution_time.loop import LoopMaxIterationsExceeded, execute_loop_gate + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + + class GateTaskInstance(SimpleNamespace): + loop_decision = RuntimeTaskInstance.loop_decision if until is False: until_callable = lambda: False @@ -4323,7 +4285,7 @@ def test_task_instance_info_has_the_decision_of_a_loop_gate(until, loop_index, e else: until_callable = None gate = LoopGateOperator(task_id="__loop_gate", until=until_callable) - ti = SimpleNamespace(id=uuid7(), map_index=-1) + ti = GateTaskInstance(id=uuid7(), map_index=-1, _loop_decision=None) context = {"loop": SimpleNamespace(index=loop_index, max_iterations=3), "ti": ti} if runs_out: @@ -4437,7 +4399,14 @@ def test_task_instance_info_of_stored_instance_without_task_does_not_raise(): def test_task_instance_info_uses_the_position_it_is_given(): - ti = SimpleNamespace(id=uuid4(), region_id=uuid4(), map_index=-1) + # The task instance's own loop and expansion differ, so the info can only match the given position + ti = SimpleNamespace( + id=uuid4(), + region_id=uuid4(), + map_index=-1, + expanded_ti_count=9, + loop=SimpleNamespace(index=9, max_iterations=9), + ) loop_region_id, forked_from_region_id = uuid4(), uuid4() position = TaskInstancePosition( in_loop=True, diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index 45ce30f648cae..286a8da14278c 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -903,6 +903,12 @@ metrics: legacy_name: "ol.extract.{event_type}.{operator_name}" name_variables: ["event_type", "operator_name"] + - name: "ol.manual_state_change.place_task_instance" + description: "Milliseconds taken to place a task instance in a manual OpenLineage state change event." + type: "timer" + legacy_name: "-" + name_variables: [] + - name: "airflow.io.load_filesystems" description: "Milliseconds taken to load filesystem implementations from providers" type: "timer" From 9ef857c58a19f5db24228b023045c21d3d098beb Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 21 Oct 2026 17:19:10 +0200 Subject: [PATCH 10/10] Avoid reloading Dags for OpenLineage events of failed task instances Placing a task instance pinned to another Dag version than the one the scheduler holds made Airflow load and deserialize that version again for every such task instance, inside the scheduler loop. Those now fall back to reading their region rows. That bounds the cost instead of only measuring it, so the timer added for it goes away before any release has it. Co-Authored-By: Claude --- .../docs/guides/task_execution_identity.rst | 4 +++- .../providers/openlineage/plugins/listener.py | 4 +--- .../providers/openlineage/utils/utils.py | 5 ++++- .../unit/openlineage/plugins/test_listener.py | 10 ---------- .../tests/unit/openlineage/utils/test_utils.py | 17 +++++++++++++++-- .../observability/metrics/metrics_template.yaml | 6 ------ 6 files changed, 23 insertions(+), 23 deletions(-) diff --git a/providers/openlineage/docs/guides/task_execution_identity.rst b/providers/openlineage/docs/guides/task_execution_identity.rst index b5175b459fb48..19ea08686e0f4 100644 --- a/providers/openlineage/docs/guides/task_execution_identity.rst +++ b/providers/openlineage/docs/guides/task_execution_identity.rst @@ -223,7 +223,9 @@ that never started has no START for them, and a gate's event has no ``loop.decis * The scheduler failing a task instance, e.g. when the executor reports a failure the worker couldn't send. The FAIL event has the same ``airflow`` run facet as a worker's, ``loop`` and ``expanded_ti_count`` - included on Airflow 3.4+. + included on Airflow 3.4+. When the task instance runs another Dag version than the one the scheduler + holds, it has no ``expanded_ti_count`` and its ``loop`` only has ``group_id``, ``index`` and + ``max_iterations``. * Marking a task instance successful, failed or skipped in the UI or the API. The API server holds the task instance without its task, so the event has its ``runId`` (the attempt's UUID inside a loop, as a worker's event of that attempt), ``job.name`` and ``ParentRunFacet``, but no ``airflow`` run facet: none of diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py index eca259e426a32..7a2a2e0cc7033 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py @@ -889,9 +889,7 @@ def _on_task_instance_manual_state_change( # pickler loses TaskGroup attributes and crashes event emission -- see # the equivalent note in `on_dag_run_running` (listener.py ~868). date = dagrun.logical_date or dagrun.run_after - # Runs in the scheduler or API server process and can query the region tables for each task instance - with Stats.timer("ol.manual_state_change.place_task_instance"): - position = get_stored_task_instance_position(ti, session=object_session(ti)) + position = get_stored_task_instance_position(ti, session=object_session(ti)) task_uuid = get_task_instance_run_id(ti, logical_date=date, position=position) parent_run_id = self.adapter.build_dag_run_id( dag_id=ti.dag_id, diff --git a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py index 2234eea60745d..898630aea789d 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py +++ b/providers/openlineage/src/airflow/providers/openlineage/utils/utils.py @@ -1679,9 +1679,12 @@ def _get_core_task_instance_position( region_id = task_instance.region_id if not isinstance(region_id, UUID) or region_id.int == 0: return None - # Reusing the Dag the task instance holds spares deserializing it again, unless it pins another version if (dag := getattr(getattr(task_instance, "task", None), "dag", None)) is None: return None + # The resolver would load and deserialize any other pinned version again, for every task instance, in the + # scheduler or API server process + if getattr(dag, "dag_version_id", None) != task_instance.dag_version_id: + return None from airflow.models.task_coordinates import TaskCoordinateResolver try: diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index c16161d784912..b513adc90b437 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -1344,16 +1344,6 @@ def body(): in_loop=True, map_index=-1, loop_index=2, loop_max_iterations=3, loop_group_id="body" ) - def test_manual_event_times_placing_its_task_instance(self, mocker): - listener, ti = self._create_listener_and_task_instance(runtime_ti=False) - listener._executor = mocker.Mock(spec=ProcessPoolExecutor) - mocker.patch.object(listener, "submit_callable", autospec=True) - stats = mocker.patch("airflow.providers.openlineage.plugins.listener.Stats", autospec=True) - - listener._on_task_instance_manual_state_change(ti, ti.dag_run, "success") - - stats.timer.assert_called_once_with("ol.manual_state_change.place_task_instance") - @pytest.mark.skip("Rendering fields is not migrated yet in Airflow 3") @patch("airflow.models.BaseOperator.render_template") def test_listener_does_not_change_task_instance(self, render_mock, mock_supervisor_comms, spy_agency): diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index b1ddb9972bf37..72caa0fafa2bf 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -4078,6 +4078,8 @@ def body(): PythonOperator.partial(task_id="outside", python_callable=list).expand(op_kwargs=[{}, {}, {}]) run = dag_maker.create_dagrun() tis = {(ti.task_id, ti.region_index): ti for ti in run.task_instances} + # As the scheduler's Dag bag sets it on the Dag it loads for the run + tis["body.plain", 1].task.dag.dag_version_id = tis["body.plain", 1].dag_version_id loop_region = session.get(DynamicRegion, tis["body.plain", 1].region_id) fork = DynamicRegion.fork_family(loop_region.id, 2, {loop_region.id: loop_region}, session=session) session.flush() @@ -4125,10 +4127,17 @@ def body(): @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4") +@pytest.mark.parametrize( + "pins_another_version", + [ + pytest.param(False, id="airflow_fails"), + pytest.param(True, id="pins_another_version"), + ], +) @mock.patch( "airflow.models.task_coordinates.TaskCoordinateResolver.for_dag", side_effect=ValueError("Dag not found") ) -def test_stored_task_instance_is_placed_from_its_rows_when_airflow_fails_to(mock_for_dag): +def test_stored_task_instance_is_placed_from_its_rows_without_airflow(mock_for_dag, pins_another_version): from airflow.sdk import task_group from airflow.sdk.definitions._internal.loop import create_loop @@ -4142,12 +4151,16 @@ def body(): ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7()) ti.region_id = uuid4() ti.region_index = 1 + dag.dag_version_id = uuid7() if pins_another_version else ti.dag_version_id session = MagicMock(spec=Session) assert get_stored_task_instance_position(ti, session=session) == TaskInstancePosition( in_loop=True, map_index=-1, loop_index=1, loop_max_iterations=3, loop_group_id="body" ) - mock_for_dag.assert_called_once_with(dag, session) + if pins_another_version: + mock_for_dag.assert_not_called() + else: + mock_for_dag.assert_called_once_with(dag, session) session.execute.assert_not_called() diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index 286a8da14278c..45ce30f648cae 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -903,12 +903,6 @@ metrics: legacy_name: "ol.extract.{event_type}.{operator_name}" name_variables: ["event_type", "operator_name"] - - name: "ol.manual_state_change.place_task_instance" - description: "Milliseconds taken to place a task instance in a manual OpenLineage state change event." - type: "timer" - legacy_name: "-" - name_variables: [] - - name: "airflow.io.load_filesystems" description: "Milliseconds taken to load filesystem implementations from providers" type: "timer"