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