diff --git a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py index 9ccc4adcb133b..8ed5d25e425a7 100644 --- a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py +++ b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py @@ -195,20 +195,18 @@ def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None: # Only the not-yet-expanded summary ti (map_index < 0) needs the broad # "depend on every upstream ti" behavior, so a rule that can be satisfied - # by a subset of upstreams (ONE_SUCCESS / ONE_FAILED / ONE_DONE / - # NONE_FAILED_MIN_ONE_SUCCESS) does not skip it before the mapped task - # group has expanded (see #34023, #39801). Once the ti is expanded, each - # instance must depend on the upstream instance(s) that share its map - # index, otherwise a single upstream failure would wrongly trigger every - # expanded instance (see #50210). The task may be nested in plain task - # groups inside the mapped one (see #39801), so check the closest mapped - # ancestor rather than only the immediate parent group. + # by a subset of upstreams (ONE_SUCCESS / ONE_FAILED / ONE_DONE) does not + # skip it before the mapped task group has expanded (see #34023). Once the + # ti is expanded, each instance must depend on the upstream instance(s) that + # share its map index, otherwise a single upstream failure would wrongly + # trigger every expanded instance (see #50210). The task may be nested in + # plain task groups inside the mapped one (see #39801), so check the closest + # mapped ancestor rather than only the immediate parent group. if map_index < 0 and task.get_closest_mapped_task_group() is not None: is_fast_triggered = task.trigger_rule in ( TR.ONE_SUCCESS, TR.ONE_FAILED, TR.ONE_DONE, - TR.NONE_FAILED_MIN_ONE_SUCCESS, ) if is_fast_triggered and upstream_id not in set( _iter_expansion_dependencies(task_group=task.task_group) @@ -422,6 +420,7 @@ def _evaluate_direct_relatives() -> Iterator[TIDepStatus]: success_setup = upstream_states.success_setup skipped_setup = upstream_states.skipped_setup + task_id_counts: list[tuple[str, int]] | None = None # Optimization: Don't need to hit the database if all upstreams are # "simple" tasks (no task or task group mapping involved). if not any(t.get_needs_expansion() for t in upstream_tasks.values()): @@ -431,7 +430,6 @@ def _evaluate_direct_relatives() -> Iterator[TIDepStatus]: # Regional selections share counts only for identical producer UUID sets; # legacy runs can share by task IDs. Expansion invalidates both forms. cache_key: tuple[str, str, frozenset[str | UUID]] | None = None - task_id_counts: list[tuple[str, int]] | None = None if task.get_closest_mapped_task_group() is None: scope: frozenset[str | UUID] = ( frozenset( @@ -459,6 +457,20 @@ def _evaluate_direct_relatives() -> Iterator[TIDepStatus]: upstream = sum(count for _, count in task_id_counts) upstream_setup = sum(c for t, c in task_id_counts if upstream_tasks[t].is_setup) + # Once an in-group upstream has expanded, the summary ti of a task in a mapped task + # group no longer sees it and would judge every index by the outside upstreams alone. + # For these two rules the expanded tis decide instead. The cost: an outside failure seen + # after the expansion gives one upstream_failed ti per index instead of one in total, + # which is why other rules still decide here. + if ( + map_index < 0 + and trigger_rule in (TR.ALL_DONE_MIN_ONE_SUCCESS, TR.NONE_FAILED_MIN_ONE_SUCCESS) + and task_id_counts is not None + and len(task_id_counts) < len(upstream_tasks) + and task.get_closest_mapped_task_group() is not None + ): + return + upstream_done = done >= upstream changed = False diff --git a/airflow-core/tests/unit/models/test_mappedoperator.py b/airflow-core/tests/unit/models/test_mappedoperator.py index a6b89e039780f..807b0aa2587f7 100644 --- a/airflow-core/tests/unit/models/test_mappedoperator.py +++ b/airflow-core/tests/unit/models/test_mappedoperator.py @@ -1956,15 +1956,16 @@ def group(value): assert states[task_id] == {0: "success", 1: "skipped"} -def test_none_failed_min_one_success_trigger_rule_expands_in_mapped_task_group(dag_maker): - """Regression test for #39801. +@pytest.mark.parametrize( + "trigger_rule", [TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS, TriggerRule.ALL_DONE_MIN_ONE_SUCCESS] +) +def test_min_one_success_trigger_rule_expands_in_mapped_task_group(dag_maker, trigger_rule): + """A ``*_MIN_ONE_SUCCESS`` task in a mapped task group runs once its in-group upstream succeeds. - A task with the ``NONE_FAILED_MIN_ONE_SUCCESS`` trigger rule inside a dynamically - expanded task group must expand and run once its upstream (also inside the group) - succeeds, instead of being skipped prematurely at its unexpanded (map_index -1) - summary ti before the task group has expanded. + It must not be skipped or marked upstream_failed prematurely at its unexpanded + (map_index -1) summary ti. """ - with dag_maker(dag_id="test_none_failed_min_one_success_in_mapped_task_group") as dag: + with dag_maker(dag_id=f"test_{trigger_rule.value}_in_mapped_task_group") as dag: @task def init(): @@ -1976,7 +1977,7 @@ def tg(message): def imsleepy(message): return message - @task(trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS) + @task(trigger_rule=trigger_rule) def imawake(message): pass @@ -1994,14 +1995,17 @@ def imawake(message): assert states["tg.imawake"] == {0: "success", 1: "success", 2: "success"} -def test_none_failed_min_one_success_trigger_rule_expands_in_nested_task_group(dag_maker): +@pytest.mark.parametrize( + "trigger_rule", [TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS, TriggerRule.ALL_DONE_MIN_ONE_SUCCESS] +) +def test_min_one_success_trigger_rule_expands_in_nested_task_group(dag_maker, trigger_rule): """Regression test for #39801, nested-group shape. - Same as above, but the ``NONE_FAILED_MIN_ONE_SUCCESS`` task lives in a plain task + Same as above, but the ``*_MIN_ONE_SUCCESS`` task lives in a plain task group nested inside the mapped one (the exact shape of the issue's reproducer), so its immediate ``task_group`` is not mapped — only an ancestor is. """ - with dag_maker(dag_id="test_none_failed_min_one_success_in_nested_task_group") as dag: + with dag_maker(dag_id=f"test_{trigger_rule.value}_in_nested_task_group") as dag: @task def init(): @@ -2015,7 +2019,7 @@ def imsleepy(message): @task_group(group_id="inner") def inner(message): - @task(trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS) + @task(trigger_rule=trigger_rule) def imawake(message): pass @@ -2035,6 +2039,70 @@ def imawake(message): assert states["tg.inner.imawake"] == {0: "success", 1: "success", 2: "success"} +@pytest.mark.parametrize( + ("trigger_rule", "outside_outcome", "inner_outcome", "expected_state"), + [ + (TriggerRule.ALL_DONE_MIN_ONE_SUCCESS, None, "skip", "skipped"), + (TriggerRule.ALL_DONE_MIN_ONE_SUCCESS, None, "fail", "upstream_failed"), + (TriggerRule.ALL_DONE_MIN_ONE_SUCCESS, "fail", "success", "success"), + (TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS, None, "fail", "upstream_failed"), + (TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS, "skip", "fail", "upstream_failed"), + ], +) +def test_min_one_success_decides_per_index_in_mapped_task_group( + dag_maker, trigger_rule, outside_outcome, inner_outcome, expected_state +): + """A ``*_MIN_ONE_SUCCESS`` task in a mapped group is evaluated per expanded index. + + The middle index of its in-group upstream ends with ``inner_outcome``; an optional + upstream outside the group ends with ``outside_outcome``. + """ + with dag_maker( + dag_id=f"test_{trigger_rule.value}_outside_{outside_outcome}_inner_{inner_outcome}" + ) as dag: + + @task + def init(): + return ["seize", "the", "day"] + + @task + def outside(): + if outside_outcome == "skip": + raise AirflowSkipException + raise ValueError("boom") + + outside_task = outside() if outside_outcome else None + + @task_group(group_id="tg") + def tg(message): + @task + def imsleepy(message): + if message == "the": + if inner_outcome == "skip": + raise AirflowSkipException + if inner_outcome == "fail": + raise ValueError("boom") + return message + + @task(trigger_rule=trigger_rule) + def imawake(message): + pass + + awake = imawake(imsleepy(message)) + if outside_task is not None: + outside_task >> awake + + tg.expand(message=init()) + + dr = dag.test() + + states: dict[str, dict[int, str | None]] = defaultdict(dict) + for ti in dr.get_task_instances(): + states[ti.task_id][ti.map_index] = ti.state + + assert states["tg.imawake"] == {0: "success", 1: expected_state, 2: "success"} + + def test_mapped_operator_retry_delay_default(dag_maker): """ Test that MappedOperator.retry_delay returns default value when not explicitly set.