[PR] Make finished TIs more lightweight when scanning task instances for scheduling [airflow]

via GitHub Mon, 17 Aug 2026 14:07:29 -0700


xBis7 opened a new pull request, #71737:
URL: https://github.com/apache/airflow/pull/71737

<!--
Thank you for contributing!

Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
https://chris.beams.io/posts/git-commit/

Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.

For user-facing UI changes, please attach before/after screenshots (or a
short
screen recording) so reviewers can assess the visual impact.

Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.

In case of an existing issue, reference it using one of the following:

* closes: #ISSUE
* related: #ISSUE
-->

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

https://github.com/apache/airflow/pull/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?

- [X] Yes (please specify the tool below)

Claude code Opus 4.8

---

* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#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](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.

--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

[PR] Make finished TIs more lightweight when scanning task instances for scheduling [airflow]

Reply via email to