[Bugfix][Elastic EP] Route only to surviving engines during scale-down (#55772)
Signed-off-by: Djordje Ramic <djoramic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
djramic committed
Sep 9, 2026 at 04:28 UTC
c6ba5e89d660a742f27c4a58e9580ed373a014bc
2 files changed
+62
tests/v1/engine/test_engine_core_client.py
+60
@@ -259,6 +259,66 @@ def test_dplb_burst_round_robins_despite_snapshot_rebinds():
259
assert sorted(client.engine_inflight.values()) == [2, 2, 2, 2]
260
261
262
+@pytest.mark.asyncio
263
+async def test_dplb_scale_down_routes_after_stale_stats_snapshot():
264
+ """A coordinator snapshot during scale-down must only route to survivors."""
265
+ import msgspec
266
+ import zmq.asyncio
267
+
268
+ client = _make_dplb_client(num_engines=4)
269
+ client.engine_ranks_managed = list(range(4))
270
+ client.engines_running = True
271
+ client.current_wave = 0
272
+ client.resources = SimpleNamespace(stats_update_task=None)
273
+ client.stats_update_address = "inproc://scale-down-stats"
274
+ client.first_req_sock_addr = "inproc://scale-down-first-request"
275
+ pause_started = asyncio.Event()
276
+
277
+ async def pause_scheduler(*args, **kwargs):
278
+ pause_started.set()
279
+ await asyncio.Future()
280
+
281
+ client._call_utility_async = pause_scheduler
282
+
283
+ with (
284
+ zmq.asyncio.Context() as ctx,
285
+ ctx.socket(zmq.XPUB) as coordinator,
286
+ ctx.socket(zmq.PAIR) as first_req_socket,
287
+ ):
288
+ client.ctx = ctx
289
+ coordinator.bind(client.stats_update_address)
290
+ first_req_socket.bind(client.first_req_sock_addr)
291
+ client._ensure_stats_update_task()
292
+ scale_task = None
293
+ try:
294
+ # Wait for subscription so the snapshot cannot be dropped.
295
+ await asyncio.wait_for(coordinator.recv(), timeout=5)
296
+ scale_task = asyncio.create_task(client._commit_scale_down_elastic_ep(2))
297
+ await asyncio.wait_for(pause_started.wait(), timeout=5)
298
+
299
+ # The coordinator still knows about all four engines while the
300
+ # client has already stopped routing to the two removed engines.
301
+ counts = [[0, 0, 0.0], [0, 0, 0.0], [1, 0, 0.0], [1, 0, 0.0]]
302
+ await coordinator.send(msgspec.msgpack.encode((counts, 1, True)))
303
+
304
+ async def wait_for_snapshot():
305
+ while client.current_wave != 1:
306
+ await asyncio.sleep(0)
307
+
308
+ await asyncio.wait_for(wait_for_snapshot(), timeout=5)
309
+ chosen = client.get_core_engine_for_request(
310
+ _make_pooling_request("scale-down-request")
311
+ )
312
+ assert chosen == client.core_engines[0]
313
+ finally:
314
+ tasks = [client.resources.stats_update_task]
315
+ if scale_task is not None:
316
+ tasks.append(scale_task)
317
+ for task in tasks:
318
+ task.cancel()
319
+ await asyncio.gather(*tasks, return_exceptions=True)
320
+
321
+
322
def test_dplb_snapshot_backpressure_overrides_inflight():
323
"""An engine reported heavily loaded by the coordinator is avoided even
324
when this client has routed nothing to it."""
vllm/v1/engine/core_client.py
+2
@@ -1914,6 +1914,8 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1914
# NOTE(yongji): Immediately stop sending requests to the removing engines.
1915
self.core_engines = old_core_engines[:new_data_parallel_size]
1916
self.lb_engines = self.lb_engines[:new_data_parallel_size]
1917
+ # Pending coordinator snapshots must use the surviving ranks too.
1918
+ self.engine_ranks_managed = self.engine_ranks_managed[:new_data_parallel_size]
1919
removed_dp_size = cur_data_parallel_size - new_data_parallel_size
1920
pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size
1921
pause_futures = [