From c621f58fffbb304960938cdeff009db3ded456f4 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Wed, 20 May 2026 20:57:08 +0300 Subject: [PATCH] fix(spend_counter): seed Redis counter via SET NX to prevent cross-pod double-seed (#27854) * fix(spend_counter): seed Redis counter via SET NX to prevent cross-pod double-seed Symptom ------- Customers on multi-pod deployments see team `spend` jump to ~2x (or N x the pod count) shortly after a Redis cache miss / TTL expiry, triggering spurious "Budget Crossed" alerts and blocked requests until the value is manually reset. Root cause ---------- `SpendCounterReseed.coalesced` warmed the primary spend counter by calling `redis.async_increment(key, value=db_spend, refresh_ttl=True)`, which lowers to Redis `INCRBYFLOAT`. That is additive, not idempotent. The per-counter `asyncio.Lock` only coalesces seeders inside one process. With N pods sharing one Redis, on a cold key (cold start, TTL expiry, manual delete) every pod independently passes its lock + Redis re-check, reads the same `db_spend`, and issues `INCRBYFLOAT db_spend`. Final value: N x db_spend. Fix --- Use `redis.async_set_cache(key, value=db_spend, nx=True)` for the seed. SET NX is atomic across pods: exactly one writer initializes the key; losers read the winner's value via `async_get_cache`. This is the same idiom already used by `coalesced_window` in the same file, so the two seed paths are now consistent. Per-request deltas continue to use `INCRBYFLOAT` (correct - additive behaviour is what we want for increments, not for initial seed). Verification ------------ Live two-process repro against the same Postgres + Redis (DB spend = 506): Unpatched: 4/4 runs -> Redis counter = ~1012 (~2 x db_spend) Patched: 12/12 runs -> Redis counter = ~506 Unit tests (`test_proxy_server.py`): - New `test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed` patches `_get_lock` to return a fresh lock per caller (otherwise the per-process lock masks the race), races two `coalesced` calls, and asserts final = 506 with exactly one of two SET NX attempts winning. - 4 existing tests updated for the new seed contract (SET NX for the seed, INCRBYFLOAT only for the per-request delta). - Full `spend_counter or reseed or budget` slice: 22 passed. Co-authored-by: Cursor * test(spend_counter): make SET NX mock atomic so loser branch is exercised Greptile flagged that `redis_set_cache` in test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed placed `await asyncio.sleep(0)` AFTER the NX membership check. Both concurrent tasks observed an empty `redis_store`, passed the guard, and both returned True - so the loser branch (else: read back winner's value) was never exercised. Fix the mock to model real atomic Redis SET NX: - Yield BEFORE the membership check so two concurrent callers interleave the way real SET NX does (first to resume runs check + write atomically and wins; second resumes after the key exists and loses). - Track set_cache return values; assert sorted([loser, winner]) so we know exactly one task wins and one loses. - Track async_get_cache calls that happen AFTER at least one SET NX has completed; assert at least one such read - that is the loser-path fallback (`current_value = float(cached)` when seeded is False). Verified by temporarily reverting the mock to the old order: the test now fails with `expected exactly one SET NX winner and one loser, got [True, True]`, exactly the failure mode Greptile described. No production code change. Co-authored-by: Cursor * test(spend_counter): mock async_set_cache to populate redis_store in concurrent read+write test `test_concurrent_read_and_write_paths_share_one_db_query` mocks `async_increment` to populate the in-memory `redis_store`, but did not mock `async_set_cache`. After the SET-NX seed change in `coalesced()`, the seed step writes via `async_set_cache(nx=True)` (default AsyncMock, no `redis_store` write), so the simulated Redis stays empty after the first reseed. The second `get_current_spend` then sees a clean Redis miss, re-enters the DB read path, and the test fails with `expected 1 DB query, got 2`. Fix: add a `redis_set_cache` side_effect that updates `redis_store` on `nx=True` (and rejects when the key already exists), matching the pattern used by the four sibling tests fixed in this branch's first commit. Pre-existing assertions are unchanged. Full `tests/test_litellm/proxy/test_proxy_server.py`: 158 passed. Co-authored-by: Cursor --------- Co-authored-by: Cursor (cherry picked from commit 0fb710400f80088fdcc1e382d3efa90a7ec895ea) --- litellm/proxy/db/spend_counter_reseed.py | 27 ++- tests/test_litellm/proxy/test_proxy_server.py | 167 ++++++++++++++++-- 2 files changed, 176 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 19ec669939..e7c5fa3f72 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -178,15 +178,28 @@ class SpendCounterReseed: if db_spend is None: return None # Warm even when 0 so subsequent reads hit cache, not DB. + # + # Seed via SET NX (cross-pod safe): only one pod initializes the + # Redis key with db_spend; concurrent seeders read the winner's + # value. INCRBYFLOAT-of-db_spend from N pods would multiply the + # counter (N x db_spend) and trigger spurious budget alerts. + current_value: float = float(db_spend) 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 = float(db_spend) + else: + cached = await spend_counter_cache.redis_cache.async_get_cache( + key=counter_key + ) + current_value = ( + float(cached) if cached is not None else float(db_spend) + ) spend_counter_cache.in_memory_cache.set_cache( key=counter_key, value=current_value, @@ -202,7 +215,7 @@ class SpendCounterReseed: ) if require_cache_warm: raise - return db_spend + return current_value @staticmethod async def window_from_spend_logs( diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e66dbcc349..f094600c43 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5098,6 +5098,7 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( 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(return_value=True) # SET NX wins counter_cache.redis_cache = fake_redis # Prisma returns spend=42.0 (authoritative) while the stale cached @@ -5134,16 +5135,131 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( 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). + # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. + # Only the per-request delta (1.5) goes through INCRBYFLOAT. + fake_redis.async_set_cache.assert_awaited_once_with( + key="spend:team:team-9", value=42.0, nx=True + ) 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 + assert writes == [("spend:team:team-9", 1.5)] finally: ps.user_api_key_cache = orig_user ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma +@pytest.mark.asyncio +async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed(): + """Two pods both observing a missing Redis counter must not both + INCRBYFLOAT the full DB spend. SpendCounterReseed.coalesced uses SET NX + so the loser reads the winner's value; final Redis = db_spend, not + 2 * db_spend. + + The per-counter asyncio.Lock is per-process, so it does NOT coordinate + across pods. We simulate two pods by patching _get_lock to return a + fresh lock per call (each "pod" has its own lock registry in real life). + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + counter_key = "spend:team:team-concurrent-seed" + redis_store: dict = {} + db_read_count = 0 + set_results: list = [] + get_after_set_count = 0 + set_completed_count = 0 + + async def redis_set_cache(key, value, nx=False, **_): + # Yield BEFORE the membership check so two concurrent callers + # interleave the way real atomic Redis SET NX does: the first + # to resume runs check + write atomically and wins; the second + # resumes after the key exists and loses. Yielding *after* the + # check would let both callers pass the empty-store check before + # either writes, so neither would ever lose. + await asyncio.sleep(0) + if nx and key in redis_store: + set_results.append(False) + return False + redis_store[key] = float(value) + set_results.append(True) + nonlocal set_completed_count + set_completed_count += 1 + return True + + async def redis_get_cache(key): + # Track reads that happen after at least one SET NX has completed + # — those are the loser-path fallback reads we want to verify. + if set_completed_count > 0: + nonlocal get_after_set_count + get_after_set_count += 1 + return redis_store.get(key) + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) + + async def slow_find_unique(**_): + nonlocal db_read_count + db_read_count += 1 + # Both pods read DB before either's SET NX lands. + await asyncio.sleep(0) + row = MagicMock() + row.spend = 506.0 + return row + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teamtable.find_unique = AsyncMock( + side_effect=slow_find_unique + ) + + pod_a = DualCache() + pod_a.redis_cache = fake_redis + pod_b = DualCache() + pod_b.redis_cache = fake_redis + + # Each "pod" has its own per-process lock registry. Patch _get_lock to + # always return a fresh lock so the two coalesced calls do not serialize + # via one in-process lock (which is what would happen across pods). + async def fresh_lock(_counter_key): + return asyncio.Lock() + + with patch.object(SpendCounterReseed, "_get_lock", side_effect=fresh_lock): + results = await asyncio.gather( + SpendCounterReseed.coalesced( + prisma_client=fake_prisma, + spend_counter_cache=pod_a, + counter_key=counter_key, + ), + SpendCounterReseed.coalesced( + prisma_client=fake_prisma, + spend_counter_cache=pod_b, + counter_key=counter_key, + ), + ) + + assert all(r == 506.0 for r in results), results + assert redis_store[counter_key] == pytest.approx(506.0), redis_store + # Both pods read the DB and both attempted SET NX; exactly one wrote + # (winner) and one was rejected (loser). + assert db_read_count == 2 + assert fake_redis.async_set_cache.await_count == 2 + nx_writes = [ + call + for call in fake_redis.async_set_cache.await_args_list + if call.kwargs.get("nx") is True + ] + assert len(nx_writes) == 2 + assert sorted(set_results) == [False, True], ( + f"expected exactly one SET NX winner and one loser, got {set_results}" + ) + # Loser path executed: after the winner's SET NX returned True, the + # losing coalesced() call falls back to async_get_cache to read the + # winner's value rather than re-seeding. + assert get_after_set_count >= 1, ( + "loser branch (else: read back winner's value) was never exercised" + ) + + @pytest.mark.asyncio async def test_reseed_spend_from_db_user_and_org_prefixes(): """User and org counters reseed from their own DB tables. @@ -5267,9 +5383,16 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis db_row = MagicMock() @@ -5297,6 +5420,7 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( where={"team_id": "team-stale-local"} ) + # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. assert redis_store[counter_key] == pytest.approx(43.5) assert counter_cache.in_memory_cache.get_cache( key=counter_key @@ -5687,14 +5811,14 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): from litellm.proxy.proxy_server import get_current_spend counter_cache = DualCache() - recorded_warms: list = [] + recorded_seeds: list = [] - async def record_increment(key, value, ttl=None, **kwargs): - recorded_warms.append({"key": key, "value": value}) - return value + async def record_set_cache(key, value, nx=False, **kwargs): + recorded_seeds.append({"key": key, "value": value, "nx": nx}) + return True fake_redis = AsyncMock() - fake_redis.async_increment = AsyncMock(side_effect=record_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=record_set_cache) fake_redis.async_get_cache = AsyncMock(return_value=None) counter_cache.redis_cache = fake_redis @@ -5719,9 +5843,9 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): 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 - assert ("spend:team_member:user-1:team-1", 362.0) in [ - (w["key"], w["value"]) for w in recorded_warms + # Counter warmed via SET NX so subsequent reads are fast. + assert ("spend:team_member:user-1:team-1", 362.0, True) in [ + (s["key"], s["value"], s["nx"]) for s in recorded_seeds ] assert counter_cache.in_memory_cache.get_cache( key="spend:team_member:user-1:team-1" @@ -5798,8 +5922,15 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -5906,9 +6037,16 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -6011,9 +6149,16 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis db_call_count = 0