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