mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Add EventDrivenCacheCoordinator and use it for global spend to prevent cache stampede
This commit is contained in:
Alexsander Hamir
parent
7d5439adda
commit
672d70b409
2 changed files with 231 additions and 32 deletions
|
|
@ -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(
|
||||
|
|
|
|||
191
litellm/proxy/common_utils/cache_coordinator.py
Normal file
191
litellm/proxy/common_utils/cache_coordinator.py
Normal file
|
|
@ -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()
|
||||
Loading…
Add table
Reference in a new issue