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