From 62eeba6c106b1eadf32e4f809269d1ef77ba9ff1 Mon Sep 17 00:00:00 2001 From: lostmartian Date: Sat, 15 Aug 2026 00:49:18 +0530 Subject: [PATCH] fix(proxy): increment global proxy spend only when cache key exists Addresses review feedback on #36953: - veria (High): a plain INCRBYFLOAT recreates a missing global spend key from zero, so an evicted/expired entry loses the current-period total and continued traffic preserves the underestimate. Add an atomic increment-if-present Lua operation to RedisCache and use it for the global proxy spend update, leaving an absent key missing so the auth-time loader restores the authoritative DB spend. - greptile (P1): the increment refreshed a stale worker-local value because DualCache.async_increment_cache accrues in-memory. Mirror the authoritative Redis result into the worker's in-memory tier instead, matching the spend-counter pattern, so budget checks never read a stale pre-reset value. Adds unit tests for increment-if-present and update_cache Redis-path behavior. --- litellm/caching/redis_cache.py | 35 ++++++++++ litellm/proxy/proxy_server.py | 35 +++++++--- .../test_litellm/caching/test_redis_cache.py | 42 ++++++++++++ .../proxy/proxy_server/test_spend_counters.py | 68 +++++++++++++++++++ 4 files changed, 172 insertions(+), 8 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 5fedfc5bcce..41e2ff94b13 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1033,6 +1033,41 @@ 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: int | 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/proxy_server.py b/litellm/proxy/proxy_server.py index 0f47478b7c5..1b3195c0a68 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3126,14 +3126,33 @@ async def update_cache( 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, - ) + # 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. + if user_api_key_cache.redis_cache is not None: + current = await user_api_key_cache.redis_cache.async_increment_if_exists( + key=GLOBAL_PROXY_SPEND_CACHE_KEY, + value=response_cost, + ttl=get_management_object_ttl(user_api_key_cache), + ) + 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( + 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. " diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 59200719197..2cd207cd6bd 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -435,6 +435,48 @@ 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_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/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 81919825a56..8bf29be240e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1332,3 +1332,71 @@ async def test_update_cache_global_proxy_spend_concurrent_increments_after_reset 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_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 + stale pre-reset value, so budget checks do not reject off stale memory.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + + redis_cache = MagicMock() + redis_cache.async_increment_if_exists = 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-mirror", + value={"user_id": "user-mirror", "spend": 0.0}, + ) + await ps.update_cache( + token=None, + user_id="user-mirror", + end_user_id=None, + team_id=None, + response_cost=1.5, + parent_otel_span=None, + ) + + redis_cache.async_increment_if_exists.assert_awaited_once_with( + key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, + value=1.5, + ttl=ps.get_management_object_ttl(cache), + ) + assert await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) == 3.5