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..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,11 +194,23 @@ 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) + # 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. + # + # 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 @@ -212,24 +225,33 @@ def get_dag_structure( .distinct() ), ) - .execution_options(yield_per=5) # balance between peak memory usage and round trips + .order_by(SerializedDagModel.id) ) - - 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, - ) - # 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) - - session.expunge(serdag) # to allow garbage collection + serdag_ids = list(session.scalars(serdag_id_query)) + # Release the request session's transaction/connection before the batched work. + session.commit() + + 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]