fix(proxy): reset a key's budget-window counters on spend reset - #38686
yassin-berriai merged 5 commits into
Conversation
|
|
…set cross-pod
/key/{id}/reset_spend already reset the key's lifetime spend counter in
Redis, but a key with its own budget_limits (an extra time-windowed cap,
e.g. a daily budget layered on top of the lifetime max_budget) kept its
window counter untouched, so the key stayed 429'd on
"ExceededBudget: Key over <duration> budget" even after the admin action
reported spend back to $0.
Force-expire each window on reset: zero its Redis counter and restart the
window from now (window_start is derived as reset_at - budget_duration,
so reset_at must float to now + duration, not the next calendar boundary
get_budget_reset_time gives key creation - that boundary can still be
in the past relative to the spend that triggered the block).
Also close a second, narrower race: _delete_cache_key_object evicted the
cached key object only on the handling pod, so another pod could keep
serving the stale pre-reset object (and re-derive the pre-reset spend
counter via its own floor-marker cache) until its own TTL expired. It now
broadcasts the eviction, matching the pattern already used for team,
team-member, customer, and tag caches.
Greptile SummaryThis PR makes key spend resets restart the key’s budget windows and propagate relevant cache and counter updates across proxy workers.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains.
|
| Filename | Overview |
|---|---|
| litellm/proxy/auth/auth_checks.py | Centralizes cross-worker key-object invalidation in the single-key cache deletion helper without introducing a blocking failure. |
| litellm/proxy/management_endpoints/key_management_endpoints.py | Persists restarted key budget windows, resets shared spend state, and evicts cached key objects after durable updates. |
| tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py | Adds mock-based regression coverage for window-counter resets, persisted boundaries, floor markers, and broadcast eviction. |
Reviews (5): Last reviewed commit: "fix(proxy): persist advanced budget-wind..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| """ | ||
| Set a Redis-backed spend counter to `value`, mirror it into the short-lived | ||
| spend_db_floor marker `_authoritative_floor_spend` reads, and broadcast both | ||
| to every worker (LIT-3803 pattern: setting, not deleting, means a worker's | ||
| own self-delivered broadcast still carries the reset value forward). | ||
|
|
||
| Without the floor marker, `_authoritative_floor_spend` can re-derive a | ||
| stale, pre-reset value from a marker another worker cached moments earlier | ||
| and raise the just-reset counter right back up via `_repair_stale_spend_counter`. | ||
| Without the broadcast, a worker that already cached the pre-reset key object | ||
| or floor marker keeps enforcing against it until its own TTL expires. | ||
| """ | ||
| from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache | ||
|
|
There was a problem hiding this comment.
The new helper docstrings repeat implementation details across several long blocks, increasing maintenance cost and becoming misleading when cache behavior changes
Context Used: CLAUDE.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
4448c42 to
cd7f877
Compare
Greptile P1: eviction ran before the window-reset DB write committed, so a request racing the reset could re-fetch and re-cache the pre-write row, pinning that pod to the stale budget_limits for the rest of its own cache TTL even after the write went through. Move the eviction to run last.
|
@greptileai Fixed: moved the cache eviction to run after the window-reset DB write. Please re-review the current head dd695e3. |
test_delete_cache_key_object_broadcasts_invalidation now asserts a real UserApiKeyCache no longer holds the evicted entry, rather than only inspecting a mock's call args. Suppress test-quality-ok on the hash_token/_check_proxy_or_team_admin_for_key/_delete_cache_key_object/ publish_auth_cache_invalidation patches: none has an HTTP boundary to fake, matching the pattern the file already uses for these same targets.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
isinstance(x, list) in the else branch still leaves Sequence[object] | str (a tuple satisfies Sequence without being a list), so json.loads() saw a possible non-str argument. Check isinstance(x, str) instead, which narrows each branch to exactly the type it needs.
|
@greptileai Pushed test-quality-gate fixes and a basedpyright type-narrowing fix since the last review. Please re-review the current head 05fda30. |
…counters Greptile P1: publishing a zeroed window counter before the new reset_at committed let a request racing the write compute window_start from the stale boundary, re-sum the historical spend log rows the reset was clearing, and put the counter right back above budget. Compute every window's new boundary, persist all of them in one DB write, then zero each window's Redis counter only once that write has landed.
|
@greptileai Fixed: window boundaries now persist before their counters zero, closing the race. Please re-review the current head 5720281. |
| await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=value, ttl=60) | ||
| except Exception as redis_err: | ||
| verbose_proxy_logger.warning( | ||
| "Failed to update spend counter %s in Redis: %s. " | ||
| "Budget checks may use stale value until counter expires.", | ||
| counter_key, | ||
| redis_err, | ||
| ) | ||
|
|
||
| floor_key: Final = f"spend_db_floor:{counter_key}" | ||
| spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) | ||
|
|
||
| await publish_auth_cache_invalidation(cache_key=counter_key, new_value=value, ttl=60) | ||
| await publish_auth_cache_invalidation(cache_key=floor_key, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) |
There was a problem hiding this comment.
If a request reaches a peer pod before it processes the floor-marker and key-object broadcasts, that pod derives the old window boundary from its cached key object, re-sums historical spend after observing the zeroed shared counter, and monotonically restores the counter. The key then continues returning 429 responses after the reset reports success.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
|
@greptileai Added Review notes in the PR description: this race matches invalidate_team_member_spend_state's existing broadcast pattern. Please re-review. |
3 participants
TLDR
Problem this solves:
How it solves it:
User Flow
Before: a proxy admin resets a key's spend because it's stuck returning 429s, but the key keeps 429ing anyway
max_budgetof $1000 and an extra daily budget window of $50 (budget_limits)POST https://litellm-domain/v1/chat/completionsstarts returning429 {"error":{"message":"ExceededBudget: Key over 1d budget...https://litellm-domain/ui/?page=api-keys(or callsPOST https://litellm-domain/key/<key_id>/reset_spend), and the response reports"spend": 0.0429 {"error":{"message":"ExceededBudget: Key over 1d budget..., even though the key's own dashboard now shows $0 spentAfter: the same reset actually unblocks the key
429onPOST https://litellm-domain/v1/chat/completionsPOST .../reset_spend), which again reports"spend": 0.0200with a real completionRelevant issues
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*,make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Setup for both runs: a local proxy against a real Postgres and Redis (
general_settings.coordination_redisconfigured, so spend counters are Redis-backed and cross-pod), with a realopenai/gpt-4o-minideployment.Before (597b4bb)
curl -X POST http://localhost:14920/key/generate -H "Authorization: Bearer $MASTER_KEY" -d '{"budget_limits": [{"budget_duration": "1d", "max_budget": 0.0000001}]}'->
{"key":"sk-KPDImrEUugFGOFpS6yi6ZA", "token":"f070eb42...e38a88", ...}curl -X POST http://localhost:14920/v1/chat/completions -H "Authorization: Bearer sk-KPDImrEUugFGOFpS6yi6ZA" -d '{"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "say hi"}], "max_tokens": 5}'->
200, a real OpenAI responsesame curl as step 2
->
429 {"error":{"message":"ExceededBudget: Key over 1d budget. Spend=$0.0000, Limit=$0.00","type":"budget_exceeded","code":"429"}}curl -X POST http://localhost:14920/key/f070eb42.../reset_spend -H "Authorization: Bearer $MASTER_KEY" -d '{"reset_to": 0}'->
{"key_hash":"f070eb42...","spend":0.0,"previous_spend":4.35e-6,"max_budget":null,"budget_reset_at":null}spend: 0.0:-> still
429 {"error":{"message":"ExceededBudget: Key over 1d budget. Spend=$0.0000, Limit=$0.00", ...}}(the bug)redis-cli GET spend:key:f070eb42...->0.0(the lifetime counter was reset)redis-cli GET spend:key:f070eb42...:window:1d->0.00000435(the daily window counter, still stale)After (4448c42)
budget_limits: [{"budget_duration": "1d", "max_budget": 0.0000001}]->
{"key":"sk-PKjmwcEnS9thQ_HWeLYzgQ", "token":"03b89980...ec6257", ...}->
200, a real OpenAI response->
429 {"error":{"message":"ExceededBudget: Key over 1d budget. Spend=$0.0000, Limit=$0.00", ...}}curl -X POST http://localhost:14921/key/03b89980.../reset_spend -H "Authorization: Bearer $MASTER_KEY" -d '{"reset_to": 0}'->
{"key_hash":"03b89980...","spend":0.0,"previous_spend":0.0,"max_budget":null,"budget_reset_at":null}redis-cli GET spend:key:03b89980...->0.0redis-cli GET spend:key:03b89980...:window:1d->0.0(now reset too)->
200, a real OpenAI response (fixed)Type
🐛 Bug Fix
Review notes
Greptile's outstanding finding (4/5): even after persisting each window's new boundary before zeroing its counter, a peer pod's key-object cache and floor marker are invalidated by a broadcast (
publish_auth_cache_invalidation), which is asynchronous pub/sub with no delivery acknowledgment. A request landing on that peer between the DB commit and the broadcast's arrival could, in principle, recompute and republish the pre-reset spend.That is true, and it is also the same tradeoff
invalidate_team_member_spend_statealready ships with (litellm/proxy/auth/auth_checks.py:2511-2512) for the LIT-3803 pattern this PR extends to keys: it sets the counter and floor marker locally, then does the identical fire-and-forget broadcast, with no synchronous cross-pod acknowledgment either. No reset path anywhere in the proxy (team, team-member, tag, customer, or key) achieves synchronous cross-pod cache coherence; all of them rely on best-effort pub/sub bounded by the receiving pod's own short cache TTLs (5s for the floor marker, 60s for the counter and key object). Closing this fully would mean adding a distributed acknowledgment or lock protocol to the whole family of reset endpoints, which is a materially larger change than a bug fix PR, and inconsistent with the pattern this PR is deliberately matching.What this PR does narrow relative to before it: previously a key's window counter was never reset at all (permanently stuck), and the key-object cache eviction was never broadcast for keys (only broadcast for bulk key deletion). Both gaps are closed to the same best-effort, sub-TTL-bounded consistency every other reset path in the proxy already has.
Caveats (if any)
Medium
/key/update's separatespendfield (an admin directly overriding a key's stored spend, distinct from the Reset Spend action) still only resets the lifetime counter, not budget windows; same gap, narrower and less commonly used path, left out of scope hereLow
Final Attestation