mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-04 23:11:19 +00:00
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 <cursoragent@cursor.com> * 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 <cursoragent@cursor.com> * 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 <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
milan-berri
@@ -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(
|
||||
|
||||
@@ -5708,6 +5708,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
|
||||
@@ -5744,16 +5745,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.
|
||||
@@ -5877,9 +5993,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()
|
||||
@@ -5907,6 +6030,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
|
||||
@@ -6297,14 +6421,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
|
||||
|
||||
@@ -6329,9 +6453,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"
|
||||
@@ -6408,8 +6532,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()
|
||||
@@ -6516,9 +6647,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()
|
||||
@@ -6621,9 +6759,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
|
||||
|
||||
Reference in New Issue
Block a user