This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun pushed a commit to branch v3-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-2-test by this push:
new 7249b6cc5b8 UI: Serve grid TI summaries from shared cached DagBag
(#65775) (#65966)
7249b6cc5b8 is described below
commit 7249b6cc5b81224441acb458e0b07313b0975ed4
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Tue Apr 28 09:58:46 2026 +0200
UI: Serve grid TI summaries from shared cached DagBag (#65775) (#65966)
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.
(cherry picked from commit be41cb4031e38050314c9c9605885edee5dd1df0)
---
.../api_fastapi/core_api/openapi/_private_ui.yaml | 6 +-
.../airflow/api_fastapi/core_api/routes/ui/grid.py | 85 ++++++++++++----------
.../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, 64 insertions(+), 52 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 0139c7e08d5..e4b4a43abf1 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
@@ -1267,9 +1267,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 2ee6ea30cb8..a4eb1198c07 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
@@ -20,14 +20,16 @@ from __future__ import annotations
import collections
from collections.abc import Generator, Sequence
from typing import TYPE_CHECKING, Annotated, Any
+from uuid import UUID
import structlog
from fastapi import Depends, HTTPException, Query, status
from fastapi.responses import StreamingResponse
from sqlalchemy import exists, select
-from sqlalchemy.orm import joinedload, load_only, selectinload
+from sqlalchemy.orm import Session, joinedload, load_only, selectinload
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.parameters import (
QueryDagRunRunTypesFilter,
@@ -70,6 +72,10 @@ from airflow.models.taskinstance import TaskInstance
from airflow.models.taskinstancehistory import TaskInstanceHistory
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"])
@@ -91,32 +97,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 | 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
@@ -344,7 +346,12 @@ def get_grid_runs(
def _build_ti_summaries(
- dag_id: str, run_id: str, task_instances: Sequence, session, serdag:
SerializedDagModel | None = None
+ dag_id: str,
+ run_id: str,
+ task_instances: Sequence,
+ session: Session,
+ *,
+ dag_bag: DBDagBag,
) -> dict:
ti_details: dict = collections.defaultdict(list)
for ti in task_instances:
@@ -356,19 +363,19 @@ def _build_ti_summaries(
"dag_version_number": getattr(ti, "version_number", None),
}
)
- if serdag is None:
- serdag = _get_serdag(
- dag_id=dag_id,
- dag_version_id=task_instances[0].dag_version_id,
- session=session,
- )
+ serdag = _get_serdag(
+ dag_bag,
+ dag_id,
+ task_instances[0].dag_version_id,
+ session,
+ )
if TYPE_CHECKING:
assert serdag
def get_node_summaries():
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,
):
@@ -430,6 +437,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:
"""
@@ -439,8 +447,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]:
@@ -449,7 +458,6 @@ def get_grid_ti_summaries_stream(
# released between yields. This prevents a slow client from holding a
# database connection open for the entire stream duration.
# See https://github.com/apache/airflow/issues/65010.
- serdag_cache: dict = {}
for run_id in run_ids or []:
with create_session(scoped=False) as session:
tis = session.execute(
@@ -468,10 +476,7 @@ def get_grid_ti_summaries_stream(
).all()
if not tis:
continue
- version_id = tis[0].dag_version_id
- if version_id not in serdag_cache:
- serdag_cache[version_id] = _get_serdag(dag_id, version_id,
session)
- summary = _build_ti_summaries(dag_id, run_id, tis, session,
serdag=serdag_cache[version_id])
+ summary = _build_ti_summaries(dag_id, run_id, tis, session,
dag_bag=dag_bag)
yield GridTISummaries.model_validate(summary).model_dump_json() +
"\n"
return StreamingResponse(content=_generate(),
media_type="application/x-ndjson")
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 55f9930a7c6..95e878c4809 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -1865,8 +1865,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 184e0468f6d..9107c9372c3 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -1865,8 +1865,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 26a0fff8f3c..76621317ac4 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -1865,8 +1865,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 099ccca4095..396d36a8f8a 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -1865,8 +1865,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 3db366caea3..03f2c29a498 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
@@ -4363,8 +4363,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
