diff --git a/airflow-core/src/airflow/api_fastapi/common/dagbag.py b/airflow-core/src/airflow/api_fastapi/common/dagbag.py index 85ce253fcb8f2..b96e6f65b061c 100644 --- a/airflow-core/src/airflow/api_fastapi/common/dagbag.py +++ b/airflow-core/src/airflow/api_fastapi/common/dagbag.py @@ -22,7 +22,7 @@ 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: @@ -30,7 +30,7 @@ from airflow.serialization.definitions.dag import SerializedDAG -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 = conf.getint("api", "dag_cache_ttl", fallback=3600) @@ -40,7 +40,11 @@ def create_dag_bag() -> DBDagBag: if cache_ttl < 0: raise ValueError("[api] dag_cache_ttl must be greater than or equal to 0") - return DBDagBag(cache_size=cache_size, cache_ttl=cache_ttl) + 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 2153f4267588f..7d28643702cea 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -99,7 +99,7 @@ ) from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion, _resolve_version_data -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 @@ -382,7 +382,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 dfa336f2fe0a5..92662d067a8b5 100644 --- a/airflow-core/src/airflow/models/dagbag.py +++ b/airflow-core/src/airflow/models/dagbag.py @@ -63,55 +63,34 @@ class DBDagBag: """ Internal class for retrieving dags from the database. - Optionally caches deserialized dags: a size limit enables LRU eviction, and a TTL enables - age-based eviction with or without a size limit. Callers that pass neither get a plain dict - that never evicts. + 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: Max cached entries. 0 or None means no size limit. - :param cache_ttl: Seconds until a cached entry expires, applied with or without a size limit. - 0 or None disables TTL. With neither a size limit nor a TTL the cache never evicts. - :raises ValueError: If ``cache_size`` or ``cache_ttl`` is negative. """ - # Callers should reject negative values with their own context; validate again defensively. - if cache_size is not None and cache_size < 0: - raise ValueError("cache_size must be greater than or equal to 0") - if cache_ttl is not None and cache_ttl < 0: - raise ValueError("cache_ttl must be greater than or equal to 0") - 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.""" + + def _on_cache_miss(self) -> None: + """Handle a Dag cache miss.""" - # A TTL applies with or without a size limit: an uncapped TTLCache is what lets - # ``dag_cache_size = 0`` mean "no size limit" rather than "no eviction at all". - size = cache_size or 0 - ttl = cache_ttl or 0 - if ttl > 0: - self._dags = TTLCache(maxsize=size or math.inf, ttl=ttl) - self._use_cache = True - elif size > 0: - self._dags = LRUCache(maxsize=size) - self._use_cache = True - - # Lock required for bounded caches: cachetools caches are NOT thread-safe - # (LRU reordering and TTL cleanup mutate internal linked lists). A plain dict needs no - # lock, so it uses nullcontext. - 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).""" @@ -121,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 @@ -145,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 @@ -160,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: @@ -178,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: @@ -214,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 @@ -258,6 +231,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 d56e96a24a3da..1d11b7bc5a32e 100644 --- a/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py +++ b/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py @@ -25,6 +25,7 @@ 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 @@ -53,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 @@ -93,20 +94,27 @@ class TestCreateDagBag: """Tests for create_dag_bag() function.""" @pytest.mark.parametrize( - ("cache_size", "cache_ttl", "expected_dags_type", "expected_maxsize"), + ("cache_size", "cache_ttl", "expected_bag_type", "expected_dags_type", "expected_maxsize"), [ - pytest.param("64", "3600", TTLCache, 64, id="default_ttl_cache"), - pytest.param("0", "3600", TTLCache, math.inf, id="size_zero_ttl_only"), - pytest.param("64", "0", LRUCache, 64, id="ttl_zero_lru_only"), - pytest.param("0", "0", dict, None, id="both_zero_no_eviction"), + 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"), ], ) - def test_create_dag_bag_cache_modes(self, cache_size, cache_ttl, expected_dags_type, expected_maxsize): + def test_create_dag_bag_cache_modes( + self, + cache_size, + cache_ttl, + expected_bag_type, + expected_dags_type, + expected_maxsize, + ): with conf_vars({("api", "dag_cache_size"): cache_size, ("api", "dag_cache_ttl"): cache_ttl}): dag_bag = create_dag_bag() + assert type(dag_bag) is expected_bag_type assert isinstance(dag_bag._dags, expected_dags_type) - assert dag_bag._use_cache is (expected_dags_type is not dict) if expected_maxsize is not None: assert dag_bag._dags.maxsize == expected_maxsize diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index fd3b0fd096438..21e49d303d6db 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -84,6 +84,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 @@ -420,8 +421,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 8dbc35ea385be..d4c5e66314b58 100644 --- a/airflow-core/tests/unit/models/test_dagbag.py +++ b/airflow-core/tests/unit/models/test_dagbag.py @@ -27,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 @@ -39,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. @@ -245,48 +265,26 @@ def make_lazy(task_ids): class TestDBDagBagCache: - """Tests for DBDagBag optional caching behavior.""" + """Tests for plain and configured DBDagBag caching behavior.""" @pytest.mark.parametrize( ("cache_size", "cache_ttl", "expected_type", "expected_maxsize"), [ - pytest.param(None, None, dict, None, id="neither_plain_dict"), - pytest.param(10, None, LRUCache, 10, id="size_only_lru"), + 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 = DBDagBag(cache_size=cache_size, cache_ttl=cache_ttl) + dag_bag = _stub_dag_bag(cache_size=cache_size, cache_ttl=cache_ttl) assert isinstance(dag_bag._dags, expected_type) - assert dag_bag._use_cache is (expected_type is not dict) 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, - None, - "cache_size must be greater than or equal to 0", - id="negative_size", - ), - pytest.param( - None, - -1, - "cache_ttl must be greater than or equal to 0", - id="negative_ttl", - ), - ], - ) - def test_rejects_negative_cache_configuration(self, cache_size, cache_ttl, expected_message): - with pytest.raises(ValueError, match=expected_message): - DBDagBag(cache_size=cache_size, cache_ttl=cache_ttl) - 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 @@ -297,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() @@ -313,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): @@ -326,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() @@ -339,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() @@ -369,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() @@ -395,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() @@ -411,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) @@ -419,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 @@ -437,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() @@ -460,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 140e6ab9e91ed..9a1a20d8658c0 100644 --- a/scripts/ci/prek/check_metrics_synced_with_the_registry.py +++ b/scripts/ci/prek/check_metrics_synced_with_the_registry.py @@ -94,19 +94,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"[^.]+(?:\.[^.]+)*" -def find_prefix_matched_registry_entries(metric_name: str, metrics_registry: dict[str, dict]) -> list[str]: - """Return the registry entry names whose name matches the static prefix of a dynamic metric name.""" - base = metric_name.split("{")[0].rstrip(".") - return [name for name in metrics_registry if name == base or name.startswith(base + ".")] +# 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: @@ -126,13 +157,11 @@ def find_registry_match(metric_name: str, metrics_registry: dict[str, dict]) -> return registry_metric_name # Dynamic metric name. - if "{" in metric_name and find_prefix_matched_registry_entries(metric_name, metrics_registry): - # 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 the 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 @@ -381,8 +410,8 @@ def compute_unused_registry_entries( registry_metric_name = find_registry_match(metric_name, metrics_registry) if registry_metric_name is None: continue - if registry_metric_name is _PREFIX_MATCHED: - used_entries.update(find_prefix_matched_registry_entries(metric_name, metrics_registry)) + if registry_metric_name is _PATTERN_MATCHED: + used_entries.update(find_pattern_matched_registry_entries(metric_name, metrics_registry)) else: used_entries.add(registry_metric_name) return sorted(set(metrics_registry) - used_entries) @@ -439,9 +468,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 e9b5e2ebe2979..0c8173437785f 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 @@ -24,13 +24,13 @@ import pytest from ci.prek import check_metrics_synced_with_the_registry from ci.prek.check_metrics_synced_with_the_registry import ( - _PREFIX_MATCHED, + _PATTERN_MATCHED, _except_handler_catches_expected_error, _is_stats_module_path, compute_unused_registry_entries, extract_metric_name_from_ast_node, extract_metric_names_from_ast_node, - find_prefix_matched_registry_entries, + find_pattern_matched_registry_entries, find_registry_match, find_stale_indirectly_emitted_metrics, get_stats_obj_name, @@ -117,9 +117,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"), ], ) @@ -222,14 +222,71 @@ def test_extract_metric_names_from_ast_node(code: str, expected_result): pytest.param( "ti.{state}", ["ti.scheduled", "ti.queued", "ti.start.{dag_id}.{task_id}"], - id="base_prefix_matches_multiple_entries", + id="pattern_matches_multiple_entries", ), - pytest.param("dagrun.duration.{state}", ["dagrun.duration.success"], id="dotted_base_prefix"), - pytest.param("non.existent.{var}", [], id="no_prefix_match_returns_empty_list"), + 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_prefix_matched_registry_entries(metric_name, expected_result): - assert find_prefix_matched_registry_entries(metric_name, METRICS_REGISTRY) == expected_result +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, or it marks the whole registry used. + + Its pattern would be a bare "any segments" regex, so every entry would fullmatch and + ``compute_unused_registry_entries`` would go permanently empty -- the check failing open. + """ + assert find_pattern_matched_registry_entries(metric_name, PREFIXED_METRICS_REGISTRY) == [] + assert find_registry_match(metric_name, PREFIXED_METRICS_REGISTRY) is None + assert compute_unused_registry_entries({metric_name}, PREFIXED_METRICS_REGISTRY) == sorted( + PREFIXED_METRICS_REGISTRY + ) # 'executor.open_slots' is in INDIRECTLY_EMITTED_METRICS, so it is never reported as unused. @@ -263,7 +320,7 @@ def test_find_prefix_matched_registry_entries(metric_name, expected_result): pytest.param( {"ti.{state}"}, ["dagrun.duration.success", "pool.open_slots", "scheduler.heartbeat", "task.duration"], - id="prefix_match_marks_all_prefix_entries_used", + id="pattern_match_marks_all_matched_entries_used", ), pytest.param( {"pool.open_slots.{my_pool}"}, 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 392b55d7c072f..edac39bb48c07 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -361,6 +361,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" @@ -395,6 +413,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."