Release DB connection before deserializing DAGs in grid structure endpoint#69832
Release DB connection before deserializing DAGs in grid structure endpoint#69832seanmuth wants to merge 3 commits into
Conversation
…point
`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#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#65010.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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) <noreply@anthropic.com>
|
@seanmuth — There are 3 unresolved review thread(s) on this PR from @ephraimbuddy @kaxil. Could you either push a fix or reply in each thread explaining why the feedback doesn't apply? Once you believe the feedback is addressed, mark the thread as resolved so the reviewer isn't re-pinged needlessly. Thanks! Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you. |
| # 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). |
There was a problem hiding this comment.
i'm not sure this comment is really adding value... at least not here because, it immediately precedes a very simple query.
maybe just move this to the commit message? it's kinda high level rationale for why this change overall. the area where it sorta makes sense to have comments is below where we're batching things and we can state succinctly the rationale, which you kinda sorta already do.
There was a problem hiding this comment.
maybe just say
# we get the ids so that we can split serialization into batches and balance round trips and mem usage
There was a problem hiding this comment.
+1 this comment is too verbose
| ), | ||
| ) | ||
| .execution_options(yield_per=5) # balance between peak memory usage and round trips | ||
| .order_by(SerializedDagModel.id) |
There was a problem hiding this comment.
Why a new order_by? I believe this induce a behavior change in the ordering later on. I wouldn't add this if not necessary.
| 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() |
There was a problem hiding this comment.
session.close() reads better here and achieve the same
6 participants
Overview
get_dag_structure(GET /ui/grid/structure/{dag_id}) streams historicalSerializedDagModelrows with ayield_perserver-side cursor while performing CPU-boundserdag.dagdeserialization and task-group merging between fetches. Because the cursor stays open across that work, the read transaction — and, under PgBouncer transaction pooling, the pooled server connection — stays pinned for the entire render of a large DAG.For DAGs with many tasks and many historical versions this holds a connection open for a long time (multiple minutes in the field). At scale that exhausts the PgBouncer pool and starves task-instance heartbeat updates, contributing to the metadata-DB contention tracked in #65712.
Why this matters
The UI API and the execution API are served by the same api-server process and share a single metadata-DB connection pool. A UI endpoint that holds a pooled connection across a multi-minute render therefore directly starves execution-API calls — task-SDK heartbeats and state updates — competing for that pool. When those heartbeat/state calls are delayed, workers time out and retry, which piles more load onto the same pool. Keeping UI read handlers from holding connections across CPU-bound work is what prevents UI traffic from degrading task execution.
Change
Materialize the (page-bounded) set of historical serialized DAGs, detach them from the session, and commit the read transaction so the connection returns to the pool before the deserialization/merge loop runs.
SerializedDagModel.dagonly reads the already-loadeddatacolumn, so deserialization works on detached instances with no further DB access. The result set is bounded by this endpoint's page of runs (a single DAG's versions), so materializing it is a safe trade for not holding the connection open across CPU-bound work.This mirrors the per-unit session-release pattern already applied to the streaming
ti_summariesendpoint in #65010.Related
Testing
structureendpoint tests pass, includingtest_structure_includes_historical_removed_task_with_proper_shape, which exercises the historical-versions loop that was changed.🤖 Generated with Claude Code — Claude Opus 4.8 (1M context)