Make finished TIs more lightweight when scanning task instances for scheduling - #71737

The goal of this PR is to increase scheduler performance.

Scheduler investigation

In our production environment, sometimes schedulers become very slow under load.

During my investigation, I added debug spans for every major step of the scheduler's iteration loop to help understand the issue. There is an open PR with the changes

#69809

I ran various tests and with the help of the debug spans from the PR, I was able to identify a possible bottleneck.

When we have multiple dags with a lot of tasks running in parallel, the scheduler spends most of the time scanning the task instances of the active dag runs to decide which can be queued. It actually looks for SCHEDULED tasks that can be moved to QUEUED state.

Once the tasks are moved to QUEUED state, then the worker can pick them up and it's no longer up to the scheduler.

The scan is consisted of 2 parts. The actual DB query that fetches the data and the SQL-alchemy part which hydrates ORM objects with the query results. Basically, it creates taskinstance objects and fills them with data from each row that we get back. According to the spans, the query is very fast and the entire time is spent on the object hydration part.

The code distinguishes between finished and unfinished tasks but only after the query. It fetches all of them from the DB and creates objects and then splits them into 2 separate lists.

Solution

Let's assume that we have a dag run with 1000 tasks. If 100 tasks have finished, then during the scan we hydrate 100 objects for finished tasks and 900 for unfinished. Similarly, if 400 tasks have finished, then we hydrate 400 objects for finished tasks and 600 for unfinished. As the dag run progresses we get more and more finished tasks.

For finished tasks, we don't really need full ORM objects because we only read 5 fields and we never modify them. Instead of creating a full taskinstance object, we can hydrate a more lightweight immutable object with only the 5 necessary fields.

Instead of having 1 query that fetches everything, creates objects and then splits them into 2 lists, we can have 2 queries, one for unfinished tasks that hydrates full objects and one for finished tasks that hydrates lightweight read-only objects.

There is a scenario with dag callbacks where we need the actual taskinstance object in which case we fetch it from the DB. Still cheaper than the current approach.

Testing

Setup

• 1 scheduler
• 3 celery workers × 18 slots = 54 execution slots

Config

• max_tis_per_query=16 (default)
• max_dagruns_per_loop_to_schedule=20
• parallelism, max_active_tasks_per_dag and default_pool_task_slot_count set to 4096 so that there isn't a limit other than the scheduler itself

Workload

• 35 identical dag runs all start in parallel
• each dag has a 1st deferrable task and then 9 sequential mapped task stages
• when the deferrable task finishes, it returns a number that determines the number of tasks that each stage expands to

The deferrable task doesn't have any special meaning. It just helps to make sure that most of the mapped tasks are running in parallel.

So for example, if the 1st task returns 10, the 1st stage will expand into 10 tasks. Once these 10 finish, then the 2nd expands into 10 tasks. Once, these finish then the 3rd expands into 10 tasks and so on. And this is running for 35 parallel dags.

I used 70 for the number of mapped tasks, which translates into 9stages * 70tasks * 35dags = 22050 tasks and another 35 deferred, 22085 in total.

The code from main takes 35 minutes to run the 22000 tasks, while the code from this patch took 28 minutes. The more the load, the bigger the difference.

Here are some gathered metrics during the test. Left is the code from main and right is the code from this patch.

Time needed for each scan (the PR improvement)

Before it reached up to 94ms while after it never exceeded 20ms.

Scheduler loop duration per iteration

Iterations became faster.

Scheduler vs worker vs worker slot utilization

Scheduler is working at 100% in both cases but before it was spending most of its resources on the ORM hydration and it wouldn't hand enough tasks to workers.

I've set max slot occupancy to 54. Before the change, at any given moment we would have up to 34% of the 54 slots occupied. After, we reach 100% occupancy and that didn't occur once but for most of the time.

Red line is the worker slot occupancy.

Time waiting for the scheduler vs workers

At the end, tasks end up waiting for workers instead of the scheduler.

Increased number of running tasks

Waiting on the scheduler (SCHEDULED -> QUEUED) vs waiting for the workers (QUEUED -> RUNNING)

Tasks wait less on the scheduler and more for a worker.

Same number of tasks, finish faster

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

Claude code Opus 4.8

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.