Coverage for gco / services / api_routes / inference_proxy.py: 100.00%
141 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +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.
10"""
12from __future__ import annotations
14import asyncio
15import os
16import re
17import secrets
18from collections.abc import AsyncIterator
19from functools import lru_cache
20from typing import Any
21from urllib.parse import quote
23import httpx
24from fastapi import APIRouter, HTTPException, Request
25from fastapi.responses import StreamingResponse
27from gco.services.inference_store import InferenceEndpointStore, get_inference_endpoint_store
29# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
30# Generated at (UTC): 2026-09-01T14:42:56Z
31# Generated from Git commit: 89b000378ed5a912a38c06f4feab2b029936ebcc
32# Flowchart(s) generated from this file:
33# * ``_resolve_upstream`` -> ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._resolve_upstream.html``
34# (PNG: ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._resolve_upstream.png``)
35# * ``_proxy`` -> ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._proxy.html``
36# (PNG: ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._proxy.png``)
37# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
38# <pyflowchart-code-diagram> END
41router = APIRouter(prefix="/inference", tags=["Inference"])
43_DNS_LABEL_RE = re.compile(r"^[a-z0-9](?:[-a-z0-9]{0,61}[a-z0-9])?$")
44_SUPPORTED_METHODS = ["GET", "HEAD", "POST"]
45_HOP_BY_HOP_HEADERS = frozenset(
46 {
47 "connection",
48 "keep-alive",
49 "proxy-authenticate",
50 "proxy-authorization",
51 "te",
52 "trailer",
53 "transfer-encoding",
54 "upgrade",
55 }
56)
57_ALLOWED_REQUEST_HEADERS = frozenset(
58 {
59 "accept",
60 "accept-encoding",
61 "cache-control",
62 "content-encoding",
63 "content-type",
64 "idempotency-key",
65 "if-match",
66 "if-none-match",
67 "prefer",
68 "range",
69 "user-agent",
70 "x-request-id",
71 }
72)
73_BLOCKED_PATH_SEGMENTS = frozenset(
74 {"admin", "debug", "docs", "instances", "metrics", "openapi.json"}
75)
76_V1_MODELS_RE = re.compile(r"^v1/models(?:/[^/]+)?$")
77_V1_GENERATION_RE = re.compile(r"^v1/(?:chat/completions|completions|embeddings|responses)$")
78_V2_MODELS_RE = re.compile(r"^v2/models(?:/[^/]+(?:/(?:config|infer|ready|stats))?)?$")
81def _bounded_timeout(name: str, default: float, minimum: float, maximum: float) -> float:
82 """Read one finite, bounded timeout from the environment."""
83 raw = os.getenv(name)
84 if raw is None:
85 return default
86 try:
87 value = float(raw)
88 except ValueError:
89 return default
90 return value if minimum <= value <= maximum else default
93@lru_cache(maxsize=1)
94def _get_inference_store() -> InferenceEndpointStore:
95 """Create one process-local DynamoDB endpoint-store client lazily."""
96 return get_inference_endpoint_store()
99def _validate_label(value: object, field: str) -> str:
100 """Return a safe Kubernetes DNS label or reject the request."""
101 if not isinstance(value, str) or _DNS_LABEL_RE.fullmatch(value) is None:
102 raise HTTPException(status_code=404, detail=f"Invalid inference {field}")
103 return value
106def _target_service(endpoint: dict[str, Any], endpoint_name: str) -> str:
107 """Resolve the only in-cluster Service this endpoint may use.
109 Plain endpoints use ``<name>``. Mooncake disaggregated/both endpoints use
110 their reconciled ``<name>-proxy`` Service. During an active canary, a
111 cryptographically unbiased request sample is routed to ``<name>-canary``.
112 Every value is derived from a validated endpoint record, never from a URL or
113 header supplied by the caller.
114 """
115 spec = endpoint.get("spec")
116 if not isinstance(spec, dict):
117 raise HTTPException(status_code=503, detail="Inference endpoint has an invalid spec")
119 mooncake = spec.get("mooncake")
120 if isinstance(mooncake, dict) and mooncake.get("mode") in {"disaggregated", "both"}:
121 return _validate_label(f"{endpoint_name}-proxy", "service")
123 canary = spec.get("canary")
124 region = os.getenv("REGION", "")
125 region_status = endpoint.get("region_status")
126 local_status = region_status.get(region, {}) if isinstance(region_status, dict) else {}
127 canary_status = local_status.get("canary") if isinstance(local_status, dict) else None
128 if isinstance(canary, dict) and isinstance(canary_status, dict):
129 try:
130 weight = int(canary.get("weight", 0))
131 ready = int(canary_status.get("replicas_ready", 0))
132 desired = int(canary_status.get("replicas_desired", 0))
133 except TypeError, ValueError:
134 weight = ready = desired = 0
135 canary_is_ready = (
136 canary_status.get("state") == "running"
137 and canary_status.get("image") == canary.get("image")
138 and desired > 0
139 and ready >= desired
140 )
141 if canary_is_ready and 1 <= weight <= 99 and secrets.randbelow(100) < weight:
142 return _validate_label(f"{endpoint_name}-canary", "service")
144 return endpoint_name
147async def _resolve_upstream(endpoint_name: str) -> tuple[str, str, str]:
148 """Resolve the authorized Service, namespace, and configured health path."""
149 endpoint_name = _validate_label(endpoint_name, "name")
150 endpoint = await asyncio.to_thread(_get_inference_store().get_endpoint, endpoint_name)
151 if not endpoint:
152 raise HTTPException(
153 status_code=404, detail=f"Inference endpoint '{endpoint_name}' not found"
154 )
156 namespace = _validate_label(endpoint.get("namespace", "gco-inference"), "namespace")
157 allowed_namespace = os.getenv("INFERENCE_NAMESPACE", "gco-inference")
158 if namespace != allowed_namespace:
159 raise HTTPException(status_code=503, detail="Inference endpoint namespace is not routable")
161 region = os.getenv("REGION", "")
162 target_regions = endpoint.get("target_regions")
163 if not region or not isinstance(target_regions, list) or region not in target_regions:
164 raise HTTPException(
165 status_code=404, detail="Inference endpoint is not deployed in this region"
166 )
168 desired_state = endpoint.get("desired_state")
169 region_status = endpoint.get("region_status")
170 local_status = region_status.get(region, {}) if isinstance(region_status, dict) else {}
171 local_state = local_status.get("state") if isinstance(local_status, dict) else None
172 if desired_state != "running" or local_state != "running":
173 raise HTTPException(
174 status_code=503, detail="Inference endpoint is not ready in this region"
175 )
177 spec = endpoint.get("spec")
178 configured_health_path = (
179 spec.get("health_check_path", "/health") if isinstance(spec, dict) else "/health"
180 )
181 if not isinstance(configured_health_path, str) or not configured_health_path.startswith("/"):
182 configured_health_path = "/health"
184 return _target_service(endpoint, endpoint_name), namespace, configured_health_path
187def _request_headers(request: Request) -> list[tuple[str, str]]:
188 """Forward only explicitly supported end-to-end model request headers."""
189 return [
190 (name.lower(), value)
191 for name, value in request.headers.items()
192 if name.lower() in _ALLOWED_REQUEST_HEADERS
193 ]
196def _response_headers(response: httpx.Response) -> dict[str, str]:
197 """Copy end-to-end response headers while dropping hop-by-hop framing."""
198 blocked = _HOP_BY_HOP_HEADERS | {"content-length"}
199 return {name: value for name, value in response.headers.items() if name.lower() not in blocked}
202def _validate_upstream_path(
203 upstream_path: str,
204 method: str,
205 configured_health_path: str = "/health",
206) -> str:
207 """Allow serving/configured-health APIs while denying privileged paths."""
208 normalized = upstream_path.strip("/")
209 segments = [segment.lower() for segment in normalized.split("/") if segment]
210 if any(segment in _BLOCKED_PATH_SEGMENTS for segment in segments):
211 raise HTTPException(status_code=404, detail="Inference path is not exposed")
213 method = method.upper()
214 configured_health = configured_health_path.strip("/")
215 if (
216 not normalized
217 or normalized == "health"
218 or (configured_health and normalized == configured_health)
219 or normalized == "info"
220 or _V1_MODELS_RE.fullmatch(normalized)
221 ) and method in {"GET", "HEAD"}:
222 return normalized
223 if (
224 _V1_GENERATION_RE.fullmatch(normalized) or normalized in {"generate", "generate_stream"}
225 ) and method == "POST":
226 return normalized
227 if _V2_MODELS_RE.fullmatch(normalized) and method in {"GET", "HEAD", "POST"}:
228 return normalized
230 raise HTTPException(status_code=404, detail="Inference path is not exposed")
233async def _close_upstream(response: httpx.Response, client: httpx.AsyncClient) -> None:
234 await response.aclose()
235 await client.aclose()
238async def _stream_response(
239 response: httpx.Response, client: httpx.AsyncClient
240) -> AsyncIterator[bytes]:
241 """Yield the upstream response and shield connection cleanup on cancellation."""
242 try:
243 async for chunk in response.aiter_raw():
244 yield chunk
245 finally:
246 cleanup = asyncio.create_task(_close_upstream(response, client))
247 await asyncio.shield(cleanup)
250async def _proxy(
251 request: Request, endpoint_name: str, upstream_path: str = ""
252) -> StreamingResponse:
253 """Forward one authenticated request to a managed in-cluster endpoint."""
254 if any(part in {".", ".."} for part in upstream_path.split("/")):
255 raise HTTPException(status_code=400, detail="Invalid inference path")
256 service_name, namespace, configured_health_path = await _resolve_upstream(endpoint_name)
257 upstream_path = _validate_upstream_path(
258 upstream_path,
259 request.method,
260 configured_health_path,
261 )
262 encoded_suffix = quote(upstream_path, safe="/:@-._~")
263 upstream_path_value = f"/{encoded_suffix}" if encoded_suffix else "/"
264 upstream_url = ( # nosemgrep: python.django.security.injection.tainted-url-host.tainted-url-host
265 # Both host labels passed the strict Kubernetes DNS-label allowlist in
266 # _resolve_upstream; callers cannot supply a URL, address, or suffix.
267 f"http://{service_name}.{namespace}.svc.cluster.local{upstream_path_value}"
268 )
270 body = await request.body()
271 timeout = httpx.Timeout(
272 connect=_bounded_timeout("INFERENCE_PROXY_CONNECT_TIMEOUT_SECONDS", 5.0, 0.1, 30.0),
273 read=_bounded_timeout("INFERENCE_PROXY_READ_TIMEOUT_SECONDS", 300.0, 1.0, 900.0),
274 write=_bounded_timeout("INFERENCE_PROXY_WRITE_TIMEOUT_SECONDS", 30.0, 1.0, 300.0),
275 pool=_bounded_timeout("INFERENCE_PROXY_POOL_TIMEOUT_SECONDS", 5.0, 0.1, 30.0),
276 )
277 client = httpx.AsyncClient(timeout=timeout, follow_redirects=False, trust_env=False)
278 try:
279 upstream_request = client.build_request(
280 request.method,
281 upstream_url,
282 params=list(request.query_params.multi_items()),
283 headers=_request_headers(request),
284 content=body,
285 )
286 response = await client.send(upstream_request, stream=True)
287 except httpx.TimeoutException as exc:
288 await client.aclose()
289 raise HTTPException(status_code=504, detail="Inference endpoint timed out") from exc
290 except httpx.HTTPError as exc:
291 await client.aclose()
292 raise HTTPException(status_code=502, detail="Inference endpoint is unavailable") from exc
294 return StreamingResponse(
295 _stream_response(response, client),
296 status_code=response.status_code,
297 headers=_response_headers(response),
298 media_type=None,
299 )
302async def proxy_inference_root(request: Request, endpoint_name: str) -> StreamingResponse:
303 """Proxy an endpoint-root request after platform authentication."""
304 return await _proxy(request, endpoint_name)
307async def proxy_inference_path(
308 request: Request,
309 endpoint_name: str,
310 upstream_path: str,
311) -> StreamingResponse:
312 """Proxy an endpoint sub-path after platform authentication."""
313 return await _proxy(request, endpoint_name, upstream_path)
316# Register one route per method rather than a single multi-method route.
317# FastAPI derives an operation's ``operationId`` from ``generate_unique_id``,
318# which appends ``list(route.methods)[0]`` — a single arbitrary member of an
319# unordered set — and computes it once per route. A route carrying GET, HEAD,
320# and POST therefore emits three OpenAPI operations sharing one operationId,
321# which violates the spec's uniqueness requirement, makes generated clients
322# collide, and raises a UserWarning on every schema build. One method per
323# route keeps each generated operationId distinct while leaving request
324# handling byte-for-byte identical.
325for _path, _endpoint in (
326 ("/{endpoint_name}", proxy_inference_root),
327 ("/{endpoint_name}/{upstream_path:path}", proxy_inference_path),
328):
329 for _method in _SUPPORTED_METHODS:
330 router.add_api_route(_path, _endpoint, methods=[_method])