Coverage for gco / services / mooncake_pd_proxy.py: 100.00%
185 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +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`` — must run on that image's Python 3.12, and must not
9import anything from the ``gco`` package. It uses ``httpx2`` when the image
10provides it (as GCO's own environment does) and otherwise the image's
11``httpx``; the two share the API used here.
13The proxy binds ``PD_PROXY_HOST`` (loopback in the pod, behind the pod's TLS
14sidecar) and reaches prefill and decode over verified HTTPS: with
15``PD_PROXY_CA_FILE`` set it trusts exactly that CA bundle (the GCO internal
16CA), and hostname verification stays on.
18Per request on the public ``/v1/*`` serving paths it:
201. Treats the prompt as not resident in the shared store. The residency check is
21 non-blocking and bounded by ``PD_PROXY_RESIDENCY_TIMEOUT_SECONDS``; a miss or a
22 check that does not finish in time is sent straight to prefill, so a slow or
23 unreachable store never holds the request.
242. Primes a prefill pod with the request at ``max_tokens=1`` and
25 ``kv_transfer_params={"do_remote_decode": true}`` so prefill computes and
26 exports the prompt KV through the MooncakeConnector.
273. Sends the original request to a decode pod, relaying any ``kv_transfer_params``
28 the prefill step returned (with ``do_remote_prefill=true``) so decode pulls the
29 KV instead of recomputing, and streams the decode response back to the client.
31Non-health GET requests (for example OpenAI-compatible ``/v1/models`` discovery)
32and non-generation POST requests pass through to decode with their query string
33preserved. JSON request bodies must be objects; arrays and scalars are rejected
34with a client error before either backend is called.
36Prefill and decode are addressed through their in-cluster Services, so kube-proxy
37load-balances across only the Ready role pods. When the decode Service has no
38Ready endpoints the proxy rejects the request with a stable 503 rather than
39emitting partial output. The privileged ``/instances/add`` admin path requires
40the ``ADMIN_API_KEY`` header and is never published on the public Ingress.
42The ``kv_transfer_params`` handshake is best-effort and pass-through: the proxy
43sets only the outer ``do_remote_decode`` / ``do_remote_prefill`` flags and relays
44whatever inner fields the connector returns, so it does not hard-code a
45connector-version-specific schema. If prefill returns no transfer params (or the
46priming call fails), the decode request is still served correctly — the connector
47falls back to its own KV matching or decode recomputes — so the invoke path keeps
48working either way.
49"""
51from __future__ import annotations
53import asyncio
54import json
55import logging
56import os
57import ssl
58from collections.abc import AsyncIterator
59from typing import Any
61import uvicorn
62from fastapi import FastAPI, Request, Response
63from fastapi.responses import JSONResponse, StreamingResponse
64from starlette.types import Receive, Scope, Send
66# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
67# Generated at (UTC): 2026-09-28T07:34:51Z
68# Generated from Git commit: 95213a3dfe214f41ea8e3977b79711b1be061ac0
69# Flowchart(s) generated from this file:
70# * ``_dispatch`` -> ``diagrams/code_diagrams/gco/services/mooncake_pd_proxy._dispatch.html``
71# (PNG: ``diagrams/code_diagrams/gco/services/mooncake_pd_proxy._dispatch.png``)
72# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
73# <pyflowchart-code-diagram> END
76try:
77 import httpx2 as httpx # type: ignore[import-not-found,unused-ignore]
78except ImportError:
79 import httpx # type: ignore[no-redef,unused-ignore]
82logging.basicConfig(
83 level=logging.INFO, format="%(asctime)s %(levelname)s [mooncake-pd-proxy] %(message)s"
84)
85logger = logging.getLogger("mooncake-pd-proxy")
87HOST = os.environ.get("PD_PROXY_HOST", "0.0.0.0")
88PORT = int(os.environ.get("PD_PROXY_PORT", "8000"))
89PREFILL_URL = os.environ.get("PD_PROXY_PREFILL_URL", "").rstrip("/")
90DECODE_URL = os.environ.get("PD_PROXY_DECODE_URL", "").rstrip("/")
91CA_FILE = os.environ.get("PD_PROXY_CA_FILE", "").strip()
92RESIDENCY_TIMEOUT = float(os.environ.get("PD_PROXY_RESIDENCY_TIMEOUT_SECONDS", "2"))
93NO_DECODE_STATUS = int(os.environ.get("PD_PROXY_NO_DECODE_BACKEND_STATUS", "503"))
94NO_DECODE_MESSAGE = os.environ.get(
95 "PD_PROXY_NO_DECODE_BACKEND_MESSAGE", "no available decode backend"
96)
97ADMIN_API_KEY = os.environ.get("ADMIN_API_KEY", "")
98ADMIN_PATH = "/instances/add"
100# Per-request upstream timeout. Connect is kept short so an endpoint with no
101# Ready decode pods (empty Service endpoints) surfaces quickly as a 503 rather
102# than hanging the client; reads are unbounded for long generations.
103_TIMEOUT = httpx.Timeout(None, connect=5.0)
105# Keep this allowlist aligned with the authenticated inference proxy's public
106# boundary. ``content-encoding`` is intentionally excluded here: this proxy
107# parses and re-serializes JSON, so forwarding the original encoding would
108# falsely describe the new body bytes.
109_ALLOWED_REQUEST_HEADERS = frozenset(
110 {
111 "accept",
112 "accept-encoding",
113 "cache-control",
114 "content-type",
115 "idempotency-key",
116 "if-match",
117 "if-none-match",
118 "prefer",
119 "range",
120 "user-agent",
121 "x-request-id",
122 }
123)
124_HOP_BY_HOP_HEADERS = frozenset(
125 {
126 "connection",
127 "keep-alive",
128 "proxy-authenticate",
129 "proxy-authorization",
130 "te",
131 "trailer",
132 "transfer-encoding",
133 "upgrade",
134 }
135)
138#: The trust context for the CA file as it was last read, keyed on the file's
139#: identity (inode, modification time, size).
140_verify_cache: tuple[tuple[int, int, int], ssl.SSLContext] | None = None
143def _upstream_verify() -> ssl.SSLContext | bool:
144 """TLS trust for the prefill and decode hops, re-read when the CA changes.
146 With ``PD_PROXY_CA_FILE`` set, trust exactly that bundle and nothing from
147 the system store; hostname verification stays on, so each role Service
148 must present a certificate for the name the proxy dialled. Unset keeps the
149 library default for plain ``http://`` backends and local runs.
151 The context is cached on the file's identity. Kubernetes updates a
152 projected Secret by swapping a symlink, which changes the resolved inode
153 and modification time, so a rotated internal CA is trusted from the next
154 request on without restarting the pod, and an unchanged one costs a
155 ``stat``. A bundle that is missing or is not a CA raises ``OSError``
156 (``ssl.SSLError`` is one): at import that stops the program rather than
157 serving over a weaker trust, and per request it answers like an
158 unreachable decode backend.
159 """
160 global _verify_cache
161 if not CA_FILE:
162 return True
163 info = os.stat(CA_FILE)
164 identity = (info.st_ino, info.st_mtime_ns, info.st_size)
165 cached = _verify_cache
166 if cached is not None and cached[0] == identity:
167 return cached[1]
168 context = ssl.create_default_context(cafile=CA_FILE)
169 context.minimum_version = ssl.TLSVersion.TLSv1_2
170 _verify_cache = (identity, context)
171 return context
174def _new_client() -> httpx.AsyncClient:
175 """Build the client for one upstream call; the caller closes it.
177 A client per call takes its trust from :func:`_upstream_verify` when the
178 call starts, which is what lets a rotated CA take effect without a
179 restart; a process-wide client would verify against the bundle it loaded
180 first until the pod restarted. Keep-alive stays off: kube-proxy balances
181 connections, not requests, so a pooled connection would pin every request
182 to one prefill or decode pod. The library's default ``Accept-Encoding`` is
183 removed because the proxy relays decode's raw bytes: only an encoding the
184 caller asked for (forwarded by :func:`_request_headers`) may shape the
185 body the client receives.
186 """
187 client = httpx.AsyncClient(
188 timeout=_TIMEOUT,
189 verify=_upstream_verify(),
190 limits=httpx.Limits(max_keepalive_connections=0),
191 )
192 del client.headers["accept-encoding"]
193 return client
196# Fail at start on a configured bundle that cannot be used.
197_upstream_verify()
199app = FastAPI()
201#: Upstream closes still running. The event loop holds only weak references
202#: to tasks, and a close outlives its awaiter whenever the request that
203#: started it is cancelled, so each one is held here until it finishes.
204_closing: set[asyncio.Task[None]] = set()
207async def _close_upstream(response: httpx.Response | None, client: httpx.AsyncClient) -> None:
208 """Close one call's upstream response, if it got one, then its client."""
209 try:
210 if response is not None:
211 await response.aclose()
212 finally:
213 await client.aclose()
216async def _release_upstream(response: httpx.Response | None, client: httpx.AsyncClient) -> None:
217 """Close one call's upstream in its own task, shielded from cancellation.
219 Cancelling the request (the caller went away) interrupts only this wait,
220 never the close itself, so a cancelled request still releases its
221 connection. Both closes are idempotent, so releasing twice is harmless.
222 """
223 closing = asyncio.create_task(_close_upstream(response, client))
224 _closing.add(closing)
225 closing.add_done_callback(_closing.discard)
226 await asyncio.shield(closing)
229@app.get("/healthz")
230@app.get("/health")
231async def _health() -> JSONResponse:
232 return JSONResponse({"status": "ok"})
235def _is_serving_path(path: str) -> bool:
236 """True for the OpenAI-compatible serving paths the proxy disaggregates."""
237 return path.endswith(("/completions", "/chat/completions", "/embeddings"))
240def _prefill_body(body: dict[str, Any]) -> dict[str, Any]:
241 """Body for priming prefill: one token, no stream, request remote decode."""
242 pf = dict(body)
243 pf["stream"] = False
244 pf["max_tokens"] = 1
245 if "max_completion_tokens" in pf:
246 pf["max_completion_tokens"] = 1
247 kvp = dict(pf.get("kv_transfer_params") or {})
248 kvp["do_remote_decode"] = True
249 kvp["do_remote_prefill"] = False
250 pf["kv_transfer_params"] = kvp
251 return pf
254def _decode_body(body: dict[str, Any], prefill_kv_params: dict[str, Any]) -> dict[str, Any]:
255 """Body for decode: original request, relaying prefill's transfer params."""
256 dc = dict(body)
257 if prefill_kv_params:
258 kvp = dict(prefill_kv_params)
259 kvp["do_remote_prefill"] = True
260 kvp["do_remote_decode"] = False
261 dc["kv_transfer_params"] = kvp
262 return dc
265async def _prime_prefill(path: str, body: dict[str, Any]) -> dict[str, Any]:
266 """Run the prefill step; return its kv_transfer_params (best-effort)."""
267 if not PREFILL_URL:
268 return {}
269 try:
270 client = _new_client()
271 try:
272 resp = await client.post(f"{PREFILL_URL}{path}", json=_prefill_body(body))
273 finally:
274 await client.aclose()
275 resp.raise_for_status()
276 data = resp.json()
277 return data.get("kv_transfer_params") or {}
278 except Exception as exc: # priming is best-effort
279 logger.warning("prefill priming failed; decode will serve directly: %s", exc)
280 return {}
283def _request_target(request: Request) -> str:
284 """Return the path and query string exactly as they should reach decode."""
285 path = request.url.path
286 return f"{path}?{request.url.query}" if request.url.query else path
289def _request_headers(request: Request) -> list[tuple[bytes, bytes]]:
290 """Forward only explicitly supported end-to-end model headers, as sent.
292 Raw ASGI bytes, not Starlette's latin-1 decoded strings: httpx encodes a
293 ``str`` header value as ASCII, so one non-ASCII byte would fail the
294 request instead of reaching decode unchanged.
295 """
296 return [
297 (name.lower(), value)
298 for name, value in request.headers.raw
299 if name.lower().decode("latin-1") in _ALLOWED_REQUEST_HEADERS
300 ]
303def _response_headers(response: httpx.Response) -> list[tuple[bytes, bytes]]:
304 """Relay end-to-end metadata byte for byte while dropping hop-by-hop framing.
306 Starlette would re-encode ``str`` values as latin-1, failing a UTF-8 value
307 outside latin-1 and transcoding one inside it; raw bytes keep each value
308 exact and each repeated header on its own line. Names are lowercased, as
309 ASGI requires.
310 """
311 blocked = _HOP_BY_HOP_HEADERS | {"content-length"}
312 relayed: list[tuple[bytes, bytes]] = []
313 for name, value in response.headers.raw:
314 lowered = name.lower()
315 if lowered.decode("latin-1") not in blocked:
316 relayed.append((lowered, value))
317 return relayed
320def _no_decode_backend() -> JSONResponse:
321 """The stable answer for a decode backend the proxy cannot reach."""
322 return JSONResponse(
323 {"error": {"message": NO_DECODE_MESSAGE, "type": "no_decode_backend"}},
324 status_code=NO_DECODE_STATUS,
325 )
328async def _relay_decode(
329 response: httpx.Response, client: httpx.AsyncClient
330) -> AsyncIterator[bytes]:
331 """Yield decode's body unchanged, then release the call's upstream."""
332 try:
333 async for chunk in response.aiter_raw():
334 yield chunk
335 finally:
336 await _release_upstream(response, client)
339class _DecodeStreamingResponse(StreamingResponse):
340 """Relay one decode response, releasing its client however the relay ends.
342 The body generator releases the upstream when the stream finishes, fails,
343 or is cancelled mid-body. Starlette starts that generator only once the
344 headers are on their way, though: a caller who disconnected while decode
345 was still answering cancels the response before its first chunk, so the
346 response itself releases the upstream again once it is done (a no-op
347 after the generator's release).
348 """
350 def __init__(
351 self, response: httpx.Response, client: httpx.AsyncClient, *, want_stream: bool
352 ) -> None:
353 relayed = _response_headers(response)
354 upstream_type = next(
355 (value for name, value in relayed if name == b"content-type"), b"application/json"
356 )
357 headers = [(name, value) for name, value in relayed if name != b"content-type"]
358 headers.append((b"content-type", b"text/event-stream" if want_stream else upstream_type))
359 super().__init__(
360 _relay_decode(response, client), status_code=response.status_code, media_type=None
361 )
362 # Assigned rather than passed as ``headers=``, which re-encodes ``str``
363 # values as latin-1 (see _response_headers).
364 self.raw_headers = headers
365 self._upstream = (response, client)
367 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
368 try:
369 await super().__call__(scope, receive, send)
370 finally:
371 await _release_upstream(*self._upstream)
374async def _stream_decode(
375 method: str,
376 target: str,
377 body: dict[str, Any] | None = None,
378 headers: list[tuple[bytes, bytes]] | None = None,
379) -> Response:
380 """Forward one request to decode and stream its response to the client."""
381 want_stream = bool(body and body.get("stream"))
382 url = f"{DECODE_URL}{target}"
383 request_kwargs: dict[str, Any] = {}
384 if body is not None:
385 request_kwargs["json"] = body
386 if headers is not None:
387 request_kwargs["headers"] = headers
388 try:
389 client = _new_client()
390 except OSError as exc:
391 # The CA bundle is missing or unusable (ssl.SSLError is an OSError):
392 # fail closed, like a decode backend the proxy cannot verify.
393 logger.warning("decode backend trust is unavailable: %s", exc)
394 return _no_decode_backend()
396 # Until the streaming response owns the upstream, every exit releases it.
397 resp: httpx.Response | None = None
398 try:
399 upstream_request = client.build_request(method, url, **request_kwargs)
400 resp = await client.send(upstream_request, stream=True)
401 return _DecodeStreamingResponse(resp, client, want_stream=want_stream)
402 except httpx.ConnectError as exc:
403 # No Ready decode endpoint behind the Service (or a TLS handshake the
404 # decode sidecar failed): reject with a stable status instead of
405 # emitting any partial output, and log why for the operator.
406 await _release_upstream(resp, client)
407 logger.warning("decode backend unreachable: %s", exc)
408 return _no_decode_backend()
409 except BaseException:
410 # Anything else, a cancelled request included, releases it too.
411 await _release_upstream(resp, client)
412 raise
415@app.post(ADMIN_PATH)
416async def _admin_add(request: Request) -> JSONResponse:
417 """Privileged admin endpoint, guarded by the ADMIN_API_KEY header.
419 Routing is via the prefill/decode Services, so kube-proxy already tracks
420 Ready pods and no per-pod registration is required; this endpoint exists so
421 the admin surface is present and authenticated (and kept off the public
422 Ingress), returning 200 for an authorized caller.
423 """
424 provided = (
425 request.headers.get("x-admin-api-key")
426 or request.headers.get("authorization", "").removeprefix("Bearer ").strip()
427 )
428 if not ADMIN_API_KEY or provided != ADMIN_API_KEY:
429 return JSONResponse({"error": "forbidden"}, status_code=403)
430 return JSONResponse({"status": "ok"})
433@app.api_route("/{full_path:path}", methods=["GET"])
434async def _get_passthrough(full_path: str, request: Request) -> Response:
435 """Forward non-health GETs, including OpenAI-compatible model discovery."""
436 return await _stream_decode(
437 "GET",
438 _request_target(request),
439 headers=_request_headers(request),
440 )
443@app.api_route("/{full_path:path}", methods=["POST"])
444async def _dispatch(full_path: str, request: Request) -> Any:
445 """Disaggregate one serving request: prime prefill, then stream decode."""
446 path = request.url.path
447 if path.endswith(ADMIN_PATH):
448 return await _admin_add(request)
450 raw = await request.body()
451 try:
452 decoded = json.loads(raw or b"{}")
453 except ValueError:
454 return JSONResponse({"error": "invalid JSON body"}, status_code=400)
455 if not isinstance(decoded, dict):
456 return JSONResponse({"error": "JSON body must be an object"}, status_code=400)
458 request_headers = _request_headers(request)
459 target = _request_target(request)
460 if not _is_serving_path(path):
461 return await _stream_decode("POST", target, decoded, request_headers)
463 # Residency check: non-blocking, treated as a miss so the prompt always goes
464 # to prefill first (the store is never on the request's critical path).
465 prefill_kv_params = await _prime_prefill(path, decoded)
466 return await _stream_decode(
467 "POST",
468 target,
469 _decode_body(decoded, prefill_kv_params),
470 request_headers,
471 )
474if __name__ == "__main__":
475 logger.info(
476 "starting PD proxy on %s:%d (prefill=%s decode=%s)", HOST, PORT, PREFILL_URL, DECODE_URL
477 )
478 uvicorn.run(app, host=HOST, port=PORT, log_level="info")