Repository navigation
fix(router): await Redis increment pipeline so budget sync doesn't overwrite memory with stale spend - #32624
fix(router): await Redis increment pipeline so budget sync doesn't overwrite memory with stale spend#32624devin-ai-integration[bot] wants to merge 1 commit into
Conversation
RouterBudgetLimiting._push_in_memory_increments_to_redis scheduled the increment pipeline with asyncio.create_task and returned immediately, so _sync_in_memory_spend_with_redis read Redis and overwrote in-memory spend before the pipeline had run. Await the pipeline so the read reflects the just-pushed increments, and snapshot the queue before awaiting so concurrent increments are not dropped.
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
|
|
Greptile SummaryThis PR fixes a race condition in
Confidence Score: 4/5The change is a targeted, well-scoped fix validated by two regression tests that reproduce the exact race ordering from the reported issue. The production change in budget_limiter.py is straightforward and correct. The only finding is dead test setup in the new test file — asyncio.create_task is patched in the fixture but never called by the fixed code. The budget_limiter fixture in tests/test_litellm/router_strategy/test_budget_limiter.py retains a now-unnecessary monkeypatch for asyncio.create_task.
|
| Filename | Overview |
|---|---|
| litellm/router_strategy/budget_limiter.py | Replaces fire-and-forget asyncio.create_task with a direct await in _push_in_memory_increments_to_redis, and snapshots the queue before clearing it so concurrent appends carry over to the next flush rather than being dropped. |
| tests/test_litellm/router_strategy/test_budget_limiter.py | New regression-test file with two focused async tests validating the fix; the budget_limiter fixture monkeypatches asyncio.create_task even though the fixed code no longer calls it, leaving dead setup that could confuse future readers. |
Reviews (1): Last reviewed commit: "fix(router): await Redis increment pipel..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @pytest.fixture | ||
| def budget_limiter(monkeypatch): | ||
| monkeypatch.setattr(asyncio, "create_task", lambda coro: coro.close()) | ||
| limiter = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) | ||
| limiter.provider_budget_config = {"openai": BudgetConfig(time_period="1d", budget_limit=100)} | ||
| return limiter |
There was a problem hiding this comment.
monkeypatch.setattr(asyncio, "create_task", ...) line is dead code now that the fix removes the asyncio.create_task call from _push_in_memory_increments_to_redis. It no longer affects the tests and may confuse future readers who wonder why the patch is needed.
| @pytest.fixture | |
| def budget_limiter(monkeypatch): | |
| monkeypatch.setattr(asyncio, "create_task", lambda coro: coro.close()) | |
| limiter = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) | |
| limiter.provider_budget_config = {"openai": BudgetConfig(time_period="1d", budget_limit=100)} | |
| return limiter | |
| @pytest.fixture | |
| def budget_limiter(): | |
| limiter = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) | |
| limiter.provider_budget_config = {"openai": BudgetConfig(time_period="1d", budget_limit=100)} | |
| return limiter |
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!
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
2 participants
Relevant issues
Fixes #32614
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Screenshots / Proof of Fix
This is an internal Redis-sync race, so the observable symptom (in-memory spend getting overwritten with a stale Redis value) depends on the background pipeline task not having completed before the sync read. That timing is not deterministic through a single curl against a live proxy, so the proof here is a regression test that reproduces the exact ordering from the issue: in-memory spend is
160, Redis is100with a pending+60increment queued, and the pipeline yields to the event loop before it applies. If the flush is fire-and-forget, the sync reads Redis before the pipeline lands and clobbers memory back to100; if it is awaited, memory stays160Before, at
60729f733e(the file restored to its pre-fix state):After, at
0f030430a7:Type
🐛 Bug Fix
Changes
RouterBudgetLimiting._push_in_memory_increments_to_redis()scheduled the Redis increment pipeline withasyncio.create_task(...)and returned right away, so awaiting it only waited for the task to be scheduled, not for Redis to actually be updated._sync_in_memory_spend_with_redis()then fetched Redis and wrote those values into the in-memory cache before the pipeline had run, overwriting fresh in-memory spend (e.g.160) with a stale Redis value (e.g.100). BecauseDualCache.async_batch_get_cache()trusts in-memory first, that stale value then flowed into later provider, deployment and tag budget checks inasync_filter_deployments(). The same inherited flush also runs in_PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event(), so virtual-key and end-user model budgets were affected tooThe fix awaits the pipeline before returning, mirroring the safer pattern already used in
BaseRoutingStrategy:The queue is snapshotted and cleared before the
awaitso increments appended by concurrent requests during the flush are not silently dropped; they simply carry over to the next flush instead of being wiped by an unconditional resetLink to Devin session: https://app.devin.ai/sessions/75159837ecdc45db81285457e5a86fde