seanmuth commented on code in PR #69832:
URL: https://github.com/apache/airflow/pull/69832#discussion_r3573746164
##########
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()
+ ),
Review Comment:
Good catch — you're right that materializing all of them regressed the
`yield_per=5` memory profile. Reworked in c58da2916e: I now fetch just the
serdag ids, then process them in batches of 5, each batch loaded and detached
in its own short-lived `create_session()` that's closed *before* the batch is
deserialized.
That keeps peak memory bounded to one batch (same as the old behaviour)
while still releasing the connection during the CPU-bound deserialization —
which is the whole point, since holding it across that work is what pins a
PgBouncer server connection for the entire render on large DAGs. It's the same
per-unit session-release pattern used for the streaming `ti_summaries` endpoint
(#65010).
##########
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.
Review Comment:
Removed in c58da2916e.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
