From 672d70b4091f41eae261aa4e54d4766bc23ca1e1 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Thu, 29 Jan 2026 15:33:30 -0800 Subject: [PATCH] Add EventDrivenCacheCoordinator and use it for global spend to prevent cache stampede --- litellm/proxy/auth/user_api_key_auth.py | 72 ++++--- .../proxy/common_utils/cache_coordinator.py | 191 ++++++++++++++++++ 2 files changed, 231 insertions(+), 32 deletions(-) create mode 100644 litellm/proxy/common_utils/cache_coordinator.py diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 7e7c7c8c90c..40ad2488ed2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -53,6 +53,7 @@ from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -211,6 +212,33 @@ def update_valid_token_with_end_user_params( return valid_token +# Reusable coordinator for global spend to prevent cache stampede +_global_spend_coordinator = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEND]") + + +async def _fetch_global_spend_with_event_coordination( + cache_key: str, + user_api_key_cache: DualCache, + prisma_client: PrismaClient, +) -> Optional[float]: + """ + Fetch global spend with event-driven coordination to prevent cache stampede. + Uses EventDrivenCacheCoordinator: first request queries DB and signals others when done. + """ + + async def _load_global_spend() -> Optional[float]: + sql_query = """SELECT SUM(spend) AS total_spend FROM "MonthlyGlobalSpend";""" + response = await prisma_client.db.query_raw(query=sql_query) + val = response[0]["total_spend"] + return float(val) if val is not None else None + + return await _global_spend_coordinator.get_or_load( + cache_key=cache_key, + cache=user_api_key_cache, + load_fn=_load_global_spend, + ) + + async def get_global_proxy_spend( litellm_proxy_admin_name: str, user_api_key_cache: DualCache, @@ -219,25 +247,14 @@ async def get_global_proxy_spend( proxy_logging_obj: ProxyLogging, ) -> Optional[float]: global_proxy_spend = None - if litellm.max_budget > 0: # user set proxy max budget - # check cache - global_proxy_spend = await user_api_key_cache.async_get_cache( - key="{}:spend".format(litellm_proxy_admin_name) + if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget + # Use event-driven coordination to prevent cache stampede + cache_key = "{}:spend".format(litellm_proxy_admin_name) + global_proxy_spend = await _fetch_global_spend_with_event_coordination( + cache_key=cache_key, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, ) - if global_proxy_spend is None and prisma_client is not None: - # get from db - sql_query = ( - """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";""" - ) - - response = await prisma_client.db.query_raw(query=sql_query) - - global_proxy_spend = response[0]["total_spend"] - - await user_api_key_cache.async_set_cache( - key="{}:spend".format(litellm_proxy_admin_name), - value=global_proxy_spend, - ) if global_proxy_spend is not None: user_info = CallInfo( user_id=litellm_proxy_admin_name, @@ -1120,21 +1137,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if ( litellm.max_budget > 0 and prisma_client is not None ): # user set proxy max budget - # check cache - global_proxy_spend = await user_api_key_cache.async_get_cache( - key="{}:spend".format(litellm_proxy_admin_name) + cache_key = "{}:spend".format(litellm_proxy_admin_name) + global_proxy_spend = await _fetch_global_spend_with_event_coordination( + cache_key=cache_key, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, ) - if global_proxy_spend is None: - # get from db - sql_query = """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";""" - - response = await prisma_client.db.query_raw(query=sql_query) - - global_proxy_spend = response[0]["total_spend"] - await user_api_key_cache.async_set_cache( - key="{}:spend".format(litellm_proxy_admin_name), - value=global_proxy_spend, - ) if global_proxy_spend is not None: call_info = CallInfo( diff --git a/litellm/proxy/common_utils/cache_coordinator.py b/litellm/proxy/common_utils/cache_coordinator.py new file mode 100644 index 00000000000..60d7e3947a6 --- /dev/null +++ b/litellm/proxy/common_utils/cache_coordinator.py @@ -0,0 +1,191 @@ +""" +Event-driven cache coordinator to prevent cache stampede. + +Use this when many requests can miss the same cache key at once (e.g. after +expiry or restart). Without coordination, they would all run the expensive +load (DB query, API call) in parallel and overload the backend. + +This module ensures only one request performs the load; the rest wait for a +signal and then read the freshly cached value. Reuse it for any cache-aside +pattern: global spend, feature flags, config, or other shared read-through data. +""" + +import asyncio +import time +from typing import Any, Awaitable, Callable, Optional, Protocol, TypeVar + +from litellm._logging import verbose_proxy_logger + +T = TypeVar("T") + + +class AsyncCacheProtocol(Protocol): + """Protocol for cache backends used by EventDrivenCacheCoordinator.""" + + async def async_get_cache(self, key: str, **kwargs: Any) -> Any: + ... + + async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> Any: + ... + + +class EventDrivenCacheCoordinator: + """ + Coordinates a single in-flight load per logical resource to prevent cache stampede. + + Pattern: + - First request: loads data (e.g. DB query), caches it, then signals waiters. + - Other requests: wait for the signal, then read from cache. + + Create one instance per resource (e.g. one for global spend, one for feature flags). + """ + + def __init__(self, log_prefix: str = "[CACHE]"): + self._lock = asyncio.Lock() + self._event: Optional[asyncio.Event] = None + self._query_in_progress = False + self._log_prefix = log_prefix + + async def _get_cached( + self, cache_key: str, cache: AsyncCacheProtocol + ) -> Optional[T]: + """Return value from cache if present, else None.""" + return await cache.async_get_cache(key=cache_key) + + def _log_cache_hit(self, value: T) -> None: + if self._log_prefix: + verbose_proxy_logger.debug( + "%s Cache hit, value: %s", self._log_prefix, value + ) + + def _log_cache_miss(self) -> None: + if self._log_prefix: + verbose_proxy_logger.debug("%s Cache miss", self._log_prefix) + + async def _claim_role(self) -> Optional[asyncio.Event]: + """ + Under lock: return event to wait on if load is in progress, else set us as loader and return None. + """ + async with self._lock: + if self._query_in_progress and self._event is not None: + if self._log_prefix: + verbose_proxy_logger.debug( + "%s Load in flight, waiting for signal", self._log_prefix + ) + return self._event + self._query_in_progress = True + self._event = asyncio.Event() + if self._log_prefix: + verbose_proxy_logger.debug( + "%s Starting load (will signal others when done)", + self._log_prefix, + ) + return None + + async def _wait_for_signal_and_get( + self, + event: asyncio.Event, + cache_key: str, + cache: AsyncCacheProtocol, + ) -> Optional[T]: + """Wait for loader to finish, then read from cache.""" + await event.wait() + if self._log_prefix: + verbose_proxy_logger.debug( + "%s Signal received, reading from cache", self._log_prefix + ) + value = await cache.async_get_cache(key=cache_key) + if value is not None and self._log_prefix: + verbose_proxy_logger.debug( + "%s Cache filled by other request, value: %s", + self._log_prefix, + value, + ) + elif value is None and self._log_prefix: + verbose_proxy_logger.debug( + "%s Signal received but cache still empty", self._log_prefix + ) + return value + + async def _load_and_cache( + self, + cache_key: str, + cache: AsyncCacheProtocol, + load_fn: Callable[[], Awaitable[T]], + ) -> Optional[T]: + """Double-check cache, run load_fn, set cache, return value. Caller must call _signal_done in finally.""" + value = await cache.async_get_cache(key=cache_key) + if value is not None: + if self._log_prefix: + verbose_proxy_logger.debug( + "%s Cache filled while acquiring lock, value: %s", + self._log_prefix, + value, + ) + return value + + if self._log_prefix: + verbose_proxy_logger.debug("%s Running load", self._log_prefix) + start = time.perf_counter() + value = await load_fn() + elapsed_ms = (time.perf_counter() - start) * 1000 + if self._log_prefix: + verbose_proxy_logger.debug( + "%s Load completed in %.2fms, result: %s", + self._log_prefix, + elapsed_ms, + value, + ) + + await cache.async_set_cache(key=cache_key, value=value) + if self._log_prefix: + verbose_proxy_logger.debug("%s Result cached", self._log_prefix) + return value + + async def _signal_done(self) -> None: + """Reset loader state and signal all waiters.""" + async with self._lock: + self._query_in_progress = False + if self._event is not None: + if self._log_prefix: + verbose_proxy_logger.debug( + "%s Signaling all waiting requests", self._log_prefix + ) + self._event.set() + self._event = None + + async def get_or_load( + self, + cache_key: str, + cache: AsyncCacheProtocol, + load_fn: Callable[[], Awaitable[T]], + ) -> Optional[T]: + """ + Return cached value or load it once and signal waiters. + + - cache_key: Key to read/write in the cache. + - cache: Object with async_get_cache(key) and async_set_cache(key, value). + - load_fn: Async callable that performs the load (e.g. DB query). No args. + Return value is cached and returned. If it raises, waiters are + still signaled so they can retry or handle empty cache. + + Returns the value from cache or from load_fn, or None if load failed or + cache was still empty after waiting. + """ + value = await self._get_cached(cache_key, cache) + if value is not None: + self._log_cache_hit(value) + return value + + self._log_cache_miss() + event_to_wait = await self._claim_role() + + if event_to_wait is not None: + return await self._wait_for_signal_and_get( + event_to_wait, cache_key, cache + ) + + try: + return await self._load_and_cache(cache_key, cache, load_fn) + finally: + await self._signal_done()