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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
