This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch v3-4-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 3cf7ff1292e104237be0caebe045454be89ba2fe 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
