From 95c235dd8934f12699658e92d14afafade5bbdcd Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 17 Aug 2026 12:50:38 +0000 Subject: [PATCH 1/6] Bound the scheduler's Dag cache and honor TTL without a size cap The scheduler cached deserialized Dags in a dict that never evicted, so every Dag version it had ever seen stayed resident and the process was eventually OOM killed. ``[api] dag_cache_size = 0`` had a related gap: it selected that same never-evicting dict and silently ignored ``[api] dag_cache_ttl``, so TTL eviction could not be enabled without also accepting a size limit. Each re-check resets a cached entry's expiry, so the TTL reclaims a version once its Dag runs finish and it stops being requested; ``dag_cache_size`` remains the only hard ceiling. closes: #69001 --- .../web-stack.rst | 11 +- airflow-core/docs/faq.rst | 59 +++++- .../newsfragments/71704.significant.rst | 12 ++ .../src/airflow/api_fastapi/common/dagbag.py | 34 ++-- .../src/airflow/config_templates/config.yml | 63 +++++-- .../src/airflow/jobs/scheduler_dagbag.py | 45 +++++ .../src/airflow/jobs/scheduler_job_runner.py | 3 +- airflow-core/src/airflow/models/dagbag.py | 94 ++++++--- .../unit/api_fastapi/common/test_dagbag.py | 36 ++-- .../tests/unit/jobs/test_scheduler_dagbag.py | 58 ++++++ airflow-core/tests/unit/models/test_dagbag.py | 178 ++++++++++++------ dev/airflow_perf/dag_bag_cache_overhead.py | 173 +++++++++++++++++ .../check_metrics_synced_with_the_registry.py | 67 +++++-- ..._check_metrics_synced_with_the_registry.py | 79 ++++++-- .../metrics/metrics_template.yaml | 24 +++ 15 files changed, 771 insertions(+), 165 deletions(-) create mode 100644 airflow-core/newsfragments/71704.significant.rst create mode 100644 airflow-core/src/airflow/jobs/scheduler_dagbag.py create mode 100644 airflow-core/tests/unit/jobs/test_scheduler_dagbag.py create mode 100755 dev/airflow_perf/dag_bag_cache_overhead.py 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..30612a391f4ab 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,47 @@ 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 serialized Dag objects for the life of the process, so as Dag versions +accumulate (see :ref:`faq:dag-version-inflation`) its memory grows until the process is +restarted or OOM killed. Configure eviction in the ``[scheduler]`` section: + +.. code-block:: ini + + [scheduler] + dag_cache_size = 0 ; recommended: no size limit + dag_cache_ttl = 3600 ; seconds before an idle cached entry expires + +Leaving ``dag_cache_size = 0`` is the recommended starting point. 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 rather than a predictable number of entries. A size limit set too low evicts +versions that are still being scheduled, costing a database fetch and a deserialization on the +next loop; the TTL instead reclaims each version once its runs finish and it stops being +requested. + +If the scheduler is still OOM killed, lower ``dag_cache_ttl`` so idle versions are +reclaimed sooner. Only set ``dag_cache_size`` to a non-zero value if you need a hard ceiling, for +example when so many Dag versions have runs in flight at once that the working set alone does not +fit. Setting both to 0 uses an unbounded dict with no eviction, matching the behavior before +3.3.2. + +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/71704.significant.rst b/airflow-core/newsfragments/71704.significant.rst new file mode 100644 index 0000000000000..2ebb1109c9b32 --- /dev/null +++ b/airflow-core/newsfragments/71704.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 0) and ``[scheduler] dag_cache_ttl`` (default 3600), so it no longer retains every Dag version it has ever seen. Each re-check resets a cached entry's expiry, so the TTL reclaims a version once its Dag runs finish and it stops being requested; memory then tracks the versions with runs in flight rather than growing for the life of the process. ``dag_cache_size`` remains the only hard ceiling. + +* 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..04c4738396a00 100644 --- a/airflow-core/src/airflow/api_fastapi/common/dagbag.py +++ b/airflow-core/src/airflow/api_fastapi/common/dagbag.py @@ -16,43 +16,37 @@ # 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__) +class APIServerDBDagBag(DBDagBag): + """ + DagBag for the API server, reporting cache activity under ``api_server.dag_bag``. -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 + :meta private: + """ - # Use unbounded dict (no eviction) if cache_size is 0 - if cache_size <= 0: - return DBDagBag(cache_size=0) + @classmethod + def from_config(cls) -> APIServerDBDagBag: + """Build an instance from the ``[api]`` cache options.""" + cache_size, cache_ttl = dag_cache_conf("api", size_fallback=64, ttl_fallback=3600) + return cls(cache_size=cache_size, cache_ttl=cache_ttl, stats_prefix="api_server.dag_bag") - # 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) +def create_dag_bag() -> DBDagBag: + """Create the API server's DagBag from the ``[api]`` cache options.""" + return APIServerDBDagBag.from_config() 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..2fb04c4a91dce 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,36 @@ 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. 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.3.2. 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.3.2 + type: integer + example: ~ + default: "0" + dag_cache_ttl: + description: | + Seconds a deserialized SerializedDAG stays in the scheduler's cache. 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 still referenced by active Dag + runs are kept, so memory tracks the concurrently active set 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.3.2 + type: integer + example: ~ + default: "3600" 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_dagbag.py b/airflow-core/src/airflow/jobs/scheduler_dagbag.py new file mode 100644 index 0000000000000..562fc1b8f9ce6 --- /dev/null +++ b/airflow-core/src/airflow/jobs/scheduler_dagbag.py @@ -0,0 +1,45 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from airflow.models.dagbag import DBDagBag, dag_cache_conf + + +class SchedulerDBDagBag(DBDagBag): + """ + DagBag for the scheduler, reporting cache activity under ``scheduler.dag_bag``. + + Defaults to no size limit (``[scheduler] dag_cache_size = 0``). 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, which no fixed count predicts well. A size limit would evict versions that are + still being scheduled, whereas each re-check resets an entry's expiry, so the TTL reclaims + each version only once its runs finish and it stops being requested. + + :meta private: + """ + + @classmethod + def from_config(cls) -> SchedulerDBDagBag: + """Build an instance from the ``[scheduler]`` cache options.""" + cache_size, cache_ttl = dag_cache_conf("scheduler", size_fallback=0, ttl_fallback=3600) + return cls( + load_op_links=False, + cache_size=cache_size, + cache_ttl=cache_ttl, + stats_prefix="scheduler.dag_bag", + ) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 3f2c7c8a7736d..32839ad0ee02d 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -71,6 +71,7 @@ from airflow.executors.executor_loader import ExecutorLoader from airflow.jobs.base_job_runner import BaseJobRunner from airflow.jobs.job import Job, JobState, perform_heartbeat +from airflow.jobs.scheduler_dagbag import SchedulerDBDagBag from airflow.models import Deadline, Log from airflow.models.asset import ( AssetActive, @@ -370,7 +371,7 @@ def __init__( if log: self._log = log - self.scheduler_dag_bag = DBDagBag(load_op_links=False) + self.scheduler_dag_bag = SchedulerDBDagBag.from_config() # 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..ddc21e4e54d72 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 @@ -49,13 +51,15 @@ 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.api_fastapi.common.dagbag import APIServerDBDagBag + + real_from_config = APIServerDBDagBag.from_config def factory(*args, **kwargs): self.dagbag_call_counter["count"] += 1 - return RealDagBag(*args, **kwargs) + return real_from_config(*args, **kwargs) - with mock.patch("airflow.api_fastapi.common.dagbag.DBDagBag", side_effect=factory): + with mock.patch.object(APIServerDBDagBag, "from_config", side_effect=factory): purge_cached_app() yield @@ -89,24 +93,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/jobs/test_scheduler_dagbag.py b/airflow-core/tests/unit/jobs/test_scheduler_dagbag.py new file mode 100644 index 0000000000000..ca9080402c7d4 --- /dev/null +++ b/airflow-core/tests/unit/jobs/test_scheduler_dagbag.py @@ -0,0 +1,58 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import math + +import pytest +from cachetools import LRUCache, TTLCache + +from airflow.jobs.scheduler_dagbag import SchedulerDBDagBag + +from tests_common.test_utils.config import conf_vars + +pytestmark = pytest.mark.db_test + + +class TestSchedulerDBDagBag: + @pytest.mark.parametrize( + ("cache_size", "cache_ttl", "expected_dags_type", "expected_maxsize"), + [ + pytest.param(None, None, TTLCache, math.inf, id="defaults_ttl_only"), + pytest.param("512", "3600", TTLCache, 512, id="size_and_ttl"), + pytest.param("512", "0", LRUCache, 512, id="ttl_zero_lru_only"), + pytest.param("0", "0", dict, None, id="both_zero_no_eviction"), + ], + ) + def test_from_config(self, cache_size, cache_ttl, expected_dags_type, expected_maxsize): + overrides = {} + if cache_size is not None: + overrides[("scheduler", "dag_cache_size")] = cache_size + if cache_ttl is not None: + overrides[("scheduler", "dag_cache_ttl")] = cache_ttl + + with conf_vars(overrides): + dag_bag = SchedulerDBDagBag.from_config() + + 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 + + def test_from_config_does_not_load_op_links(self): + """Operator extra links are an API-server concern; the scheduler must not deserialize them.""" + assert SchedulerDBDagBag.from_config().load_op_links is False diff --git a/airflow-core/tests/unit/models/test_dagbag.py b/airflow-core/tests/unit/models/test_dagbag.py index 79668d4fe54f4..5153d5d388213 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 APIServerDBDagBag +from airflow.jobs.scheduler_dagbag import SchedulerDBDagBag from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion from airflow.models.dagbag import DBDagBag, _CacheEntry @@ -38,6 +41,25 @@ pytestmark = pytest.mark.db_test +STATS_PATH = "airflow.models.dagbag.stats" + +# The hooks live on the base and build their names from the ``stats_prefix`` each component passes +# in, so the shared plumbing is exercised once per component. +METRIC_SOURCES = [ + pytest.param(APIServerDBDagBag, "api_server.dag_bag", id="api_server"), + pytest.param(SchedulerDBDagBag, "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 +268,39 @@ def make_lazy(task_ids): class TestDBDagBagCache: """Tests for DBDagBag optional caching behavior.""" - def test_no_caching_by_default(self): - """Test that DBDagBag uses a simple dict without caching by default.""" - dag_bag = DBDagBag() - assert dag_bag._use_cache is False - assert isinstance(dag_bag._dags, dict) - - def test_lru_cache_enabled_with_cache_size(self): - """Test that LRU cache is enabled when cache_size is provided.""" - dag_bag = DBDagBag(cache_size=10) - assert dag_bag._use_cache is True - assert isinstance(dag_bag._dags, LRUCache) - - def test_ttl_cache_enabled_with_cache_size_and_ttl(self): - """Test that TTL cache is enabled when both cache_size and cache_ttl are provided.""" - dag_bag = DBDagBag(cache_size=10, cache_ttl=60) - assert dag_bag._use_cache is True - assert isinstance(dag_bag._dags, TTLCache) - - def test_zero_cache_size_uses_unbounded_dict(self): - """Test that cache_size=0 uses unbounded dict (same as no caching).""" - dag_bag = DBDagBag(cache_size=0, cache_ttl=60) - assert dag_bag._use_cache is False - assert isinstance(dag_bag._dags, dict) + @pytest.mark.parametrize( + ("cache_size", "cache_ttl", "expected_type", "expected_maxsize"), + [ + pytest.param(None, None, dict, None, id="no_args_plain_dict"), + pytest.param(0, 0, dict, None, id="both_zero_plain_dict"), + pytest.param(0, None, dict, None, id="size_zero_no_ttl_plain_dict"), + pytest.param(10, None, LRUCache, 10, id="size_only_lru"), + pytest.param(10, 0, LRUCache, 10, id="ttl_zero_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(None, 60, TTLCache, math.inf, id="ttl_only_size_none_uncapped"), + pytest.param(-1, 60, TTLCache, math.inf, id="negative_size_clamped_to_uncapped"), + pytest.param(10, -1, LRUCache, 10, id="negative_ttl_clamped_to_lru"), + pytest.param(-1, -1, dict, None, id="both_negative_plain_dict"), + ], + ) + 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_uncapped_ttl_cache_accepts_entries(self): + """A size of 0 must mean "no limit", not cachetools' zero-capacity cache.""" + dag_bag = _stub_dag_bag(cache_size=0, cache_ttl=60) + for i in range(200): + dag_bag._dags[f"version_{i}"] = MagicMock() + assert len(dag_bag._dags) == 200 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 +327,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 +340,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 +353,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 +383,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 +409,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 +422,65 @@ 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(("dag_bag_cls", "prefix"), METRIC_SOURCES) + def test_stats_prefix_expands_to_registered_metrics(self, dag_bag_cls, 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 dag_bag_cls.from_config()._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() + + @pytest.mark.parametrize(("dag_bag_cls", "prefix"), METRIC_SOURCES) + def test_cache_hit_metric_emitted(self, dag_bag_cls, prefix): + dag_bag = dag_bag_cls(cache_size=10, cache_ttl=60, stats_prefix=prefix) 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"{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) + @pytest.mark.parametrize(("dag_bag_cls", "prefix"), METRIC_SOURCES) + def test_cache_miss_metric_emitted(self, dag_bag_cls, prefix): + dag_bag = dag_bag_cls(cache_size=10, cache_ttl=60, stats_prefix=prefix) mock_session = MagicMock() # Set up a DB result so _get_dag reaches the miss metric path @@ -421,29 +491,31 @@ 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"{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) + @pytest.mark.parametrize(("dag_bag_cls", "prefix"), METRIC_SOURCES) + def test_cache_clear_metric_emitted(self, dag_bag_cls, prefix): + dag_bag = dag_bag_cls(cache_size=10, cache_ttl=60, stats_prefix=prefix) 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"{prefix}.cache_clear") + mock_stats.gauge.assert_called_with(f"{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) + @pytest.mark.parametrize(("dag_bag_cls", "prefix"), METRIC_SOURCES) + def test_cache_size_gauge_emitted(self, dag_bag_cls, prefix): + dag_bag = dag_bag_cls(cache_size=10, cache_ttl=60, stats_prefix=prefix) 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"{prefix}.cache_size", 1, rate=0.1) diff --git a/dev/airflow_perf/dag_bag_cache_overhead.py b/dev/airflow_perf/dag_bag_cache_overhead.py new file mode 100755 index 0000000000000..33259a5b7402d --- /dev/null +++ b/dev/airflow_perf/dag_bag_cache_overhead.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "cachetools>=6.0.0", +# ] +# /// +""" +Measure the per-lookup overhead ``DBDagBag`` caching adds to the scheduler. + +The scheduler calls ``DBDagBag.get_dag_for_run`` once per Dag run per scheduling +loop. Today its ``_dags`` is a plain dict guarded by ``nullcontext``. Giving the +scheduler a bounded cache swaps that for a cachetools mapping guarded by an +``RLock``, so this isolates the two costs that change: + +* mapping bookkeeping -- ``dict`` vs ``LRUCache`` vs ``TTLCache`` +* lock -- ``nullcontext`` vs an uncontended ``RLock`` + +It replays ``_get_dag``'s cache-hit path (the overwhelmingly common case): a +guarded ``.get()``, a ``time.monotonic()`` freshness comparison, and -- once the +revalidation window has elapsed -- a guarded write-back. No database, no +deserialization; those dominate a real miss and would drown the signal. + +Run it with ``uv run``, which resolves ``cachetools`` from the inline script metadata above:: + + uv run dev/airflow_perf/dag_bag_cache_overhead.py + uv run dev/airflow_perf/dag_bag_cache_overhead.py --versions 5000 --lookups 200000 +""" + +from __future__ import annotations + +import argparse +import statistics +import time +from contextlib import nullcontext +from threading import RLock +from typing import TYPE_CHECKING, Any, NamedTuple + +from cachetools import LRUCache, TTLCache + +if TYPE_CHECKING: + from collections.abc import MutableMapping + + +class _CacheEntry(NamedTuple): + """Mirror of ``airflow.models.dagbag._CacheEntry`` (a real Dag stands in as ``object()``).""" + + dag: Any + dag_hash: str + last_validated: float + + +def _replay_hit_path( + dags: MutableMapping[str, _CacheEntry], + lock: Any, + version_ids: list[str], + lookups: int, + revalidation_interval: int, +) -> None: + """Replay ``_get_dag``'s cache-hit path ``lookups`` times, round-robin over ``version_ids``.""" + n = len(version_ids) + for i in range(lookups): + version_id = version_ids[i % n] + with lock: + cached = dags.get(version_id) + if cached is None: + continue + now = time.monotonic() + if now - cached.last_validated < revalidation_interval: + continue + # Past the revalidation window: _get_dag re-confirms the hash (a DB round-trip we + # deliberately skip) and writes the entry back. On a TTLCache this __setitem__ also + # resets the entry's expiry, which is why hot entries never age out. + with lock: + current = dags.get(version_id) + if current is not None: + dags[version_id] = current._replace(last_validated=now) + + +def _build(kind: str, versions: int) -> MutableMapping[str, _CacheEntry]: + cache: MutableMapping[str, _CacheEntry] + if kind == "dict": + cache = {} + elif kind == "lru": + cache = LRUCache(maxsize=versions) + elif kind == "ttl": + cache = TTLCache[str, _CacheEntry](maxsize=versions, ttl=3600) + elif kind == "ttl-uncapped": + cache = TTLCache[str, _CacheEntry](maxsize=float("inf"), ttl=3600) + else: + raise ValueError(f"unknown mapping kind: {kind}") + return cache + + +def _time_once( + kind: str, locked: bool, versions: int, lookups: int, revalidation_interval: int, stale: bool +) -> float: + dags = _build(kind, versions) + lock = RLock() if locked else nullcontext() + version_ids = [f"version-{i}" for i in range(versions)] + for version_id in version_ids: + dags[version_id] = _CacheEntry(object(), "hash", time.monotonic()) + + # Seeding ``last_validated`` in the past would only make the FIRST visit to each version + # stale, because the write-back refreshes it -- 1 write-back per version, not per lookup. + # A zero-length revalidation window keeps every lookup past the window instead. + effective_interval = 0 if stale else revalidation_interval + start = time.perf_counter() + _replay_hit_path(dags, lock, version_ids, lookups, effective_interval) + return time.perf_counter() - start + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--versions", type=int, default=1000, help="distinct Dag versions held") + parser.add_argument("--lookups", type=int, default=200_000, help="get_dag_for_run calls to replay") + parser.add_argument("--repeat", type=int, default=5, help="timed runs per configuration") + parser.add_argument( + "--revalidation-interval", + type=int, + default=30, + help="[core] min_serialized_dag_update_interval", + ) + args = parser.parse_args() + + configs = [ + ("dict + nullcontext (scheduler today)", "dict", False), + ("dict + RLock", "dict", True), + ("LRUCache + RLock", "lru", True), + ("TTLCache + RLock", "ttl", True), + ("TTLCache uncapped + RLock (proposed)", "ttl-uncapped", True), + ("TTLCache uncapped + nullcontext", "ttl-uncapped", False), + ] + + for stale in (False, True): + branch = "write-back on every lookup" if stale else "all within revalidation window" + print(f"\n{branch} -- {args.lookups:,} lookups over {args.versions:,} versions") + print(f"{'configuration':<40} {'median':>10} {'ns/lookup':>12}") + print("-" * 64) + baseline_ns = None + for label, kind, locked in configs: + timings = [ + _time_once(kind, locked, args.versions, args.lookups, args.revalidation_interval, stale) + for _ in range(args.repeat) + ] + median = statistics.median(timings) + per_lookup_ns = median / args.lookups * 1e9 + if baseline_ns is None: + baseline_ns = per_lookup_ns + delta = "baseline" + else: + delta = f"{per_lookup_ns - baseline_ns:+.0f} ns" + print(f"{label:<40} {median:>9.3f}s {per_lookup_ns:>9.0f} ns {delta}") + + +if __name__ == "__main__": + main() 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." From dfca4564db3f66cb9ef286df7e56ec346d99a8fa Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 18 Aug 2026 13:42:54 +0000 Subject: [PATCH 2/6] Address Ash's comments - Move the scheduler's DagBag factory into `scheduler_job_runner` as `_create_scheduler_dag_bag`, dropping the `airflow/jobs/scheduler_dagbag.py` module and the `SchedulerDBDagBag` subclass it held. - Collapse `APIServerDBDagBag` back into the existing `create_dag_bag`, so both components resolve their config section and metric prefix in a plain factory next to where they build the bag. - Drop `airflow-core/tests/unit/jobs/test_scheduler_dagbag.py`. - Trim `test_cache_selection` from 11 cases to the 5 distinct branches of the mapping selection, and drop `test_uncapped_ttl_cache_accepts_entries`, which only asserted that cachetools honours an unbounded `maxsize`. - Un-parametrize the four cache-metric tests over the two components. The subclasses never overrode the `_stat_*` hooks, so both runs exercised the same base code; `test_stats_prefix_expands_to_registered_metrics` still pins each component's real prefix against the metrics registry. - Drop `dev/airflow_perf/dag_bag_cache_overhead.py`. An in-memory dict lookup is orders of magnitude cheaper than the DB load and deserialize it guards, so the harness has no long-term value in the repo. --- .../src/airflow/api_fastapi/common/dagbag.py | 19 +- .../src/airflow/jobs/scheduler_dagbag.py | 45 ----- .../src/airflow/jobs/scheduler_job_runner.py | 24 ++- .../unit/api_fastapi/common/test_dagbag.py | 8 +- .../tests/unit/jobs/test_scheduler_dagbag.py | 58 ------ airflow-core/tests/unit/models/test_dagbag.py | 64 +++---- dev/airflow_perf/dag_bag_cache_overhead.py | 173 ------------------ 7 files changed, 50 insertions(+), 341 deletions(-) delete mode 100644 airflow-core/src/airflow/jobs/scheduler_dagbag.py delete mode 100644 airflow-core/tests/unit/jobs/test_scheduler_dagbag.py delete mode 100755 dev/airflow_perf/dag_bag_cache_overhead.py diff --git a/airflow-core/src/airflow/api_fastapi/common/dagbag.py b/airflow-core/src/airflow/api_fastapi/common/dagbag.py index 04c4738396a00..000894a644440 100644 --- a/airflow-core/src/airflow/api_fastapi/common/dagbag.py +++ b/airflow-core/src/airflow/api_fastapi/common/dagbag.py @@ -30,23 +30,10 @@ from airflow.serialization.definitions.dag import SerializedDAG -class APIServerDBDagBag(DBDagBag): - """ - DagBag for the API server, reporting cache activity under ``api_server.dag_bag``. - - :meta private: - """ - - @classmethod - def from_config(cls) -> APIServerDBDagBag: - """Build an instance from the ``[api]`` cache options.""" - cache_size, cache_ttl = dag_cache_conf("api", size_fallback=64, ttl_fallback=3600) - return cls(cache_size=cache_size, cache_ttl=cache_ttl, stats_prefix="api_server.dag_bag") - - def create_dag_bag() -> DBDagBag: - """Create the API server's DagBag from the ``[api]`` cache options.""" - return APIServerDBDagBag.from_config() + """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/jobs/scheduler_dagbag.py b/airflow-core/src/airflow/jobs/scheduler_dagbag.py deleted file mode 100644 index 562fc1b8f9ce6..0000000000000 --- a/airflow-core/src/airflow/jobs/scheduler_dagbag.py +++ /dev/null @@ -1,45 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -from __future__ import annotations - -from airflow.models.dagbag import DBDagBag, dag_cache_conf - - -class SchedulerDBDagBag(DBDagBag): - """ - DagBag for the scheduler, reporting cache activity under ``scheduler.dag_bag``. - - Defaults to no size limit (``[scheduler] dag_cache_size = 0``). 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, which no fixed count predicts well. A size limit would evict versions that are - still being scheduled, whereas each re-check resets an entry's expiry, so the TTL reclaims - each version only once its runs finish and it stops being requested. - - :meta private: - """ - - @classmethod - def from_config(cls) -> SchedulerDBDagBag: - """Build an instance from the ``[scheduler]`` cache options.""" - cache_size, cache_ttl = dag_cache_conf("scheduler", size_fallback=0, ttl_fallback=3600) - return cls( - load_op_links=False, - cache_size=cache_size, - cache_ttl=cache_ttl, - stats_prefix="scheduler.dag_bag", - ) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 32839ad0ee02d..64d54202ff024 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -71,7 +71,6 @@ from airflow.executors.executor_loader import ExecutorLoader from airflow.jobs.base_job_runner import BaseJobRunner from airflow.jobs.job import Job, JobState, perform_heartbeat -from airflow.jobs.scheduler_dagbag import SchedulerDBDagBag from airflow.models import Deadline, Log from airflow.models.asset import ( AssetActive, @@ -100,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 @@ -299,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 no size limit (``dag_cache_size = 0``). 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, which no fixed count predicts well. A size limit would evict versions that are still + being scheduled, whereas each re-check resets an entry's expiry, so the TTL reclaims each + version only once its runs finish and it stops being requested. + """ + cache_size, cache_ttl = dag_cache_conf("scheduler", size_fallback=0, ttl_fallback=3600) + 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. @@ -371,7 +389,7 @@ def __init__( if log: self._log = log - self.scheduler_dag_bag = SchedulerDBDagBag.from_config() + 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/tests/unit/api_fastapi/common/test_dagbag.py b/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py index ddc21e4e54d72..c2924536cbe9e 100644 --- a/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py +++ b/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py @@ -51,15 +51,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.api_fastapi.common.dagbag import APIServerDBDagBag - - real_from_config = APIServerDBDagBag.from_config + from airflow.models.dagbag import DBDagBag as RealDagBag def factory(*args, **kwargs): self.dagbag_call_counter["count"] += 1 - return real_from_config(*args, **kwargs) + return RealDagBag(*args, **kwargs) - with mock.patch.object(APIServerDBDagBag, "from_config", side_effect=factory): + with mock.patch("airflow.api_fastapi.common.dagbag.DBDagBag", side_effect=factory): purge_cached_app() yield diff --git a/airflow-core/tests/unit/jobs/test_scheduler_dagbag.py b/airflow-core/tests/unit/jobs/test_scheduler_dagbag.py deleted file mode 100644 index ca9080402c7d4..0000000000000 --- a/airflow-core/tests/unit/jobs/test_scheduler_dagbag.py +++ /dev/null @@ -1,58 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -from __future__ import annotations - -import math - -import pytest -from cachetools import LRUCache, TTLCache - -from airflow.jobs.scheduler_dagbag import SchedulerDBDagBag - -from tests_common.test_utils.config import conf_vars - -pytestmark = pytest.mark.db_test - - -class TestSchedulerDBDagBag: - @pytest.mark.parametrize( - ("cache_size", "cache_ttl", "expected_dags_type", "expected_maxsize"), - [ - pytest.param(None, None, TTLCache, math.inf, id="defaults_ttl_only"), - pytest.param("512", "3600", TTLCache, 512, id="size_and_ttl"), - pytest.param("512", "0", LRUCache, 512, id="ttl_zero_lru_only"), - pytest.param("0", "0", dict, None, id="both_zero_no_eviction"), - ], - ) - def test_from_config(self, cache_size, cache_ttl, expected_dags_type, expected_maxsize): - overrides = {} - if cache_size is not None: - overrides[("scheduler", "dag_cache_size")] = cache_size - if cache_ttl is not None: - overrides[("scheduler", "dag_cache_ttl")] = cache_ttl - - with conf_vars(overrides): - dag_bag = SchedulerDBDagBag.from_config() - - 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 - - def test_from_config_does_not_load_op_links(self): - """Operator extra links are an API-server concern; the scheduler must not deserialize them.""" - assert SchedulerDBDagBag.from_config().load_op_links is False diff --git a/airflow-core/tests/unit/models/test_dagbag.py b/airflow-core/tests/unit/models/test_dagbag.py index 5153d5d388213..2918e8fb7cdc3 100644 --- a/airflow-core/tests/unit/models/test_dagbag.py +++ b/airflow-core/tests/unit/models/test_dagbag.py @@ -25,8 +25,8 @@ import time_machine from cachetools import LRUCache, TTLCache -from airflow.api_fastapi.common.dagbag import APIServerDBDagBag -from airflow.jobs.scheduler_dagbag import SchedulerDBDagBag +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 @@ -43,11 +43,10 @@ STATS_PATH = "airflow.models.dagbag.stats" -# The hooks live on the base and build their names from the ``stats_prefix`` each component passes -# in, so the shared plumbing is exercised once per component. +# Each component's own factory, paired with the prefix it is expected to wire up. METRIC_SOURCES = [ - pytest.param(APIServerDBDagBag, "api_server.dag_bag", id="api_server"), - pytest.param(SchedulerDBDagBag, "scheduler.dag_bag", id="scheduler"), + 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") @@ -271,17 +270,11 @@ class TestDBDagBagCache: @pytest.mark.parametrize( ("cache_size", "cache_ttl", "expected_type", "expected_maxsize"), [ - pytest.param(None, None, dict, None, id="no_args_plain_dict"), - pytest.param(0, 0, dict, None, id="both_zero_plain_dict"), - pytest.param(0, None, dict, None, id="size_zero_no_ttl_plain_dict"), + 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, 0, LRUCache, 10, id="ttl_zero_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(None, 60, TTLCache, math.inf, id="ttl_only_size_none_uncapped"), - pytest.param(-1, 60, TTLCache, math.inf, id="negative_size_clamped_to_uncapped"), - pytest.param(10, -1, LRUCache, 10, id="negative_ttl_clamped_to_lru"), - pytest.param(-1, -1, dict, None, id="both_negative_plain_dict"), ], ) def test_cache_selection(self, cache_size, cache_ttl, expected_type, expected_maxsize): @@ -291,13 +284,6 @@ def test_cache_selection(self, cache_size, cache_ttl, expected_type, expected_ma if expected_maxsize is not None: assert dag_bag._dags.maxsize == expected_maxsize - def test_uncapped_ttl_cache_accepts_entries(self): - """A size of 0 must mean "no limit", not cachetools' zero-capacity cache.""" - dag_bag = _stub_dag_bag(cache_size=0, cache_ttl=60) - for i in range(200): - dag_bag._dags[f"version_{i}"] = MagicMock() - assert len(dag_bag._dags) == 200 - def test_clear_cache_with_caching(self): """Test clear_cache() with caching enabled.""" dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) @@ -422,8 +408,8 @@ 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 - @pytest.mark.parametrize(("dag_bag_cls", "prefix"), METRIC_SOURCES) - def test_stats_prefix_expands_to_registered_metrics(self, dag_bag_cls, prefix): + @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 @@ -432,7 +418,7 @@ def test_stats_prefix_expands_to_registered_metrics(self, dag_bag_cls, prefix): """ from airflow._shared.observability.metrics.metrics_registry import MetricsRegistry - assert dag_bag_cls.from_config()._stats_prefix == prefix + 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 @@ -465,9 +451,8 @@ def test_uncached_bag_emits_no_metrics_without_a_stats_prefix(self): mock_stats.incr.assert_not_called() mock_stats.gauge.assert_not_called() - @pytest.mark.parametrize(("dag_bag_cls", "prefix"), METRIC_SOURCES) - def test_cache_hit_metric_emitted(self, dag_bag_cls, prefix): - dag_bag = dag_bag_cls(cache_size=10, cache_ttl=60, stats_prefix=prefix) + 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) @@ -476,11 +461,10 @@ def test_cache_hit_metric_emitted(self, dag_bag_cls, prefix): with patch(STATS_PATH) as mock_stats: dag_bag._get_dag("test_version", mock_session) - mock_stats.incr.assert_called_with(f"{prefix}.cache_hit") + mock_stats.incr.assert_called_with(f"{STUB_PREFIX}.cache_hit") - @pytest.mark.parametrize(("dag_bag_cls", "prefix"), METRIC_SOURCES) - def test_cache_miss_metric_emitted(self, dag_bag_cls, prefix): - dag_bag = dag_bag_cls(cache_size=10, cache_ttl=60, stats_prefix=prefix) + 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 @@ -494,22 +478,20 @@ def test_cache_miss_metric_emitted(self, dag_bag_cls, prefix): with patch(STATS_PATH) as mock_stats: dag_bag._get_dag("uncached_version", mock_session) - mock_stats.incr.assert_any_call(f"{prefix}.cache_miss") + mock_stats.incr.assert_any_call(f"{STUB_PREFIX}.cache_miss") - @pytest.mark.parametrize(("dag_bag_cls", "prefix"), METRIC_SOURCES) - def test_cache_clear_metric_emitted(self, dag_bag_cls, prefix): - dag_bag = dag_bag_cls(cache_size=10, cache_ttl=60, stats_prefix=prefix) + def test_cache_clear_metric_emitted(self): + dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) dag_bag._dags["test_version"] = MagicMock() with patch(STATS_PATH) as mock_stats: dag_bag.clear_cache() - mock_stats.incr.assert_called_with(f"{prefix}.cache_clear") - mock_stats.gauge.assert_called_with(f"{prefix}.cache_size", 0, rate=1.0) + 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) - @pytest.mark.parametrize(("dag_bag_cls", "prefix"), METRIC_SOURCES) - def test_cache_size_gauge_emitted(self, dag_bag_cls, prefix): - dag_bag = dag_bag_cls(cache_size=10, cache_ttl=60, stats_prefix=prefix) + 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() @@ -518,4 +500,4 @@ def test_cache_size_gauge_emitted(self, dag_bag_cls, prefix): with patch(STATS_PATH) as mock_stats: dag_bag._read_dag(mock_serdag) - mock_stats.gauge.assert_called_with(f"{prefix}.cache_size", 1, rate=0.1) + mock_stats.gauge.assert_called_with(f"{STUB_PREFIX}.cache_size", 1, rate=0.1) diff --git a/dev/airflow_perf/dag_bag_cache_overhead.py b/dev/airflow_perf/dag_bag_cache_overhead.py deleted file mode 100755 index 33259a5b7402d..0000000000000 --- a/dev/airflow_perf/dag_bag_cache_overhead.py +++ /dev/null @@ -1,173 +0,0 @@ -#!/usr/bin/env python -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# /// script -# requires-python = ">=3.10" -# dependencies = [ -# "cachetools>=6.0.0", -# ] -# /// -""" -Measure the per-lookup overhead ``DBDagBag`` caching adds to the scheduler. - -The scheduler calls ``DBDagBag.get_dag_for_run`` once per Dag run per scheduling -loop. Today its ``_dags`` is a plain dict guarded by ``nullcontext``. Giving the -scheduler a bounded cache swaps that for a cachetools mapping guarded by an -``RLock``, so this isolates the two costs that change: - -* mapping bookkeeping -- ``dict`` vs ``LRUCache`` vs ``TTLCache`` -* lock -- ``nullcontext`` vs an uncontended ``RLock`` - -It replays ``_get_dag``'s cache-hit path (the overwhelmingly common case): a -guarded ``.get()``, a ``time.monotonic()`` freshness comparison, and -- once the -revalidation window has elapsed -- a guarded write-back. No database, no -deserialization; those dominate a real miss and would drown the signal. - -Run it with ``uv run``, which resolves ``cachetools`` from the inline script metadata above:: - - uv run dev/airflow_perf/dag_bag_cache_overhead.py - uv run dev/airflow_perf/dag_bag_cache_overhead.py --versions 5000 --lookups 200000 -""" - -from __future__ import annotations - -import argparse -import statistics -import time -from contextlib import nullcontext -from threading import RLock -from typing import TYPE_CHECKING, Any, NamedTuple - -from cachetools import LRUCache, TTLCache - -if TYPE_CHECKING: - from collections.abc import MutableMapping - - -class _CacheEntry(NamedTuple): - """Mirror of ``airflow.models.dagbag._CacheEntry`` (a real Dag stands in as ``object()``).""" - - dag: Any - dag_hash: str - last_validated: float - - -def _replay_hit_path( - dags: MutableMapping[str, _CacheEntry], - lock: Any, - version_ids: list[str], - lookups: int, - revalidation_interval: int, -) -> None: - """Replay ``_get_dag``'s cache-hit path ``lookups`` times, round-robin over ``version_ids``.""" - n = len(version_ids) - for i in range(lookups): - version_id = version_ids[i % n] - with lock: - cached = dags.get(version_id) - if cached is None: - continue - now = time.monotonic() - if now - cached.last_validated < revalidation_interval: - continue - # Past the revalidation window: _get_dag re-confirms the hash (a DB round-trip we - # deliberately skip) and writes the entry back. On a TTLCache this __setitem__ also - # resets the entry's expiry, which is why hot entries never age out. - with lock: - current = dags.get(version_id) - if current is not None: - dags[version_id] = current._replace(last_validated=now) - - -def _build(kind: str, versions: int) -> MutableMapping[str, _CacheEntry]: - cache: MutableMapping[str, _CacheEntry] - if kind == "dict": - cache = {} - elif kind == "lru": - cache = LRUCache(maxsize=versions) - elif kind == "ttl": - cache = TTLCache[str, _CacheEntry](maxsize=versions, ttl=3600) - elif kind == "ttl-uncapped": - cache = TTLCache[str, _CacheEntry](maxsize=float("inf"), ttl=3600) - else: - raise ValueError(f"unknown mapping kind: {kind}") - return cache - - -def _time_once( - kind: str, locked: bool, versions: int, lookups: int, revalidation_interval: int, stale: bool -) -> float: - dags = _build(kind, versions) - lock = RLock() if locked else nullcontext() - version_ids = [f"version-{i}" for i in range(versions)] - for version_id in version_ids: - dags[version_id] = _CacheEntry(object(), "hash", time.monotonic()) - - # Seeding ``last_validated`` in the past would only make the FIRST visit to each version - # stale, because the write-back refreshes it -- 1 write-back per version, not per lookup. - # A zero-length revalidation window keeps every lookup past the window instead. - effective_interval = 0 if stale else revalidation_interval - start = time.perf_counter() - _replay_hit_path(dags, lock, version_ids, lookups, effective_interval) - return time.perf_counter() - start - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--versions", type=int, default=1000, help="distinct Dag versions held") - parser.add_argument("--lookups", type=int, default=200_000, help="get_dag_for_run calls to replay") - parser.add_argument("--repeat", type=int, default=5, help="timed runs per configuration") - parser.add_argument( - "--revalidation-interval", - type=int, - default=30, - help="[core] min_serialized_dag_update_interval", - ) - args = parser.parse_args() - - configs = [ - ("dict + nullcontext (scheduler today)", "dict", False), - ("dict + RLock", "dict", True), - ("LRUCache + RLock", "lru", True), - ("TTLCache + RLock", "ttl", True), - ("TTLCache uncapped + RLock (proposed)", "ttl-uncapped", True), - ("TTLCache uncapped + nullcontext", "ttl-uncapped", False), - ] - - for stale in (False, True): - branch = "write-back on every lookup" if stale else "all within revalidation window" - print(f"\n{branch} -- {args.lookups:,} lookups over {args.versions:,} versions") - print(f"{'configuration':<40} {'median':>10} {'ns/lookup':>12}") - print("-" * 64) - baseline_ns = None - for label, kind, locked in configs: - timings = [ - _time_once(kind, locked, args.versions, args.lookups, args.revalidation_interval, stale) - for _ in range(args.repeat) - ] - median = statistics.median(timings) - per_lookup_ns = median / args.lookups * 1e9 - if baseline_ns is None: - baseline_ns = per_lookup_ns - delta = "baseline" - else: - delta = f"{per_lookup_ns - baseline_ns:+.0f} ns" - print(f"{label:<40} {median:>9.3f}s {per_lookup_ns:>9.0f} ns {delta}") - - -if __name__ == "__main__": - main() From ef09e82d74d8bbf95861fc456d44d9a24c5fd6d3 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 18 Aug 2026 13:52:40 +0000 Subject: [PATCH 3/6] Default the scheduler Dag cache to a bounded LRU of 1024 versions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A TTL cannot cap the cache on its own. Each re-check re-arms an entry's expiry, so a TTL reclaims a Dag version only once its runs finish and it stops being requested — that bounds memory by the concurrently active set, which no fixed number predicts, rather than outright. A default that leaves memory dependent on request patterns is the wrong default for the OOM this PR set out to fix. A size limit is the only hard ceiling, so the scheduler now defaults to `dag_cache_size = 1024` with `dag_cache_ttl = 0`. 1024 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; the `scheduler.dag_bag.cache_miss` metric is what tells an operator to raise it. The TTL-only and no-eviction modes both remain reachable by configuration. --- airflow-core/docs/faq.rst | 37 ++++++++++--------- .../newsfragments/71704.significant.rst | 2 +- .../src/airflow/config_templates/config.yml | 30 +++++++++------ .../src/airflow/jobs/scheduler_job_runner.py | 12 +++--- airflow-core/tests/unit/models/test_dagbag.py | 10 +++++ 5 files changed, 54 insertions(+), 37 deletions(-) diff --git a/airflow-core/docs/faq.rst b/airflow-core/docs/faq.rst index 30612a391f4ab..79ece11b23c70 100644 --- a/airflow-core/docs/faq.rst +++ b/airflow-core/docs/faq.rst @@ -773,28 +773,29 @@ See :ref:`config:api__server_type`, :ref:`config:api__worker_refresh_interval`, How to prevent scheduler memory growth? ---------------------------------------- -The scheduler caches serialized Dag objects for the life of the process, so as Dag versions -accumulate (see :ref:`faq:dag-version-inflation`) its memory grows until the process is -restarted or OOM killed. Configure eviction in the ``[scheduler]`` section: +The scheduler caches deserialized Dag objects, so before 3.3.2 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 = 0 ; recommended: no size limit - dag_cache_ttl = 3600 ; seconds before an idle cached entry expires - -Leaving ``dag_cache_size = 0`` is the recommended starting point. 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 rather than a predictable number of entries. A size limit set too low evicts -versions that are still being scheduled, costing a database fetch and a deserialization on the -next loop; the TTL instead reclaims each version once its runs finish and it stops being -requested. - -If the scheduler is still OOM killed, lower ``dag_cache_ttl`` so idle versions are -reclaimed sooner. Only set ``dag_cache_size`` to a non-zero value if you need a hard ceiling, for -example when so many Dag versions have runs in flight at once that the working set alone does not -fit. Setting both to 0 uses an unbounded dict with no eviction, matching the behavior before -3.3.2. + dag_cache_size = 1024 ; 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 1024 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; 1024 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.3.2. 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 diff --git a/airflow-core/newsfragments/71704.significant.rst b/airflow-core/newsfragments/71704.significant.rst index 2ebb1109c9b32..9ef3c2d947488 100644 --- a/airflow-core/newsfragments/71704.significant.rst +++ b/airflow-core/newsfragments/71704.significant.rst @@ -1,4 +1,4 @@ -``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 0) and ``[scheduler] dag_cache_ttl`` (default 3600), so it no longer retains every Dag version it has ever seen. Each re-check resets a cached entry's expiry, so the TTL reclaims a version once its Dag runs finish and it stops being requested; memory then tracks the versions with runs in flight rather than growing for the life of the process. ``dag_cache_size`` remains the only hard ceiling. +``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 1024) and ``[scheduler] dag_cache_ttl`` (default 0), so it no longer retains every Dag version it has ever seen: the cache is capped at 1024 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 diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 2fb04c4a91dce..43c8e31d24971 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2667,23 +2667,29 @@ scheduler: dag_cache_size: description: | Max number of deserialized SerializedDAG objects the scheduler keeps in memory, keyed by - Dag version ID. 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.3.2. The cache then grows with the number of Dag versions - the scheduler has ever seen, which on a long-running scheduler is unbounded. + 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.3.2; 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.3.2 type: integer example: ~ - default: "0" + default: "1024" dag_cache_ttl: description: | - Seconds a deserialized SerializedDAG stays in the scheduler's cache. Set to 0 to disable - TTL, leaving eviction to the ``dag_cache_size`` LRU policy. + 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 this only reclaims Dag versions - that stop being requested for the whole interval. Versions still referenced by active Dag - runs are kept, so memory tracks the concurrently active set and ``dag_cache_size`` remains - the only hard ceiling. + 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 @@ -2693,7 +2699,7 @@ scheduler: version_added: 3.3.2 type: integer example: ~ - default: "3600" + 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 64d54202ff024..6a24e35b86124 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -302,13 +302,13 @@ def _create_scheduler_dag_bag() -> DBDagBag: """ Build the scheduler's DagBag from the ``[scheduler]`` cache options. - Defaults to no size limit (``dag_cache_size = 0``). 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, which no fixed count predicts well. A size limit would evict versions that are still - being scheduled, whereas each re-check resets an entry's expiry, so the TTL reclaims each - version only once its runs finish and it stops being requested. + Defaults to an LRU cache of 1024 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. 1024 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=0, ttl_fallback=3600) + cache_size, cache_ttl = dag_cache_conf("scheduler", size_fallback=1024, ttl_fallback=0) return DBDagBag( load_op_links=False, cache_size=cache_size, diff --git a/airflow-core/tests/unit/models/test_dagbag.py b/airflow-core/tests/unit/models/test_dagbag.py index 2918e8fb7cdc3..45b22d5df4ab5 100644 --- a/airflow-core/tests/unit/models/test_dagbag.py +++ b/airflow-core/tests/unit/models/test_dagbag.py @@ -284,6 +284,16 @@ def test_cache_selection(self, cache_size, cache_ttl, expected_type, expected_ma 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) + assert dag_bag._dags.maxsize == 1024 + def test_clear_cache_with_caching(self): """Test clear_cache() with caching enabled.""" dag_bag = _stub_dag_bag(cache_size=10, cache_ttl=60) From 80a08bb805183bceac2d0ca5763c93ac494b0dff Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 18 Aug 2026 16:08:47 +0000 Subject: [PATCH 4/6] Mark the scheduler Dag cache options as added in 3.4.0 These options are new on main, which is 3.4.0. Claiming 3.3.2 would advertise them as available in a patch release that never carried them, sending anyone on 3.3.x looking for settings they cannot configure. The surrounding docs described the pre-cache behaviour as ending at the same wrong version, so they move together. The `[api]` pair keeps 3.2.2: those options already shipped, and that value is the correction this PR makes to their previously mis-stated 3.3.0. --- airflow-core/docs/faq.rst | 4 ++-- airflow-core/src/airflow/config_templates/config.yml | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/airflow-core/docs/faq.rst b/airflow-core/docs/faq.rst index 79ece11b23c70..2b4ce6a206497 100644 --- a/airflow-core/docs/faq.rst +++ b/airflow-core/docs/faq.rst @@ -773,7 +773,7 @@ See :ref:`config:api__server_type`, :ref:`config:api__worker_refresh_interval`, How to prevent scheduler memory growth? ---------------------------------------- -The scheduler caches deserialized Dag objects, so before 3.3.2 its memory grew with every Dag +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: @@ -795,7 +795,7 @@ that are still being scheduled and costs a database fetch and a deserialization 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.3.2. +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 diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 43c8e31d24971..122d76849fd07 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2674,10 +2674,10 @@ scheduler: 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.3.2; the cache then grows with the number + 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.3.2 + version_added: 3.4.0 type: integer example: ~ default: "1024" @@ -2696,7 +2696,7 @@ scheduler: 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.3.2 + version_added: 3.4.0 type: integer example: ~ default: "0" From f5ec4222dac8d76bbb691b69e211f81dc9da2bce Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 19 Aug 2026 02:17:02 +0000 Subject: [PATCH 5/6] Lower the scheduler Dag cache default from 1024 to 512 versions Ash flagged 1024 as possibly too aggressive a default. 512 still sits above the versions-with-runs-in-flight working set of a typical deployment while roughly halving the worst-case memory footprint of the default cache. --- airflow-core/docs/faq.rst | 6 +++--- airflow-core/newsfragments/71704.significant.rst | 2 +- airflow-core/src/airflow/config_templates/config.yml | 2 +- airflow-core/src/airflow/jobs/scheduler_job_runner.py | 6 +++--- airflow-core/tests/unit/models/test_dagbag.py | 2 +- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/airflow-core/docs/faq.rst b/airflow-core/docs/faq.rst index 2b4ce6a206497..5caccc5b68a3c 100644 --- a/airflow-core/docs/faq.rst +++ b/airflow-core/docs/faq.rst @@ -780,13 +780,13 @@ OOM killed. The cache is now bounded by default. Tune it in the ``[scheduler]`` .. code-block:: ini [scheduler] - dag_cache_size = 1024 ; max cached versions, evicting least recently used + 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 1024 Dag versions and evict the least recently used one beyond +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; 1024 is meant to sit above that for a typical deployment. +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 diff --git a/airflow-core/newsfragments/71704.significant.rst b/airflow-core/newsfragments/71704.significant.rst index 9ef3c2d947488..5d8d7b46a6253 100644 --- a/airflow-core/newsfragments/71704.significant.rst +++ b/airflow-core/newsfragments/71704.significant.rst @@ -1,4 +1,4 @@ -``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 1024) and ``[scheduler] dag_cache_ttl`` (default 0), so it no longer retains every Dag version it has ever seen: the cache is capped at 1024 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. +``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 diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 122d76849fd07..3c91c16391799 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2680,7 +2680,7 @@ scheduler: version_added: 3.4.0 type: integer example: ~ - default: "1024" + default: "512" dag_cache_ttl: description: | Seconds a deserialized SerializedDAG stays in the scheduler's cache. Defaults to 0, which diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 6a24e35b86124..47bf35b5c9ef4 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -302,13 +302,13 @@ def _create_scheduler_dag_bag() -> DBDagBag: """ Build the scheduler's DagBag from the ``[scheduler]`` cache options. - Defaults to an LRU cache of 1024 versions with no TTL. A size limit is the only hard ceiling: + 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. 1024 is meant to sit above the versions-with-runs-in-flight working set of a typical + 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=1024, ttl_fallback=0) + cache_size, cache_ttl = dag_cache_conf("scheduler", size_fallback=512, ttl_fallback=0) return DBDagBag( load_op_links=False, cache_size=cache_size, diff --git a/airflow-core/tests/unit/models/test_dagbag.py b/airflow-core/tests/unit/models/test_dagbag.py index 45b22d5df4ab5..bc8a81c779a2e 100644 --- a/airflow-core/tests/unit/models/test_dagbag.py +++ b/airflow-core/tests/unit/models/test_dagbag.py @@ -292,7 +292,7 @@ def test_scheduler_defaults_to_a_bounded_lru_cache(self): """ dag_bag = _create_scheduler_dag_bag() assert isinstance(dag_bag._dags, LRUCache) - assert dag_bag._dags.maxsize == 1024 + assert dag_bag._dags.maxsize == 512 def test_clear_cache_with_caching(self): """Test clear_cache() with caching enabled.""" From 451e63415082037dfd60e7b4c77610a424543b4d Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 19 Aug 2026 05:25:58 +0000 Subject: [PATCH 6/6] Renumber the newsfragment to this PR The scheduler cache bound it originally accompanied now ships separately in #71704; what remains here is the configuration, metric namespacing, and the `[api]` TTL fix, so the entry belongs to this PR's number. --- .../{71704.significant.rst => 71813.significant.rst} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename airflow-core/newsfragments/{71704.significant.rst => 71813.significant.rst} (100%) diff --git a/airflow-core/newsfragments/71704.significant.rst b/airflow-core/newsfragments/71813.significant.rst similarity index 100% rename from airflow-core/newsfragments/71704.significant.rst rename to airflow-core/newsfragments/71813.significant.rst