BranchPythonOperator skips downstream tasks for all mapped instances in TaskGroup mapping

[matthieucx]

Apache Airflow version

2.5.1, 2.6.0

What happened

Hello!

When using a branching operator in a mapped task group, skipped tasks will be for all mapped instances of the task_group.

Here is an example DAG exhibiting the issue.

When the BranchOperator sets a downstream task as "skipped", it will also do so retroactively.
If branch_a is selected and has time to run before the first time where branch_b is selected, it will do without issue. However, the status of that instance will still be set to skipped. Any subsequent choice of "branch_a" will be skipped.

Logs for such a case below (obtained using the DAG below.).

I am running Airflow v2.5.1.

What you think should happen instead

branch_a selected:

[2023-04-26, 13:58:09 UTC] {python.py:177} INFO - Done. Returned value was: showcase_branching_issues.branch_a
[2023-04-26, 13:58:09 UTC] {python.py:211} INFO - Branch callable return showcase_branching_issues.branch_a
[2023-04-26, 13:58:09 UTC] {skipmixin.py:155} INFO - Following branch showcase_branching_issues.branch_a
[2023-04-26, 13:58:09 UTC] {skipmixin.py:211} INFO - Skipping tasks ['showcase_branching_issues.branch_b']
[2023-04-26, 13:58:09 UTC] {taskinstance.py:1318} INFO - Marking task as SUCCESS. dag_id=branching_issue, task_id=showcase_branching_issues.branch_int, map_index=0, execution_date=20230426T135806, start_date=20230426T135809, end_date=20230426T135809
[2023-04-26, 13:58:09 UTC] {local_task_job.py:208} INFO - Task exited with return code 0
[2023-04-26, 13:58:09 UTC] {taskinstance.py:2578} INFO - 2 downstream tasks scheduled from follow-on schedule check

branch_a running:

[2023-04-26, 13:58:10 UTC] {python.py:177} INFO - Done. Returned value was: None
[2023-04-26, 13:58:10 UTC] {taskinstance.py:1318} INFO - Marking task as SUCCESS. dag_id=branching_issue, task_id=showcase_branching_issues.branch_a, map_index=0, execution_date=20230426T135806, start_date=20230426T135810, end_date=20230426T135810
[2023-04-26, 13:58:10 UTC] {local_task_job.py:208} INFO - Task exited with return code 0
[2023-04-26, 13:58:10 UTC] {taskinstance.py:2578} INFO - 0 downstream tasks scheduled from follow-on schedule check

branch_b selected:

[2023-04-26, 13:58:14 UTC] {python.py:177} INFO - Done. Returned value was: showcase_branching_issues.branch_b
[2023-04-26, 13:58:14 UTC] {python.py:211} INFO - Branch callable return showcase_branching_issues.branch_b
[2023-04-26, 13:58:14 UTC] {skipmixin.py:155} INFO - Following branch showcase_branching_issues.branch_b
[2023-04-26, 13:58:14 UTC] {skipmixin.py:211} INFO - Skipping tasks ['showcase_branching_issues.branch_a']
[2023-04-26, 13:58:14 UTC] {taskinstance.py:1318} INFO - Marking task as SUCCESS. dag_id=branching_issue, task_id=showcase_branching_issues.branch_int, map_index=1, execution_date=20230426T135806, start_date=20230426T135809, end_date=20230426T135814
[2023-04-26, 13:58:14 UTC] {local_task_job.py:208} INFO - Task exited with return code 0
[2023-04-26, 13:58:14 UTC] {taskinstance.py:2578} INFO - 0 downstream tasks scheduled from follow-on schedule check

All branch_a and branch_b instances set to skipped, no task_b instance ran.

---

Branch selection and "skipped" status should be relative to a particular task_group instance.

How to reproduce

Here is a minimal example DAG which showcases the issue:

from datetime import datetime

from airflow.decorators import dag, task, task_group


@dag(
    dag_id="branching_issue",
    schedule=None,
    start_date=datetime(2021, 1, 1),
)
def BranchingIssue():
    @task
    def branch_b():
        pass

    @task
    def branch_a():
        pass

    @task
    def initiate_dynamic_mapping():
        import random

        random_len = random.randint(1, 10)
        return [i for i in range(random_len)]

    @task.branch
    def branch_int(k):
        import time

        branch = "showcase_branching_issues."
        if k % 2 == 0:
            branch += "branch_a"
        else:
            time.sleep(5)
            branch += "branch_b"
        return branch

    @task_group
    def showcase_branching_issues(k):
        selected_branch = branch_int(k)
        selected_branch >> [branch_a(), branch_b()]

    list_k = initiate_dynamic_mapping()
    showcase_branching_issues.expand(k=list_k)


dag = BranchingIssue()

Operating System

Ubuntu 22.04.1 LTS

Versions of Apache Airflow Providers

No response

Deployment

Docker-Compose

Deployment details

No response

Anything else

I tried searching for related issues or fixes in newer/upcoming releases but found nothing, please excuse me if I missed something.

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[matthieucx]

@josh-fell Just tested it on the new 2.6.0 release, the behavior is the same. Would it be possible to update the label ?

[josh-fell]

@josh-fell Just tested it on the new 2.6.0 release, the behavior is the same. Would it be possible to update the label ?

Thanks for the heads up! Updated.

[seeholza]

Can confirm this still exists in 2.6.2 and I independently hit this bug using the following structure.

import time
from datetime import datetime

from airflow import DAG
from airflow.decorators import task, task_group
from airflow.operators.empty import EmptyOperator
from airflow.utils.trigger_rule import TriggerRule

branches = [1, 2, 3]

with DAG(
    dag_id="test-map-branching",
    schedule=None,
    catchup=False,
    start_date=datetime(2023, 1, 1),
    default_args={"retries": 0},
) as dag:

    @task_group(group_id="branch_run")
    def mapped_group(branch_number):

        @task.branch(dag=dag)
        def choose_branch(branch_number, **kwargs):
            print(branch_number)

            if branch_number == 1:
                return "branch_run.first"
            else:
                return "branch_run.other"

        choose_tech_branch_result = choose_branch(branch_number)

        @task(dag=dag)
        def first(branch_number, **kwargs):
            print(branch_number)

        first_result = first(branch_number)
        choose_tech_branch_result >> first_result

        @task(dag=dag)
        def other(branch_number, **kwargs):
            print(branch_number)
            if branch_number == 3:
                time.sleep(5)
            return branch_number

        other_result = other(branch_number)
        choose_tech_branch_result >> other_result

        @task(dag=dag, trigger_rule=TriggerRule.NONE_FAILED)
        def finalize_map(branch_number, **kwargs):
            print(branch_number)
            return branch_number

        finalize_map_result = finalize_map(branch_number)
        [first_result, other_result] >> finalize_map_result

        return finalize_map_result

    @task(dag=dag)
    def finalize_all(group_results):
        print(group_results)

    group_results = mapped_group.expand(branch_number=branches)
    finalizer_result = finalize_all(group_results)

This DAG always skips ALL downstream tasks of a BranchOperator

I have never seen either of tasks first or second being run. Indeed, changing the trigger rule on finalize_map to TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS makes this DAG fail consistently.