You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "kacpermuda (via GitHub)" <gi...@apache.org> on 2026/10/14 16:52:20 UTC
[PR] Improve task loop information and tests for OpenLineage provider [airflow]
kacpermuda opened a new pull request, #74752:
URL: https://github.com/apache/airflow/pull/74752
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
For user-facing UI changes, please attach before/after screenshots (or a short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
Task loops are new, and OpenLineage events need to describe them from day one. Inside a loop the iteration of a
task was invisible: `taskInstance.map_index` is only set for mapped tasks and the runId is the attempt UUID, so
iterations of one task looked identical, and nothing verified that runIds stay consistent across START and
COMPLETE, retries, map indexes and iterations.
- `taskInstance` facet gets `loop_index` and `loop_max_iterations`. Loop groups expose `max_iterations`,
`has_until`, `gate_task_id` and `terminal_task_id` in `task_group` and in `AirflowJobFacet.taskGroups`.
- The manual state change path passes its computed position to the facet, and `TaskInstanceInfo` no longer
raises for a stored task instance it can't place.
- `OpenLineageTestOperator` checks runId integrity by default: one job and attempt per runId, a complete
START and COMPLETE/FAIL lifecycle, and one runId per attempt.
- New example Dags for retries, a rescheduled sensor, and loops with a plain task, a retried loop task and a
retried mapped loop task (Airflow 3.4+).
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (Claude Code)
Generated-by: Claude Code (Sonnet 5.5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315008221
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1123,9 +1123,67 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
Review Comment:
Trimmed in 7441d2d: `TaskInstancePosition` now only describes its own fields and points to `TaskInstanceInfo`, which keeps the full explanation for readers of the code, while the schema keeps short per-field descriptions for consumers.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -2049,6 +2224,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")
Review Comment:
Added in 7441d2d, to the `task_group` definition in `AirflowRunFacet.json` and to `taskGroups` in `AirflowJobFacet.json`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315004733
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1489,28 +1587,38 @@ 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
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
Review Comment:
Good catch, fixed in 7441d2d: `_find_enclosing_loop` assumes a loop again when the walk stops on a group of an unknown type, and logs a warning, so the worst case is an extra runId rather than merged iterations.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
Review Comment:
Fixed in 7441d2d: when the query only fills in mapped tasks' iterations, it runs in a savepoint, and a failure only leaves `loop_index` unset and logs a warning. The event still goes out and the caller's transaction stays usable. Task instances without their task still raise as before this PR.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
+ )
)
)
}
for task_instance in unplaced:
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,
Review Comment:
I checked this against core: it never creates a region for a mapped task group. Every task that needs expansion, a task in a mapped task group included, gets its own region with `node_id == task_id` (core's log addressing classifies regions with the same rule), so `g.t` takes the first branch with its map index. 7441d2d makes the parsing check the shapes it expects (a loop's region is top-level, a mapped task's region sits directly in its loop's) and places anything else as a loop without its iteration, with a warning. `test_region_shapes_of_airflow_are_the_ones_openlineage_places` builds a real Dag run with these shapes, so a change in core fails the provider's tests.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
+ )
)
)
}
for task_instance in unplaced:
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
Review Comment:
Airflow doesn't allow that nesting today: the SDK rejects expanding a task inside a mapped task group, and inside a loop each task of a mapped task group gets its region directly under the loop's. 7441d2d states this, uses `parent_region_index` only when the parent region is the task's own loop and top-level, and pins the SDK restrictions in `test_region_shapes_stay_limited_by_the_task_sdk`.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1135,11 +1193,44 @@ 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")
+ # 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}
Review Comment:
Changed in 7441d2d: `reason` is only added when set. (On the wire it was already absent, as event serialization drops `None` values.)
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix task loop information and tests for OpenLineage provider [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #74752:
URL: https://github.com/apache/airflow/pull/74752#issuecomment-6362108166
**Scope (non-blocking).** The title and description are about loops, but the diff also adds generic facet fields (`state`, `max_tries`, bundle name/version for stored task instances, the `rendered_map_index` rework, `retry_delay`, job-facet `trigger_rule`/`depends_on_past`/`retries`) and three non-loop system-test Dags (retry, sensor reschedule, task states), ~11k lines in total. Splitting into (1) loop/region/runId, (2) generic facet fields, (3) system-test expansion would let each be reviewed on its own. The problems in the inline comments sit in (1) (scheduler placement) and (3) (Airflow 3.0 expectations).
Minor:
- `sort_by_execution` is defined identically in `example_openlineage_task_loop_dag.py` and `example_openlineage_task_loop_groups_dag.py`; the "every task event's identity: ..." paragraph is pasted into five Dag docstrings.
- `test_stored_task_instance_in_loop_has_loop_index_and_max_iterations` (test_utils.py:4804) never asserts `max_iterations`.
- A scheduler-sent FAIL for a mapped task without `map_index_template` no longer carries `taskInstance.rendered_map_index` (`"2"` before, absent now). That is parity with the worker, but it is consumer-visible and could use a changelog note.
- `openlineage_manual_lineage_dag.json`: the added START/COMPLETE entries only assert `try_number`; they exist so the integrity check sees the lifecycle, which is fine, but they add no field coverage.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315004733
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1489,28 +1587,38 @@ 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
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
Review Comment:
Good catch, fixed in 7441d2d: `_find_enclosing_loop` assumes a loop again when the walk stops on a group of an unknown type, and logs a warning, so the worst case is an extra runId rather than merged iterations.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix task loop information and tests for OpenLineage provider [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #74752:
URL: https://github.com/apache/airflow/pull/74752#issuecomment-6362108410
At `listener.py:916` (`get_task_documentation(task)`, not part of this diff): **Scheduler FAIL for mapped tasks is dropped (pre-existing, but it hides the new mapped fields).** On the scheduler `ti.task` is a `SerializedMappedOperator` for a mapped TI, and `get_task_documentation(task)` reads `operator.doc`, which on that class is `attrs.field(init=False)` with no default. It raises `AttributeError: ... object has no attribute 'doc'` (reproduced on a serialized Dag; a plain task returns `(None, None)`). The hook wrapper swallows it, so no event is sent for any mapped task instance the scheduler fails, and only after the `get_stored_task_instance_position` call above has already run its queries. A heartbeat kill of a mapped-in-loop slot on a live 3.4.0 stack logged exactly this error.
So the mapped part of "FAIL events sent by the scheduler carry the same fields as worker ones" (`expanded_ti_count`, mapped-in-loop `loop.index` from the region rows) can't reach an event today, and the unit tests can't notice because `test_manual_event_...` uses mocks.
Fix: `getattr(operator, "doc", None)` for the `doc*` attributes (or give the serialized class defaults), and add a listener test that goes through a real serialized mapped task.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on PR #74752:
URL: https://github.com/apache/airflow/pull/74752#issuecomment-6321838594
Finishing up last test fixes and additions, should be ready for re-review soon
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4314212153
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1489,28 +1587,38 @@ 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
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
Review Comment:
**Unrecognized group types now count as "not in a loop" instead of "in a loop"**
This function walks up the task's parent groups looking for a loop group. The walk stops as soon as it meets something that isn't a `TaskGroup` / `SerializedTaskGroup`. The old `_is_task_in_loop` then returned `group is not None`: if the walk stopped on some unexpected non-`None` object, it assumed the task *was* in a loop. The new code returns `None` in that case, so it now assumes the task is *not* in a loop.
Why the direction of this fallback matters:
- When a task is in a loop, the listener uses the task instance's UUID (`ti.id`) as the OpenLineage runId. Each iteration gets a new UUID, so iterations stay separate.
- When it's not in a loop, the runId is built from `(dag_id, task_id, try_number, logical_date, map_index)`. Call that the "computed runId": the same inputs always give the same runId.
Every iteration of a looped task has the same values for all five inputs. If loop detection wrongly says "not in a loop", iteration 0 and iteration 1 emit START/COMPLETE under one runId, and consumers merge their lineage into a single run.
I'm not sure which group types can reach this point in practice (maybe a group from a different SDK version, a proxy, or a subclass that doesn't inherit from these two classes). But if one does, the old fallback gave a harmless extra UUID while the new one silently merges runs. Keeping `return group is not None`, or at least logging when the walk stops on an unexpected type, seems safer.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
Review Comment:
**A new database query in the manual state change path can make the whole event disappear**
Before this PR, a stored task instance whose task object was attached got its position computed purely in Python, with no query. This line adds a second group, `without_iteration`: mapped tasks inside a loop. Their iteration number only lives on the `DynamicRegion` row (the record Airflow keeps for each loop iteration or mapped expansion), so this now queries the database for them.
That code runs from `_on_task_instance_manual_state_change`, the handler for when the scheduler fails a task (e.g. `handle_failure` after a heartbeat timeout) or a user marks a task failed or successful in the UI. So this query runs inside the scheduler/API server process.
The scenario I'm worried about:
1. A mapped task inside a loop gets marked failed by the scheduler.
2. `get_orm_task_instance_positions` now runs the `DynamicRegion` query for it.
3. The query raises (connection hiccup, the TI's session already detached or closed, etc.).
4. The broad `try/except` in `_on_task_instance_manual_state_change` catches the error and returns, so **no FAIL event is emitted for that task at all**.
All of this is only to fill `loop_index`, which is descriptive: the runId doesn't depend on it. `_add_runtime_loop_iteration` states this exact rule in its comment ("failing to read it must not cost the event its run id"). Could the `without_iteration` lookup get its own `try/except`, so a failure only leaves `loop_index` unset?
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
+ )
)
)
}
for task_instance in unplaced:
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,
Review Comment:
**Tasks in a mapped task group (with no loop at all) get reported as loop iterations**
This `else` branch handles a stored task instance without its task object attached. That happens for the asset-event source and for the manual state change path. Here the code reads the task's `DynamicRegion` row and assumes that if the row's `node_id` isn't the task itself, the row must belong to an enclosing loop.
But `DynamicRegion` rows are also created for mapped task groups (`@task_group` + `.expand()`), not just for loops. Example:
```python
@task_group
def g(x):
t(x)
g.expand(x=[1, 2, 3]) # no loop anywhere
```
For task `g.t`, map index 1:
1. Its region row has `node_id = "g"`, which isn't `"g.t"`, so we land in this `else`.
2. We set `in_loop=True`, `map_index=-1`, `loop_index=1`, `loop_group_id="g"`.
So the event says "iteration 1 of loop `g`" for a task that isn't in a loop, drops its real `map_index` (1), and uses the UUID runId. The same task's START/COMPLETE from the worker uses the computed runId (built from dag/task/try/date/map_index), so the manual-state FAIL and the worker's START end up with different runIds.
Checking the region's kind (loop vs. mapped group) before treating it as a loop would avoid this, if `DynamicRegion` records that.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
+ )
)
)
}
for task_instance in unplaced:
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
Review Comment:
**`parent_region_index` is the loop iteration only if the loop is the direct parent**
For a mapped task in a loop, this takes `region.parent_region_index`: the index of the mapped task's region inside its *parent* region. That equals the loop iteration only when the parent region is the loop itself.
If there's another mapped level in between, for example:
```
loop L
└── mapped task group G (.expand)
└── mapped task t (.expand)
```
then `t`'s region's parent is G's region, so `parent_region_index` is G's map index, not L's iteration. In iteration 0 of the loop, `t` could report `loop_index=1` just because it sits in G's second expansion.
I haven't checked whether Airflow allows this nesting. If it doesn't, a short comment stating that assumption would help. If it does, this probably needs to walk up the parent regions until it reaches the loop's region.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1135,11 +1193,44 @@ 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")
+ # 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}
Review Comment:
**`loop_decision.reason` is emitted as `null`, but the schema says it's a string**
`reason` comes from what the loop's `until` callable returns. If `until` returns a plain `True`/`False`, or the loop has no `until` and simply stops at `max_iterations`, `decision.reason` is `None`, and the gate task's COMPLETE event contains:
```json
"loop_decision": {"stop": true, "reason": null}
```
`AirflowRunFacet.json` declares `reason` as `{"type": "string"}`, so a consumer that validates events against the schema will flag or reject this event. Only adding `reason` when it isn't `None` would match the schema, and the example Dag docs, which describe this case as having "no reason".
##########
providers/openlineage/tests/system/openlineage/operator.py:
##########
@@ -351,33 +479,44 @@ 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)
Review Comment:
**One failed template match produces a second, misleading runId error**
In this loop, an event key goes into `events_by_key` only if `_check_key` passes. If it raises, the key is recorded in `errors` and left out of `events_by_key`. Then `check_run_id_integrity` runs over the partial `events_by_key` anyway.
Example: the test expects `dag.task.event.start` and `dag.task.event.complete`.
1. The COMPLETE event exists but doesn't match its template (say a facet value differs), so it goes to `errors`.
2. `events_by_key` now holds only the START for that runId.
3. The integrity check sees a runId with a START and no COMPLETE, and adds: `runId X of TASK dag.task should have one START and one COMPLETE/FAIL`.
The person debugging the test now sees two failures. The second one points to a broken runId lifecycle that doesn't exist, which sends them in the wrong direction. Skipping the integrity check when template errors already exist (or skipping only the runIds those keys touched) would keep the output focused on the real problem.
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -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)
Review Comment:
**The rule for choosing a task's runId is now copied in six places**
Before, `get_regional_task_instance_run_id` decided between "use the task instance UUID (in a loop)" and "use the computed runId from dag/task/try/date/map_index (not in a loop)". With it removed, the same `if position.in_loop: str(ti.id) else build_task_instance_run_id(...)` block is written out separately in `on_task_instance_running`, `on_task_instance_success`, `on_task_instance_failed`, and `on_task_instance_skipped`. The manual state change path and `macros.py` (the `lineage_run_id` macro) have it too, so that's six copies.
These copies must stay identical. If someone later changes how the runId is derived (e.g. includes `region_id`) and misses one, START and COMPLETE for the same attempt get different runIds. That's exactly what the new integrity check in the test operator is designed to catch, but only for Dags that the system tests cover. Keeping one helper, e.g. `get_task_instance_run_id(ti, position, logical_date)`, that all six call sites use would make this impossible to get out of sync.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1135,11 +1193,44 @@ 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)
Review Comment:
**The loop iteration lookup, and its warning, run twice per event**
The listener already computes the task instance's position (including the loop iteration) with `get_task_instance_position`, and passes it in here as `self._position`. This line then calls `_add_runtime_loop_iteration` on it a second time.
Usually that's a no-op, because `loop_index` is already set. But when the first attempt couldn't find the iteration (the `loop` key is missing from the template context, or reading the context raised), `loop_index` is still `None`, so the second call:
- calls `task_instance.get_template_context()` again. Per the comment in `_add_runtime_loop_iteration`, that re-validates the Dag params each time, and `get_dag_run_dag_and_task_from_ti` had already fetched the context once, so that's a third time;
- if it fails again, logs "OpenLineage failed to read the loop iteration..." a second time for the same event.
Using a position passed in by the caller as-is, i.e. only calling `_add_runtime_loop_iteration` when `self._position` is `None`, would avoid the repeat work and the duplicate warning.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1123,9 +1123,67 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
Review Comment:
**The same loop/region explanation is written out in several places**
This docstring grew to about 60 lines, and much of it (what a region is, how clearing a loop forks it, how try numbers relate to iterations) is also explained in:
- the `TaskInstancePosition` docstring,
- the `region_id` / `region_index` / `try_number` descriptions in `AirflowRunFacet.json`,
- the example Dag docstring.
Airflow's `AGENTS.md` asks to keep explanations "short, precise and at the source of truth" and not repeat them at each call site. With four copies, a future behavior change will likely update one or two and leave the rest outdated. I'd suggest keeping the full explanation in one place (the facet schema seems the natural one, since consumers read it) and having the others refer to it.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -2049,6 +2224,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")
Review Comment:
**These new loop-group fields aren't documented in any schema**
For loop groups, this adds `gate_task_id`, `has_until`, `max_iterations` and `terminal_task_id` to the task group output. That output ends up in the run facet's `task_group` and in `AirflowJobFacet.taskGroups`.
The `taskInstance` additions in this PR (`loop_index`, `loop_group_id`, etc.) were all added to `AirflowRunFacet.json` with descriptions. These four weren't: neither the `task_group` definition in `AirflowRunFacet.json` nor `AirflowJobFacet.json` mentions them. Validation still passes because those objects allow `additionalProperties`. But anyone reading the schema, or generating types or docs from it, won't know these fields exist or what they mean (e.g. what the "gate" and "terminal" tasks are). Could they get schema entries alongside the `taskInstance` ones?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4277074953
##########
providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json:
##########
@@ -196,6 +196,14 @@
"map_index": {
"type": "integer"
},
+ "loop_index": {
+ "description": "Zero-based iteration of the loop the task runs in. Absent outside a loop.",
+ "type": "integer"
+ },
+ "loop_max_iterations": {
+ "description": "Maximum number of iterations of the loop the task runs in. Absent outside a loop.",
+ "type": "integer"
+ },
Review Comment:
Just having `loop_index` will mean this needs to change if/when nested loops are allowed in Airflow.
It also doesn't support mapped tasks inside a loop, which is supported today. This is why Airflow has a "dynamic_region" concept.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315007454
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -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)
Review Comment:
Fixed in 7441d2d: one `get_task_instance_run_id` helper, used by the four worker hooks, the manual state change path and the `lineage_run_id` macro.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4360121186
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -2042,17 +2391,49 @@ def _get_task_groups_details(dag: DAG | SerializedDAG, edge_map: dict[str, tuple
result = {}
for tg_id, tg in dag.task_group_dict.items():
task_edges, group_edges = edge_map.get(tg_id, ({}, {}))
+ is_mapped = _is_mapped_task_group(tg)
result[tg_id] = {
"parent_group": tg.parent_group.group_id if tg.parent_group else None,
+ # Independent, so a group could be both
+ "is_mapped": is_mapped,
+ "is_loop": hasattr(tg, "gate_task_id"),
Review Comment:
Done: `is_loop`, and the loop fields with it, now come from the same `isinstance` check as `_find_enclosing_loop`, with a test for a group that only looks like a loop.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1542,56 +1682,228 @@ def _require_task_instance_position(
return position
+def _add_runtime_loop_iteration(
+ task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition
+) -> TaskInstancePosition:
+ if isinstance(task_instance, TaskInstance):
+ return position
+ # 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:
+ return position
+ return attrs.evolve(
+ position,
+ loop_index=loop.index,
+ 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_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))
+
+
+@attrs.define
+class _TaskCoordinate:
+ """Where Airflow stores a task instance, as its ``TaskCoordinateResolver`` reads it."""
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ region_id: UUID
+ region_index: int
+ dag_version_id: UUID | None
+
+
+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 on."""
+ if session is None or not AIRFLOW_V_3_4_PLUS:
+ return None
+ region_id = task_instance.region_id
+ if not isinstance(region_id, UUID) or region_id.int == 0:
+ return None
+ dag = getattr(getattr(task_instance, "task", None), "dag", None)
+ if (dag_version_id := getattr(dag, "dag_version_id", None)) is None:
+ return None
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ # Placed with the Dag version the task instance holds, which its task facet describes too, even when it pins
+ # another one: the resolver would load and deserialize that version for every task instance in the scheduler
+ # or API server process, while the region rows it reads don't depend on the version.
+ coordinate = _TaskCoordinate(
+ dag_id=task_instance.dag_id,
+ run_id=task_instance.run_id,
+ task_id=task_instance.task_id,
+ region_id=region_id,
+ region_index=task_instance.region_index,
+ dag_version_id=dag_version_id,
+ )
+ 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(
+ coordinate.dag_id, coordinate.run_id, coordinate.task_id, dag_version_id=dag_version_id
+ )
+ loop = resolver.get_loop_context(coordinate, task=task)
+ expanded_ti_count = resolver.get_expanded_ti_count(coordinate, 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_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+ 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]:
"""
- 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:
+ # 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
+ 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).where(
- DynamicRegion.id.in_({task_instance.region_id for task_instance in unplaced})
+ parent = aliased(DynamicRegion)
Review Comment:
Dropped. The region-row fallback is back to main's `in_loop` and `map_index`; the parent join, the second pass and their tests are gone.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1542,56 +1682,228 @@ def _require_task_instance_position(
return position
+def _add_runtime_loop_iteration(
+ task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition
+) -> TaskInstancePosition:
+ if isinstance(task_instance, TaskInstance):
+ return position
+ # 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:
+ return position
+ return attrs.evolve(
+ position,
+ loop_index=loop.index,
+ 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_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))
+
+
+@attrs.define
+class _TaskCoordinate:
+ """Where Airflow stores a task instance, as its ``TaskCoordinateResolver`` reads it."""
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ region_id: UUID
+ region_index: int
+ dag_version_id: UUID | None
+
+
+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 on."""
+ if session is None or not AIRFLOW_V_3_4_PLUS:
+ return None
+ region_id = task_instance.region_id
+ if not isinstance(region_id, UUID) or region_id.int == 0:
+ return None
+ dag = getattr(getattr(task_instance, "task", None), "dag", None)
+ if (dag_version_id := getattr(dag, "dag_version_id", None)) is None:
+ return None
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ # Placed with the Dag version the task instance holds, which its task facet describes too, even when it pins
+ # another one: the resolver would load and deserialize that version for every task instance in the scheduler
+ # or API server process, while the region rows it reads don't depend on the version.
+ coordinate = _TaskCoordinate(
+ dag_id=task_instance.dag_id,
+ run_id=task_instance.run_id,
+ task_id=task_instance.task_id,
+ region_id=region_id,
+ region_index=task_instance.region_index,
+ dag_version_id=dag_version_id,
+ )
+ try:
+ # The savepoint keeps the caller's transaction usable after a failure
+ with session.begin_nested():
Review Comment:
Placing a scheduler-sent FAIL is query-free again when the task is attached: no resolver, savepoint or `count(*)`. Those events leave out `expanded_ti_count` and the loop run fields, which the schema now says. I also dropped `taskInstance.note`, a lazy load per failure on the scheduler's session that never reached the manual state changes it was added for.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
Review Comment:
Fixed in #75458, for 3.4.0, no need to fix here now.
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -3924,38 +4273,901 @@ 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(),
+ dag_id="run_id_dag",
+ task_id="body.work",
+ try_number=1,
+ region_id=uuid4(),
+ map_index=-1,
+ task=dag.get_task("body.work"),
+ )
+ logical_date = timezone.datetime(2024, 1, 1)
+
+ run_id = get_task_instance_run_id(ti, logical_date=logical_date)
+
+ 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
+
+
+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")
+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
+
+ 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"), loop=None
+ )
- assert get_regional_task_instance_run_id(ti) == (str(ti.id) if in_loop else None)
+ 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="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,
+ ),
+ )
+
+ assert get_task_instance_position(ti) == TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ expanded_ti_count=2,
+ loop_index=2,
+ 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,
+ )
+
+
+@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4")
+@pytest.mark.parametrize("holds_another_version", [False, True])
+def test_stored_task_instance_is_placed_as_airflow_places_it_for_the_worker(
+ dag_maker, session, holds_another_version
+):
+ 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}
+ # As the scheduler's Dag bag sets it on the Dag it loads for the run, which can be a later version than the
+ # one the task instances pin
+ tis["body.plain", 1].task.dag.dag_version_id = (
+ uuid7() if holds_another_version else 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()
+ 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_trajectory": 1, "loop_group_id": "body"}
+
+ # 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_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"] == {
+ "index": 2,
+ "region_id": str(fork.id),
+ "trajectory": 2,
+ "forked_from_region_id": str(loop_region.id),
+ "resumes_from_index": 2,
+ }
@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Dag versions require Airflow 3")
+@pytest.mark.parametrize("load_fails", [False, True])
+@patch("airflow.providers.openlineage.utils.utils.log", autospec=True)
+def test_stored_task_instance_info_reads_its_dag_version_without_risking_the_event(
Review Comment:
The savepoints are gone, and that test with them. Added a db test through `on_task_instance_failed` with a real loop task instance, asserting runId == `taskInstance.id` == `ti.id` and `loop.index`; passing `session=None` in the listener fails it.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1542,56 +1682,228 @@ def _require_task_instance_position(
return position
+def _add_runtime_loop_iteration(
+ task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition
+) -> TaskInstancePosition:
+ if isinstance(task_instance, TaskInstance):
+ return position
+ # 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:
+ return position
+ return attrs.evolve(
+ position,
+ loop_index=loop.index,
+ 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_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))
+
+
+@attrs.define
+class _TaskCoordinate:
+ """Where Airflow stores a task instance, as its ``TaskCoordinateResolver`` reads it."""
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ region_id: UUID
+ region_index: int
+ dag_version_id: UUID | None
+
+
+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 on."""
+ if session is None or not AIRFLOW_V_3_4_PLUS:
+ return None
+ region_id = task_instance.region_id
+ if not isinstance(region_id, UUID) or region_id.int == 0:
+ return None
+ dag = getattr(getattr(task_instance, "task", None), "dag", None)
+ if (dag_version_id := getattr(dag, "dag_version_id", None)) is None:
+ return None
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ # Placed with the Dag version the task instance holds, which its task facet describes too, even when it pins
+ # another one: the resolver would load and deserialize that version for every task instance in the scheduler
+ # or API server process, while the region rows it reads don't depend on the version.
Review Comment:
Fixed differently: the resolver path is gone, and a task instance whose held Dag version isn't its own is placed from its region rows, which don't depend on the version (one query, only then). The worker also counts an attempt as looped when Airflow sent it a loop context, for a worker that parsed a newer file. A db test through `on_task_instance_failed` uses two versions that differ in loop membership.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
- "dag_bundle_version": lambda ti: (
- ti.bundle_instance.version if hasattr(ti, "bundle_instance") else None
+ "dag_version_id": lambda ti: (
+ str(dag_version_id) if (dag_version_id := getattr(ti, "dag_version_id", None)) else None
),
- "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None,
+ "dag_bundle_version": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "version"),
+ "dag_bundle_name": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "name"),
Review Comment:
The Dag version isn't read at all anymore. A stored task instance's bundle name and version come from the Dag and the Dag run the event is built from, the same sources the worker's workload uses, without a query.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1461,6 +1556,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.
Review Comment:
Leaving it for a follow-up PR: it needs a new per-task-instance list in `airflowState` (map index, loop iteration, state, duration), which I'd rather design on its own. I'll link it in the TODO.
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -3924,38 +4273,901 @@ 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(),
+ dag_id="run_id_dag",
+ task_id="body.work",
+ try_number=1,
+ region_id=uuid4(),
+ map_index=-1,
+ task=dag.get_task("body.work"),
+ )
+ logical_date = timezone.datetime(2024, 1, 1)
+
+ run_id = get_task_instance_run_id(ti, logical_date=logical_date)
+
+ 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
+
+
+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")
+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
+
+ 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"), loop=None
+ )
- assert get_regional_task_instance_run_id(ti) == (str(ti.id) if in_loop else None)
+ 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="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,
+ ),
+ )
+
+ assert get_task_instance_position(ti) == TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ expanded_ti_count=2,
+ loop_index=2,
+ 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,
+ )
+
+
+@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4")
+@pytest.mark.parametrize("holds_another_version", [False, True])
+def test_stored_task_instance_is_placed_as_airflow_places_it_for_the_worker(
+ dag_maker, session, holds_another_version
+):
+ 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}
+ # As the scheduler's Dag bag sets it on the Dag it loads for the run, which can be a later version than the
+ # one the task instances pin
+ tis["body.plain", 1].task.dag.dag_version_id = (
+ uuid7() if holds_another_version else 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()
+ 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_trajectory": 1, "loop_group_id": "body"}
+
+ # 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_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"] == {
+ "index": 2,
+ "region_id": str(fork.id),
+ "trajectory": 2,
+ "forked_from_region_id": str(loop_region.id),
+ "resumes_from_index": 2,
+ }
@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Dag versions require Airflow 3")
+@pytest.mark.parametrize("load_fails", [False, True])
+@patch("airflow.providers.openlineage.utils.utils.log", autospec=True)
+def test_stored_task_instance_info_reads_its_dag_version_without_risking_the_event(
+ mock_log, dag_maker, session, load_fails
+):
+ with dag_maker(session=session):
+ PythonOperator(task_id="work", python_callable=list)
+ ti = dag_maker.create_dagrun().task_instances[0]
+ session.expire(ti, ["dag_version"])
+
+ if load_fails:
+ with patch.object(
+ TaskInstance,
+ "dag_version",
+ new_callable=PropertyMock,
+ side_effect=RuntimeError("Lost connection"),
+ ):
+ info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1)))
+ else:
+ info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1)))
+
+ expected_bundle_name = None if load_fails else "dag_maker"
+ assert (info["dag_bundle_name"], info["dag_version_id"]) == (expected_bundle_name, str(ti.dag_version_id))
+ assert mock_log.warning.call_count == (
+ 2 if load_fails else 0
+ ) # Once for the bundle name, once for its version
+ # The failure was rolled back to its savepoint, so the caller's transaction is still usable
+ assert session.execute(select(literal(1))).scalar() == 1
+
+
+@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4")
+@pytest.mark.parametrize("task_id", ["body.plain", "body.inner.mapped"])
+def test_stored_and_runtime_task_instances_identify_an_execution_alike(dag_maker, session, task_id):
+ """
+ The scheduler's events about a stored task instance must tell everything the worker's tell about it.
+
+ A field missing on one side only would leave consumers without it for that kind of event, while the other
+ side's events made them rely on it.
+ """
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with dag_maker(serialized=True):
+
+ @task_group
+ def body():
+ plain = EmptyOperator(task_id="plain")
+ with TaskGroup("inner"):
+ plain >> PythonOperator.partial(task_id="mapped", python_callable=list).expand(
+ op_kwargs=[{}, {}]
+ )
+
+ create_loop(body, max_iterations=3)
+ run = dag_maker.create_dagrun()
+ stored = next(ti for ti in run.task_instances if ti.task_id == task_id and ti.region_index in (0, 1))
+ stored.task.dag.dag_version_id = stored.dag_version_id
+ stored.state, stored.try_number, stored.max_tries = TaskInstanceState.UP_FOR_RETRY, 1, 2
+ resolver = TaskCoordinateResolver.for_dag(stored.task.dag, session)
+ mapped = stored.task.get_needs_expansion()
+ # What the server sends the worker for this execution, on the task of the Dag the worker parses
+ runtime = SimpleNamespace(
+ id=stored.id,
+ dag_id=stored.dag_id,
+ run_id=stored.run_id,
+ task_id=stored.task_id,
+ try_number=stored.try_number,
+ max_tries=stored.max_tries,
+ state=stored.state,
+ map_index=stored.region_index if mapped else -1,
+ rendered_map_index=None,
+ region_id=stored.region_id,
+ region_index=stored.region_index,
+ dag_version_id=stored.dag_version_id,
+ bundle_instance=SimpleNamespace(
+ name=stored.dag_version.bundle_name, version=stored.dag_version.bundle_version
+ ),
+ loop=resolver.get_loop_context(stored, task=stored.task),
+ expanded_ti_count=resolver.get_expanded_ti_count(stored, task=stored.task),
+ task=dag_maker.dag.get_task(task_id),
+ )
+
+ stored_info = dict(
+ TaskInstanceInfo(stored, position=get_stored_task_instance_position(stored, session=session))
+ )
+ runtime_info = dict(TaskInstanceInfo(runtime, position=get_task_instance_position(runtime)))
+
+ # Only stored task instances have these, which describe the attempt's resources rather than which it is
+ resources = {"duration", "pool", "queued_dttm", "note", "log_url"}
+ assert {k: v for k, v in stored_info.items() if k not in resources} == {
+ k: v for k, v in runtime_info.items() if k not in resources
+ }
+ assert ("expanded_ti_count" in runtime_info) is mapped
+ assert runtime_info["loop"]["trajectory"] == 1
+ identity = ("task_id", "task_group", "loop", "retries", "retry_delay")
+ assert {k: dict(TaskInfo(stored.task))[k] for k in identity} == {
+ k: dict(TaskInfo(runtime.task))[k] for k in identity
+ }
+
+
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4")
+@pytest.mark.parametrize(
+ "dag_has_version", [pytest.param(True, id="airflow_fails"), pytest.param(False, id="no_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_without_airflow(mock_for_dag, dag_has_version):
+ 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
+ dag.dag_version_id = ti.dag_version_id if dag_has_version else None
+ session = MagicMock(spec=Session)
+
+ assert get_stored_task_instance_position(ti, session=session) == TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_index=1, loop_group_id="body"
+ )
+ if dag_has_version:
+ mock_for_dag.assert_called_once_with(dag, session)
+ else:
+ mock_for_dag.assert_not_called()
+ 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
+ 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_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)
+
+
@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_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.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))
+
+ assert info["map_index"] == (1 if mapped else None)
+ if in_loop:
+ assert info["loop"] == {
+ "index": 2,
+ }
+ else:
+ assert "loop" not in info
+
+
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4")
+@pytest.mark.parametrize("nested", [False, True])
+def test_task_info_has_the_loop_around_the_task(nested):
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ def ready():
+ return True
+
+ with DAG("nested_loop_dag") as dag:
+ with TaskGroup("outer"):
+
+ @task_group
+ def body():
+ if nested:
+
+ @task_group
+ def inner():
+ EmptyOperator(task_id="work")
+
+ inner()
+ else:
+ EmptyOperator(task_id="work")
+
+ create_loop(body, max_iterations=3, until=ready)
+ task_id = "outer.body.inner.work" if nested else "outer.body.work"
+
+ info = dict(TaskInfo(dag.get_task(task_id)))
+
+ assert info["loop"] == {
+ "group_id": "outer.body",
+ "max_iterations": 3,
+ "has_until": True,
+ "gate_task_id": dag.task_group_dict["outer.body"].gate_task_id,
+ "terminal_task_id": task_id,
+ }
+ # task_group is the task's own group and never describes the loop; its unique id tells whether it is the loop
+ assert (info["task_group"]["group_id"], info["task_group"]["unique_group_id"]) == (
+ ("inner", "outer.body.inner") if nested else ("body", "outer.body")
+ )
+ assert (info["task_group"]["unique_group_id"] == info["loop"]["group_id"]) is not nested
+ assert "gate_task_id" not in info["task_group"]
+
+
+@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}, False, id="no_until_continues"),
+ 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):
+ 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
+ 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 = GateTaskInstance(id=uuid7(), map_index=-1, _loop_decision=None)
+ 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" 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_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 = [_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)
+
+ assert positions == {
+ ti.id: TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ loop_index=2 if with_session else None,
+ 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" not in info
+
+
+def test_task_instance_info_uses_the_position_it_is_given():
+ # 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,
+ map_index=1,
+ expanded_ti_count=4,
+ loop_index=2,
+ 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["expanded_ti_count"]) == (1, 4)
+ assert info["loop"] == {
+ "index": 2,
+ "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")
+def test_task_group_info_has_no_loop_attributes():
+ 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_attributes = {"max_iterations", "has_until", "terminal_task_id", "gate_task_id"}
+ loop_info = dict(TaskGroupInfo(loop))
+ assert loop_info["group_id"] == "body"
+ assert not loop_attributes & set(loop_info)
+ assert not loop_attributes & set(TaskGroupInfo(plain_group))
+
+
+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_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
+
+ 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", 1): TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_index=1, loop_group_id="body"
+ ),
+ ("body.last", 1): TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_index=1, loop_group_id="body"
+ ),
+ **{
+ (task_id, map_index): TaskInstancePosition(
+ in_loop=True, map_index=map_index, loop_index=1, loop_group_id="body"
+ )
+ 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):
Review Comment:
Added `match=`. With the region parsing back to main's, only one rejected shape matters to our code, a mapped task in a mapped task group, so the test pins that one.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
- "dag_bundle_version": lambda ti: (
- ti.bundle_instance.version if hasattr(ti, "bundle_instance") else None
+ "dag_version_id": lambda ti: (
+ str(dag_version_id) if (dag_version_id := getattr(ti, "dag_version_id", None)) else None
),
- "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None,
+ "dag_bundle_version": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "version"),
+ "dag_bundle_name": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "name"),
}
+ @staticmethod
+ def get_rendered_map_index(ti: TaskInstance | RuntimeTaskInstance) -> str | None:
+ """Return the label a mapped task rendered with ``map_index_template``, as the worker sends it."""
+ if AIRFLOW_V_3_0_PLUS and isinstance(ti, TaskInstance):
+ # Its rendered_map_index falls back to the map index when no label was rendered, the worker's doesn't
+ return safe_getattr(ti, "_rendered_map_index")
+ return getattr(ti, "rendered_map_index", None)
+
+ @staticmethod
+ def get_bundle_info(ti: TaskInstance | RuntimeTaskInstance, key: str) -> str | None:
+ """Return the name or version of the Dag bundle the attempt runs from."""
+ if (bundle := getattr(ti, "bundle_instance", None)) is not None: # A runtime task instance
+ return getattr(bundle, key, None)
+ if not isinstance(ti, TaskInstance):
+ return None
+ # Only for a FAIL the scheduler sends itself, e.g. for a missed heartbeat: UI and API state changes build no
+ # airflow facet, and the worker's runtime task instance has its bundle above
+ return getattr(TaskInstanceInfo.get_stored_dag_version(ti), f"bundle_{key}", None)
+
+ @staticmethod
+ def get_stored_dag_version(ti: TaskInstance) -> Any:
+ """Return the Dag version of a stored task instance, or None when it can't be read, as on Airflow 2."""
+ from sqlalchemy import inspect
+
+ try:
+ state: InstanceState = inspect(ti)
+ if "dag_version" not in state.unloaded or state.session is None:
+ # Loaded already, or not in a session to load it from, so reading it doesn't query
+ return safe_getattr(ti, "dag_version")
+ # The event must still be sent without it, and the savepoint keeps the caller's transaction usable
+ with state.session.begin_nested():
Review Comment:
Agreed, removed. Scheduler-sent events no longer read the Dag version, loop context, expansion size or region rows through savepoints: the task instance is placed from its attached task with no query, as on main. Only one without its task, or holding another Dag version, reads its region rows, in one plain query as main did. I didn't use `create_session()` instead: in the scheduler thread it returns the scheduler's own scoped session, and would commit and close it.
##########
providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py:
##########
@@ -0,0 +1,164 @@
+# 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 and one stopped early by `until`
+(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`, `state`, `map_index` and `loop` of every event: `loop.index` is the
+ iteration, counted from 1, and `map_index` the mapped position, and a plain task has no `loop`
+ - task's `loop`, the loop around the task: its `group_id`, `max_iterations`, `has_until`, `gate_task_id`
+ and `terminal_task_id`, also for a task of a task group inside a loop, and absent for the plain task
+ - `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
Review Comment:
Added the invariants to `check_run_id_integrity` (runId is `taskInstance.id` in a loop, `loop.region_id` against `region_id` for mapped and unmapped tasks, one `loop.region_id` per loop run) with unit tests, and trimmed the docstring to what the JSON asserts. A Dag that clears a gate and checks the rerun (trajectory 2, fork fields) will come in a follow-up PR with the manual state change e2e OL dags.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
- "dag_bundle_version": lambda ti: (
- ti.bundle_instance.version if hasattr(ti, "bundle_instance") else None
+ "dag_version_id": lambda ti: (
+ str(dag_version_id) if (dag_version_id := getattr(ti, "dag_version_id", None)) else None
),
- "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None,
+ "dag_bundle_version": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "version"),
+ "dag_bundle_name": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "name"),
}
+ @staticmethod
+ def get_rendered_map_index(ti: TaskInstance | RuntimeTaskInstance) -> str | None:
+ """Return the label a mapped task rendered with ``map_index_template``, as the worker sends it."""
+ if AIRFLOW_V_3_0_PLUS and isinstance(ti, TaskInstance):
+ # Its rendered_map_index falls back to the map index when no label was rendered, the worker's doesn't
+ return safe_getattr(ti, "_rendered_map_index")
+ return getattr(ti, "rendered_map_index", None)
+
+ @staticmethod
+ def get_bundle_info(ti: TaskInstance | RuntimeTaskInstance, key: str) -> str | None:
+ """Return the name or version of the Dag bundle the attempt runs from."""
+ if (bundle := getattr(ti, "bundle_instance", None)) is not None: # A runtime task instance
+ return getattr(bundle, key, None)
+ if not isinstance(ti, TaskInstance):
+ return None
+ # Only for a FAIL the scheduler sends itself, e.g. for a missed heartbeat: UI and API state changes build no
+ # airflow facet, and the worker's runtime task instance has its bundle above
+ return getattr(TaskInstanceInfo.get_stored_dag_version(ti), f"bundle_{key}", None)
+
+ @staticmethod
+ def get_stored_dag_version(ti: TaskInstance) -> Any:
+ """Return the Dag version of a stored task instance, or None when it can't be read, as on Airflow 2."""
+ from sqlalchemy import inspect
+
+ try:
+ state: InstanceState = inspect(ti)
+ if "dag_version" not in state.unloaded or state.session is None:
+ # Loaded already, or not in a session to load it from, so reading it doesn't query
+ return safe_getattr(ti, "dag_version")
+ # The event must still be sent without it, and the savepoint keeps the caller's transaction usable
+ with state.session.begin_nested():
+ return ti.dag_version
+ except Exception:
+ log.warning(
+ "OpenLineage failed to read the Dag version of task instance %s", getattr(ti, "id", None)
+ )
+ log.debug("Exception details:", exc_info=True)
+ return 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 or
+ # with get_task_instance_position; it is then used as is.
+ self._position = position
+ super().__init__(obj)
+
def _extend_fields(self) -> None:
- map_index = get_task_instance_map_index(self.obj)
+ 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
+ self.rendered_map_index = self.get_rendered_map_index(self.obj) if map_index != -1 else None
self._fields.extend(("map_index", "rendered_map_index"))
+ 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:
Review Comment:
Done: a stored task instance's `id` is only emitted from 3.4 on, a runtime one on every 3.x, and the schema says so.
##########
providers/openlineage/tests/system/openlineage/operator.py:
##########
@@ -351,33 +479,47 @@ 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 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)},
+ ):
+ message = "\n".join(f" {problem}" for problem in problems)
+ if self.fail_fast:
+ raise ValueError(f"runId integrity check failed:\n{message}")
+ errors["<runId integrity>"] = f"\n{message}"
finally:
if self.clear_variables:
Review Comment:
Added a parametrized `execute()` test (default, `fail_fast`, `check_run_ids=False`, declared repeated STARTs). Each mutation you listed now fails it, and the mocks are autospec'd.
##########
providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json:
##########
@@ -196,6 +225,73 @@
"map_index": {
"type": "integer"
},
+ "id": {
+ "description": "Id of the task instance attempt (Airflow 3+).",
+ "type": "string"
+ },
+ "state": {
+ "type": "string"
+ },
+ "max_tries": {
+ "type": "integer"
+ },
+ "rendered_map_index": {
+ "type": "string"
+ },
+ "expanded_ti_count": {
+ "description": "Number of instances in the mapped task's expansion (Airflow 3.4+).",
+ "type": "integer"
+ },
+ "region_id": {
+ "description": "Region of the loop iteration or mapped expansion (Airflow 3.4+).",
+ "type": "string"
+ },
+ "region_index": {
+ "description": "Position in the region (Airflow 3.4+).",
+ "type": "integer"
+ },
Review Comment:
Described each field in the schema, including what `region_index` stands for, `loop.region_id` after a clear, `trajectory`, `resumes_from_index` and `expanded_ti_count` per pass. I kept `region_id` and `region_index` on purpose: together they are what Airflow's REST API and UI select a task instance in a region by, and `region_index` is the only iteration a scheduler-sent FAIL placed from region rows has.
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -2011,6 +2126,201 @@ 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
Review Comment:
Moved the imports to module level under `if AIRFLOW_V_3_4_PLUS:`, dropped the re-imports, added an `_add_loop()` helper for the repeated loop Dag, and removed the overlapping placement tests with the code they covered.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
- "dag_bundle_version": lambda ti: (
- ti.bundle_instance.version if hasattr(ti, "bundle_instance") else None
+ "dag_version_id": lambda ti: (
+ str(dag_version_id) if (dag_version_id := getattr(ti, "dag_version_id", None)) else None
),
- "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None,
+ "dag_bundle_version": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "version"),
+ "dag_bundle_name": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "name"),
}
+ @staticmethod
+ def get_rendered_map_index(ti: TaskInstance | RuntimeTaskInstance) -> str | None:
+ """Return the label a mapped task rendered with ``map_index_template``, as the worker sends it."""
+ if AIRFLOW_V_3_0_PLUS and isinstance(ti, TaskInstance):
+ # Its rendered_map_index falls back to the map index when no label was rendered, the worker's doesn't
+ return safe_getattr(ti, "_rendered_map_index")
+ return getattr(ti, "rendered_map_index", None)
+
+ @staticmethod
+ def get_bundle_info(ti: TaskInstance | RuntimeTaskInstance, key: str) -> str | None:
+ """Return the name or version of the Dag bundle the attempt runs from."""
+ if (bundle := getattr(ti, "bundle_instance", None)) is not None: # A runtime task instance
+ return getattr(bundle, key, None)
+ if not isinstance(ti, TaskInstance):
+ return None
+ # Only for a FAIL the scheduler sends itself, e.g. for a missed heartbeat: UI and API state changes build no
+ # airflow facet, and the worker's runtime task instance has its bundle above
+ return getattr(TaskInstanceInfo.get_stored_dag_version(ti), f"bundle_{key}", None)
+
+ @staticmethod
+ def get_stored_dag_version(ti: TaskInstance) -> Any:
+ """Return the Dag version of a stored task instance, or None when it can't be read, as on Airflow 2."""
+ from sqlalchemy import inspect
Review Comment:
`inspect`, `aliased` and `TaskCoordinateResolver` went with their code, `MappedTaskGroup` is a module-level compat import, and the 3.4-only loop classes stay inline with a comment saying why.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1542,56 +1682,228 @@ def _require_task_instance_position(
return position
+def _add_runtime_loop_iteration(
+ task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition
+) -> TaskInstancePosition:
+ if isinstance(task_instance, TaskInstance):
+ return position
+ # 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:
+ return position
+ return attrs.evolve(
+ position,
+ loop_index=loop.index,
+ 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_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))
+
+
+@attrs.define
+class _TaskCoordinate:
+ """Where Airflow stores a task instance, as its ``TaskCoordinateResolver`` reads it."""
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ region_id: UUID
+ region_index: int
+ dag_version_id: UUID | None
+
+
+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 on."""
+ if session is None or not AIRFLOW_V_3_4_PLUS:
+ return None
+ region_id = task_instance.region_id
+ if not isinstance(region_id, UUID) or region_id.int == 0:
+ return None
+ dag = getattr(getattr(task_instance, "task", None), "dag", None)
+ if (dag_version_id := getattr(dag, "dag_version_id", None)) is None:
+ return None
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ # Placed with the Dag version the task instance holds, which its task facet describes too, even when it pins
+ # another one: the resolver would load and deserialize that version for every task instance in the scheduler
+ # or API server process, while the region rows it reads don't depend on the version.
+ coordinate = _TaskCoordinate(
+ dag_id=task_instance.dag_id,
+ run_id=task_instance.run_id,
+ task_id=task_instance.task_id,
+ region_id=region_id,
+ region_index=task_instance.region_index,
+ dag_version_id=dag_version_id,
+ )
+ 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(
+ coordinate.dag_id, coordinate.run_id, coordinate.task_id, dag_version_id=dag_version_id
+ )
+ loop = resolver.get_loop_context(coordinate, task=task)
+ expanded_ti_count = resolver.get_expanded_ti_count(coordinate, 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_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+ 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]:
"""
- 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:
+ # 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
+ 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).where(
- DynamicRegion.id.in_({task_instance.region_id for task_instance in unplaced})
+ 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"),
Review Comment:
Agreed. Two remain: placing from the task, and main's region fallback (`in_loop`, `map_index`). The schema says which fields scheduler-sent events lack. Core exposing loop and expansion accessors on a stored task instance would be the clean way to bring them back later.
##########
providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py:
##########
@@ -0,0 +1,73 @@
+# 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
+ - task's `retries` and `retry_delay` on every event
+ - taskInstance's `state`: up_for_retry on the FAIL of the first try, success on the COMPLETE of the retry
+ - each try has its own `runId` (the `runId` check of OpenLineageTestOperator is enabled by default)
+ - every task event's identity: taskInstance's `id`, `dag_version_id` (Airflow 3.1+) and bundle, the absence
+ of fields that don't apply, task's `task_id`, `trigger_rule`, `retries` and `retry_delay`, the Dag and Dag run,
+ and the parent run, plus each task in the Dag's job facet
+"""
+
+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)},
Review Comment:
Done: `retries` and `retry_delay` are on `flaky_task` only. `OpenLineageTestOperator` now also defaults to `retries=0` itself, so a Dag's `default_args` can't make a check retry; the other Dags already set 0.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #74752:
URL: https://github.com/apache/airflow/pull/74752#issuecomment-6232789143
<!-- uv-lock-rebase-notice -->
`uv.lock` on `main` just moved via [#74899](https://redirect.github.com/apache/airflow/pull/74899) ("Cap `pydantic-ai-skills` below 2 for the Common AI skills extra"), commit [`c7cec68`](https://redirect.github.com/apache/airflow/commit/c7cec687bb6e42f2ed12ed506b4cabf84cff30be) and this PR currently conflicts.
Quickest fix:
```bash
git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease
```
_Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future `uv.lock` bumps._
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315005637
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
+ )
)
)
}
for task_instance in unplaced:
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,
Review Comment:
I checked this against core: it never creates a region for a mapped task group. Every task that needs expansion, a task in a mapped task group included, gets its own region with `node_id == task_id` (core's log addressing classifies regions with the same rule), so `g.t` takes the first branch with its map index. 7441d2d makes the parsing check the shapes it expects (a loop's region is top-level, a mapped task's region sits directly in its loop's) and places anything else as a loop without its iteration, with a warning. `test_region_shapes_of_airflow_are_the_ones_openlineage_places` builds a real Dag run with these shapes, so a change in core fails the provider's tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4325728407
##########
providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py:
##########
@@ -0,0 +1,151 @@
+# 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 and one stopped early by `until`
+(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 `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
+ 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
+ - 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 LoopDecision, 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:
Review Comment:
In #74773
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4295839313
##########
providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json:
##########
@@ -196,6 +196,14 @@
"map_index": {
"type": "integer"
},
+ "loop_index": {
+ "description": "Zero-based iteration of the loop the task runs in. Absent outside a loop.",
+ "type": "integer"
+ },
+ "loop_max_iterations": {
+ "description": "Maximum number of iterations of the loop the task runs in. Absent outside a loop.",
+ "type": "integer"
+ },
Review Comment:
Thanks, extended this with region_id. Also looking into exposing a bit more data from core/sdk so that OL can show it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315006152
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
+ )
)
)
}
for task_instance in unplaced:
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
Review Comment:
Airflow doesn't allow that nesting today: the SDK rejects expanding a task inside a mapped task group, and inside a loop each task of a mapped task group gets its region directly under the loop's. 7441d2d states this, uses `parent_region_index` only when the parent region is the task's own loop and top-level, and pins the SDK restrictions in `test_region_shapes_stay_limited_by_the_task_sdk`.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1135,11 +1193,44 @@ 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")
+ # 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}
Review Comment:
Changed in 7441d2d: `reason` is only added when set. (On the wire it was already absent, as event serialization drops `None` values.)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4325815647
##########
providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py:
##########
@@ -0,0 +1,151 @@
+# 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 and one stopped early by `until`
+(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 `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
+ 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
+ - 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 LoopDecision, 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:
Review Comment:
Yeah, that's a good call, was confused in the UI. Will adjust the test (here or in the followup PR if there are no other nits in this PR)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4314212153
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1489,28 +1587,38 @@ 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
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
Review Comment:
`_is_task_in_loop` used to fall back to `return group is not None`; this now falls back to `return None`. A task with a real region whose group chain isn't a recognized `TaskGroup`/`SerializedTaskGroup` (some deserialized/proxied group type) is therefore treated as *not* in a loop.
The consequence is worse than before: it gets the deterministic runId from (dag, task, try, date, map_index), which is identical for every iteration, so all iterations share one runId and their lineage merges. Erring on the "in loop" side (UUID runId) seems safer here.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
Review Comment:
This adds a synchronous `DynamicRegion` query for mapped tasks in a loop (`without_iteration`), only to fill the descriptive `loop_index`. On the manual state change path this runs on the scheduler (`handle_failure`, mark-failed from UI), where such a TI was previously placed with zero queries.
If the query raises, the outer `try` in `_on_task_instance_manual_state_change` swallows it and no FAIL/COMPLETE is emitted at all — which contradicts the rule stated in `_add_runtime_loop_iteration` that reading the iteration must not cost the event its run id. Could the `loop_index` lookup be isolated so its failure only drops `loop_index`?
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
+ )
)
)
}
for task_instance in unplaced:
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,
Review Comment:
This branch assumes a found region with `node_id != task_id` is a loop's region. But `DynamicRegion` also stores mapped task group expansions.
E.g. task `g.t` in a mapped task group `g` that is not in any loop, placed without its task (asset-event source / manual-state path): region `node_id='g' != 'g.t'`, so it becomes `in_loop=True, map_index=-1, loop_index=<its map index>, loop_group_id='g'`. The event then reports a loop iteration that doesn't exist, drops the real `map_index`, and uses a UUID runId that doesn't match the deterministic one its worker events use.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
+ )
)
)
}
for task_instance in unplaced:
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
Review Comment:
`region.parent_region_index` of a mapped task's own region is the loop iteration only if the region's direct parent is the loop region. If the mapped task sits inside another mapped region within the loop (e.g. loop → mapped task group → mapped task), this is the enclosing group's map index, not the iteration. I haven't checked whether that nesting is allowed — if it isn't, maybe worth a comment/assertion stating the assumption.
##########
providers/openlineage/tests/system/openlineage/operator.py:
##########
@@ -351,33 +479,44 @@ 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)
Review Comment:
`events_by_key` only collects keys that passed `_check_key`, but `check_run_id_integrity` still runs on that partial set. So if e.g. `dag.task.event.complete` fails template matching, integrity additionally reports "runId X ... should have one START and one COMPLETE/FAIL" — every real failure produces a second, misleading one. Maybe skip the integrity check (or the affected runIds) when key checks already failed?
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -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)
Review Comment:
With `get_regional_task_instance_run_id` removed, the `if position.in_loop: str(ti.id) else build_task_instance_run_id(...)` rule is now repeated in on_running/on_success/on_failure/on_skipped, plus the manual state change path and `macros.py` — six copies. Any divergence makes START and COMPLETE carry different runIds (exactly what the new integrity check hunts for). A single helper taking `(ti, position, logical_date)` would keep the rule in one place.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1135,11 +1193,44 @@ 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)
Review Comment:
When the listener already passed a position, this still runs `_add_runtime_loop_iteration` on it; whenever `loop_index` is still `None` that calls `get_template_context()` again (re-running `process_params`) and logs the "failed to read the loop iteration" warning a second time. Could a caller-supplied position be used as-is, or the already-fetched context be passed in?
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1135,11 +1193,44 @@ 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")
+ # 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}
Review Comment:
`reason` is always included and is `None` when `until` returns a plain bool (or there is no `until`), so the gate emits `"loop_decision": {"stop": true, "reason": null}`. The schema types `reason` as `string`, so strict validators will reject this. Omitting the key when there's no reason would match the schema (and the example docs saying "no reason").
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1123,9 +1123,67 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
Review Comment:
This docstring is ~60 lines, and the trajectory/clear/region semantics are also explained in `TaskInstancePosition`, the `AirflowRunFacet.json` descriptions and the example Dag docstring. AGENTS.md asks to keep explanations "short, precise and at the source of truth" — could this be trimmed to one canonical place so the copies don't drift?
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -2049,6 +2224,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")
Review Comment:
These loop-group fields (`gate_task_id`, `has_until`, `max_iterations`, `terminal_task_id`) are now emitted in `task_group` and `AirflowJobFacet.taskGroups`, but neither the `task_group` definition in `AirflowRunFacet.json` nor `AirflowJobFacet.json` describes them — only the `taskInstance` additions were documented. They pass via `additionalProperties`, but consumers generating types/docs won't see them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315006986
##########
providers/openlineage/tests/system/openlineage/operator.py:
##########
@@ -351,33 +479,44 @@ 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)
Review Comment:
Agreed, fixed in 7441d2d: the runId check is skipped, with a log line, once any key failed its template.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -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)
Review Comment:
Fixed in 7441d2d: one `get_task_instance_run_id` helper, used by the four worker hooks, the manual state change path and the `lineage_run_id` macro.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1135,11 +1193,44 @@ 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)
Review Comment:
Fixed in 7441d2d: a position passed in is used as is; the facet only looks up the iteration itself when no position was given.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1123,9 +1123,67 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
Review Comment:
Trimmed in 7441d2d: `TaskInstancePosition` now only describes its own fields and points to `TaskInstanceInfo`, which keeps the full explanation for readers of the code, while the schema keeps short per-field descriptions for consumers.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -2049,6 +2224,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")
Review Comment:
Added in 7441d2d, to the `task_group` definition in `AirflowRunFacet.json` and to `taskGroups` in `AirflowJobFacet.json`.
---
Drafted-by: Claude Code (Opus 5.5) (no human review before posting)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4325722508
##########
providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py:
##########
@@ -0,0 +1,151 @@
+# 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 and one stopped early by `until`
+(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 `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
+ 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
+ - 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 LoopDecision, 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:
Review Comment:
Oh, Brent is changing this to be 1-based, not 0-based!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix task loop information and tests for OpenLineage provider [airflow]
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4350955107
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
- "dag_bundle_version": lambda ti: (
- ti.bundle_instance.version if hasattr(ti, "bundle_instance") else None
+ "dag_version_id": lambda ti: (
+ str(dag_version_id) if (dag_version_id := getattr(ti, "dag_version_id", None)) else None
),
- "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None,
+ "dag_bundle_version": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "version"),
+ "dag_bundle_name": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "name"),
}
+ @staticmethod
+ def get_rendered_map_index(ti: TaskInstance | RuntimeTaskInstance) -> str | None:
+ """Return the label a mapped task rendered with ``map_index_template``, as the worker sends it."""
+ if AIRFLOW_V_3_0_PLUS and isinstance(ti, TaskInstance):
+ # Its rendered_map_index falls back to the map index when no label was rendered, the worker's doesn't
+ return safe_getattr(ti, "_rendered_map_index")
+ return getattr(ti, "rendered_map_index", None)
+
+ @staticmethod
+ def get_bundle_info(ti: TaskInstance | RuntimeTaskInstance, key: str) -> str | None:
+ """Return the name or version of the Dag bundle the attempt runs from."""
+ if (bundle := getattr(ti, "bundle_instance", None)) is not None: # A runtime task instance
+ return getattr(bundle, key, None)
+ if not isinstance(ti, TaskInstance):
+ return None
+ # Only for a FAIL the scheduler sends itself, e.g. for a missed heartbeat: UI and API state changes build no
+ # airflow facet, and the worker's runtime task instance has its bundle above
+ return getattr(TaskInstanceInfo.get_stored_dag_version(ti), f"bundle_{key}", None)
+
+ @staticmethod
+ def get_stored_dag_version(ti: TaskInstance) -> Any:
+ """Return the Dag version of a stored task instance, or None when it can't be read, as on Airflow 2."""
+ from sqlalchemy import inspect
+
+ try:
+ state: InstanceState = inspect(ti)
+ if "dag_version" not in state.unloaded or state.session is None:
+ # Loaded already, or not in a session to load it from, so reading it doesn't query
+ return safe_getattr(ti, "dag_version")
+ # The event must still be sent without it, and the savepoint keeps the caller's transaction usable
+ with state.session.begin_nested():
Review Comment:
When the scheduler fails a task instance itself (e.g. its worker died), `TaskInstance.handle_failure` calls `notify_failure` → `on_task_instance_failed` **in the middle of its unit of work**: it has just `session.add`-ed a `Log` row and set `ti.state` / `ti.end_date`, and nothing is flushed yet. This listener then reads the Dag version, loop context and regions on that same session (`object_session(ti)`), each inside `session.begin_nested()`.
`begin_nested()` always flushes the session before it emits `SAVEPOINT`, whether or not autoflush is on (in SQLAlchemy 2.0 this is `SessionTransaction._take_snapshot` → `session.flush()`). Airflow's sessions use `autoflush=False`, so before this PR the provider's reads never wrote anything. Now building the facet forces the scheduler's half-finished changes to the DB, before `prepare_db_for_next_try` runs.
If that flush fails (a lock timeout, an integrity error on the pending `Log`, …), the error is raised by `begin_nested()` itself, before the savepoint exists. The `except Exception` here logs it as "OpenLineage failed to read the Dag version" (the real error is only at debug level), and the session stays in the "rolled back due to a previous exception during flush" state. The scheduler's next statement then fails with `PendingRollbackError`, and the cause is hidden in an OpenLineage debug log. So "the savepoint keeps the caller's transaction usable" holds only for errors in the read itself, not for the flush it triggers.
The same pattern is in `_get_core_task_instance_position` and `get_orm_task_instance_positions`.
Could these reads use their own short session (`with create_session() as s:`, as `_place_asset_event_sources` already does)? Everything they read (`dag_version`, `dynamic_region`, the TI count) is already committed, so they don't need to see the scheduler's uncommitted changes. That also removes the need for the savepoints.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -2042,17 +2391,49 @@ def _get_task_groups_details(dag: DAG | SerializedDAG, edge_map: dict[str, tuple
result = {}
for tg_id, tg in dag.task_group_dict.items():
task_edges, group_edges = edge_map.get(tg_id, ({}, {}))
+ is_mapped = _is_mapped_task_group(tg)
result[tg_id] = {
"parent_group": tg.parent_group.group_id if tg.parent_group else None,
+ # Independent, so a group could be both
+ "is_mapped": is_mapped,
+ "is_loop": hasattr(tg, "gate_task_id"),
Review Comment:
(optional) Nothing triggers this today, because the job facet is built from the serialized Dag, where every `SerializedLoopTaskGroup` has `gate_task_id`. But everywhere else in the PR a loop is recognised by `isinstance(…, (LoopTaskGroup, SerializedLoopTaskGroup))` (`_find_enclosing_loop`). Using the same test here keeps `taskGroups[...].is_loop` consistent with `tasks[...].loop_group`, and doesn't treat any group that happens to have a `gate_task_id` attribute as a loop. For example:
```python
def _is_loop_task_group(task_group: Any) -> bool:
from airflow.sdk.definitions._internal.loop import LoopTaskGroup
from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup
return isinstance(task_group, (LoopTaskGroup, SerializedLoopTaskGroup))
```
(guarded by `AIRFLOW_V_3_4_PLUS`, like `get_task_loop_info`).
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1542,56 +1682,228 @@ def _require_task_instance_position(
return position
+def _add_runtime_loop_iteration(
+ task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition
+) -> TaskInstancePosition:
+ if isinstance(task_instance, TaskInstance):
+ return position
+ # 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:
+ return position
+ return attrs.evolve(
+ position,
+ loop_index=loop.index,
+ 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_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))
+
+
+@attrs.define
+class _TaskCoordinate:
+ """Where Airflow stores a task instance, as its ``TaskCoordinateResolver`` reads it."""
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ region_id: UUID
+ region_index: int
+ dag_version_id: UUID | None
+
+
+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 on."""
+ if session is None or not AIRFLOW_V_3_4_PLUS:
+ return None
+ region_id = task_instance.region_id
+ if not isinstance(region_id, UUID) or region_id.int == 0:
+ return None
+ dag = getattr(getattr(task_instance, "task", None), "dag", None)
+ if (dag_version_id := getattr(dag, "dag_version_id", None)) is None:
+ return None
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ # Placed with the Dag version the task instance holds, which its task facet describes too, even when it pins
+ # another one: the resolver would load and deserialize that version for every task instance in the scheduler
+ # or API server process, while the region rows it reads don't depend on the version.
+ coordinate = _TaskCoordinate(
+ dag_id=task_instance.dag_id,
+ run_id=task_instance.run_id,
+ task_id=task_instance.task_id,
+ region_id=region_id,
+ region_index=task_instance.region_index,
+ dag_version_id=dag_version_id,
+ )
+ 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(
+ coordinate.dag_id, coordinate.run_id, coordinate.task_id, dag_version_id=dag_version_id
+ )
+ loop = resolver.get_loop_context(coordinate, task=task)
+ expanded_ti_count = resolver.get_expanded_ti_count(coordinate, 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_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+ 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]:
"""
- 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:
+ # 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
+ 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).where(
- DynamicRegion.id.in_({task_instance.region_id for task_instance in unplaced})
+ parent = aliased(DynamicRegion)
Review Comment:
This function now joins each region to its parent, and assumes a loop region is top-level and a mapped region sits directly under it, in order to fill `loop_index` / `loop_group_id`. That is a second, hand-written version of what core's `DynamicRegion.find_physical_coordinate` / `TaskCoordinateResolver.get_loop_context` do.
As far as I can trace, those new fields only matter on one path:
- `_place_asset_event_sources` passes only TIs with no task attached (`unplaced`). It then reads just `position.in_loop` and `position.map_index` to build a `runId`, and those come out the same as before this PR for every region shape.
- `get_stored_task_instance_position` gets here when `_get_core_task_instance_position` returned `None`. With a session on 3.4+, for a task in a loop, that only happens when the resolver **raised**. (No region means the task isn't in a loop, and `dag_version_id` is set on a serialized Dag.)
So the `without_iteration` list, the parent/grandparent join and the "unexpected shape" branches only fill in loop facet fields after core's own placement has already failed, using a copy of core's rules that can fall out of sync if region layout changes (forks, nesting). Would you consider dropping that part and letting the fallback report `in_loop` / `map_index` only, as it did before? If you want to keep the fallback complete, it could call `load_region_ancestry` + `DynamicRegion.find_physical_coordinate` instead of reimplementing them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315007871
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1135,11 +1193,44 @@ 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)
Review Comment:
Fixed in 7441d2d: a position passed in is used as is; the facet only looks up the iteration itself when no position was given.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315005213
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1570,26 +1724,47 @@ 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)}
Review Comment:
Fixed in 7441d2d: when the query only fills in mapped tasks' iterations, it runs in a savepoint, and a failure only leaves `loop_index` unset and logs a warning. The event still goes out and the caller's transaction stays usable. Task instances without their task still raise as before this PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Improve task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4315006986
##########
providers/openlineage/tests/system/openlineage/operator.py:
##########
@@ -351,33 +479,44 @@ 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)
Review Comment:
Agreed, fixed in 7441d2d: the runId check is skipped, with a log line, once any key failed its template.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix task loop information and tests for OpenLineage provider [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4358221307
##########
providers/openlineage/tests/system/openlineage/operator.py:
##########
@@ -351,33 +479,47 @@ 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 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)},
+ ):
+ message = "\n".join(f" {problem}" for problem in problems)
+ if self.fail_fast:
+ raise ValueError(f"runId integrity check failed:\n{message}")
+ errors["<runId integrity>"] = f"\n{message}"
finally:
if self.clear_variables:
Review Comment:
Nothing tests that `execute()` actually runs this. The PR's headline is that system tests check runId integrity by default, but the only test of `execute()` (`test_operator_skips_the_run_id_check_when_a_template_fails`) covers the skip branch; the other 45 tests call `check_run_id_integrity` directly. On a copy, each of these leaves 46/46 passing: `elif self.check_run_ids:` -> `elif False:`, the default changed to `check_run_ids=False`, `repeated_start_keys=set()`, removing the `fail_fast` raise, and replacing `errors["<runId integrity>"] = ...` with `pass`. So the check can be silently disabled in every system test.
One `execute()` test with `_check_key` returning events whose START and COMPLETE runIds differ would bind it: default -> `Event check failures`, `fail_fast` -> raises, `check_run_ids=False` -> passes. (`mock.patch("...time.sleep")` and `patch.object(operator, "_check_key")` in that test file also lack `autospec`.)
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
- "dag_bundle_version": lambda ti: (
- ti.bundle_instance.version if hasattr(ti, "bundle_instance") else None
+ "dag_version_id": lambda ti: (
+ str(dag_version_id) if (dag_version_id := getattr(ti, "dag_version_id", None)) else None
),
- "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None,
+ "dag_bundle_version": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "version"),
+ "dag_bundle_name": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "name"),
}
+ @staticmethod
+ def get_rendered_map_index(ti: TaskInstance | RuntimeTaskInstance) -> str | None:
+ """Return the label a mapped task rendered with ``map_index_template``, as the worker sends it."""
+ if AIRFLOW_V_3_0_PLUS and isinstance(ti, TaskInstance):
+ # Its rendered_map_index falls back to the map index when no label was rendered, the worker's doesn't
+ return safe_getattr(ti, "_rendered_map_index")
+ return getattr(ti, "rendered_map_index", None)
+
+ @staticmethod
+ def get_bundle_info(ti: TaskInstance | RuntimeTaskInstance, key: str) -> str | None:
+ """Return the name or version of the Dag bundle the attempt runs from."""
+ if (bundle := getattr(ti, "bundle_instance", None)) is not None: # A runtime task instance
+ return getattr(bundle, key, None)
+ if not isinstance(ti, TaskInstance):
+ return None
+ # Only for a FAIL the scheduler sends itself, e.g. for a missed heartbeat: UI and API state changes build no
+ # airflow facet, and the worker's runtime task instance has its bundle above
+ return getattr(TaskInstanceInfo.get_stored_dag_version(ti), f"bundle_{key}", None)
+
+ @staticmethod
+ def get_stored_dag_version(ti: TaskInstance) -> Any:
+ """Return the Dag version of a stored task instance, or None when it can't be read, as on Airflow 2."""
+ from sqlalchemy import inspect
+
+ try:
+ state: InstanceState = inspect(ti)
+ if "dag_version" not in state.unloaded or state.session is None:
+ # Loaded already, or not in a session to load it from, so reading it doesn't query
+ return safe_getattr(ti, "dag_version")
+ # The event must still be sent without it, and the savepoint keeps the caller's transaction usable
+ with state.session.begin_nested():
+ return ti.dag_version
+ except Exception:
+ log.warning(
+ "OpenLineage failed to read the Dag version of task instance %s", getattr(ti, "id", None)
+ )
+ log.debug("Exception details:", exc_info=True)
+ return 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 or
+ # with get_task_instance_position; it is then used as is.
+ self._position = position
+ super().__init__(obj)
+
def _extend_fields(self) -> None:
- map_index = get_task_instance_map_index(self.obj)
+ 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
+ self.rendered_map_index = self.get_rendered_map_index(self.obj) if map_index != -1 else None
self._fields.extend(("map_index", "rendered_map_index"))
+ 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:
Review Comment:
`taskInstance.id` is not the id of the attempt that failed for a stored `TaskInstance` on Airflow 3.0-3.3, so "Airflow 3+" is wrong here. Those versions call `prepare_db_for_next_try` (which does `self.id = uuid7()`) *before* `on_task_instance_failed` in `fetch_handle_failure_context` (3.0.6 `taskinstance.py:1945` then `:1952`; 3.3.2 `:1917` then `:1922`). 3.4 reverses the order (`notify_failure` at `:2727`, rotation after). So a scheduler-sent FAIL for an attempt that still has retries carries the *next* attempt's id (state `up_for_retry`), which is what the next START will carry. Probe against 3.2.2 and 3.3.2 trees: START id `01a164b1-5f93-...`, id in the FAIL facet `01a164b1-5fb4-...`. The runId is unaffected on those versions, but this field exists to say which execution sent the event.
Suggest emitting `id` for a stored `TaskInstance` only on 3.4+ (`AIRFLOW_V_3_4_PLUS or not isinstance(self.obj, TaskInstance)`), keeping it for `RuntimeTaskInstance` on every 3.x version, and saying so in the `id` description in `AirflowRunFacet.json`. Nothing tests this: the unit and system tests only run against the installed Airflow version.
##########
providers/openlineage/tests/system/openlineage/example_openlineage_task_states_dag.py:
##########
@@ -0,0 +1,76 @@
+# 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 tasks whose events only taskInstance's `state` tells apart (requires Airflow 3+).
+
+It checks:
+ - a task that skips itself emits START with `state` running, and COMPLETE with `state` skipped, which
+ a successful task's COMPLETE would carry as success
+ (Airflow 3.0 calls no listener when a task skips itself, so it emits no COMPLETE there)
+ - a task that raises AirflowFailException with retries left emits START with `state` running, and FAIL
+ with `state` failed at `try_number` 1 and `max_tries` 2: it isn't retried although
+ `try_number <= max_tries`
+ - taskInstance's `dag_version_id` of every event, from Airflow 3.1 on
+ - every task event's identity: taskInstance's `id`, `dag_version_id` (Airflow 3.1+) and bundle, the absence
+ of fields that don't apply, task's `task_id`, `trigger_rule`, `retries` and `retry_delay`, the Dag and Dag run,
+ and the parent run, plus each task in the Dag's job facet
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.common.compat.sdk import AirflowFailException, AirflowSkipException
+from airflow.sdk import task
+
+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_states_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:
+
+ @task
+ def skip_itself():
+ raise AirflowSkipException("Intentional skip")
+
+ @task(retries=2)
+ def fail_without_retry():
+ raise AirflowFailException("Intentional failure that skips the retries left")
+
+ check_events = OpenLineageTestOperator(
+ task_id="check_events",
+ file_path=get_expected_event_file_path(DAG_ID),
+ trigger_rule="all_done",
+ )
Review Comment:
This will fail on Airflow 3.0. `check_run_ids` defaults to `True`, and for a TASK runId `check_run_id_integrity` requires exactly one COMPLETE/FAIL (`operator.py:357`, `job_type != "DAG" and not terminals`). The `__af3_0` template just made `skip_itself`'s COMPLETE `$absent` (3.0 calls no listener when a task skips itself), so once the templates match its runId group holds only a START and `check_events` fails. Running the PR's own `check_run_id_integrity` over events built from that template gives:
`runId skip_itself-uuid-try1 of TASK `openlineage_task_states_dag.skip_itself` should have one START (...) and one COMPLETE/FAIL, got: ['START']`
`openlineage_task_states_dag.json` (3.1+) passes. 5ee168897b0 fixed the e2e failure for the missing COMPLETE; this is the next failure on the same job. Either pass `check_run_ids=False` on 3.0 for this Dag, or have the checker treat a terminal key declared `$absent` as an intentionally partial lifecycle (with a unit test for START + `$absent` terminals). A scan of the other expected-event files found no other START-without-terminal job.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1542,56 +1682,228 @@ def _require_task_instance_position(
return position
+def _add_runtime_loop_iteration(
+ task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition
+) -> TaskInstancePosition:
+ if isinstance(task_instance, TaskInstance):
+ return position
+ # 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:
+ return position
+ return attrs.evolve(
+ position,
+ loop_index=loop.index,
+ 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_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))
+
+
+@attrs.define
+class _TaskCoordinate:
+ """Where Airflow stores a task instance, as its ``TaskCoordinateResolver`` reads it."""
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ region_id: UUID
+ region_index: int
+ dag_version_id: UUID | None
+
+
+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 on."""
+ if session is None or not AIRFLOW_V_3_4_PLUS:
+ return None
+ region_id = task_instance.region_id
+ if not isinstance(region_id, UUID) or region_id.int == 0:
+ return None
+ dag = getattr(getattr(task_instance, "task", None), "dag", None)
+ if (dag_version_id := getattr(dag, "dag_version_id", None)) is None:
+ return None
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ # Placed with the Dag version the task instance holds, which its task facet describes too, even when it pins
+ # another one: the resolver would load and deserialize that version for every task instance in the scheduler
+ # or API server process, while the region rows it reads don't depend on the version.
+ coordinate = _TaskCoordinate(
+ dag_id=task_instance.dag_id,
+ run_id=task_instance.run_id,
+ task_id=task_instance.task_id,
+ region_id=region_id,
+ region_index=task_instance.region_index,
+ dag_version_id=dag_version_id,
+ )
+ try:
+ # The savepoint keeps the caller's transaction usable after a failure
+ with session.begin_nested():
Review Comment:
**Scheduler cost.** `get_stored_task_instance_position` now runs for every scheduler-sent FAIL (`process_executor_events` and the heartbeat purge), and this path costs a SAVEPOINT, a fresh `TaskCoordinateResolver` + `DBDagBag`, region-ancestry selects and, for mapped tasks, a `count(*)` over the whole expansion. That is 4 statements for a loop body and 5-6 for a mapped task in a loop (your own `assert_queries_count(4)` / `(5)`), where main needed 0 when the task is attached.
It isn't only loops: any mapped task instance has a non-sentinel region, so every mapped TI the scheduler fails goes through here, and the `count(*)` makes N failures scan O(N^2) rows (0.35 -> 1.17 ms per call from N=200 to N=2000 on SQLite; PostgreSQL not measured). N is the number of task instances the scheduler fails in one pass, e.g. a worker pod lost with 500 running mapped TIs is ~3000 statements inside the scheduler loop's transaction, with nothing cached between calls (`for_dag` builds a new `DBDagBag` each time).
Cheap options: skip the resolver when the task is not in a loop (`get_orm_task_instance_positions` already places an attached-task mapped TI with 0 queries), omit `expanded_ti_count` on scheduler-sent events or memoize it per `(dag_id, run_id, task_id, region_id)` for the pass, and share one resolver per pass.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1542,56 +1682,228 @@ def _require_task_instance_position(
return position
+def _add_runtime_loop_iteration(
+ task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition
+) -> TaskInstancePosition:
+ if isinstance(task_instance, TaskInstance):
+ return position
+ # 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:
+ return position
+ return attrs.evolve(
+ position,
+ loop_index=loop.index,
+ 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_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))
+
+
+@attrs.define
+class _TaskCoordinate:
+ """Where Airflow stores a task instance, as its ``TaskCoordinateResolver`` reads it."""
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ region_id: UUID
+ region_index: int
+ dag_version_id: UUID | None
+
+
+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 on."""
+ if session is None or not AIRFLOW_V_3_4_PLUS:
+ return None
+ region_id = task_instance.region_id
+ if not isinstance(region_id, UUID) or region_id.int == 0:
+ return None
+ dag = getattr(getattr(task_instance, "task", None), "dag", None)
+ if (dag_version_id := getattr(dag, "dag_version_id", None)) is None:
+ return None
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ # Placed with the Dag version the task instance holds, which its task facet describes too, even when it pins
+ # another one: the resolver would load and deserialize that version for every task instance in the scheduler
+ # or API server process, while the region rows it reads don't depend on the version.
Review Comment:
**Loop / runId.** The scheduler places this task instance with the Dag version it *holds* (`ti.task.dag`), not the one it *pinned* (`ti.dag_version_id`), and the reasoning in this comment is incomplete: the region rows don't depend on the version, but loop membership does. `enclosing_loop(task)` comes from that version's group definition, so if the held version no longer has the task in a loop, `get_loop_context` returns `None` without raising, you get `in_loop=False`, and `get_task_instance_run_id` builds a coordinate hash instead of `str(ti.id)`.
When can the two differ: `get_dag_for_run` resolves to the *latest* version for any run without a `bundle_version` (`DBDagBag._version_from_dag_run`, i.e. the default local bundle), while the worker was placed with `ti.dag_version_id` (`ti_run`). A Dag edit that removes or renames the loop between the worker's START and a scheduler-detected failure (heartbeat timeout, executor-reported failure) is enough.
Reproduced with two serialized versions of one Dag (`body.work` in a loop, then `body` as a plain `TaskGroup`):
- pinned version: `in_loop=True, loop_index=1`, runId == `ti.id` (`01a164ed-d412-...`)
- held (latest) version: `in_loop=False`, runId `018cc251-f400-757a-...`
So the FAIL lands on a different run than the START, which stays open in the lineage backend. The loop facet fields (`max_iterations`, `terminal_task_id`) are also read from the held version even when the runId agrees.
The commit message of 456a27fd158 says task instances pinned to another version "fall back to reading their region rows"; the code doesn't do that. And `test_stored_task_instance_is_placed_as_airflow_places_it_for_the_worker[holds_another_version=True]` only reassigns `dag_version_id` on a Dag with identical content, so it can't see this: passing `task_instance.dag_version_id` instead of `dag.dag_version_id` to `get_task` leaves the whole suite green (`DBDagBag.get_dag` is called 0 times unmutated, 1 time mutated).
Suggestion: build the coordinate from `task_instance` (it already satisfies `TaskCoordinate`, and it is what `ti_run` passes) and only when `task_instance.dag_version_id != dag.dag_version_id` resolve against the pinned version. That is the one case where paying the Dag load is justified. The test then needs two versions that differ in loop membership, plus a spy on `DBDagBag.get_dag` (`assert_queries_count` is an upper bound and the identity map hides the repeated select).
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -3924,38 +4273,901 @@ 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(),
+ dag_id="run_id_dag",
+ task_id="body.work",
+ try_number=1,
+ region_id=uuid4(),
+ map_index=-1,
+ task=dag.get_task("body.work"),
+ )
+ logical_date = timezone.datetime(2024, 1, 1)
+
+ run_id = get_task_instance_run_id(ti, logical_date=logical_date)
+
+ 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
+
+
+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")
+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
+
+ 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"), loop=None
+ )
- assert get_regional_task_instance_run_id(ti) == (str(ti.id) if in_loop else None)
+ 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="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,
+ ),
+ )
+
+ assert get_task_instance_position(ti) == TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ expanded_ti_count=2,
+ loop_index=2,
+ 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,
+ )
+
+
+@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4")
+@pytest.mark.parametrize("holds_another_version", [False, True])
+def test_stored_task_instance_is_placed_as_airflow_places_it_for_the_worker(
+ dag_maker, session, holds_another_version
+):
+ 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}
+ # As the scheduler's Dag bag sets it on the Dag it loads for the run, which can be a later version than the
+ # one the task instances pin
+ tis["body.plain", 1].task.dag.dag_version_id = (
+ uuid7() if holds_another_version else 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()
+ 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_trajectory": 1, "loop_group_id": "body"}
+
+ # 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_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"] == {
+ "index": 2,
+ "region_id": str(fork.id),
+ "trajectory": 2,
+ "forked_from_region_id": str(loop_region.id),
+ "resumes_from_index": 2,
+ }
@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Dag versions require Airflow 3")
+@pytest.mark.parametrize("load_fails", [False, True])
+@patch("airflow.providers.openlineage.utils.utils.log", autospec=True)
+def test_stored_task_instance_info_reads_its_dag_version_without_risking_the_event(
+ mock_log, dag_maker, session, load_fails
+):
+ with dag_maker(session=session):
+ PythonOperator(task_id="work", python_callable=list)
+ ti = dag_maker.create_dagrun().task_instances[0]
+ session.expire(ti, ["dag_version"])
+
+ if load_fails:
+ with patch.object(
+ TaskInstance,
+ "dag_version",
+ new_callable=PropertyMock,
+ side_effect=RuntimeError("Lost connection"),
+ ):
+ info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1)))
+ else:
+ info = dict(TaskInstanceInfo(ti, position=TaskInstancePosition(in_loop=False, map_index=-1)))
+
+ expected_bundle_name = None if load_fails else "dag_maker"
+ assert (info["dag_bundle_name"], info["dag_version_id"]) == (expected_bundle_name, str(ti.dag_version_id))
+ assert mock_log.warning.call_count == (
+ 2 if load_fails else 0
+ ) # Once for the bundle name, once for its version
+ # The failure was rolled back to its savepoint, so the caller's transaction is still usable
+ assert session.execute(select(literal(1))).scalar() == 1
+
+
+@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4")
+@pytest.mark.parametrize("task_id", ["body.plain", "body.inner.mapped"])
+def test_stored_and_runtime_task_instances_identify_an_execution_alike(dag_maker, session, task_id):
+ """
+ The scheduler's events about a stored task instance must tell everything the worker's tell about it.
+
+ A field missing on one side only would leave consumers without it for that kind of event, while the other
+ side's events made them rely on it.
+ """
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with dag_maker(serialized=True):
+
+ @task_group
+ def body():
+ plain = EmptyOperator(task_id="plain")
+ with TaskGroup("inner"):
+ plain >> PythonOperator.partial(task_id="mapped", python_callable=list).expand(
+ op_kwargs=[{}, {}]
+ )
+
+ create_loop(body, max_iterations=3)
+ run = dag_maker.create_dagrun()
+ stored = next(ti for ti in run.task_instances if ti.task_id == task_id and ti.region_index in (0, 1))
+ stored.task.dag.dag_version_id = stored.dag_version_id
+ stored.state, stored.try_number, stored.max_tries = TaskInstanceState.UP_FOR_RETRY, 1, 2
+ resolver = TaskCoordinateResolver.for_dag(stored.task.dag, session)
+ mapped = stored.task.get_needs_expansion()
+ # What the server sends the worker for this execution, on the task of the Dag the worker parses
+ runtime = SimpleNamespace(
+ id=stored.id,
+ dag_id=stored.dag_id,
+ run_id=stored.run_id,
+ task_id=stored.task_id,
+ try_number=stored.try_number,
+ max_tries=stored.max_tries,
+ state=stored.state,
+ map_index=stored.region_index if mapped else -1,
+ rendered_map_index=None,
+ region_id=stored.region_id,
+ region_index=stored.region_index,
+ dag_version_id=stored.dag_version_id,
+ bundle_instance=SimpleNamespace(
+ name=stored.dag_version.bundle_name, version=stored.dag_version.bundle_version
+ ),
+ loop=resolver.get_loop_context(stored, task=stored.task),
+ expanded_ti_count=resolver.get_expanded_ti_count(stored, task=stored.task),
+ task=dag_maker.dag.get_task(task_id),
+ )
+
+ stored_info = dict(
+ TaskInstanceInfo(stored, position=get_stored_task_instance_position(stored, session=session))
+ )
+ runtime_info = dict(TaskInstanceInfo(runtime, position=get_task_instance_position(runtime)))
+
+ # Only stored task instances have these, which describe the attempt's resources rather than which it is
+ resources = {"duration", "pool", "queued_dttm", "note", "log_url"}
+ assert {k: v for k, v in stored_info.items() if k not in resources} == {
+ k: v for k, v in runtime_info.items() if k not in resources
+ }
+ assert ("expanded_ti_count" in runtime_info) is mapped
+ assert runtime_info["loop"]["trajectory"] == 1
+ identity = ("task_id", "task_group", "loop", "retries", "retry_delay")
+ assert {k: dict(TaskInfo(stored.task))[k] for k in identity} == {
+ k: dict(TaskInfo(runtime.task))[k] for k in identity
+ }
+
+
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4")
+@pytest.mark.parametrize(
+ "dag_has_version", [pytest.param(True, id="airflow_fails"), pytest.param(False, id="no_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_without_airflow(mock_for_dag, dag_has_version):
+ 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
+ dag.dag_version_id = ti.dag_version_id if dag_has_version else None
+ session = MagicMock(spec=Session)
+
+ assert get_stored_task_instance_position(ti, session=session) == TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_index=1, loop_group_id="body"
+ )
+ if dag_has_version:
+ mock_for_dag.assert_called_once_with(dag, session)
+ else:
+ mock_for_dag.assert_not_called()
+ 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
+ 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_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)
+
+
@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_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.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))
+
+ assert info["map_index"] == (1 if mapped else None)
+ if in_loop:
+ assert info["loop"] == {
+ "index": 2,
+ }
+ else:
+ assert "loop" not in info
+
+
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4")
+@pytest.mark.parametrize("nested", [False, True])
+def test_task_info_has_the_loop_around_the_task(nested):
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ def ready():
+ return True
+
+ with DAG("nested_loop_dag") as dag:
+ with TaskGroup("outer"):
+
+ @task_group
+ def body():
+ if nested:
+
+ @task_group
+ def inner():
+ EmptyOperator(task_id="work")
+
+ inner()
+ else:
+ EmptyOperator(task_id="work")
+
+ create_loop(body, max_iterations=3, until=ready)
+ task_id = "outer.body.inner.work" if nested else "outer.body.work"
+
+ info = dict(TaskInfo(dag.get_task(task_id)))
+
+ assert info["loop"] == {
+ "group_id": "outer.body",
+ "max_iterations": 3,
+ "has_until": True,
+ "gate_task_id": dag.task_group_dict["outer.body"].gate_task_id,
+ "terminal_task_id": task_id,
+ }
+ # task_group is the task's own group and never describes the loop; its unique id tells whether it is the loop
+ assert (info["task_group"]["group_id"], info["task_group"]["unique_group_id"]) == (
+ ("inner", "outer.body.inner") if nested else ("body", "outer.body")
+ )
+ assert (info["task_group"]["unique_group_id"] == info["loop"]["group_id"]) is not nested
+ assert "gate_task_id" not in info["task_group"]
+
+
+@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}, False, id="no_until_continues"),
+ 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):
+ 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
+ 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 = GateTaskInstance(id=uuid7(), map_index=-1, _loop_decision=None)
+ 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" 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_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 = [_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)
+
+ assert positions == {
+ ti.id: TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ loop_index=2 if with_session else None,
+ 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" not in info
+
+
+def test_task_instance_info_uses_the_position_it_is_given():
+ # 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,
+ map_index=1,
+ expanded_ti_count=4,
+ loop_index=2,
+ 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["expanded_ti_count"]) == (1, 4)
+ assert info["loop"] == {
+ "index": 2,
+ "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")
+def test_task_group_info_has_no_loop_attributes():
+ 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_attributes = {"max_iterations", "has_until", "terminal_task_id", "gate_task_id"}
+ loop_info = dict(TaskGroupInfo(loop))
+ assert loop_info["group_id"] == "body"
+ assert not loop_attributes & set(loop_info)
+ assert not loop_attributes & set(TaskGroupInfo(plain_group))
+
+
+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_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
+
+ 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", 1): TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_index=1, loop_group_id="body"
+ ),
+ ("body.last", 1): TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_index=1, loop_group_id="body"
+ ),
+ **{
+ (task_id, map_index): TaskInstancePosition(
+ in_loop=True, map_index=map_index, loop_index=1, loop_group_id="body"
+ )
+ 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):
Review Comment:
`pytest.raises(NotImplementedError)` has no `match`, so any unrelated `NotImplementedError` raised while building the Dag keeps this tripwire green. The three real messages are distinct ("Nested loops are not supported", "A loop cannot be inside a mapped task group", "operator expansion in an expanded task group is not yet supported"); parametrize the expected message and pass `match=`.
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -2011,6 +2126,201 @@ 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
Review Comment:
This PR adds ~43 function-body imports to this file (and 6 to test_listener.py), none with a reason. The same `task_group`/`create_loop` pair is repeated ~17 times; `DynamicRegion` is re-imported at line 4381 although it is already imported at line 121 under `if AIRFLOW_V_3_4_PLUS:`; `task` is re-imported at 2233, 5022 and 5091 although it is a top-level import; and `create_scheduler_dag` is imported inline at 2236 and 2312 while test_listener.py imports it at module level. Extend the existing `if AIRFLOW_V_3_4_PLUS:` block at the top instead.
Also `create_loop(body, max_iterations=3)` with the same body is repeated ~15 times (1 in the base file); a small helper returning the Dag, with optional mapped/nested variants, would remove ~150 lines, and the four mapped-in-loop placement tests overlap (only the real-database ones measure the code rather than the stub).
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1461,6 +1556,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.
Review Comment:
The title is "Fix task loop information", and `airflowState` (`tasksState`/`tasksDuration`) is still keyed by `task_id`, so the loop passes and mapped slots of one task collapse into one entry (the last by `(region_id, region_index)` wins). The TODO records it. Either fix it here, or link an issue in the TODO so it is tracked: as it stands the Dag-run-level facet is the one place that still reports a single state for every iteration of a loop.
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -3924,38 +4273,901 @@ 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(),
+ dag_id="run_id_dag",
+ task_id="body.work",
+ try_number=1,
+ region_id=uuid4(),
+ map_index=-1,
+ task=dag.get_task("body.work"),
+ )
+ logical_date = timezone.datetime(2024, 1, 1)
+
+ run_id = get_task_instance_run_id(ti, logical_date=logical_date)
+
+ 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
+
+
+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")
+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
+
+ 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"), loop=None
+ )
- assert get_regional_task_instance_run_id(ti) == (str(ti.id) if in_loop else None)
+ 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="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,
+ ),
+ )
+
+ assert get_task_instance_position(ti) == TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ expanded_ti_count=2,
+ loop_index=2,
+ 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,
+ )
+
+
+@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires Airflow 3.4")
+@pytest.mark.parametrize("holds_another_version", [False, True])
+def test_stored_task_instance_is_placed_as_airflow_places_it_for_the_worker(
+ dag_maker, session, holds_another_version
+):
+ 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}
+ # As the scheduler's Dag bag sets it on the Dag it loads for the run, which can be a later version than the
+ # one the task instances pin
+ tis["body.plain", 1].task.dag.dag_version_id = (
+ uuid7() if holds_another_version else 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()
+ 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_trajectory": 1, "loop_group_id": "body"}
+
+ # 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_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"] == {
+ "index": 2,
+ "region_id": str(fork.id),
+ "trajectory": 2,
+ "forked_from_region_id": str(loop_region.id),
+ "resumes_from_index": 2,
+ }
@pytest.mark.db_test
+@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Dag versions require Airflow 3")
+@pytest.mark.parametrize("load_fails", [False, True])
+@patch("airflow.providers.openlineage.utils.utils.log", autospec=True)
+def test_stored_task_instance_info_reads_its_dag_version_without_risking_the_event(
Review Comment:
This test doesn't bind the savepoint. The injected failure is a Python `RuntimeError` from a `PropertyMock`, never a database error, so "the caller's transaction is still usable" holds with or without `begin_nested()`. Removing it from `get_stored_dag_version` leaves test_utils.py green, and so does removing it from `_get_core_task_instance_position` (the resolver-failure test uses `MagicMock(spec=Session)` and a mocked resolver). The savepoint matters when a statement fails on PostgreSQL and poisons the scheduler's `process_executor_events` transaction, so it needs a real failing statement (a `@pytest.mark.backend("postgres")` db test), or at least an assertion that `begin_nested` was entered for a real session.
Related gap: nothing runs a scheduler-style FAIL through `OpenLineageListener` with a real session. `test_manual_event_passes_its_stored_position_to_the_airflow_run_facet` uses a session-less TI, so changing `session=object_session(ti)` to `session=None` at `listener.py:892` leaves 383 tests green, and "scheduler FAIL carries the same fields as the worker" is pinned only at the `utils` level. One db test through `on_task_instance_failed` with a real loop TI, asserting `taskInstance.loop.trajectory`, `expanded_ti_count` and `runId == taskInstance.id`, would cover it.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
Review Comment:
**Loop (stored vs public index).** `log_url` is read from the stored task instance, and `TaskInstance.log_url` builds its `/mapped/<n>` segment from `self.region_index`, which for a non-mapped loop body is the *iteration*. On a live 3.4.0 stack, a scheduler-sent FAIL (heartbeat kill) of `body.plain` at iteration 1 had `.../tasks/body.plain/mapped/1?try_number=1`, where the worker's START has `.../tasks/body.plain?try_number=1`; the first link doesn't exist. This is the same synonym trap you already handle for `map_index`, and it contradicts "FAIL events sent by the scheduler carry the same fields as worker ones".
Build it from the placed `position.map_index` for a stored TI (or drop the segment when it is -1). A unit test comparing the stored and runtime `log_url` of a loop body would pin it.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
- "dag_bundle_version": lambda ti: (
- ti.bundle_instance.version if hasattr(ti, "bundle_instance") else None
+ "dag_version_id": lambda ti: (
+ str(dag_version_id) if (dag_version_id := getattr(ti, "dag_version_id", None)) else None
),
- "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None,
+ "dag_bundle_version": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "version"),
+ "dag_bundle_name": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "name"),
}
+ @staticmethod
+ def get_rendered_map_index(ti: TaskInstance | RuntimeTaskInstance) -> str | None:
+ """Return the label a mapped task rendered with ``map_index_template``, as the worker sends it."""
+ if AIRFLOW_V_3_0_PLUS and isinstance(ti, TaskInstance):
+ # Its rendered_map_index falls back to the map index when no label was rendered, the worker's doesn't
+ return safe_getattr(ti, "_rendered_map_index")
+ return getattr(ti, "rendered_map_index", None)
+
+ @staticmethod
+ def get_bundle_info(ti: TaskInstance | RuntimeTaskInstance, key: str) -> str | None:
+ """Return the name or version of the Dag bundle the attempt runs from."""
+ if (bundle := getattr(ti, "bundle_instance", None)) is not None: # A runtime task instance
+ return getattr(bundle, key, None)
+ if not isinstance(ti, TaskInstance):
+ return None
+ # Only for a FAIL the scheduler sends itself, e.g. for a missed heartbeat: UI and API state changes build no
+ # airflow facet, and the worker's runtime task instance has its bundle above
+ return getattr(TaskInstanceInfo.get_stored_dag_version(ti), f"bundle_{key}", None)
+
+ @staticmethod
+ def get_stored_dag_version(ti: TaskInstance) -> Any:
+ """Return the Dag version of a stored task instance, or None when it can't be read, as on Airflow 2."""
+ from sqlalchemy import inspect
Review Comment:
Imports added inside function bodies with no stated reason: `from sqlalchemy import inspect` (here), `from sqlalchemy.orm import aliased` (line 1818) and `from airflow.sdk.definitions.taskgroup import MappedTaskGroup` with its `except ImportError` fallback (line 2420). sqlalchemy is already imported at module level (`from sqlalchemy.orm.exc import DetachedInstanceError`), and `SerializedMappedOperator` (lines 83-88) shows the module's top-level try/except pattern for a compat import. `TaskCoordinateResolver` (line 1735) is genuinely 3.4+-only, but needs a one-line comment saying so.
##########
providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py:
##########
@@ -0,0 +1,164 @@
+# 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 and one stopped early by `until`
+(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`, `state`, `map_index` and `loop` of every event: `loop.index` is the
+ iteration, counted from 1, and `map_index` the mapped position, and a plain task has no `loop`
+ - task's `loop`, the loop around the task: its `group_id`, `max_iterations`, `has_until`, `gate_task_id`
+ and `terminal_task_id`, also for a task of a task group inside a loop, and absent for the plain task
+ - `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
Review Comment:
**Loop fields the e2e doesn't actually bind.** The docstring says `loop.region_id` is the loop region "of the iteration, shared by the tasks of one iteration, mapped ones included". Nothing asserts that, and it isn't accurate: `loop.region_id` is the loop's physical region, shared by *all* iterations until a clear forks it (same value for iterations 1-3 in a captured run). It is a mapped task's `taskInstance.region_id` (its expansion region) that differs from it.
In both loop JSONs `region_id`, `loop.region_id` and `taskInstance.id` are only `{{ is_uuid(result) }}`, and `match()` can't compare two fields, so wrong values pass. On a live 3.4.0 scheduler, replacing `loop_region_id=loop.region_id` with the task's own `region_id` (in `_add_runtime_loop_iteration`) still ended `success`. So did a constant `loop.region_id`, a constant `taskInstance.id`, and forcing `trajectory` to 1 or `forked_from_region_id`/`resumes_from_index` to `None` (those are only asserted absent).
"no loop is cleared" also means `trajectory > 1`, the fork fields, and the `region_id` term this PR adds to `Attempt` (it only matters across a clear) are exercised by no system test, only by unit tests with hand-built objects. Clearing the gate through the API on the same stack produced trajectory 2 with the fork fields set and `check_run_id_integrity` reported no problems, so a Dag step that clears a gate after the first pass is feasible.
Suggest loop invariants in `check_run_id_integrity`: for a TI with `loop` and no `map_index`, `loop.region_id == region_id`; mapped in a loop, `loop.region_id != region_id`; `loop.region_id` constant per (loop group, trajectory); `taskInstance.id == runId` when `loop` is set. Plus one e2e Dag that clears a gate. Otherwise trim the claims from the docstring. The job-facet bullet overclaims too: only 4 of the 10 tasks assert `task_group`/`loop_group`/`is_mapped`/`parse_time_mapped_count`.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1542,56 +1682,228 @@ def _require_task_instance_position(
return position
+def _add_runtime_loop_iteration(
+ task_instance: TaskInstance | RuntimeTaskInstance, position: TaskInstancePosition
+) -> TaskInstancePosition:
+ if isinstance(task_instance, TaskInstance):
+ return position
+ # 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:
+ return position
+ return attrs.evolve(
+ position,
+ loop_index=loop.index,
+ 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_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))
+
+
+@attrs.define
+class _TaskCoordinate:
+ """Where Airflow stores a task instance, as its ``TaskCoordinateResolver`` reads it."""
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ region_id: UUID
+ region_index: int
+ dag_version_id: UUID | None
+
+
+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 on."""
+ if session is None or not AIRFLOW_V_3_4_PLUS:
+ return None
+ region_id = task_instance.region_id
+ if not isinstance(region_id, UUID) or region_id.int == 0:
+ return None
+ dag = getattr(getattr(task_instance, "task", None), "dag", None)
+ if (dag_version_id := getattr(dag, "dag_version_id", None)) is None:
+ return None
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ # Placed with the Dag version the task instance holds, which its task facet describes too, even when it pins
+ # another one: the resolver would load and deserialize that version for every task instance in the scheduler
+ # or API server process, while the region rows it reads don't depend on the version.
+ coordinate = _TaskCoordinate(
+ dag_id=task_instance.dag_id,
+ run_id=task_instance.run_id,
+ task_id=task_instance.task_id,
+ region_id=region_id,
+ region_index=task_instance.region_index,
+ dag_version_id=dag_version_id,
+ )
+ 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(
+ coordinate.dag_id, coordinate.run_id, coordinate.task_id, dag_version_id=dag_version_id
+ )
+ loop = resolver.get_loop_context(coordinate, task=task)
+ expanded_ti_count = resolver.get_expanded_ti_count(coordinate, 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_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+ 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]:
"""
- 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:
+ # 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
+ 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).where(
- DynamicRegion.id.in_({task_instance.region_id for task_instance in unplaced})
+ 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"),
Review Comment:
**Design question, non-blocking.** There are now three implementations of "where does this task instance sit": `_get_task_instance_position` (task definition + stored index), `_get_core_task_instance_position` (Airflow's `TaskCoordinateResolver`) and this region-row parser with the self-join and the `without_iteration` second pass. They duplicate parts of the resolver, `loop_position` and `DynamicRegion.find_physical_coordinate`. I compared them on linear, forked, mapped-in-loop, mapped-group-member and placeholder rows and they agree on `in_loop`, `map_index`, `loop_index` and `loop_group_id`. But the fallback never has `loop.region_id`, `trajectory`, `forked_from_region_id`, `resumes_from_index` or `expanded_ti_count`, so when the resolver fails, the scheduler FAIL of a forked-region task instance carries less than its START, and the schema doesn't say those fields can be absent.
`loop_index` from region rows only reaches an event when the resolver raised: the API-server path has no `ti.task`, so it builds no `airflow` facet, and the asset-event path only needs `in_loop` and `map_index`. The `elif region.grandparent_region_id is None` branch and the "unexpected shape" warning are bound only by hand-built rows (flipping the branch to `elif True:` fails one synthetic-row test and not `test_region_shapes_of_airflow_are_the_ones_openlineage_places`), and `loop_group_id` is read once in `src` and asserted on ~15 test lines.
Would it be enough for the fallback to be (`in_loop`, `map_index`) as before this PR, with the facet schema saying which `loop` fields can be missing on scheduler-sent events? That drops the parent join, the second pass and ~4 tests. Related: should core expose loop/expansion accessors on a stored TI, as #75132 did for the runtime TI, so the provider isn't hand-building a `_TaskCoordinate` just to override `dag_version_id`?
##########
providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py:
##########
@@ -0,0 +1,73 @@
+# 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
+ - task's `retries` and `retry_delay` on every event
+ - taskInstance's `state`: up_for_retry on the FAIL of the first try, success on the COMPLETE of the retry
+ - each try has its own `runId` (the `runId` check of OpenLineageTestOperator is enabled by default)
+ - every task event's identity: taskInstance's `id`, `dag_version_id` (Airflow 3.1+) and bundle, the absence
+ of fields that don't apply, task's `task_id`, `trigger_rule`, `retries` and `retry_delay`, the Dag and Dag run,
+ and the parent run, plus each task in the Dag's job facet
+"""
+
+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)},
Review Comment:
`default_args` also gives `check_events` `retries=1`. On a template mismatch the operator's `finally` deletes the Variables, the retry then runs and fails with "variable does not exist (no events emitted)", and the e2e log shows only the latest attempt, so the real mismatch from attempt 1 is hidden (a lane observed `check_events` start twice in a run where nothing was emitted). Put `retries=1, retry_delay=...` on `flaky_task` and leave the Dag default at 0; the expected JSON only lists `flaky_task` events, so it doesn't change.
##########
providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json:
##########
@@ -196,6 +225,73 @@
"map_index": {
"type": "integer"
},
+ "id": {
+ "description": "Id of the task instance attempt (Airflow 3+).",
+ "type": "string"
+ },
+ "state": {
+ "type": "string"
+ },
+ "max_tries": {
+ "type": "integer"
+ },
+ "rendered_map_index": {
+ "type": "string"
+ },
+ "expanded_ti_count": {
+ "description": "Number of instances in the mapped task's expansion (Airflow 3.4+).",
+ "type": "integer"
+ },
+ "region_id": {
+ "description": "Region of the loop iteration or mapped expansion (Airflow 3.4+).",
+ "type": "string"
+ },
+ "region_index": {
+ "description": "Position in the region (Airflow 3.4+).",
+ "type": "integer"
+ },
Review Comment:
**Docs / design question, non-blocking.** The facet schema is now the only description of these fields (the guide was dropped from this PR), and several mean different things depending on the task:
- `region_index` is the 1-based loop iteration for a non-mapped loop body but the 0-based map slot for a mapped task (probe: `body.plain` 1,2; `body.mapped` 0,1). It is described only as "Position in the region".
- `region_id` is the loop's region for a loop body but the mapped task's *expansion* region for a mapped task, so a mapped task in a loop carries two different ids (`region_id` and `loop.region_id`). After a clear, `loop.region_id` is the fork's region, not the family root.
- `loop.index` (1-based), `loop.trajectory` (1 = original, n+1 = nth fork) and `loop.resumes_from_index` have no description. `expanded_ti_count` counts the live (non-REMOVED) instances of that expansion, so it is per loop pass, and is 1 for a not-yet-expanded placeholder.
Is exposing the physical storage coordinate (`region_id`/`region_index`) intended? `region_index` always duplicates `loop.index` or `map_index`, and `loop.index` + `map_index` + `loop.region_id` already carry the same information. If you keep them, descriptions like the above would save consumers from reverse-engineering the storage model.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1125,21 +1125,105 @@ def team_name(cls, dagrun: DagRun) -> str | None:
class TaskInstanceInfo(InfoJsonEncodable):
"""Defines encoding TaskInstance object to JSON."""
- includes = ["duration", "log_url", "pool", "queued_dttm", "try_number"]
+ includes = ["duration", "log_url", "max_tries", "pool", "queued_dttm", "state", "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
- "dag_bundle_version": lambda ti: (
- ti.bundle_instance.version if hasattr(ti, "bundle_instance") else None
+ "dag_version_id": lambda ti: (
+ str(dag_version_id) if (dag_version_id := getattr(ti, "dag_version_id", None)) else None
),
- "dag_bundle_name": lambda ti: ti.bundle_instance.name if hasattr(ti, "bundle_instance") else None,
+ "dag_bundle_version": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "version"),
+ "dag_bundle_name": lambda ti: TaskInstanceInfo.get_bundle_info(ti, "name"),
Review Comment:
`dag_bundle_version` and `dag_bundle_name` each call `get_bundle_info`, which calls `get_stored_dag_version`. For a stored TI the Dag version is therefore read twice per event, and when the read fails each call opens its own savepoint and logs the same warning ("OpenLineage failed to read the Dag version of task instance ...") twice (probe with `begin_nested` forced to raise: 2 calls, 2 warnings). Read it once per `TaskInstanceInfo`, e.g. set both bundle fields in `_extend_fields`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix task loop information and tests for OpenLineage provider [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4358287099
##########
providers/openlineage/tests/system/openlineage/example_openlineage_task_states_dag.py:
##########
@@ -0,0 +1,76 @@
+# 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 tasks whose events only taskInstance's `state` tells apart (requires Airflow 3+).
+
+It checks:
+ - a task that skips itself emits START with `state` running, and COMPLETE with `state` skipped, which
+ a successful task's COMPLETE would carry as success
+ (Airflow 3.0 calls no listener when a task skips itself, so it emits no COMPLETE there)
+ - a task that raises AirflowFailException with retries left emits START with `state` running, and FAIL
+ with `state` failed at `try_number` 1 and `max_tries` 2: it isn't retried although
+ `try_number <= max_tries`
+ - taskInstance's `dag_version_id` of every event, from Airflow 3.1 on
+ - every task event's identity: taskInstance's `id`, `dag_version_id` (Airflow 3.1+) and bundle, the absence
+ of fields that don't apply, task's `task_id`, `trigger_rule`, `retries` and `retry_delay`, the Dag and Dag run,
+ and the parent run, plus each task in the Dag's job facet
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.common.compat.sdk import AirflowFailException, AirflowSkipException
+from airflow.sdk import task
+
+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_states_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:
+
+ @task
+ def skip_itself():
+ raise AirflowSkipException("Intentional skip")
+
+ @task(retries=2)
+ def fail_without_retry():
+ raise AirflowFailException("Intentional failure that skips the retries left")
+
+ check_events = OpenLineageTestOperator(
+ task_id="check_events",
+ file_path=get_expected_event_file_path(DAG_ID),
+ trigger_rule="all_done",
+ )
Review Comment:
Yup, already spotted that and fixed locally, thanks !
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org