diff --git a/airflow-core/src/airflow/api_fastapi/common/dagbag.py b/airflow-core/src/airflow/api_fastapi/common/dagbag.py index d87aca49a524a..b96e6f65b061c 100644 --- a/airflow-core/src/airflow/api_fastapi/common/dagbag.py +++ b/airflow-core/src/airflow/api_fastapi/common/dagbag.py @@ -16,43 +16,35 @@ # under the License. from __future__ import annotations -import logging from typing import TYPE_CHECKING, Annotated from fastapi import Depends, HTTPException, Request, status from sqlalchemy.orm import Session from airflow.configuration import conf -from airflow.models.dagbag import DBDagBag +from airflow.models.dagbag import CachedDBDagBag, DBDagBag from airflow.models.serialized_dag import SerializedDagModel if TYPE_CHECKING: from airflow.models.dagrun import DagRun from airflow.serialization.definitions.dag import SerializedDAG -log = logging.getLogger(__name__) - -def create_dag_bag() -> DBDagBag: +def create_dag_bag() -> CachedDBDagBag: """Create DagBag with configurable LRU+TTL caching for API server usage.""" cache_size = conf.getint("api", "dag_cache_size", fallback=64) - cache_ttl_config = conf.getint("api", "dag_cache_ttl", fallback=3600) + cache_ttl = conf.getint("api", "dag_cache_ttl", fallback=3600) if cache_size < 0: - log.warning("dag_cache_size must be >= 0, using unbounded dict") - cache_size = 0 - if cache_ttl_config < 0: - log.warning("dag_cache_ttl must be >= 0, disabling TTL") - cache_ttl_config = 0 - - # Use unbounded dict (no eviction) if cache_size is 0 - if cache_size <= 0: - return DBDagBag(cache_size=0) - - # Disable TTL if cache_ttl is 0 - cache_ttl: int | None = cache_ttl_config if cache_ttl_config > 0 else None - - return DBDagBag(cache_size=cache_size, cache_ttl=cache_ttl) + raise ValueError("[api] dag_cache_size must be greater than or equal to 0") + if cache_ttl < 0: + raise ValueError("[api] dag_cache_ttl must be greater than or equal to 0") + + return CachedDBDagBag( + cache_size=cache_size, + cache_ttl=cache_ttl, + stats_prefix="api_server.dag_bag", + ) def dag_bag_from_app(request: Request) -> DBDagBag: diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index d03b1d456f466..f9dfcf4c34b43 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -98,7 +98,7 @@ ) from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion -from airflow.models.dagbag import DBDagBag +from airflow.models.dagbag import CachedDBDagBag, DBDagBag from airflow.models.dagbundle import DagBundleModel from airflow.models.dagrun import DagRun from airflow.models.dagwarning import DagWarning, DagWarningType @@ -360,7 +360,12 @@ def __init__( if log: self._log = log - self.scheduler_dag_bag = DBDagBag(load_op_links=False, cache_size=SCHEDULER_DAG_CACHE_SIZE) + self.scheduler_dag_bag = CachedDBDagBag( + load_op_links=False, + cache_size=SCHEDULER_DAG_CACHE_SIZE, + cache_ttl=0, + stats_prefix="scheduler.dag_bag", + ) # Set of (dag_id, asset_name, asset_uri) tuples for trigger policies that # are permanently unreachable for the rollup window's cardinality — the diff --git a/airflow-core/src/airflow/models/dagbag.py b/airflow-core/src/airflow/models/dagbag.py index c4bd8eceea102..386e35e1bb31a 100644 --- a/airflow-core/src/airflow/models/dagbag.py +++ b/airflow-core/src/airflow/models/dagbag.py @@ -18,6 +18,7 @@ from __future__ import annotations import hashlib +import math import time from collections.abc import MutableMapping from contextlib import nullcontext @@ -62,44 +63,34 @@ class DBDagBag: """ Internal class for retrieving dags from the database. - Optionally supports LRU+TTL caching when cache_size is provided. - The scheduler uses this without caching, while the API server can - enable caching via configuration. + Deserialized Dags are retained in an unbounded dictionary. Use :class:`CachedDBDagBag` when + the caller needs configurable eviction, thread safety, and cache metrics. :meta private: """ - def __init__( - self, - load_op_links: bool = True, - cache_size: int | None = None, - cache_ttl: int | None = None, - ) -> None: + def __init__(self, load_op_links: bool = True) -> None: """ Initialize DBDagBag. :param load_op_links: Should the extra operator link be loaded when de-serializing the DAG? - :param cache_size: Size of LRU cache. If None or 0, uses unbounded dict (no eviction). - :param cache_ttl: Time-to-live for cache entries in seconds. If None or 0, no TTL (LRU only). """ self.load_op_links = load_op_links self._dags: MutableMapping[UUID | str, _CacheEntry] = {} - self._use_cache = False - self._revalidation_interval = conf.getint("core", "min_serialized_dag_update_interval") + self._lock: RLock | nullcontext = nullcontext() + + def _on_cache_hit(self) -> None: + """Handle a Dag cache hit.""" - # Initialize bounded cache if cache_size is provided and > 0 - if cache_size and cache_size > 0: - if cache_ttl and cache_ttl > 0: - self._dags = TTLCache(maxsize=cache_size, ttl=cache_ttl) - else: - self._dags = LRUCache(maxsize=cache_size) - self._use_cache = True + def _on_cache_miss(self) -> None: + """Handle a Dag cache miss.""" - # Lock required for bounded caches: cachetools caches are NOT thread-safe - # (LRU reordering and TTL cleanup mutate internal linked lists). - # nullcontext for unbounded dict avoids lock overhead in the scheduler path. - self._lock: RLock | nullcontext = RLock() if self._use_cache else nullcontext() + def _on_cache_clear(self) -> None: + """Handle the Dag cache being cleared.""" + + def _on_cache_size(self, *, rate: float = 1.0) -> None: + """Handle a change in the Dag cache size.""" def _read_dag(self, serdag: SerializedDagModel) -> SerializedDAG | None: """Read and cache a SerializedDAG (with its ``dag_hash`` for staleness detection).""" @@ -109,9 +100,7 @@ def _read_dag(self, serdag: SerializedDagModel) -> SerializedDAG | None: return None with self._lock: self._dags[serdag.dag_version_id] = _CacheEntry(dag, serdag.dag_hash, time.monotonic()) - cache_size = len(self._dags) - if self._use_cache: - stats.gauge("api_server.dag_bag.cache_size", cache_size, rate=0.1) + self._on_cache_size(rate=0.1) return dag @staticmethod @@ -133,8 +122,7 @@ def _get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG | # [core] min_serialized_dag_update_interval, so an entry validated within that window # cannot have gone stale yet -- serve it without touching the DB. if now - cached.last_validated < self._revalidation_interval: - if self._use_cache: - stats.incr("api_server.dag_bag.cache_hit") + self._on_cache_hit() return cached.dag # Past the window: a version may have been updated in place (same dag_version_id, new # content + new dag_hash) by SerializedDagModel.write_dag, so confirm the cached copy @@ -148,8 +136,7 @@ def _get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG | current = self._dags.get(version_id) if current is not None and current.dag_hash == cached.dag_hash: self._dags[version_id] = current._replace(last_validated=now) - if self._use_cache: - stats.incr("api_server.dag_bag.cache_hit") + self._on_cache_hit() return cached.dag # Stale (updated in place) or the version no longer exists: drop and reload below. with self._lock: @@ -166,12 +153,11 @@ def _get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG | # served without an extra hash check, consistent with the policy above. Only emit the miss # metric after confirming no other thread cached it, to avoid counting a single lookup as # both a miss and a hit. - if self._use_cache: - with self._lock: - if (cached := self._dags.get(version_id)) is not None: - stats.incr("api_server.dag_bag.cache_hit") - return cached.dag - stats.incr("api_server.dag_bag.cache_miss") + with self._lock: + if (cached := self._dags.get(version_id)) is not None: + self._on_cache_hit() + return cached.dag + self._on_cache_miss() return self._read_dag(serdag) def get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG | None: @@ -202,9 +188,8 @@ def clear_cache(self) -> int: count = len(self._dags) self._dags.clear() - if self._use_cache: - stats.incr("api_server.dag_bag.cache_clear") - stats.gauge("api_server.dag_bag.cache_size", 0) + self._on_cache_clear() + self._on_cache_size() return count @staticmethod @@ -243,6 +228,56 @@ def get_latest_version_of_dag(self, dag_id: str, *, session: Session) -> Seriali return self._read_dag(serdag) +class CachedDBDagBag(DBDagBag): + """Retrieve Dags through a configurable, thread-safe cache that emits component metrics.""" + + def __init__( + self, + load_op_links: bool = True, + *, + cache_size: int, + cache_ttl: int, + stats_prefix: str, + ) -> None: + """ + Initialize CachedDBDagBag. + + :param load_op_links: Should the extra operator link be loaded when de-serializing the DAG? + :param cache_size: Maximum cached entries. Zero means no size limit. + :param cache_ttl: Seconds until a cached entry expires. Zero disables TTL. + :param stats_prefix: Metric namespace for this component's cache. + :raises ValueError: If the metrics namespace is empty. + """ + if not stats_prefix: + raise ValueError("CachedDBDagBag requires a stats_prefix") + + super().__init__(load_op_links=load_op_links) + + if cache_ttl > 0: + self._dags = TTLCache(maxsize=cache_size or math.inf, ttl=cache_ttl) + elif cache_size > 0: + self._dags = LRUCache(maxsize=cache_size) + + # Configured caches are shared across component threads. cachetools caches need this for + # linked-list mutations, and the unbounded dict needs it for the double-checked load path. + self._lock = RLock() + self._stats_prefix = stats_prefix + + def _on_cache_hit(self) -> None: + stats.incr(f"{self._stats_prefix}.cache_hit") + + def _on_cache_miss(self) -> None: + stats.incr(f"{self._stats_prefix}.cache_miss") + + def _on_cache_clear(self) -> None: + stats.incr(f"{self._stats_prefix}.cache_clear") + + def _on_cache_size(self, *, rate: float = 1.0) -> None: + with self._lock: + size = len(self._dags) + stats.gauge(f"{self._stats_prefix}.cache_size", size, rate=rate) + + def generate_md5_hash(context): bundle_name = context.get_current_parameters()["bundle_name"] relative_fileloc = context.get_current_parameters()["relative_fileloc"] diff --git a/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py b/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py index 48c6f706ba7e2..1d11b7bc5a32e 100644 --- a/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py +++ b/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py @@ -16,14 +16,19 @@ # under the License. from __future__ import annotations +import math +import re from unittest import mock import pytest from cachetools import LRUCache, TTLCache from airflow.api_fastapi.app import purge_cached_app +from airflow.api_fastapi.common.dagbag import create_dag_bag +from airflow.models.dagbag import CachedDBDagBag from airflow.sdk import BaseOperator +from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import clear_db_dags, clear_db_runs, clear_db_serialized_dags pytestmark = pytest.mark.db_test @@ -49,13 +54,13 @@ def patch_dagbag_once_before_app(self): """Patch DagBag once before app is created, and reset counter.""" self.dagbag_call_counter["count"] = 0 - from airflow.models.dagbag import DBDagBag as RealDagBag + from airflow.models.dagbag import CachedDBDagBag as RealDagBag def factory(*args, **kwargs): self.dagbag_call_counter["count"] += 1 return RealDagBag(*args, **kwargs) - with mock.patch("airflow.api_fastapi.common.dagbag.DBDagBag", side_effect=factory): + with mock.patch("airflow.api_fastapi.common.dagbag.CachedDBDagBag", side_effect=factory): purge_cached_app() yield @@ -89,24 +94,48 @@ class TestCreateDagBag: """Tests for create_dag_bag() function.""" @pytest.mark.parametrize( - ("cache_size", "cache_ttl", "expected_use_cache", "expected_dags_type"), + ("cache_size", "cache_ttl", "expected_bag_type", "expected_dags_type", "expected_maxsize"), [ - pytest.param(64, 3600, True, TTLCache, id="default_ttl_cache"), - pytest.param(0, 3600, False, dict, id="size_zero_unbounded"), - pytest.param(64, 0, True, LRUCache, id="ttl_zero_lru_only"), + pytest.param("64", "3600", CachedDBDagBag, TTLCache, 64, id="default_ttl_cache"), + pytest.param("0", "3600", CachedDBDagBag, TTLCache, math.inf, id="size_zero_ttl_only"), + pytest.param("64", "0", CachedDBDagBag, LRUCache, 64, id="ttl_zero_lru_only"), + pytest.param("0", "0", CachedDBDagBag, dict, None, id="both_zero_no_eviction"), ], ) - @mock.patch("airflow.api_fastapi.common.dagbag.conf") def test_create_dag_bag_cache_modes( - self, mock_conf, cache_size, cache_ttl, expected_use_cache, expected_dags_type + self, + cache_size, + cache_ttl, + expected_bag_type, + expected_dags_type, + expected_maxsize, ): - from airflow.api_fastapi.common.dagbag import create_dag_bag + with conf_vars({("api", "dag_cache_size"): cache_size, ("api", "dag_cache_ttl"): cache_ttl}): + dag_bag = create_dag_bag() - mock_conf.getint.side_effect = lambda section, key, fallback: { - "dag_cache_size": cache_size, - "dag_cache_ttl": cache_ttl, - }.get(key, fallback) - - dag_bag = create_dag_bag() - assert dag_bag._use_cache is expected_use_cache + assert type(dag_bag) is expected_bag_type assert isinstance(dag_bag._dags, expected_dags_type) + if expected_maxsize is not None: + assert dag_bag._dags.maxsize == expected_maxsize + + @pytest.mark.parametrize( + ("cache_size", "cache_ttl", "expected_message"), + [ + pytest.param( + "-1", + "3600", + "[api] dag_cache_size must be greater than or equal to 0", + id="negative_size", + ), + pytest.param( + "64", + "-1", + "[api] dag_cache_ttl must be greater than or equal to 0", + id="negative_ttl", + ), + ], + ) + def test_create_dag_bag_rejects_negative_config(self, cache_size, cache_ttl, expected_message): + with conf_vars({("api", "dag_cache_size"): cache_size, ("api", "dag_cache_ttl"): cache_ttl}): + with pytest.raises(ValueError, match=re.escape(expected_message)): + create_dag_bag() diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 419b459bc3619..e01641f4f39eb 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -77,6 +77,7 @@ ) from airflow.models.dag import DagModel, get_last_dagrun, infer_automated_data_interval from airflow.models.dag_version import DagVersion +from airflow.models.dagbag import CachedDBDagBag from airflow.models.dagbundle import DagBundleModel from airflow.models.dagrun import DagRun from airflow.models.dagwarning import DagWarning @@ -416,8 +417,11 @@ def test_scheduler_dag_bag_is_bounded(self): job_runner = SchedulerJobRunner(Job()) + assert isinstance(job_runner.scheduler_dag_bag, CachedDBDagBag) assert isinstance(job_runner.scheduler_dag_bag._dags, LRUCache) assert job_runner.scheduler_dag_bag._dags.maxsize == SCHEDULER_DAG_CACHE_SIZE + # Reported separately from the API server's cache, not folded into it. + assert job_runner.scheduler_dag_bag._stats_prefix == "scheduler.dag_bag" @pytest.mark.parametrize( "heartrate", diff --git a/airflow-core/tests/unit/models/test_dagbag.py b/airflow-core/tests/unit/models/test_dagbag.py index 79668d4fe54f4..d4c5e66314b58 100644 --- a/airflow-core/tests/unit/models/test_dagbag.py +++ b/airflow-core/tests/unit/models/test_dagbag.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import math import time from concurrent.futures import ThreadPoolExecutor from unittest.mock import MagicMock, patch @@ -26,7 +27,7 @@ from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion -from airflow.models.dagbag import DBDagBag, _CacheEntry +from airflow.models.dagbag import CachedDBDagBag, DBDagBag, _CacheEntry from airflow.models.dagbundle import DagBundleModel from airflow.models.serialized_dag import SerializedDagModel from airflow.providers.standard.operators.empty import EmptyOperator @@ -38,6 +39,26 @@ pytestmark = pytest.mark.db_test +STATS_PATH = "airflow.models.dagbag.stats" + +CACHE_METRIC_SUFFIXES = ("cache_hit", "cache_miss", "cache_clear", "cache_size") + +# Every namespace a component can report under. CachedDBDagBag builds names from the prefix each +# component passes in, so the shared plumbing is exercised once per component. +METRIC_PREFIXES = ["api_server.dag_bag", "scheduler.dag_bag"] + +STUB_PREFIX = "test.dag_bag" + + +def _stub_dag_bag(*, cache_size: int, cache_ttl: int = 0) -> CachedDBDagBag: + """Build a configured cache with a test-only metric prefix.""" + return CachedDBDagBag( + cache_size=cache_size, + cache_ttl=cache_ttl, + stats_prefix=STUB_PREFIX, + ) + + # This file previously contained tests for DagBag functionality, but those tests # have been moved to airflow-core/tests/unit/dag_processing/test_dagbag.py to match # the source code reorganization where DagBag moved from models to dag_processing. @@ -244,35 +265,26 @@ def make_lazy(task_ids): class TestDBDagBagCache: - """Tests for DBDagBag optional caching behavior.""" - - def test_no_caching_by_default(self): - """Test that DBDagBag uses a simple dict without caching by default.""" - dag_bag = DBDagBag() - assert dag_bag._use_cache is False - assert isinstance(dag_bag._dags, dict) - - def test_lru_cache_enabled_with_cache_size(self): - """Test that LRU cache is enabled when cache_size is provided.""" - dag_bag = DBDagBag(cache_size=10) - assert dag_bag._use_cache is True - assert isinstance(dag_bag._dags, LRUCache) - - def test_ttl_cache_enabled_with_cache_size_and_ttl(self): - """Test that TTL cache is enabled when both cache_size and cache_ttl are provided.""" - dag_bag = DBDagBag(cache_size=10, cache_ttl=60) - assert dag_bag._use_cache is True - assert isinstance(dag_bag._dags, TTLCache) - - def test_zero_cache_size_uses_unbounded_dict(self): - """Test that cache_size=0 uses unbounded dict (same as no caching).""" - dag_bag = DBDagBag(cache_size=0, cache_ttl=60) - assert dag_bag._use_cache is False - assert isinstance(dag_bag._dags, dict) + """Tests for plain and configured DBDagBag caching behavior.""" + + @pytest.mark.parametrize( + ("cache_size", "cache_ttl", "expected_type", "expected_maxsize"), + [ + pytest.param(10, 0, LRUCache, 10, id="size_only_lru"), + pytest.param(10, 60, TTLCache, 10, id="size_and_ttl_bounded_ttl"), + pytest.param(0, 60, TTLCache, math.inf, id="ttl_only_uncapped"), + pytest.param(0, 0, dict, None, id="no_eviction"), + ], + ) + def test_cache_selection(self, cache_size, cache_ttl, expected_type, expected_maxsize): + dag_bag = _stub_dag_bag(cache_size=cache_size, cache_ttl=cache_ttl) + assert isinstance(dag_bag._dags, expected_type) + if expected_maxsize is not None: + assert dag_bag._dags.maxsize == expected_maxsize def test_clear_cache_with_caching(self): """Test clear_cache() with caching enabled.""" - dag_bag = DBDagBag(cache_size=10, cache_ttl=60) + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) mock_dag = MagicMock() dag_bag._dags["version_1"] = mock_dag @@ -283,6 +295,46 @@ def test_clear_cache_with_caching(self): assert count == 2 assert len(dag_bag._dags) == 0 + @pytest.mark.parametrize("prefix", METRIC_PREFIXES) + def test_stats_prefix_expands_to_registered_metrics(self, prefix): + """Every name a component can emit must exist in the metrics registry. + + The registry prek check sees only the ``{_stats_prefix}.`` template, so it can + verify the suffixes but not the prefix each component supplies. This pins the expanded + names so a renamed or misspelled prefix cannot ship unregistered. + """ + from airflow._shared.observability.metrics.metrics_registry import MetricsRegistry + + registry = MetricsRegistry() + missing = [ + name for suffix in CACHE_METRIC_SUFFIXES if registry.get(name := f"{prefix}.{suffix}") is None + ] + assert not missing + + def test_api_server_reports_under_its_own_namespace(self): + from airflow.api_fastapi.common.dagbag import create_dag_bag + + assert create_dag_bag()._stats_prefix == "api_server.dag_bag" + + def test_cached_bag_requires_non_empty_stats_prefix(self): + """A cache with no namespace to report under must fail at wiring time, not mid-request.""" + with pytest.raises(ValueError, match="requires a stats_prefix"): + CachedDBDagBag(cache_size=10, cache_ttl=60, stats_prefix="") + + def test_plain_bag_emits_no_metrics(self): + """The unbounded base implementation does not report component cache metrics.""" + dag_bag = DBDagBag() + mock_serdag = MagicMock() + mock_serdag.dag_version_id = "test_version_1" + mock_serdag.dag = MagicMock() + + with patch(STATS_PATH) as mock_stats: + dag_bag._read_dag(mock_serdag) + dag_bag.clear_cache() + + mock_stats.incr.assert_not_called() + mock_stats.gauge.assert_not_called() + def test_clear_cache_without_caching(self): """Test clear_cache() without caching enabled.""" dag_bag = DBDagBag() @@ -299,7 +351,7 @@ def test_ttl_cache_expiry(self): """Test that cached DAGs expire after TTL.""" # TTLCache defaults to time.monotonic which time_machine cannot control. # Use time.time as the timer so time_machine can advance it. - dag_bag = DBDagBag(cache_size=10, cache_ttl=1) + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=1) dag_bag._dags = TTLCache(maxsize=10, ttl=1, timer=time.time) with time_machine.travel("2025-01-01 00:00:00", tick=False): @@ -312,7 +364,7 @@ def test_ttl_cache_expiry(self): def test_lru_eviction(self): """Test that LRU eviction works when cache is full.""" - dag_bag = DBDagBag(cache_size=2) + dag_bag = _stub_dag_bag(cache_size=2) dag_bag._dags["version_1"] = MagicMock() dag_bag._dags["version_2"] = MagicMock() @@ -325,7 +377,7 @@ def test_lru_eviction(self): def test_thread_safety_with_caching(self): """Test concurrent access doesn't cause race conditions with caching enabled.""" - dag_bag = DBDagBag(cache_size=100, cache_ttl=60) + dag_bag = _stub_dag_bag(cache_size=100, cache_ttl=60) errors = [] mock_session = MagicMock() @@ -355,7 +407,7 @@ def access_cache(i): def test_read_dag_stores_in_bounded_cache(self): """Test that _read_dag stores DAG in bounded cache when cache_size > 0.""" - dag_bag = DBDagBag(cache_size=10, cache_ttl=60) + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) mock_sdm = MagicMock() mock_sdm.dag = MagicMock() @@ -381,7 +433,7 @@ def test_read_dag_stores_in_unbounded_dict(self): def test_iter_all_latest_version_dags_does_not_cache(self): """Test that iter_all_latest_version_dags does not cache to prevent thrashing.""" - dag_bag = DBDagBag(cache_size=10, cache_ttl=60) + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) mock_session = MagicMock() mock_sdm = MagicMock() @@ -397,7 +449,7 @@ def test_iter_all_latest_version_dags_does_not_cache(self): @patch("airflow.models.dagbag.stats") def test_cache_hit_metric_emitted(self, mock_stats): """Test that cache hit metric is emitted when caching is enabled.""" - dag_bag = DBDagBag(cache_size=10, cache_ttl=60) + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) mock_session = MagicMock() # last_validated=0.0 forces revalidation; the hash matches, so it counts as a hit. dag_bag._dags["test_version"] = _CacheEntry(MagicMock(), "hash1", 0.0) @@ -405,12 +457,12 @@ def test_cache_hit_metric_emitted(self, mock_stats): dag_bag._get_dag("test_version", mock_session) - mock_stats.incr.assert_called_with("api_server.dag_bag.cache_hit") + mock_stats.incr.assert_called_with(f"{STUB_PREFIX}.cache_hit") @patch("airflow.models.dagbag.stats") def test_cache_miss_metric_emitted(self, mock_stats): """Test that cache miss metric is emitted when DAG is found in DB but not in cache.""" - dag_bag = DBDagBag(cache_size=10, cache_ttl=60) + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) mock_session = MagicMock() # Set up a DB result so _get_dag reaches the miss metric path @@ -423,22 +475,22 @@ def test_cache_miss_metric_emitted(self, mock_stats): dag_bag._get_dag("uncached_version", mock_session) - mock_stats.incr.assert_any_call("api_server.dag_bag.cache_miss") + mock_stats.incr.assert_any_call(f"{STUB_PREFIX}.cache_miss") @patch("airflow.models.dagbag.stats") def test_cache_clear_metric_emitted(self, mock_stats): """Test that cache clear metric is emitted when caching is enabled.""" - dag_bag = DBDagBag(cache_size=10, cache_ttl=60) + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) dag_bag._dags["test_version"] = MagicMock() dag_bag.clear_cache() - mock_stats.incr.assert_called_with("api_server.dag_bag.cache_clear") + mock_stats.incr.assert_called_with(f"{STUB_PREFIX}.cache_clear") @patch("airflow.models.dagbag.stats") def test_cache_size_gauge_emitted(self, mock_stats): """Test that cache size gauge is emitted when a DAG is cached.""" - dag_bag = DBDagBag(cache_size=10, cache_ttl=60) + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) mock_serdag = MagicMock() mock_serdag.dag_version_id = "test_version_1" mock_serdag.dag = MagicMock() @@ -446,4 +498,4 @@ def test_cache_size_gauge_emitted(self, mock_stats): dag_bag._read_dag(mock_serdag) - mock_stats.gauge.assert_called_with("api_server.dag_bag.cache_size", 1, rate=0.1) + mock_stats.gauge.assert_called_with(f"{STUB_PREFIX}.cache_size", 1, rate=0.1) diff --git a/scripts/ci/prek/check_metrics_synced_with_the_registry.py b/scripts/ci/prek/check_metrics_synced_with_the_registry.py index 94b5bc254de28..239fbbccf0fe1 100644 --- a/scripts/ci/prek/check_metrics_synced_with_the_registry.py +++ b/scripts/ci/prek/check_metrics_synced_with_the_registry.py @@ -73,13 +73,50 @@ def normalize_metric_name(registry_metric_name: str) -> str: "{job_name}_start" → "*_start" "pool.open_slots" → "pool.open_slots" """ - return re.sub(r"\{[^}]+\}", "*", registry_metric_name) + return _VARIABLE_RE.sub("*", registry_metric_name) -# Sentinel returned when a dynamic metric name is partially matched based on a common prefix. +# Sentinel returned when a dynamic metric name is structurally matched against registry entries. # For dynamic metric names that include variables, the check can't find an exact match with a registry -# entry or its type. So, a partially matched prefix is good enough and type checking is skipped. -_PREFIX_MATCHED = "__prefix_matched__" +# entry or its type. So, a structural match is good enough and type checking is skipped. +_PATTERN_MATCHED = "__pattern_matched__" + +# A ``{variable}`` stands for one or more dot-separated segments, so one pattern covers both a +# single-segment substitution (``{state}`` -> ``running``) and a multi-segment one +# (``{stats_prefix}`` -> ``api_server.dag_bag``). +_VARIABLE_SEGMENTS = r"[^.]+(?:\.[^.]+)*" + +# The ``{variable}`` placeholder itself, shared by name normalization and pattern compilation. +_VARIABLE_RE = re.compile(r"\{[^}]+\}") + + +def compile_dynamic_metric_pattern(metric_name: str) -> re.Pattern[str]: + """Compile a ``{variable}``-containing metric name into a regex over its static parts. + + Matching on the whole shape rather than only the prefix before the first variable means a + variable may sit anywhere in the name, including at the start or between static parts:: + + "ti.{state}" matches "ti.running" + "{stats_prefix}.cache_hit" matches "api_server.dag_bag.cache_hit" + "{prefix}.foo.{state}.duration" matches "a.b.foo.success.duration" + """ + literals = _VARIABLE_RE.split(metric_name) + return re.compile(_VARIABLE_SEGMENTS.join(re.escape(literal) for literal in literals)) + + +def find_pattern_matched_registry_entries(metric_name: str, metrics_registry: dict[str, dict]) -> list[str]: + """Return the registry entry names a dynamic metric name structurally matches.""" + literals = _VARIABLE_RE.split(metric_name) + if len(literals) == 1: + # Static name: the exact and normalized lookups already had their chance. + return [] + if not any(literals): + # Nothing but variables, e.g. ``{name}`` or ``{prefix}{suffix}``. The pattern would be a + # bare "any segments" regex matching every entry, which marks the whole registry used and + # silently disables the unused-entry check. Match nothing so the name is reported missing. + return [] + pattern = compile_dynamic_metric_pattern(metric_name) + return [name for name in metrics_registry if pattern.fullmatch(name)] def find_registry_match(metric_name: str, metrics_registry: dict[str, dict]) -> str | None: @@ -99,16 +136,11 @@ def find_registry_match(metric_name: str, metrics_registry: dict[str, dict]) -> return registry_metric_name # Dynamic metric name. - if "{" in metric_name: - base = metric_name.split("{")[0].rstrip(".") - for registry_metric_name in metrics_registry: - if registry_metric_name == base or registry_metric_name.startswith(base + "."): - # Metric prefix matches the prefix of a dynamic registry entry. - # If the static part before the first variable, matches an exact registry entry name, - # or a dotted-prefix of one, then τηε name is considered covered and - # _PREFIX_MATCHED is returned. The type check must be skipped because - # the resulting metric name with all variables expanded, cannot be determined. - return _PREFIX_MATCHED + if find_pattern_matched_registry_entries(metric_name, metrics_registry): + # The name's static parts line up with at least one registry entry, so it is considered + # covered and _PATTERN_MATCHED is returned. The type check must be skipped because the + # resulting metric name with all variables expanded cannot be determined. + return _PATTERN_MATCHED # All checks for matching failed. return None @@ -249,9 +281,9 @@ def main() -> None: metrics_with_type_mismatch: dict[str, list[tuple[MetricCall, str, str]]] = {} for name, calls in code_metrics.items(): registry_metric_name = find_registry_match(name, metrics_registry) - if registry_metric_name is None or registry_metric_name is _PREFIX_MATCHED: + if registry_metric_name is None or registry_metric_name is _PATTERN_MATCHED: # If None, then it's reported as missing, no need for type check. - # If _PREFIX_MATCHED, then the exact entry can't be determined. Skip the type check. + # If _PATTERN_MATCHED, then the exact entry can't be determined. Skip the type check. continue registry_type = metrics_registry[registry_metric_name].get("type", "").lower() mismatched = [ diff --git a/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py b/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py index 98f7f1ee2567d..e93d8780af1e9 100644 --- a/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py +++ b/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py @@ -22,8 +22,9 @@ import pytest from ci.prek.check_metrics_synced_with_the_registry import ( - _PREFIX_MATCHED, + _PATTERN_MATCHED, extract_metric_name_from_ast_node, + find_pattern_matched_registry_entries, find_registry_match, get_stats_obj_name, normalize_metric_name, @@ -108,9 +109,9 @@ def test_normalize_metric_name(metric_name, expected_result): # In this case, the legacy name of 'task.duration', is 'dag.{dag_id}.{task_id}.duration'. # Once normalized, both will be 'dag.*.*.duration' and there should be a match. pytest.param("dag.{x}.{y}.duration", "task.duration", id="legacy_name_match_different_structure"), - pytest.param("ti.{state}", _PREFIX_MATCHED, id="prefix_match_returns_sentinel"), - pytest.param("dagrun.duration.{state}", _PREFIX_MATCHED, id="prefix_match_dotted_base"), - pytest.param("non.existent.{var}", None, id="dynamic_metric_no_prefix_match_returns_none"), + pytest.param("ti.{state}", _PATTERN_MATCHED, id="pattern_match_returns_sentinel"), + pytest.param("dagrun.duration.{state}", _PATTERN_MATCHED, id="pattern_match_dotted_static_part"), + pytest.param("non.existent.{var}", None, id="dynamic_metric_no_pattern_match_returns_none"), pytest.param("non.existent", None, id="static_metric_not_in_registry_returns_none"), ], ) @@ -180,6 +181,72 @@ def test_extract_metric_name_from_ast_node(code: str, expected_result): assert extract_metric_name_from_ast_node(node) == expected_result +@pytest.mark.parametrize( + "metric_name, expected_result", + [ + pytest.param( + "ti.{state}", + ["ti.scheduled", "ti.queued", "ti.start.{dag_id}.{task_id}"], + id="pattern_matches_multiple_entries", + ), + pytest.param("dagrun.duration.{state}", ["dagrun.duration.success"], id="dotted_static_part"), + pytest.param("non.existent.{var}", [], id="no_pattern_match_returns_empty_list"), + ], +) +def test_find_pattern_matched_registry_entries(metric_name, expected_result): + assert find_pattern_matched_registry_entries(metric_name, METRICS_REGISTRY) == expected_result + + +# A registry whose entries share a suffix but differ in how many segments precede it, which is the +# shape produced by a metric name built from a per-component prefix. +PREFIXED_METRICS_REGISTRY = { + "api_server.dag_bag.cache_hit": {"name": "api_server.dag_bag.cache_hit", "type": "counter"}, + "scheduler.dag_bag.cache_hit": {"name": "scheduler.dag_bag.cache_hit", "type": "counter"}, + "pool.open_slots": {"name": "pool.open_slots", "type": "gauge"}, + "a.b.foo.success.duration": {"name": "a.b.foo.success.duration", "type": "timer"}, +} + + +@pytest.mark.parametrize( + "metric_name, expected_result", + [ + pytest.param( + "{stats_prefix}.cache_hit", + ["api_server.dag_bag.cache_hit", "scheduler.dag_bag.cache_hit"], + id="leading_variable_spans_multiple_segments", + ), + pytest.param( + "{prefix}.foo.{state}.duration", + ["a.b.foo.success.duration"], + id="variables_around_a_static_middle", + ), + pytest.param("{stats_prefix}.cache_miss", [], id="unregistered_suffix_matches_nothing"), + pytest.param("{prefix}.open_slots", ["pool.open_slots"], id="single_segment_prefix"), + ], +) +def test_find_pattern_matched_registry_entries_with_variable_prefix(metric_name, expected_result): + """A variable anywhere in the name resolves, which static-prefix matching could not do.""" + assert find_pattern_matched_registry_entries(metric_name, PREFIXED_METRICS_REGISTRY) == expected_result + + +def test_pattern_match_does_not_cross_static_parts(): + """The static parts must line up, so a name is not matched just because it shares a suffix.""" + assert find_pattern_matched_registry_entries("{prefix}.bar.duration", PREFIXED_METRICS_REGISTRY) == [] + + +@pytest.mark.parametrize( + "metric_name", + [ + pytest.param("{variable}", id="single_variable"), + pytest.param("{prefix}{suffix}", id="adjacent_variables"), + ], +) +def test_all_variable_name_matches_nothing(metric_name): + """A name with no static part must not match every registry entry.""" + assert find_pattern_matched_registry_entries(metric_name, PREFIXED_METRICS_REGISTRY) == [] + assert find_registry_match(metric_name, PREFIXED_METRICS_REGISTRY) is None + + @pytest.fixture def code_to_py_file(tmp_path): """Write python source code to a tmp file and return its path.""" diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index 6c51e32ff7798..b191d85ba34d3 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -345,6 +345,24 @@ metrics: legacy_name: "-" name_variables: [] + - name: "scheduler.dag_bag.cache_hit" + description: "Number of cache hits when retrieving SerializedDAG from DBDagBag in the scheduler" + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "scheduler.dag_bag.cache_miss" + description: "Number of cache misses when retrieving SerializedDAG from DBDagBag in the scheduler" + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "scheduler.dag_bag.cache_clear" + description: "Number of times the DBDagBag cache was cleared in the scheduler" + type: "counter" + legacy_name: "-" + name_variables: [] + - name: "connection_test.success" description: "Number of worker-dispatched connection tests that completed successfully." type: "counter" @@ -379,6 +397,12 @@ metrics: legacy_name: "-" name_variables: [] + - name: "scheduler.dag_bag.cache_size" + description: "Current number of SerializedDAG objects cached in the scheduler's DBDagBag" + type: "gauge" + legacy_name: "-" + name_variables: [] + - name: "connection_test.active" description: "Number of connection tests currently in flight (``queued`` + ``running``), sampled by the scheduler each tick."