kacpermuda commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4360121186
##########
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:
Done: `is_loop`, and the loop fields with it, now come from the same
`isinstance` check as `_find_enclosing_loop`, with a test for a group that only
looks like a loop.
##########
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:
Dropped. The region-row fallback is back to main's `in_loop` and
`map_index`; the parent join, the second pass and their tests are gone.
##########
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():
Review Comment:
Placing a scheduler-sent FAIL is query-free again when the task is attached:
no resolver, savepoint or `count(*)`. Those events leave out
`expanded_ti_count` and the loop run fields, which the schema now says. I also
dropped `taskInstance.note`, a lazy load per failure on the scheduler's session
that never reached the manual state changes it was added for.
##########
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
Review Comment:
Fixed in #75458, for 3.4.0, no need to fix here now.
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -3924,38 +4273,901 @@ def body():
create_loop(body, max_iterations=2)
else:
body()
- ti = SimpleNamespace(id=uuid4(), region_id=uuid4(),
task=dag.get_task("body.work"))
+ ti = SimpleNamespace(
+ id=uuid4(),
+ dag_id="run_id_dag",
+ task_id="body.work",
+ try_number=1,
+ region_id=uuid4(),
+ map_index=-1,
+ task=dag.get_task("body.work"),
+ )
+ logical_date = timezone.datetime(2024, 1, 1)
+
+ run_id = get_task_instance_run_id(ti, logical_date=logical_date)
+
+ assert run_id == (
+ str(ti.id)
+ if in_loop
+ else build_task_instance_ol_run_id(
+ dag_id="run_id_dag", task_id="body.work", try_number=1,
logical_date=logical_date, map_index=-1
+ )
+ )
+ assert get_task_instance_map_index(ti) == -1
+
+
+def test_task_instance_run_id_uses_the_map_index_of_the_given_position():
+ ti = SimpleNamespace(id=uuid4(), dag_id="dag", task_id="task",
try_number=2)
+ logical_date = timezone.datetime(2024, 1, 1)
+
+ run_id = get_task_instance_run_id(
+ ti, logical_date=logical_date,
position=TaskInstancePosition(in_loop=False, map_index=3)
+ )
+
+ assert run_id == build_task_instance_ol_run_id(
+ dag_id="dag", task_id="task", try_number=2, logical_date=logical_date,
map_index=3
+ )
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
+def
test_runtime_task_instance_position_without_its_loop_context_has_no_iteration():
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("position_dag") as dag:
+
+ @task_group
+ def body():
+ EmptyOperator(task_id="work")
+
+ create_loop(body, max_iterations=3)
+ ti = SimpleNamespace(
+ id=uuid4(), region_id=uuid4(), map_index=-1,
task=dag.get_task("body.work"), loop=None
+ )
- assert get_regional_task_instance_run_id(ti) == (str(ti.id) if in_loop
else None)
+ assert get_task_instance_position(ti) == TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_group_id="body"
+ )
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
+def
test_runtime_task_instance_position_reads_its_loop_and_expansion_from_the_task_instance():
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("position_dag") as dag:
+
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="work",
python_callable=list).expand(op_kwargs=[{}, {}])
+
+ create_loop(body, max_iterations=3)
+ loop_region_id, forked_from_region_id = uuid4(), uuid4()
+ ti = SimpleNamespace(
+ id=uuid4(),
+ region_id=uuid4(),
+ map_index=1,
+ task=dag.get_task("body.work"),
+ expanded_ti_count=2,
+ loop=SimpleNamespace(
+ index=2,
+ max_iterations=3,
+ region_id=loop_region_id,
+ trajectory=2,
+ forked_from_region_id=forked_from_region_id,
+ resumes_from_index=1,
+ ),
+ )
+
+ assert get_task_instance_position(ti) == TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ expanded_ti_count=2,
+ loop_index=2,
+ loop_group_id="body",
+ loop_region_id=loop_region_id,
+ loop_trajectory=2,
+ loop_forked_from_region_id=forked_from_region_id,
+ loop_resumes_from_index=1,
+ )
+
+
[email protected]_test
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
[email protected]("holds_another_version", [False, True])
+def test_stored_task_instance_is_placed_as_airflow_places_it_for_the_worker(
+ dag_maker, session, holds_another_version
+):
+ from airflow.models.dynamic_region import DynamicRegion
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with dag_maker(serialized=True):
+
+ @task_group
+ def body():
+ EmptyOperator(task_id="plain") >> PythonOperator.partial(
+ task_id="mapped", python_callable=list
+ ).expand(op_kwargs=[{}, {}])
+
+ create_loop(body, max_iterations=3)
+ PythonOperator.partial(task_id="outside",
python_callable=list).expand(op_kwargs=[{}, {}, {}])
+ run = dag_maker.create_dagrun()
+ tis = {(ti.task_id, ti.region_index): ti for ti in run.task_instances}
+ # As the scheduler's Dag bag sets it on the Dag it loads for the run,
which can be a later version than the
+ # one the task instances pin
+ tis["body.plain", 1].task.dag.dag_version_id = (
+ uuid7() if holds_another_version else tis["body.plain",
1].dag_version_id
+ )
+ loop_region = session.get(DynamicRegion, tis["body.plain", 1].region_id)
+ fork = DynamicRegion.fork_family(loop_region.id, 2, {loop_region.id:
loop_region}, session=session)
+ session.flush()
+ rerun = TaskInstance(
+ task=tis["body.plain", 1].task, run_id=run.run_id,
dag_version_id=tis["body.plain", 1].dag_version_id
+ )
+ rerun.region_id, rerun.region_index = fork.id, 2
+ session.add(rerun)
+ session.flush()
+ in_loop = {"in_loop": True, "loop_trajectory": 1, "loop_group_id": "body"}
+
+ # Placing reuses the Dag the task instance holds: its region rows, then
the count of a mapped expansion
+ with assert_queries_count(4):
+ plain = get_stored_task_instance_position(tis["body.plain", 1],
session=session)
+ with assert_queries_count(5):
+ mapped = get_stored_task_instance_position(tis["body.mapped", 1],
session=session)
+ outside = get_stored_task_instance_position(tis["outside", 2],
session=session)
+ rerun_position = get_stored_task_instance_position(rerun, session=session)
+
+ assert plain == TaskInstancePosition(map_index=-1, loop_index=1,
loop_region_id=loop_region.id, **in_loop)
+ assert mapped == TaskInstancePosition(
+ map_index=1, expanded_ti_count=2, loop_index=1,
loop_region_id=loop_region.id, **in_loop
+ )
+ assert outside == TaskInstancePosition(in_loop=False, map_index=2,
expanded_ti_count=3)
+ assert rerun_position == TaskInstancePosition(
+ in_loop=True,
+ map_index=-1,
+ loop_index=2,
+ loop_group_id="body",
+ loop_region_id=fork.id,
+ loop_trajectory=2,
+ loop_forked_from_region_id=loop_region.id,
+ loop_resumes_from_index=2,
+ )
+ assert dict(TaskInstanceInfo(rerun, position=rerun_position))["loop"] == {
+ "index": 2,
+ "region_id": str(fork.id),
+ "trajectory": 2,
+ "forked_from_region_id": str(loop_region.id),
+ "resumes_from_index": 2,
+ }
@pytest.mark.db_test
[email protected](not AIRFLOW_V_3_0_PLUS, reason="Dag versions require
Airflow 3")
[email protected]("load_fails", [False, True])
+@patch("airflow.providers.openlineage.utils.utils.log", autospec=True)
+def
test_stored_task_instance_info_reads_its_dag_version_without_risking_the_event(
Review Comment:
The savepoints are gone, and that test with them. Added a db test through
`on_task_instance_failed` with a real loop task instance, asserting runId ==
`taskInstance.id` == `ti.id` and `loop.index`; passing `session=None` in the
listener fails it.
##########
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.
Review Comment:
Fixed differently: the resolver path is gone, and a task instance whose held
Dag version isn't its own is placed from its region rows, which don't depend on
the version (one query, only then). The worker also counts an attempt as looped
when Airflow sent it a loop context, for a worker that parsed a newer file. A
db test through `on_task_instance_failed` uses two versions that differ in loop
membership.
##########
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"),
Review Comment:
The Dag version isn't read at all anymore. A stored task instance's bundle
name and version come from the Dag and the Dag run the event is built from, the
same sources the worker's workload uses, without a query.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1461,6 +1556,8 @@ def get_task_duration(ti):
return {
"airflowState": AirflowStateRunFacet(
dagRunState=dag_run_state,
+ # TODO: keyed by task_id only, so the mapped instances and
loop iterations of one task collapse
+ # into a single entry (the last one by region wins). Needs
the map index and loop iteration here.
Review Comment:
Leaving it for a follow-up PR: it needs a new per-task-instance list in
`airflowState` (map index, loop iteration, state, duration), which I'd rather
design on its own. I'll link it in the TODO.
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -3924,38 +4273,901 @@ def body():
create_loop(body, max_iterations=2)
else:
body()
- ti = SimpleNamespace(id=uuid4(), region_id=uuid4(),
task=dag.get_task("body.work"))
+ ti = SimpleNamespace(
+ id=uuid4(),
+ dag_id="run_id_dag",
+ task_id="body.work",
+ try_number=1,
+ region_id=uuid4(),
+ map_index=-1,
+ task=dag.get_task("body.work"),
+ )
+ logical_date = timezone.datetime(2024, 1, 1)
+
+ run_id = get_task_instance_run_id(ti, logical_date=logical_date)
+
+ assert run_id == (
+ str(ti.id)
+ if in_loop
+ else build_task_instance_ol_run_id(
+ dag_id="run_id_dag", task_id="body.work", try_number=1,
logical_date=logical_date, map_index=-1
+ )
+ )
+ assert get_task_instance_map_index(ti) == -1
+
+
+def test_task_instance_run_id_uses_the_map_index_of_the_given_position():
+ ti = SimpleNamespace(id=uuid4(), dag_id="dag", task_id="task",
try_number=2)
+ logical_date = timezone.datetime(2024, 1, 1)
+
+ run_id = get_task_instance_run_id(
+ ti, logical_date=logical_date,
position=TaskInstancePosition(in_loop=False, map_index=3)
+ )
+
+ assert run_id == build_task_instance_ol_run_id(
+ dag_id="dag", task_id="task", try_number=2, logical_date=logical_date,
map_index=3
+ )
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
+def
test_runtime_task_instance_position_without_its_loop_context_has_no_iteration():
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("position_dag") as dag:
+
+ @task_group
+ def body():
+ EmptyOperator(task_id="work")
+
+ create_loop(body, max_iterations=3)
+ ti = SimpleNamespace(
+ id=uuid4(), region_id=uuid4(), map_index=-1,
task=dag.get_task("body.work"), loop=None
+ )
- assert get_regional_task_instance_run_id(ti) == (str(ti.id) if in_loop
else None)
+ assert get_task_instance_position(ti) == TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_group_id="body"
+ )
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
+def
test_runtime_task_instance_position_reads_its_loop_and_expansion_from_the_task_instance():
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("position_dag") as dag:
+
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="work",
python_callable=list).expand(op_kwargs=[{}, {}])
+
+ create_loop(body, max_iterations=3)
+ loop_region_id, forked_from_region_id = uuid4(), uuid4()
+ ti = SimpleNamespace(
+ id=uuid4(),
+ region_id=uuid4(),
+ map_index=1,
+ task=dag.get_task("body.work"),
+ expanded_ti_count=2,
+ loop=SimpleNamespace(
+ index=2,
+ max_iterations=3,
+ region_id=loop_region_id,
+ trajectory=2,
+ forked_from_region_id=forked_from_region_id,
+ resumes_from_index=1,
+ ),
+ )
+
+ assert get_task_instance_position(ti) == TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ expanded_ti_count=2,
+ loop_index=2,
+ loop_group_id="body",
+ loop_region_id=loop_region_id,
+ loop_trajectory=2,
+ loop_forked_from_region_id=forked_from_region_id,
+ loop_resumes_from_index=1,
+ )
+
+
[email protected]_test
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
[email protected]("holds_another_version", [False, True])
+def test_stored_task_instance_is_placed_as_airflow_places_it_for_the_worker(
+ dag_maker, session, holds_another_version
+):
+ from airflow.models.dynamic_region import DynamicRegion
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with dag_maker(serialized=True):
+
+ @task_group
+ def body():
+ EmptyOperator(task_id="plain") >> PythonOperator.partial(
+ task_id="mapped", python_callable=list
+ ).expand(op_kwargs=[{}, {}])
+
+ create_loop(body, max_iterations=3)
+ PythonOperator.partial(task_id="outside",
python_callable=list).expand(op_kwargs=[{}, {}, {}])
+ run = dag_maker.create_dagrun()
+ tis = {(ti.task_id, ti.region_index): ti for ti in run.task_instances}
+ # As the scheduler's Dag bag sets it on the Dag it loads for the run,
which can be a later version than the
+ # one the task instances pin
+ tis["body.plain", 1].task.dag.dag_version_id = (
+ uuid7() if holds_another_version else tis["body.plain",
1].dag_version_id
+ )
+ loop_region = session.get(DynamicRegion, tis["body.plain", 1].region_id)
+ fork = DynamicRegion.fork_family(loop_region.id, 2, {loop_region.id:
loop_region}, session=session)
+ session.flush()
+ rerun = TaskInstance(
+ task=tis["body.plain", 1].task, run_id=run.run_id,
dag_version_id=tis["body.plain", 1].dag_version_id
+ )
+ rerun.region_id, rerun.region_index = fork.id, 2
+ session.add(rerun)
+ session.flush()
+ in_loop = {"in_loop": True, "loop_trajectory": 1, "loop_group_id": "body"}
+
+ # Placing reuses the Dag the task instance holds: its region rows, then
the count of a mapped expansion
+ with assert_queries_count(4):
+ plain = get_stored_task_instance_position(tis["body.plain", 1],
session=session)
+ with assert_queries_count(5):
+ mapped = get_stored_task_instance_position(tis["body.mapped", 1],
session=session)
+ outside = get_stored_task_instance_position(tis["outside", 2],
session=session)
+ rerun_position = get_stored_task_instance_position(rerun, session=session)
+
+ assert plain == TaskInstancePosition(map_index=-1, loop_index=1,
loop_region_id=loop_region.id, **in_loop)
+ assert mapped == TaskInstancePosition(
+ map_index=1, expanded_ti_count=2, loop_index=1,
loop_region_id=loop_region.id, **in_loop
+ )
+ assert outside == TaskInstancePosition(in_loop=False, map_index=2,
expanded_ti_count=3)
+ assert rerun_position == TaskInstancePosition(
+ in_loop=True,
+ map_index=-1,
+ loop_index=2,
+ loop_group_id="body",
+ loop_region_id=fork.id,
+ loop_trajectory=2,
+ loop_forked_from_region_id=loop_region.id,
+ loop_resumes_from_index=2,
+ )
+ assert dict(TaskInstanceInfo(rerun, position=rerun_position))["loop"] == {
+ "index": 2,
+ "region_id": str(fork.id),
+ "trajectory": 2,
+ "forked_from_region_id": str(loop_region.id),
+ "resumes_from_index": 2,
+ }
@pytest.mark.db_test
[email protected](not AIRFLOW_V_3_0_PLUS, reason="Dag versions require
Airflow 3")
[email protected]("load_fails", [False, True])
+@patch("airflow.providers.openlineage.utils.utils.log", autospec=True)
+def
test_stored_task_instance_info_reads_its_dag_version_without_risking_the_event(
+ mock_log, dag_maker, session, load_fails
+):
+ with dag_maker(session=session):
+ PythonOperator(task_id="work", python_callable=list)
+ ti = dag_maker.create_dagrun().task_instances[0]
+ session.expire(ti, ["dag_version"])
+
+ if load_fails:
+ with patch.object(
+ TaskInstance,
+ "dag_version",
+ new_callable=PropertyMock,
+ side_effect=RuntimeError("Lost connection"),
+ ):
+ info = dict(TaskInstanceInfo(ti,
position=TaskInstancePosition(in_loop=False, map_index=-1)))
+ else:
+ info = dict(TaskInstanceInfo(ti,
position=TaskInstancePosition(in_loop=False, map_index=-1)))
+
+ expected_bundle_name = None if load_fails else "dag_maker"
+ assert (info["dag_bundle_name"], info["dag_version_id"]) ==
(expected_bundle_name, str(ti.dag_version_id))
+ assert mock_log.warning.call_count == (
+ 2 if load_fails else 0
+ ) # Once for the bundle name, once for its version
+ # The failure was rolled back to its savepoint, so the caller's
transaction is still usable
+ assert session.execute(select(literal(1))).scalar() == 1
+
+
[email protected]_test
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
[email protected]("task_id", ["body.plain", "body.inner.mapped"])
+def
test_stored_and_runtime_task_instances_identify_an_execution_alike(dag_maker,
session, task_id):
+ """
+ The scheduler's events about a stored task instance must tell everything
the worker's tell about it.
+
+ A field missing on one side only would leave consumers without it for that
kind of event, while the other
+ side's events made them rely on it.
+ """
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with dag_maker(serialized=True):
+
+ @task_group
+ def body():
+ plain = EmptyOperator(task_id="plain")
+ with TaskGroup("inner"):
+ plain >> PythonOperator.partial(task_id="mapped",
python_callable=list).expand(
+ op_kwargs=[{}, {}]
+ )
+
+ create_loop(body, max_iterations=3)
+ run = dag_maker.create_dagrun()
+ stored = next(ti for ti in run.task_instances if ti.task_id == task_id and
ti.region_index in (0, 1))
+ stored.task.dag.dag_version_id = stored.dag_version_id
+ stored.state, stored.try_number, stored.max_tries =
TaskInstanceState.UP_FOR_RETRY, 1, 2
+ resolver = TaskCoordinateResolver.for_dag(stored.task.dag, session)
+ mapped = stored.task.get_needs_expansion()
+ # What the server sends the worker for this execution, on the task of the
Dag the worker parses
+ runtime = SimpleNamespace(
+ id=stored.id,
+ dag_id=stored.dag_id,
+ run_id=stored.run_id,
+ task_id=stored.task_id,
+ try_number=stored.try_number,
+ max_tries=stored.max_tries,
+ state=stored.state,
+ map_index=stored.region_index if mapped else -1,
+ rendered_map_index=None,
+ region_id=stored.region_id,
+ region_index=stored.region_index,
+ dag_version_id=stored.dag_version_id,
+ bundle_instance=SimpleNamespace(
+ name=stored.dag_version.bundle_name,
version=stored.dag_version.bundle_version
+ ),
+ loop=resolver.get_loop_context(stored, task=stored.task),
+ expanded_ti_count=resolver.get_expanded_ti_count(stored,
task=stored.task),
+ task=dag_maker.dag.get_task(task_id),
+ )
+
+ stored_info = dict(
+ TaskInstanceInfo(stored,
position=get_stored_task_instance_position(stored, session=session))
+ )
+ runtime_info = dict(TaskInstanceInfo(runtime,
position=get_task_instance_position(runtime)))
+
+ # Only stored task instances have these, which describe the attempt's
resources rather than which it is
+ resources = {"duration", "pool", "queued_dttm", "note", "log_url"}
+ assert {k: v for k, v in stored_info.items() if k not in resources} == {
+ k: v for k, v in runtime_info.items() if k not in resources
+ }
+ assert ("expanded_ti_count" in runtime_info) is mapped
+ assert runtime_info["loop"]["trajectory"] == 1
+ identity = ("task_id", "task_group", "loop", "retries", "retry_delay")
+ assert {k: dict(TaskInfo(stored.task))[k] for k in identity} == {
+ k: dict(TaskInfo(runtime.task))[k] for k in identity
+ }
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
[email protected](
+ "dag_has_version", [pytest.param(True, id="airflow_fails"),
pytest.param(False, id="no_version")]
+)
[email protected](
+ "airflow.models.task_coordinates.TaskCoordinateResolver.for_dag",
side_effect=ValueError("Dag not found")
+)
+def
test_stored_task_instance_is_placed_from_its_rows_without_airflow(mock_for_dag,
dag_has_version):
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("fallback_dag") as dag:
+
+ @task_group
+ def body():
+ EmptyOperator(task_id="work")
+
+ create_loop(body, max_iterations=3)
+ ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7())
+ ti.region_id = uuid4()
+ ti.region_index = 1
+ dag.dag_version_id = ti.dag_version_id if dag_has_version else None
+ session = MagicMock(spec=Session)
+
+ assert get_stored_task_instance_position(ti, session=session) ==
TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_index=1, loop_group_id="body"
+ )
+ if dag_has_version:
+ mock_for_dag.assert_called_once_with(dag, session)
+ else:
+ mock_for_dag.assert_not_called()
+ session.execute.assert_not_called()
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
+def
test_stored_mapped_task_in_loop_is_placed_without_iteration_when_its_region_query_fails():
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("stored_mapped_loop_dag") as dag:
+
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="work",
python_callable=list).expand(op_kwargs=[{}, {}])
+
+ create_loop(body, max_iterations=3)
+ ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7())
+ ti.region_id = uuid4()
+ ti.region_index = 1
+ session = MagicMock(spec=Session)
+ session.execute.side_effect = OSError("connection lost")
+
+ positions = get_orm_task_instance_positions([ti], session=session)
+
+ assert positions == {ti.id: TaskInstancePosition(in_loop=True,
map_index=1, loop_group_id="body")}
+ session.begin_nested.assert_called_once_with()
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Region identity requires
Airflow 3.4")
+def
test_stored_task_instance_without_task_still_raises_when_its_region_query_fails():
+ ti = TaskInstance(task=MagicMock(), run_id="run_id",
dag_version_id=uuid7())
+ ti.region_id = uuid4()
+ ti.region_index = 1
+ ti.task = None
+ session = MagicMock(spec=Session)
+ session.execute.side_effect = OSError("connection lost")
+
+ with pytest.raises(OSError, match="connection lost"):
+ get_orm_task_instance_positions([ti], session=session)
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4")
+def test_task_in_a_task_group_of_unknown_type_is_assumed_in_a_loop():
+ task = SimpleNamespace(task_id="work",
task_group=SimpleNamespace(group_id="unknown"))
+ ti = SimpleNamespace(id=uuid4(), region_id=uuid4(), map_index=-1,
task=task)
+
+ assert get_task_instance_position(ti) ==
TaskInstancePosition(in_loop=True, map_index=-1)
+
+
@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region identity requires
Airflow 3.4")
+def test_stored_task_instance_without_task_or_region_has_no_loop_index():
+ ti = TaskInstance(task=MagicMock(), run_id="run_id",
dag_version_id=uuid7())
+ ti.dag_id = "dag_id"
+ ti.task_id = "task_id"
+ ti.region_id = uuid4()
+ ti.region_index = 1
+ ti.task = None
+ session = MagicMock(spec=Session)
+ session.execute.return_value = []
+
+ assert get_orm_task_instance_positions([ti], session=session) == {
+ ti.id: TaskInstancePosition(in_loop=True, map_index=-1,
loop_index=None)
+ }
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
@pytest.mark.parametrize("mapped", [False, True])
-def
test_stored_regional_instance_without_its_task_is_placed_from_its_region(dag_maker,
session, mapped):
[email protected]("in_loop", [False, True])
+def test_runtime_task_instance_info_has_loop_only_inside_a_loop(in_loop,
mapped):
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("loop_index_dag") as dag:
+
+ @task_group
+ def body():
+ if mapped:
+ PythonOperator.partial(task_id="work",
python_callable=list).expand(op_kwargs=[{}, {}])
+ else:
+ EmptyOperator(task_id="work")
+
+ if in_loop:
+ create_loop(body, max_iterations=3)
+ else:
+ body()
+ ti = SimpleNamespace(
+ id=uuid4(),
+ region_id=uuid4(),
+ map_index=1 if mapped else -1,
+ task=dag.get_task("body.work"),
+ )
+ if in_loop:
+ ti.loop = SimpleNamespace(
+ index=2,
+ max_iterations=3,
+ region_id=None,
+ trajectory=None,
+ forked_from_region_id=None,
+ resumes_from_index=None,
+ )
+
+ info = dict(TaskInstanceInfo(ti))
+
+ assert info["map_index"] == (1 if mapped else None)
+ if in_loop:
+ assert info["loop"] == {
+ "index": 2,
+ }
+ else:
+ assert "loop" not in info
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4")
[email protected]("nested", [False, True])
+def test_task_info_has_the_loop_around_the_task(nested):
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ def ready():
+ return True
+
+ with DAG("nested_loop_dag") as dag:
+ with TaskGroup("outer"):
+
+ @task_group
+ def body():
+ if nested:
+
+ @task_group
+ def inner():
+ EmptyOperator(task_id="work")
+
+ inner()
+ else:
+ EmptyOperator(task_id="work")
+
+ create_loop(body, max_iterations=3, until=ready)
+ task_id = "outer.body.inner.work" if nested else "outer.body.work"
+
+ info = dict(TaskInfo(dag.get_task(task_id)))
+
+ assert info["loop"] == {
+ "group_id": "outer.body",
+ "max_iterations": 3,
+ "has_until": True,
+ "gate_task_id": dag.task_group_dict["outer.body"].gate_task_id,
+ "terminal_task_id": task_id,
+ }
+ # task_group is the task's own group and never describes the loop; its
unique id tells whether it is the loop
+ assert (info["task_group"]["group_id"],
info["task_group"]["unique_group_id"]) == (
+ ("inner", "outer.body.inner") if nested else ("body", "outer.body")
+ )
+ assert (info["task_group"]["unique_group_id"] == info["loop"]["group_id"])
is not nested
+ assert "gate_task_id" not in info["task_group"]
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4")
[email protected](
+ ("until", "loop_index", "expected", "runs_out"),
+ [
+ pytest.param(None, 1, {"stop": False}, False, id="no_until_continues"),
+ pytest.param(None, 3, {"stop": True}, False,
id="no_until_stops_at_last"),
+ pytest.param("done", 1, {"stop": True, "reason": "done"}, False,
id="until_stops_with_reason"),
+ pytest.param(False, 3, {"stop": False}, True,
id="until_false_at_last"),
+ ],
+)
+def test_task_instance_info_has_the_decision_of_a_loop_gate(until, loop_index,
expected, runs_out):
+ from airflow.sdk import LoopDecision
+ from airflow.sdk.definitions._internal.loop import LoopGateOperator
+ from airflow.sdk.execution_time.loop import LoopMaxIterationsExceeded,
execute_loop_gate
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+ class GateTaskInstance(SimpleNamespace):
+ loop_decision = RuntimeTaskInstance.loop_decision
+
+ if until is False:
+ until_callable = lambda: False
+ elif until is not None:
+ until_callable = lambda: LoopDecision(stop=True, reason=until)
+ else:
+ until_callable = None
+ gate = LoopGateOperator(task_id="__loop_gate", until=until_callable)
+ ti = GateTaskInstance(id=uuid7(), map_index=-1, _loop_decision=None)
+ context = {"loop": SimpleNamespace(index=loop_index, max_iterations=3),
"ti": ti}
+
+ if runs_out:
+ with pytest.raises(LoopMaxIterationsExceeded):
+ execute_loop_gate(gate, context)
+ else:
+ execute_loop_gate(gate, context)
+ info = dict(TaskInstanceInfo(ti,
position=TaskInstancePosition(in_loop=False, map_index=-1)))
+
+ assert info["loop"]["decision"] == expected
+
+
+def test_task_instance_info_has_no_loop_decision_for_other_tasks():
+ ti = SimpleNamespace(id=uuid7(), map_index=-1)
+
+ info = dict(TaskInstanceInfo(ti,
position=TaskInstancePosition(in_loop=False, map_index=-1)))
+
+ assert "loop" not in info
+
+
[email protected](not AIRFLOW_V_3_0_PLUS, reason="Attempt UUIDs require
Airflow 3")
[email protected]("stored", [False, True])
[email protected]("in_region", [False, True])
+def test_task_instance_info_has_attempt_id_and_region_only_in_a_region(stored,
in_region):
+ region_id = uuid4() if in_region else UUID(int=0)
+ if stored:
+ ti = TaskInstance(task=MagicMock(), run_id="run_id",
dag_version_id=uuid7())
+ else:
+ ti = SimpleNamespace(id=uuid7(), map_index=-1)
+ ti.region_id = region_id
+ ti.region_index = 2 if in_region else -1
+
+ info = dict(TaskInstanceInfo(ti,
position=TaskInstancePosition(in_loop=False, map_index=-1)))
+
+ assert info["id"] == str(ti.id)
+ if in_region:
+ assert (info["region_id"], info["region_index"]) == (str(region_id), 2)
+ else:
+ assert "region_id" not in info
+ assert "region_index" not in info
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
+def test_stored_task_instance_in_loop_has_loop_index_and_max_iterations():
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("stored_loop_dag") as dag:
+
+ @task_group
+ def body():
+ EmptyOperator(task_id="work")
+
+ create_loop(body, max_iterations=3)
+ ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7())
+ ti.region_id = uuid4()
+ ti.region_index = 2
+
+ assert get_orm_task_instance_positions([ti], session=None) == {
+ ti.id: TaskInstancePosition(in_loop=True, map_index=-1, loop_index=2,
loop_group_id="body")
+ }
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
[email protected]("with_session", [True, False])
+def
test_stored_mapped_task_in_loop_gets_its_iteration_from_its_region(with_session):
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("stored_mapped_loop_dag") as dag:
+
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="work",
python_callable=list).expand(op_kwargs=[{}, {}])
+
+ create_loop(body, max_iterations=3)
+ ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7())
+ ti.region_id = uuid4()
+ ti.region_index = 1
+ session = MagicMock(spec=Session)
+ session.execute.return_value = [_region_row(ti.region_id, "body.work",
parent=(uuid4(), 2, "body"))]
+
+ positions = get_orm_task_instance_positions([ti], session=session if
with_session else None)
+
+ assert positions == {
+ ti.id: TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ loop_index=2 if with_session else None,
+ loop_group_id="body",
+ )
+ }
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Region identity requires
Airflow 3.4")
+def test_task_instance_info_of_stored_instance_without_task_does_not_raise():
+ ti = TaskInstance(task=MagicMock(), run_id="run_id",
dag_version_id=uuid7())
+ ti.dag_id = "dag_id"
+ ti.task_id = "task_id"
+ ti.region_id = uuid4()
+ ti.region_index = 2
+ ti.task = None
+
+ info = dict(TaskInstanceInfo(ti))
+
+ assert info["map_index"] is None
+ assert "loop" not in info
+
+
+def test_task_instance_info_uses_the_position_it_is_given():
+ # The task instance's own loop and expansion differ, so the info can only
match the given position
+ ti = SimpleNamespace(
+ id=uuid4(),
+ region_id=uuid4(),
+ map_index=-1,
+ expanded_ti_count=9,
+ loop=SimpleNamespace(index=9, max_iterations=9),
+ )
+ loop_region_id, forked_from_region_id = uuid4(), uuid4()
+ position = TaskInstancePosition(
+ in_loop=True,
+ map_index=1,
+ expanded_ti_count=4,
+ loop_index=2,
+ loop_group_id="body",
+ loop_region_id=loop_region_id,
+ loop_trajectory=2,
+ loop_forked_from_region_id=forked_from_region_id,
+ loop_resumes_from_index=2,
+ )
+
+ info = dict(TaskInstanceInfo(ti, position=position))
+
+ assert (info["map_index"], info["expanded_ti_count"]) == (1, 4)
+ assert info["loop"] == {
+ "index": 2,
+ "region_id": str(loop_region_id),
+ "trajectory": 2,
+ "forked_from_region_id": str(forked_from_region_id),
+ "resumes_from_index": 2,
+ }
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4")
+def test_task_group_info_has_no_loop_attributes():
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("loop_group_dag"):
+
+ @task_group
+ def body():
+ EmptyOperator(task_id="work")
+
+ loop = create_loop(body, max_iterations=3)
+
+ @task_group
+ def plain():
+ EmptyOperator(task_id="work")
+
+ plain_group = plain()
+
+ loop_attributes = {"max_iterations", "has_until", "terminal_task_id",
"gate_task_id"}
+ loop_info = dict(TaskGroupInfo(loop))
+ assert loop_info["group_id"] == "body"
+ assert not loop_attributes & set(loop_info)
+ assert not loop_attributes & set(TaskGroupInfo(plain_group))
+
+
+def _region_row(region_id, node_id, *, parent=None,
grandparent_region_id=None):
+ """Build a row of the region query: ``parent`` is ``(parent_region_id,
parent_region_index, parent_node_id)``."""
+ parent_region_id, parent_region_index, parent_node_id = parent or (None,
None, None)
+ return SimpleNamespace(
+ id=region_id,
+ node_id=node_id,
+ parent_region_id=parent_region_id,
+ parent_region_index=parent_region_index,
+ parent_node_id=parent_node_id,
+ grandparent_region_id=grandparent_region_id,
+ )
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Region identity requires
Airflow 3.4")
[email protected](
+ ("own_region", "parent", "grandparent_region_id", "map_index"),
+ [
+ pytest.param(True, (uuid4(), 2, "outer"), uuid4(), 1,
id="mapped_task_two_levels_deep"),
+ pytest.param(False, (uuid4(), 2, "outer"), None, -1,
id="enclosing_region_not_top_level"),
+ pytest.param(False, None, None, -1, id="region_missing"),
+ ],
+)
+def
test_stored_task_instance_in_an_unexpected_region_shape_is_placed_without_iteration(
+ own_region, parent, grandparent_region_id, map_index
+):
+ ti = TaskInstance(task=MagicMock(), run_id="run_id",
dag_version_id=uuid7())
+ ti.task_id = "work"
+ ti.region_id = uuid4()
+ ti.region_index = 1
+ ti.task = None
+ session = MagicMock(spec=Session)
+ rows = []
+ if parent is not None or own_region:
+ rows = [
+ _region_row(
+ ti.region_id,
+ "work" if own_region else "group",
+ parent=parent,
+ grandparent_region_id=grandparent_region_id,
+ )
+ ]
+ session.execute.return_value = rows
+
+ assert get_orm_task_instance_positions([ti], session=session) == {
+ ti.id: TaskInstancePosition(in_loop=True, map_index=map_index)
+ }
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loop iteration requires
Airflow 3.4")
[email protected](
+ ("parent", "grandparent_region_id"),
+ [
+ pytest.param((uuid4(), 2, "another_loop"), None,
id="parent_is_not_its_loop"),
+ pytest.param((uuid4(), 2, "body"), uuid4(),
id="loop_region_not_top_level"),
+ ],
+)
+def
test_stored_mapped_task_in_loop_keeps_no_iteration_from_an_unexpected_region_shape(
+ parent, grandparent_region_id
+):
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with DAG("stored_mapped_loop_dag") as dag:
+
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="work",
python_callable=list).expand(op_kwargs=[{}, {}])
+
+ create_loop(body, max_iterations=3)
+ ti = TaskInstance(task=dag.get_task("body.work"), dag_version_id=uuid7())
+ ti.region_id = uuid4()
+ ti.region_index = 1
+ session = MagicMock(spec=Session)
+ session.execute.return_value = [
+ _region_row(ti.region_id, "body.work", parent=parent,
grandparent_region_id=grandparent_region_id)
+ ]
+
+ assert get_orm_task_instance_positions([ti], session=session) == {
+ ti.id: TaskInstancePosition(in_loop=True, map_index=1,
loop_group_id="body")
+ }
+
+
[email protected]_test
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4")
[email protected]("task_attached", [False, True])
+def test_region_shapes_of_airflow_are_the_ones_openlineage_places(dag_maker,
session, task_attached):
+ """
+ Pin the regions Airflow creates, which ``get_orm_task_instance_positions``
relies on to place task instances.
+
+ If this fails, Airflow changed which regions it creates or how they nest:
update the region parsing in
+ ``get_orm_task_instance_positions`` and its comment before changing these
expectations.
+ """
+ from airflow.sdk import task, task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ with dag_maker(serialized=True) as dag:
+
+ @task_group
+ def body():
+ plain = EmptyOperator(task_id="plain")
+ mapped = PythonOperator.partial(task_id="mapped",
python_callable=list).expand(op_kwargs=[{}, {}])
+
+ @task_group
+ def mapped_group(x):
+ @task
+ def member(x):
+ return x
+
+ member(x)
+
+ last = EmptyOperator(task_id="last")
+ plain >> mapped >> mapped_group.expand(x=[1, 2]) >> last
+
+ create_loop(body, max_iterations=2)
+
+ @task_group
+ def top_group(x):
+ @task
+ def top_member(x):
+ return x
+
+ top_member(x)
+
+ top_group.expand(x=[1, 2])
+ run = dag_maker.create_dagrun()
+ tis = [ti for ti in run.task_instances if ti.task_id != "body.__loop_gate"]
+ for ti in tis:
+ ti.task = dag.get_task(ti.task_id) if task_attached else None
+
+ positions = get_orm_task_instance_positions(tis, session=session)
+
+ assert {(ti.task_id, ti.region_index): positions[ti.id] for ti in tis} == {
+ ("body.plain", 1): TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_index=1, loop_group_id="body"
+ ),
+ ("body.last", 1): TaskInstancePosition(
+ in_loop=True, map_index=-1, loop_index=1, loop_group_id="body"
+ ),
+ **{
+ (task_id, map_index): TaskInstancePosition(
+ in_loop=True, map_index=map_index, loop_index=1,
loop_group_id="body"
+ )
+ for task_id in ("body.mapped", "body.mapped_group.member")
+ for map_index in (0, 1)
+ },
+ **{
+ ("top_group.top_member", map_index):
TaskInstancePosition(in_loop=False, map_index=map_index)
+ for map_index in (0, 1)
+ },
+ }
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4")
[email protected]("shape", ["nested_loop", "loop_in_mapped_group",
"expand_in_mapped_group"])
+def test_region_shapes_stay_limited_by_the_task_sdk(shape):
+ """
+ Pin the Dag shapes the task SDK rejects, which
``get_orm_task_instance_positions`` relies on.
+
+ If one of them stops raising, a task can sit two region levels deep, which
OpenLineage places as a loop
+ without its iteration: support the new shape in
``get_orm_task_instance_positions``, then update this test.
+ """
+ from airflow.sdk import task, task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
+
+ def define_dag():
+ with DAG(f"{shape}_dag"):
+
+ @task_group
+ def inner():
+ EmptyOperator(task_id="work")
+
+ if shape == "nested_loop":
+
+ @task_group
+ def outer():
+ create_loop(inner, max_iterations=2)
+
+ create_loop(outer, max_iterations=2)
+ elif shape == "loop_in_mapped_group":
+
+ @task_group
+ def mapped_outer(x):
+ create_loop(inner, max_iterations=2)
+
+ mapped_outer.expand(x=[1, 2])
+ else:
+
+ @task_group
+ def mapped_outer(x):
+ @task
+ def member(y):
+ return y
+
+ member.expand(y=[x])
+
+ mapped_outer.expand(x=[1, 2])
+
+ with pytest.raises(NotImplementedError):
Review Comment:
Added `match=`. With the region parsing back to main's, only one rejected
shape matters to our code, a mapped task in a mapped task group, so the test
pins that one.
##########
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:
Agreed, removed. Scheduler-sent events no longer read the Dag version, loop
context, expansion size or region rows through savepoints: the task instance is
placed from its attached task with no query, as on main. Only one without its
task, or holding another Dag version, reads its region rows, in one plain query
as main did. I didn't use `create_session()` instead: in the scheduler thread
it returns the scheduler's own scoped session, and would commit and close it.
##########
providers/openlineage/tests/system/openlineage/example_openlineage_task_loop_dag.py:
##########
@@ -0,0 +1,164 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+DAG with a plain task and task loops, some of their tasks retrying and one
stopped early by `until`
+(requires Airflow 3.4+).
+
+It checks, for a plain task, tasks, a mapped task and a mapped task group in a
loop, and loop gates:
+ - taskInstance's `try_number`, `state`, `map_index` and `loop` of every
event: `loop.index` is the
+ iteration, counted from 1, and `map_index` the mapped position, and a
plain task has no `loop`
+ - task's `loop`, the loop around the task: its `group_id`,
`max_iterations`, `has_until`, `gate_task_id`
+ and `terminal_task_id`, also for a task of a task group inside a loop,
and absent for the plain task
+ - `loop.region_id`, the loop region of the iteration, shared by the tasks
of one iteration, mapped ones
+ included, `loop.trajectory` 1, and no `loop.forked_from_region_id` nor
`loop.resumes_from_index`, as
+ no loop is cleared
Review Comment:
Added the invariants to `check_run_id_integrity` (runId is `taskInstance.id`
in a loop, `loop.region_id` against `region_id` for mapped and unmapped tasks,
one `loop.region_id` per loop run) with unit tests, and trimmed the docstring
to what the JSON asserts. A Dag that clears a gate and checks the rerun
(trajectory 2, fork fields) will come in a follow-up PR with the manual state
change e2e OL dags.
##########
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():
+ return ti.dag_version
+ except Exception:
+ log.warning(
+ "OpenLineage failed to read the Dag version of task instance
%s", getattr(ti, "id", None)
+ )
+ log.debug("Exception details:", exc_info=True)
+ return 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 or
+ # with get_task_instance_position; it is then used as is.
+ self._position = position
+ super().__init__(obj)
+
def _extend_fields(self) -> None:
- map_index = get_task_instance_map_index(self.obj)
+ position = self._position
+ if position is None:
+ # 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 = _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.rendered_map_index = self.get_rendered_map_index(self.obj) if
map_index != -1 else None
self._fields.extend(("map_index", "rendered_map_index"))
+ if position and position.expanded_ti_count is not None:
+ self.expanded_ti_count = position.expanded_ti_count
+ self._fields.append("expanded_ti_count")
+ if AIRFLOW_V_3_0_PLUS and (ti_id := getattr(self.obj, "id", None)) is
not None:
Review Comment:
Done: a stored task instance's `id` is only emitted from 3.4 on, a runtime
one on every 3.x, and the schema says so.
##########
providers/openlineage/tests/system/openlineage/operator.py:
##########
@@ -351,33 +479,47 @@ 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)
except ValueError as e:
if self.fail_fast:
raise
errors[key] = str(e)
+ if self.check_run_ids and errors:
+ # A key that failed its template is missing from
events_by_key, so its runIds would look incomplete
+ log.info("Skipping the runId integrity check until the event
templates match")
+ elif self.check_run_ids:
+ if problems := check_run_id_integrity(
+ events_by_key,
+ repeated_start_keys={key for key in events_by_key if
self._expects_repeated_events(key)},
+ ):
+ message = "\n".join(f" {problem}" for problem in
problems)
+ if self.fail_fast:
+ raise ValueError(f"runId integrity check
failed:\n{message}")
+ errors["<runId integrity>"] = f"\n{message}"
finally:
if self.clear_variables:
Review Comment:
Added a parametrized `execute()` test (default, `fail_fast`,
`check_run_ids=False`, declared repeated STARTs). Each mutation you listed now
fails it, and the mocks are autospec'd.
##########
providers/openlineage/src/airflow/providers/openlineage/facets/AirflowRunFacet.json:
##########
@@ -196,6 +225,73 @@
"map_index": {
"type": "integer"
},
+ "id": {
+ "description": "Id of the task instance attempt (Airflow 3+).",
+ "type": "string"
+ },
+ "state": {
+ "type": "string"
+ },
+ "max_tries": {
+ "type": "integer"
+ },
+ "rendered_map_index": {
+ "type": "string"
+ },
+ "expanded_ti_count": {
+ "description": "Number of instances in the mapped task's expansion
(Airflow 3.4+).",
+ "type": "integer"
+ },
+ "region_id": {
+ "description": "Region of the loop iteration or mapped expansion
(Airflow 3.4+).",
+ "type": "string"
+ },
+ "region_index": {
+ "description": "Position in the region (Airflow 3.4+).",
+ "type": "integer"
+ },
Review Comment:
Described each field in the schema, including what `region_index` stands
for, `loop.region_id` after a clear, `trajectory`, `resumes_from_index` and
`expanded_ti_count` per pass. I kept `region_id` and `region_index` on purpose:
together they are what Airflow's REST API and UI select a task instance in a
region by, and `region_index` is the only iteration a scheduler-sent FAIL
placed from region rows has.
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -2011,6 +2126,201 @@ def test_get_task_groups_details():
assert result == expected
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Loops require Airflow 3.4")
+def test_get_task_groups_details_has_loop_attributes_only_for_a_loop():
+ from airflow.sdk import task_group
+ from airflow.sdk.definitions._internal.loop import create_loop
Review Comment:
Moved the imports to module level under `if AIRFLOW_V_3_4_PLUS:`, dropped
the re-imports, added an `_add_loop()` helper for the repeated loop Dag, and
removed the overlapping placement tests with the code they covered.
##########
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
Review Comment:
`inspect`, `aliased` and `TaskCoordinateResolver` went with their code,
`MappedTaskGroup` is a module-level compat import, and the 3.4-only loop
classes stay inline with a comment saying why.
##########
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)
+ query = (
+ select(
+ DynamicRegion.id,
+ DynamicRegion.node_id,
+ DynamicRegion.parent_region_id,
+ DynamicRegion.parent_region_index,
+ parent.node_id.label("parent_node_id"),
+ parent.parent_region_id.label("grandparent_region_id"),
Review Comment:
Agreed. Two remain: placing from the task, and main's region fallback
(`in_loop`, `map_index`). The schema says which fields scheduler-sent events
lack. Core exposing loop and expansion accessors on a stored task instance
would be the clean way to bring them back later.
##########
providers/openlineage/tests/system/openlineage/example_openlineage_retry_dag.py:
##########
@@ -0,0 +1,73 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+DAG with a task that fails on its first try and succeeds on the retry.
+
+It checks:
+ - the first try emits START and FAIL, the retry emits START and COMPLETE
+ - taskInstance's `try_number` is 1 for the first try and 2 for the retry,
and `max_tries` is 1 for both
+ - task's `retries` and `retry_delay` on every event
+ - taskInstance's `state`: up_for_retry on the FAIL of the first try,
success on the COMPLETE of the retry
+ - each try has its own `runId` (the `runId` check of
OpenLineageTestOperator is enabled by default)
+ - every task event's identity: taskInstance's `id`, `dag_version_id`
(Airflow 3.1+) and bundle, the absence
+ of fields that don't apply, task's `task_id`, `trigger_rule`, `retries`
and `retry_delay`, the Dag and Dag run,
+ and the parent run, plus each task in the Dag's job facet
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+
+from airflow import DAG
+
+try:
+ from airflow.sdk import task
+except ImportError:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+
+from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT
+from system.openlineage.expected_events import get_expected_event_file_path
+from system.openlineage.operator import OpenLineageTestOperator
+
+DAG_ID = "openlineage_retry_dag"
+
+with DAG(
+ dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT,
+ dag_id=DAG_ID,
+ start_date=datetime(2021, 1, 1),
+ schedule=None,
+ catchup=False,
+ default_args={"retries": 1, "retry_delay": timedelta(seconds=1)},
Review Comment:
Done: `retries` and `retry_delay` are on `flaky_task` only.
`OpenLineageTestOperator` now also defaults to `retries=0` itself, so a Dag's
`default_args` can't make a check retry; the other Dags already set 0.
--
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]
