Repository navigation
fix(scheduler): remove a request's queue entry once it stops waiting - #43061
mateo-berri merged 15 commits into
Conversation
Requests admitted while a healthy deployment existed never left the priority queue, and neither did requests cancelled while waiting. The stale entries blocked later requests during cooldown and, with Redis, made add_request raise a TypeError on queues read back as JSON lists Both scheduling paths now share one polling helper that removes the entry in a finally block, whether the request was admitted, timed out or cancelled Related to BerriAI#43059
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
…ge check The coverage script only counts direct calls in test files. The helper is exercised through prioritized acompletion and atext_completion in test_router.py, like the other allowlisted entries
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
…clean up cancelled enqueues poll() raised on an empty queue before checking for healthy deployments. With the cleanup now rewriting the queue after every admission, a concurrent write from another replica can erase a waiting request's entry, and that request then failed while a deployment was healthy. poll() now admits as soon as a deployment is healthy and only reads the queue during cooldown add_request also moved inside the try block, so a request cancelled while its queue write is in flight still has its entry removed
The router passes a healthy-deployments callable into the scheduler, so tests inject a Scheduler directly instead of replacing the router's scheduler attribute. The cancelled-mid-enqueue test moves to test_scheduler.py with the other scheduler tests, and poll() takes the deployments as a Sequence since it only checks whether any are healthy
|
mcp-integration failure is an unrelated MCP TCP timing flake. It passed in the previous two runs and locally, so it needs a maintainer rerun |
VANDRANKI
left a comment
VANDRANKI
left a comment
There was a problem hiding this comment.
Community review, does not clear the merge gate.
I read the full diff (router.py, scheduler.py, and both test files) and traced the actual bug rather than trusting the title.
The real bug is in Scheduler.poll(), not just in the caller's timeout-cleanup path. The old code:
if len(health_deployments) == 0:
if queue[0][1] == id:
heapq.heappop(queue)
...
return True
else:
return False
return TrueWhen health_deployments is non-empty, execution falls straight to the final return True without ever touching the queue. So every request admitted via "a healthy deployment exists" (presumably the common case, not the edge case) was returning True and proceeding to make its actual API call while its FlowItem stayed permanently stuck in the scheduler's queue (Redis-backed or in-memory DualCache depending on config). The old caller code (schedule_acompletion/_schedule_factory) only ever called scheduler.remove_request(...) explicitly on the timeout branch; the success branch relied entirely on poll()'s internal heapq.heappop, which never runs on the healthy-deployment fast path. That's an unbounded queue leak under normal healthy operation, not just a cancellation edge case.
There's a second, separate leak this also fixes: the old polling loop in the router methods was a bare while loop with no try/finally. If the awaiting task was cancelled mid-poll (client disconnect, upstream cancellation), asyncio.CancelledError propagates straight out, skipping both the success block and the else: remove_request(...) timeout-cleanup block, leaking the entry.
The fix consolidates both call sites into Scheduler.wait_for_turn(), wrapping add_request and the poll loop in try: ... finally: await self.remove_request(...), which runs on every exit path: success (whether or not poll() itself already popped the item), timeout, or cancellation. I checked remove_request() is safe to call redundantly ([item for item in queue if item[1] != request_id], a no-op if the id is already gone), so calling it unconditionally in finally doesn't double-remove or error when poll() already popped the entry via the priority-order path.
I traced the new tests against this:
test_prioritized_request_leaves_queue_when_it_stops_waiting(parametrized cancel/timeout) directly exercises both leak scenarios: a cooled-down deployment forces the second request to sit in queue behind a manually-inserted head-of-queue item, then either cancels the waiting task or lets it time out, and asserts the queue ends up as exactly[(0, "head-of-queue")]in both cases, i.e., the request's own entry is gone either way.test_admitted_prioritized_request_does_not_block_later_request_during_cooldownchecks the router's cooldown/rate-limit path still raisesRouterRateLimitErrorcorrectly through the refactored_wait_for_scheduler_turn.
One gap I want to flag honestly: test_wait_for_turn_removes_entry_when_cancelled_mid_enqueue in the new test_scheduler.py sets up a scenario where cancellation happens while still inside add_request (via a subclass that pauses forever right after enqueueing), asserts the entry is present, cancels, and asserts CancelledError is raised — but it never asserts the queue is actually empty afterward, despite the test's name claiming to verify removal. I traced the production code myself and the finally block should still run in this case since the cancellation lands inside the try, so I believe the behavior is correct, but this particular test doesn't actually prove it. Worth tightening before merge, not a blocker on its own since the other test (test_prioritized_request_leaves_queue_when_it_stops_waiting) does verify post-cancellation queue state for the main call path.
This is a real, meaningful fix for an unbounded scheduler-queue leak, correctly scoped, with one test that could use a stronger final assertion.
|
Offering a regression suite — no obligation, and no competing PR: I had a minimal fix for the same issue (#43060) and closed it in favour of this one, since Before closing I built tests that fail on
Two things I hit that may be useful for your
I measured both the failures and the fix through Happy to push these as a commit to your branch or paste them into the thread, whichever you prefer. |
…emove_admitted_requests # Conflicts: # tests/test_litellm/test_router.py
|
@VANDRANKI really appreciate the thorough review. For the mid-enqueue test, the queue-is-empty check is on its last line, in case it got cut off |
|
You're right, thanks for pointing that out. I just went back and checked the full test body and |
…ove_admitted_requests
…ed again during cleanup A second cancel, or an anyio cancel scope that re-cancels on every await, interrupted remove_request mid-write and left the entry in Redis.
|
@apex-mochen Thanks, the shield point was real. Added asyncio.shield plus a double-cancel regression test in 0c9ab36. Cleanup can't raise, since DualCache swallows Redis errors |
…emove_admitted_requests # Conflicts: # tests/code_coverage_tests/router_code_coverage.py
|
Thanks for confirming this and for adding the double-cancel regression test. For the record, the asyncio.shield concern came from the cancellation-cleanup review on #43060, and I’m glad it was useful in hardening this PR. If LiteLLM keeps contributor or review acknowledgments in the PR description, a brief note there would be appreciated. No commit-history change is needed. |
|
Hi @RachelHuangZW, thanks again for incorporating the cancellation-cleanup point — specifically asyncio.shield and the double-cancel regression. Since that review feedback directly influenced the final implementation, would you mind adding a brief acknowledgment in the PR description, e.g.: Thanks to @apex-mochen for identifying the cancellation-cleanup race and suggesting shielded cleanup / regression coverage. |
…emove_admitted_requests # Conflicts: # litellm/router.py # litellm/scheduler.py
|
@apex-mochen Done, added the acknowledgment to the description under Relevant issues. Sorry for the slow reply, and thanks again for the review |
…ove_admitted_requests # Conflicts: # litellm/router.py # tests/unit/test_scheduler.py
|
bugbot run |
|
bugbot run |
|
|
||
|
|
||
|
|
||
| class _JsonRoundTripRedisCache: |
There was a problem hiding this comment.
_JsonRoundTripRedisCache and _ExpiringCache, violating the repository directive against doubles for its own code. This requirement must be satisfied before merging
Context Used: AGENTS.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!
There was a problem hiding this comment.
These stand in for the Redis server behind Scheduler's redis_cache seam, like main's StubRedisCache; RedisCache offers no client injection, and tests assert queue state
There was a problem hiding this comment.
That’s a fair clarification. These are focused cache-seam fakes rather than mocks changing Scheduler behavior: _JsonRoundTripRedisCache reproduces Redis’s JSON list round-trip, _ExpiringCache models the queue key disappearing, and the existing StubRedisCache test establishes the same pattern in this file.
Since RedisCache does not expose client injection, requiring a live Redis server for these unit-level state assertions would make the tests slower and less deterministic. I’m withdrawing the finding; the real-Redis behavior is appropriately covered by the live proxy verification.
…ances Adds tests/integration/routing/test_priority_scheduler_queue_cleanup.py: 27 cells against the real proxy (two workers, Postgres, Redis, a scripted upstream) for every prioritized surface (/v1/chat/completions, /v1/completions, /queue/chat/completions, streaming and not, OpenAI SDK sync and async, raw httpx), the non-integer priority pass-through, the in-memory queue on one proxy, two proxies sharing a Redis queue (served requests leave no entry, a dead replica's entry is skipped or expires, a waiter that times out or disconnects removes only itself), a Redis outage mid burst and a SIGKILLed worker. Each cell asserts the caller's response, the upstream's requests by marker and the Redis queue contents. On the merge base the cross-instance cells fail with the list-of-lists TypeError and the in-memory cell with 408s behind the leaked entry; on this branch every cell passes
|
bugbot run |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 0be0db2. Configure here.
mateo-berri
left a comment
mateo-berri
left a comment
There was a problem hiding this comment.
LGTM. Thanks for the contribution!
4 participants
TLDR
Problem this solves:
How it solves it:
Scheduler.wait_for_turnshared by both scheduling pathsfinallyblockpoll()admits on a healthy deployment before reading the queueasyncio.shield, so a second cancel during cleanup can't cut it shortget_queue()validates every entry read from the cache into a(priority, request_id)tuple, so the heap compares like with like whatever Redis handed backpoll()puts a request whose queue entry vanished (the key expired, or a concurrent writer erased it) back at its priority, so it keeps its place instead of failing withIncorrectly setup. Queue is invalidpoll()takes theFlowItemnow, since putting the entry back needs the priorityUser Flow
Before: an app sending prioritized chat requests through a proxy backed by Redis sees about half of them fail while the model is healthy
redis_hostandredis_portunderrouter_settingsand starts the proxy"priority": 1, once a second'<' not supported between instances of 'tuple' and 'list', then the cycle repeatsAfter: the same requests all succeed
redis_hostandredis_portunderrouter_settingsand starts the proxy"priority": 1, once a secondRelevant issues
Fixes #43059
#43060 was closed in favor of this PR. Thanks to @apex-mochen for identifying the cancellation cleanup race and suggesting the shielded cleanup and its regression coverage
Linear ticket
Resolves LIT-9076
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
uv run pytest tests/unit/<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
Both legs drive the same two-replica rig: two proxies with 2 uvicorn workers each (
--num_workers 2 --timeout_worker_healthcheck 60 --telemetry False), sharing one local Redis and running without a database, each booted from the commit named in its heading. Requests alternate between instance A and instance B so a queue entry written through one worker is read back through another. Every call goes to real OpenAIgpt-5.4-nano(nomock_response),qa-cooldowncarries an invalid OpenAI key so its only deployment cools down afterallowed_fails: 1, and$LITELLM_MASTER_KEYis exported from the proxy's.env. The0at the top of an output is theredis-cli delreply (the key had already expired, see the observations)proxy_config.yaml, identical on both sides apart from the Redis port:Before (31b90eb)
Instance A on
127.0.0.1:42184, instance B on127.0.0.1:42403, Redis on48170. Every curl carries-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json"Steady prioritized chat on /v1/chat/completions, 16 requests one per second, alternating A, B
pong, 11 of 16 HTTP 500 with the TypeError, on both instances)Streaming prioritized chat on /v1/chat/completions, 6 requests, alternating A, B
pong, 5 of 6 HTTP 500 with the TypeError as a plain JSON body before any SSE frame)Steady prioritized text completion on /v1/completions, 16 requests one per second, alternating A, B
choices[0].text=pong, 11 of 16 HTTP 500 with the TypeError)Cooldown: a priority 1 call, then a priority 2 probe on the other instance
8 concurrent prioritized chat requests fired at once, 4 per instance
Cooldown: a priority 2 request behind a stale head entry
A worker killed mid-wait leaves its entry at the head of the queue. The seed below writes such an entry the way the scheduler itself writes the key (
EX 5), after 6 s of idle so no worker holds a fresh in-memory copy, then sends a prioritized request behind it whileqa-cooldownis still in cooldown from the case aboveAfter (c835b73)
Instance A on
127.0.0.1:32015, instance B on127.0.0.1:30458, Redis on47831. Same headers on every curlSteady prioritized chat on /v1/chat/completions, 16 requests one per second, alternating A, B
pong)Streaming prioritized chat on /v1/chat/completions, 6 requests, alternating A, B
pongwith HTTP 200)Steady prioritized text completion on /v1/completions, 16 requests one per second, alternating A, B
choices[0].text=pong)Cooldown: a priority 1 call, then a priority 2 probe on the other instance
All deployments for selected model are in cooldownin 0.02 s and 0.03 s, no 30 s hang and no poll timeout)Cooldown: a priority 2 request behind a stale head entry
Same seed as the Before case: an entry a killed worker would leave at the head, written the way the scheduler writes the key (
EX 5), after 6 s of idle, withqa-cooldownstill in cooldownAll deployments for selected model are in cooldownafter 5.07 s and 5.06 s, no 500 and no poll timeout)8 concurrent prioritized chat requests fired at once, 4 per instance
pong, each in 0.60 to 0.93 s)Observations from the two legs, each with whether this PR causes it, makes it worse, or leaves it alone:
Last updated: c835b73. QA and /live-pr-risk ran on c835b73, the last commit with a product diff
/live-pr-risk
Rig: two proxies per side, each
--num_workers 2, sharing one Redis and one forwarding recorder in front ofhttps://api.openai.comthat logs every upstream request with its status, requests alternating between the two instances; a base side at 31b90eb, a head side at c835b73, and a mixed side with instance A at 31b90eb and instance B at c835b73 on one Redis, the shape of a rolling deploy. The same cells ran on every side in the same order:/queue/chat/completionsnon-stream and stream,/v1/chat/completionswith and withoutpriority, streaming withpriority,/v1/completionswith and withoutpriority,priorityas a string and as a list, 8 concurrent prioritized chat requests, three calls that tripqa-cooldowninto cooldown, a prioritized call the client cancels after 2 s (non-stream and stream), then the cooled-down model through the chat, queue and text routes, a healthy model with and withoutpriority,fastest_responseover the healthy and the cooled-down model withpriority, and the cooled-down model once more. Config:openai/gpt-5.4-nano,timeout: 30,num_retries: 0,allowed_fails: 1,cooldown_time: 180, defaultpolling_interval, no databaseBreaking
None observed. Every cell that fails on the base and passes on the head fails with one of the two symptoms this PR targets:
'<' not supported between instances of 'tuple' and 'list'(500 on/v1/chat/completionsand/v1/completions, 400auth_erroron/queue/chat/completionsstreaming, 7 of the 8 concurrent requests, see the Before proof above) or a leaked entry blocking the cooldown queue (the cancelled request still waiting when the client gives up at 2 s, the next request on that instance answeringIncorrectly setup. Queue is invalid. Queue=[]after about 7 s). On the head those cells answer 200, or an immediate 429All deployments for selected model are in cooldownonce the model is cooled down. Every cell that passes on the base passes identically on the head: same status, samex-litellm-*header set with ids, spend and cost normalized, same body shape, and the same request body key set recorded in front of OpenAIBackward incompatible
Redis traffic per prioritized request changed. The head writes the queue twice per request (the add, then the removal that leaves
[]) and no longer reads it while a healthy deployment exists; the base wrote once and read on every poll, and never removed anything. In the MONITOR capture below the base did 16 GET and 11 SET onscheduler:queue:*over its run and the head 8 GET and 57 SET while serving every prioritized request the base failed. This is the fix's mechanism (the removal write is what clears the queue), landed in 629b9ef, and it stays listed under Low caveats; c835b73 adds one write only when a poll during cooldown finds its own entry goneScheduler.get_queue()now returns validated(priority, request_id)tuples instead of whatever the cache held (lists after a Redis round trip), andScheduler.poll()returnsTruewithout touching the queue while any deployment is healthy. Both are called only fromlitellm/router.py(_wait_for_scheduler_turnthroughScheduler.wait_for_turn) and from this PR's tests; no proxy route, dashboard page orlitellm/proxy/clientcommand reads scheduler state.Scheduler.poll()also takes theFlowIteminstead ofidandmodel_name, and during cooldown it puts a request whose entry vanished back at its priority instead of raisingIncorrectly setup. Queue is invalid, so a caller that relied on that exception as a signal no longer gets it;wait_for_turnis its only caller. Landed in 629b9ef, e6cdc06, 1313444 and c835b73No HTTP contract change was observed: statuses, error types and the
x-litellm-*header set match the base on every cell that passes on both sidesRegression risk
Router.abatch_completion_one_model_multiple_requestswithpriorityis reachable from the SDK only: it fans out Nacompletioncalls that now each enqueue and remove. The per-call path is the one verified live, the fan-out is not, and exercising it takes an SDK script against RedisA
polling_intervalother than the default, and atimeoutchanged at runtime through/config/update(needs Postgres): the wait loop reads both, and only the defaults ran hereRedis becoming unreachable while a request waits, and a worker SIGTERM while a request waits (the shielded removal runs during cancellation): neither was induced
Pre-existing, identical on the base and the head, and not made worse by this PR: with
router_settings.default_priorityset, every/v1/chat/completionsrequest fails 500Router.schedule_acompletion() missing 1 required positional argument: 'priority'(400 on/queue/chat/completions), explicit integerpriorityincluded, becauseschedule_acompletionre-entersacompletion, which re-applies the default;/v1/completionsis unaffected. Tracked internally for its own fix. A non-integerpriority("1",[1]) skips the scheduler and is forwarded to OpenAI, which answers 400Unknown parameter: 'priority'; the recorder shows theprioritykey in the upstream body 2 times on the base and 2 times on the headDependency graph
Router.acompletionwith an integerpriority->schedule_acompletion->_wait_for_scheduler_turn->Scheduler.wait_for_turn: verified live (chat, chat stream, 8 concurrent, cooldown, client cancel) and tested intests/test_litellm/test_router.py/queue/chat/completions->router.schedule_acompletion: verified live, non-stream, stream and during cooldown/v1/completionswithpriority->atext_completion->_schedule_factory->_wait_for_scheduler_turn: verified livefastest_responsewithpriority->abatch_completion_fastest_response->acompletionper model: verified live, the healthy model answers while the other is cooled downScheduler.add_request,poll,remove_request,get_queue,save_queue->DualCache-> Redis keyscheduler:queue:<model>with a 5 s TTL: verified live through MONITOR (head: every add is followed by a[]removal; base: entries are never removed) and tested intests/unit/test_scheduler.pywith a JSON round-trip stub. The vanished-entry path inpollis verified live by the stale head entry case in the proof above (the seeded key expires while the request waits) and tested with a cache the test clears mid-waittuple/listerror, and its next prioritized request on the cooled-down modelIncorrectly setup. Queue is invalid. Queue=[]after 7.5 s). Neither replica made the other worseRouter.abatch_completion_one_model_multiple_requestswithpriority: untested, SDK onlyRouter.timeoutandpolling_intervalas read by the wait loop: defaults onlytests/llm_translation,tests/local_testingandtests/integration: none drives the scheduler orpriority, so nothing there asserts the old behaviorui/litellm-dashboard/src) andlitellm/proxy/client: no consumer of scheduler stateNot verified
The dependents marked untested or defaults only above:
abatch_completion_one_model_multiple_requestswithpriority, a non-defaultpolling_interval, a runtimetimeoutchange through/config/update, a Redis outage mid-wait, a SIGTERM mid-waitA merged tree:
mainhad moved 15 commits past the merge base when the rig ran (31b90ebb47...a9b9700790) and none of them toucheslitellm/scheduler.py,litellm/router.py,litellm/caching/,litellm/router_utils/or the queue route inlitellm/proxy/proxy_server.py, so the rig ran on the PR tip rather than on a merged treeService-logger metrics around the scheduler were not read; Redis operations were counted instead. The rig ran without Postgres
Redis MONITOR, scheduler:queue operations per side
Base (31b90eb), instances A and B:
Head (e6cdc06), instances A and B:
Mixed (A at 31b90eb, B at e6cdc06):
/audit
Merge base 31b90eb and head c835b73, both two-worker proxies with their own Postgres, Redis and scripted upstream on the shared QA box; the head selection ran twice with
--seed 1 --order-seed 1. The 16 rig-proxy rows ran in lanesseq2/base-r3,seq6/head-r6a,seq6/head-r6b, the owned-proxy rows inseq5/base-r5,seq5/head-r5a,seq5/head-r5band P8 inseq7/base-r7,seq7/head-r7a,seq7/head-r7b, one lane after another. Aredon the merge base is the bug this PR fixes, shown by the same cell;P3andJ2are red or green on the merge base by timing (its leaked entries and the list-of-lists reads race), and every head run is green/v1/chat/completions, OpenAI SDK sync,priority: 1, healthytests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_sdk_chat_with_priority_is_served_once_and_billedtests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_sdk_chat_stream_with_priority_is_served_once_and_billedtests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_async_sdk_chat_with_priority_is_served_oncepriorityJSON key ("not-an-int"then1, last wins)tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_raw_chat_with_duplicate_priority_keys_takes_the_last_one/v1/completions, raw httpx,priority: 1tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_other_prioritized_endpoints_are_served_once_and_billed[completions]/v1/completionsstream, consumedtests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_other_prioritized_endpoints_are_served_once_and_billed[completions-stream]/queue/chat/completionsraw httpx non-streamtests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_other_prioritized_endpoints_are_served_once_and_billed[queue]/queue/chat/completionsstream, consumedtests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_other_prioritized_endpoints_are_served_once_and_billed[queue-stream]priority: "1"(string)tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_non_integer_priority_bypasses_the_scheduler_and_is_forwarded[string]priority: [1]tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_non_integer_priority_bypasses_the_scheduler_and_is_forwarded[list]priority: ""tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_non_integer_priority_bypasses_the_scheduler_and_is_forwarded[empty]priority: <5 KB string>tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_non_integer_priority_bypasses_the_scheduler_and_is_forwarded[5kb]priority: 0(falsy, never scheduled: pre-existing quirk, observation only)tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_non_integer_priority_bypasses_the_scheduler_and_is_forwarded[zero]priority: 1unauthenticatedtests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_unauthenticated_prioritized_request_never_reaches_the_upstream/v1/messageswithpriority: 1(no scheduler on this path)tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_priority_on_unscheduled_routes_is_served_once[messages]/v1/responseswithpriority: 1(no scheduler on this path)tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_priority_on_unscheduled_routes_is_served_once[responses]priority: 1healthy, trip the deployment, B then Cpriority: 2during cooldowntests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_in_memory_queue_forgets_served_requests_before_a_cooldownfirst, B atsecond, C atfirst, all healthytests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_served_request_leaves_no_queue_entry_for_the_other_instancefirst, then 12 concurrent prioritized requests (both instances x chat, chat stream, completions, completions stream, queue, queue stream), then one more atsecondtests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_concurrent_prioritized_burst_across_instances_and_endpoints_is_served/v1/completionsSDK during cooldown with a fresh queuetests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_fresh_queue_during_a_cooldown_is_refused_at_oncefirst, trip, Bpriority: 2atsecondtests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_served_request_on_one_instance_does_not_block_a_waiter_on_the_other[[1, "ghost"]]seeded in Redis, trip, Bpriority: 2via/queue/chat/completionsatsecondtests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_waiter_behind_an_expired_dead_replica_entry_is_re_enqueued_and_proceedspriority: 2waits behind it via/v1/chat/completions, the key isPERSISTed once B's entry showstests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_waiter_behind_a_live_entry_times_out_and_removes_only_itselfpriority: 2waits behind the ghost at/v1/chat/completions, the key is rewritten to B's entry alonetests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_waiter_proceeds_once_the_entry_ahead_is_clearedhttp.clientkeep-alive connection tosecondcarries the trip, the 429 confirmation and then Bpriority: 2at/v1/chat/completions, so all three land on the worker that recorded the cooldown (a fresh connection can reach the other worker, whose cooldown view lags and admits B at once, the pre-existing per-worker convergence P3 also shows); the key isPERSISTed once B's entry shows, then the client closes the connection (cancel_on_disconnect: true)tests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_disconnected_waiter_is_removed_from_the_queuetests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_prioritized_requests_survive_a_redis_outagetests/integration/routing/test_priority_scheduler_queue_cleanup.py::test_sibling_worker_keeps_serving_prioritized_requests_after_a_worker_is_killedType
Bug Fix
Caveats (if any)
Low
router_settings.default_priorityset, every/v1/chat/completionsrequest fails 500Router.schedule_acompletion() missing 1 required positional argument: 'priority'on the merge base and on this PR alike--timeout_worker_healthcheckderived from the harness readiness budget, because uvicorn's default 5 s worker ping respawned memory-starved workers in a loop on the loaded QA box and the cell never saw a ready proxy; the harness's other owned-proxy files keep the defaultScheduler'sredis_cacheconstructor seam the way main's existing test does; real Redis is covered by the proxy proof aboveRedisCacheoffers no client injection, so a real cache over fakeredis would need a monkeypatch, which the test rules rank worsemake checkat the tip failslint-e2e-basedpyrightontests/e2e/llm_translation/realtime/realtime_client.py:366(InputAudioBufferAppendis not insend()'s event union)send()after bd70e6d (fix(bedrock): forward each Nova Sonic assistant sentence once over the realtime API #44987) started sending audio frames, and each PR was green at its own merge baseFinal Attestation
Note
Medium Risk
Changes core router scheduling and Redis-backed priority queue semantics (extra writes, different poll behavior); well-covered by integration tests but affects all integer-
prioritytraffic on Redis proxies.Overview
Fixes prioritized request scheduling by centralizing queue wait logic and ensuring queue entries are always removed when a request finishes, times out, or is cancelled.
Router —
schedule_acompletionand_schedule_factoryno longer duplicate enqueue/poll loops; both call new_wait_for_scheduler_turn, which delegates toScheduler.wait_for_turn. Response headers and priority-on-exception behavior are unchanged.Scheduler — Adds
wait_for_turn(enqueue, poll until admitted or timeout) with shieldedremove_requestin afinallyso cleanup completes even on cancel.pollnow takes aFlowItem, returns immediately when deployments are healthy (no queue read), and during cooldown re-enqueues requests whose Redis/in-memory entry disappeared instead of erroring.get_queuenormalizes cached JSON lists to(priority, request_id)tuples to fix tuple/list comparison failures on Redis-backed proxies.Tests — Large integration suite for multi-instance Redis queue behavior plus expanded unit/router tests for cancellation, cooldown, and stale entries.
Reviewed by Cursor Bugbot for commit 0be0db2. Bugbot is set up for automated code reviews on this repo. Configure here.