EmptyOperator in dynamically mapped TaskGroups does not respect upstream dependencies

[seeholza]

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 .

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.

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.

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

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

[pragnareddye]

@Flinz @potiuk I want to pick this up. Could you give me a potential direction in which I should start exploring this problem?

[potiuk]

@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.

[seeholza]

@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.

[seeholza]

@pragnareddye here is a more stable test, that will fail reliably if the order of execution is not respected in each branch.

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

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.

[tirkarthi]

@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.

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,

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])

[seeholza]

@tirkarthi yes that fixed it, thanks for the pull request!