Repository: apache/airflow

Issue:
EmptyOperator in dynamically mapped TaskGroups does not respect upstream dependencies
### Apache Airflow version

2.6.2

### What happened

When using an EmptyOperator in dynamically mapped TaskGroups (https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/dynamic-task-mapping.html#mapping-over-a-task-group), the EmptyOperator of all branches starts as soon as the first upstream task dependency of the EmptyOperator **in any branch** completes. This causes downstream tasks of the EmptyOperator to start prematurely in all branches, breaking depth-first execution of the mapped TaskGroup.

I have provided a test for this behavior below, by introducing an artificial wait time in a `variable_task`, followed by an `EmptyOperator` in `checkpoint` and a `final` dependent task . 
![image](https://github.com/apache/airflow/assets/97735/e9d202fa-9b79-4766-b778-a8682a891050)


Running this test, during the execution I see this: The `checkpoint` and `final` tasks are already complete, while the upstream `variable_task` in the group is still running.
![image](https://github.com/apache/airflow/assets/97735/ad335ab5-ee91-4e91-805b-69b58e9bcd99)
I have measured the difference of time when of each the branches' `final` tasks execute, and compared them, to cause a failure condition, which you can see failing here in the `assert_branch_waited` task.

By using just a regular Task, one gets the correct behavior.

### What you think should happen instead

In each branch separately, the `EmptyOperator` should wait for its dependency to complete, before it starts. This would be the same behavior as using a regular `Task` for `checkpoint`.

### How to reproduce

Here are test cases in two dags, one with an `EmptyOperator`, showing incorrect behavior, one with a `Task` in sequence instead of the `EmptyOperator`, that has correct behavior.

```python
import time
from datetime import datetime

from airflow import DAG
from airflow.decorators import task, task_group
from airflow.models import TaskInstance
from airflow.operators.empty import EmptyOperator

branches = [1, 2]
seconds_difference_expected = 10

for use_empty_operator in [False, True]:
    dag_id = "test-mapped-group"
    if use_empty_operator:
        dag_id += "-with-emptyoperator"
    else:
        dag_id += "-no-emptyoperator"

    with DAG(
        dag_id=dag_id,
        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):
            """Branch 2 will take > `seconds_difference_expected` seconds, branch 1 will be immediately executed"""

            @task(dag=dag)
            def variable_task(branch_number):
                """Waits `seconds_difference_expected` seconds for branch 2"""
                if branch_number == 2:
                    time.sleep(seconds_difference_expected)
                return branch_number

            variable_task_result = variable_task(branch_number)

            if use_empty_operator:
                # emptyoperator as a checkpoint
                checkpoint_result = EmptyOperator(task_id="checkpoint")
            else:

                @task
                def checkpoint():
                    pass

                checkpoint_result = checkpoint()

            @task(dag=dag)
            def final(ti: TaskInstance = None):
                """Return the time at the task execution"""
                return datetime.now()

            final_result = final()
            variable_task_result >> checkpoint_result >> final_result

            return final_result

        @task(dag=dag)
        def assert_branch_waited(times):
            """Check that the difference of the start times of the final step in each branch
            are at least `seconds_difference_expected`, i.e. the branch waited for all steps
            """
            seconds_difference = abs(times[1] - times[0]).total_seconds()
            if not seconds_difference >= seconds_difference_expected:
                raise RuntimeError(
                    "Branch 2 completed too fast with respect to branch 1: "
                    + f"observed [seconds difference]: {seconds_difference}; "
                    + f"expected [seconds difference]: >= {seconds_difference_expected}"
                )

        mapping_results = mapped_group.expand(branch_number=branches)
        assert_branch_waited(mapping_results)
```

### Operating System

Debian GNU/Linux 11 (bullseye) on docker (official image)

### Versions of Apache Airflow Providers

_No response_

### Deployment

Docker-Compose

### Deployment details

_No response_

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
@flinz @potiuk I want to pick this up. Could you give me a potential direction in which I should start exploring this problem? 
> @flinz @potiuk I want to pick this up. Could you give me a potential direction in which I should start exploring this problem?

Nothing special from my side. Just marked it as good-first issue as it seemed like it - but maybe @flinz as the author will.
@pragnareddye you can use my test case as provided above to check for effectiveness of any fix. One note on this: I would recommend switching the `time.now()` comparison, which is used to measure the times at which the `final` task was running, to something more stable. I had tried using `ti.start_date`, but in practice that was sometimes a fraction of a second shorter than expected (e.g. `9.985.. s` instead of `>=10`) and lead to failure of the test, even when using `Task` instead of `EmptyOperator`. I guess this is interesting behavior in itself, since I would expect the scheduler to have to wait at least the time that the upstream task `variable_task` completes (i.e. 10s) before even scheduling `checkpoint` or `final`. 

From my side, I sadly have not been able to look under the hood of the scheduler enough to be able to tackle this. 
I assume its the scheduler though, and related to task dependencies - since the EmptyOperator seems to be firing off too early, even if dependencies are not finished, in the branch of the mapped task.  

@potiuk if you have any pointers on how to debug the scheduler itself (if my assumptions here are correct) I would be happy to look as well. I am seeing I would have to go through https://github.com/apache/airflow/blob/main/BREEZE.rst to setup local environments probably.
@pragnareddye here is a more stable test, that will fail reliably if the order of execution is not respected in each branch.
```python
import time
from datetime import datetime
from airflow import DAG
from airflow.decorators import task, task_group
from airflow.models import TaskInstance
from airflow.operators.empty import EmptyOperator
from airflow.api.common.trigger_dag import trigger_dag

# seconds for branch 2 to be delayed artificially
wait_time = 5 

branches = [1,2]

for use_empty_operator in [False, True]:
    dag_id = "test-mapped-group"
    if use_empty_operator:
        dag_id += "-with-emptyoperator"
    else:
        dag_id += "-no-emptyoperator"

    with DAG(
        dag_id=dag_id,
        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):
            """Branch 2 will take > `wait_time` seconds, branch 1 will be immediately executed"""

            @task(dag=dag)
            def variable_task(branch_number):
                """Waits `wait_time` seconds for branch 2"""
                if branch_number == 2:
                    time.sleep(wait_time)
                return datetime.now()

            variable_task_result = variable_task(branch_number)

            if use_empty_operator:
                # emptyoperator as a checkpoint
                checkpoint_result = EmptyOperator(task_id="checkpoint")
            else:

                @task
                def checkpoint():
                    pass

                checkpoint_result = checkpoint()

            @task(dag=dag)
            def final(ti: TaskInstance = None):
                """Return the time at the task execution"""
                return datetime.now()

            final_result = final()
            variable_task_result >> checkpoint_result >> final_result

            return {"dt_final": final_result, "dt_variable_task": variable_task_result}

        @task(dag=dag)
        def assert_branch_waited(times):
            """Check that the difference of the times to reach end of `final` in each branch
            are later than the time of the end of `variable_task`, i.e. the later job in each
            branch waited for all steps to complete
            """
            dt_final = times["dt_final"]
            dt_variable_task = times["dt_variable_task"]
            
            for k in range(len(branches)):

                # this should be positive, since we assume that final runs after
                # variable_task in each execution, since they depend on each other
                # through the intermediate checkpoint
                seconds_difference = (dt_final[k] - dt_variable_task[k]).total_seconds()
                
                if seconds_difference < 0:
                    raise RuntimeError(
                        f"Branch {k+1} had broken execution order: "
                        + f"`variable_task` ended {-seconds_difference} seconds AFTER "
                        + f"`final` task"
                    )

        mapping_results = mapped_group.expand(branch_number=branches)
        assert_branch_waited(mapping_results)
```

Besides, I had a successful shot following the breeze tutorial and was able to setup the dev environment. Currently I am stuck since this bug does not really appear similarly in sequential-like execution, so I am having a hard time reproducing it with `dag.test()`, which would allow me to attach my debugger. However, just looking at DEBUG logs I do see something interesting:

As expected, I do see `checkpoint` not being scheduled due to unmet dependencies, while `variable_task` is still running:
```
[2023-07-04T14:36:29.480+0000] {taskinstance.py:1125} DEBUG - Dependencies not met for <TaskInstance: test-mapped-group-with-emptyoperator.branch_run.checkpoint manual__2023-07-04T14:36:28.641629+00:00 map_index=1 [None]>, dependency 'Trigger Rule' FAILED: Task's trigger rule 'all_success' requires all upstream tasks to have succeeded, but found 1 non-success(es). [...] upstream_task_ids={'branch_run.variable_task'}
...
[2023-07-04T14:36:29.482+0000] {taskinstance.py:1125} DEBUG - Dependencies not met for <TaskInstance: test-mapped-group-with-emptyoperator.branch_run.checkpoint manual__2023-07-04T14:36:28.641629+00:00 map_index=0 [None]>, dependency 'Trigger Rule' FAILED: Task's trigger rule 'all_success' requires all upstream tasks to have succeeded, but found 1 non-success(es). [...] upstream_task_ids={'branch_run.variable_task'}
```

Then, once the first branch of `variable_task` finishes this changes to checkpoint being run for the branch with map_index=0, as expected
```
[2023-07-04T14:36:29.702+0000] {dagrun.py:762} DEBUG - number of scheduleable tasks for <DagRun test-mapped-group-with-emptyoperator @ 2023-07-04 14:36:28.641629+00:00: manual__2023-07-04T14:36:28.641629+00:00, state:running, queued_at: 2023-07-04 14:36:28.651891+00:00. externally triggered: True>: 5 task(s)
...
[2023-07-04T14:36:29.711+0000] {taskinstance.py:1125} DEBUG - Dependencies not met for <TaskInstance: test-mapped-group-with-emptyoperator.branch_run.checkpoint manual__2023-07-04T14:36:28.641629+00:00 map_index=1 [None]>, dependency 'Trigger Rule' FAILED: Task's trigger rule 'all_success' requires all upstream tasks to have succeeded, but found 1 non-success(es).  [...] upstream_task_ids={'branch_run.variable_task'}
...
[2023-07-04T14:36:29.712+0000] {taskinstance.py:1135} DEBUG - Dependencies all met for dep_context=None ti=<TaskInstance: test-mapped-group-with-emptyoperator.branch_run.checkpoint manual__2023-07-04T14:36:28.641629+00:00 map_index=0 [None]>
```

Then in the immediately subsequent scheduling loop, all EmptyOperators have been discarded (schedulable tasks drops to 3 tasks) and all `final` steps meet the scheduling criteria, which is unexpected. I indeed am missing any dependency `'Trigger Rule'` appearing in the checks here for the `final` steps downstream of the empty operator.
```
[2023-07-04T14:36:30.766+0000] {dagrun.py:762} DEBUG - number of scheduleable tasks for <DagRun test-mapped-group-with-emptyoperator @ 2023-07-04 14:36:28.641629+00:00: manual__2023-07-04T14:36:28.641629+00:00, state:running, queued_at: 2023-07-04 14:36:28.651891+00:00. externally triggered: True>: 3 task(s)
...
[2023-07-04T14:36:30.767+0000] {taskinstance.py:1144} DEBUG - <TaskInstance: test-mapped-group-with-emptyoperator.branch_run.final manual__2023-07-04T14:36:28.641629+00:00 map_index=1 [None]> dependency 'Not In Retry Period' PASSED: True, The task instance was not marked for retrying.
[2023-07-04T14:36:30.768+0000] {taskinstance.py:1144} DEBUG - <TaskInstance: test-mapped-group-with-emptyoperator.branch_run.final manual__2023-07-04T14:36:28.641629+00:00 map_index=1 [None]> dependency 'Previous Dagrun State' PASSED: True, The task did not have depends_on_past set.
[2023-07-04T14:36:30.768+0000] {taskinstance.py:1135} DEBUG - Dependencies all met for dep_context=None ti=<TaskInstance: test-mapped-group-with-emptyoperator.branch_run.final manual__2023-07-04T14:36:28.641629+00:00 map_index=1 [None]>
[2023-07-04T14:36:30.769+0000] {taskinstance.py:1144} DEBUG - <TaskInstance: test-mapped-group-with-emptyoperator.branch_run.final manual__2023-07-04T14:36:28.641629+00:00 map_index=0 [None]> dependency 'Not In Retry Period' PASSED: True, The task instance was not marked for retrying.
[2023-07-04T14:36:30.770+0000] {taskinstance.py:1144} DEBUG - <TaskInstance: test-mapped-group-with-emptyoperator.branch_run.final manual__2023-07-04T14:36:28.641629+00:00 map_index=0 [None]> dependency 'Previous Dagrun State' PASSED: True, The task did not have depends_on_past set.
[2023-07-04T14:36:30.770+0000] {taskinstance.py:1135} DEBUG - Dependencies all met for dep_context=None ti=<TaskInstance: test-mapped-group-with-emptyoperator.branch_run.final manual__2023-07-04T14:36:28.641629+00:00 map_index=0 [None]>
```
Full logs attached in: [scheduler_logs_debug.csv](https://github.com/apache/airflow/files/11950560/scheduler_logs_debug.csv)

In summary, I am suspecting task dependencies for EmptyOperator (maybe in TaskFlow?) to be the culprit here, but I am unsure how to continue to debug this.


@flinz I looked into this and can you please try below patch. The issue IMO is that EmptyOperator is never executed and marked as success. In the below sample dag when one of the branch of t1 without sleep succeeds then the corresponding EmptyOperator needs to be marked as success. In below code the task_id of t2 is appended to dummy_ti_ids without the map_index since `group.t2` is the task_id all tasks of the group get marked as success triggering t3. The fix would be to append map_index of the corresponding EmptyOperator too so that only that is marked as success with others waiting for the upstream task with sleep to complete.

```patch
diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index 8f3b3a3301..b0b56f3c6c 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -1338,7 +1338,7 @@ class DagRun(Base, LoggingMixin):
                 and not ti.task.on_success_callback
                 and not ti.task.outlets
             ):
-                dummy_ti_ids.append(ti.task_id)
+                dummy_ti_ids.append((ti.task_id, ti.map_index))
             else:
                 schedulable_ti_ids.append((ti.task_id, ti.map_index))
 
@@ -1369,7 +1369,7 @@ class DagRun(Base, LoggingMixin):
                     .where(
                         TI.dag_id == self.dag_id,
                         TI.run_id == self.run_id,
-                        TI.task_id.in_(dummy_ti_ids_chunk),
+                        tuple_in_condition((TI.task_id, TI.map_index), dummy_ti_ids_chunk),
                     )
                     .values(
                         state=TaskInstanceState.SUCCESS,
```

```python
from datetime import datetime

from airflow import DAG
from airflow.models.param import Param
from airflow.operators.empty import EmptyOperator
from airflow.decorators import task, task_group


from datetime import timedelta

from airflow.utils import timezone
from airflow.models.baseoperator import BaseOperator
from custom_trigger import DateTimeTrigger


with DAG(
    dag_id="mapped_task_group",
    start_date=datetime(2021, 1, 1),
    catchup=False,
    schedule_interval=None,
) as dag:

    @task_group
    def group(i):

        @task
        def t1(i):
            if i == 2 or i == 3:
                import time
                time.sleep(20)
            return i

        t2 = EmptyOperator(task_id='t2')

        @task
        def t3(i):
            return i

        t1(i) >> t2 >> t3(i)

    group.expand(i=[1, 2, 3])
```

Resolving patch:
diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -1338,7 +1338,7 @@ def schedule_tis(
                 and not ti.task.on_success_callback
                 and not ti.task.outlets
             ):
-                dummy_ti_ids.append(ti.task_id)
+                dummy_ti_ids.append((ti.task_id, ti.map_index))
             else:
                 schedulable_ti_ids.append((ti.task_id, ti.map_index))
 
@@ -1369,7 +1369,7 @@ def schedule_tis(
                     .where(
                         TI.dag_id == self.dag_id,
                         TI.run_id == self.run_id,
-                        TI.task_id.in_(dummy_ti_ids_chunk),
+                        tuple_in_condition((TI.task_id, TI.map_index), dummy_ti_ids_chunk),
                     )
                     .values(
                         state=TaskInstanceState.SUCCESS,


Test patch:
diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py
--- a/tests/models/test_dagrun.py
+++ b/tests/models/test_dagrun.py
@@ -1811,6 +1811,52 @@ def tg(x):
     ]
 
 
+def test_mapped_task_group_empty_operator(dag_maker, session):
+    """
+    Test that dynamic task inside a dynamic task group only marks
+    the corresponding downstream EmptyOperator as success.
+    """
+
+    literal = [1, 2, 3]
+
+    with dag_maker(session=session) as dag:
+
+        @task_group
+        def tg(x):
+            @task
+            def t1(x):
+                return x
+
+            t2 = EmptyOperator(task_id="t2")
+
+            @task
+            def t3(x):
+                return x
+
+            t1(x) >> t2 >> t3(x)
+
+        tg.expand(x=literal)
+
+    dr = dag_maker.create_dagrun()
+
+    t2_task = dag.get_task("tg.t2")
+    t2_0 = dr.get_task_instance(task_id="tg.t2", map_index=0)
+    t2_0.refresh_from_task(t2_task)
+    assert t2_0.state is None
+
+    t2_1 = dr.get_task_instance(task_id="tg.t2", map_index=1)
+    t2_1.refresh_from_task(t2_task)
+    assert t2_1.state is None
+
+    dr.schedule_tis([t2_0])
+
+    t2_0 = dr.get_task_instance(task_id="tg.t2", map_index=0)
+    assert t2_0.state == TaskInstanceState.SUCCESS
+
+    t2_1 = dr.get_task_instance(task_id="tg.t2", map_index=1)
+    assert t2_1.state is None
+
+
 def test_ti_scheduling_mapped_zero_length(dag_maker, session):
     with dag_maker(session=session):
         task = BaseOperator(task_id="task_1")