diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index d0e8f08eef9..5fedfc5bcce 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1033,41 +1033,6 @@ class RedisCache(BaseCache): result = result.decode() return float(result) - @_redis_circuit_breaker_guard - async def async_increment_if_exists( - self, - key: str, - value: float, - ttl: float | None = None, - ) -> float | None: - """Atomically increment ``key`` by ``value`` only when the key already - exists, returning the new value; an absent key is left missing and - ``None`` is returned so callers can fall back to loading the - authoritative total. The GET/increment/SET runs in a single Lua call, - so it is race-free across callers and pods, and refreshing the TTL on - a present key keeps the accumulator alive while it is being used. - """ - _redis_client: Final = self.init_async_client() - _used_ttl: Final = self.get_ttl(ttl=ttl) - key = self.check_and_fix_namespace(key=key) - lua: Final = ( - "local cur = redis.call('GET', KEYS[1]) " - "if cur == false then return nil end " - "local new = tonumber(cur) + tonumber(ARGV[1]) " - "redis.call('SET', KEYS[1], new) " - "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " - "return new" - ) - result = cast( # cast-ok: redis eval returns an untyped scalar - "str | bytes | int | float | None", - await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))), - ) - if result is None: - return None - if isinstance(result, bytes): - result = result.decode() - return float(result) - async def flush_cache_buffer(self): print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}") await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a850c50248c..ae4bd24ebf8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -494,6 +494,10 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use # 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 + async def _fetch_global_spend_with_event_coordination( cache_key: str, @@ -522,14 +526,49 @@ async def _fetch_global_spend_with_event_coordination( ) 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:{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 ) - return await _global_spend_coordinator.get_or_load( + 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: + return None + + 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: + await spend_cache.async_set_cache(key=cache_key, value=db_spend) + numeric_value = db_spend + + return numeric_value async def get_global_proxy_spend( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1b3195c0a68..5cbdbaaad67 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3126,25 +3126,21 @@ async def update_cache( if response_cost is None: return try: - # Atomically accumulate the shared global-proxy scalar without the - # read-modify-write race with ResetBudgetJob. A plain INCRBYFLOAT - # would recreate a missing key from zero, so increment-if-present - # leaves an absent key missing and lets the auth-time loader - # restore the authoritative DB spend instead. + # Atomically accumulate the shared global-proxy scalar in Redis. + # The returned value is authoritative across workers, so mirror it + # into this worker's in-memory tier to avoid accumulating a stale + # pre-reset local value. if user_api_key_cache.redis_cache is not None: - current = await user_api_key_cache.redis_cache.async_increment_if_exists( + current = 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, + ) + await user_api_key_cache.in_memory_cache.async_set_cache( + key=GLOBAL_PROXY_SPEND_CACHE_KEY, + value=current, ) - if current is not None: - # Mirror the authoritative Redis value into this worker's - # in-memory tier so a stale pre-reset value is never - # accumulated or read back. - await user_api_key_cache.in_memory_cache.async_set_cache( - key=GLOBAL_PROXY_SPEND_CACHE_KEY, - value=current, - ) else: # No Redis: single-process, so an in-memory increment is coherent. await user_api_key_cache.async_increment_cache( diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 5e4a83cfe46..59200719197 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -435,66 +435,6 @@ def test_increment_cache_namespaces_key( mock_client.incr.assert_called_once_with(name=expected, amount=1) -@pytest.mark.asyncio -async def test_async_increment_if_exists_increments_present_key( - monkeypatch, redis_no_ping -): - """A present key is atomically incremented and its TTL refreshed.""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace=None) - mock_redis_instance = AsyncMock() - mock_redis_instance.eval = AsyncMock(return_value="5.0") - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - result = await redis_cache.async_increment_if_exists(key="k", value=2.0, ttl=60) - - assert result == 5.0 - lua, num_keys, key, value, ttl = mock_redis_instance.eval.await_args.args - assert num_keys == 1 - assert key == "k" - assert value == "2.0" - assert ttl == "60" - assert "if cur == false then return nil end" in lua - - -@pytest.mark.asyncio -async def test_async_increment_if_exists_decodes_bytes_result( - monkeypatch, redis_no_ping -): - """Redis returns bytes; the method decodes them to a float.""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace=None) - mock_redis_instance = AsyncMock() - mock_redis_instance.eval = AsyncMock(return_value=b"7.0") - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - result = await redis_cache.async_increment_if_exists(key="k", value=2.0) - - assert result == 7.0 - - -@pytest.mark.asyncio -async def test_async_increment_if_exists_leaves_missing_key_absent( - monkeypatch, redis_no_ping -): - """An absent key stays missing: the method returns None and does not create it.""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace=None) - mock_redis_instance = AsyncMock() - mock_redis_instance.eval = AsyncMock(return_value=None) - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - result = await redis_cache.async_increment_if_exists(key="k", value=2.0) - - assert result is None - - @pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) def test_delete_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") 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 1555bc34a06..ee8c8024001 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 @@ -6134,11 +6134,14 @@ async def test_global_proxy_spend_reads_shared_redis_scalar_when_configured(): 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 + proxy_budget_row = MagicMock() + proxy_budget_row.spend = 5.0 prisma_client = MagicMock() prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=AssertionError("a Redis hit must not hit the DB") + return_value=proxy_budget_row ) result = await _fetch_global_spend_with_event_coordination( @@ -6149,6 +6152,43 @@ async def test_global_proxy_spend_reads_shared_redis_scalar_when_configured(): 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() + + +@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, + ) + 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_cache = AsyncMock() + 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 + redis_cache.async_get_cache.assert_awaited_once_with(key="default_user_id:spend") + redis_cache.async_set_cache.assert_awaited_once_with( + key="default_user_id:spend", value=9.0 + ) @pytest.mark.asyncio 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 277e6e083fd..70a1d6f80cf 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1334,40 +1334,6 @@ async def test_update_cache_global_proxy_spend_concurrent_increments_after_reset assert final_value == pytest.approx(0.5) -@pytest.mark.asyncio -async def test_update_cache_global_proxy_spend_leaves_missing_key_absent(monkeypatch): - """An absent global spend key must not be recreated from zero by an - increment; the auth-time loader restores the authoritative DB spend.""" - from unittest.mock import AsyncMock, MagicMock - - from litellm.caching.caching import DualCache - - redis_cache = MagicMock() - redis_cache.async_increment_if_exists = AsyncMock(return_value=None) - 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-missing", - value={"user_id": "user-missing", "spend": 0.0}, - ) - await ps.update_cache( - token=None, - user_id="user-missing", - end_user_id=None, - team_id=None, - response_cost=1.0, - parent_otel_span=None, - ) - - redis_cache.async_increment_if_exists.assert_awaited_once_with( - key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, - value=1.0, - ttl=ps.get_management_object_ttl(cache), - ) - assert await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) is None - - @pytest.mark.asyncio async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monkeypatch): """In-memory mirrors the authoritative Redis result instead of accruing a @@ -1377,7 +1343,7 @@ async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monke from litellm.caching.caching import DualCache redis_cache = MagicMock() - redis_cache.async_increment_if_exists = AsyncMock(return_value=3.5) + 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) @@ -1394,10 +1360,11 @@ async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monke parent_otel_span=None, ) - redis_cache.async_increment_if_exists.assert_awaited_once_with( + 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, ) assert await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) == 3.5 @@ -1410,9 +1377,7 @@ async def test_update_cache_global_proxy_spend_swallows_redis_failure(monkeypatc from litellm.caching.caching import DualCache redis_cache = MagicMock() - redis_cache.async_increment_if_exists = AsyncMock( - side_effect=RuntimeError("redis down") - ) + 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) @@ -1429,8 +1394,9 @@ async def test_update_cache_global_proxy_spend_swallows_redis_failure(monkeypatc parent_otel_span=None, ) - redis_cache.async_increment_if_exists.assert_awaited_once_with( + 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, )