Skip to content

Fix mapped task group short circuit skipping - #74188

Closed
anzinmhd wants to merge 3 commits into
apache:mainfrom
anzinmhd:fix/74175-mapped-short-circuit
Closed

anzinmhd wants to merge 3 commits into
apache:mainfrom
anzinmhd:fix/74175-mapped-short-circuit

Conversation

@anzinmhd

@anzinmhd anzinmhd commented Oct 4, 2026 •

Copy link
Copy Markdown
Contributor

What is the issue?

When a ShortCircuitOperator is used inside a mapped task group, transitive downstream tasks can incorrectly execute even when the short-circuit condition is False.

The ShortCircuitOperator records all downstream tasks that should be skipped in its SkipMixin XCom. However, NotPreviouslySkippedDep only checks the direct upstream relatives of a task.

As a result, a transitive downstream task with a permissive trigger rule such as all_done can run even though it was explicitly recorded as skipped by the upstream ShortCircuitOperator.

For example:

group.gate
    |
group.a
    |
group.b (all_done)

For a short-circuited map index, group.a is skipped correctly, but group.b can incorrectly run.

What does this PR change?

Update NotPreviouslySkippedDep to inspect all upstream relatives instead of only direct upstream relatives.

This allows the dependency to find the upstream SkipMixin task and honor its recorded skip decision for transitive downstream tasks.

Performance consideration: This change traverses upstream relatives rather than only direct upstream tasks so that SkipMixin decisions can be honored by transitive downstream tasks. This adds ancestor traversal during dependency evaluation, particularly for wide or deep DAGs, but is necessary to correctly handle transitive skip decisions.

Tests

Added regression tests covering:

  • A ShortCircuitOperator inside a mapped task group.
  • A transitive downstream task using TriggerRule.ALL_DONE.
  • Skip decisions for a specific map index.
  • Ensuring the affected map index is skipped.
  • Ensuring the unaffected map index is not skipped.
  • A Branch → Join case ensuring transitive traversal does not incorrectly skip the join task.

Validation

The following checks were run successfully:

  • Focused NotPreviouslySkippedDep unit tests.
  • Ruff checks for the modified files.
  • Ruff formatting checks for the modified files.
  • git diff --check.
  • Manual reproduction of the reported issue.

The manual reproduction now results in:

Map 0:
  gate -> SUCCESS
  a    -> SUCCESS
  b    -> SUCCESS

Map 1:
  gate -> SUCCESS
  a    -> SKIPPED
  b    -> SKIPPED

Closes: #74175


Was generative AI tooling used to co-author this PR?
  • Yes (ChatGPT)

  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@boring-cyborg

boring-cyborg Bot commented Oct 4, 2026

Copy link
Copy Markdown

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
Here are some useful points:

  • Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
  • In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example Dag that shows how users should use it.
  • Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
  • Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
  • Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
  • Be sure to read the Airflow Coding style.
  • Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
    Apache Airflow is a community-driven project and together we are making it better 🚀.
    In case of doubts contact the developers at:
    Mailing List: dev@airflow.apache.org
    Slack: https://s.apache.org/airflow-slack

@shahar1 shahar1 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for picking this up, and for the clear write-up. The analysis in the description is right: NotPreviouslySkippedDep only reads direct parents, so a transitive all_done task inside a mapped group never sees the short-circuit decision. The new test is a valid regression test (it fails on main). But switching to get_flat_relatives as-is regresses the branching case, so this needs one more iteration before it can land.

Blocking — followed only lists direct children, so every branch + join Dag now skips the join (not_previously_skipped_dep.py:49)

Flag any existing test modified to accommodate new behavior — this may indicate a behavioral regression rather than a genuine fix.

(No existing test was modified here. The problem is the opposite: the existing suite has no transitive-branch case, so it stays green while the behaviour changes. Quoting the rule for the regression it describes.)

SkipMixin.skip_all_except writes {"followed": [...]} containing only the direct downstream tasks it followed (task-sdk/src/airflow/sdk/bases/skipmixin.py, follow_task_ids = [t.task_id for t in downstream_tasks if ...]). The dep's followed branch then says "skip anything not in followed". With get_flat_relatives, a branch operator becomes an upstream of every transitive descendant, none of which are in followed, so they all get set to SKIPPED:

branch --> task1 --> join   (none_failed_min_one_success)
   \-----> task2 --/

With this PR applied, after branch follows task1:

NotPreviouslySkippedDep.get_dep_statuses(join) ==
  [TIDepStatus(dep_name='Not Previously Skipped', passed=False,
               reason='Skipping because of previous XCom result from parent task branch')]

so join is skipped. On main the same probe returns [] and join runs. This is the most common branching pattern in Airflow (example_branch_operator is exactly this shape), so it would break on every scheduler loop, not only in mapped groups.

Fix that keeps your mapped-group fix and does not regress branching: walk the flat ancestors, but only honour the followed rule for direct parents. The explicit skipped list (what ShortCircuitOperator writes) is safe to honour from any ancestor because it already enumerates the exact task ids:

        direct_upstream_ids = {t.task_id for t in ti.task.get_direct_relatives(upstream=True)}
        upstream = ti.task.get_flat_relatives(upstream=True)

and in the should_skip block:

                if (
                    parent.task_id in direct_upstream_ids
                    and XCOM_SKIPMIXIN_FOLLOWED in prev_result
                    and ti.task_id not in prev_result[XCOM_SKIPMIXIN_FOLLOWED]
                ):

I verified locally: with that change your new test, the full test_not_previously_skipped_dep.py suite, and the branch + join probe all pass (10/10). Please also add the branch + join case as a test in the same file so this stays covered; the probe I used:

def test_branch_join_transitive_not_skipped(session, dag_maker):
    with dag_maker("probe_branch_join", schedule=None, start_date=pendulum.datetime(2020, 1, 1), session=session):
        branch = BranchPythonOperator(task_id="branch", python_callable=lambda: "task1")
        task1 = EmptyOperator(task_id="task1")
        task2 = EmptyOperator(task_id="task2")
        join = EmptyOperator(task_id="join", trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)
        branch >> [task1, task2]
        task1 >> join
        task2 >> join

    dr = dag_maker.create_dagrun(run_type=DagRunType.MANUAL, state=State.RUNNING)
    tis = {ti.task_id: ti for ti in dr.task_instances}
    run_task_instance(tis["branch"], branch)
    session.flush()

    dep = NotPreviouslySkippedDep()
    assert dep.is_met(tis["task1"], session=session)
    assert not dep.is_met(tis["task2"], session=session)
    assert list(dep.get_dep_statuses(tis["join"], DepContext(), session=session)) == []
    assert dep.is_met(tis["join"], session=session)
    assert tis["join"].state != State.SKIPPED

Smaller observations

  • not_previously_skipped_dep.py:38 — the class docstring still says "any of the task's direct upstream relatives". Update it to match the new traversal.
  • not_previously_skipped_dep.py:49 — get_flat_relatives is a BFS over all ancestors and this dep runs for every task instance the scheduler examines. For wide or deep Dags that is noticeably more work per scheduling pass than the previous direct-parent lookup. Not a blocker for this fix, but worth a sentence in the PR description so the trade-off is on record.
  • test_not_previously_skipped_dep.py:266 — ruff format wants a second blank line before the new test (and reflows the >> chain). The static-checks job will fail on this as-is; run prek run ruff-format --from-ref main locally.

Worth a second look from

This change touches the skip dependency that #73959 just reworked for mapped task groups; folks with the most context here:

  • @kaxil — authored the most recent change to this file (#73959, mapped-group short-circuit fix, 2026-10-01)

They have not been notified; asking for an extra pass is the maintainer's call, and optional.


This review was drafted by an AI-assisted tool and
confirmed by an Apache Airflow maintainer. After you've
addressed the points above and pushed an update, an Apache Airflow
maintainer — a real person — will take the next look
at the PR. The findings cite the project's review criteria;
if you think one of them is mis-applied, please reply on the
PR and a maintainer will weigh in.

More on how Apache Airflow handles maintainer review:
Contributing guide.

from airflow.utils.state import TaskInstanceState

upstream = ti.task.get_direct_relatives(upstream=True)
upstream = ti.task.get_flat_relatives(upstream=True)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blocking: this regresses every branch + join Dag.

SkipMixin.skip_all_except stores {"followed": [...]} with only the direct children it followed. The followed rule below ("skip anything not in followed") now fires for every transitive descendant of a branch operator, so a join downstream of task1 gets SKIPPED:

branch --> task1 --> join   (none_failed_min_one_success)
   \-----> task2 --/

Verified locally on this branch: get_dep_statuses(join) returns passed=False, reason='Skipping because of previous XCom result from parent task branch'; on main it returns [].

Keep the flat walk, but only apply the followed rule to direct parents (the explicit skipped list is safe from any ancestor):

Suggested change
upstream = ti.task.get_flat_relatives(upstream=True)
direct_upstream_ids = {t.task_id for t in ti.task.get_direct_relatives(upstream=True)}
upstream = ti.task.get_flat_relatives(upstream=True)

and guard the followed branch with parent.task_id in direct_upstream_ids and .... With that, your new test, the existing suite and a branch + join probe all pass (10/10). Full probe test in the review body; please add it to this file.

Also worth noting in the description: get_flat_relatives walks all ancestors on every dep evaluation, which is more scheduler work than the previous direct lookup on large Dags.

assert dep.is_met(tis[("group.child", 0)], session=session)
assert tis[("group.child", 0)].state != State.SKIPPED

def test_parent_in_mapped_task_group_skips_transitive_downstream(session, dag_maker):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: ruff format wants two blank lines before this def (and reflows the >> chain below). The static-checks job will fail as-is; prek run ruff-format --from-ref main fixes it.

@anzinmhd

anzinmhd commented Oct 5, 2026

Copy link
Copy Markdown
Contributor Author

@kaxil, the maintainer suggested getting a second look from you since this PR touches the skip-dependency changes from #73959. I've updated NotPreviouslySkippedDep to handle transitive skip decisions while keeping the followed behavior limited to direct parents, and added a Branch → Join regression test. Would appreciate your review when you have a chance.

@anzinmhd

anzinmhd commented Oct 5, 2026

Copy link
Copy Markdown
Contributor Author

@shahar1, I’ve addressed the requested changes from your review and pushed the updates.

The fix now walks transitive upstream relatives while applying the followed rule only to direct parents, and I’ve also added the Branch → Join regression test and updated the NotPreviouslySkippedDep docstring. The requested formatting and validation checks have also been completed.

Could you please take another look when you have a chance? Thanks!

@anzinmhd
anzinmhd requested a review from shahar1 October 5, 2026 14:05
@eladkal eladkal added this to the Airflow 3.4.0 milestone Oct 5, 2026
@eladkal eladkal added the type:bug-fix Changelog: Bug Fixes label Oct 5, 2026
@kaxil

kaxil commented Oct 5, 2026

Copy link
Copy Markdown
Member

Closing in favour of #74283.

The cause you found is the right one. The problem is get_flat_relatives runs for every task in every Dag, so each task below any SkipMixin task does its own XCom query on each scheduling pass. I measured a plain unmapped Dag (a short-circuit and 1999 tasks after it): 1999 queries per pass instead of 1, 0.64 s instead of 0.029 s on SQLite.

#74283 only reads skip decisions for tasks inside a mapped task group, with one query per pass, and your branch + join case is one of its tests. If it doesn't land we can reopen this.

@kaxil kaxil closed this Oct 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

type:bug-fix Changelog: Bug Fixes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Short-circuit in a mapped task group lets later all_done and none_failed tasks run

4 participants