diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index ace4455828..3912c1c67a 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -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: :: 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 diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py index b026b5537a..9a92b1fdbc 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py @@ -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