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 be41cb4031e UI: Serve grid TI summaries from shared cached DagBag
(#65775)
be41cb4031e is described below
commit be41cb4031e38050314c9c9605885edee5dd1df0
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Mon Apr 27 17:54:17 2026 +0200
UI: Serve grid TI summaries from shared cached DagBag (#65775)
The grid TI summaries streaming endpoint used a per-request dict to
dedupe serialized Dag loads within a single response. Switch to the
app-wide DBDagBag so the serialized Dag is cached across requests (via
the LRU+TTL configured by [api] dag_cache_size / dag_cache_ttl) and
the in-memory representation is the lightweight SerializedDAG instead
of the SerializedDagModel ORM row holding the full JSON blob.
---
.../api_fastapi/core_api/openapi/_private_ui.yaml | 6 +-
.../airflow/api_fastapi/core_api/routes/ui/grid.py | 82 ++++++++++------------
.../ui/openapi-gen/queries/ensureQueryData.ts | 5 +-
.../src/airflow/ui/openapi-gen/queries/prefetch.ts | 5 +-
.../src/airflow/ui/openapi-gen/queries/queries.ts | 5 +-
.../src/airflow/ui/openapi-gen/queries/suspense.ts | 5 +-
.../ui/openapi-gen/requests/services.gen.ts | 5 +-
7 files changed, 55 insertions(+), 58 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
index ebf0975307a..364ecd38f95 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
@@ -1427,9 +1427,11 @@ paths:
progressively without waiting for all runs to complete.
- The serialized Dag structure is loaded once and reused for all runs
that
+ The serialized Dag structure is served from the app-wide ``DBDagBag``
cache
- share the same ``dag_version_id``, avoiding repeated deserialization.'
+ (keyed by ``dag_version_id``), which avoids repeated deserialization
across
+
+ runs of the same version *and* across requests.'
operationId: get_grid_ti_summaries_stream
security:
- OAuth2PasswordBearer: []
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 d266dcf76d7..d3ddf8a2a52 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
@@ -19,6 +19,7 @@ from __future__ import annotations
from collections.abc import Generator, Iterable
from typing import TYPE_CHECKING, Annotated, Any
+from uuid import UUID
import structlog
from fastapi import Depends, HTTPException, Query, status
@@ -27,6 +28,7 @@ from sqlalchemy import exists, select
from sqlalchemy.orm import Session, joinedload, load_only
from airflow.api_fastapi.auth.managers.models.resource_details import
DagAccessEntity
+from airflow.api_fastapi.common.dagbag import DagBagDep
from airflow.api_fastapi.common.db.common import SessionDep, paginated_select
from airflow.api_fastapi.common.db.dag_runs import attach_dag_versions_to_runs
from airflow.api_fastapi.common.parameters import (
@@ -69,6 +71,10 @@ from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.taskinstance import TaskInstance
from airflow.utils.session import create_session
+if TYPE_CHECKING:
+ from airflow.models.dagbag import DBDagBag
+ from airflow.serialization.definitions.dag import SerializedDAG
+
log = structlog.get_logger(logger_name=__name__)
grid_router = AirflowRouter(prefix="/grid", tags=["Grid"])
@@ -90,32 +96,28 @@ def _get_latest_serdag(dag_id, session):
return serdag
-def _get_serdag(dag_id, dag_version_id, session) -> SerializedDagModel | None:
- # this is a simplification - we account for structure based on the first
task
- version = session.scalar(
- select(DagVersion)
- .where(DagVersion.id == dag_version_id)
- .options(joinedload(DagVersion.serialized_dag))
+def _get_serdag(
+ dag_bag: DBDagBag,
+ dag_id: str,
+ dag_version_id: UUID | str | None,
+ session: Session,
+) -> SerializedDAG | None:
+ """Resolve the serialized Dag for a grid TI summary via the shared
(cached) ``DBDagBag``."""
+ if dag_version_id is not None:
+ serdag = dag_bag.get_dag(dag_version_id, session=session)
+ if serdag is None:
+ log.error("No serialized dag found", dag_id=dag_id,
version_id=dag_version_id)
+ return serdag
+
+ # Fallback: pre-3.0 upgrade — pick the oldest DagVersion for this dag_id.
+ oldest_version_id = session.scalar(
+ select(DagVersion.id).where(DagVersion.dag_id ==
dag_id).order_by(DagVersion.id).limit(1)
)
- if not version:
- version = session.scalar(
- select(DagVersion)
- .where(
- DagVersion.dag_id == dag_id,
- )
- .options(joinedload(DagVersion.serialized_dag))
- .order_by(DagVersion.id) # ascending cus this is mostly for
pre-3.0 upgrade
- .limit(1)
- )
- if not version:
+ if oldest_version_id is None:
return None
- if not (serdag := version.serialized_dag):
- log.error(
- "No serialized dag found",
- dag_id=dag_id,
- version_id=version.id,
- version_number=version.version_number,
- )
+ serdag = dag_bag.get_dag(oldest_version_id, session=session)
+ if serdag is None:
+ log.error("No serialized dag found", dag_id=dag_id,
version_id=oldest_version_id)
return serdag
@@ -356,12 +358,12 @@ def _build_ti_summaries(
task_instances: Iterable[Any],
session: Session,
*,
- serdag: SerializedDagModel | None = None,
- serdag_cache: dict[Any, SerializedDagModel | None] | None = None,
+ dag_bag: DBDagBag,
) -> dict[str, Any] | None:
ti_details: dict[str, GridNodeAgg] = {}
dag_version_id = None
for ti in task_instances:
+ # this is a simplification - we account for structure based on the
first task
dag_version_id = dag_version_id or ti.dag_version_id
summary = ti_details.get(ti.task_id)
if summary is None:
@@ -374,28 +376,15 @@ def _build_ti_summaries(
)
if not ti_details:
return None
- if serdag is None:
- if serdag_cache is not None:
- if dag_version_id not in serdag_cache:
- serdag_cache[dag_version_id] = _get_serdag(
- dag_id=dag_id,
- dag_version_id=dag_version_id,
- session=session,
- )
- serdag = serdag_cache[dag_version_id]
- else:
- serdag = _get_serdag(
- dag_id=dag_id,
- dag_version_id=dag_version_id,
- session=session,
- )
+
+ serdag = _get_serdag(dag_bag, dag_id, dag_version_id, session)
if TYPE_CHECKING:
assert serdag
def get_node_summaries() -> Iterable[dict[str, Any]]:
yielded_task_ids: set[str] = set()
for node, _ in _find_aggregates(
- node=serdag.dag.task_group,
+ node=serdag.task_group,
parent_node=None,
ti_details=ti_details,
):
@@ -457,6 +446,7 @@ def _build_ti_summaries(
)
def get_grid_ti_summaries_stream(
dag_id: str,
+ dag_bag: DagBagDep,
run_ids: Annotated[list[str] | None, Query()] = None,
) -> StreamingResponse:
"""
@@ -466,8 +456,9 @@ def get_grid_ti_summaries_stream(
run's task instances have been processed, so the client can render columns
progressively without waiting for all runs to complete.
- The serialized Dag structure is loaded once and reused for all runs that
- share the same ``dag_version_id``, avoiding repeated deserialization.
+ The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+ (keyed by ``dag_version_id``), which avoids repeated deserialization across
+ runs of the same version *and* across requests.
"""
def _generate() -> Generator[str, None, None]:
@@ -476,7 +467,6 @@ def get_grid_ti_summaries_stream(
# database connection open for the entire stream duration.
# See https://github.com/apache/airflow/issues/65010.
- serdag_cache: dict[Any, SerializedDagModel | None] = {}
for run_id in run_ids or []:
with create_session(scoped=False) as session:
tis = session.execute(
@@ -499,7 +489,7 @@ def get_grid_ti_summaries_stream(
run_id,
tis,
session,
- serdag_cache=serdag_cache,
+ dag_bag=dag_bag,
)
if summary is None:
continue
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
index 2808db355cf..1958854733a 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -1903,8 +1903,9 @@ export const ensureUseGridServiceGetGridRunsData =
(queryClient: QueryClient, {
* run's task instances have been processed, so the client can render columns
* progressively without waiting for all runs to complete.
*
-* The serialized Dag structure is loaded once and reused for all runs that
-* share the same ``dag_version_id``, avoiding repeated deserialization.
+* The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+* (keyed by ``dag_version_id``), which avoids repeated deserialization across
+* runs of the same version *and* across requests.
* @param data The data for the request.
* @param data.dagId
* @param data.runIds
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
index e1782283f52..71436325670 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -1903,8 +1903,9 @@ export const prefetchUseGridServiceGetGridRuns =
(queryClient: QueryClient, { da
* run's task instances have been processed, so the client can render columns
* progressively without waiting for all runs to complete.
*
-* The serialized Dag structure is loaded once and reused for all runs that
-* share the same ``dag_version_id``, avoiding repeated deserialization.
+* The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+* (keyed by ``dag_version_id``), which avoids repeated deserialization across
+* runs of the same version *and* across requests.
* @param data The data for the request.
* @param data.dagId
* @param data.runIds
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
index 4c2bf1a86c6..f7c589343fc 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -1903,8 +1903,9 @@ export const useGridServiceGetGridRuns = <TData =
Common.GridServiceGetGridRunsD
* run's task instances have been processed, so the client can render columns
* progressively without waiting for all runs to complete.
*
-* The serialized Dag structure is loaded once and reused for all runs that
-* share the same ``dag_version_id``, avoiding repeated deserialization.
+* The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+* (keyed by ``dag_version_id``), which avoids repeated deserialization across
+* runs of the same version *and* across requests.
* @param data The data for the request.
* @param data.dagId
* @param data.runIds
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
index a6d1d118db4..2a29df44c2a 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -1903,8 +1903,9 @@ export const useGridServiceGetGridRunsSuspense = <TData =
Common.GridServiceGetG
* run's task instances have been processed, so the client can render columns
* progressively without waiting for all runs to complete.
*
-* The serialized Dag structure is loaded once and reused for all runs that
-* share the same ``dag_version_id``, avoiding repeated deserialization.
+* The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+* (keyed by ``dag_version_id``), which avoids repeated deserialization across
+* runs of the same version *and* across requests.
* @param data The data for the request.
* @param data.dagId
* @param data.runIds
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index 1bb87669cee..a92d59a7be8 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -4415,8 +4415,9 @@ export class GridService {
* run's task instances have been processed, so the client can render
columns
* progressively without waiting for all runs to complete.
*
- * The serialized Dag structure is loaded once and reused for all runs that
- * share the same ``dag_version_id``, avoiding repeated deserialization.
+ * The serialized Dag structure is served from the app-wide ``DBDagBag``
cache
+ * (keyed by ``dag_version_id``), which avoids repeated deserialization
across
+ * runs of the same version *and* across requests.
* @param data The data for the request.
* @param data.dagId
* @param data.runIds
