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?