This is an automated email from the ASF dual-hosted git repository.
henry3260 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 43129ecb633 Fix clearing with upstream and downstream selecting
unrelated tasks (#74548)
43129ecb633 is described below
commit 43129ecb63355c8244df1ac67b392bc8a775f2fa
Author: Kevin Yang <[email protected]>
AuthorDate: Tue Oct 13 10:19:27 2026 -0400
Fix clearing with upstream and downstream selecting unrelated tasks (#74548)
* Fix clearing with upstream and downstream selecting unrelated tasks
Clearing a task in a specific Dag run with both upstream and downstream
enabled resolved the downstream relatives from the selected tasks plus
their upstream tasks, so every other descendant of those upstream tasks
was selected too. In a Dag whose branches share a root task this
cleared the whole run, and selecting one map index of a mapped task
cleared all of its map indexes.
Each direction has to be resolved from the requested tasks alone, as
clearing without a specific Dag run already does.
* Pass the requested tasks to the relative lookup explicitly
Both directions must be resolved from the same requested tasks.
Passing the sets in makes that visible where the lookups are made,
instead of relying on a helper that reads sets which are updated
after it returns.
* Fix clearing in Dags with a loop selecting unrelated tasks
Any Dag that contains a loop is cleared through the loop-aware
selection, which still resolved downstream relatives from the selected
tasks plus their upstream tasks. Clearing with both upstream and
downstream therefore still selected unrelated tasks in those Dags.
For a task inside a loop it also reached the gate of every iteration,
so the clear replaced the selected iteration and re-ran the loop from
its first iteration, and marking the task with both options changed
every iteration.
* Give the upstream and downstream mapped loop clear its own test
The test name says clearing from the predecessor selects every mapped
occurrence in the loop. That still holds downstream only, but with
upstream and downstream the other map index is now deliberately left
out, so that case no longer fits the name.
---
.../core_api/routes/public/task_instances.py | 33 ++++---
airflow-core/src/airflow/models/taskinstance.py | 5 +-
.../core_api/routes/public/test_task_instances.py | 103 +++++++++++++++++++++
airflow-core/tests/unit/models/test_cleartasks.py | 85 +++++++++++++++--
4 files changed, 204 insertions(+), 22 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
index f50c86b812a..88715c445b2 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
@@ -18,7 +18,7 @@
from __future__ import annotations
from collections import defaultdict
-from collections.abc import Sequence
+from collections.abc import Collection, Iterable, Sequence
from typing import Annotated, Literal, cast
from uuid import UUID
@@ -145,6 +145,7 @@ from airflow.models.taskinstance import (
TaskInstance as TI,
apply_loop_clear_scope,
clear_task_instances,
+ find_relevant_relatives,
select_loop_clear_scope,
)
from airflow.ti_deps.dep_context import DepContext
@@ -1090,19 +1091,20 @@ def post_clear_task_instances(
# Unmapped tasks are expressed in their task_ids (without map_indexes)
normal_task_ids = {t for t in task_markers_to_clear if not
isinstance(t, tuple)}
- def _collect_relatives(run_id: str, direction: Literal["upstream",
"downstream"]) -> None:
- from airflow.models.taskinstance import find_relevant_relatives
-
- relevant_relatives = find_relevant_relatives(
- normal_task_ids,
- mapped_tasks_tuples,
+ def _find_relatives(
+ run_id: str,
+ normal_tasks: Iterable[str],
+ mapped_tasks: Iterable[tuple[str, int]],
+ direction: Literal["upstream", "downstream"],
+ ) -> Collection[str | tuple[str, int]]:
+ return find_relevant_relatives(
+ normal_tasks,
+ mapped_tasks,
dag=dag,
run_id=run_id,
direction=direction,
session=session,
)
- normal_task_ids.update(t for t in relevant_relatives if not
isinstance(t, tuple))
- mapped_tasks_tuples.update(t for t in relevant_relatives if
isinstance(t, tuple))
# We can't easily calculate upstream/downstream map indexes when not
# working for a specific dag run. It's possible by looking at the runs
@@ -1119,10 +1121,19 @@ def post_clear_task_instances(
)
normal_task_ids.update(partial_dag.task_dict)
else:
+ # Both directions must start from the requested tasks only.
Merging upstream results
+ # before the downstream pass would also select every descendant of
those upstream tasks.
+ relatives: set[str | tuple[str, int]] = set()
if upstream:
- _collect_relatives(dag_run_id, "upstream")
+ relatives.update(
+ _find_relatives(dag_run_id, normal_task_ids,
mapped_tasks_tuples, "upstream")
+ )
if downstream:
- _collect_relatives(dag_run_id, "downstream")
+ relatives.update(
+ _find_relatives(dag_run_id, normal_task_ids,
mapped_tasks_tuples, "downstream")
+ )
+ normal_task_ids.update(t for t in relatives if not isinstance(t,
tuple))
+ mapped_tasks_tuples.update(t for t in relatives if isinstance(t,
tuple))
task_markers_to_clear = [
*normal_task_ids,
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index fe68cfe0e48..e594ba5407a 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -785,10 +785,13 @@ def select_loop_clear_scope(
matched_relatives[key] = tuple(matches)
retry.update((other.id, other) for other in matched_relatives[key])
+ # Both directions start from the selection only. Walking downstream from
upstream relatives
+ # would also select their other descendants, including gates whose later
passes then get archived.
+ seeds = tuple(retry.values())
for is_upstream, enabled in ((True, upstream), (False, downstream)):
if not enabled:
continue
- for ti in tuple(retry.values()):
+ for ti in seeds:
task = resolver.get_task(dag_id, run_id, ti.task_id,
dag_version_id=ti.dag_version_id)
add_relatives(ti, task.get_flat_relatives(upstream=is_upstream))
if include_setups_and_teardowns:
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index 9c0e5e6d664..952ed417fd9 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -5443,6 +5443,53 @@ class
TestPostClearTaskInstances(TestTaskInstanceEndpoint):
assert response.status_code == 200
assert response.json()["total_entries"] == group_size
+ @pytest.mark.parametrize(
+ ("mapped", "task_id", "expected"),
+ [
+ pytest.param(False, "a", [("a", -1), ("b", -1), ("root", -1)],
id="unmapped"),
+ pytest.param(True, ["a", 0], [("a", 0), ("b", -1), ("root", -1)],
id="mapped"),
+ ],
+ )
+ @pytest.mark.parametrize("with_loop", [False, True], ids=["without-loop",
"with-unrelated-loop"])
+ def test_clear_with_upstream_and_downstream_excludes_unrelated_tasks(
+ self, test_client, dag_maker, session, mapped, task_id, expected,
with_loop
+ ):
+ """Upstream + downstream must not include other descendants of the
upstream tasks."""
+
+ @task_group
+ def body():
+ BaseOperator(task_id="member")
+
+ dag_id = "clear_upstream_and_downstream"
+ with dag_maker(session=session, dag_id=dag_id,
start_date=DEFAULT_DATETIME_1, serialized=True):
+ if with_loop:
+ create_loop(body, max_iterations=2)
+ root = BaseOperator(task_id="root")
+ if mapped:
+ a = MockOperator.partial(task_id="a").expand(arg2=["x", "y",
"z"])
+ else:
+ a = BaseOperator(task_id="a")
+ b = BaseOperator(task_id="b")
+ other = BaseOperator(task_id="other")
+ root >> a >> b
+ root >> other
+ dr = dag_maker.create_dagrun(run_id="run1",
logical_date=DEFAULT_DATETIME_1)
+
+ response = test_client.post(
+ f"/dags/{dag_id}/clearTaskInstances",
+ json={
+ "dry_run": True,
+ "only_failed": False,
+ "dag_run_id": dr.run_id,
+ "task_ids": [task_id],
+ "include_upstream": True,
+ "include_downstream": True,
+ },
+ )
+ assert response.status_code == 200
+ response_tis = response.json()["task_instances"]
+ assert sorted((ti["task_id"], ti["map_index"]) for ti in response_tis)
== expected
+
class TestGetTaskInstanceTries(TestTaskInstanceEndpoint):
def test_history_without_dag_version_serializes_null(self, test_client,
session):
@@ -5900,6 +5947,62 @@ class
TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
)
assert current_last.state == (State.FAILED if dry_run else None)
+ def test_loop_mark_with_upstream_and_downstream_excludes_unrelated_tasks(
+ self, test_client, dag_maker, session
+ ):
+ @task_group
+ def body():
+ prepare = MockOperator(task_id="prepare")
+ combine = MockOperator(task_id="combine")
+ prepare >> MockOperator(task_id="process") >> combine
+ prepare >> MockOperator(task_id="audit") >> combine
+
+ with dag_maker("mark_loop_upstream_and_downstream", serialized=True):
+ start = MockOperator(task_id="start")
+ loop = create_loop(body, max_iterations=3)
+ start >> MockOperator(task_id="report")
+ start >> loop >> MockOperator(task_id="finish")
+ dr = dag_maker.create_dagrun()
+ root =
session.scalar(select(DynamicRegion).where(DynamicRegion.node_id ==
loop.group_id))
+ for index in (1, 2):
+ for loop_task in loop.iter_tasks():
+ session.add(
+ TaskInstance(
+ task=loop_task,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=root.id,
+ region_index=index,
+ )
+ )
+ session.flush()
+ for ti in dr.get_task_instances(session=session):
+ ti.set_state(State.SUCCESS, session=session)
+ session.commit()
+
+ response = test_client.patch(
+
f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.process/dry_run",
+ json={
+ "new_state": "failed",
+ "region_id": str(root.id),
+ "region_index": 1,
+ "include_upstream": True,
+ "include_downstream": True,
+ },
+ )
+
+ assert response.status_code == 200, response.text
+ assert {
+ (item["task_id"], item.get("region_index")) for item in
response.json()["task_instances"]
+ } == {
+ ("start", None),
+ ("body.prepare", 1),
+ ("body.process", 1),
+ ("body.combine", 1),
+ (loop.gate_task_id, 1),
+ ("finish", None),
+ }
+
def test_loop_mark_without_coordinates_requires_a_region_and_index(self,
test_client, loop_instances):
dr, loop, root, tis = loop_instances
diff --git a/airflow-core/tests/unit/models/test_cleartasks.py
b/airflow-core/tests/unit/models/test_cleartasks.py
index f004638ea60..95d52e56301 100644
--- a/airflow-core/tests/unit/models/test_cleartasks.py
+++ b/airflow-core/tests/unit/models/test_cleartasks.py
@@ -2164,6 +2164,50 @@ def
test_loop_upstream_scope_stays_in_iteration_before_downstream_expansion(loop
}
+def
test_loop_upstream_and_downstream_scope_excludes_other_descendants_of_upstream_tasks(dag_maker,
session):
+ @task_group
+ def body():
+ prepare = EmptyOperator(task_id="prepare")
+ combine = EmptyOperator(task_id="combine")
+ prepare >> EmptyOperator(task_id="process") >> combine
+ prepare >> EmptyOperator(task_id="audit") >> combine
+
+ with dag_maker(serialized=True):
+ start = EmptyOperator(task_id="start")
+ loop = create_loop(body, max_iterations=3)
+ start >> EmptyOperator(task_id="report")
+ start >> loop >> EmptyOperator(task_id="finish")
+ dr = dag_maker.create_dagrun()
+ root = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id ==
loop.group_id))
+ for index in (1, 2):
+ for loop_task in loop.iter_tasks():
+ session.add(
+ TaskInstance(
+ task=loop_task,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=root.id,
+ region_index=index,
+ )
+ )
+ session.flush()
+ tis = {(ti.task_id, ti.region_index): ti for ti in
dr.get_task_instances(session=session)}
+
+ scope = select_loop_clear_scope([tis["body.process", 1]], upstream=True,
downstream=True, session=session)
+
+ assert {key for key, ti in tis.items() if ti.id in scope.retry_ids} == {
+ ("start", -1),
+ ("body.prepare", 1),
+ ("body.process", 1),
+ ("body.combine", 1),
+ (loop.gate_task_id, 1),
+ ("finish", -1),
+ }
+ assert {key for key, ti in tis.items() if ti.id in scope.archive_ids} == {
+ (loop_task.task_id, 2) for loop_task in loop.iter_tasks()
+ }
+
+
def test_outside_seed_upstream_selects_all_current_loop_occurrences(loop_run,
session):
dr, dag, loop, root, tis, iteration = loop_run
outside = next(ti for ti in tis if ti.task_id == "outside")
@@ -2174,8 +2218,8 @@ def
test_outside_seed_upstream_selects_all_current_loop_occurrences(loop_run, se
assert scope.archive_ids == {ti.id for ti in tis if ti.task_id !=
"outside" and iteration(ti) > 0}
[email protected]("upstream", [False, True])
-def test_outside_predecessor_selects_all_mapped_loop_occurrences(dag_maker,
session, upstream):
[email protected]
+def mapped_loop_after_predecessor(dag_maker, session):
@task_group
def body():
PythonOperator.partial(task_id="process",
python_callable=list).expand(op_kwargs=[{}, {}])
@@ -2208,23 +2252,44 @@ def
test_outside_predecessor_selects_all_mapped_loop_occurrences(dag_maker, sess
)
)
session.flush()
- tis = dr.get_task_instances(session=session)
- seed = next(
- ti
- for ti in tis
- if ti.task_id == ("body.process" if upstream else "before")
- and ti.region_index == (0 if upstream else -1)
+ return loop, child, dr.get_task_instances(session=session)
+
+
+def
test_outside_predecessor_selects_all_mapped_loop_occurrences(mapped_loop_after_predecessor,
session):
+ loop, child, tis = mapped_loop_after_predecessor
+ before = next(ti for ti in tis if ti.task_id == "before")
+
+ scope = select_loop_clear_scope(
+ [before],
+ upstream=False,
+ downstream=True,
+ later_loop_iterations=False,
+ session=session,
)
+ assert scope.retry_ids == {ti.id for ti in tis}
+ assert not scope.archive_ids
+
+
+def test_loop_upstream_and_downstream_scope_excludes_other_map_indexes(
+ mapped_loop_after_predecessor, session
+):
+ loop, child, tis = mapped_loop_after_predecessor
+ seed = next(ti for ti in tis if ti.region_id == child.id and
ti.region_index == 0)
+
scope = select_loop_clear_scope(
[seed],
- upstream=upstream,
+ upstream=True,
downstream=True,
later_loop_iterations=False,
session=session,
)
- assert scope.retry_ids == {ti.id for ti in tis}
+ assert scope.retry_ids == {seed.id} | {
+ ti.id
+ for ti in tis
+ if ti.task_id == "before" or (ti.task_id == loop.gate_task_id and
ti.region_index == 1)
+ }
assert not scope.archive_ids
