From 3c3d7fd6345b62dd3e7fc7e82bf34d39e29fce5d Mon Sep 17 00:00:00 2001 From: lostmartian Date: Sat, 15 Aug 2026 23:32:40 +0530 Subject: [PATCH] fix(proxy): track global proxy spend with the shared spend-counter machinery The previous iterations hand-rolled a separate cache path for the global proxy spend scalar, and each variant reintroduced a race the shared spend counter system was built to avoid. This routes global proxy spend through that machinery instead: - Write: _init_and_increment_spend_counter seeds a cold counter from the authoritative litellm-proxy-budget row (cross-pod SET NX) and then increments atomically, keyed spend:user:litellm-proxy-budget. - Read: get_current_spend reads Redis-first, reseeds from the DB on a clean miss, and only floors to the DB row when the freshly loaded fallback exceeds the counter, repaired with the monotonic async_set_max. - Reset: the existing user-reset path already zeroes the counter for the proxy-budget row; the job now also invalidates the cached row-spend hint so no worker acts on a pre-reset fallback. Removes the bespoke coordinator, per-worker DB-floor marker, and delete-on- reset invalidation. --- litellm/constants.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 138 ++++------- .../proxy/common_utils/reset_budget_job.py | 29 +-- litellm/proxy/proxy_server.py | 38 +-- .../proxy/auth/test_user_api_key_auth.py | 224 +++--------------- .../common_utils/test_reset_budget_job.py | 58 +---- .../proxy/proxy_server/test_spend_counters.py | 152 ++++-------- tests/test_litellm/proxy/test_proxy_server.py | 75 ++---- 8 files changed, 181 insertions(+), 535 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6449834d6a4..9835507af1e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1449,7 +1449,7 @@ LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE: Final = int( ) LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id" LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget" -GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend" +GLOBAL_PROXY_SPEND_COUNTER_KEY: Final = f"spend:user:{LITELLM_PROXY_BUDGET_NAME}" ########################### CLI SSO AUTHENTICATION CONSTANTS ########################### LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli" diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index fdc86307d5e..b0af758a349 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -23,7 +23,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, + GLOBAL_PROXY_SPEND_COUNTER_KEY, LITELLM_PROXY_BUDGET_NAME, LITELLM_PROXY_MASTER_KEY_ALIAS, ) @@ -76,7 +76,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, @@ -491,103 +490,57 @@ 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: Final = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEND]") - -# Cache the authoritative DB spend read for a few seconds per worker so the -# global-proxy budget floor check does not query the DB on every request. -_GLOBAL_PROXY_SPEND_DB_FLOOR_TTL_SECONDS: Final = 5 -_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX: Final = "global_proxy_spend_db_floor:" +async def _read_global_proxy_spend_hint_from_cache(user_api_key_cache: UserApiKeyCache) -> float | None: + redis_cache: Final = user_api_key_cache.redis_cache + if redis_cache is not None: + try: + cached: Final = await redis_cache.async_get_cache(key=LITELLM_PROXY_BUDGET_NAME) + except Exception as e: # noqa: BLE001 # hint read failure must not break auth; in-memory fallback below + verbose_proxy_logger.debug("Global proxy spend hint: Redis read failed, falling back to in-memory: %s", e) + else: + return float(cached) if cached is not None else None + cached_in_memory: Final = user_api_key_cache.in_memory_cache.get_cache(key=LITELLM_PROXY_BUDGET_NAME) + return float(cached_in_memory) if cached_in_memory is not None else None -def invalidate_global_proxy_spend_db_floor_marker(user_api_key_cache: UserApiKeyCache, cache_key: str) -> None: - """Drop the worker-local DB-floor marker for ``cache_key``. - - ResetBudgetJob calls this after zeroing the aggregate DB row so a stale - pre-reset marker is not used to write previous-window spend back into - Redis during the next auth-time floor check. - """ - marker_key: Final = f"{_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX}{cache_key}" - user_api_key_cache.in_memory_cache.delete_cache(key=marker_key) - - -async def _fetch_global_spend_with_event_coordination( - cache_key: str, +async def _load_global_proxy_spend_hint( user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, ) -> float | None: + """Spend of the proxy-budget aggregate user row, cached in the management cache. + + This is only the fallback/staleness signal for get_current_spend; the + enforced value is the shared spend counter. Read Redis-first so the reset + job's invalidation is visible to every worker immediately. """ - Fetch global spend with event-driven coordination to prevent cache stampede. - Uses EventDrivenCacheCoordinator: first request queries DB and signals others when done. - - Reads the proxy budget aggregate user row, which accrues proxy-wide spend - per request and is zeroed by ResetBudgetJob every ``litellm.budget_duration``. - - With Redis configured, the shared scalar is read straight from Redis: the - global spend is incremented by every worker, so a per-worker in-memory copy - can go stale across resets and falsely reject requests at auth time. Reading - the Redis value directly keeps the budget check on the shared total. On a - Redis miss the coordinator loads the authoritative DB spend once and - repopulates Redis. Without Redis (single process), the normal DualCache - coordinator path is used. - """ - - async def _load_global_spend() -> float | None: - proxy_budget_row: Final = 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 - - async def _authoritative_db_spend() -> float | None: - # Cached per-worker for a few seconds to avoid a DB query per request. - marker_key: Final = f"{_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX}{cache_key}" - cached = user_api_key_cache.in_memory_cache.get_cache(key=marker_key) - if cached is not None: - return float(cached) - db_spend = await _load_global_spend() - if db_spend is not None: - user_api_key_cache.in_memory_cache.set_cache( - key=marker_key, - value=db_spend, - ttl=_GLOBAL_PROXY_SPEND_DB_FLOOR_TTL_SECONDS, - ) - return db_spend - - spend_cache: Final = ( - user_api_key_cache.redis_cache if user_api_key_cache.redis_cache is not None else user_api_key_cache + cached: Final = await _read_global_proxy_spend_hint_from_cache(user_api_key_cache) + if cached is not None: + return cached + proxy_budget_row: Final = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": LITELLM_PROXY_BUDGET_NAME} ) - value = await _global_spend_coordinator.get_or_load( - cache_key=cache_key, - cache=spend_cache, # pyright: ignore[reportArgumentType] - load_fn=_load_global_spend, - ) - if value is None: + if proxy_budget_row is None: return None + spend: Final = float(proxy_budget_row.spend or 0.0) + await user_api_key_cache.async_set_cache(key=LITELLM_PROXY_BUDGET_NAME, value=spend) + return spend - try: - numeric_value = float(value) - except (TypeError, ValueError): - numeric_value = 0.0 - # Floor to the authoritative DB spend when Redis is in use. A Redis counter - # that was clobbered by a reset race, restarted from an old RDB snapshot, - # or recreated from zero after expiry can under-report current-period spend; - # capping the value to the DB row prevents requests from leaking past the - # global cap. Without Redis the coordinator already loaded the DB row. - if user_api_key_cache.redis_cache is not None: - db_spend = await _authoritative_db_spend() - if db_spend is not None and db_spend > numeric_value: - # Use a monotonic write so a concurrent increment that already - # pushed Redis above the DB value is not overwritten by a repair - # carrying a slightly-stale DB total. Use the resulting Redis value - # so the enforced total reflects any concurrent increment. - repaired_value: Final[float | None] = await user_api_key_cache.redis_cache.async_set_max( - key=cache_key, value=db_spend - ) - numeric_value = repaired_value if repaired_value is not None else db_spend - user_api_key_cache.in_memory_cache.set_cache(key=cache_key, value=numeric_value) +async def _get_global_proxy_spend_value( + user_api_key_cache: UserApiKeyCache, + prisma_client: PrismaClient, +) -> float: + from litellm.proxy.proxy_server import get_current_spend - return numeric_value + hint: Final = await _load_global_proxy_spend_hint( + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + ) + return await get_current_spend( + counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY, + fallback_spend=hint if hint is not None else 0.0, + max_budget=litellm.max_budget, + ) async def get_global_proxy_spend( @@ -599,10 +552,7 @@ async def get_global_proxy_spend( ) -> float | None: 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: Final = GLOBAL_PROXY_SPEND_CACHE_KEY - global_proxy_spend = await _fetch_global_spend_with_event_coordination( - cache_key=cache_key, + global_proxy_spend = await _get_global_proxy_spend_value( user_api_key_cache=user_api_key_cache, prisma_client=prisma_client, ) @@ -2112,10 +2062,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: Final = 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, + global_proxy_spend = await _get_global_proxy_spend_value( user_api_key_cache=user_api_key_cache, prisma_client=prisma_client, ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 6b9a2904867..801ce643828 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -11,7 +11,6 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.constants import ( - GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME, RESET_BUDGET_JOB_BATCH_SIZE, RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN, @@ -262,28 +261,14 @@ class ResetBudgetJob: @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 read reloads the - zeroed DB row. Reads and writes are atomic now (increment-if-present on - Redis), so deleting cannot resurrect a stale pre-reset value, and it - cannot clobber a concurrent increment the way a plain SET to zero could. - """ - try: - from litellm.proxy.auth.user_api_key_auth import ( - invalidate_global_proxy_spend_db_floor_marker, - ) - from litellm.proxy.proxy_server import user_api_key_cache + """Drop the cached proxy-budget row spend after the row is reset. - await user_api_key_cache.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) - invalidate_global_proxy_spend_db_floor_marker( - user_api_key_cache=user_api_key_cache, - cache_key=GLOBAL_PROXY_SPEND_CACHE_KEY, - ) - except Exception as e: # noqa: BLE001 # cache invalidation failure must not break budget reset - verbose_proxy_logger.warning( - "Failed to invalidate global proxy spend cache after budget reset: %s", - e, - ) + Auth uses it only as the fallback signal for the shared spend counter's + DB floor, so every worker must reload the zeroed row instead of acting + on a pre-reset value. The counter itself is zeroed by the standard + ``_invalidate_spend_counter`` path for the proxy-budget user row. + """ + await ResetBudgetJob._invalidate_user_api_key_cache_entry(LITELLM_PROXY_BUDGET_NAME) @staticmethod async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7780a68c83f..cb4ba0890a7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -235,7 +235,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, MONTHLY_SPEND_REPORT_JOB_ID, @@ -292,7 +292,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, ) @@ -3123,27 +3122,16 @@ async def update_cache( ### UPDATE GLOBAL PROXY SPEND ### async def _update_global_proxy_spend() -> None: - if response_cost is None: + if response_cost is None or litellm.max_budget <= 0: return try: - # Atomically accumulate the shared global-proxy scalar in Redis. - # Authorization reads this Redis scalar directly, so there is no - # per-worker in-memory copy that can go stale across resets. - if user_api_key_cache.redis_cache is not None: - await user_api_key_cache.redis_cache.async_increment( - key=GLOBAL_PROXY_SPEND_CACHE_KEY, - value=response_cost, - ttl=get_management_object_ttl(user_api_key_cache), - refresh_ttl=True, - ) - else: - # No Redis: single-process, so an in-memory increment is coherent. - await user_api_key_cache.async_increment_cache( - key=GLOBAL_PROXY_SPEND_CACHE_KEY, - value=response_cost, - ttl=get_management_object_ttl(user_api_key_cache), - refresh_ttl=True, - ) + # Same spend-counter path as every other budget entity: seed a cold + # counter from the authoritative DB row (SET NX), then increment. + await _init_and_increment_spend_counter( + counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY, + source_cache_key=LITELLM_PROXY_BUDGET_NAME, + increment=response_cost, + ) except Exception as e: # noqa: BLE001 # spend update must not fail the request verbose_proxy_logger.warning( "Spend tracking - failed to update global proxy spend in cache. " @@ -8613,11 +8601,9 @@ class ProxyStartupEvent: ) -> None: """Warm global spend cache once at startup to reduce impact of first wave of requests.""" try: - cache_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY - await _fetch_global_spend_with_event_coordination( - cache_key=cache_key, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, + await get_current_spend( + counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY, + fallback_spend=0.0, ) 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 abef9530783..0faac49da5b 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 @@ -6068,16 +6068,14 @@ class TestCheckKeyModelBudgetWithFallback: @pytest.mark.asyncio -async def test_global_proxy_spend_reads_resettable_proxy_budget_row(): +async def test_global_proxy_spend_hint_reads_resettable_proxy_budget_row(): """Regression for the global proxy budget ignoring budget_duration - (LIT-4309 / gh#31292): the enforced global spend must be loaded from the + (LIT-4309 / gh#31292): the fallback 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.auth.user_api_key_auth import _load_global_proxy_spend_hint from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache proxy_budget_row = MagicMock() @@ -6088,8 +6086,7 @@ async def test_global_proxy_spend_reads_resettable_proxy_budget_row(): side_effect=AssertionError("global spend must not be loaded from the fixed-30d MonthlyGlobalSpend view") ) - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", + result = await _load_global_proxy_spend_hint( user_api_key_cache=UserApiKeyCache(), prisma_client=prisma_client, ) @@ -6101,19 +6098,16 @@ 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(): - """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, - ) +async def test_global_proxy_spend_hint_none_when_proxy_budget_row_missing(): + """Before the startup upsert creates the aggregate row, the fallback hint + is None so the counter falls back to 0 rather than raising.""" + from litellm.proxy.auth.user_api_key_auth import _load_global_proxy_spend_hint from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache prisma_client = MagicMock() prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", + result = await _load_global_proxy_spend_hint( user_api_key_cache=UserApiKeyCache(), prisma_client=prisma_client, ) @@ -6122,209 +6116,61 @@ async def test_global_proxy_spend_none_when_proxy_budget_row_missing(): @pytest.mark.asyncio -async def test_global_proxy_spend_reads_shared_redis_scalar_when_configured(): - """With Redis, the global spend read goes straight to the shared Redis - scalar instead of a per-worker in-memory copy, which can go stale across - resets and falsely reject requests at auth time.""" - from litellm.proxy.auth.user_api_key_auth import ( - _fetch_global_spend_with_event_coordination, - ) +async def test_global_proxy_spend_hint_reads_redis_first(): + """The hint is read from the shared Redis tier first so a reset + invalidation is visible to every worker immediately; a per-worker + in-memory copy is only a Redis-outage fallback.""" + from litellm.proxy.auth.user_api_key_auth import _load_global_proxy_spend_hint from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache cache = UserApiKeyCache() redis_cache = MagicMock() redis_cache.async_get_cache = AsyncMock(return_value=8.5) - redis_cache.async_set_cache = AsyncMock() cache.redis_cache = redis_cache + cache.in_memory_cache.set_cache(key="litellm-proxy-budget", value=999.0) - 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 + side_effect=AssertionError("DB must not be queried when Redis has the hint") ) - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", + result = await _load_global_proxy_spend_hint( user_api_key_cache=cache, prisma_client=prisma_client, ) assert result == 8.5 - redis_cache.async_get_cache.assert_awaited_once_with(key="default_user_id:spend") - redis_cache.async_set_cache.assert_not_called() + redis_cache.async_get_cache.assert_awaited_once_with(key="litellm-proxy-budget") @pytest.mark.asyncio -async def test_global_proxy_spend_floors_stale_low_redis_to_db(): - """A Redis counter that is stale-low (e.g., recreated from zero after a - reset or Redis restart) must be floored to the authoritative DB spend so - requests cannot leak past the global cap.""" - from litellm.proxy.auth.user_api_key_auth import ( - _fetch_global_spend_with_event_coordination, - ) +async def test_global_proxy_spend_value_uses_shared_spend_counter(monkeypatch): + """The enforced global proxy spend comes from the shared spend counter + (spend:user:litellm-proxy-budget) via get_current_spend, with the DB row + spend as the fallback signal and the configured cap as max_budget.""" + import litellm + from litellm.proxy.auth.user_api_key_auth import _get_global_proxy_spend_value from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - cache = UserApiKeyCache() - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value="2.0") - redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value) - cache.redis_cache = redis_cache - proxy_budget_row = MagicMock() proxy_budget_row.spend = 9.0 prisma_client = MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=proxy_budget_row - ) + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=proxy_budget_row) - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", - user_api_key_cache=cache, + get_current_spend = AsyncMock(return_value=9.0) + monkeypatch.setattr("litellm.proxy.proxy_server.get_current_spend", get_current_spend) + monkeypatch.setattr(litellm, "max_budget", 100.0) + + result = await _get_global_proxy_spend_value( + user_api_key_cache=UserApiKeyCache(), prisma_client=prisma_client, ) assert result == 9.0 - redis_cache.async_get_cache.assert_awaited_once_with(key="default_user_id:spend") - redis_cache.async_set_max.assert_awaited_once_with( - key="default_user_id:spend", value=9.0 - ) - - -@pytest.mark.asyncio -async def test_global_proxy_spend_floor_does_not_overwrite_concurrent_increment(): - """When Redis already holds a higher value than the DB floor (because a - concurrent request incremented it after our read), the repair must use a - monotonic write so the concurrent increment is not discarded.""" - 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 - - cache = UserApiKeyCache() - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value="2.0") - redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value) - cache.redis_cache = redis_cache - - proxy_budget_row = MagicMock() - proxy_budget_row.spend = 1.0 - prisma_client = MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=proxy_budget_row - ) - - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", - user_api_key_cache=cache, - prisma_client=prisma_client, - ) - - # DB (1.0) is not greater than Redis (2.0), so no repair is attempted. - assert result == 2.0 - redis_cache.async_set_max.assert_not_called() - - -@pytest.mark.asyncio -async def test_global_proxy_spend_marker_cached_db_value_used_without_db_query(): - """The DB-floor marker should be reused for subsequent reads so we do not - query Postgres on every request.""" - from litellm.proxy.auth.user_api_key_auth import ( - _fetch_global_spend_with_event_coordination, - _GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX, - ) - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - - cache = UserApiKeyCache() - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value="2.0") - redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value) - cache.redis_cache = redis_cache - - # Seed the marker cache with a DB value greater than Redis. - marker_key = f"{_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX}default_user_id:spend" - cache.in_memory_cache.set_cache(key=marker_key, value=9.0) - - prisma_client = MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=AssertionError("DB must not be queried when marker is present") - ) - - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", - user_api_key_cache=cache, - prisma_client=prisma_client, - ) - - assert result == 9.0 - prisma_client.db.litellm_usertable.find_unique.assert_not_called() - redis_cache.async_set_max.assert_awaited_once_with( - key="default_user_id:spend", value=9.0 - ) - - -@pytest.mark.asyncio -async def test_global_proxy_spend_floor_falls_back_when_set_max_returns_none(): - """If the monotonic Redis repair returns None (e.g. connection hiccup), - the enforced total must still fall back to the authoritative DB floor.""" - 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 - - cache = UserApiKeyCache() - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value="2.0") - redis_cache.async_set_max = AsyncMock(return_value=None) - cache.redis_cache = redis_cache - - proxy_budget_row = MagicMock() - proxy_budget_row.spend = 9.0 - prisma_client = MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=proxy_budget_row - ) - - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", - user_api_key_cache=cache, - prisma_client=prisma_client, - ) - - assert result == 9.0 - - -@pytest.mark.asyncio -async def test_global_proxy_spend_treats_unparseable_redis_as_zero(): - """A corrupted/non-numeric Redis value is treated as 0.0 and then floored - to the DB total so enforcement remains safe.""" - 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 - - cache = UserApiKeyCache() - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value="not-a-number") - redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value) - cache.redis_cache = redis_cache - - proxy_budget_row = MagicMock() - proxy_budget_row.spend = 7.0 - prisma_client = MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=proxy_budget_row - ) - - result = await _fetch_global_spend_with_event_coordination( - cache_key="default_user_id:spend", - user_api_key_cache=cache, - prisma_client=prisma_client, - ) - - assert result == 7.0 - redis_cache.async_set_max.assert_awaited_once_with( - key="default_user_id:spend", value=7.0 + get_current_spend.assert_awaited_once_with( + counter_key="spend:user:litellm-proxy-budget", + fallback_spend=9.0, + max_budget=100.0, ) 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 3e5134c9faf..88959627ae2 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 @@ -863,48 +863,6 @@ def test_reset_budget_windows_uses_is_not_null_filter(monkeypatch): assert "budget_limits IS NOT NULL" in team_query -def test_invalidate_global_proxy_spend_cache_deletes_key_and_db_floor_marker(monkeypatch): - """ResetBudgetJob must drop both the shared Redis spend key and the worker- - local DB-floor marker. Otherwise a stale pre-reset marker can be used to - write previous-window spend back into Redis at the next auth check.""" - from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY - from litellm.proxy.auth.user_api_key_auth import ( - _GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX, - ) - - user_api_key_cache = MagicMock() - user_api_key_cache.async_delete_cache = AsyncMock() - user_api_key_cache.in_memory_cache.delete_cache = MagicMock() - - fake_module = types.ModuleType("litellm.proxy.proxy_server") - fake_module.user_api_key_cache = user_api_key_cache - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=MagicMock()) - asyncio.run(job._invalidate_global_proxy_spend_cache()) - - user_api_key_cache.async_delete_cache.assert_awaited_once_with( - key=GLOBAL_PROXY_SPEND_CACHE_KEY - ) - user_api_key_cache.in_memory_cache.delete_cache.assert_called_once_with( - key=f"{_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX}{GLOBAL_PROXY_SPEND_CACHE_KEY}" - ) - - -def test_invalidate_global_proxy_spend_cache_swallows_cache_failure(monkeypatch): - """A failure to invalidate the cache must not break the budget reset job.""" - user_api_key_cache = MagicMock() - user_api_key_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis down")) - - fake_module = types.ModuleType("litellm.proxy.proxy_server") - fake_module.user_api_key_cache = user_api_key_cache - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=MagicMock()) - # Should not raise. - asyncio.run(job._invalidate_global_proxy_spend_cache()) - - def test_reset_budget_windows_resets_expired_key_window(monkeypatch): """A key whose window's `reset_at` has passed gets an update with a new `reset_at` in the future, and the in-memory spend counter is cleared.""" @@ -1134,8 +1092,9 @@ def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock def test_reset_budget_for_proxy_budget_row_resets_global_spend_cache( reset_budget_job, mock_prisma_client, monkeypatch ): - """Resetting the proxy-wide budget aggregate row invalidates the cached - global-spend accumulator so the next read reloads the zeroed DB row.""" + """Resetting the proxy-wide budget aggregate row zeroes its spend counter + and invalidates the cached row-spend hint so the next auth reloads the + zeroed DB row instead of acting on a pre-reset value.""" counter_cache = _make_counter_invalidation_job(monkeypatch) now = datetime.now(timezone.utc) @@ -1155,16 +1114,19 @@ def test_reset_budget_for_proxy_budget_row_resets_global_spend_cache( asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:user:litellm-proxy-budget", value=0.0, ttl=60 + ) counter_cache.user_api_key_cache.async_delete_cache.assert_any_call( - key="default_user_id:spend" + key="litellm-proxy-budget" ) def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache( reset_budget_job, mock_prisma_client, monkeypatch ): - """The global-spend accumulator must only be invalidated when the proxy - budget aggregate row itself resets, not on every user reset.""" + """The proxy-budget hint must only be invalidated when the proxy budget + aggregate row itself resets, not on every user reset.""" counter_cache = _make_counter_invalidation_job(monkeypatch) now = datetime.now(timezone.utc) @@ -1185,7 +1147,7 @@ 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" + call.kwargs.get("key") == "litellm-proxy-budget" for call in counter_cache.user_api_key_cache.async_delete_cache.call_args_list ) assert not any( 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 065f3c83692..a0456769be4 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1273,133 +1273,79 @@ async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkey @pytest.mark.asyncio -async def test_update_cache_global_proxy_spend_starts_from_zero_after_reset( - monkeypatch, -): - """Global proxy spend starts from zero after the cache key is reset.""" +async def test_update_cache_global_proxy_spend_uses_spend_counter(monkeypatch): + """Global proxy spend goes through the shared spend-counter path so cold + counters seed from the authoritative DB row and increments stay atomic.""" + from unittest.mock import AsyncMock + from litellm.caching.caching import DualCache - cache = DualCache(default_in_memory_ttl=300) - monkeypatch.setattr(ps, "user_api_key_cache", cache) + init_and_increment = AsyncMock() + monkeypatch.setattr(ps, "_init_and_increment_spend_counter", init_and_increment) + monkeypatch.setattr(ps.litellm, "max_budget", 100.0) + monkeypatch.setattr(ps, "user_api_key_cache", DualCache(default_in_memory_ttl=300)) - await cache.async_set_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, value=50.0) - await cache.async_delete_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) - await cache.async_set_cache( - key="user-lit", - value={"user_id": "user-lit", "spend": 0.0}, - ) await ps.update_cache( token=None, - user_id="user-lit", - end_user_id=None, - team_id=None, - response_cost=1.05, - parent_otel_span=None, - ) - - final_value = await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) - assert final_value == 1.05 - - -@pytest.mark.asyncio -async def test_update_cache_global_proxy_spend_concurrent_increments_after_reset( - monkeypatch, -): - """Concurrent spend updates after a reset accumulate from zero.""" - from litellm.caching.caching import DualCache - - cache = DualCache(default_in_memory_ttl=300) - monkeypatch.setattr(ps, "user_api_key_cache", cache) - - await cache.async_set_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, value=100.0) - await cache.async_delete_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) - await cache.async_set_cache( - key="user-concurrent", - value={"user_id": "user-concurrent", "spend": 0.0}, - ) - - async def _update(cost: float) -> None: - await ps.update_cache( - token=None, - user_id="user-concurrent", - end_user_id=None, - team_id=None, - response_cost=cost, - parent_otel_span=None, - ) - - await asyncio.gather(*[_update(0.01) for _ in range(50)]) - - final_value = await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) - assert final_value == pytest.approx(0.5) - - -@pytest.mark.asyncio -async def test_update_cache_global_proxy_spend_does_not_mirror_to_in_memory( - monkeypatch, -): - """Authorization reads the shared Redis scalar directly, so the increment - path must not maintain a per-worker in-memory copy that could go stale - across resets and confuse budget checks.""" - from unittest.mock import AsyncMock, MagicMock - - from litellm.caching.caching import DualCache - - redis_cache = MagicMock() - redis_cache.async_increment = AsyncMock(return_value=3.5) - cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=300) - monkeypatch.setattr(ps, "user_api_key_cache", cache) - - await cache.async_set_cache( - key="user-no-mirror", - value={"user_id": "user-no-mirror", "spend": 0.0}, - ) - await ps.update_cache( - token=None, - user_id="user-no-mirror", + user_id=None, end_user_id=None, team_id=None, response_cost=1.5, parent_otel_span=None, ) - redis_cache.async_increment.assert_awaited_once_with( - key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, - value=1.5, - ttl=ps.get_management_object_ttl(cache), - refresh_ttl=True, + init_and_increment.assert_awaited_once_with( + counter_key="spend:user:litellm-proxy-budget", + source_cache_key="litellm-proxy-budget", + increment=1.5, ) - assert cache.in_memory_cache.get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) is None @pytest.mark.asyncio -async def test_update_cache_global_proxy_spend_swallows_redis_failure(monkeypatch): - """A Redis failure on the global spend increment must not fail the request.""" - from unittest.mock import AsyncMock, MagicMock +async def test_update_cache_global_proxy_spend_skipped_without_max_budget(monkeypatch): + """No global cap configured means the proxy-budget row is never written by + the DB flusher either, so the counter must not accumulate spend it could + not be reseeded from.""" + from unittest.mock import AsyncMock from litellm.caching.caching import DualCache - redis_cache = MagicMock() - redis_cache.async_increment = AsyncMock(side_effect=RuntimeError("redis down")) - cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=300) - monkeypatch.setattr(ps, "user_api_key_cache", cache) + init_and_increment = AsyncMock() + monkeypatch.setattr(ps, "_init_and_increment_spend_counter", init_and_increment) + monkeypatch.setattr(ps.litellm, "max_budget", 0) + monkeypatch.setattr(ps, "user_api_key_cache", DualCache(default_in_memory_ttl=300)) - await cache.async_set_cache( - key="user-redis-down", - value={"user_id": "user-redis-down", "spend": 0.0}, - ) await ps.update_cache( token=None, - user_id="user-redis-down", + user_id=None, + end_user_id=None, + team_id=None, + response_cost=1.5, + parent_otel_span=None, + ) + + init_and_increment.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_cache_global_proxy_spend_swallows_failure(monkeypatch): + """A failure on the global spend increment must not fail the request.""" + from unittest.mock import AsyncMock + + from litellm.caching.caching import DualCache + + init_and_increment = AsyncMock(side_effect=RuntimeError("redis down")) + monkeypatch.setattr(ps, "_init_and_increment_spend_counter", init_and_increment) + monkeypatch.setattr(ps.litellm, "max_budget", 100.0) + monkeypatch.setattr(ps, "user_api_key_cache", DualCache(default_in_memory_ttl=300)) + + await ps.update_cache( + token=None, + user_id=None, end_user_id=None, team_id=None, response_cost=1.0, parent_otel_span=None, ) - redis_cache.async_increment.assert_awaited_once_with( - key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, - value=1.0, - ttl=ps.get_management_object_ttl(cache), - refresh_ttl=True, - ) + init_and_increment.assert_awaited_once() diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 9cb5e4da286..18c8c3eb6af 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5273,65 +5273,38 @@ async def test_spend_tracking_never_writes_the_auth_object_back(): @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. The scalar is updated via an atomic cache increment, which - remains shared across pods and avoids resurrecting a stale pre-reset value. + The proxy-wide spend estimate must live in the shared spend counter + (``spend:user:litellm-proxy-budget``), which seeds from the authoritative + DB row when cold and increments atomically. A per-pod cache key would let + traffic spread across replicas exceed the proxy budget by roughly a factor + of the replica count within a cache TTL. """ from litellm.caching.caching import DualCache - admin_name = litellm.proxy.proxy_server.litellm_proxy_admin_name - global_key = f"{admin_name}:spend" - - async def fake_get(key, **kwargs): - if key == "user-lit": - return {"user_id": "user-lit", "spend": 1.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) + original_max_budget = litellm.max_budget + init_and_increment = AsyncMock() + setattr(litellm.proxy.proxy_server, "user_api_key_cache", DualCache(default_in_memory_ttl=300)) + setattr(litellm, "max_budget", 100.0) + setattr(litellm.proxy.proxy_server, "_init_and_increment_spend_counter", init_and_increment) 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: - with patch.object(cache, "async_increment_cache", new=AsyncMock()) as mock_increment: - 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, - ) + await litellm.proxy.proxy_server.update_cache( + token=None, + user_id=None, + 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 == [] - - increment_calls = [ - c - for c in mock_increment.await_args_list - if (c.args[0] if c.args else c.kwargs.get("key")) == global_key - ] - assert len(increment_calls) == 1 - call = increment_calls[0] - value = call.args[1] if len(call.args) >= 2 else call.kwargs.get("value") - assert value == 5.0 - assert call.kwargs.get("refresh_ttl") is True + init_and_increment.assert_awaited_once_with( + counter_key="spend:user:litellm-proxy-budget", + source_cache_key="litellm-proxy-budget", + increment=5.0, + ) finally: setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache) + setattr(litellm, "max_budget", original_max_budget) @pytest.mark.asyncio