diff --git a/litellm/constants.py b/litellm/constants.py index 78bfc6501e8..a131585d5b8 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1420,7 +1420,7 @@ LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE = int( ) LITELLM_PROXY_ADMIN_NAME = "default_user_id" LITELLM_PROXY_BUDGET_NAME = "litellm-proxy-budget" -GLOBAL_PROXY_SPEND_CACHE_KEY = f"{LITELLM_PROXY_ADMIN_NAME}:spend" +GLOBAL_PROXY_SPEND_COUNTER_KEY = f"spend:user:{LITELLM_PROXY_BUDGET_NAME}" ########################### CLI SSO AUTHENTICATION CONSTANTS ########################### LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli" diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f4d07c1a674..bccaf427c72 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -24,8 +24,7 @@ import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm.constants import ( - GLOBAL_PROXY_SPEND_CACHE_KEY, - LITELLM_PROXY_BUDGET_NAME, + GLOBAL_PROXY_SPEND_COUNTER_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, ) from litellm.integrations.otel.model.config import is_otel_v2_enabled @@ -76,7 +75,6 @@ from litellm.proxy.auth.resolvers import CredentialRef, Principal from litellm.proxy.auth.resolvers.store import IdentityStore from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs -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, @@ -492,33 +490,21 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use 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: UserApiKeyCache, - prisma_client: PrismaClient, -) -> Optional[float]: +async def _fetch_global_proxy_spend() -> float: """ - Fetch global spend with event-driven coordination to prevent cache stampede. - Uses EventDrivenCacheCoordinator: first request queries DB and signals others when done. + Read proxy-wide spend from the shared spend counter. - Reads the proxy budget aggregate user row, which accrues proxy-wide spend - per request and is zeroed by ResetBudgetJob every ``litellm.budget_duration``. + The counter is incremented atomically per request (``increment_spend_counters``) + so concurrent requests cannot lose an increment, and it reseeds from the + "litellm-proxy-budget" user row, which the spend writer accrues into and + ResetBudgetJob zeroes every ``litellm.budget_duration``. """ + from litellm.proxy.proxy_server import get_current_spend - async def _load_global_spend() -> Optional[float]: - proxy_budget_row = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": LITELLM_PROXY_BUDGET_NAME} - ) - return float(proxy_budget_row.spend) if proxy_budget_row is not None else None - - return await _global_spend_coordinator.get_or_load( - cache_key=cache_key, - cache=user_api_key_cache, # pyright: ignore[reportArgumentType] - load_fn=_load_global_spend, + return await get_current_spend( + counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY, + fallback_spend=0.0, + max_budget=litellm.max_budget, ) @@ -531,13 +517,7 @@ async def get_global_proxy_spend( ) -> Optional[float]: global_proxy_spend = None 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 = GLOBAL_PROXY_SPEND_CACHE_KEY - 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, - ) + global_proxy_spend = await _fetch_global_proxy_spend() if global_proxy_spend is not None: user_info = CallInfo( user_id=litellm_proxy_admin_name, @@ -1975,13 +1955,8 @@ async def _user_api_key_auth_builder( global_proxy_spend = None if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget - cache_key = GLOBAL_PROXY_SPEND_CACHE_KEY with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"): - 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, - ) + global_proxy_spend = await _fetch_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 deleted file mode 100644 index 2f09e92f1ba..00000000000 --- a/litellm/proxy/common_utils/cache_coordinator.py +++ /dev/null @@ -1,195 +0,0 @@ -""" -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. - - Matches ``DualCache`` / ``UserApiKeyCache`` call shapes (explicit optional params - before ``**kwargs``), not only ``(key, **kwargs)``, so overloads validate. - """ - - async def async_get_cache( - self, - key: str, - parent_otel_span: Any = None, - local_only: bool = False, - **kwargs: Any, - ) -> Any: ... - - async def async_set_cache( - self, - key: str, - value: Any, - local_only: bool = False, - **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). - - Args: - log_prefix: Prefix for debug log messages. - """ - - 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[Any]: - """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: Optional[T] = 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: - result = await self._load_and_cache(cache_key, cache, load_fn) - return result - finally: - await self._signal_done() diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1210192416f..23a5b8f9c53 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -6,7 +6,6 @@ from typing import Any, Callable, List, Literal, Optional, Union import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME from litellm.proxy._types import ( LiteLLM_BudgetTableFull, LiteLLM_EndUserTable, @@ -99,14 +98,6 @@ class ResetBudgetJob: except Exception as e: verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e) - @staticmethod - async def _invalidate_global_proxy_spend_cache() -> None: - """Drop the cached global-proxy spend accumulator after the proxy - budget aggregate row is reset, so the next auth-time load reads the - zeroed row instead of a stale (potentially never-expiring) counter. - """ - await ResetBudgetJob._invalidate_user_api_key_cache_entry(GLOBAL_PROXY_SPEND_CACHE_KEY) - @staticmethod async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None: """Drop a stale management-cache entry so the next read fetches from DB. @@ -562,8 +553,6 @@ class ResetBudgetJob: user_id = getattr(u, "user_id", None) if user_id: await self._invalidate_spend_counter(f"spend:user:{user_id}") - if user_id == LITELLM_PROXY_BUDGET_NAME: - await self._invalidate_global_proxy_spend_cache() end_time = time.time() if len(failed_users) > 0: # If any users failed to reset diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9023f5fb23f..de99047e44b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -232,7 +232,7 @@ from litellm.constants import ( DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, - GLOBAL_PROXY_SPEND_CACHE_KEY, + GLOBAL_PROXY_SPEND_COUNTER_KEY, LITELLM_PROXY_ADMIN_NAME, LITELLM_PROXY_BUDGET_NAME, PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, @@ -286,7 +286,6 @@ from litellm.proxy.auth.model_checks import ( get_team_models, ) from litellm.proxy.auth.user_api_key_auth import ( - _fetch_global_spend_with_event_coordination, user_api_key_auth, user_api_key_auth_websocket, ) @@ -1073,7 +1072,6 @@ async def proxy_startup_event(app: FastAPI): ProxyStartupEvent._add_proxy_budget_to_db() asyncio.create_task( ProxyStartupEvent._warm_global_spend_cache( - user_api_key_cache=user_api_key_cache, prisma_client=prisma_client, ) ) @@ -2492,6 +2490,18 @@ async def increment_spend_counters( increment=cost, ) + async def _global_proxy_scope() -> None: + # The proxy-wide accumulator is the "litellm-proxy-budget" user row that + # db_spend_update_writer also accrues into, so it shares the user counter + # key space and reseeds from that row like any other user counter. + if GLOBAL_PROXY_SPEND_COUNTER_KEY in reserved_counter_keys: + return + await _init_and_increment_spend_counter( + counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY, + source_cache_key=LITELLM_PROXY_BUDGET_NAME, + increment=cost, + ) + scope_coros = tuple( coro for coro in ( @@ -2499,6 +2509,7 @@ async def increment_spend_counters( _team_scope(team_id) if team_id is not None else None, _team_member_scope(user_id, team_id) if user_id is not None and team_id is not None else None, _user_scope(user_id) if user_id is not None else None, + _global_proxy_scope() if litellm.max_budget > 0 and user_id != LITELLM_PROXY_BUDGET_NAME else None, _increment_end_user_and_tag_spend_counters( end_user_id=end_user_id, tags=tags, @@ -2904,14 +2915,6 @@ async def update_cache( CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_UserTable), ) ) - ## UPDATE GLOBAL PROXY ## - global_proxy_spend = await user_api_key_cache.async_get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) - if global_proxy_spend is None: - # do nothing if not in cache - return - elif response_cost is not None and global_proxy_spend is not None: - increment = global_proxy_spend + response_cost - values_to_update_in_cache.append((GLOBAL_PROXY_SPEND_CACHE_KEY, increment)) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update user spend in cache. " @@ -3072,27 +3075,15 @@ async def update_cache( if tags is not None: await _update_tag_cache() - global_proxy_spend_key = GLOBAL_PROXY_SPEND_CACHE_KEY - local_object_updates = tuple((k, v) for k, v in values_to_update_in_cache if k != global_proxy_spend_key) - shared_scalar_updates = tuple((k, v) for k, v in values_to_update_in_cache if k == global_proxy_spend_key) - - if local_object_updates: + if values_to_update_in_cache: asyncio.create_task( user_api_key_cache.async_set_cache_pipeline( - cache_list=list(local_object_updates), + cache_list=values_to_update_in_cache, ttl=get_management_object_ttl(user_api_key_cache), litellm_parent_otel_span=parent_otel_span, local_only=True, ) ) - if shared_scalar_updates: - asyncio.create_task( - user_api_key_cache.async_set_cache_pipeline( - cache_list=list(shared_scalar_updates), - ttl=get_management_object_ttl(user_api_key_cache), - litellm_parent_otel_span=parent_otel_span, - ) - ) def run_ollama_serve(): @@ -7977,16 +7968,14 @@ class ProxyStartupEvent: @classmethod async def _warm_global_spend_cache( cls, - user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, ) -> None: - """Warm global spend cache once at startup to reduce impact of first wave of requests.""" + """Warm global spend counter once at startup to reduce impact of first wave of requests.""" try: - cache_key = GLOBAL_PROXY_SPEND_CACHE_KEY - await _fetch_global_spend_with_event_coordination( - cache_key=cache_key, - user_api_key_cache=user_api_key_cache, + await SpendCounterReseed.coalesced( prisma_client=prisma_client, + spend_counter_cache=spend_counter_cache, + counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY, ) except Exception as e: verbose_proxy_logger.debug("Global spend cache warm-up at startup skipped or failed: %s", e) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index affaaa3fbf4..904cb57de9f 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5082,17 +5082,16 @@ class TestCheckKeyModelBudgetWithFallback: @pytest.mark.asyncio -async def test_global_proxy_spend_reads_resettable_proxy_budget_row(): +async def test_global_proxy_spend_reads_resettable_proxy_budget_row(monkeypatch): """Regression for the global proxy budget ignoring budget_duration (LIT-4309 / gh#31292): the enforced global spend must be loaded from the "litellm-proxy-budget" user row, which the spend writer increments per request and ResetBudgetJob zeroes every budget_duration. It must NOT be loaded from the MonthlyGlobalSpend view, whose window is hardcoded to a trailing 30 days and never resets on the configured duration.""" - from litellm.proxy.auth.user_api_key_auth import ( - _fetch_global_spend_with_event_coordination, - ) - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + import litellm.proxy.proxy_server as ps + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.user_api_key_auth import _fetch_global_proxy_spend proxy_budget_row = MagicMock() proxy_budget_row.spend = 42.5 @@ -5101,12 +5100,11 @@ async def test_global_proxy_spend_reads_resettable_proxy_budget_row(): prisma_client.db.query_raw = AsyncMock( side_effect=AssertionError("global spend must not be loaded from the fixed-30d MonthlyGlobalSpend view") ) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "spend_counter_cache", DualCache()) + monkeypatch.setattr(litellm, "max_budget", 100.0) - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", - user_api_key_cache=UserApiKeyCache(), - prisma_client=prisma_client, - ) + result = await _fetch_global_proxy_spend() assert result == 42.5 prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with( @@ -5115,24 +5113,46 @@ async def test_global_proxy_spend_reads_resettable_proxy_budget_row(): @pytest.mark.asyncio -async def test_global_proxy_spend_none_when_proxy_budget_row_missing(): +async def test_global_proxy_spend_zero_when_proxy_budget_row_missing(monkeypatch): """Before the startup upsert creates the aggregate row, enforcement must - see None (no cap applied) rather than raising.""" - from litellm.proxy.auth.user_api_key_auth import ( - _fetch_global_spend_with_event_coordination, - ) - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + see no recorded spend (no cap applied) rather than raising.""" + import litellm.proxy.proxy_server as ps + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.user_api_key_auth import _fetch_global_proxy_spend prisma_client = MagicMock() prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "spend_counter_cache", DualCache()) + monkeypatch.setattr(litellm, "max_budget", 100.0) - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", - user_api_key_cache=UserApiKeyCache(), - prisma_client=prisma_client, - ) + result = await _fetch_global_proxy_spend() - assert result is None + assert result == 0.0 + + +@pytest.mark.asyncio +async def test_global_proxy_spend_reads_shared_counter_not_the_db_row(monkeypatch): + """Once the counter is warm, the enforced value is the shared counter that + every pod increments atomically, so spend recorded by other pods since the + last DB flush is visible without another read of the lagging DB row.""" + import litellm.proxy.proxy_server as ps + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.user_api_key_auth import _fetch_global_proxy_spend + + proxy_budget_row = MagicMock() + proxy_budget_row.spend = 5.0 + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=proxy_budget_row) + counter_cache = DualCache() + counter_cache.in_memory_cache.set_cache(key="spend:user:litellm-proxy-budget", value=9.0) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(litellm, "max_budget", 100.0) + + result = await _fetch_global_proxy_spend() + + assert result == 9.0 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index f04d6f3cf5a..ddb6dbdb34e 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1367,14 +1367,14 @@ def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:user:alice", value=0.0, ttl=60) -def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_cache( +def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_counter( reset_budget_job, mock_prisma_client, monkeypatch ): """Regression for LIT-4309: resetting the proxy-wide budget aggregate row - ("litellm-proxy-budget") must also drop the cached global-spend - accumulator ("{admin}:spend") that _global_proxy_budget_check enforces - against. Without the invalidation, the cached value survives the DB reset - and the global cap keeps blocking requests for the whole next window.""" + ("litellm-proxy-budget") must also zero the shared spend counter that + _global_proxy_budget_check enforces against. Without the invalidation, the + cached value survives the DB reset and the global cap keeps blocking + requests for the whole next window.""" counter_cache = _make_counter_invalidation_job(monkeypatch) now = datetime.now(timezone.utc) @@ -1394,13 +1394,15 @@ def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_cache( asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) - counter_cache.user_api_key_cache.async_delete_cache.assert_any_call(key="default_user_id:spend") + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:user:litellm-proxy-budget", value=0.0, ttl=60 + ) -def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache( +def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_counter( reset_budget_job, mock_prisma_client, monkeypatch ): - """The global-spend accumulator must only be dropped when the proxy + """The global-spend accumulator must only be zeroed when the proxy budget aggregate row itself resets, not on every user reset.""" counter_cache = _make_counter_invalidation_job(monkeypatch) @@ -1422,8 +1424,8 @@ def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache( asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) assert not any( - call.kwargs.get("key") == "default_user_id:spend" - for call in counter_cache.user_api_key_cache.async_delete_cache.call_args_list + call.kwargs.get("key") == "spend:user:litellm-proxy-budget" + for call in counter_cache.in_memory_cache.set_cache.call_args_list ) diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 51980342a1d..aadbb34dcae 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -506,6 +506,141 @@ async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): } +class _AtomicRedisCounter: + """Redis double with INCRBYFLOAT semantics and a read barrier. + + ``async_get_cache`` releases only once ``expected`` callers have read the + counter, so every concurrent caller observes the same pre-increment value. + A read-add-write implementation collapses to a single increment under that + barrier; an atomic increment does not. + """ + + def __init__(self, expected_readers: int): + self.values: dict[str, float] = {} + self.barrier = asyncio.Barrier(expected_readers) + + async def async_get_cache(self, key, **kwargs): + if key.startswith("spend:"): + await self.barrier.wait() + return self.values.get(key) + + async def async_increment(self, *, key, value, refresh_ttl=True, **kwargs): + self.values[key] = self.values.get(key, 0.0) + value + return self.values[key] + + +@pytest.mark.asyncio +async def test_increment_spend_counters_global_proxy_counter_keeps_concurrent_costs( + monkeypatch, +): + """Regression for gh#35567: the proxy-wide accumulator enforced against + ``litellm.max_budget`` must be an atomic shared increment. Two concurrent + completions that both read the counter before either writes must leave the + counter at the sum of their costs, not at one of them.""" + import litellm + + redis = _AtomicRedisCounter(expected_readers=2) + fake_cache = _make_spend_counter_cache(redis_get_value=None) + fake_cache.redis_cache.async_get_cache = redis.async_get_cache + fake_cache.redis_cache.async_increment = redis.async_increment + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache(get_value=None)) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) + monkeypatch.setattr(litellm, "max_budget", 100.0) + + await asyncio.gather( + ps.increment_spend_counters(token=None, team_id=None, user_id=None, response_cost=1.0), + ps.increment_spend_counters(token=None, team_id=None, user_id=None, response_cost=1.0), + ) + + assert redis.values["spend:user:litellm-proxy-budget"] == 2.0 + + +@pytest.mark.asyncio +async def test_increment_spend_counters_skips_global_proxy_counter_without_max_budget( + monkeypatch, +): + """No global cap configured means no proxy-wide row to accrue into, so the + counter must not be created (and must not cost an extra Redis round trip).""" + import litellm + + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=1.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache(get_value=None)) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) + monkeypatch.setattr(litellm, "max_budget", 0.0) + + await ps.increment_spend_counters(token=None, team_id=None, user_id="u1", response_cost=1.0) + + incremented_keys = {call.kwargs["key"] for call in fake_cache.redis_cache.async_increment.call_args_list} + assert incremented_keys == {"spend:user:u1"} + + +@pytest.mark.asyncio +async def test_increment_spend_counters_does_not_double_count_proxy_budget_user( + monkeypatch, +): + """The proxy-wide accumulator lives in the user counter key space, so a + request attributed to that user id must increment it once, not twice.""" + import litellm + + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=1.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache(get_value=None)) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) + monkeypatch.setattr(litellm, "max_budget", 100.0) + + await ps.increment_spend_counters( + token=None, team_id=None, user_id="litellm-proxy-budget", response_cost=1.0 + ) + + incremented_keys = [call.kwargs["key"] for call in fake_cache.redis_cache.async_increment.call_args_list] + assert incremented_keys == ["spend:user:litellm-proxy-budget"] + + +@pytest.mark.asyncio +async def test_update_cache_does_not_write_the_global_spend_scalar(monkeypatch): + """gh#35567: the cached-object refresh must stay per-pod. Any shared write + of the proxy-wide spend from here is a read-add-write that loses concurrent + increments; the atomic counter owns that value now.""" + import litellm + from litellm.models.user import LiteLLM_UserTable + from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec + + cached = { + "u1": CacheCodec.serialize(LiteLLM_UserTable(user_id="u1", spend=0.0), model_type=LiteLLM_UserTable), + "default_user_id:spend": 10.0, + } + fake_user_cache = _make_user_api_key_cache(get_side_effect=lambda key, **kwargs: cached.get(key)) + monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) + monkeypatch.setattr(litellm, "max_budget", 100.0) + + await ps.update_cache( + token=None, + user_id="u1", + end_user_id=None, + team_id=None, + response_cost=1.0, + parent_otel_span=None, + tags=None, + ) + await asyncio.sleep(0) + + written_keys = [ + key + for call in fake_user_cache.async_set_cache_pipeline.call_args_list + for key, _ in call.kwargs["cache_list"] + ] + assert written_keys == ["u1"] + assert all( + call.kwargs.get("local_only") is True + for call in fake_user_cache.async_set_cache_pipeline.call_args_list + ) + + @pytest.mark.asyncio async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch): """Counters already reserved by a budget reservation are skipped, every diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ab64224fdc3..75e47186def 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4952,74 +4952,6 @@ async def test_spend_tracking_never_writes_the_auth_object_back(): setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache) -@pytest.mark.asyncio -async def test_update_cache_global_proxy_spend_scalar_stays_shared(): - """ - The proxy-wide spend estimate must keep flowing to Redis when the spend - writeback goes per-pod: the global max_budget check reads the - ``{litellm_proxy_admin_name}:spend`` cache entry between authoritative DB - reloads, so keeping it pod-local would let traffic spread across replicas - exceed the proxy budget by roughly a factor of the replica count within a - cache TTL. Sharing this scalar is safe because it carries no limits or - permissions, so it cannot resurrect an invalidated auth blob. - """ - from litellm.caching.caching import DualCache - - admin_name = litellm.proxy.proxy_server.litellm_proxy_admin_name - global_key = "{}:spend".format(admin_name) - - async def fake_get(key, **kwargs): - if key == "user-lit": - return {"user_id": "user-lit", "spend": 1.0} - if key == global_key: - return 10.0 - return None - - original_cache = litellm.proxy.proxy_server.user_api_key_cache - cache = DualCache(default_in_memory_ttl=300) - setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) - try: - with patch.object( - cache, "async_get_cache", new=AsyncMock(side_effect=fake_get) - ): - with patch.object( - cache, "async_set_cache_pipeline", new=AsyncMock() - ) as mock_set_cache: - await litellm.proxy.proxy_server.update_cache( - token=None, - user_id="user-lit", - end_user_id=None, - team_id=None, - response_cost=5.0, - parent_otel_span=None, - ) - - pending = [ - t for t in asyncio.all_tasks() if t is not asyncio.current_task() - ] - if pending: - await asyncio.wait(pending, timeout=5) - - calls = mock_set_cache.await_args_list - local_keys = [ - k - for c in calls - if c.kwargs.get("local_only") is True - for k, _ in c.kwargs["cache_list"] - ] - shared_keys = [ - k - for c in calls - if c.kwargs.get("local_only") is not True - for k, _ in c.kwargs["cache_list"] - ] - assert "user-lit" in local_keys - assert global_key not in local_keys - assert shared_keys == [global_key] - finally: - setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache) - - @pytest.mark.asyncio async def test_init_sso_settings_in_db(): """