Load each serialized Dag once per triggerer pickup batch

airflow-core/src/airflow/jobs/triggerer_job_runner.py

@@ -59,7 +59,7 @@
 from airflow.executors.workloads.task import TaskInstanceDTO
 from airflow.jobs.base_job_runner import BaseJobRunner
 from airflow.jobs.job import perform_heartbeat
-from airflow.models.dagbag import DBDagBag
+from airflow.models.dagbag import BatchSerializedDagResolver, DBDagBag
 from airflow.models.task_coordinates import TaskCoordinateResolver
 from airflow.models.trigger import Trigger
 from airflow.observability.metrics import stats_utils
@@ -890,6 +890,8 @@ def _create_workload(
         render_log_fname: Callable[..., str],
         session: Session,
         log_context: TaskLogContext | None = None,
+        *,
+        serdag_resolver: BatchSerializedDagResolver | None = None,
     ) -> workloads.RunTrigger | None:
         if trigger.task_instance is None:
             watched_assets: dict[str, str] | None = None
@@ -948,16 +950,17 @@ def _create_workload(
         )
 
         dag_run = trigger.task_instance.get_dagrun(session=session)
-        serialized_dag_model = dag_bag.get_serialized_dag_model_for_run(
-            dag_run, session=session
-        ) or dag_bag.get_serialized_dag_model(
-            version_id=trigger.task_instance.dag_version_id, session=session
+        if serdag_resolver is None:
+            serdag_resolver = BatchSerializedDagResolver(dag_bag)
+        resolved_serdag = serdag_resolver.resolve(
+            dag_run, fallback_version_id=trigger.task_instance.dag_version_id, session=session
         )
 
-        if serialized_dag_model:
+        if resolved_serdag is not None:
+            serialized_dag_model, serialized_dag = resolved_serdag
             task = None
             try:
-                task = serialized_dag_model.dag.get_task(trigger.task_instance.task_id)
+                task = serialized_dag.get_task(trigger.task_instance.task_id)
             except TaskNotFound:
                 log.warning(
                     "Task not found in resolved Dag version; building plain workload",
@@ -1009,6 +1012,7 @@ def fetch_trigger_assignments(self, trigger_ids: set[int]) -> dict[int, int | No
     def build_trigger_workloads(self, new_trigger_ids: set[int]) -> list[workloads.RunTrigger]:
         """Build workloads for new trigger IDs."""
         dag_bag = DBDagBag()
+        serdag_resolver = BatchSerializedDagResolver(dag_bag)
         with create_session() as session:
             new_triggers = self.fetch_trigger_details(new_trigger_ids, session=session)
             log_contexts = prepare_task_log_contexts(
@@ -1062,6 +1066,7 @@ def render_log_fname(*, ti):
                         if new_trigger_orm.task_instance is not None
                         else None
                     ),
+                    serdag_resolver=serdag_resolver,
                 ):
                     to_create.append(workload)
 

airflow-core/src/airflow/models/dagbag.py

@@ -36,7 +36,7 @@
 from airflow.models.dag_version import DagVersion
 
 if TYPE_CHECKING:
-    from collections.abc import Generator
+    from collections.abc import Callable, Generator
 
     from sqlalchemy.orm import Session
 
@@ -193,29 +193,32 @@ def clear_cache(self) -> int:
         return count
 
     @staticmethod
-    def _version_from_dag_run(dag_run: DagRun, *, session: Session) -> UUID | None:
+    def _resolve_version_for_run(
+        dag_run: DagRun, *, get_latest_version_id: Callable[[str], UUID | None]
+    ) -> UUID | None:
+        """Take the latest-version lookup as a callable, so callers resolving many runs can memoize it."""
         # A run with no version of its own can only resolve to the latest. Runs carried over from
         # Airflow 2 are like this, as are runs whose version `airflow db clean` has since deleted --
         # the latter keep their bundle version, so they would otherwise resolve to nothing at all.
         if not dag_run.bundle_version or not dag_run.created_dag_version_id:
-            if dag_version := DagVersion.get_latest_version(dag_id=dag_run.dag_id, session=session):
-                return dag_version.id
+            if (latest_version_id := get_latest_version_id(dag_run.dag_id)) is not None:
+                return latest_version_id
 
         return dag_run.created_dag_version_id
 
+    @staticmethod
+    def _version_from_dag_run(dag_run: DagRun, *, session: Session) -> UUID | None:
+        def get_latest_version_id(dag_id: str) -> UUID | None:
+            dag_version = DagVersion.get_latest_version(dag_id=dag_id, session=session)
+            return dag_version.id if dag_version else None
+
+        return DBDagBag._resolve_version_for_run(dag_run, get_latest_version_id=get_latest_version_id)
+
     def get_dag_for_run(self, dag_run: DagRun, session: Session) -> SerializedDAG | None:
         if version_id := self._version_from_dag_run(dag_run=dag_run, session=session):
             return self._get_dag(version_id=version_id, session=session)
         return None
 
-    def get_serialized_dag_model_for_run(
-        self, dag_run: DagRun, *, session: Session
-    ) -> SerializedDagModel | None:
-        """Return the SerializedDagModel for the dag version this dag run resolves to."""
-        if version_id := self._version_from_dag_run(dag_run=dag_run, session=session):
-            return self.get_serialized_dag_model(version_id=version_id, session=session)
-        return None
-
     def iter_all_latest_version_dags(self, *, session: Session) -> Generator[SerializedDAG, None, None]:
         """
         Walk through all latest version dags available in the database.
@@ -289,6 +292,55 @@ def _on_cache_size(self, *, rate: float = 1.0) -> None:
         stats.gauge(f"{self._stats_prefix}.cache_size", size, rate=rate)
 
 
+class BatchSerializedDagResolver:
+    """
+    Memoize Dag version resolution and serialized Dag loading for one batch of lookups.
+
+    Entries are never revalidated: a Dag's latest version is read once, and a version found missing
+    or rewritten in place mid-batch is not re-read, so use a new instance per batch.
+
+    :meta private:
+    """
+
+    def __init__(self, dag_bag: DBDagBag) -> None:
+        self._dag_bag = dag_bag
+        self._latest_version_ids: dict[str, UUID | None] = {}
+        self._serdags: dict[UUID, tuple[SerializedDagModel, SerializedDAG] | None] = {}
+
+    def _get_latest_version_id(self, dag_id: str, *, session: Session) -> UUID | None:
+        if dag_id not in self._latest_version_ids:
+            dag_version = DagVersion.get_latest_version(dag_id=dag_id, session=session)
+            self._latest_version_ids[dag_id] = dag_version.id if dag_version else None
+        return self._latest_version_ids[dag_id]
+
+    def _get_serdag(
+        self, version_id: UUID, *, session: Session
+    ) -> tuple[SerializedDagModel, SerializedDAG] | None:
+        if version_id not in self._serdags:
+            model = self._dag_bag.get_serialized_dag_model(version_id=version_id, session=session)
+            self._serdags[version_id] = (model, model.dag) if model is not None else None
+        return self._serdags[version_id]
+
+    def resolve(
+        self, dag_run: DagRun, *, fallback_version_id: UUID, session: Session
+    ) -> tuple[SerializedDagModel, SerializedDAG] | None:
+        """
+        Return the serialized Dag model and its deserialized Dag for ``dag_run``.
+
+        Falls back to ``fallback_version_id`` when the run's version has no serialized Dag.
+        """
+        version_id = DBDagBag._resolve_version_for_run(
+            dag_run,
+            get_latest_version_id=lambda dag_id: self._get_latest_version_id(dag_id, session=session),
+        )
+        resolved = None
+        if version_id is not None:
+            resolved = self._get_serdag(version_id, session=session)
+        if resolved is None:
+            resolved = self._get_serdag(fallback_version_id, session=session)
+        return resolved
+
+
 def generate_md5_hash(context):
     bundle_name = context.get_current_parameters()["bundle_name"]
     relative_fileloc = context.get_current_parameters()["relative_fileloc"]

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -74,7 +74,7 @@
 )
 from airflow.models import Connection, DagModel, DagRun, Trigger, Variable
 from airflow.models.dag_version import DagVersion
-from airflow.models.dagbag import DBDagBag
+from airflow.models.dagbag import BatchSerializedDagResolver, DBDagBag
 from airflow.models.dagbundle import DagBundleModel
 from airflow.models.dynamic_region import DynamicRegion
 from airflow.models.serialized_dag import SerializedDagModel
@@ -110,6 +110,7 @@
     _ResponseFrame,
 )
 from airflow.sdk.execution_time.context import AssetStateStoreAccessors
+from airflow.serialization.definitions.dag import SerializedDAG
 from airflow.serialization.serialized_objects import LazyDeserializedDAG
 from airflow.triggers.base import BaseEventTrigger, BaseTrigger, TriggerEvent
 from airflow.triggers.shared_stream import SharedStreamProducer
@@ -860,17 +861,18 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
         return_value=mocker.Mock(spec=TaskInstanceDTO),
     )
 
-    dag_bag = mocker.Mock()
-    serialized_dag_model = mocker.Mock()
+    serialized_dag = mocker.Mock(spec=SerializedDAG)
     task = mocker.Mock(start_from_trigger=False)
-    serialized_dag_model.dag.get_task.return_value = task
-    dag_bag.get_serialized_dag_model_for_run.return_value = serialized_dag_model
+    serialized_dag.get_task.return_value = task
+    serdag_resolver = mocker.Mock(spec=BatchSerializedDagResolver)
+    serdag_resolver.resolve.return_value = (mocker.Mock(spec=SerializedDagModel), serialized_dag)
 
     render_log_fname = mocker.Mock(return_value="/logs/ti")
 
     jobless_supervisor._create_workload(
         trigger=trigger,
-        dag_bag=dag_bag,
+        dag_bag=mocker.Mock(spec=DBDagBag),
+        serdag_resolver=serdag_resolver,
         render_log_fname=render_log_fname,
         session=mocker.Mock(),
     )
@@ -934,6 +936,36 @@ def test_create_workload_resolves_serialized_dag_from_run(jobless_supervisor, mo
     dag_bag.get_serialized_dag_model.assert_called_once_with(version_id=expected_version, session=session)
 
 
+def test_build_trigger_workloads_batches_serdag_resolution(session, supervisor_builder, dag_maker, mocker):
+    """
+    Hydrating many triggers of one run resolves the Dag version and loads the serialized Dag once
+    per batch, not once per trigger, and still builds every workload.
+    """
+    with dag_maker(serialized=True):
+        for i in range(3):
+            EmptyOperator(task_id=f"task_{i}")
+    run = dag_maker.create_dagrun()
+    assert run.bundle_version is None  # unpinned: every trigger resolves through the latest version
+
+    trigger_ids = set()
+    for ti in run.task_instances:
+        trigger_orm = Trigger.from_object(TimeDeltaTrigger(datetime.timedelta(days=7)))
+        session.add(trigger_orm)
+        session.flush()
+        ti.trigger_id = trigger_orm.id
+        trigger_ids.add(trigger_orm.id)
+    session.commit()
+
+    get_latest_version = mocker.spy(DagVersion, "get_latest_version")
+    get_serialized_dag_model = mocker.spy(DBDagBag, "get_serialized_dag_model")
+
+    built = supervisor_builder().build_trigger_workloads(trigger_ids)
+
+    assert {workload.id for workload in built} == trigger_ids
+    assert get_latest_version.call_count == 1
+    assert get_serialized_dag_model.call_count == 1
+
+
 def test_load_triggers_survives_task_missing_from_resolved_dag_version(supervisor_builder, session, caplog):
     """
     A deferred TI's task may be missing from the Dag version an unpinned run resolves to

airflow-core/tests/unit/models/test_dagbag.py

@@ -19,17 +19,20 @@
 import math
 import re
 import time
+import uuid
 from concurrent.futures import ThreadPoolExecutor
 from unittest.mock import MagicMock, patch
 
 import pytest
 import time_machine
 from cachetools import LRUCache, TTLCache
+from sqlalchemy.orm import Session
 
 from airflow.models.dag import DagModel
 from airflow.models.dag_version import DagVersion
-from airflow.models.dagbag import CachedDBDagBag, DBDagBag, _CacheEntry
+from airflow.models.dagbag import BatchSerializedDagResolver, CachedDBDagBag, DBDagBag, _CacheEntry
 from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dagrun import DagRun
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.sdk import DAG
@@ -556,3 +559,98 @@ def test_cache_size_gauge_emitted(self, mock_stats):
         dag_bag._read_dag(mock_serdag)
 
         mock_stats.gauge.assert_called_with(f"{STUB_PREFIX}.cache_size", 1, rate=0.1)
+
+
+def _make_dag_run_mock(mocker, *, dag_id="test_dag", bundle_version=None, created_dag_version_id=None):
+    dag_run = mocker.Mock(spec=DagRun)
+    dag_run.dag_id = dag_id
+    dag_run.bundle_version = bundle_version
+    dag_run.created_dag_version_id = created_dag_version_id
+    return dag_run
+
+
+class TestBatchSerializedDagResolver:
+    def test_memoizes_latest_version_and_serialized_dag(self, mocker):
+        """Lookups of one Dag share one latest-version lookup, one fetch and one deserialization."""
+        latest_version = uuid.uuid4()
+        get_latest_version = mocker.patch.object(
+            DagVersion,
+            "get_latest_version",
+            autospec=True,
+            return_value=mocker.Mock(spec=DagVersion, id=latest_version),
+        )
+        dag_bag = DBDagBag()
+        serialized_dag_model = mocker.Mock(spec=SerializedDagModel)
+        deserialize_dag = mocker.PropertyMock(return_value=mocker.Mock(spec=SerializedDAG))
+        type(serialized_dag_model).dag = deserialize_dag
+        get_serialized_dag_model = mocker.patch.object(
+            dag_bag, "get_serialized_dag_model", autospec=True, return_value=serialized_dag_model
+        )
+        resolver = BatchSerializedDagResolver(dag_bag)
+        session = mocker.Mock(spec=Session)
+
+        results = [
+            resolver.resolve(_make_dag_run_mock(mocker), fallback_version_id=uuid.uuid4(), session=session)
+            for _ in range(3)
+        ]
+
+        assert all(result == (serialized_dag_model, deserialize_dag.return_value) for result in results)
+        get_latest_version.assert_called_once_with(dag_id="test_dag", session=session)
+        get_serialized_dag_model.assert_called_once_with(version_id=latest_version, session=session)
+        deserialize_dag.assert_called_once_with()
+
+    def test_falls_back_to_ti_version(self, mocker):
+        """A run resolving to a version with no serialized Dag falls back to the TI's own version."""
+        run_version = uuid.uuid4()
+        ti_version = uuid.uuid4()
+        dag_bag = DBDagBag()
+        serialized_dag_model = mocker.Mock(spec=SerializedDagModel)
+        get_serialized_dag_model = mocker.patch.object(
+            dag_bag,
+            "get_serialized_dag_model",
+            autospec=True,
+            side_effect=lambda version_id, session: (
+                serialized_dag_model if version_id == ti_version else None
+            ),
+        )
+        resolver = BatchSerializedDagResolver(dag_bag)
+
+        dag_run = _make_dag_run_mock(mocker, bundle_version="v1", created_dag_version_id=run_version)
+        resolved = resolver.resolve(
+            dag_run, fallback_version_id=ti_version, session=mocker.Mock(spec=Session)
+        )
+
+        assert resolved == (serialized_dag_model, serialized_dag_model.dag)
+        assert [call.kwargs["version_id"] for call in get_serialized_dag_model.call_args_list] == [
+            run_version,
+            ti_version,
+        ]
+
+    def test_looks_up_latest_version_once_per_dag(self, mocker):
+        """Runs of different Dags in one batch each get their own latest-version snapshot, taken once."""
+        latest_versions = {"dag_a": uuid.uuid4(), "dag_b": uuid.uuid4()}
+        get_latest_version = mocker.patch.object(
+            DagVersion,
+            "get_latest_version",
+            autospec=True,
+            side_effect=lambda dag_id, session: mocker.Mock(spec=DagVersion, id=latest_versions[dag_id]),
+        )
+        dag_bag = DBDagBag()
+        mocker.patch.object(
+            dag_bag,
+            "get_serialized_dag_model",
+            autospec=True,
+            return_value=mocker.Mock(spec=SerializedDagModel),
+        )
+        resolver = BatchSerializedDagResolver(dag_bag)
+        session = mocker.Mock(spec=Session)
+
+        for dag_id in ["dag_a", "dag_a", "dag_b", "dag_a"]:
+            resolver.resolve(
+                _make_dag_run_mock(mocker, dag_id=dag_id),
+                fallback_version_id=uuid.uuid4(),
+                session=session,
+            )
+
+        assert get_latest_version.call_count == 2
+        assert {call.kwargs["dag_id"] for call in get_latest_version.call_args_list} == {"dag_a", "dag_b"}