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