[Performance][EPD] Reduce Python proxy serialization overhead (#56657)
Signed-off-by: Tianyu Guo <guoty@inferact.ai>
Tianyu Guo committed
Sep 14, 2026 at 06:07 UTC
319cc5ef19946d34c2e66cbbec5bda29d0bfa328
2 files changed
+160
-25
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+32
-20
@@ -33,9 +33,10 @@ from collections.abc import AsyncIterator
33
from typing import Any
34
35
import aiohttp
36
+import msgspec
37
import uvicorn
38
from fastapi import FastAPI, HTTPException, Request
38
-from fastapi.responses import JSONResponse, StreamingResponse
39
+from fastapi.responses import JSONResponse, Response, StreamingResponse
40
41
###############################################################################
42
# FastAPI app & global state
@@ -264,7 +265,7 @@ async def fanout_encoder_primer(
265
for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)):
266
# Derive a *child* request id: <parent>:<index>:<random-short>
267
child_req_id = f"{req_id}:{idx}:{uuid.uuid4().hex[:6]}"
267
- headers = {"x-request-id": child_req_id}
268
+ headers = {"x-request-id": child_req_id, "Content-Type": "application/json"}
269
270
# With --no-rewrite the decoder still receives the raw image and derives
271
# the cache key by hashing it, so the encoder must do the same -- passing
@@ -307,7 +308,7 @@ async def fanout_encoder_primer(
308
tasks.append(
309
encode_session.post(
310
f"{target_url}/v1/chat/completions",
310
- json=encoder_req,
311
+ data=msgspec.json.encode(encoder_req),
312
headers=headers,
313
)
314
)
@@ -347,7 +348,7 @@ async def fanout_encoder_primer(
348
# The encoder reports each mm_hash's metadata (e.g. the grid) here,
349
# keyed by the same uuid this proxy assigned above.
350
try:
350
- params = (await r.json()).get("ec_transfer_params") or {}
351
+ params = msgspec.json.decode(await r.read()).get("ec_transfer_params") or {}
352
except Exception:
353
logger.warning("[%s] Could not read encoder metadata #%d", req_id, idx)
354
params = {}
@@ -407,7 +408,7 @@ async def maybe_prefill(
408
409
prefill_response = await process_prefill_stage(req_data, p_url, req_id)
410
# for nixl connector to facilitate kv transfer...
410
- prefill_response_json = await prefill_response.json()
411
+ prefill_response_json = msgspec.json.decode(await prefill_response.read())
412
kv_transfer_params = prefill_response_json.get("kv_transfer_params", {})
413
if kv_transfer_params:
414
return {**req_data, "kv_transfer_params": kv_transfer_params}
@@ -439,10 +440,12 @@ async def process_prefill_stage(
440
if "stream_options" in prefill_request:
441
del prefill_request["stream_options"]
442
442
- headers = {"x-request-id": req_id}
443
+ headers = {"x-request-id": req_id, "Content-Type": "application/json"}
444
try:
445
prefill_response = await prefill_session.post(
445
- f"{p_url}/v1/chat/completions", json=prefill_request, headers=headers
446
+ f"{p_url}/v1/chat/completions",
447
+ data=msgspec.json.encode(prefill_request),
448
+ headers=headers,
449
)
450
prefill_response.raise_for_status()
451
@@ -598,7 +601,7 @@ async def forward_non_stream(
601
d_url: str,
602
consumer_zmq: str | None,
603
dp_rank: int | None = None,
601
-) -> dict:
604
+) -> Response:
605
try:
606
for attempt in range(DECODE_RETRIES + 1):
607
_t0 = time.perf_counter()
@@ -608,12 +611,14 @@ async def forward_non_stream(
611
_t2 = time.perf_counter()
612
613
logger.info("[%s] Forwarding to decode: %s", req_id, d_url)
611
- headers = {"x-request-id": req_id}
614
+ headers = {"x-request-id": req_id, "Content-Type": "application/json"}
615
if dp_rank is not None:
616
headers["X-data-parallel-rank"] = str(dp_rank)
617
618
async with decode_session.post(
616
- f"{d_url}/v1/chat/completions", json=prepared, headers=headers
619
+ f"{d_url}/v1/chat/completions",
620
+ data=msgspec.json.encode(prepared),
621
+ headers=headers,
622
) as resp:
623
if resp.status >= 400:
624
detail = await resp.text()
@@ -637,7 +642,7 @@ async def forward_non_stream(
642
detail,
643
)
644
raise HTTPException(status_code=resp.status, detail=detail)
640
- out = await resp.json()
645
+ out = await resp.read()
646
_t3 = time.perf_counter()
647
logger.info(
648
"STAGE %s encode=%.1f rewrite=%.1f decode=%.1f total=%.1f "
@@ -649,7 +654,15 @@ async def forward_non_stream(
654
(_t3 - _t0) * 1e3,
655
attempt,
656
)
652
- return out
657
+ return Response(
658
+ content=out,
659
+ status_code=resp.status,
660
+ headers={
661
+ "Content-Type": resp.headers.get(
662
+ "Content-Type", "application/json"
663
+ )
664
+ },
665
+ )
666
raise HTTPException(status_code=500, detail="Decode failed after re-encoding")
667
668
except HTTPException:
@@ -667,7 +680,7 @@ async def forward_stream(
680
d_url: str,
681
consumer_zmq: str | None,
682
dp_rank: int | None = None,
670
-) -> AsyncIterator[str]:
683
+) -> AsyncIterator[bytes]:
684
try:
685
for attempt in range(DECODE_RETRIES + 1):
686
_t0 = time.perf_counter()
@@ -677,14 +690,14 @@ async def forward_stream(
690
_t2 = time.perf_counter()
691
692
logger.info("[%s] Starting streaming from decode: %s", req_id, d_url)
680
- headers = {"x-request-id": req_id}
693
+ headers = {"x-request-id": req_id, "Content-Type": "application/json"}
694
if dp_rank is not None:
695
headers["X-data-parallel-rank"] = str(dp_rank)
696
697
_first = None
698
async with decode_session.post(
699
f"{d_url}/v1/chat/completions",
687
- json=prepared,
700
+ data=msgspec.json.encode(prepared),
701
headers=headers,
702
) as resp:
703
# Retry only before the first chunk: once anything reached the
@@ -701,11 +714,11 @@ async def forward_stream(
714
)
715
continue
716
resp.raise_for_status()
704
- async for chunk in resp.content.iter_chunked(1024):
717
+ async for chunk in resp.content.iter_any():
718
if chunk:
719
if _first is None:
720
_first = time.perf_counter()
708
- yield chunk.decode("utf-8", errors="ignore")
721
+ yield chunk
722
_t3 = time.perf_counter()
723
724
logger.info(
@@ -739,7 +752,7 @@ async def forward_stream(
752
@app.post("/v1/chat/completions")
753
async def chat_completions(request: Request):
754
try:
742
- req_data = await request.json()
755
+ req_data = msgspec.json.decode(await request.body())
756
req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
757
758
e_urls = app.state.e_urls # we want the full list for fan-out
@@ -764,10 +777,9 @@ async def chat_completions(request: Request):
777
),
778
media_type="text/event-stream",
779
)
767
- result = await forward_non_stream(
780
+ return await forward_non_stream(
781
req_data, req_id, e_urls, p_url, d_url, consumer_zmq, dp_rank
782
)
770
- return JSONResponse(content=result)
783
784
except HTTPException:
785
raise
tests/v1/ec_connector/unit/test_epd_proxy_retry.py
+128
-5
@@ -10,7 +10,11 @@ import asyncio
10
import importlib.util
11
from pathlib import Path
12
13
+import httpx
14
+import msgspec
15
import pytest
16
+from aiohttp import web
17
+from aiohttp.test_utils import TestServer
18
19
PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py"
20
@@ -29,8 +33,8 @@ class _Response:
33
def __init__(self, params):
34
self._params = params
35
32
- async def json(self):
33
- return {"kv_transfer_params": self._params}
36
+ async def read(self):
37
+ return msgspec.json.encode({"kv_transfer_params": self._params})
38
39
40
def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch):
@@ -62,8 +66,8 @@ class _EncoderResponse:
66
self.status = 200
67
self._params = params
68
65
- async def json(self):
66
- return {"ec_transfer_params": self._params}
69
+ async def read(self):
70
+ return msgspec.json.encode({"ec_transfer_params": self._params})
71
72
async def text(self):
73
return ""
@@ -75,7 +79,7 @@ class _EncoderSession:
79
def __init__(self, replies):
80
self._replies = list(replies)
81
78
- async def post(self, url, json=None, headers=None):
82
+ async def post(self, url, data=None, headers=None):
83
return _EncoderResponse(self._replies.pop(0))
84
85
@@ -147,6 +151,125 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
151
assert params["ec_items"][0]["transfer_id"]
152
153
154
+@pytest.mark.asyncio
155
+@pytest.mark.parametrize("stream", [False, True])
156
+@pytest.mark.parametrize("prefill", [False, True])
157
+async def test_http_roundtrip_preserves_payload_and_response_bytes(
158
+ proxy, monkeypatch, stream, prefill
159
+):
160
+ """Exercise real HTTP hops, rewrite and a decode retry without model servers."""
161
+ seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []}
162
+ prefix = b'data: {"content":"'
163
+ # Split a multibyte character at the old 1024-byte forwarding boundary.
164
+ expected = (
165
+ prefix + b"x" * (1023 - len(prefix)) + '图片🌍"}\n\ndata: [DONE]\n\n'.encode()
166
+ if stream
167
+ else '{ "choices": [{"message": {"content": "图片🌍"}}] }\n'.encode()
168
+ )
169
+
170
+ async def backend(request):
171
+ stage = request.match_info["stage"]
172
+ assert request.content_type == "application/json"
173
+ body = await request.json()
174
+ seen[stage].append(body)
175
+ if stage == "encode":
176
+ item = body["messages"][0]["content"][0]
177
+ return web.json_response(
178
+ {
179
+ "ec_transfer_params": {
180
+ item["uuid"]: {
181
+ "metadata": {"image_grid_thw": [[1, 2, 2]]},
182
+ "peer_port": len(seen[stage]),
183
+ }
184
+ }
185
+ }
186
+ )
187
+ if stage == "prefill":
188
+ assert body["max_tokens"] == 1 and body["stream"] is False
189
+ return web.json_response(
190
+ {"kv_transfer_params": {"remote_block_ids": [len(seen[stage])]}}
191
+ )
192
+ if len(seen[stage]) == 1:
193
+ return web.Response(status=500, text="retry")
194
+ return web.Response(
195
+ body=expected,
196
+ content_type="text/event-stream" if stream else "application/json",
197
+ )
198
+
199
+ backend_app = web.Application()
200
+ backend_app.router.add_post("/{stage}/v1/chat/completions", backend)
201
+ async with TestServer(backend_app) as server:
202
+ for stage, field in [
203
+ ("encode", "e_urls"),
204
+ ("prefill", "p_urls"),
205
+ ("decode", "d_urls"),
206
+ ]:
207
+ urls = [str(server.make_url(f"/{stage}"))]
208
+ if stage == "prefill" and not prefill:
209
+ urls = []
210
+ monkeypatch.setattr(proxy.app.state, field, urls, raising=False)
211
+ monkeypatch.setattr(proxy.app.state, "d_ec_urls", [], raising=False)
212
+ monkeypatch.setattr(proxy.app.state, "ec_consumer_dp_size", 1, raising=False)
213
+ monkeypatch.setattr(proxy, "DECODE_RETRIES", 1)
214
+ monkeypatch.setattr(proxy, "NO_REWRITE", False)
215
+ item = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJj"}}
216
+ body = {
217
+ "model": "test",
218
+ "messages": [
219
+ {
220
+ "role": "user",
221
+ "content": [
222
+ {"type": "text", "text": "描述图片🌍"},
223
+ item,
224
+ item,
225
+ ],
226
+ }
227
+ ],
228
+ "stream": stream,
229
+ "temperature": 0.01,
230
+ "max_tokens": 32,
231
+ "seed": 42,
232
+ "structured_outputs": {"choice": ["A", "B"]},
233
+ }
234
+ await proxy.on_startup()
235
+ try:
236
+ async with httpx.AsyncClient(
237
+ transport=httpx.ASGITransport(app=proxy.app), base_url="http://proxy"
238
+ ) as client:
239
+ response = await client.post("/v1/chat/completions", json=body)
240
+ assert response.status_code == 200
241
+ assert response.content == expected
242
+ assert response.headers["content-type"].startswith(
243
+ "text/event-stream" if stream else "application/json"
244
+ )
245
+ finally:
246
+ await proxy.on_shutdown()
247
+
248
+ assert len(seen["encode"]) == 4
249
+ assert len(seen["decode"]) == 2
250
+ final = seen["decode"][-1]
251
+ for key in [
252
+ "model",
253
+ "stream",
254
+ "temperature",
255
+ "max_tokens",
256
+ "seed",
257
+ "structured_outputs",
258
+ ]:
259
+ assert final[key] == body[key]
260
+ content = final["messages"][0]["content"]
261
+ assert content[0] == body["messages"][0]["content"][0]
262
+ for reference in content[1:]:
263
+ assert reference == {
264
+ "type": "image_embeds",
265
+ "image_embeds": {"image_grid_thw": [1, 2, 2]},
266
+ "uuid": proxy.content_uuid(item),
267
+ }
268
+ assert final["ec_transfer_params"][proxy.content_uuid(item)]["peer_port"] in (3, 4)
269
+ if prefill:
270
+ assert final["kv_transfer_params"] == {"remote_block_ids": [2]}
271
+
272
+
273
@pytest.mark.parametrize("server_keep_alive", ["0.1", "1", "5", "2", "30"])
274
def test_pooled_connections_are_retired_before_the_server_closes_them(
275
proxy, monkeypatch, server_keep_alive