[Bugfix][Frontend] Keep in-flight requests on the same DP engine (#59017)
Signed-off-by: Benedikt Falk <bfalk@nvidia.com>
Benedikt Falk committed
Sep 29, 2026 at 01:52 UTC
c68eb98b69f4863d94318fbe9dedf654be0c5945
2 files changed
+17
tests/v1/engine/test_engine_core_client.py
+13
@@ -248,6 +248,19 @@ def test_dplb_non_late_interaction_still_uses_lb():
248
assert client.lb_engines[1][0] == 1
249
250
251
+def test_dplb_in_flight_request_stays_on_same_engine():
252
+ client = _make_dplb_client()
253
+ request = make_request(SamplingParams(max_tokens=1))
254
+ engine = client.core_engines[2]
255
+ client.reqs_in_flight[request.request_id] = engine
256
+ client.engine_inflight[engine] = 1
257
+
258
+ chosen_engine = client.get_core_engine_for_request(request)
259
+
260
+ assert chosen_engine == engine
261
+ assert client.engine_inflight[engine] == 1
262
+
263
+
264
def test_dplb_burst_round_robins_despite_snapshot_rebinds():
265
"""A stats snapshot rebind wipes the optimistic lb_engines increments;
266
the exact in-flight floor must keep a burst spreading round-robin."""
vllm/v1/engine/core_client.py
+4
@@ -1587,6 +1587,10 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1587
) // client_count
1588
1589
def get_core_engine_for_request(self, request: EngineCoreRequest) -> EngineIdentity:
1590
+ if (engine := self.reqs_in_flight.get(request.request_id)) is not None:
1591
+ # This is for streaming-input session affinity.
1592
+ return engine
1593
+
1594
# Engines are in rank order.
1595
if (eng_index := request.data_parallel_rank) is None and (
1596
eng_index := get_late_interaction_engine_index(