ashb commented on code in PR #74752:
URL: https://github.com/apache/airflow/pull/74752#discussion_r4358221307
##########
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:
Nothing tests that `execute()` actually runs this. The PR's headline is that
system tests check runId integrity by default, but the only test of `execute()`
(`test_operator_skips_the_run_id_check_when_a_template_fails`) covers the skip
branch; the other 45 tests call `check_run_id_integrity` directly. On a copy,
each of these leaves 46/46 passing: `elif self.check_run_ids:` -> `elif
False:`, the default changed to `check_run_ids=False`,
`repeated_start_keys=set()`, removing the `fail_fast` raise, and replacing
`errors["<runId integrity>"] = ...` with `pass`. So the check can be silently
disabled in every system test.
One `execute()` test with `_check_key` returning events whose START and
COMPLETE runIds differ would bind it: default -> `Event check failures`,
`fail_fast` -> raises, `check_run_ids=False` -> passes.
(`mock.patch("...time.sleep")` and `patch.object(operator, "_check_key")` in
that test file also lack `autospec`.)
##########
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:
`taskInstance.id` is not the id of the attempt that failed for a stored
`TaskInstance` on Airflow 3.0-3.3, so "Airflow 3+" is wrong here. Those
versions call `prepare_db_for_next_try` (which does `self.id = uuid7()`)
*before* `on_task_instance_failed` in `fetch_handle_failure_context` (3.0.6
`taskinstance.py:1945` then `:1952`; 3.3.2 `:1917` then `:1922`). 3.4 reverses
the order (`notify_failure` at `:2727`, rotation after). So a scheduler-sent
FAIL for an attempt that still has retries carries the *next* attempt's id
(state `up_for_retry`), which is what the next START will carry. Probe against
3.2.2 and 3.3.2 trees: START id `01a164b1-5f93-...`, id in the FAIL facet
`01a164b1-5fb4-...`. The runId is unaffected on those versions, but this field
exists to say which execution sent the event.
Suggest emitting `id` for a stored `TaskInstance` only on 3.4+
(`AIRFLOW_V_3_4_PLUS or not isinstance(self.obj, TaskInstance)`), keeping it
for `RuntimeTaskInstance` on every 3.x version, and saying so in the `id`
description in `AirflowRunFacet.json`. Nothing tests this: the unit and system
tests only run against the installed Airflow version.
##########
providers/openlineage/tests/system/openlineage/example_openlineage_task_states_dag.py:
##########
@@ -0,0 +1,76 @@
+# 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 tasks whose events only taskInstance's `state` tells apart (requires
Airflow 3+).
+
+It checks:
+ - a task that skips itself emits START with `state` running, and COMPLETE
with `state` skipped, which
+ a successful task's COMPLETE would carry as success
+ (Airflow 3.0 calls no listener when a task skips itself, so it emits no
COMPLETE there)
+ - a task that raises AirflowFailException with retries left emits START
with `state` running, and FAIL
+ with `state` failed at `try_number` 1 and `max_tries` 2: it isn't
retried although
+ `try_number <= max_tries`
+ - taskInstance's `dag_version_id` of every event, from Airflow 3.1 on
+ - 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
+
+from airflow import DAG
+from airflow.providers.common.compat.sdk import AirflowFailException,
AirflowSkipException
+from airflow.sdk import task
+
+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_task_states_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": 0},
+) as dag:
+
+ @task
+ def skip_itself():
+ raise AirflowSkipException("Intentional skip")
+
+ @task(retries=2)
+ def fail_without_retry():
+ raise AirflowFailException("Intentional failure that skips the retries
left")
+
+ check_events = OpenLineageTestOperator(
+ task_id="check_events",
+ file_path=get_expected_event_file_path(DAG_ID),
+ trigger_rule="all_done",
+ )
Review Comment:
This will fail on Airflow 3.0. `check_run_ids` defaults to `True`, and for a
TASK runId `check_run_id_integrity` requires exactly one COMPLETE/FAIL
(`operator.py:357`, `job_type != "DAG" and not terminals`). The `__af3_0`
template just made `skip_itself`'s COMPLETE `$absent` (3.0 calls no listener
when a task skips itself), so once the templates match its runId group holds
only a START and `check_events` fails. Running the PR's own
`check_run_id_integrity` over events built from that template gives:
`runId skip_itself-uuid-try1 of TASK
`openlineage_task_states_dag.skip_itself` should have one START (...) and one
COMPLETE/FAIL, got: ['START']`
`openlineage_task_states_dag.json` (3.1+) passes. 5ee168897b0 fixed the e2e
failure for the missing COMPLETE; this is the next failure on the same job.
Either pass `check_run_ids=False` on 3.0 for this Dag, or have the checker
treat a terminal key declared `$absent` as an intentionally partial lifecycle
(with a unit test for START + `$absent` terminals). A scan of the other
expected-event files found no other START-without-terminal job.
##########
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:
**Scheduler cost.** `get_stored_task_instance_position` now runs for every
scheduler-sent FAIL (`process_executor_events` and the heartbeat purge), and
this path costs a SAVEPOINT, a fresh `TaskCoordinateResolver` + `DBDagBag`,
region-ancestry selects and, for mapped tasks, a `count(*)` over the whole
expansion. That is 4 statements for a loop body and 5-6 for a mapped task in a
loop (your own `assert_queries_count(4)` / `(5)`), where main needed 0 when the
task is attached.
It isn't only loops: any mapped task instance has a non-sentinel region, so
every mapped TI the scheduler fails goes through here, and the `count(*)` makes
N failures scan O(N^2) rows (0.35 -> 1.17 ms per call from N=200 to N=2000 on
SQLite; PostgreSQL not measured). N is the number of task instances the
scheduler fails in one pass, e.g. a worker pod lost with 500 running mapped TIs
is ~3000 statements inside the scheduler loop's transaction, with nothing
cached between calls (`for_dag` builds a new `DBDagBag` each time).
Cheap options: skip the resolver when the task is not in a loop
(`get_orm_task_instance_positions` already places an attached-task mapped TI
with 0 queries), omit `expanded_ti_count` on scheduler-sent events or memoize
it per `(dag_id, run_id, task_id, region_id)` for the pass, and share one
resolver per pass.
##########
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:
**Loop / runId.** The scheduler places this task instance with the Dag
version it *holds* (`ti.task.dag`), not the one it *pinned*
(`ti.dag_version_id`), and the reasoning in this comment is incomplete: the
region rows don't depend on the version, but loop membership does.
`enclosing_loop(task)` comes from that version's group definition, so if the
held version no longer has the task in a loop, `get_loop_context` returns
`None` without raising, you get `in_loop=False`, and `get_task_instance_run_id`
builds a coordinate hash instead of `str(ti.id)`.
When can the two differ: `get_dag_for_run` resolves to the *latest* version
for any run without a `bundle_version` (`DBDagBag._version_from_dag_run`, i.e.
the default local bundle), while the worker was placed with `ti.dag_version_id`
(`ti_run`). A Dag edit that removes or renames the loop between the worker's
START and a scheduler-detected failure (heartbeat timeout, executor-reported
failure) is enough.
Reproduced with two serialized versions of one Dag (`body.work` in a loop,
then `body` as a plain `TaskGroup`):
- pinned version: `in_loop=True, loop_index=1`, runId == `ti.id`
(`01a164ed-d412-...`)
- held (latest) version: `in_loop=False`, runId `018cc251-f400-757a-...`
So the FAIL lands on a different run than the START, which stays open in the
lineage backend. The loop facet fields (`max_iterations`, `terminal_task_id`)
are also read from the held version even when the runId agrees.
The commit message of 456a27fd158 says task instances pinned to another
version "fall back to reading their region rows"; the code doesn't do that. And
`test_stored_task_instance_is_placed_as_airflow_places_it_for_the_worker[holds_another_version=True]`
only reassigns `dag_version_id` on a Dag with identical content, so it can't
see this: passing `task_instance.dag_version_id` instead of
`dag.dag_version_id` to `get_task` leaves the whole suite green
(`DBDagBag.get_dag` is called 0 times unmutated, 1 time mutated).
Suggestion: build the coordinate from `task_instance` (it already satisfies
`TaskCoordinate`, and it is what `ti_run` passes) and only when
`task_instance.dag_version_id != dag.dag_version_id` resolve against the pinned
version. That is the one case where paying the Dag load is justified. The test
then needs two versions that differ in loop membership, plus a spy on
`DBDagBag.get_dag` (`assert_queries_count` is an upper bound and the identity
map hides the repeated select).
##########
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:
`pytest.raises(NotImplementedError)` has no `match`, so any unrelated
`NotImplementedError` raised while building the Dag keeps this tripwire green.
The three real messages are distinct ("Nested loops are not supported", "A loop
cannot be inside a mapped task group", "operator expansion in an expanded task
group is not yet supported"); parametrize the expected message and pass
`match=`.
##########
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:
This PR adds ~43 function-body imports to this file (and 6 to
test_listener.py), none with a reason. The same `task_group`/`create_loop` pair
is repeated ~17 times; `DynamicRegion` is re-imported at line 4381 although it
is already imported at line 121 under `if AIRFLOW_V_3_4_PLUS:`; `task` is
re-imported at 2233, 5022 and 5091 although it is a top-level import; and
`create_scheduler_dag` is imported inline at 2236 and 2312 while
test_listener.py imports it at module level. Extend the existing `if
AIRFLOW_V_3_4_PLUS:` block at the top instead.
Also `create_loop(body, max_iterations=3)` with the same body is repeated
~15 times (1 in the base file); a small helper returning the Dag, with optional
mapped/nested variants, would remove ~150 lines, and the four mapped-in-loop
placement tests overlap (only the real-database ones measure the code rather
than the stub).
##########
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:
The title is "Fix task loop information", and `airflowState`
(`tasksState`/`tasksDuration`) is still keyed by `task_id`, so the loop passes
and mapped slots of one task collapse into one entry (the last by `(region_id,
region_index)` wins). The TODO records it. Either fix it here, or link an issue
in the TODO so it is tracked: as it stands the Dag-run-level facet is the one
place that still reports a single state for every iteration of a loop.
##########
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:
This test doesn't bind the savepoint. The injected failure is a Python
`RuntimeError` from a `PropertyMock`, never a database error, so "the caller's
transaction is still usable" holds with or without `begin_nested()`. Removing
it from `get_stored_dag_version` leaves test_utils.py green, and so does
removing it from `_get_core_task_instance_position` (the resolver-failure test
uses `MagicMock(spec=Session)` and a mocked resolver). The savepoint matters
when a statement fails on PostgreSQL and poisons the scheduler's
`process_executor_events` transaction, so it needs a real failing statement (a
`@pytest.mark.backend("postgres")` db test), or at least an assertion that
`begin_nested` was entered for a real session.
Related gap: nothing runs a scheduler-style FAIL through
`OpenLineageListener` with a real session.
`test_manual_event_passes_its_stored_position_to_the_airflow_run_facet` uses a
session-less TI, so changing `session=object_session(ti)` to `session=None` at
`listener.py:892` leaves 383 tests green, and "scheduler FAIL carries the same
fields as the worker" is pinned only at the `utils` level. One db test through
`on_task_instance_failed` with a real loop TI, asserting
`taskInstance.loop.trajectory`, `expanded_ti_count` and `runId ==
taskInstance.id`, would cover it.
##########
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:
**Loop (stored vs public index).** `log_url` is read from the stored task
instance, and `TaskInstance.log_url` builds its `/mapped/<n>` segment from
`self.region_index`, which for a non-mapped loop body is the *iteration*. On a
live 3.4.0 stack, a scheduler-sent FAIL (heartbeat kill) of `body.plain` at
iteration 1 had `.../tasks/body.plain/mapped/1?try_number=1`, where the
worker's START has `.../tasks/body.plain?try_number=1`; the first link doesn't
exist. This is the same synonym trap you already handle for `map_index`, and it
contradicts "FAIL events sent by the scheduler carry the same fields as worker
ones".
Build it from the placed `position.map_index` for a stored TI (or drop the
segment when it is -1). A unit test comparing the stored and runtime `log_url`
of a loop body would pin it.
##########
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:
Imports added inside function bodies with no stated reason: `from sqlalchemy
import inspect` (here), `from sqlalchemy.orm import aliased` (line 1818) and
`from airflow.sdk.definitions.taskgroup import MappedTaskGroup` with its
`except ImportError` fallback (line 2420). sqlalchemy is already imported at
module level (`from sqlalchemy.orm.exc import DetachedInstanceError`), and
`SerializedMappedOperator` (lines 83-88) shows the module's top-level
try/except pattern for a compat import. `TaskCoordinateResolver` (line 1735) is
genuinely 3.4+-only, but needs a one-line comment saying so.
##########
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:
**Loop fields the e2e doesn't actually bind.** The docstring says
`loop.region_id` is the loop region "of the iteration, shared by the tasks of
one iteration, mapped ones included". Nothing asserts that, and it isn't
accurate: `loop.region_id` is the loop's physical region, shared by *all*
iterations until a clear forks it (same value for iterations 1-3 in a captured
run). It is a mapped task's `taskInstance.region_id` (its expansion region)
that differs from it.
In both loop JSONs `region_id`, `loop.region_id` and `taskInstance.id` are
only `{{ is_uuid(result) }}`, and `match()` can't compare two fields, so wrong
values pass. On a live 3.4.0 scheduler, replacing
`loop_region_id=loop.region_id` with the task's own `region_id` (in
`_add_runtime_loop_iteration`) still ended `success`. So did a constant
`loop.region_id`, a constant `taskInstance.id`, and forcing `trajectory` to 1
or `forked_from_region_id`/`resumes_from_index` to `None` (those are only
asserted absent).
"no loop is cleared" also means `trajectory > 1`, the fork fields, and the
`region_id` term this PR adds to `Attempt` (it only matters across a clear) are
exercised by no system test, only by unit tests with hand-built objects.
Clearing the gate through the API on the same stack produced trajectory 2 with
the fork fields set and `check_run_id_integrity` reported no problems, so a Dag
step that clears a gate after the first pass is feasible.
Suggest loop invariants in `check_run_id_integrity`: for a TI with `loop`
and no `map_index`, `loop.region_id == region_id`; mapped in a loop,
`loop.region_id != region_id`; `loop.region_id` constant per (loop group,
trajectory); `taskInstance.id == runId` when `loop` is set. Plus one e2e Dag
that clears a gate. Otherwise trim the claims from the docstring. The job-facet
bullet overclaims too: only 4 of the 10 tasks assert
`task_group`/`loop_group`/`is_mapped`/`parse_time_mapped_count`.
##########
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:
**Design question, non-blocking.** There are now three implementations of
"where does this task instance sit": `_get_task_instance_position` (task
definition + stored index), `_get_core_task_instance_position` (Airflow's
`TaskCoordinateResolver`) and this region-row parser with the self-join and the
`without_iteration` second pass. They duplicate parts of the resolver,
`loop_position` and `DynamicRegion.find_physical_coordinate`. I compared them
on linear, forked, mapped-in-loop, mapped-group-member and placeholder rows and
they agree on `in_loop`, `map_index`, `loop_index` and `loop_group_id`. But the
fallback never has `loop.region_id`, `trajectory`, `forked_from_region_id`,
`resumes_from_index` or `expanded_ti_count`, so when the resolver fails, the
scheduler FAIL of a forked-region task instance carries less than its START,
and the schema doesn't say those fields can be absent.
`loop_index` from region rows only reaches an event when the resolver
raised: the API-server path has no `ti.task`, so it builds no `airflow` facet,
and the asset-event path only needs `in_loop` and `map_index`. The `elif
region.grandparent_region_id is None` branch and the "unexpected shape" warning
are bound only by hand-built rows (flipping the branch to `elif True:` fails
one synthetic-row test and not
`test_region_shapes_of_airflow_are_the_ones_openlineage_places`), and
`loop_group_id` is read once in `src` and asserted on ~15 test lines.
Would it be enough for the fallback to be (`in_loop`, `map_index`) as before
this PR, with the facet schema saying which `loop` fields can be missing on
scheduler-sent events? That drops the parent join, the second pass and ~4
tests. Related: should core expose loop/expansion accessors on a stored TI, as
#75132 did for the runtime TI, so the provider isn't hand-building a
`_TaskCoordinate` just to override `dag_version_id`?
##########
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:
`default_args` also gives `check_events` `retries=1`. On a template mismatch
the operator's `finally` deletes the Variables, the retry then runs and fails
with "variable does not exist (no events emitted)", and the e2e log shows only
the latest attempt, so the real mismatch from attempt 1 is hidden (a lane
observed `check_events` start twice in a run where nothing was emitted). Put
`retries=1, retry_delay=...` on `flaky_task` and leave the Dag default at 0;
the expected JSON only lists `flaky_task` events, so it doesn't change.
##########
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:
**Docs / design question, non-blocking.** The facet schema is now the only
description of these fields (the guide was dropped from this PR), and several
mean different things depending on the task:
- `region_index` is the 1-based loop iteration for a non-mapped loop body
but the 0-based map slot for a mapped task (probe: `body.plain` 1,2;
`body.mapped` 0,1). It is described only as "Position in the region".
- `region_id` is the loop's region for a loop body but the mapped task's
*expansion* region for a mapped task, so a mapped task in a loop carries two
different ids (`region_id` and `loop.region_id`). After a clear,
`loop.region_id` is the fork's region, not the family root.
- `loop.index` (1-based), `loop.trajectory` (1 = original, n+1 = nth fork)
and `loop.resumes_from_index` have no description. `expanded_ti_count` counts
the live (non-REMOVED) instances of that expansion, so it is per loop pass, and
is 1 for a not-yet-expanded placeholder.
Is exposing the physical storage coordinate (`region_id`/`region_index`)
intended? `region_index` always duplicates `loop.index` or `map_index`, and
`loop.index` + `map_index` + `loop.region_id` already carry the same
information. If you keep them, descriptions like the above would save consumers
from reverse-engineering the storage model.
##########
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:
`dag_bundle_version` and `dag_bundle_name` each call `get_bundle_info`,
which calls `get_stored_dag_version`. For a stored TI the Dag version is
therefore read twice per event, and when the read fails each call opens its own
savepoint and logs the same warning ("OpenLineage failed to read the Dag
version of task instance ...") twice (probe with `begin_nested` forced to
raise: 2 calls, 2 warnings). Read it once per `TaskInstanceInfo`, e.g. set both
bundle fields in `_extend_fields`.
--
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]
