[Bugfix][Router] loadaware: score bursts against live load, not the arrival snapshot

[ibrahimnd2000]

Problem

Under a burst, loadaware places almost every request on one endpoint.

route_general_request snapshots request_stats before it awaits route_request. LoadAwareRouter.route_request then awaits the controller lookup (and the instance-map refresh). A request only counts as in flight once process_request calls on_new_request. So when many requests arrive together, every one of them is scored against the same pre-burst load. The load term is equal for all endpoints, and whichever endpoint holds even a short shared prefix (for example a common system prompt) wins every placement.

Impact: a GLM-5.x deployment with 3 vLLM replicas ran a benchmark cell of 128 concurrent ~32k-token prompts.

• The opening burst was placed 95 / 30 / 22.
• The overloaded replica needed ~130 s to prefill its queue, so a large share of the burst waited past the client's first-token timeout.

Fix

• After the last await in route_request, re-read request_stats from request.app.state.request_stats_monitor (new helper LoadAwareRouter.live_request_stats). This applies to the select path and to the no-cache fallback path.
• Why this is enough: nothing awaits between the placement decision and on_new_request. After route_request returns, route_general_request runs only synchronous code until await anext(stream_generator), which runs process_request synchronously up to on_new_request. So each decision sees every request placed before it.
• Without a reachable monitor (e.g. unit tests without an app), the caller's snapshot is used, as before.

Tests

New tests in src/tests/test_loadaware_router.py route 128 concurrent requests through route_request with a real RequestStatsMonitor. The controller round-trip is mocked with 1–20 ms of latency, and a 512-of-32,000-token prefix is cached on one of three endpoints. As in the real request path, each placement is followed by on_new_request.

	Placement over 3 endpoints
Stale snapshot (old behaviour)	128 / 0 / 0
With this fix	43 / 43 / 42

Also added: unit tests for live_request_stats, covering the live read and the fallback to the snapshot without a monitor.

pytest src/tests/test_loadaware_router.py src/tests/test_kvaware_router.py src/tests/test_request_stats.py
47 passed

Four of the new tests fail on main. pre-commit (black, isort, ruff, codespell) passes.

Validation

In the deployment above, with the fix, a live 30-request burst with a shared cached system prompt was placed 11 / 10 / 9. The router's per-endpoint in-flight counts matched the engines' running requests.

Notes

• This complements [Bugfix][Router] Fix leaked in-flight counters in RequestStatsMonitor #1072, which fixed in-flight counters leaking on errors and client disconnects. That fix keeps the counts correct over time; this one keeps them current within a burst.
• "loadaware: KV lookup for chat completions" (separate PR) adds a lookup-failure fallback in route_request. Once both are merged, that fallback should also use live_request_stats. Both PRs add a test section at the same place in test_loadaware_router.py, so whichever lands second needs a trivial rebase.

[gemini-code-assist]

Code Review

This pull request introduces a live_request_stats method to fetch up-to-date request statistics from the monitor during routing, rather than relying on stale snapshots taken at the request's arrival. This prevents concurrent request bursts from being routed to the same endpoint due to outdated load information. Corresponding unit tests have been added to verify burst routing behavior and live stats retrieval. There are no review comments, so no further feedback is provided.

[chrikrah]

@ibrahimnd2000 I would merge this, and the no-await argument checks out. At 7085a91, deleting both live_request_stats call sites fails test_a_burst_is_spread_by_live_load, where your branch passes all 47.

non-blocking: the live read on the no-cache path at routing_logic.py:775 is pinned by nothing. Deleting only that line leaves all 47 passing, and that path carries the opening burst of a cold fleet.

$ PYTHONPATH=src python -m pytest src/tests/test_loadaware_router.py src/tests/test_kvaware_router.py src/tests/test_request_stats.py -q
47 passed                     # at 7085a91
4 failed, 43 passed           # routing_logic.py reverted to 691dffa, tests kept
1 failed, 46 passed           # only the two call sites deleted
47 passed                     # only line 775 deleted

$ PYTHONPATH=src python -m pytest src/tests/test_loadaware_router.py -q -s -k cold_burst   # with the case below
PLACEMENT [43, 43, 42]        # at 7085a91: 1 passed
PLACEMENT [128]               # line 775 deleted: 1 failed
# not run: the rest of src/tests, and no real controller or engines

The cold-burst case

-async def route_burst(router, monitor):
+async def route_burst(router, monitor, layout_info=None):
@@ query_manager
+        if layout_info is not None:
+            return LookupRet(layout_info)
         return LookupRet({INST_A: (LOCAL, SHARED_PREFIX_TOKENS)})
@@ end of file
+@pytest.mark.asyncio
+async def test_a_cold_burst_is_spread_by_live_load(stats_monitor):
+    """Nothing cached anywhere: the fallback must also see live load."""
+    from uhashring import HashRing
+
+    router = burst_router()
+    router.session_key = "x-user-id"
+    router.hash_ring = HashRing()
+    placement = await route_burst(router, stats_monitor, layout_info={})
+    print("PLACEMENT", sorted(placement.values(), reverse=True))
+    assert max(placement.values()) <= BURST // 2

non-blocking: #1107 adds a third fallback_url call, on the lookup-failure path, still with the arrival snapshot. Reading live_request_stats inside fallback_url would cover all three in either merge order.

@ruizhang0101, you are on #1107 too, so whichever lands second needs this. @ibrahimnd2000, would you add the cold-burst case above as a second test?

[ibrahimnd2000]

@chrikrah thanks for the careful review and the counter-checks. Both points are addressed in 431d783:

• Live read inside fallback_url: moved the live_request_stats read from the no-cache branch into fallback_url, so all three fallbacks score against the live load: no cache info, no endpoint selected, and the lookup-failure path [Bugfix][Router] loadaware: KV lookup for chat completions, fall back when the lookup fails #1107 adds, in either merge order. The select path keeps its own read.
• Cold-burst test: added test_a_cold_burst_is_spread_by_live_load as you wrote it, plus a sum(...) == BURST check like the warm-burst test.

$ PYTHONPATH=src python -m pytest src/tests/test_loadaware_router.py src/tests/test_kvaware_router.py src/tests/test_request_stats.py -q
48 passed                     # at 431d783
1 failed, 47 passed           # live read removed from fallback_url: test_a_cold_burst_is_spread_by_live_load fails

The rest of src/tests (excluding e2e) passes too: 241 vs 240 at 7085a91. test_multipart_proxy.py was not collected locally (no httpx in my env).