This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit e6662d0a59ea5277e0244bdae48cd670aa90bf95
Author: Ephraim Anierobi <splendidzig...@gmail.com>
AuthorDate: Wed Nov 1 21:37:15 2023 +0100
Fix pre-mature evaluation of tasks in mapped task group (#34337)
* Fix pre-mature evaluation of tasks in mapped task group
Getting the relevant upstream indexes of a task instance in a mapped task
group
should only be done when the task has expanded. If the task has not
expanded yet,
we should return None so that the task can wait for the upstreams before
trying
to run.
This issue is more noticeable when the trigger rule is ONE_FAILED because
then,
the task instance is marked as SKIPPED.
This commit fixes this issue.
closes: https://github.com/apache/airflow/issues/34023
* fixup! Fix pre-mature evaluation of tasks in mapped task group
* fixup! fixup! Fix pre-mature evaluation of tasks in mapped task group
* fixup! fixup! fixup! Fix pre-mature evaluation of tasks in mapped task
group
* Fix tests
(cherry picked from commit 69938fd163045d750b8c218500d79bc89858f9c1)
---
airflow/ti_deps/deps/trigger_rule_dep.py | 18 +++++++++++
tests/models/test_mappedoperator.py | 4 +--
tests/ti_deps/deps/test_trigger_rule_dep.py | 47 +++++++++++++++++++++++++----
3 files changed, 61 insertions(+), 8 deletions(-)
diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py
b/airflow/ti_deps/deps/trigger_rule_dep.py
index 7fbc540adf..3e151cb49c 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -28,6 +28,7 @@ from airflow.models.taskinstance import PAST_DEPENDS_MET
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.deps.base_ti_dep import BaseTIDep, TIDepStatus
from airflow.utils.state import TaskInstanceState
+from airflow.utils.task_group import MappedTaskGroup
from airflow.utils.trigger_rule import TriggerRule as TR
if TYPE_CHECKING:
@@ -131,6 +132,20 @@ class TriggerRuleDep(BaseTIDep):
"""
return ti.task.get_mapped_ti_count(ti.run_id, session=session)
+ def _iter_expansion_dependencies() -> Iterator[str]:
+ from airflow.models.mappedoperator import MappedOperator
+
+ if isinstance(ti.task, MappedOperator):
+ for op in ti.task.iter_mapped_dependencies():
+ yield op.task_id
+ task_group = ti.task.task_group
+ if task_group and task_group.iter_mapped_task_groups():
+ yield from (
+ op.task_id
+ for tg in task_group.iter_mapped_task_groups()
+ for op in tg.iter_mapped_dependencies()
+ )
+
@functools.lru_cache
def _get_relevant_upstream_map_indexes(upstream_id: str) -> int |
range | None:
"""Get the given task's map indexes relevant to the current ti.
@@ -141,6 +156,9 @@ class TriggerRuleDep(BaseTIDep):
"""
if TYPE_CHECKING:
assert isinstance(ti.task.dag, DAG)
+ if isinstance(ti.task.task_group, MappedTaskGroup):
+ if upstream_id not in set(_iter_expansion_dependencies()):
+ return None
try:
expanded_ti_count = _get_expanded_ti_count()
except (NotFullyPopulated, NotMapped):
diff --git a/tests/models/test_mappedoperator.py
b/tests/models/test_mappedoperator.py
index 5509634d17..bbadc57bd0 100644
--- a/tests/models/test_mappedoperator.py
+++ b/tests/models/test_mappedoperator.py
@@ -1303,8 +1303,8 @@ class TestMappedSetupTeardown:
states = self.get_states(dr)
expected = {
"file_transforms.my_setup": {0: "success", 1: "failed", 2:
"skipped"},
- "file_transforms.my_work": {0: "success", 1: "upstream_failed", 2:
"skipped"},
- "file_transforms.my_teardown": {0: "success", 1:
"upstream_failed", 2: "skipped"},
+ "file_transforms.my_work": {2: "upstream_failed", 1:
"upstream_failed", 0: "upstream_failed"},
+ "file_transforms.my_teardown": {2: "success", 1: "success", 0:
"success"},
}
assert states == expected
diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py
b/tests/ti_deps/deps/test_trigger_rule_dep.py
index 2168215ed3..3df5f893be 100644
--- a/tests/ti_deps/deps/test_trigger_rule_dep.py
+++ b/tests/ti_deps/deps/test_trigger_rule_dep.py
@@ -1164,19 +1164,23 @@ def
test_upstream_in_mapped_group_triggers_only_relevant(dag_maker, session):
tis = _one_scheduling_decision_iteration()
assert sorted(tis) == [("tg.t1", 0), ("tg.t1", 1), ("tg.t1", 2)]
- # After running the first t1, the first t2 becomes immediately available.
+ # After running the first t1, the remaining t1 must be run before t2 is
available.
tis["tg.t1", 0].run()
tis = _one_scheduling_decision_iteration()
- assert sorted(tis) == [("tg.t1", 1), ("tg.t1", 2), ("tg.t2", 0)]
+ assert sorted(tis) == [("tg.t1", 1), ("tg.t1", 2)]
- # Similarly for the subsequent t2 instances.
+ # After running all t1, t2 is available.
+ tis["tg.t1", 1].run()
tis["tg.t1", 2].run()
tis = _one_scheduling_decision_iteration()
- assert sorted(tis) == [("tg.t1", 1), ("tg.t2", 0), ("tg.t2", 2)]
+ assert sorted(tis) == [("tg.t2", 0), ("tg.t2", 1), ("tg.t2", 2)]
- # But running t2 partially does not make t3 available.
- tis["tg.t1", 1].run()
+ # Similarly for t2 instances. They both have to complete before t3 is
available
tis["tg.t2", 0].run()
+ tis = _one_scheduling_decision_iteration()
+ assert sorted(tis) == [("tg.t2", 1), ("tg.t2", 2)]
+
+ # But running t2 partially does not make t3 available.
tis["tg.t2", 2].run()
tis = _one_scheduling_decision_iteration()
assert sorted(tis) == [("tg.t2", 1)]
@@ -1406,3 +1410,34 @@ class TestTriggerRuleDepSetupConstraint:
(status,) = self.get_dep_statuses(dr, "w2",
flag_upstream_failed=True, session=session)
assert status.reason.startswith("All setup tasks must complete
successfully")
assert self.get_ti(dr, "w2").state == expected
+
+
+def
test_mapped_tasks_in_mapped_task_group_waits_for_upstreams_to_complete(dag_maker,
session):
+ """Test that one failed trigger rule works well in mapped task group"""
+ with dag_maker() as dag:
+
+ @dag.task
+ def t1():
+ return [1, 2, 3]
+
+ @task_group("tg1")
+ def tg1(a):
+ @dag.task()
+ def t2(a):
+ return a
+
+ @dag.task(trigger_rule=TriggerRule.ONE_FAILED)
+ def t3(a):
+ return a
+
+ t2(a) >> t3(a)
+
+ t = t1()
+ tg1.expand(a=t)
+
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance(task_id="t1")
+ ti.run()
+ dr.task_instance_scheduling_decisions()
+ ti3 = dr.get_task_instance(task_id="tg1.t3")
+ assert not ti3.state
