From a82707787a3d97d461e5fddde43277df3f6f5578 Mon Sep 17 00:00:00 2001 From: seanmuth Date: Mon, 13 Jul 2026 10:09:13 -0500 Subject: [PATCH 1/3] Release DB connection before deserializing DAGs in grid structure endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `get_dag_structure` (`GET /ui/grid/structure/{dag_id}`) streamed historical `SerializedDagModel` rows with a `yield_per` server-side cursor while running CPU-bound `serdag.dag` deserialization and task-group merging between fetches. That keeps the read transaction — and, under PgBouncer transaction pooling, the pooled server connection — pinned for the entire render of a large DAG. At scale this holds connections open for minutes, exhausting the PgBouncer pool and starving task-instance heartbeats, contributing to the contention tracked in apache/airflow#65712. Materialize the (page-bounded) historical serialized DAGs, detach them, and commit the read transaction so the connection returns to the pool *before* the deserialization/merge. `SerializedDagModel.dag` only reads the already-loaded `data` column, so it works on detached instances. This mirrors the per-unit session-release pattern already applied to the streaming `ti_summaries` endpoint in apache/airflow#65010. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../api_fastapi/core_api/routes/ui/grid.py | 58 +++++++++++-------- 1 file changed, 34 insertions(+), 24 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py index a16f9336fd57e..b74ab2c030117 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py @@ -193,29 +193,41 @@ def get_dag_structure( _merge_node_dicts(merged_nodes, nodes) del latest_dag - # Process serdags one by one and merge immediately to reduce memory usage. - # Use yield_per() for streaming results and expunge each serdag after processing - # to allow garbage collection and prevent memory buildup in the session identity map. - serdags_query = ( - select(SerializedDagModel) - .where( - # Even though dag_id is filtered in base_query, - # adding this line here can improve the performance of this endpoint - SerializedDagModel.dag_id == dag_id, - SerializedDagModel.id != latest_serdag_id, - SerializedDagModel.dag_version_id.in_( - select(TaskInstance.dag_version_id) - .join(TaskInstance.dag_run) - .where( - DagRun.id.in_(run_ids), - ) - .distinct() - ), - ) - .execution_options(yield_per=5) # balance between peak memory usage and round trips + # Fetch the historical serialized Dags for the visible runs, then detach them + # and release the DB connection *before* the CPU-bound deserialization + merge + # below. Deserializing (``serdag.dag``) and merging a large DAG's task groups is + # expensive; doing it while a server-side cursor (``yield_per``) is open keeps the + # transaction — and, under PgBouncer transaction pooling, the pooled server + # connection — pinned for the entire render. At scale that exhausts the pool and + # starves task-instance heartbeats, so the connection must not be held across the + # deserialization. See https://github.com/apache/airflow/issues/65712. + # + # ``SerializedDagModel.dag`` only reads the already-loaded ``data`` column, so it + # works on detached instances after the session is released. The result set is + # bounded by this endpoint's page of runs (one DAG's versions), so materializing + # it is a safe trade for not holding the connection open. + serdags_query = select(SerializedDagModel).where( + # Even though dag_id is filtered in base_query, + # adding this line here can improve the performance of this endpoint + SerializedDagModel.dag_id == dag_id, + SerializedDagModel.id != latest_serdag_id, + SerializedDagModel.dag_version_id.in_( + select(TaskInstance.dag_version_id) + .join(TaskInstance.dag_run) + .where( + DagRun.id.in_(run_ids), + ) + .distinct() + ), ) - - for serdag in session.scalars(serdags_query): + serdags = session.scalars(serdags_query).all() + for serdag in serdags: + session.expunge(serdag) # detach so `.dag` deserializes without the session + # End the read transaction so the connection returns to the pool during the + # deserialization/merge below instead of sitting idle-in-transaction. + session.commit() + + for serdag in serdags: filtered_dag = serdag.dag # Apply the same filtering to historical Dag versions if root: @@ -229,8 +241,6 @@ def get_dag_structure( nodes = [task_group_to_dict_grid(x) for x in task_group_sort(filtered_dag.task_group)] _merge_node_dicts(merged_nodes, nodes) - session.expunge(serdag) # to allow garbage collection - return [GridNodeResponse(**n) for n in merged_nodes] From b78a209ff983358ed92f0f752b8d0418bc378856 Mon Sep 17 00:00:00 2001 From: seanmuth Date: Mon, 13 Jul 2026 10:11:36 -0500 Subject: [PATCH 2/3] Add newsfragment for #69832 Co-Authored-By: Claude Opus 4.8 (1M context) --- airflow-core/newsfragments/69832.bugfix.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 airflow-core/newsfragments/69832.bugfix.rst diff --git a/airflow-core/newsfragments/69832.bugfix.rst b/airflow-core/newsfragments/69832.bugfix.rst new file mode 100644 index 0000000000000..bf1d0a8bfaacf --- /dev/null +++ b/airflow-core/newsfragments/69832.bugfix.rst @@ -0,0 +1 @@ +Release the metadata DB connection in the grid ``structure`` endpoint before deserializing historical DAG versions, so it is not held open across CPU-bound work and does not exhaust the connection pool on large DAGs. From c58da2916e4739dc6593872e8187cafefca846b3 Mon Sep 17 00:00:00 2001 From: seanmuth Date: Mon, 13 Jul 2026 14:58:50 -0500 Subject: [PATCH 3/3] Batch historical serdag loading to keep memory bounded Address review: instead of materializing all historical SerializedDagModel rows at once (which regressed the yield_per=5 memory profile), fetch their ids and process them in batches of 5, each batch loaded and detached in its own short-lived session that is closed before the batch is deserialized. This keeps peak memory bounded to one batch (matching the previous behaviour) while still releasing the DB connection during the CPU-bound deserialization. Also drop the newsfragment per review. Co-Authored-By: Claude Opus 4.8 (1M context) --- airflow-core/newsfragments/69832.bugfix.rst | 1 - .../api_fastapi/core_api/routes/ui/grid.py | 98 +++++++++++-------- 2 files changed, 55 insertions(+), 44 deletions(-) delete mode 100644 airflow-core/newsfragments/69832.bugfix.rst diff --git a/airflow-core/newsfragments/69832.bugfix.rst b/airflow-core/newsfragments/69832.bugfix.rst deleted file mode 100644 index bf1d0a8bfaacf..0000000000000 --- a/airflow-core/newsfragments/69832.bugfix.rst +++ /dev/null @@ -1 +0,0 @@ -Release the metadata DB connection in the grid ``structure`` endpoint before deserializing historical DAG versions, so it is not held open across CPU-bound work and does not exhaust the connection pool on large DAGs. diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py index b74ab2c030117..b21261ae27959 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py @@ -69,6 +69,7 @@ from airflow.models.deadline import Deadline from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance, TaskInstanceNote +from airflow.utils.helpers import chunks from airflow.utils.session import create_session if TYPE_CHECKING: @@ -193,53 +194,64 @@ def get_dag_structure( _merge_node_dicts(merged_nodes, nodes) del latest_dag - # Fetch the historical serialized Dags for the visible runs, then detach them - # and release the DB connection *before* the CPU-bound deserialization + merge - # below. Deserializing (``serdag.dag``) and merging a large DAG's task groups is - # expensive; doing it while a server-side cursor (``yield_per``) is open keeps the - # transaction — and, under PgBouncer transaction pooling, the pooled server - # connection — pinned for the entire render. At scale that exhausts the pool and - # starves task-instance heartbeats, so the connection must not be held across the - # deserialization. See https://github.com/apache/airflow/issues/65712. + # Identify the historical serialized Dags for the visible runs (ids only, cheap), + # then process them in small batches so we never hold a DB connection open across + # the CPU-bound deserialization/merge below. # - # ``SerializedDagModel.dag`` only reads the already-loaded ``data`` column, so it - # works on detached instances after the session is released. The result set is - # bounded by this endpoint's page of runs (one DAG's versions), so materializing - # it is a safe trade for not holding the connection open. - serdags_query = select(SerializedDagModel).where( - # Even though dag_id is filtered in base_query, - # adding this line here can improve the performance of this endpoint - SerializedDagModel.dag_id == dag_id, - SerializedDagModel.id != latest_serdag_id, - SerializedDagModel.dag_version_id.in_( - select(TaskInstance.dag_version_id) - .join(TaskInstance.dag_run) - .where( - DagRun.id.in_(run_ids), - ) - .distinct() - ), + # Deserializing (``serdag.dag``) and merging a large DAG's task groups is + # expensive. Streaming the rows with a server-side cursor (``yield_per``) and + # deserializing between fetches keeps the transaction — and, under PgBouncer + # transaction pooling, the pooled server connection — pinned for the entire + # render; at scale that exhausts the pool and starves task-instance heartbeats. + # See https://github.com/apache/airflow/issues/65712. + # + # Each batch is loaded and detached in its own short-lived session that is closed + # before the batch is deserialized, so the connection is released during the CPU + # work while peak memory stays bounded to one batch (``SerializedDagModel.dag`` + # only reads the already-loaded ``data`` column, so it works on detached rows). + serdag_id_query = ( + select(SerializedDagModel.id) + .where( + # Even though dag_id is filtered in base_query, + # adding this line here can improve the performance of this endpoint + SerializedDagModel.dag_id == dag_id, + SerializedDagModel.id != latest_serdag_id, + SerializedDagModel.dag_version_id.in_( + select(TaskInstance.dag_version_id) + .join(TaskInstance.dag_run) + .where( + DagRun.id.in_(run_ids), + ) + .distinct() + ), + ) + .order_by(SerializedDagModel.id) ) - serdags = session.scalars(serdags_query).all() - for serdag in serdags: - session.expunge(serdag) # detach so `.dag` deserializes without the session - # End the read transaction so the connection returns to the pool during the - # deserialization/merge below instead of sitting idle-in-transaction. + serdag_ids = list(session.scalars(serdag_id_query)) + # Release the request session's transaction/connection before the batched work. session.commit() - for serdag in serdags: - filtered_dag = serdag.dag - # Apply the same filtering to historical Dag versions - if root: - filtered_dag = filtered_dag.partial_subset( - task_ids=root, - include_upstream=include_upstream, - include_downstream=include_downstream, - depth=depth, - ) - # Merge immediately instead of collecting all Dags in memory - nodes = [task_group_to_dict_grid(x) for x in task_group_sort(filtered_dag.task_group)] - _merge_node_dicts(merged_nodes, nodes) + for serdag_id_batch in chunks(serdag_ids, 5): # balance memory usage and round trips + with create_session(scoped=False) as batch_session: + serdags = batch_session.scalars( + select(SerializedDagModel).where(SerializedDagModel.id.in_(serdag_id_batch)) + ).all() + for serdag in serdags: + batch_session.expunge(serdag) # detach so `.dag` deserializes without the session + # Connection is released here; deserialize + merge this batch outside the transaction. + for serdag in serdags: + filtered_dag = serdag.dag + # Apply the same filtering to historical Dag versions + if root: + filtered_dag = filtered_dag.partial_subset( + task_ids=root, + include_upstream=include_upstream, + include_downstream=include_downstream, + depth=depth, + ) + # Merge immediately instead of collecting all Dags in memory + nodes = [task_group_to_dict_grid(x) for x in task_group_sort(filtered_dag.task_group)] + _merge_node_dicts(merged_nodes, nodes) return [GridNodeResponse(**n) for n in merged_nodes]