diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index 9a1bd51721..267176fd3e 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -51,6 +51,11 @@ encode_session: aiohttp.ClientSession | None = None prefill_session: aiohttp.ClientSession | None = None decode_session: aiohttp.ClientSession | None = None +# Cursor for round-robin encoder assignment, shared across requests so the +# fan-out doesn't restart from e_urls[0] every time. +encoder_rr_idx = 0 +encoder_rr_lock = asyncio.Lock() + ############################################################################### # Utils ############################################################################### @@ -58,6 +63,21 @@ decode_session: aiohttp.ClientSession | None = None MM_TYPES = {"image_url", "audio_url", "input_audio"} + +def encoder_rr_assignment( + e_urls: list[str], start: int, count: int +) -> tuple[list[str], int]: + """Assign `count` items to encoder URLs starting from cursor `start`. + + Returns the per-item URL list and the cursor value the next call should + start from, so the assignment is contiguous across calls instead of + restarting at e_urls[0] every time. + """ + urls = [e_urls[(start + i) % len(e_urls)] for i in range(count)] + next_start = (start + count) % len(e_urls) + return urls, next_start + + # Diagnostic switch: forward the original request to the decoder so the # only difference from the rewrite path is the rewrite itself. NO_REWRITE = False @@ -189,8 +209,14 @@ async def fanout_encoder_primer( item_uuids: dict[int, str] = {} item_meta: dict[int, dict] = {} - # Round-robin over encode servers to distribute load a bit - url_cycle = (e_urls[i % len(e_urls)] for i in range(len(mm_items))) + # Round-robin over encode servers to distribute load a bit. The cursor + # persists across requests so fan-out doesn't restart at e_urls[0] every + # time (which would hot-spot the first encoder for single-item requests). + global encoder_rr_idx + async with encoder_rr_lock: + url_cycle, encoder_rr_idx = encoder_rr_assignment( + e_urls, encoder_rr_idx, len(mm_items) + ) for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)): # Derive a *child* request id: :: diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py new file mode 100644 index 0000000000..48123235a7 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder fan-out fairness in the disaggregated EPD proxy. + +Exercises the REAL ``encoder_rr_assignment`` helper loaded from its +``examples/`` path -- no routing logic is re-implemented here, so a future +change to that function is what these tests exercise. + +Regression target: the fan-out cursor used to restart at e_urls[0] on every +incoming request (``e_urls[i % len(e_urls)] for i in range(len(mm_items))``), +so single-item requests -- the common case -- always hit the first encoder +instance and left the rest idle. ``encoder_rr_assignment`` threads a cursor +across calls so consecutive requests keep advancing through every URL. +""" + +import importlib.util +from collections import Counter +from pathlib import Path + +import pytest + +PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" + + +def _load_proxy_module(): + path = Path(__file__).parents[4] / PROXY_REL + spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def assign(): + return _load_proxy_module().encoder_rr_assignment + + +def _drive(assign, e_urls, counts): + """Feed a sequence of request item-counts through the cursor, as + fanout_encoder_primer does one request at a time.""" + cursor = 0 + all_urls = [] + for count in counts: + urls, cursor = assign(e_urls, cursor, count) + assert len(urls) == count + all_urls.append(urls) + return all_urls + + +@pytest.mark.parametrize("n_urls", [1, 2, 3, 5]) +def test_full_url_space_is_covered_uniformly(assign, n_urls): + e_urls = [f"E{i}" for i in range(n_urls)] + # Three full cycles of single-item requests. + urls = _drive(assign, e_urls, counts=[1] * (n_urls * 3)) + hits = Counter(u for req in urls for u in req) + assert set(hits) == set(e_urls) + assert max(hits.values()) == min(hits.values()) + + +def test_single_item_requests_rotate_through_all_encoders(assign): + # Exact scenario from the bug report: 3 encoders, one image/request. + e_urls = ["E0", "E1", "E2"] + routed = _drive(assign, e_urls, counts=[1] * 6) + assert [urls[0] for urls in routed] == ["E0", "E1", "E2", "E0", "E1", "E2"] + + +def test_cursor_stays_contiguous_across_varying_item_counts(assign): + e_urls = ["E0", "E1", "E2"] + routed = _drive(assign, e_urls, counts=[2, 1, 3, 1]) + assert routed == [ + ["E0", "E1"], + ["E2"], + ["E0", "E1", "E2"], + ["E0"], + ] + + +def test_single_encoder_always_resolves_to_it(assign): + urls, next_cursor = assign(["E0"], 0, 4) + assert urls == ["E0"] * 4 + assert next_cursor == 0