This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bcf1dbf41a0 Report Dag cache metrics under each component's own
namespace (#71815)
bcf1dbf41a0 is described below
commit bcf1dbf41a0593e44398c04b3d2aee7eeabe6ed5
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Fri Aug 21 15:06:47 2026 +0800
Report Dag cache metrics under each component's own namespace (#71815)
* Report Dag cache metrics under each component's own namespace
Every DBDagBag emitted its cache counters under api_server.dag_bag.*, so
once the
scheduler gained a cache its traffic was silently counted against the API
server's series. An operator reading cache_hit or cache_size could not tell
the
two components apart, and the API server's numbers became wrong rather than
merely incomplete.
Each caller now supplies the namespace it reports under, and a cache built
without one fails at construction instead of emitting a partially-formed
metric
name mid-request.
The registry check matched dynamic metric names only by the static prefix
ahead
of their first variable, which cannot express a name assembled from a caller
supplied prefix. It now matches on all static parts wherever the variable
sits.
* Scope Dag cache configuration to cached database Dag bags
DBDagBag is used by callers that do not configure cache eviction or
metrics. Keeping those concerns behind a dedicated subtype prevents optional
constructor combinations and makes metric ownership explicit.
* Preserve no-eviction Dag cache configuration
* Avoid redundant Dag cache checks
---
.../src/airflow/api_fastapi/common/dagbag.py | 10 +-
.../src/airflow/jobs/scheduler_job_runner.py | 9 +-
airflow-core/src/airflow/models/dagbag.py | 127 ++++++++++++---------
.../tests/unit/api_fastapi/common/test_dagbag.py | 26 +++--
airflow-core/tests/unit/jobs/test_scheduler_job.py | 4 +
airflow-core/tests/unit/models/test_dagbag.py | 120 ++++++++++++-------
.../prek/check_metrics_synced_with_the_registry.py | 67 ++++++++---
.../test_check_metrics_synced_with_the_registry.py | 79 +++++++++++--
.../observability/metrics/metrics_template.yaml | 24 ++++
9 files changed, 329 insertions(+), 137 deletions(-)
diff --git a/airflow-core/src/airflow/api_fastapi/common/dagbag.py
b/airflow-core/src/airflow/api_fastapi/common/dagbag.py
index 85ce253fcb8..b96e6f65b06 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 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:
@@ -30,7 +30,7 @@ if TYPE_CHECKING:
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 d744de41816..6b08dab3f02 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -100,7 +100,7 @@ from airflow.models.connection_test import (
)
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
@@ -383,7 +383,12 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
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 dfa336f2fe0..92662d067a8 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 @@ class DBDagBag:
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 @@ class DBDagBag:
# [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 @@ class DBDagBag:
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 @@ class DBDagBag:
# 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 @@ class DBDagBag:
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 @@ class DBDagBag:
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 d56e96a24a3..1d11b7bc5a3 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 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
@@ -53,13 +54,13 @@ class TestDagBagSingleton:
"""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 88fcb0555ce..9f14dee0b78 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.connection_test import (
)
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 @@ class TestSchedulerJob:
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 8dbc35ea385..d4c5e66314b 100644
--- a/airflow-core/tests/unit/models/test_dagbag.py
+++ b/airflow-core/tests/unit/models/test_dagbag.py
@@ -27,7 +27,7 @@ from cachetools import LRUCache, TTLCache
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 @@ from tests_common.test_utils import db
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 @@ class TestDBDagBag:
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 @@ class TestDBDagBagCache:
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}.<suffix>``
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 @@ class TestDBDagBagCache:
"""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 @@ class TestDBDagBagCache:
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 @@ class TestDBDagBagCache:
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 @@ class TestDBDagBagCache:
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 @@ class TestDBDagBagCache:
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 @@ class TestDBDagBagCache:
@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 @@ class TestDBDagBagCache:
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 @@ class TestDBDagBagCache:
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 @@ class TestDBDagBagCache:
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 140e6ab9e91..9a1a20d8658 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 e9b5e2ebe29..0c817343778 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 @@ from unittest import mock
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"},
+}
+
+
[email protected](
+ "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) == []
+
+
[email protected](
+ "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 392b55d7c07..edac39bb48c 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."
