This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun 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 b278b622736 Release DB connection before deserializing DAGs in grid
structure endpoint (#69832)
b278b622736 is described below
commit b278b6227369a37763970aa2eb6b7364d719d0bd
Author: Sean Muth <[email protected]>
AuthorDate: Fri Aug 14 11:18:15 2026 -0500
Release DB connection before deserializing DAGs in grid structure endpoint
(#69832)
* Release DB connection before deserializing DAGs in grid structure endpoint
`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) <[email protected]>
* Add newsfragment for #69832
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
* 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) <[email protected]>
* Address review: trim comment, drop unnecessary order_by, use
session.close()
- Comment above serdag_id_query was too verbose for a simple query; the
rationale belongs in the commit message. Trimmed per dstandish's
suggestion.
- Dropped .order_by(SerializedDagModel.id): the original code had no
explicit
ordering, and _merge_node_dicts is sensitive to processing order for nodes
that differ across historical Dag versions, so this restores exact prior
behavior instead of introducing an unreviewed ordering guarantee.
- Replaced session.commit() with session.close(): this session has done no
writes (pure read), so close() is the more accurate call and, per review,
reads clearer while releasing the connection the same way.
---------
Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
---
.../airflow/api_fastapi/core_api/routes/ui/grid.py | 84 +++++++++++-----------
1 file changed, 44 insertions(+), 40 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 48bd5971461..e4f9b02f7d1 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
@@ -73,6 +73,7 @@ from airflow.models.dagrun import DagRun, DagRunNote
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:
@@ -215,47 +216,50 @@ def get_dag_structure(
_merge_node_dicts(merged_nodes, nodes)
del latest_dag, latest_group_dict
- # 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
- )
-
- for serdag in session.scalars(serdags_query):
- 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,
+ # we get the ids so that we can split serialization into batches and
balance round trips and mem usage
+ 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),
)
- # Merge immediately instead of collecting all Dags in memory
- filtered_group_dict = filtered_dag.task_group.get_task_group_dict()
- nodes = [
- task_group_to_dict_grid(x, group_dict=filtered_group_dict)
- for x in task_group_sort(filtered_dag.task_group,
filtered_group_dict)
- ]
- _merge_node_dicts(merged_nodes, nodes)
-
- session.expunge(serdag) # to allow garbage collection
+ .distinct()
+ ),
+ )
+ serdag_ids = list(session.scalars(serdag_id_query))
+ # Release the request session's transaction/connection before the batched
work.
+ session.close()
+
+ 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
+ filtered_group_dict = filtered_dag.task_group.get_task_group_dict()
+ nodes = [
+ task_group_to_dict_grid(x, group_dict=filtered_group_dict)
+ for x in task_group_sort(filtered_dag.task_group,
filtered_group_dict)
+ ]
+ _merge_node_dicts(merged_nodes, nodes)
return [GridNodeResponse(**n) for n in merged_nodes]
