iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Performance][EPD] Reduce Python proxy serialization overhead (#56657)

Browse Source 
Signed-off-by: Tianyu Guo <guoty@inferact.ai>
This commit is contained in:
Tianyu Guo
2026-09-13 15:07:12 -07:00
committed by GitHub
parent b23433088b
commit 319cc5ef19
2 changed files with 160 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+32 -20
View File
@@ -33,9 +33,10 @@ from collections.abc import AsyncIterator
from typing import Any
import aiohttp
import msgspec
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.responses import JSONResponse, Response, StreamingResponse
###############################################################################
# FastAPI app & global state
@@ -264,7 +265,7 @@ async def fanout_encoder_primer(
for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)):
# Derive a *child* request id: <parent>:<index>:<random-short>
child_req_id = f"{req_id}:{idx}:{uuid.uuid4().hex[:6]}"
headers = {"x-request-id": child_req_id}
headers = {"x-request-id": child_req_id, "Content-Type": "application/json"}
# With --no-rewrite the decoder still receives the raw image and derives
# the cache key by hashing it, so the encoder must do the same -- passing
@@ -307,7 +308,7 @@ async def fanout_encoder_primer(
tasks.append(
encode_session.post(
f"{target_url}/v1/chat/completions",
json=encoder_req,
data=msgspec.json.encode(encoder_req),
headers=headers,
)
)
@@ -347,7 +348,7 @@ async def fanout_encoder_primer(
# The encoder reports each mm_hash's metadata (e.g. the grid) here,
# keyed by the same uuid this proxy assigned above.
try:
params = (await r.json()).get("ec_transfer_params") or {}
params = msgspec.json.decode(await r.read()).get("ec_transfer_params") or {}
except Exception:
logger.warning("[%s] Could not read encoder metadata #%d", req_id, idx)
params = {}
@@ -407,7 +408,7 @@ async def maybe_prefill(
prefill_response = await process_prefill_stage(req_data, p_url, req_id)
# for nixl connector to facilitate kv transfer...
prefill_response_json = await prefill_response.json()
prefill_response_json = msgspec.json.decode(await prefill_response.read())
kv_transfer_params = prefill_response_json.get("kv_transfer_params", {})
if kv_transfer_params:
return {**req_data, "kv_transfer_params": kv_transfer_params}
@@ -439,10 +440,12 @@ async def process_prefill_stage(
if "stream_options" in prefill_request:
del prefill_request["stream_options"]
headers = {"x-request-id": req_id}
headers = {"x-request-id": req_id, "Content-Type": "application/json"}
try:
prefill_response = await prefill_session.post(
f"{p_url}/v1/chat/completions", json=prefill_request, headers=headers
f"{p_url}/v1/chat/completions",
data=msgspec.json.encode(prefill_request),
headers=headers,
)
prefill_response.raise_for_status()
@@ -598,7 +601,7 @@ async def forward_non_stream(
d_url: str,
consumer_zmq: str | None,
dp_rank: int | None = None,
) -> dict:
) -> Response:
try:
for attempt in range(DECODE_RETRIES + 1):
_t0 = time.perf_counter()
@@ -608,12 +611,14 @@ async def forward_non_stream(
_t2 = time.perf_counter()
logger.info("[%s] Forwarding to decode: %s", req_id, d_url)
headers = {"x-request-id": req_id}
headers = {"x-request-id": req_id, "Content-Type": "application/json"}
if dp_rank is not None:
headers["X-data-parallel-rank"] = str(dp_rank)
async with decode_session.post(
f"{d_url}/v1/chat/completions", json=prepared, headers=headers
f"{d_url}/v1/chat/completions",
data=msgspec.json.encode(prepared),
headers=headers,
) as resp:
if resp.status >= 400:
detail = await resp.text()
@@ -637,7 +642,7 @@ async def forward_non_stream(
detail,
)
raise HTTPException(status_code=resp.status, detail=detail)
out = await resp.json()
out = await resp.read()
_t3 = time.perf_counter()
logger.info(
"STAGE %s encode=%.1f rewrite=%.1f decode=%.1f total=%.1f "
@@ -649,7 +654,15 @@ async def forward_non_stream(
(_t3 - _t0) * 1e3,
attempt,
)
return out
return Response(
content=out,
status_code=resp.status,
headers={
"Content-Type": resp.headers.get(
"Content-Type", "application/json"
)
},
)
raise HTTPException(status_code=500, detail="Decode failed after re-encoding")
except HTTPException:
@@ -667,7 +680,7 @@ async def forward_stream(
d_url: str,
consumer_zmq: str | None,
dp_rank: int | None = None,
) -> AsyncIterator[str]:
) -> AsyncIterator[bytes]:
try:
for attempt in range(DECODE_RETRIES + 1):
_t0 = time.perf_counter()
@@ -677,14 +690,14 @@ async def forward_stream(
_t2 = time.perf_counter()
logger.info("[%s] Starting streaming from decode: %s", req_id, d_url)
headers = {"x-request-id": req_id}
headers = {"x-request-id": req_id, "Content-Type": "application/json"}
if dp_rank is not None:
headers["X-data-parallel-rank"] = str(dp_rank)
_first = None
async with decode_session.post(
f"{d_url}/v1/chat/completions",
json=prepared,
data=msgspec.json.encode(prepared),
headers=headers,
) as resp:
# Retry only before the first chunk: once anything reached the
@@ -701,11 +714,11 @@ async def forward_stream(
)
continue
resp.raise_for_status()
async for chunk in resp.content.iter_chunked(1024):
async for chunk in resp.content.iter_any():
if chunk:
if _first is None:
_first = time.perf_counter()
yield chunk.decode("utf-8", errors="ignore")
yield chunk
_t3 = time.perf_counter()
logger.info(
@@ -739,7 +752,7 @@ async def forward_stream(
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
try:
req_data = await request.json()
req_data = msgspec.json.decode(await request.body())
req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
e_urls = app.state.e_urls # we want the full list for fan-out
@@ -764,10 +777,9 @@ async def chat_completions(request: Request):
),
media_type="text/event-stream",
)
result = await forward_non_stream(
return await forward_non_stream(
req_data, req_id, e_urls, p_url, d_url, consumer_zmq, dp_rank
)
return JSONResponse(content=result)
except HTTPException:
raise
tests/v1/ec_connector/unit/test_epd_proxy_retry.py
+128 -5
View File
@@ -10,7 +10,11 @@ import asyncio
import importlib.util
from pathlib import Path
import httpx
import msgspec
import pytest
from aiohttp import web
from aiohttp.test_utils import TestServer
PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py"
@@ -29,8 +33,8 @@ class _Response:
def __init__(self, params):
self._params = params
async def json(self):
return {"kv_transfer_params": self._params}
async def read(self):
return msgspec.json.encode({"kv_transfer_params": self._params})
def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch):
@@ -62,8 +66,8 @@ class _EncoderResponse:
self.status = 200
self._params = params
async def json(self):
return {"ec_transfer_params": self._params}
async def read(self):
return msgspec.json.encode({"ec_transfer_params": self._params})
async def text(self):
return ""
@@ -75,7 +79,7 @@ class _EncoderSession:
def __init__(self, replies):
self._replies = list(replies)
async def post(self, url, json=None, headers=None):
async def post(self, url, data=None, headers=None):
return _EncoderResponse(self._replies.pop(0))
@@ -147,6 +151,125 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
assert params["ec_items"][0]["transfer_id"]
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.parametrize("prefill", [False, True])
async def test_http_roundtrip_preserves_payload_and_response_bytes(
proxy, monkeypatch, stream, prefill
):
"""Exercise real HTTP hops, rewrite and a decode retry without model servers."""
seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []}
prefix = b'data: {"content":"'
# Split a multibyte character at the old 1024-byte forwarding boundary.
expected = (
prefix + b"x" * (1023 - len(prefix)) + '图片🌍"}\n\ndata: [DONE]\n\n'.encode()
if stream
else '{ "choices": [{"message": {"content": "图片🌍"}}] }\n'.encode()
)
async def backend(request):
stage = request.match_info["stage"]
assert request.content_type == "application/json"
body = await request.json()
seen[stage].append(body)
if stage == "encode":
item = body["messages"][0]["content"][0]
return web.json_response(
{
"ec_transfer_params": {
item["uuid"]: {
"metadata": {"image_grid_thw": [[1, 2, 2]]},
"peer_port": len(seen[stage]),
}
}
}
)
if stage == "prefill":
assert body["max_tokens"] == 1 and body["stream"] is False
return web.json_response(
{"kv_transfer_params": {"remote_block_ids": [len(seen[stage])]}}
)
if len(seen[stage]) == 1:
return web.Response(status=500, text="retry")
return web.Response(
body=expected,
content_type="text/event-stream" if stream else "application/json",
)
backend_app = web.Application()
backend_app.router.add_post("/{stage}/v1/chat/completions", backend)
async with TestServer(backend_app) as server:
for stage, field in [
("encode", "e_urls"),
("prefill", "p_urls"),
("decode", "d_urls"),
]:
urls = [str(server.make_url(f"/{stage}"))]
if stage == "prefill" and not prefill:
urls = []
monkeypatch.setattr(proxy.app.state, field, urls, raising=False)
monkeypatch.setattr(proxy.app.state, "d_ec_urls", [], raising=False)
monkeypatch.setattr(proxy.app.state, "ec_consumer_dp_size", 1, raising=False)
monkeypatch.setattr(proxy, "DECODE_RETRIES", 1)
monkeypatch.setattr(proxy, "NO_REWRITE", False)
item = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJj"}}
body = {
"model": "test",
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "描述图片🌍"},
item,
item,
],
}
],
"stream": stream,
"temperature": 0.01,
"max_tokens": 32,
"seed": 42,
"structured_outputs": {"choice": ["A", "B"]},
}
await proxy.on_startup()
try:
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=proxy.app), base_url="http://proxy"
) as client:
response = await client.post("/v1/chat/completions", json=body)
assert response.status_code == 200
assert response.content == expected
assert response.headers["content-type"].startswith(
"text/event-stream" if stream else "application/json"
)
finally:
await proxy.on_shutdown()
assert len(seen["encode"]) == 4
assert len(seen["decode"]) == 2
final = seen["decode"][-1]
for key in [
"model",
"stream",
"temperature",
"max_tokens",
"seed",
"structured_outputs",
]:
assert final[key] == body[key]
content = final["messages"][0]["content"]
assert content[0] == body["messages"][0]["content"][0]
for reference in content[1:]:
assert reference == {
"type": "image_embeds",
"image_embeds": {"image_grid_thw": [1, 2, 2]},
"uuid": proxy.content_uuid(item),
}
assert final["ec_transfer_params"][proxy.content_uuid(item)]["peer_port"] in (3, 4)
if prefill:
assert final["kv_transfer_params"] == {"remote_block_ids": [2]}
@pytest.mark.parametrize("server_keep_alive", ["0.1", "1", "5", "2", "30"])
def test_pooled_connections_are_retired_before_the_server_closes_them(
proxy, monkeypatch, server_keep_alive