Repository navigation
Fix scheduler memory spike when a Dag run has a large conf - #73907
Conversation
a0a6952 to
923edde
Compare
|
Could someone please take a look at this PR? |
|
Rebased on main to resolve the conflicts from the dynamic-regions changes. The state re-check in Drafted-by: Claude Code (Fable 5.1); reviewed by @luc-pimentel before posting |
923edde to
9d98eea
Compare
|
This PR makes a huge difference in production. My Scheduler was struggling badly with 2-5 MB Confs, so I created a cluster policy based on this PR as a temporary backport, which helped address the issue completly. This is a very valuable improvement, and I highly recommend merging it. |
ashb
left a comment
ashb
left a comment
There was a problem hiding this comment.
I think this is okay, but I'd like a second pair of eyes from someone else familiar with the code base on it before we merge.
TaskInstance.dag_run is a joined eager load, so a task-instance query that also selects dag_run.conf returns and decodes one copy of the conf per row. A scheduling pass does that twice: fetch_task_instances reads every task instance of the run, and _get_ready_tis re-reads the ones still waiting on their upstreams. A large conf on a run with many mapped task instances multiplies scheduler memory this way.
TaskInstance.dag_run is a joined eager load, so every row of a task-instance query also carries the run's conf. The scheduler fetches all task instances of each running Dag run on every loop, and re-reads the ones still waiting on their upstreams, so a large conf was decoded once per task instance: a 1 MB conf on a run with 500 mapped task instances took the scheduler from 150 MB to over 3 GB. Neither query needs the conf up front: the few callers that read it can load it on access, and the re-check only compares states.
|
Rebased to resolve the conflict with #71737; conf deferral stays on the ORM query for unfinished task instances, preserving the new lightweight finished-task query. Drafted-by: OpenCode 1.19.0 (GPT-6 Astra); reviewed by @luc-pimentel before posting |
9d98eea to
af38159
Compare
|
Hi maintainer, this PR was merged without a milestone set.
|
Backport failed to create: v3-4-test. View the failure log Run detailsNote: As of Merging PRs targeted for Airflow 3.X In matter of doubt please ask in #release-management Slack channel.
You can attempt to backport this manually by running: cherry_picker ce01a73 v3-4-testThis should apply the commit to the v3-4-test branch and leave the commit in conflict state marking After you have resolved the conflicts, you can continue the backport process by running: cherry_picker --continueIf you don't have cherry-picker installed, see the installation guide. |
apache#73907) * Add tests for the run conf loaded per task instance when scheduling TaskInstance.dag_run is a joined eager load, so a task-instance query that also selects dag_run.conf returns and decodes one copy of the conf per row. A scheduling pass does that twice: fetch_task_instances reads every task instance of the run, and _get_ready_tis re-reads the ones still waiting on their upstreams. A large conf on a run with many mapped task instances multiplies scheduler memory this way. * Fix scheduler memory spike when a Dag run has a large conf TaskInstance.dag_run is a joined eager load, so every row of a task-instance query also carries the run's conf. The scheduler fetches all task instances of each running Dag run on every loop, and re-reads the ones still waiting on their upstreams, so a large conf was decoded once per task instance: a 1 MB conf on a run with 500 mapped task instances took the scheduler from 150 MB to over 3 GB. Neither query needs the conf up front: the few callers that read it can load it on access, and the re-check only compares states. (cherry picked from commit ce01a73)
|
Any chance this makes it into 3.4.0 ? |
apache#73907) * Add tests for the run conf loaded per task instance when scheduling TaskInstance.dag_run is a joined eager load, so a task-instance query that also selects dag_run.conf returns and decodes one copy of the conf per row. A scheduling pass does that twice: fetch_task_instances reads every task instance of the run, and _get_ready_tis re-reads the ones still waiting on their upstreams. A large conf on a run with many mapped task instances multiplies scheduler memory this way. * Fix scheduler memory spike when a Dag run has a large conf TaskInstance.dag_run is a joined eager load, so every row of a task-instance query also carries the run's conf. The scheduler fetches all task instances of each running Dag run on every loop, and re-reads the ones still waiting on their upstreams, so a large conf was decoded once per task instance: a 1 MB conf on a run with 500 mapped task instances took the scheduler from 150 MB to over 3 GB. Neither query needs the conf up front: the few callers that read it can load it on access, and the re-check only compares states. (cherry picked from commit ce01a73)
#73907) (#76252) * Add tests for the run conf loaded per task instance when scheduling TaskInstance.dag_run is a joined eager load, so a task-instance query that also selects dag_run.conf returns and decodes one copy of the conf per row. A scheduling pass does that twice: fetch_task_instances reads every task instance of the run, and _get_ready_tis re-reads the ones still waiting on their upstreams. A large conf on a run with many mapped task instances multiplies scheduler memory this way. * Fix scheduler memory spike when a Dag run has a large conf TaskInstance.dag_run is a joined eager load, so every row of a task-instance query also carries the run's conf. The scheduler fetches all task instances of each running Dag run on every loop, and re-reads the ones still waiting on their upstreams, so a large conf was decoded once per task instance: a 1 MB conf on a run with 500 mapped task instances took the scheduler from 150 MB to over 3 GB. Neither query needs the conf up front: the few callers that read it can load it on access, and the re-check only compares states. (cherry picked from commit ce01a73) Co-authored-by: Lucas Pimentel <[email protected]>
#73907) (#76252) * Add tests for the run conf loaded per task instance when scheduling TaskInstance.dag_run is a joined eager load, so a task-instance query that also selects dag_run.conf returns and decodes one copy of the conf per row. A scheduling pass does that twice: fetch_task_instances reads every task instance of the run, and _get_ready_tis re-reads the ones still waiting on their upstreams. A large conf on a run with many mapped task instances multiplies scheduler memory this way. * Fix scheduler memory spike when a Dag run has a large conf TaskInstance.dag_run is a joined eager load, so every row of a task-instance query also carries the run's conf. The scheduler fetches all task instances of each running Dag run on every loop, and re-reads the ones still waiting on their upstreams, so a large conf was decoded once per task instance: a 1 MB conf on a run with 500 mapped task instances took the scheduler from 150 MB to over 3 GB. Neither query needs the conf up front: the few callers that read it can load it on access, and the re-check only compares states. (cherry picked from commit ce01a73) Co-authored-by: Lucas Pimentel <[email protected]>
A Dag run with a large conf made the scheduler's memory grow with the number of task instances in the run.
TaskInstance.dag_runis a joined eager load, so every row of a task-instance query also carries the run's conf. On every loop the scheduler reads all task instances of each running Dag run throughDagRun.fetch_task_instancesinairflow-core/src/airflow/models/dagrun.py, and re-reads the ones still waiting on their upstreams in_get_ready_tis, so it decoded the conf once per task instance.fetch_task_instancesnow defers the conf, which still loads if something reads it, and the re-check, which only compares states, no longer joins the run. The other task-instance queries in the scheduling loop return a bounded number of rows (the critical section, executor events) or a few rows per task instance (the mapped upstream dependency check), so they are left as they are. #73022 changed these two queries and four others; it was closed under the open PR limit before it was reviewed.Reproduced with
airflow standalone(SQLite, LocalExecutor): one run of a Dag with a single mapped task, triggered through the REST API. The 1 MB conf is about 8,400 small JSON records. Scheduler RSS, idle → peak:A task mapped over an upstream task's output behaves the same: 206 → 3,355 MB on main, 206 → 328 MB with this change. The remaining ~110 MB is the same for 50 and 500 task instances, and the 500-instance run with the 1 MB conf takes 196 s instead of 376 s.
The two new tests in
airflow-core/tests/unit/models/test_dagrun.pyfail without this change; reverting either query's change fails at least one of them. Locally, the airflow-coremodels,jobs,ti_depsand execution API unit tests and the other tests that call these functions pass, apart from four that also fail without this change (two needgraphviz, onemoto, andtest_still_in_retry_period). The prek pre-commit hooks, includingmypy-airflow-core, and the manual stage pass.closes: #71267
Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Opus 5.5) following the guidelines