[Bugfix][EPD] Fix encoder round-robin fan-out (#52491)
Signed-off-by: AnkitNakhawa <ankitnakhawa@gmail.com>
Ankit Nakhawa committed
Aug 20, 2026 at 03:30 UTC
d66300a1baa7779c68c7dfa4e51eee2502b48017
2 files changed
+110
-2
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+28
-2
@@ -51,6 +51,11 @@ encode_session: aiohttp.ClientSession | None = None
51
prefill_session: aiohttp.ClientSession | None = None
52
decode_session: aiohttp.ClientSession | None = None
53
54
+# Cursor for round-robin encoder assignment, shared across requests so the
55
+# fan-out doesn't restart from e_urls[0] every time.
56
+encoder_rr_idx = 0
57
+encoder_rr_lock = asyncio.Lock()
58
+
59
###############################################################################
60
# Utils
61
###############################################################################
@@ -58,6 +63,21 @@ decode_session: aiohttp.ClientSession | None = None
63
64
MM_TYPES = {"image_url", "audio_url", "input_audio"}
65
66
+
67
+def encoder_rr_assignment(
68
+ e_urls: list[str], start: int, count: int
69
+) -> tuple[list[str], int]:
70
+ """Assign `count` items to encoder URLs starting from cursor `start`.
71
+
72
+ Returns the per-item URL list and the cursor value the next call should
73
+ start from, so the assignment is contiguous across calls instead of
74
+ restarting at e_urls[0] every time.
75
+ """
76
+ urls = [e_urls[(start + i) % len(e_urls)] for i in range(count)]
77
+ next_start = (start + count) % len(e_urls)
78
+ return urls, next_start
79
+
80
+
81
# Diagnostic switch: forward the original request to the decoder so the
82
# only difference from the rewrite path is the rewrite itself.
83
NO_REWRITE = False
@@ -189,8 +209,14 @@ async def fanout_encoder_primer(
209
item_uuids: dict[int, str] = {}
210
item_meta: dict[int, dict] = {}
211
192
- # Round-robin over encode servers to distribute load a bit
193
- url_cycle = (e_urls[i % len(e_urls)] for i in range(len(mm_items)))
212
+ # Round-robin over encode servers to distribute load a bit. The cursor
213
+ # persists across requests so fan-out doesn't restart at e_urls[0] every
214
+ # time (which would hot-spot the first encoder for single-item requests).
215
+ global encoder_rr_idx
216
+ async with encoder_rr_lock:
217
+ url_cycle, encoder_rr_idx = encoder_rr_assignment(
218
+ e_urls, encoder_rr_idx, len(mm_items)
219
+ )
220
221
for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)):
222
# Derive a *child* request id: <parent>:<index>:<random-short>
tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py
new
+82
@@ -0,0 +1,82 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Encoder fan-out fairness in the disaggregated EPD proxy.
4
+
5
+Exercises the REAL ``encoder_rr_assignment`` helper loaded from its
6
+``examples/`` path -- no routing logic is re-implemented here, so a future
7
+change to that function is what these tests exercise.
8
+
9
+Regression target: the fan-out cursor used to restart at e_urls[0] on every
10
+incoming request (``e_urls[i % len(e_urls)] for i in range(len(mm_items))``),
11
+so single-item requests -- the common case -- always hit the first encoder
12
+instance and left the rest idle. ``encoder_rr_assignment`` threads a cursor
13
+across calls so consecutive requests keep advancing through every URL.
14
+"""
15
+
16
+import importlib.util
17
+from collections import Counter
18
+from pathlib import Path
19
+
20
+import pytest
21
+
22
+PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py"
23
+
24
+
25
+def _load_proxy_module():
26
+ path = Path(__file__).parents[4] / PROXY_REL
27
+ spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path)
28
+ assert spec is not None and spec.loader is not None
29
+ module = importlib.util.module_from_spec(spec)
30
+ spec.loader.exec_module(module)
31
+ return module
32
+
33
+
34
+@pytest.fixture(scope="module")
35
+def assign():
36
+ return _load_proxy_module().encoder_rr_assignment
37
+
38
+
39
+def _drive(assign, e_urls, counts):
40
+ """Feed a sequence of request item-counts through the cursor, as
41
+ fanout_encoder_primer does one request at a time."""
42
+ cursor = 0
43
+ all_urls = []
44
+ for count in counts:
45
+ urls, cursor = assign(e_urls, cursor, count)
46
+ assert len(urls) == count
47
+ all_urls.append(urls)
48
+ return all_urls
49
+
50
+
51
+@pytest.mark.parametrize("n_urls", [1, 2, 3, 5])
52
+def test_full_url_space_is_covered_uniformly(assign, n_urls):
53
+ e_urls = [f"E{i}" for i in range(n_urls)]
54
+ # Three full cycles of single-item requests.
55
+ urls = _drive(assign, e_urls, counts=[1] * (n_urls * 3))
56
+ hits = Counter(u for req in urls for u in req)
57
+ assert set(hits) == set(e_urls)
58
+ assert max(hits.values()) == min(hits.values())
59
+
60
+
61
+def test_single_item_requests_rotate_through_all_encoders(assign):
62
+ # Exact scenario from the bug report: 3 encoders, one image/request.
63
+ e_urls = ["E0", "E1", "E2"]
64
+ routed = _drive(assign, e_urls, counts=[1] * 6)
65
+ assert [urls[0] for urls in routed] == ["E0", "E1", "E2", "E0", "E1", "E2"]
66
+
67
+
68
+def test_cursor_stays_contiguous_across_varying_item_counts(assign):
69
+ e_urls = ["E0", "E1", "E2"]
70
+ routed = _drive(assign, e_urls, counts=[2, 1, 3, 1])
71
+ assert routed == [
72
+ ["E0", "E1"],
73
+ ["E2"],
74
+ ["E0", "E1", "E2"],
75
+ ["E0"],
76
+ ]
77
+
78
+
79
+def test_single_encoder_always_resolves_to_it(assign):
80
+ urls, next_cursor = assign(["E0"], 0, 4)
81
+ assert urls == ["E0"] * 4
82
+ assert next_cursor == 0