This is an automated email from the ASF dual-hosted git repository.

pierrejeambrun pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/v3-3-test by this push:
     new 728605cfb5b [v3-3-test] Release DB connection before deserializing 
DAGs in grid structure endpoint (#69832) (#71626)
728605cfb5b is described below

commit 728605cfb5b770e998e4cae25404a4267807dd12
Author: github-actions[bot] 
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Mon Aug 17 12:59:58 2026 +0200

    [v3-3-test] Release DB connection before deserializing DAGs in grid 
structure endpoint (#69832) (#71626)
    
    * 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.
    
    
    
    * Add newsfragment for #69832
    
    
    
    * 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.
    
    
    
    * 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.
    
    ---------
    (cherry picked from commit b278b6227369a37763970aa2eb6b7364d719d0bd)
    
    Co-authored-by: Sean Muth <[email protected]>
    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 7622b13bff0..e0a2b95aea0 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.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:
@@ -198,47 +199,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]
 

Reply via email to