Fix topological sort

[pierrejeambrun]

Fixes: #55899
Closes #56321

Similarly to https://github.com/apache/airflow/blob/main/task-sdk/src/airflow/sdk/definitions/taskgroup.py#L561, we need an exit condition if the taskgroup is found in usorted graph.

Adjusted test, which indeed were not in the correct topological order.

Testing dag code:

from __future__ import annotations

import datetime

import pendulum

from airflow.sdk import dag, task, task_group


@task
def get_nums() -> list[int]:
    return [1, 2, 4]


@task
def times_2(n: int) -> int:
    return n * 2


@task_group(group_id="process_number")
def process_number(n: int):
    value = times_2(n)
    return value


@task
def log_success() -> None:
    print("Processed successful!")


@dag(
    schedule=None,
    catchup=False,
    start_date=pendulum.datetime(2025, 4, 1, tz="Europe/Copenhagen"),
    dagrun_timeout=datetime.timedelta(minutes=30),
    dag_id="55899_bug",
)
def test():
    nums = get_nums()
    processed = process_number.expand(n=nums)
    processed >> log_success()


test()

Before

After

[dheerajturaga]

I tried to fix this in #56321
however seems like your implementation is simpler. let me know if I need to close mine

[pierrejeambrun]

Yes, I Think this implementation is simpler and should be favored, also it is more in lined with the sdk implementation

[pierrejeambrun]

Cc: @kaxil

[dheerajturaga]

Yes, I Think this implementation is simpler and should be favored, also it is more in lined with the sdk implementation

Agreed! Thanks for fixing this!

[github-actions]

Backport successfully created: v3-1-test

Status	Branch	Result
✅	v3-1-test