Coverage for gco / services / api_routes / inference_proxy.py: 100.00%
181 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""Authenticated reverse proxy for managed inference endpoints.
3All public inference traffic terminates at the dedicated inference-proxy
4service, whose ``AuthenticationMiddleware`` validates the Lambda proxy's
5short-lived HMAC envelope (timestamp, nonce, method, target, and body digest).
6The service then forwards to one strictly derived in-cluster Service name. This
7keeps model traffic out of the manifest processor and removes the historical
8direct ALB target groups that allowed callers to bypass API Gateway through
9Global Accelerator.
11The upstream hop is verified HTTPS on port 8443: every managed model pod runs
12a TLS sidecar serving the ``gco-inference`` wildcard certificate, and this
13proxy trusts only the GCO internal CA (see :mod:`gco.services.internal_tls`).
15Every proxied request builds its own upstream client and closes it once its
16stream ends. A process-wide client would pool nothing: kube-proxy balances
17connections, not requests, so keep-alive stays off to spread requests across
18the Ready model replicas, and each request opens its own connection either
19way. A client per request also takes its TLS context from
20:func:`~gco.services.internal_tls.internal_ssl_context` when the request
21arrives. That cache hands back the same context until the projected
22``ca.crt`` changes and a fresh one afterwards, so a rotated internal CA (for
23example the post-Helm pass re-issuing ``inference-proxy-tls`` during an
24in-place upgrade) is trusted from the next request on, without restarting
25the pod; a long-lived client would keep verifying against the bundle it
26loaded first and answer 502 until a restart.
27"""
29from __future__ import annotations
31import asyncio
32import logging
33import os
34import re
35import secrets
36from collections.abc import AsyncIterator
37from functools import lru_cache
38from typing import Any
39from urllib.parse import quote
41import httpx2
42from fastapi import APIRouter, HTTPException, Request
43from fastapi.responses import StreamingResponse
44from starlette.types import Receive, Scope, Send
46from gco.services import internal_tls, tracing
47from gco.services.inference_store import InferenceEndpointStore, get_inference_endpoint_store
49# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
50# Generated at (UTC): 2026-10-25T13:00:05Z
51# Generated from Git commit: de508a874e7ac267180e26a4c6fc56c1cd0878c2
52# Flowchart(s) generated from this file:
53# * ``_resolve_upstream`` -> ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._resolve_upstream.html``
54# (PNG: ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._resolve_upstream.png``)
55# * ``_proxy`` -> ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._proxy.html``
56# (PNG: ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._proxy.png``)
57# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
58# <pyflowchart-code-diagram> END
61router = APIRouter(prefix="/inference", tags=["Inference"])
62logger = logging.getLogger(__name__)
64_DNS_LABEL_RE = re.compile(r"^[a-z0-9](?:[-a-z0-9]{0,61}[a-z0-9])?$")
65_SUPPORTED_METHODS = ["GET", "HEAD", "POST"]
66_HOP_BY_HOP_HEADERS = frozenset(
67 {
68 "connection",
69 "keep-alive",
70 "proxy-authenticate",
71 "proxy-authorization",
72 "te",
73 "trailer",
74 "transfer-encoding",
75 "upgrade",
76 }
77)
78_ALLOWED_REQUEST_HEADERS = frozenset(
79 {
80 "accept",
81 "accept-encoding",
82 "cache-control",
83 "content-encoding",
84 "content-type",
85 "idempotency-key",
86 "if-match",
87 "if-none-match",
88 "prefer",
89 "range",
90 "user-agent",
91 "x-request-id",
92 }
93)
94_BLOCKED_PATH_SEGMENTS = frozenset(
95 {"admin", "debug", "docs", "instances", "metrics", "openapi.json"}
96)
97_V1_MODELS_RE = re.compile(r"^v1/models(?:/[^/]+)?$")
98# ``systemone`` is vLLM's structured decisions API (v0.32.0+): typed choice,
99# noul and score questions over a state, answered from the model's label
100# probabilities. Like the generation calls it is POST-only.
101_V1_GENERATION_RE = re.compile(
102 r"^v1/(?:chat/completions|completions|embeddings|responses|systemone)$"
103)
104_V2_MODELS_RE = re.compile(r"^v2/models(?:/[^/]+(?:/(?:config|infer|ready|stats))?)?$")
107def _bounded_timeout(name: str, default: float, minimum: float, maximum: float) -> float:
108 """Read one finite, bounded timeout from the environment."""
109 raw = os.getenv(name)
110 if raw is None:
111 return default
112 try:
113 value = float(raw)
114 except ValueError:
115 return default
116 return value if minimum <= value <= maximum else default
119@lru_cache(maxsize=1)
120def _get_inference_store() -> InferenceEndpointStore:
121 """Create one process-local DynamoDB endpoint-store client lazily."""
122 return get_inference_endpoint_store()
125#: Upstream cleanups still running. The event loop holds only weak references
126#: to tasks, and a cleanup outlives its awaiter whenever the request that
127#: started it is cancelled, so each one is held here until it finishes.
128_cleanup_tasks: set[asyncio.Task[None]] = set()
131def _new_upstream_client() -> httpx2.AsyncClient:
132 """Build the client for one proxied request; the caller releases it.
134 The TLS context is looked up now, per request, which is what lets a
135 rotated internal CA take effect without a restart (see the module
136 docstring); a lookup whose CA file is unchanged costs one ``stat``, not a
137 CA load. No keep-alive and no connection cap: the client carries exactly
138 one request, and its connection closes as soon as that response ends.
139 ``verify`` and ``limits`` sit on the transport because a client given
140 ``transport=`` ignores its own. The default ``Accept-Encoding`` is removed
141 because the proxy relays raw upstream bytes: only an encoding the caller
142 asked for may reach the model. Raises
143 :class:`~gco.services.internal_tls.InternalTLSError` while the CA bundle
144 is missing or invalid, before anything exists that would need closing.
145 """
146 transport = httpx2.AsyncHTTPTransport(
147 verify=internal_tls.internal_ssl_context(),
148 limits=httpx2.Limits(max_connections=None, max_keepalive_connections=0),
149 trust_env=False,
150 )
151 client = httpx2.AsyncClient(
152 transport=tracing.wrap_async_transport(transport),
153 follow_redirects=False,
154 trust_env=False,
155 )
156 del client.headers["accept-encoding"]
157 return client
160def _validate_label(value: object, field: str) -> str:
161 """Return a safe Kubernetes DNS label or reject the request."""
162 if not isinstance(value, str) or _DNS_LABEL_RE.fullmatch(value) is None:
163 raise HTTPException(status_code=404, detail=f"Invalid inference {field}")
164 return value
167def _target_service(endpoint: dict[str, Any], endpoint_name: str) -> str:
168 """Resolve the only in-cluster Service this endpoint may use.
170 Plain endpoints use ``<name>``. Mooncake disaggregated/both endpoints use
171 their reconciled ``<name>-proxy`` Service. During an active canary, a
172 cryptographically unbiased request sample is routed to ``<name>-canary``.
173 Every value is derived from a validated endpoint record, never from a URL or
174 header supplied by the caller.
175 """
176 spec = endpoint.get("spec")
177 if not isinstance(spec, dict):
178 raise HTTPException(status_code=503, detail="Inference endpoint has an invalid spec")
180 mooncake = spec.get("mooncake")
181 if isinstance(mooncake, dict) and mooncake.get("mode") in {"disaggregated", "both"}:
182 return _validate_label(f"{endpoint_name}-proxy", "service")
184 canary = spec.get("canary")
185 region = os.getenv("REGION", "")
186 region_status = endpoint.get("region_status")
187 local_status = region_status.get(region, {}) if isinstance(region_status, dict) else {}
188 canary_status = local_status.get("canary") if isinstance(local_status, dict) else None
189 if isinstance(canary, dict) and isinstance(canary_status, dict):
190 try:
191 weight = int(canary.get("weight", 0))
192 ready = int(canary_status.get("replicas_ready", 0))
193 desired = int(canary_status.get("replicas_desired", 0))
194 except TypeError, ValueError:
195 weight = ready = desired = 0
196 canary_is_ready = (
197 canary_status.get("state") == "running"
198 and canary_status.get("image") == canary.get("image")
199 and desired > 0
200 and ready >= desired
201 )
202 if canary_is_ready and 1 <= weight <= 99 and secrets.randbelow(100) < weight:
203 return _validate_label(f"{endpoint_name}-canary", "service")
205 return endpoint_name
208async def _resolve_upstream(endpoint_name: str) -> tuple[str, str, str]:
209 """Resolve the authorized Service, namespace, and configured health path."""
210 endpoint_name = _validate_label(endpoint_name, "name")
211 endpoint = await asyncio.to_thread(_get_inference_store().get_endpoint, endpoint_name)
212 if not endpoint:
213 raise HTTPException(
214 status_code=404, detail=f"Inference endpoint '{endpoint_name}' not found"
215 )
217 namespace = _validate_label(endpoint.get("namespace", "gco-inference"), "namespace")
218 allowed_namespace = os.getenv("INFERENCE_NAMESPACE", "gco-inference")
219 if namespace != allowed_namespace:
220 raise HTTPException(status_code=503, detail="Inference endpoint namespace is not routable")
222 region = os.getenv("REGION", "")
223 target_regions = endpoint.get("target_regions")
224 if not region or not isinstance(target_regions, list) or region not in target_regions:
225 raise HTTPException(
226 status_code=404, detail="Inference endpoint is not deployed in this region"
227 )
229 desired_state = endpoint.get("desired_state")
230 region_status = endpoint.get("region_status")
231 local_status = region_status.get(region, {}) if isinstance(region_status, dict) else {}
232 local_state = local_status.get("state") if isinstance(local_status, dict) else None
233 # Readiness is decided per region. ``deploying`` is the endpoint-wide
234 # summary while any target region still converges, and every spec update
235 # (scale, set-topology, configure-store, update-image) re-enters it until
236 # the monitor has observed all regions running again; a region whose own
237 # status is ``running`` has Ready pods and keeps serving throughout.
238 # Stopped and deleted endpoints never serve.
239 if desired_state not in ("running", "deploying") or local_state != "running":
240 raise HTTPException(
241 status_code=503, detail="Inference endpoint is not ready in this region"
242 )
244 spec = endpoint.get("spec")
245 configured_health_path = (
246 spec.get("health_check_path", "/health") if isinstance(spec, dict) else "/health"
247 )
248 if not isinstance(configured_health_path, str) or not configured_health_path.startswith("/"):
249 configured_health_path = "/health"
251 return _target_service(endpoint, endpoint_name), namespace, configured_health_path
254def _request_headers(request: Request) -> list[tuple[bytes, bytes]]:
255 """Forward only explicitly supported end-to-end model request headers, as sent.
257 The raw ASGI bytes are forwarded rather than Starlette's latin-1 decoded
258 strings: httpx2 encodes a ``str`` header value as ASCII, so one non-ASCII
259 byte (a UTF-8 ``user-agent``, say) would fail the request with a 500
260 instead of reaching the model unchanged.
261 """
262 return [
263 (name.lower(), value)
264 for name, value in request.headers.raw
265 if name.lower().decode("latin-1") in _ALLOWED_REQUEST_HEADERS
266 ]
269def _response_headers(response: httpx2.Response) -> list[tuple[bytes, bytes]]:
270 """Copy end-to-end response headers byte for byte, dropping hop-by-hop framing.
272 Starlette encodes ``str`` header values as latin-1, so relaying httpx2's
273 decoded values failed any response carrying a UTF-8 value outside latin-1
274 with a 500 and silently transcoded one inside it. The raw bytes keep every
275 value exact and every repeated header on its own line; names are
276 lowercased, as ASGI requires.
277 """
278 blocked = _HOP_BY_HOP_HEADERS | {"content-length"}
279 relayed: list[tuple[bytes, bytes]] = []
280 for name, value in response.headers.raw:
281 lowered = name.lower()
282 if lowered.decode("latin-1") not in blocked:
283 relayed.append((lowered, value))
284 return relayed
287def _validate_upstream_path(
288 upstream_path: str,
289 method: str,
290 configured_health_path: str = "/health",
291) -> str:
292 """Allow serving/configured-health APIs while denying privileged paths."""
293 normalized = upstream_path.strip("/")
294 segments = [segment.lower() for segment in normalized.split("/") if segment]
295 if any(segment in _BLOCKED_PATH_SEGMENTS for segment in segments):
296 raise HTTPException(status_code=404, detail="Inference path is not exposed")
298 method = method.upper()
299 configured_health = configured_health_path.strip("/")
300 # ``server_info`` is SGLang's read-only identity document (the launcher
301 # arguments the running server resolved, including ``model_path`` and
302 # ``revision``); it is what ``gco inference models`` reads for that runtime.
303 if (
304 not normalized
305 or normalized == "health"
306 or (configured_health and normalized == configured_health)
307 or normalized == "server_info"
308 or _V1_MODELS_RE.fullmatch(normalized)
309 ) and method in {"GET", "HEAD"}:
310 return normalized
311 # ``generate`` is SGLang's native generation API; streaming is a body flag
312 # on the same path.
313 if (_V1_GENERATION_RE.fullmatch(normalized) or normalized == "generate") and method == "POST":
314 return normalized
315 if _V2_MODELS_RE.fullmatch(normalized) and method in {"GET", "HEAD", "POST"}:
316 return normalized
318 raise HTTPException(status_code=404, detail="Inference path is not exposed")
321async def _close_upstream(response: httpx2.Response | None, client: httpx2.AsyncClient) -> None:
322 """Close one request's upstream response, if it got one, then its client.
324 Both closes are idempotent, so releasing an upstream twice is harmless.
325 """
326 try:
327 if response is not None:
328 await response.aclose()
329 finally:
330 await client.aclose()
333async def _release_upstream(response: httpx2.Response | None, client: httpx2.AsyncClient) -> None:
334 """Close one request's upstream in its own task, shielded from cancellation.
336 Cancelling the request (the caller went away, or shutdown outlived its
337 grace period) interrupts only this wait, never the close itself, so a
338 cancelled request still releases its model connection.
339 """
340 cleanup = asyncio.create_task(_close_upstream(response, client))
341 _cleanup_tasks.add(cleanup)
342 cleanup.add_done_callback(_cleanup_tasks.discard)
343 await asyncio.shield(cleanup)
346async def _stream_response(
347 response: httpx2.Response, client: httpx2.AsyncClient
348) -> AsyncIterator[bytes]:
349 """Yield the upstream body unchanged, then release the request's upstream."""
350 try:
351 async for chunk in response.aiter_raw():
352 yield chunk
353 finally:
354 await _release_upstream(response, client)
357class _UpstreamStreamingResponse(StreamingResponse):
358 """Relay one upstream response, releasing its client however the relay ends.
360 The body generator releases the upstream as soon as the stream finishes,
361 fails, or is cancelled mid-body. Starlette starts that generator only once
362 the headers are on their way, though: when the caller disconnected while
363 the model was still answering, the response is cancelled before its first
364 chunk and the generator never runs, so the response itself releases the
365 upstream again once it is done (a no-op after the generator's release).
366 """
368 def __init__(self, response: httpx2.Response, client: httpx2.AsyncClient) -> None:
369 relayed_headers = _response_headers(response)
370 super().__init__(
371 _stream_response(response, client),
372 status_code=response.status_code,
373 media_type=None,
374 )
375 # Assigned rather than passed as ``headers=``, which takes ``str``
376 # values and re-encodes them as latin-1 (see _response_headers).
377 self.raw_headers = relayed_headers
378 self._upstream = (response, client)
380 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
381 try:
382 await super().__call__(scope, receive, send)
383 finally:
384 await _release_upstream(*self._upstream)
387async def _proxy(
388 request: Request, endpoint_name: str, upstream_path: str = ""
389) -> StreamingResponse:
390 """Forward one authenticated request to a managed in-cluster endpoint."""
391 if any(part in {".", ".."} for part in upstream_path.split("/")):
392 raise HTTPException(status_code=400, detail="Invalid inference path")
393 service_name, namespace, configured_health_path = await _resolve_upstream(endpoint_name)
394 upstream_path = _validate_upstream_path(
395 upstream_path,
396 request.method,
397 configured_health_path,
398 )
399 encoded_suffix = quote(upstream_path, safe="/:@-._~")
400 upstream_path_value = f"/{encoded_suffix}" if encoded_suffix else "/"
401 upstream_url = ( # nosemgrep: python.django.security.injection.tainted-url-host.tainted-url-host
402 # Both host labels passed the strict Kubernetes DNS-label allowlist in
403 # _resolve_upstream; callers cannot supply a URL, address, or suffix.
404 # Port 8443 is the model pod's TLS sidecar, the only port its Service
405 # publishes.
406 f"https://{service_name}.{namespace}.svc.cluster.local:8443{upstream_path_value}"
407 )
409 body = await request.body()
410 timeout = httpx2.Timeout(
411 connect=_bounded_timeout("INFERENCE_PROXY_CONNECT_TIMEOUT_SECONDS", 5.0, 0.1, 30.0),
412 read=_bounded_timeout("INFERENCE_PROXY_READ_TIMEOUT_SECONDS", 300.0, 1.0, 900.0),
413 write=_bounded_timeout("INFERENCE_PROXY_WRITE_TIMEOUT_SECONDS", 30.0, 1.0, 300.0),
414 pool=_bounded_timeout("INFERENCE_PROXY_POOL_TIMEOUT_SECONDS", 5.0, 0.1, 30.0),
415 )
416 try:
417 client = _new_upstream_client()
418 except internal_tls.InternalTLSError as exc:
419 logger.error("Inference upstream TLS trust is unavailable: %s", exc)
420 raise HTTPException(status_code=502, detail="Inference endpoint is unavailable") from exc
422 # Until the streaming response owns the upstream, every exit releases it.
423 response: httpx2.Response | None = None
424 try:
425 upstream_request = client.build_request(
426 request.method,
427 upstream_url,
428 params=list(request.query_params.multi_items()),
429 headers=_request_headers(request),
430 content=body,
431 timeout=timeout,
432 )
433 response = await client.send(upstream_request, stream=True)
434 return _UpstreamStreamingResponse(response, client)
435 except httpx2.TimeoutException as exc:
436 await _release_upstream(response, client)
437 raise HTTPException(status_code=504, detail="Inference endpoint timed out") from exc
438 except httpx2.HTTPError as exc:
439 await _release_upstream(response, client)
440 logger.warning(
441 "Inference upstream request failed: service=%s namespace=%s error_type=%s error=%s",
442 service_name,
443 namespace,
444 type(exc).__name__,
445 exc,
446 )
447 raise HTTPException(status_code=502, detail="Inference endpoint is unavailable") from exc
448 except BaseException:
449 # Anything else, a cancelled request included, releases it too.
450 await _release_upstream(response, client)
451 raise
454async def proxy_inference_root(request: Request, endpoint_name: str) -> StreamingResponse:
455 """Proxy an endpoint-root request after platform authentication."""
456 return await _proxy(request, endpoint_name)
459async def proxy_inference_path(
460 request: Request,
461 endpoint_name: str,
462 upstream_path: str,
463) -> StreamingResponse:
464 """Proxy an endpoint sub-path after platform authentication."""
465 return await _proxy(request, endpoint_name, upstream_path)
468# Register one route per method rather than a single multi-method route.
469# FastAPI derives an operation's ``operationId`` from ``generate_unique_id``,
470# which appends ``list(route.methods)[0]`` — a single arbitrary member of an
471# unordered set — and computes it once per route. A route carrying GET, HEAD,
472# and POST therefore emits three OpenAPI operations sharing one operationId,
473# which violates the spec's uniqueness requirement, makes generated clients
474# collide, and raises a UserWarning on every schema build. One method per
475# route keeps each generated operationId distinct while leaving request
476# handling byte-for-byte identical.
477for _path, _endpoint in (
478 ("/{endpoint_name}", proxy_inference_root),
479 ("/{endpoint_name}/{upstream_path:path}", proxy_inference_path),
480):
481 for _method in _SUPPORTED_METHODS:
482 router.add_api_route(_path, _endpoint, methods=[_method])