fix(proxy): make spend counter reseed idempotent across pods via SET NX - #28590
Conversation
SpendCounterReseed.coalesced() re-seeded the Redis spend counter from DB using async_increment() (Redis INCRBYFLOAT), which is additive, not idempotent. The per-process asyncio.Lock added in BerriAI#26459 collapses duplicate reseeds within a single pod, but provides no cross-pod coordination. With N proxy pods racing the same cold or expired counter, Redis ends up at approximately N * db_spend and subsequent get_current_spend() calls in user_api_key_auth -> common_checks raise BudgetExceededError on entities well under their configured limit. Switch the reseed write to async_set_cache(..., nx=True) and, on a lost race, read back the winner's value. Mirrors the existing pattern already used by coalesced_window() in the same file. Single-pod behavior is unchanged (INCRBYFLOAT from absent is equivalent to SET); only the broken multi-pod path is fixed. Tests: - Adds test_coalesced_reseed_idempotent_under_concurrent_multi_pod_reseed simulating N pods with independent locks and a shared in-memory Redis to assert the counter converges to db_spend, not N * db_spend. - Adds test_coalesced_reseed_toctou_fallback_uses_increment_when_winner_key_vanishes covering the narrow window where SET NX returns False but the follow-up GET returns None (winner's key expired in between), so the fallback to async_increment still warms the cache. - Existing reseed test stubs are extended with the async_set_cache(nx=True) side-effect so they continue to model the new idempotent reseed path correctly. No assertion semantics weakened.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR fixes a multi-pod spend counter inflation bug where N pods racing to reseed a cold Redis counter each called
Confidence Score: 4/5Safe to merge — the core fix correctly prevents multi-pod spend inflation and the test coverage is thorough, with one minor residual edge case to be aware of. The SET NX approach is sound and mirrors the already-established pattern in coalesced_window(). The only remaining concern is the TOCTOU fallback: if the winner's key expires between a loser's failed NX and the loser's GET, multiple losers still fall through to INCRBYFLOAT, which can re-inflate the counter in a very narrow window. This is documented and astronomically unlikely in practice, but it is a remaining imperfection in the fix rather than a complete solution. The TOCTOU fallback block in litellm/proxy/db/spend_counter_reseed.py (lines 200–207) is the one area worth a second look.
|
| Filename | Overview |
|---|---|
| litellm/proxy/db/spend_counter_reseed.py | Switches the reseed write from INCRBYFLOAT to SET NX to prevent multi-pod spend inflation; the TOCTOU fallback back to INCRBYFLOAT is very narrow but still theoretically reachable |
| tests/test_litellm/proxy/test_proxy_server.py | Test stubs updated to model the new SET NX path; existing assertions restructured (not weakened) and two new tests added for multi-pod race and the TOCTOU edge case |
Reviews (1): Last reviewed commit: "fix(proxy): make spend counter reseed id..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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, | ||
| ) | ||
| ) |
There was a problem hiding this comment.
When N pods all lose the SET NX race (key was set by another pod) and then the winner's key expires before each loser's GET returns, all N losers fall into the async_increment(db_spend) branch. Since INCRBYFLOAT on a non-existent key starts from 0, the final counter ends up at (N-1) * db_spend. This is identical to the race the fix is designed to prevent, just in a much narrower window — the key must expire within the round-trip of a single Redis call. The PR description acknowledges this, and the existing coalesced_window() already uses the same pattern, so this is a deliberate trade-off rather than a new regression.
|
🤖 litellm-agent: This PR is currently BLOCKED from merge. Score: 3/5 ❌ Why blocked:
Details: Score docked for: 1 PR-related CI failure (This PR will be auto-closed as it lacks a screenshot for proof of fix. Please include one in the PR description. Add the Fix the issues above and push an update — the bot will re-review automatically.
|
|
Closing as superseded by #27854, which was merged into
Closing to avoid duplicate review. |
1 participant
Reopens #28248 (auto-closed when base branch `shin_agent_oss_staging_05_19_2026` was deleted after that staging branch merged into main). Same diff, rebased onto stable `litellm_oss_staging`.
Summary
Multi-pod proxy deployments could over-inflate spend counters: when a cold Redis counter expired, N pods racing to reseed all called `INCRBYFLOAT +db_spend`, producing `N × db_spend`. The fix uses Redis `SET NX` so only the first pod seeds the cached value; everyone else falls back to the normal increment path.
Test plan