fix schedule_downstream_tasks bug

[luoyuliuyin]

closes: #42581

Problem Description

The trigger_rule of task_one_success is one_success. When the upstream node of task_one_success has not yet run, task_one_success is skipped. According to the semantics of one_success, task_one_success should be able to run.

In this scenario, Airflow turns on the schedule_after_task_execution parameter, which means that after the upstream node finishes running, it will try to schedule the downstream node in the current worker.

This problem may occur when task_1 runs faster than task_run. More specifically, it occurs when task_1 finishes running and successfully schedules downstream tasks in the current worker.

Related Code

Below is the code in question


When task_1 is finished, it will try to schedule downstream tasks. First, a partial dag will be generated.

partial_dag = task.dag.partial_subset(
                task.downstream_task_ids,
                include_downstream=True,
                include_upstream=False,
                include_direct_upstream=True,
            )

task => "task_1"
task.downstream_task_ids => "task_2"

include_downstream=True => ["task_2"]
include_upstream=False => ["task_2"]
include_direct_upstream=True => ["task_2", "task_skip", "task_one_success", "task_1"]

So the final partial_dag is ["task_2", "task_skip", "task_one_success", "task_1"]

This partial_dag is incomplete because task_one_success's other upstream node task_run is not in it.Specifically, the include_upstream parameter should not be false

Solution

The correct subgraph division should be as follows, include_upstream=True:

partial_dag = task.dag.partial_subset(
                task.downstream_task_ids,
                include_downstream=True,
                include_upstream=True,
                include_direct_upstream=True,
            )

task => "task_1"
task.downstream_task_ids => "task_2"

include_downstream=True => ["task_2"]
include_upstream=True =>["task_2", "task_skip", "task_one_success", "task_1", "task_run", "branch"]
include_direct_upstream=True => ["task_2", "task_skip", "task_one_success", "task_1", "task_run", "branch"]

So the final partial_dag is ["task_2", "task_skip", "task_one_success", "task_1", "task_run", "branch"]

The final partial_dag should be as follows:

Subgraph pruning will only be performed when the schedule_after_task_execution parameter is turned on. Normal scheduler scheduling will not have this problem.

[romsharon98]

Can you add test that prevent regression?

[luoyuliuyin]

Can you add test that prevent regression?

Thank you for the suggestion. I've added test to prevent regression. Please check the latest commit.

[luoyuliuyin]

Adding include_upstream=True is still unsafe. Here is a bad case where task_one_success is ignored again. Given that partial_subset has little impact on performance, I recommend removing partial_subset completely.

task => "task_0"
task.downstream_task_ids => "task_1"

include_downstream=True => ["task_1", "task_2"]
include_upstream=True =>["task_0", "task_1", "task_2"]
include_direct_upstream=True => ["task_0", "task_1", "task_2", "task_one_success", "task_skip"]

So the final partial_dag is ["task_0", "task_1", "task_2", "task_one_success", "task_skip"]

[shahar1]

Also, DB tests currently fail

[shahar1]

LGTM - I'm ok with merging it after resolving my last nitpick.
@potiuk / @uranusjr / @ephraimbuddy - any objections?

[potiuk]

I have completely no problem. The use of partial_subset caused problems with serialization already for a number of users so removing it seems like a good idaa, and performance-wise it should not be problematic.

@ashb @uranusjr ? Do You have any objections?

[coding-dragon520]

0 downstream tasks scheduled from follow-on schedule check, Actually, my downstream tasks can be scheduled, Production problems have labor to help see, This problem occurs occasionally, the upstream node succeeds, but the downstream node is not scheduled

[shahar1]

0 downstream tasks scheduled from follow-on schedule check, Actually, my downstream tasks can be scheduled, Production problems have labor to help see, This problem occurs occasionally, the upstream node succeeds, but the downstream node is not scheduled

Thanks for reporting! Could you please create a GitHub issue with a minimal example to reproduce it (considering the latest Airflow version)?