diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 19ec66993908..bd3f192ab432 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -178,15 +178,35 @@ async def coalesced( if db_spend is None: return None # Warm even when 0 so subsequent reads hit cache, not DB. + # Use SET NX rather than INCRBYFLOAT so concurrent reseeds across + # pods (where the per-process singleflight lock above does not + # coordinate) are idempotent: only the first writer seeds db_spend, + # losers read the winner's value. Mirrors coalesced_window() below. try: if spend_counter_cache.redis_cache is not None: - current_value = ( - await spend_counter_cache.redis_cache.async_increment( - key=counter_key, - value=db_spend, - refresh_ttl=True, - ) + seeded = await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, + value=db_spend, + nx=True, ) + if seeded: + current_value = db_spend + else: + current_cached_value = ( + await spend_counter_cache.redis_cache.async_get_cache( + key=counter_key + ) + ) + if current_cached_value is None: + current_value = ( + await spend_counter_cache.redis_cache.async_increment( + key=counter_key, + value=db_spend, + refresh_ttl=True, + ) + ) + else: + current_value = float(current_cached_value) spend_counter_cache.in_memory_cache.set_cache( key=counter_key, value=current_value, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6d10d2a63533..6f7097e248f0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5699,15 +5699,30 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( from litellm.caching.dual_cache import DualCache counter_cache = DualCache() + redis_store: dict = {} + recorded_seeds: list = [] recorded_increments: list = [] + async def record_set_cache(key, value, **kwargs): + nx = kwargs.get("nx", False) + if nx and key in redis_store: + return False + redis_store[key] = value + recorded_seeds.append({"key": key, "value": value}) + return True + async def record_increment(key, value, ttl=None, **kwargs): + redis_store[key] = (redis_store.get(key) or 0.0) + value recorded_increments.append({"key": key, "value": value, "ttl": ttl}) - return value + return redis_store[key] + + async def redis_get(key, **_): + return redis_store.get(key) fake_redis = AsyncMock() fake_redis.async_increment = AsyncMock(side_effect=record_increment) - fake_redis.async_get_cache = AsyncMock(return_value=None) # counter missing + fake_redis.async_set_cache = AsyncMock(side_effect=record_set_cache) + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) counter_cache.redis_cache = fake_redis # Prisma returns spend=42.0 (authoritative) while the stale cached @@ -5744,10 +5759,14 @@ async def record_increment(key, value, ttl=None, **kwargs): fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( where={"team_id": "team-9"} ) - # Two increments keyed on the counter: seed ($42) then request ($1.50). - writes = [(c["key"], c["value"]) for c in recorded_increments] - assert ("spend:team:team-9", 42.0) in writes - assert ("spend:team:team-9", 1.5) in writes + # Seed write goes through SET NX (idempotent across pods); the + # request-cost write goes through INCRBYFLOAT. + assert ("spend:team:team-9", 42.0) in [ + (s["key"], s["value"]) for s in recorded_seeds + ] + assert ("spend:team:team-9", 1.5) in [ + (c["key"], c["value"]) for c in recorded_increments + ] finally: ps.user_api_key_cache = orig_user ps.spend_counter_cache = orig_counter @@ -5873,12 +5892,23 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): redis_store: dict = {} + async def redis_get(key, **_): + return redis_store.get(key) + + async def redis_set_cache(key, value, **kwargs): + nx = kwargs.get("nx", False) + if nx and key in redis_store: + return False + redis_store[key] = value + return True + async def redis_increment(key, value, **_): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] fake_redis = AsyncMock() - fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) counter_cache.redis_cache = fake_redis @@ -6297,15 +6327,23 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): from litellm.proxy.proxy_server import get_current_spend counter_cache = DualCache() + redis_store: dict = {} recorded_warms: list = [] - async def record_increment(key, value, ttl=None, **kwargs): + async def redis_get(key, **_): + return redis_store.get(key) + + async def record_set_cache(key, value, **kwargs): + nx = kwargs.get("nx", False) + if nx and key in redis_store: + return False + redis_store[key] = value recorded_warms.append({"key": key, "value": value}) - return value + return True fake_redis = AsyncMock() - fake_redis.async_increment = AsyncMock(side_effect=record_increment) - fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) + fake_redis.async_set_cache = AsyncMock(side_effect=record_set_cache) counter_cache.redis_cache = fake_redis # DB has authoritative spend=362.0; caller hands us stale fallback=30.0 @@ -6329,7 +6367,7 @@ async def record_increment(key, value, ttl=None, **kwargs): f"expected DB reseed to return 362.0, got {spend} " f"(fallback would have returned 30.0 and caused bypass)" ) - # Counter warmed so subsequent reads are fast + # Counter warmed via SET NX so subsequent reads are fast assert ("spend:team_member:user-1:team-1", 362.0) in [ (w["key"], w["value"]) for w in recorded_warms ] @@ -6404,11 +6442,19 @@ async def slow_find_unique(**kwargs): async def redis_get(key, **_): return redis_store.get(key) + async def redis_set_cache(key, value, **kwargs): + nx = kwargs.get("nx", False) + if nx and key in redis_store: + return False + redis_store[key] = value + return True + async def redis_increment(key, value, **_): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) counter_cache.redis_cache = fake_redis @@ -6438,6 +6484,135 @@ async def redis_increment(key, value, **_): ps.prisma_client = orig_prisma +@pytest.mark.asyncio +async def test_coalesced_reseed_idempotent_under_concurrent_multi_pod_reseed( + monkeypatch, +): + """ + Regression: multiple pods racing to reseed the same Redis spend counter + must not multiply db_spend. Before this fix, the reseed used INCRBYFLOAT + and produced N * db_spend across N pods because the per-process + asyncio.Lock added for singleflight provides no cross-pod coordination. + The reseed now uses SET NX so only the first writer seeds db_spend and + losers read the winner's value, mirroring coalesced_window(). + + Simulates N pods by giving each concurrent caller its own asyncio.Lock + (defeating the per-process singleflight, which does not exist across pods + in real deployments). + """ + import asyncio as _asyncio + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + # 1) Defeat the per-process singleflight so each caller models a distinct pod. + async def _fresh_lock(_counter_key): + return _asyncio.Lock() + + monkeypatch.setattr(SpendCounterReseed, "_get_lock", staticmethod(_fresh_lock)) + + # 2) Shared "Redis" backed by a dict + internal lock (single Redis instance + # shared across pods). + class _SharedFakeRedis: + def __init__(self) -> None: + self.store: dict = {} + self._lock = _asyncio.Lock() + + async def async_get_cache(self, key, **_kwargs): + async with self._lock: + return self.store.get(key) + + async def async_set_cache(self, key, value, **kwargs): + nx = kwargs.get("nx", False) + async with self._lock: + if nx and key in self.store: + return False + self.store[key] = value + return True + + async def async_increment(self, key, value, **_kwargs): + async with self._lock: + self.store[key] = float(self.store.get(key, 0.0)) + float(value) + return self.store[key] + + fake_redis = _SharedFakeRedis() + spend_cache = DualCache() + spend_cache.redis_cache = fake_redis # type: ignore[assignment] + + # 3) DB always reports the same authoritative spend. + db_spend = 100.0 + + async def _fake_from_db(_prisma, _counter_key): + return db_spend + + monkeypatch.setattr(SpendCounterReseed, "from_db", staticmethod(_fake_from_db)) + + counter_key = "spend:user:test-user-multi-pod" + n_pods = 5 + results = await _asyncio.gather( + *[ + SpendCounterReseed.coalesced( + prisma_client=None, + spend_counter_cache=spend_cache, + counter_key=counter_key, + ) + for _ in range(n_pods) + ] + ) + + assert all( + r == db_spend for r in results + ), f"all callers should see db_spend={db_spend}, got {results}" + assert fake_redis.store[counter_key] == db_spend, ( + f"Redis counter must equal db_spend after concurrent reseed, " + f"got {fake_redis.store[counter_key]} (would be " + f"{n_pods * db_spend} under non-idempotent INCRBYFLOAT)" + ) + + +@pytest.mark.asyncio +async def test_coalesced_reseed_toctou_fallback_uses_increment_when_winner_key_vanishes(): + """ + Edge case: SET NX returns False (another pod seeded first), but the + follow-up GET returns None (the winner's key already expired between + NX and GET). In this narrow window, the loser falls back to + INCRBYFLOAT so the cache is still warmed. Matches coalesced_window()'s + behaviour for the same TOCTOU window. + """ + from unittest.mock import AsyncMock + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_set_cache = AsyncMock(return_value=False) + fake_redis.async_increment = AsyncMock(return_value=42.0) + + spend_cache = DualCache() + spend_cache.redis_cache = fake_redis # type: ignore[assignment] + + async def _fake_from_db(_prisma, _counter_key): + return 42.0 + + monkeypatch_target = SpendCounterReseed + original_from_db = monkeypatch_target.from_db + monkeypatch_target.from_db = staticmethod(_fake_from_db) # type: ignore[method-assign] + try: + result = await SpendCounterReseed.coalesced( + prisma_client=None, + spend_counter_cache=spend_cache, + counter_key="spend:user:test-toctou", + ) + finally: + monkeypatch_target.from_db = staticmethod(original_from_db) # type: ignore[method-assign] + + assert result == 42.0 + fake_redis.async_set_cache.assert_awaited_once() + fake_redis.async_get_cache.assert_awaited() + fake_redis.async_increment.assert_awaited_once() + + @pytest.mark.asyncio async def test_get_current_spend_uses_db_zero_over_stale_fallback(): """ @@ -6512,12 +6687,20 @@ async def slow_find_unique(**kwargs): async def redis_get(key, **_): return redis_store.get(key) + async def redis_set_cache(key, value, **kwargs): + nx = kwargs.get("nx", False) + if nx and key in redis_store: + return False + redis_store[key] = value + return True + async def redis_increment(key, value, **_): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) counter_cache.redis_cache = fake_redis @@ -6617,12 +6800,20 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): async def redis_get(key, **_): return redis_store.get(key) + async def redis_set_cache(key, value, **kwargs): + nx = kwargs.get("nx", False) + if nx and key in redis_store: + return False + redis_store[key] = value + return True + async def redis_increment(key, value, **_): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) counter_cache.redis_cache = fake_redis