This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new a264c38fdd6 [v3-4-test] Fix scheduler memory spike when a Dag run has
a large conf (#73907) (#76252)
a264c38fdd6 is described below
commit a264c38fdd66c0239e01357ddd72e9631a4bcbb7
Author: Shahar Epstein <[email protected]>
AuthorDate: Fri Nov 6 09:05:28 2026 +0200
[v3-4-test] Fix scheduler memory spike when a Dag run has a large conf
(#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 ce01a733e93f3262ef6f19f3d82632da72175925)
Co-authored-by: Lucas Pimentel
<[email protected]>
---
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 dbb0c76f55a..d5b8d600ca4 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -62,6 +62,7 @@ from sqlalchemy.orm import (
Mapped,
declared_attr,
joinedload,
+ lazyload,
mapped_column,
relationship,
synonym,
@@ -988,9 +989,11 @@ class DagRun(Base, LoggingMixin):
session: Session = NEW_SESSION,
) -> list[TI]:
"""Return the task instances for this dag run."""
+ # 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.
tis = (
select(TI)
- .options(joinedload(TI.dag_run))
+ .options(joinedload(TI.dag_run).defer(DagRun.conf))
.where(
TI.dag_id == dag_id,
TI.run_id == run_id,
@@ -1831,7 +1834,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 1c3f1208916..7a63545926b 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
@@ -114,7 +115,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
@@ -132,6 +133,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):
@@ -1122,6 +1125,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
