Hello, In our production environment, schedulers are often very slow under load and my infrastructure team reports extended db locking times.
After investigation and testing, I've come up with an optimization which has proved to improve scheduler performance greatly. Investigation We already have traces and spans for dag runs to help users understand what's going on under their tasks and possibly optimize them. So I thought why not do the same for Airflow internal operations such as the scheduler loop. I've got an open PR that adds a span for every major step of the scheduler's loop iteration. PR: Add optional debug spans for the scheduler loop <https://github.com/apache/airflow/pull/69809> By using the spans in the above PR, I was able to pin-point the performance bottleneck in the part of examining the task instances for scheduling for a particular dag run. The link below points to the exact part of the code where the issue lies. It's where it fetches all the tasks from the DB and then hydrates them into ORM objects. https://github.com/apache/airflow/pull/69809/changes#diff-aa0338f81a481ca6bc69703521c531d593558d4347a958dc00e5f9689d6802a1R1002-R1007 According to the spans, 10% of that time is spent on the query, which is very fast and 90% is spent in the ORM object hydration. Each iteration of the scheduler loop is linear and all operations are taking place in a sequence. We can't get to operation 2 unless operation 1 finishes. And so, the scheduler scans for tasks that can be queued (operation 1) and only after the scan has finished, it sends the tasks to the workers (operation 2). The more time it takes to scan the tasks, the bigger the interval at which tasks are set to QUEUED and picked by the workers. To explain it in another way, if operation 1 takes 5 minutes, then we will run operation 2 every 5 minutes but if it takes 10 minutes, then we will run operation 2 every 10 minutes. To give you an idea of how heavy the scan is, I ran a test with multiple dags, reaching up to 22.000 tasks. For the time needed to execute all these tasks, 44% of the entire scheduler work across all iterations was spent just on hydrating ORM task objects. Sometimes tasks are running for a while and the scan doesn't do any work, but it's still a very important operation that allows the scheduler to pick up dag run changes quickly. We shouldn't skip it, but we can optimize it. Proposed approach The scan fetches all tasks for a dag run from the DB, hydrates ORM objects and then splits them into 2 lists of finished and unfinished tasks. For finished tasks, we don't need full TaskInstance objects because we never modify them and we only ever read 5 fields from them. Essentially we are wasting computing creating heavy objects we don't need. If instead of generating the full TaskInstance object for every finished task, we just create an immutable lightweight object with only the needed fields, performance increases greatly. Instead of 1 query and then a split, we make the split upfront by having 2 queries, one for finished tasks which will hydrate the lightweight objects and one for unfinished tasks which will hydrate the full TaskInstance objects. Everything else stays the same. Based on gathered metrics, I can see that with the current code in main, the workers always have available slots to run tasks and are mostly waiting on the scheduler. With the improvement, the scheduler sets tasks to QUEUED way faster than the workers can handle. The tasks sit in the queue waiting for minutes and the workers are actually becoming the bottleneck. Here is an open PR with the changes. The PR description contains more info regarding testing and screenshots from gathered metrics which display the improvement clearly. PR: Make finished TIs more lightweight when scanning task instances for scheduling <https://github.com/apache/airflow/pull/71737> Any feedback is highly appreciated! This proved to be an improvement in every scenario that I tested. When there isn't much load in the system, the difference isn't noticeable. If people could also give it a try to make sure that it doesn't introduce a regression, that would be great! Thanks, Christos
