This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ce01a733e93 Fix scheduler memory spike when a Dag run has a large conf
(#73907)
ce01a733e93 is described below
commit ce01a733e93f3262ef6f19f3d82632da72175925
Author: Lucas Pimentel <[email protected]>
AuthorDate: Sun Nov 1 16:39:36 2026 -0300
Fix scheduler memory spike when a Dag run has a large conf (#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.
---
airflow-core/src/airflow/models/dagrun.py | 10 ++++--
airflow-core/tests/unit/models/test_dagrun.py | 51 ++++++++++++++++++++++++++-
2 files changed, 58 insertions(+), 3 deletions(-)
diff --git a/airflow-core/src/airflow/models/dagrun.py
b/airflow-core/src/airflow/models/dagrun.py
index 6a1ea5eb6f5..6b8236805f1 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -63,6 +63,7 @@ from sqlalchemy.orm import (
Mapped,
declared_attr,
joinedload,
+ lazyload,
mapped_column,
relationship,
synonym,
@@ -1064,8 +1065,10 @@ class DagRun(Base, LoggingMixin):
If ``columns`` is provided then the result is plain rows with the
values of these columns, otherwise the result is ORM objects.
"""
+ # The joined run comes back on every row. Its conf can be large, so it
loads on first
+ # access instead of being decoded once per task instance.
stmt = (
- (select(*columns) if columns else
select(TI).options(joinedload(TI.dag_run)))
+ (select(*columns) if columns else
select(TI).options(joinedload(TI.dag_run).defer(DagRun.conf)))
.where(TI.dag_id == dag_id, TI.run_id == run_id)
.order_by(TI.task_id, TI.region_id, TI.region_index)
)
@@ -1975,7 +1978,10 @@ class DagRun(Base, LoggingMixin):
# Check if any ti changed state
if old_states:
- fresh_tis =
session.scalars(select(TI).where(TI.id.in_(old_states))).all()
+ # Joining the run here would decode its conf once per row; only
the states are needed.
+ fresh_tis = session.scalars(
+
select(TI).options(lazyload(TI.dag_run)).where(TI.id.in_(old_states))
+ ).all()
changed_tis = any(ti.state != old_states[ti.id] for ti in
fresh_tis)
return ready_tis, changed_tis, expansion_happened
diff --git a/airflow-core/tests/unit/models/test_dagrun.py
b/airflow-core/tests/unit/models/test_dagrun.py
index 7097d93da05..cf162e0487d 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -18,6 +18,7 @@
from __future__ import annotations
import datetime
+import re
from collections import defaultdict
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
@@ -116,7 +117,7 @@ from airflow.utils.state import DagRunState,
DagSchedulingState, State, TaskInst
from airflow.utils.types import DagRunTriggeredByType, DagRunType
from tests_common.test_utils import db
-from tests_common.test_utils.asserts import assert_queries_count
+from tests_common.test_utils.asserts import assert_queries_count,
capture_orm_selects
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.dag import sync_dag_to_db
from tests_common.test_utils.db import clear_db_dags, clear_db_runs,
clear_db_serialized_dags, clear_db_xcom
@@ -134,6 +135,8 @@ pytestmark = [pytest.mark.db_test,
pytest.mark.need_serialized_dag]
TI = TaskInstance
DEFAULT_DATE = pendulum.instance(_DEFAULT_DATE)
+# A joined eager load of TaskInstance.dag_run aliases the table, hence the
optional suffix.
+SELECTS_RUN_CONF = re.compile(r"\bdag_run(?:_\d+)?\.conf\b")
def get_callback_ti(ti):
@@ -1132,6 +1135,52 @@ class TestDagRun:
ti = dag_run.get_task_instance("test_short_circuit_false")
assert ti is None
+ def test_fetch_task_instances_does_not_select_the_run_conf(self,
dag_maker, session):
+ """
+ The run's conf must not come back on every task-instance row.
+
+ ``TaskInstance.dag_run`` is a joined eager load, so a query that
selected ``dag_run.conf``
+ would return and decode one copy of the conf per task instance.
+ """
+ conf = {"payload": "x" * 1024}
+ with dag_maker(session=session):
+ EmptyOperator(task_id="t1")
+ EmptyOperator(task_id="t2")
+ dag_run = dag_maker.create_dagrun(conf=conf)
+ dag_id, run_id = dag_run.dag_id, dag_run.run_id
+ session.expunge_all()
+
+ with capture_orm_selects("task_instance") as statements:
+ tis = DagRun.fetch_task_instances(dag_id=dag_id, run_id=run_id,
session=session)
+
+ assert [ti.task_id for ti in tis] == ["t1", "t2"]
+ assert statements
+ assert not any(SELECTS_RUN_CONF.search(sql) for sql in statements),
statements
+ # The run is still joined, and its conf loads when read.
+ assert tis[0].dag_run.run_id == run_id
+ assert tis[0].dag_run.conf == conf
+
+ def test_scheduling_decisions_do_not_select_the_run_conf(self, dag_maker,
session):
+ """
+ A scheduling pass must not load the run's conf once per task instance.
+
+ ``downstream`` waits for ``upstream``, so besides fetching the run's
task instances the pass
+ re-reads the waiting ones to see whether their state changed.
+ """
+ with dag_maker(session=session):
+ EmptyOperator(task_id="upstream") >>
EmptyOperator(task_id="downstream")
+ dag_run = dag_maker.create_dagrun(conf={"payload": "x" * 1024})
+
+ with capture_orm_selects("task_instance") as statements:
+ decision =
dag_run.task_instance_scheduling_decisions(session=session)
+
+ assert [ti.task_id for ti in decision.schedulable_tis] == ["upstream"]
+ assert statements
+ selecting_conf = [sql for sql in statements if
SELECTS_RUN_CONF.search(sql)]
+ assert not selecting_conf, (
+ f"{len(selecting_conf)} of {len(statements)} task-instance queries
select the run conf"
+ )
+
def test_get_latest_runs(self, dag_maker, session):
with dag_maker(
dag_id="test_latest_runs_1", schedule=datetime.timedelta(days=1),
start_date=DEFAULT_DATE
