diff --git a/airflow-core/docs/administration-and-deployment/web-stack.rst b/airflow-core/docs/administration-and-deployment/web-stack.rst index f043789d5e553..55146ce9a7316 100644 --- a/airflow-core/docs/administration-and-deployment/web-stack.rst +++ b/airflow-core/docs/administration-and-deployment/web-stack.rst @@ -142,8 +142,8 @@ The following configuration options are available in the ``[api]`` section: - ``server_type``: ``uvicorn`` (default) or ``gunicorn`` - ``worker_refresh_interval``: Seconds between worker refresh cycles (0 = disabled, default) - ``worker_refresh_batch_size``: Number of workers to refresh per cycle (default: 1) -- ``dag_cache_size``: Max cached SerializedDAG versions in the API server (default: 64, 0 = unbounded) -- ``dag_cache_ttl``: TTL in seconds for cached DAGs (default: 3600, 0 = LRU only) +- ``dag_cache_size``: Max cached SerializedDAG versions in the API server (default: 64, 0 = no size limit) +- ``dag_cache_ttl``: Idle timeout in seconds for cached Dags (default: 3600, 0 = no TTL; both 0 = no eviction) When to Use Gunicorn ^^^^^^^^^^^^^^^^^^^^ @@ -188,9 +188,10 @@ For example, to trigger a rolling restart of the API server pods: kubectl rollout restart deployment airflow-api-server -The API server also supports bounded DAG caching via ``dag_cache_size`` and -``dag_cache_ttl``, which limits memory consumed by cached SerializedDAG objects. -This reduces memory growth from DAG version accumulation regardless of server type. +The API server also evicts cached SerializedDAG objects via ``dag_cache_size`` and +``dag_cache_ttl``, which reduces memory growth from Dag version accumulation regardless of +server type. Note that only ``dag_cache_size`` caps memory outright: each re-check resets a +cached entry's expiry, so ``dag_cache_ttl`` reclaims only the versions that stop being requested. In many Kubernetes environments, relying solely on Kubernetes OOM kills or crash restarts is not recommended, as memory growth may not always trigger an diff --git a/airflow-core/docs/faq.rst b/airflow-core/docs/faq.rst index ce89ecac393ce..5caccc5b68a3c 100644 --- a/airflow-core/docs/faq.rst +++ b/airflow-core/docs/faq.rst @@ -717,22 +717,29 @@ The API server caches serialized Dag objects in memory. Over time, as Dag versio There are two complementary approaches: -**1. Bounded DAG caching (available since Airflow 3.3.0)** +**1. Dag cache eviction (available since Airflow 3.2.2)** -The API server supports LRU+TTL caching that bounds how many serialized Dag versions are kept -in memory. Configure this in the ``[api]`` section: +The API server can evict cached serialized Dag versions by size, by age, or both. Configure +this in the ``[api]`` section: .. code-block:: ini [api] - dag_cache_size = 64 ; max cached versions (0 = unbounded, pre-3.2 behavior) - dag_cache_ttl = 3600 ; seconds before a cached entry expires (0 = LRU only) + dag_cache_size = 64 ; max cached versions (0 = no size limit) + dag_cache_ttl = 3600 ; seconds before an idle cached entry expires (0 = no TTL) + +``dag_cache_size`` is the only hard ceiling on memory. Each re-check resets a cached entry's +expiry, so ``dag_cache_ttl`` reclaims only the Dag versions that stop being requested for the +whole interval; versions that keep being requested are kept. It therefore trims the versions +that accumulate over time, but memory still tracks the set of versions in active use. Setting +both to 0 uses an unbounded dict with no eviction, matching the behavior before 3.2.2. The cache is keyed by Dag version ID. After a Dag is updated, the API server may serve the previous version until the cached entry expires (controlled by ``dag_cache_ttl``). See :ref:`config:api__dag_cache_size` and :ref:`config:api__dag_cache_ttl` for the full -configuration reference. +configuration reference. The scheduler has the same pair of options under ``[scheduler]`` +(see :ref:`faq:scheduler-memory-growth`). **2. Gunicorn with rolling worker restarts (available since Airflow 3.2.0)** @@ -758,9 +765,48 @@ See :ref:`config:api__server_type`, :ref:`config:api__worker_refresh_interval`, .. note:: Worker recycling handles memory growth from *any* source, not just the Dag cache. - For production deployments, using both bounded caching and gunicorn worker recycling + For production deployments, using both cache eviction and gunicorn worker recycling provides the best results. +.. _faq:scheduler-memory-growth: + +How to prevent scheduler memory growth? +---------------------------------------- + +The scheduler caches deserialized Dag objects, so before 3.4.0 its memory grew with every Dag +version it had ever seen (see :ref:`faq:dag-version-inflation`) until the process was restarted or +OOM killed. The cache is now bounded by default. Tune it in the ``[scheduler]`` section: + +.. code-block:: ini + + [scheduler] + dag_cache_size = 512 ; max cached versions, evicting least recently used + dag_cache_ttl = 0 ; seconds before an idle cached entry expires (0 = no TTL) + +The defaults cap the cache at 512 Dag versions and evict the least recently used one beyond +that, so memory cannot grow with the number of versions the scheduler has ever seen. The +scheduler reaches the cache through the Dag version of each active Dag run, so its working set is +the versions with runs in flight; 512 is meant to sit above that for a typical deployment. + +If the scheduler is still OOM killed, lower ``dag_cache_size``. Raise it if you have more Dag +versions with runs in flight than the limit, since a limit below the working set evicts versions +that are still being scheduled and costs a database fetch and a deserialization on the next loop +— watch ``scheduler.dag_bag.cache_miss`` to tell the two apart. Setting ``dag_cache_size = 0`` +switches to no size limit, leaving eviction to ``dag_cache_ttl``, which bounds memory by the +concurrently active set rather than outright: each re-check resets an entry's expiry, so a TTL +reclaims a version only once its runs finish and it stops being requested. Setting both to 0 uses +an unbounded dict with no eviction, matching the behavior before 3.4.0. + +Neither option affects how quickly the scheduler picks up a Dag change. A Dag update that creates +a new version is seen immediately, because the new version is a different cache key. A version +rewritten in place is re-checked against its current hash once +:ref:`config:core__min_serialized_dag_update_interval` has elapsed since the entry was last +validated, so that option, not ``dag_cache_ttl``, bounds how long a rewritten version can be +served stale. The same applies to the API server. + +See :ref:`config:scheduler__dag_cache_size` and :ref:`config:scheduler__dag_cache_ttl` for the +full configuration reference. + MySQL and MySQL variant Databases ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/airflow-core/newsfragments/71813.significant.rst b/airflow-core/newsfragments/71813.significant.rst new file mode 100644 index 0000000000000..5d8d7b46a6253 --- /dev/null +++ b/airflow-core/newsfragments/71813.significant.rst @@ -0,0 +1,12 @@ +``dag_cache_ttl`` is now honored when ``dag_cache_size = 0``, which previously disabled eviction and silently ignored the TTL. An ``[api] dag_cache_size = 0`` deployment now evicts on the default ``dag_cache_ttl = 3600``; set ``[api] dag_cache_ttl = 0`` to keep the old behavior. The scheduler gained the same options, ``[scheduler] dag_cache_size`` (default 512) and ``[scheduler] dag_cache_ttl`` (default 0), so it no longer retains every Dag version it has ever seen: the cache is capped at 512 versions and evicts the least recently used one beyond that. Raise ``[scheduler] dag_cache_size`` if more Dag versions have runs in flight than the limit, since a limit below that working set evicts versions still being scheduled and costs a re-fetch on the next loop. Set it to 0 for no size limit, leaving eviction to ``[scheduler] dag_cache_ttl``, which bounds memory by the concurrently active set rather than outright because each re-check resets a cached entry's expiry. + +* Types of change + + * [ ] Dag changes + * [x] Config changes + * [ ] API changes + * [ ] CLI changes + * [x] Behaviour changes + * [ ] Plugin changes + * [ ] Dependency changes + * [ ] Code interface changes diff --git a/airflow-core/src/airflow/api_fastapi/common/dagbag.py b/airflow-core/src/airflow/api_fastapi/common/dagbag.py index d87aca49a524a..000894a644440 100644 --- a/airflow-core/src/airflow/api_fastapi/common/dagbag.py +++ b/airflow-core/src/airflow/api_fastapi/common/dagbag.py @@ -16,43 +16,24 @@ # 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 DBDagBag, dag_cache_conf 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: - """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) - - 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) + """Build the API server's DagBag from the ``[api]`` cache options.""" + cache_size, cache_ttl = dag_cache_conf("api", size_fallback=64, ttl_fallback=3600) + return DBDagBag(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/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 6d2438fd21362..3c91c16391799 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -1714,26 +1714,37 @@ api: default: "False" dag_cache_size: description: | - Size of the LRU cache for SerializedDAG objects in the API server. - Set to 0 to use an unbounded dict (no eviction, matching pre-3.2 behavior). + Max number of deserialized SerializedDAG objects the API server keeps in memory. + Set to 0 for no size limit, leaving eviction to ``dag_cache_ttl``. Set both this and + ``dag_cache_ttl`` to 0 to use an unbounded dict with no eviction, matching the behavior + before 3.2.2. + The cache is keyed by Dag version ID, so lookups by Dag ID (e.g., viewing a Dag's details) always query the database for the latest version, but the deserialized result is cached for subsequent version-specific lookups. - version_added: 3.3.0 + version_added: 3.2.2 type: integer example: ~ default: "64" dag_cache_ttl: description: | - Time-to-live (seconds) for cached SerializedDAG objects in the API server. - After this time, cached DAGs will be re-fetched from the database on next access. - Set to 0 to disable TTL (cache entries will only be evicted by LRU policy). - - Note: After a DAG is updated, the API server may serve the previous version - until the cached entry expires. Lower values reduce staleness but increase - database load. - version_added: 3.3.0 + Seconds a deserialized SerializedDAG stays in the API server's cache, applied whether or + not ``dag_cache_size`` sets a size limit. After this time the Dag is re-fetched from the + database on next access. Set to 0 to disable TTL, leaving eviction to the + ``dag_cache_size`` LRU policy. + + Note that each re-check resets an entry's expiry, so this only reclaims Dag versions + that stop being requested for the whole interval. Versions that keep being requested are + kept, so memory tracks the set of Dag versions in active use and ``dag_cache_size`` + remains the only hard ceiling. + + Note: this does not govern staleness. A Dag update that creates a new version is picked + up immediately, because the new version is a different cache key. A version rewritten in + place is re-checked against its current ``dag_hash`` once + ``[core] min_serialized_dag_update_interval`` has elapsed since the entry was last + validated, so that option bounds how long a rewritten version can be served stale. + version_added: 3.2.2 type: integer example: ~ default: "3600" @@ -2653,6 +2664,42 @@ scheduler: type: integer example: ~ default: "60" + dag_cache_size: + description: | + Max number of deserialized SerializedDAG objects the scheduler keeps in memory, keyed by + Dag version ID. Once full, the least recently used version is evicted, so this is a hard + ceiling on how much the cache can grow no matter how many Dag versions accumulate. + + Raise it if the scheduler has more Dag versions with runs in flight than this; a limit + below the working set evicts versions that are still being scheduled, costing a database + fetch and a deserialization on the next loop. Set to 0 for no size limit, leaving eviction + to ``dag_cache_ttl``. Set both this and ``dag_cache_ttl`` to 0 to use an unbounded dict + with no eviction, matching the behavior before 3.4.0; the cache then grows with the number + of Dag versions the scheduler has ever seen, which on a long-running scheduler is + unbounded. + version_added: 3.4.0 + type: integer + example: ~ + default: "512" + dag_cache_ttl: + description: | + Seconds a deserialized SerializedDAG stays in the scheduler's cache. Defaults to 0, which + disables TTL and leaves eviction to the ``dag_cache_size`` LRU policy. + + Note that each re-check resets an entry's expiry, so a TTL only reclaims Dag versions that + stop being requested for the whole interval. Versions still referenced by active Dag runs + are kept, so ``dag_cache_size`` remains the only hard ceiling and a TTL alone + (``dag_cache_size = 0``) bounds memory by the concurrently active set rather than outright. + + Note: this does not govern staleness. A Dag update that creates a new version is picked + up immediately, because the new version is a different cache key. A version rewritten in + place is re-checked against its current ``dag_hash`` once + ``[core] min_serialized_dag_update_interval`` has elapsed since the entry was last + validated, so that option bounds how long a rewritten version can be served stale. + version_added: 3.4.0 + type: integer + example: ~ + default: "0" pool_metrics_interval: description: | How often (in seconds) should pool usage stats be sent to StatsD (if statsd_on is enabled) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 3f2c7c8a7736d..47bf35b5c9ef4 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 DBDagBag, dag_cache_conf from airflow.models.dagbundle import DagBundleModel from airflow.models.dagrun import DagRun from airflow.models.dagwarning import DagWarning, DagWarningType @@ -298,6 +298,25 @@ def _get_current_dr_task_concurrency(states: Iterable[TaskInstanceState]) -> Sub ) +def _create_scheduler_dag_bag() -> DBDagBag: + """ + Build the scheduler's DagBag from the ``[scheduler]`` cache options. + + Defaults to an LRU cache of 512 versions with no TTL. A size limit is the only hard ceiling: + each re-check resets an entry's expiry, so a TTL reclaims a version only once its runs finish + and it stops being requested, bounding memory by the concurrently active set rather than + outright. 512 is meant to sit above the versions-with-runs-in-flight working set of a typical + deployment, so eviction costs a re-fetch only where that working set is genuinely larger. + """ + cache_size, cache_ttl = dag_cache_conf("scheduler", size_fallback=512, ttl_fallback=0) + return DBDagBag( + load_op_links=False, + cache_size=cache_size, + cache_ttl=cache_ttl, + stats_prefix="scheduler.dag_bag", + ) + + class SchedulerJobRunner(BaseJobRunner, LoggingMixin): """ SchedulerJobRunner runs for a specific time interval and schedules jobs that are ready to run. @@ -370,7 +389,7 @@ def __init__( if log: self._log = log - self.scheduler_dag_bag = DBDagBag(load_op_links=False) + self.scheduler_dag_bag = _create_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 5c0556974ff3b..1a84fa543e63b 100644 --- a/airflow-core/src/airflow/models/dagbag.py +++ b/airflow-core/src/airflow/models/dagbag.py @@ -18,6 +18,8 @@ from __future__ import annotations import hashlib +import logging +import math import time from collections.abc import MutableMapping from contextlib import nullcontext @@ -43,6 +45,31 @@ from airflow.models.serialized_dag import SerializedDagModel from airflow.serialization.definitions.dag import SerializedDAG +log = logging.getLogger(__name__) + + +def dag_cache_conf(section: str, size_fallback: int, ttl_fallback: int) -> tuple[int, int]: + """ + Read and validate a component's ``dag_cache_size`` / ``dag_cache_ttl`` pair. + + Both options treat 0 as "no limit of this kind": no size limit, and no TTL respectively. + Negative values are meaningless, so they are clamped to 0 with a warning rather than + silently selecting the most memory-hungry cache mode. + + :param section: Config section holding the options, e.g. ``api`` or ``scheduler``. + :param size_fallback: ``dag_cache_size`` default if the option is absent. + :param ttl_fallback: ``dag_cache_ttl`` default if the option is absent. + """ + size = conf.getint(section, "dag_cache_size", fallback=size_fallback) + ttl = conf.getint(section, "dag_cache_ttl", fallback=ttl_fallback) + if size < 0: + log.warning("[%s] dag_cache_size must be >= 0, using 0 (no size limit)", section) + size = 0 + if ttl < 0: + log.warning("[%s] dag_cache_ttl must be >= 0, disabling TTL", section) + ttl = 0 + return size, ttl + class _CacheEntry(NamedTuple): """A cached deserialized DAG plus the metadata needed to detect staleness on lookup.""" @@ -62,9 +89,9 @@ 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. + 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. :meta private: """ @@ -74,13 +101,19 @@ def __init__( load_op_links: bool = True, cache_size: int | None = None, cache_ttl: int | None = None, + *, + stats_prefix: str | None = None, ) -> 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). + :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. + :param stats_prefix: Metric namespace for this component's cache, e.g. ``scheduler.dag_bag``. + The ``_stat_*`` hooks append the suffixes, so every component emits the same set of + names. Required when a cache is enabled, unused otherwise. """ self.load_op_links = load_op_links self._dags: MutableMapping[UUID | str, _CacheEntry] = {} @@ -88,19 +121,38 @@ def __init__( self._revalidation_interval = conf.getint("core", "min_serialized_dag_update_interval") - # 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) + size = max(cache_size or 0, 0) + ttl = max(cache_ttl or 0, 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). - # nullcontext for unbounded dict avoids lock overhead in the scheduler path. + # (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() + if self._use_cache and not stats_prefix: + # Caching without a namespace would emit metrics under a partial name. Fail here, at + # wiring time, rather than from the first emission mid-request or mid-scheduling-loop. + raise ValueError("a cached DBDagBag needs stats_prefix to namespace its cache metrics") + self._stats_prefix = stats_prefix + + def _stat_cache_hit(self) -> None: + stats.incr(f"{self._stats_prefix}.cache_hit") + + def _stat_cache_miss(self) -> None: + stats.incr(f"{self._stats_prefix}.cache_miss") + + def _stat_cache_clear(self) -> None: + stats.incr(f"{self._stats_prefix}.cache_clear") + + def _stat_cache_size(self, size: int, *, rate: float = 1.0) -> None: + stats.gauge(f"{self._stats_prefix}.cache_size", size, rate=rate) + def _read_dag(self, serdag: SerializedDagModel) -> SerializedDAG | None: """Read and cache a SerializedDAG (with its ``dag_hash`` for staleness detection).""" serdag.load_op_links = self.load_op_links @@ -109,9 +161,9 @@ 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) + cache_size = len(self._dags) if self._use_cache else 0 if self._use_cache: - stats.gauge("api_server.dag_bag.cache_size", cache_size, rate=0.1) + self._stat_cache_size(cache_size, rate=0.1) return dag @staticmethod @@ -134,7 +186,7 @@ def _get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG | # 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._stat_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 @@ -149,7 +201,7 @@ def _get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG | 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._stat_cache_hit() return cached.dag # Stale (updated in place) or the version no longer exists: drop and reload below. with self._lock: @@ -169,9 +221,9 @@ def _get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG | 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") + self._stat_cache_hit() return cached.dag - stats.incr("api_server.dag_bag.cache_miss") + self._stat_cache_miss() return self._read_dag(serdag) def get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG | None: @@ -203,8 +255,8 @@ def clear_cache(self) -> int: 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._stat_cache_clear() + self._stat_cache_size(0) return count @staticmethod 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..c2924536cbe9e 100644 --- a/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py +++ b/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import math from unittest import mock import pytest @@ -24,6 +25,7 @@ from airflow.api_fastapi.app import purge_cached_app 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 @@ -89,24 +91,22 @@ 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_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", 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("-1", "3600", TTLCache, math.inf, id="negative_size_clamped"), ], ) - @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 - ): + def test_create_dag_bag_cache_modes(self, cache_size, cache_ttl, expected_dags_type, expected_maxsize): from airflow.api_fastapi.common.dagbag import 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) + with conf_vars({("api", "dag_cache_size"): cache_size, ("api", "dag_cache_ttl"): cache_ttl}): + dag_bag = create_dag_bag() - dag_bag = create_dag_bag() - assert dag_bag._use_cache is expected_use_cache 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/models/test_dagbag.py b/airflow-core/tests/unit/models/test_dagbag.py index 79668d4fe54f4..bc8a81c779a2e 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 @@ -24,6 +25,8 @@ import time_machine from cachetools import LRUCache, TTLCache +from airflow.api_fastapi.common.dagbag import create_dag_bag +from airflow.jobs.scheduler_job_runner import _create_scheduler_dag_bag from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion from airflow.models.dagbag import DBDagBag, _CacheEntry @@ -38,6 +41,24 @@ pytestmark = pytest.mark.db_test +STATS_PATH = "airflow.models.dagbag.stats" + +# Each component's own factory, paired with the prefix it is expected to wire up. +METRIC_SOURCES = [ + pytest.param(create_dag_bag, "api_server.dag_bag", id="api_server"), + pytest.param(_create_scheduler_dag_bag, "scheduler.dag_bag", id="scheduler"), +] + +CACHE_METRIC_SUFFIXES = ("cache_hit", "cache_miss", "cache_clear", "cache_size") + +STUB_PREFIX = "test.dag_bag" + + +def _stub_dag_bag(**kwargs) -> DBDagBag: + """Drive the base's caching directly, standing in for a component's own prefix.""" + return DBDagBag(stats_prefix=STUB_PREFIX, **kwargs) + + # 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. @@ -246,33 +267,36 @@ 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 + @pytest.mark.parametrize( + ("cache_size", "cache_ttl", "expected_type", "expected_maxsize"), + [ + pytest.param(None, None, dict, None, id="neither_plain_dict"), + pytest.param(-1, -1, dict, None, id="negatives_clamped_to_plain_dict"), + pytest.param(10, None, 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"), + ], + ) + 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) + assert dag_bag._use_cache is (expected_type is not dict) + if expected_maxsize is not None: + assert dag_bag._dags.maxsize == expected_maxsize + + def test_scheduler_defaults_to_a_bounded_lru_cache(self): + """The scheduler's default has to cap memory outright. + + A TTL alone would not: each re-check resets an entry's expiry, so it bounds the cache by + the concurrently active set rather than by a fixed ceiling. + """ + dag_bag = _create_scheduler_dag_bag() 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) + assert dag_bag._dags.maxsize == 512 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 @@ -299,7 +323,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 +336,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 +349,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 +379,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 +405,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() @@ -394,23 +418,63 @@ def test_iter_all_latest_version_dags_does_not_cache(self): # Cache should be empty -- iter doesn't cache to prevent thrashing assert len(dag_bag._dags) == 0 - @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) + @pytest.mark.parametrize(("create_bag", "prefix"), METRIC_SOURCES) + def test_stats_prefix_expands_to_registered_metrics(self, create_bag, 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 + + assert create_bag()._stats_prefix == prefix + registry = MetricsRegistry() + missing = [ + name for suffix in CACHE_METRIC_SUFFIXES if registry.get(name := f"{prefix}.{suffix}") is None + ] + assert not missing + + @pytest.mark.parametrize( + ("cache_size", "cache_ttl"), + [ + pytest.param(10, 60, id="ttl_cache"), + pytest.param(10, 0, id="lru_cache"), + ], + ) + def test_caching_without_a_stats_prefix_is_rejected_at_construction(self, cache_size, cache_ttl): + """A cache with no namespace to report under must fail at wiring time, not mid-request.""" + with pytest.raises(ValueError, match="needs stats_prefix"): + DBDagBag(cache_size=cache_size, cache_ttl=cache_ttl) + + def test_uncached_bag_emits_no_metrics_without_a_stats_prefix(self): + """Without a cache there is nothing to report, so no namespace is needed.""" + 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_cache_hit_metric_emitted(self): + 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) mock_session.scalar.return_value = "hash1" - dag_bag._get_dag("test_version", mock_session) + with patch(STATS_PATH) as 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) + def test_cache_miss_metric_emitted(self): + 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 @@ -421,29 +485,29 @@ def test_cache_miss_metric_emitted(self, mock_stats): mock_dag_version.serialized_dag = mock_serdag mock_session.get.return_value = mock_dag_version - dag_bag._get_dag("uncached_version", mock_session) + with patch(STATS_PATH) as 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) + def test_cache_clear_metric_emitted(self): + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) dag_bag._dags["test_version"] = MagicMock() - dag_bag.clear_cache() + with patch(STATS_PATH) as mock_stats: + 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") + mock_stats.gauge.assert_called_with(f"{STUB_PREFIX}.cache_size", 0, rate=1.0) - @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) + def test_cache_size_gauge_emitted(self): + 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() mock_serdag.load_op_links = True - dag_bag._read_dag(mock_serdag) + with patch(STATS_PATH) as 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."