diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bf760a92d88..d00614cc8ed 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -262,11 +262,21 @@ 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 load reads the - zeroed row instead of a stale (potentially never-expiring) counter. + """Zero the cached global-proxy spend accumulator after the proxy + budget aggregate row is reset so atomic increments start from a known floor. """ - await ResetBudgetJob._invalidate_user_api_key_cache_entry(GLOBAL_PROXY_SPEND_CACHE_KEY) + try: + from litellm.proxy.proxy_server import user_api_key_cache + + await user_api_key_cache.async_set_cache( + key=GLOBAL_PROXY_SPEND_CACHE_KEY, + value=0.0, + ) + except Exception as e: # noqa: BLE001 # cache zero failure must not break budget reset + verbose_proxy_logger.warning( + "Failed to zero global proxy spend cache after budget reset: %s", + e, + ) @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 b1b5a7ffbe5..0f47478b7c5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3079,7 +3079,7 @@ async def update_cache( ### UPDATE USER SPEND ### async def _update_user_cache(): - ## UPDATE CACHE FOR USER ID + GLOBAL PROXY + ## UPDATE CACHE FOR USER ID if response_cost is None: return user_ids: Final = [user_id] @@ -3110,14 +3110,6 @@ async def update_cache( CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_UserTable), ) ) - ## UPDATE GLOBAL PROXY ## - global_proxy_spend: Final = 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: Final = 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. " @@ -3129,6 +3121,29 @@ async def update_cache( traceback.format_exc(), ) + ### UPDATE GLOBAL PROXY SPEND ### + async def _update_global_proxy_spend() -> None: + if response_cost is None: + return + try: + # Atomic increment avoids the read-modify-write race with + # ResetBudgetJob and starts from 0 when the key is missing. + 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, + ) + 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. " + "Budget enforcement may use stale spend values. " + "response_cost=%s - %s\n%s", + response_cost, + str(e), + traceback.format_exc(), + ) + ### UPDATE END-USER SPEND ### async def _update_end_user_cache(): if end_user_id is None or response_cost is None: @@ -3272,6 +3287,8 @@ async def update_cache( if user_id is not None: await _update_user_cache() + await _update_global_proxy_spend() + if end_user_id is not None: await _update_end_user_cache() @@ -3281,27 +3298,15 @@ async def update_cache( if tags is not None: await _update_tag_cache() - global_proxy_spend_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY - local_object_updates: Final = tuple((k, v) for k, v in values_to_update_in_cache if k != global_proxy_spend_key) - shared_scalar_updates: Final = 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=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(): 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 608dc8cb5c8..a0e9c157025 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 @@ -1030,6 +1030,7 @@ def _make_counter_invalidation_job(monkeypatch): user_api_key_cache = MagicMock() user_api_key_cache.async_delete_cache = AsyncMock() + user_api_key_cache.async_set_cache = AsyncMock() fake_module = types.ModuleType("litellm.proxy.proxy_server") fake_module.spend_counter_cache = spend_counter_cache @@ -1088,14 +1089,11 @@ 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_resets_global_spend_cache( 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.""" + """Resetting the proxy-wide budget aggregate row zeroes the cached + global-spend accumulator so the next period starts from a known floor.""" counter_cache = _make_counter_invalidation_job(monkeypatch) now = datetime.now(timezone.utc) @@ -1115,13 +1113,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.user_api_key_cache.async_set_cache.assert_any_call( + key="default_user_id:spend", value=0.0 + ) 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 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) @@ -1146,6 +1146,10 @@ def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache( call.kwargs.get("key") == "default_user_id:spend" for call in counter_cache.user_api_key_cache.async_delete_cache.call_args_list ) + assert not any( + call.kwargs.get("key") == "default_user_id:spend" + for call in counter_cache.user_api_key_cache.async_set_cache.call_args_list + ) def test_reset_budget_for_teams_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): 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..81919825a56 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1270,3 +1270,65 @@ async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkey ) assert result is None + + +@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.""" + 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=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) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 918d39646b0..9cb5e4da286 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5278,19 +5278,17 @@ async def test_update_cache_global_proxy_spend_scalar_stays_shared(): ``{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. + cache TTL. The scalar is updated via an atomic cache increment, which + remains shared across pods and avoids resurrecting a stale pre-reset value. """ from litellm.caching.caching import DualCache admin_name = litellm.proxy.proxy_server.litellm_proxy_admin_name - global_key = "{}:spend".format(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} - if key == global_key: - return 10.0 return None original_cache = litellm.proxy.proxy_server.user_api_key_cache @@ -5299,27 +5297,39 @@ async def test_update_cache_global_proxy_spend_scalar_stays_shared(): 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, - ) + 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, + ) - pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] - if pending: - await asyncio.wait(pending, timeout=5) + 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] + 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 finally: setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)