Coverage for gco / services / mooncake_pd_proxy.py: 100.00%
117 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Mooncake prefill-decode (PD) proxy for disaggregated inference endpoints.
3This is the program the ``{name}-proxy`` pod runs. It is shipped to the pod as a
4ConfigMap (the monitor reads this file's own source and mounts it at
5``/etc/pd-proxy/mooncake_pd_proxy.py``), and the proxy container runs it with
6``python3 /etc/pd-proxy/mooncake_pd_proxy.py``. It therefore must depend only on
7what the upstream ``vllm/vllm-openai`` image already ships — ``fastapi``,
8``uvicorn`` and ``httpx`` — and must not import anything from the ``gco``
9package.
11Per request on the public ``/v1/*`` serving paths it:
131. Treats the prompt as not resident in the shared store. The residency check is
14 non-blocking and bounded by ``PD_PROXY_RESIDENCY_TIMEOUT_SECONDS``; a miss or a
15 check that does not finish in time is sent straight to prefill, so a slow or
16 unreachable store never holds the request.
172. Primes a prefill pod with the request at ``max_tokens=1`` and
18 ``kv_transfer_params={"do_remote_decode": true}`` so prefill computes and
19 exports the prompt KV through the MooncakeConnector.
203. Sends the original request to a decode pod, relaying any ``kv_transfer_params``
21 the prefill step returned (with ``do_remote_prefill=true``) so decode pulls the
22 KV instead of recomputing, and streams the decode response back to the client.
24Non-health GET requests (for example OpenAI-compatible ``/v1/models`` discovery)
25and non-generation POST requests pass through to decode with their query string
26preserved. JSON request bodies must be objects; arrays and scalars are rejected
27with a client error before either backend is called.
29Prefill and decode are addressed through their in-cluster Services, so kube-proxy
30load-balances across only the Ready role pods. When the decode Service has no
31Ready endpoints the proxy rejects the request with a stable 503 rather than
32emitting partial output. The privileged ``/instances/add`` admin path requires
33the ``ADMIN_API_KEY`` header and is never published on the public Ingress.
35The ``kv_transfer_params`` handshake is best-effort and pass-through: the proxy
36sets only the outer ``do_remote_decode`` / ``do_remote_prefill`` flags and relays
37whatever inner fields the connector returns, so it does not hard-code a
38connector-version-specific schema. If prefill returns no transfer params (or the
39priming call fails), the decode request is still served correctly — the connector
40falls back to its own KV matching or decode recomputes — so the invoke path keeps
41working either way.
42"""
44from __future__ import annotations
46import json
47import logging
48import os
49from collections.abc import AsyncIterator
50from typing import Any
52import httpx
53import uvicorn
54from fastapi import FastAPI, Request, Response
55from fastapi.responses import JSONResponse, StreamingResponse
57# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
58# Generated at (UTC): 2026-08-30T12:00:00Z
59# Generated from Git commit: affbf6eccf3773dc3cfeac202e2cc6cbf92d4fc7
60# Flowchart(s) generated from this file:
61# * ``_dispatch`` -> ``diagrams/code_diagrams/gco/services/mooncake_pd_proxy._dispatch.html``
62# (PNG: ``diagrams/code_diagrams/gco/services/mooncake_pd_proxy._dispatch.png``)
63# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
64# <pyflowchart-code-diagram> END
67logging.basicConfig(
68 level=logging.INFO, format="%(asctime)s %(levelname)s [mooncake-pd-proxy] %(message)s"
69)
70logger = logging.getLogger("mooncake-pd-proxy")
72PORT = int(os.environ.get("PD_PROXY_PORT", "8000"))
73PREFILL_URL = os.environ.get("PD_PROXY_PREFILL_URL", "").rstrip("/")
74DECODE_URL = os.environ.get("PD_PROXY_DECODE_URL", "").rstrip("/")
75RESIDENCY_TIMEOUT = float(os.environ.get("PD_PROXY_RESIDENCY_TIMEOUT_SECONDS", "2"))
76NO_DECODE_STATUS = int(os.environ.get("PD_PROXY_NO_DECODE_BACKEND_STATUS", "503"))
77NO_DECODE_MESSAGE = os.environ.get(
78 "PD_PROXY_NO_DECODE_BACKEND_MESSAGE", "no available decode backend"
79)
80ADMIN_API_KEY = os.environ.get("ADMIN_API_KEY", "")
81ADMIN_PATH = "/instances/add"
83# Per-request upstream timeout. Connect is kept short so an endpoint with no
84# Ready decode pods (empty Service endpoints) surfaces quickly as a 503 rather
85# than hanging the client; reads are unbounded for long generations.
86_TIMEOUT = httpx.Timeout(None, connect=5.0)
88# Keep this allowlist aligned with the authenticated inference proxy's public
89# boundary. ``content-encoding`` is intentionally excluded here: this proxy
90# parses and re-serializes JSON, so forwarding the original encoding would
91# falsely describe the new body bytes.
92_ALLOWED_REQUEST_HEADERS = frozenset(
93 {
94 "accept",
95 "accept-encoding",
96 "cache-control",
97 "content-type",
98 "idempotency-key",
99 "if-match",
100 "if-none-match",
101 "prefer",
102 "range",
103 "user-agent",
104 "x-request-id",
105 }
106)
107_HOP_BY_HOP_HEADERS = frozenset(
108 {
109 "connection",
110 "keep-alive",
111 "proxy-authenticate",
112 "proxy-authorization",
113 "te",
114 "trailer",
115 "transfer-encoding",
116 "upgrade",
117 }
118)
120app = FastAPI()
121_client = httpx.AsyncClient(timeout=_TIMEOUT)
124@app.get("/healthz")
125@app.get("/health")
126async def _health() -> JSONResponse:
127 return JSONResponse({"status": "ok"})
130def _is_serving_path(path: str) -> bool:
131 """True for the OpenAI-compatible serving paths the proxy disaggregates."""
132 return path.endswith(("/completions", "/chat/completions", "/embeddings"))
135def _prefill_body(body: dict[str, Any]) -> dict[str, Any]:
136 """Body for priming prefill: one token, no stream, request remote decode."""
137 pf = dict(body)
138 pf["stream"] = False
139 pf["max_tokens"] = 1
140 if "max_completion_tokens" in pf:
141 pf["max_completion_tokens"] = 1
142 kvp = dict(pf.get("kv_transfer_params") or {})
143 kvp["do_remote_decode"] = True
144 kvp["do_remote_prefill"] = False
145 pf["kv_transfer_params"] = kvp
146 return pf
149def _decode_body(body: dict[str, Any], prefill_kv_params: dict[str, Any]) -> dict[str, Any]:
150 """Body for decode: original request, relaying prefill's transfer params."""
151 dc = dict(body)
152 if prefill_kv_params:
153 kvp = dict(prefill_kv_params)
154 kvp["do_remote_prefill"] = True
155 kvp["do_remote_decode"] = False
156 dc["kv_transfer_params"] = kvp
157 return dc
160async def _prime_prefill(path: str, body: dict[str, Any]) -> dict[str, Any]:
161 """Run the prefill step; return its kv_transfer_params (best-effort)."""
162 if not PREFILL_URL:
163 return {}
164 try:
165 resp = await _client.post(f"{PREFILL_URL}{path}", json=_prefill_body(body))
166 resp.raise_for_status()
167 data = resp.json()
168 return data.get("kv_transfer_params") or {}
169 except Exception as exc: # noqa: BLE001 - priming is best-effort
170 logger.warning("prefill priming failed; decode will serve directly: %s", exc)
171 return {}
174def _request_target(request: Request) -> str:
175 """Return the path and query string exactly as they should reach decode."""
176 path = request.url.path
177 return f"{path}?{request.url.query}" if request.url.query else path
180def _request_headers(request: Request) -> list[tuple[str, str]]:
181 """Forward only explicitly supported end-to-end model headers."""
182 return [
183 (name.lower(), value)
184 for name, value in request.headers.items()
185 if name.lower() in _ALLOWED_REQUEST_HEADERS
186 ]
189def _response_headers(response: httpx.Response) -> dict[str, str]:
190 """Relay end-to-end metadata while dropping hop-by-hop framing."""
191 blocked = _HOP_BY_HOP_HEADERS | {"content-length"}
192 return {name: value for name, value in response.headers.items() if name.lower() not in blocked}
195async def _stream_decode(
196 method: str,
197 target: str,
198 body: dict[str, Any] | None = None,
199 headers: list[tuple[str, str]] | None = None,
200) -> Response:
201 """Forward one request to decode and stream its response to the client."""
202 want_stream = bool(body and body.get("stream"))
203 url = f"{DECODE_URL}{target}"
204 request_kwargs: dict[str, Any] = {}
205 if body is not None:
206 request_kwargs["json"] = body
207 if headers is not None:
208 request_kwargs["headers"] = headers
209 upstream_request = _client.build_request(method, url, **request_kwargs)
210 try:
211 resp = await _client.send(upstream_request, stream=True)
212 except httpx.ConnectError:
213 # No Ready decode endpoint behind the Service: reject with a stable
214 # status instead of emitting any partial output.
215 return JSONResponse(
216 {"error": {"message": NO_DECODE_MESSAGE, "type": "no_decode_backend"}},
217 status_code=NO_DECODE_STATUS,
218 )
220 async def _body_iter() -> AsyncIterator[bytes]:
221 try:
222 async for chunk in resp.aiter_raw():
223 yield chunk
224 finally:
225 await resp.aclose()
227 response_headers = _response_headers(resp)
228 response_headers["content-type"] = (
229 "text/event-stream"
230 if want_stream
231 else response_headers.get("content-type", "application/json")
232 )
233 return StreamingResponse(
234 _body_iter(),
235 status_code=resp.status_code,
236 headers=response_headers,
237 media_type=None,
238 )
241@app.post(ADMIN_PATH)
242async def _admin_add(request: Request) -> JSONResponse:
243 """Privileged admin endpoint, guarded by the ADMIN_API_KEY header.
245 Routing is via the prefill/decode Services, so kube-proxy already tracks
246 Ready pods and no per-pod registration is required; this endpoint exists so
247 the admin surface is present and authenticated (and kept off the public
248 Ingress), returning 200 for an authorized caller.
249 """
250 provided = (
251 request.headers.get("x-admin-api-key")
252 or request.headers.get("authorization", "").removeprefix("Bearer ").strip()
253 )
254 if not ADMIN_API_KEY or provided != ADMIN_API_KEY:
255 return JSONResponse({"error": "forbidden"}, status_code=403)
256 return JSONResponse({"status": "ok"})
259@app.api_route("/{full_path:path}", methods=["GET"])
260async def _get_passthrough(full_path: str, request: Request) -> Response:
261 """Forward non-health GETs, including OpenAI-compatible model discovery."""
262 return await _stream_decode(
263 "GET",
264 _request_target(request),
265 headers=_request_headers(request),
266 )
269@app.api_route("/{full_path:path}", methods=["POST"])
270async def _dispatch(full_path: str, request: Request) -> Any:
271 """Disaggregate one serving request: prime prefill, then stream decode."""
272 path = request.url.path
273 if path.endswith(ADMIN_PATH):
274 return await _admin_add(request)
276 raw = await request.body()
277 try:
278 decoded = json.loads(raw or b"{}")
279 except ValueError:
280 return JSONResponse({"error": "invalid JSON body"}, status_code=400)
281 if not isinstance(decoded, dict):
282 return JSONResponse({"error": "JSON body must be an object"}, status_code=400)
284 request_headers = _request_headers(request)
285 target = _request_target(request)
286 if not _is_serving_path(path):
287 return await _stream_decode("POST", target, decoded, request_headers)
289 # Residency check: non-blocking, treated as a miss so the prompt always goes
290 # to prefill first (the store is never on the request's critical path).
291 prefill_kv_params = await _prime_prefill(path, decoded)
292 return await _stream_decode(
293 "POST",
294 target,
295 _decode_body(decoded, prefill_kv_params),
296 request_headers,
297 )
300if __name__ == "__main__":
301 logger.info("starting PD proxy on :%d (prefill=%s decode=%s)", PORT, PREFILL_URL, DECODE_URL)
302 uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")