[Bugfix][EPD] Preserve media processing options in encoder requests (#56786)
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
jiangkuaixue123 committed
Sep 14, 2026 at 18:05 UTC
1678b396270406c27fcab8f5b86b21fd305ac605
2 files changed
+45
-7
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+9
-2
@@ -278,7 +278,6 @@ async def fanout_encoder_primer(
278
item_transfer_ids[idx] = transfer_id
279
280
encoder_req = {
281
- # You *may* need to keep additional fields
281
"model": orig_request.get("model"),
282
"messages": [
283
{
@@ -292,6 +291,14 @@ async def fanout_encoder_primer(
291
# once the prompt is encoded and its embeddings are published.
292
"stream": False,
293
}
294
+ for key in (
295
+ "mm_processor_kwargs",
296
+ "media_io_kwargs",
297
+ "priority",
298
+ "session_id",
299
+ ):
300
+ if key in orig_request:
301
+ encoder_req[key] = orig_request[key]
302
if consumer_zmq is not None:
303
# No mm_hash here on purpose. The encoder's own
304
# `mm_features[i].identifier` is derived from the uuid *and* the
@@ -378,7 +385,7 @@ async def fanout_encoder_primer(
385
# connector's own handle on the published embedding (for NIXL,
386
# peer_host/peer_port/size_bytes). The decoder's connector
387
# looks it up by mm_hash on the request, so carry it through.
381
- ec_params[item_uuids.get(idx, ec_mm_hash)] = reported
388
+ ec_params[ec_mm_hash] = reported
389
if NO_REWRITE and consumer_zmq is not None:
390
ec_params.setdefault("ec_items", []).append(
391
{"mm_hash": ec_mm_hash, "transfer_id": item_transfer_ids[idx]}
tests/v1/ec_connector/unit/test_epd_proxy_retry.py
+36
-5
@@ -116,7 +116,9 @@ def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch
116
117
first, _, _ = asyncio.run(proxy.prepare_for_decode(body, *args))
118
reported = first["ec_transfer_params"]
119
- assert [handle] == [value for key, value in reported.items() if key != "ec_items"]
119
+ assert {key: value for key, value in reported.items() if key != "ec_items"} == {
120
+ "encoder-side-hash": handle
121
+ }
122
assert "ec_transfer_params" not in body
123
124
# Attempt 2's encode reports nothing: decode must be told nothing.
@@ -154,10 +156,22 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
156
@pytest.mark.asyncio
157
@pytest.mark.parametrize("stream", [False, True])
158
@pytest.mark.parametrize("prefill", [False, True])
159
+@pytest.mark.parametrize(
160
+ "request_options",
161
+ [
162
+ {},
163
+ {
164
+ "mm_processor_kwargs": {"max_pixels": 262144},
165
+ "media_io_kwargs": {"image": {"image_mode": "RGB"}},
166
+ "priority": -3,
167
+ "session_id": "session-123",
168
+ },
169
+ ],
170
+)
171
async def test_http_roundtrip_preserves_payload_and_response_bytes(
158
- proxy, monkeypatch, stream, prefill
172
+ proxy, monkeypatch, stream, prefill, request_options
173
):
160
- """Exercise real HTTP hops, rewrite and a decode retry without model servers."""
174
+ """Preserve media and scheduling context across HTTP hops and retries."""
175
seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []}
176
prefix = b'data: {"content":"'
177
# Split a multibyte character at the old 1024-byte forwarding boundary.
@@ -174,10 +188,13 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
188
seen[stage].append(body)
189
if stage == "encode":
190
item = body["messages"][0]["content"][0]
191
+ mm_hash = item["uuid"]
192
+ if body.get("mm_processor_kwargs"):
193
+ mm_hash += "-processed"
194
return web.json_response(
195
{
196
"ec_transfer_params": {
180
- item["uuid"]: {
197
+ mm_hash: {
198
"metadata": {"image_grid_thw": [[1, 2, 2]]},
199
"peer_port": len(seen[stage]),
200
}
@@ -230,6 +247,7 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
247
"max_tokens": 32,
248
"seed": 42,
249
"structured_outputs": {"choice": ["A", "B"]},
250
+ **request_options,
251
}
252
await proxy.on_startup()
253
try:
@@ -247,6 +265,15 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
265
266
assert len(seen["encode"]) == 4
267
assert len(seen["decode"]) == 2
268
+ for requests in seen.values():
269
+ for forwarded in requests:
270
+ for key in (
271
+ "mm_processor_kwargs",
272
+ "media_io_kwargs",
273
+ "priority",
274
+ "session_id",
275
+ ):
276
+ assert forwarded.get(key) == body.get(key)
277
final = seen["decode"][-1]
278
for key in [
279
"model",
@@ -265,7 +292,11 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
292
"image_embeds": {"image_grid_thw": [1, 2, 2]},
293
"uuid": proxy.content_uuid(item),
294
}
268
- assert final["ec_transfer_params"][proxy.content_uuid(item)]["peer_port"] in (3, 4)
295
+ mm_hash = proxy.content_uuid(item)
296
+ if request_options:
297
+ mm_hash += "-processed"
298
+ assert set(final["ec_transfer_params"]) - {"ec_items"} == {mm_hash}
299
+ assert final["ec_transfer_params"][mm_hash]["peer_port"] in (3, 4)
300
if prefill:
301
assert final["kv_transfer_params"] == {"remote_block_ids": [2]}
302