Coverage for gco/services/api_routes/inference_proxy.py: 100.00%

1"""Authenticated reverse proxy for managed inference endpoints. 

2 

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""" 

11 

12from __future__ import annotations 

13 

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 

22 

23import httpx 

24from fastapi import APIRouter, HTTPException, Request 

25from fastapi.responses import StreamingResponse 

26 

27from gco.services.inference_store import InferenceEndpointStore, get_inference_endpoint_store 

28 

29# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit 

30# Generated at (UTC): 2026-07-18T01:03:40Z 

31# Flowchart(s) generated from this file: 

32# * ``_resolve_upstream`` -> ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._resolve_upstream.html`` 

33# (PNG: ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._resolve_upstream.png``) 

34# * ``_proxy`` -> ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._proxy.html`` 

35# (PNG: ``diagrams/code_diagrams/gco/services/api_routes/inference_proxy._proxy.png``) 

36# Regenerate with ``python diagrams/code_diagrams/generate.py``. 

37# <pyflowchart-code-diagram> END 

38 

39 

40router = APIRouter(prefix="/inference", tags=["Inference"]) 

41 

42_DNS_LABEL_RE = re.compile(r"^[a-z0-9](?:[-a-z0-9]{0,61}[a-z0-9])?$") 

43_SUPPORTED_METHODS = ["GET", "HEAD", "POST"] 

44_HOP_BY_HOP_HEADERS = frozenset( 

45 { 

46 "connection", 

47 "keep-alive", 

48 "proxy-authenticate", 

49 "proxy-authorization", 

50 "te", 

51 "trailer", 

52 "transfer-encoding", 

53 "upgrade", 

54 } 

55) 

56_ALLOWED_REQUEST_HEADERS = frozenset( 

57 { 

58 "accept", 

59 "accept-encoding", 

60 "cache-control", 

61 "content-encoding", 

62 "content-type", 

63 "idempotency-key", 

64 "if-match", 

65 "if-none-match", 

66 "prefer", 

67 "range", 

68 "user-agent", 

69 "x-request-id", 

70 } 

71) 

72_BLOCKED_PATH_SEGMENTS = frozenset( 

73 {"admin", "debug", "docs", "instances", "metrics", "openapi.json"} 

74) 

75_V1_MODELS_RE = re.compile(r"^v1/models(?:/[^/]+)?$") 

76_V1_GENERATION_RE = re.compile(r"^v1/(?:chat/completions|completions|embeddings|responses)$") 

77_V2_MODELS_RE = re.compile(r"^v2/models(?:/[^/]+(?:/(?:config|infer|ready|stats))?)?$") 

78 

79 

80def _bounded_timeout(name: str, default: float, minimum: float, maximum: float) -> float: 

81 """Read one finite, bounded timeout from the environment.""" 

82 raw = os.getenv(name) 

83 if raw is None: 

84 return default 

85 try: 

86 value = float(raw) 

87 except ValueError: 

88 return default 

89 return value if minimum <= value <= maximum else default 

90 

91 

92@lru_cache(maxsize=1) 

93def _get_inference_store() -> InferenceEndpointStore: 

94 """Create one process-local DynamoDB endpoint-store client lazily.""" 

95 return get_inference_endpoint_store() 

96 

97 

98def _validate_label(value: object, field: str) -> str: 

99 """Return a safe Kubernetes DNS label or reject the request.""" 

100 if not isinstance(value, str) or _DNS_LABEL_RE.fullmatch(value) is None: 

101 raise HTTPException(status_code=404, detail=f"Invalid inference {field}") 

102 return value 

103 

104 

105def _target_service(endpoint: dict[str, Any], endpoint_name: str) -> str: 

106 """Resolve the only in-cluster Service this endpoint may use. 

107 

108 Plain endpoints use ``<name>``. Mooncake disaggregated/both endpoints use 

109 their reconciled ``<name>-proxy`` Service. During an active canary, a 

110 cryptographically unbiased request sample is routed to ``<name>-canary``. 

111 Every value is derived from a validated endpoint record, never from a URL or 

112 header supplied by the caller. 

113 """ 

114 spec = endpoint.get("spec") 

115 if not isinstance(spec, dict): 

116 raise HTTPException(status_code=503, detail="Inference endpoint has an invalid spec") 

117 

118 mooncake = spec.get("mooncake") 

119 if isinstance(mooncake, dict) and mooncake.get("mode") in {"disaggregated", "both"}: 

120 return _validate_label(f"{endpoint_name}-proxy", "service") 

121 

122 canary = spec.get("canary") 

123 region = os.getenv("REGION", "") 

124 region_status = endpoint.get("region_status") 

125 local_status = region_status.get(region, {}) if isinstance(region_status, dict) else {} 

126 canary_status = local_status.get("canary") if isinstance(local_status, dict) else None 

127 if isinstance(canary, dict) and isinstance(canary_status, dict): 

128 try: 

129 weight = int(canary.get("weight", 0)) 

130 ready = int(canary_status.get("replicas_ready", 0)) 

131 desired = int(canary_status.get("replicas_desired", 0)) 

132 except TypeError, ValueError: 

133 weight = ready = desired = 0 

134 canary_is_ready = ( 

135 canary_status.get("state") == "running" 

136 and canary_status.get("image") == canary.get("image") 

137 and desired > 0 

138 and ready >= desired 

139 ) 

140 if canary_is_ready and 1 <= weight <= 99 and secrets.randbelow(100) < weight: 

141 return _validate_label(f"{endpoint_name}-canary", "service") 

142 

143 return endpoint_name 

144 

145 

146async def _resolve_upstream(endpoint_name: str) -> tuple[str, str, str]: 

147 """Resolve the authorized Service, namespace, and configured health path.""" 

148 endpoint_name = _validate_label(endpoint_name, "name") 

149 endpoint = await asyncio.to_thread(_get_inference_store().get_endpoint, endpoint_name) 

150 if not endpoint: 

151 raise HTTPException( 

152 status_code=404, detail=f"Inference endpoint '{endpoint_name}' not found" 

153 ) 

154 

155 namespace = _validate_label(endpoint.get("namespace", "gco-inference"), "namespace") 

156 allowed_namespace = os.getenv("INFERENCE_NAMESPACE", "gco-inference") 

157 if namespace != allowed_namespace: 

158 raise HTTPException(status_code=503, detail="Inference endpoint namespace is not routable") 

159 

160 region = os.getenv("REGION", "") 

161 target_regions = endpoint.get("target_regions") 

162 if not region or not isinstance(target_regions, list) or region not in target_regions: 

163 raise HTTPException( 

164 status_code=404, detail="Inference endpoint is not deployed in this region" 

165 ) 

166 

167 desired_state = endpoint.get("desired_state") 

168 region_status = endpoint.get("region_status") 

169 local_status = region_status.get(region, {}) if isinstance(region_status, dict) else {} 

170 local_state = local_status.get("state") if isinstance(local_status, dict) else None 

171 if desired_state != "running" or local_state != "running": 

172 raise HTTPException( 

173 status_code=503, detail="Inference endpoint is not ready in this region" 

174 ) 

175 

176 spec = endpoint.get("spec") 

177 configured_health_path = ( 

178 spec.get("health_check_path", "/health") if isinstance(spec, dict) else "/health" 

179 ) 

180 if not isinstance(configured_health_path, str) or not configured_health_path.startswith("/"): 

181 configured_health_path = "/health" 

182 

183 return _target_service(endpoint, endpoint_name), namespace, configured_health_path 

184 

185 

186def _request_headers(request: Request) -> list[tuple[str, str]]: 

187 """Forward only explicitly supported end-to-end model request headers.""" 

188 return [ 

189 (name.lower(), value) 

190 for name, value in request.headers.items() 

191 if name.lower() in _ALLOWED_REQUEST_HEADERS 

192 ] 

193 

194 

195def _response_headers(response: httpx.Response) -> dict[str, str]: 

196 """Copy end-to-end response headers while dropping hop-by-hop framing.""" 

197 blocked = _HOP_BY_HOP_HEADERS | {"content-length"} 

198 return {name: value for name, value in response.headers.items() if name.lower() not in blocked} 

199 

200 

201def _validate_upstream_path( 

202 upstream_path: str, 

203 method: str, 

204 configured_health_path: str = "/health", 

205) -> str: 

206 """Allow serving/configured-health APIs while denying privileged paths.""" 

207 normalized = upstream_path.strip("/") 

208 segments = [segment.lower() for segment in normalized.split("/") if segment] 

209 if any(segment in _BLOCKED_PATH_SEGMENTS for segment in segments): 

210 raise HTTPException(status_code=404, detail="Inference path is not exposed") 

211 

212 method = method.upper() 

213 configured_health = configured_health_path.strip("/") 

214 if ( 

215 not normalized 

216 or normalized == "health" 

217 or (configured_health and normalized == configured_health) 

218 or _V1_MODELS_RE.fullmatch(normalized) 

219 ) and method in {"GET", "HEAD"}: 

220 return normalized 

221 if ( 

222 _V1_GENERATION_RE.fullmatch(normalized) or normalized in {"generate", "generate_stream"} 

223 ) and method == "POST": 

224 return normalized 

225 if _V2_MODELS_RE.fullmatch(normalized) and method in {"GET", "HEAD", "POST"}: 

226 return normalized 

227 

228 raise HTTPException(status_code=404, detail="Inference path is not exposed") 

229 

230 

231async def _close_upstream(response: httpx.Response, client: httpx.AsyncClient) -> None: 

232 await response.aclose() 

233 await client.aclose() 

234 

235 

236async def _stream_response( 

237 response: httpx.Response, client: httpx.AsyncClient 

238) -> AsyncIterator[bytes]: 

239 """Yield the upstream response and shield connection cleanup on cancellation.""" 

240 try: 

241 async for chunk in response.aiter_raw(): 

242 yield chunk 

243 finally: 

244 cleanup = asyncio.create_task(_close_upstream(response, client)) 

245 await asyncio.shield(cleanup) 

246 

247 

248async def _proxy( 

249 request: Request, endpoint_name: str, upstream_path: str = "" 

250) -> StreamingResponse: 

251 """Forward one authenticated request to a managed in-cluster endpoint.""" 

252 if any(part in {".", ".."} for part in upstream_path.split("/")): 

253 raise HTTPException(status_code=400, detail="Invalid inference path") 

254 service_name, namespace, configured_health_path = await _resolve_upstream(endpoint_name) 

255 upstream_path = _validate_upstream_path( 

256 upstream_path, 

257 request.method, 

258 configured_health_path, 

259 ) 

260 encoded_suffix = quote(upstream_path, safe="/:@-._~") 

261 upstream_path_value = f"/{encoded_suffix}" if encoded_suffix else "/" 

262 upstream_url = ( # nosemgrep: python.django.security.injection.tainted-url-host.tainted-url-host 

263 # Both host labels passed the strict Kubernetes DNS-label allowlist in 

264 # _resolve_upstream; callers cannot supply a URL, address, or suffix. 

265 f"http://{service_name}.{namespace}.svc.cluster.local{upstream_path_value}" 

266 ) 

267 

268 body = await request.body() 

269 timeout = httpx.Timeout( 

270 connect=_bounded_timeout("INFERENCE_PROXY_CONNECT_TIMEOUT_SECONDS", 5.0, 0.1, 30.0), 

271 read=_bounded_timeout("INFERENCE_PROXY_READ_TIMEOUT_SECONDS", 300.0, 1.0, 900.0), 

272 write=_bounded_timeout("INFERENCE_PROXY_WRITE_TIMEOUT_SECONDS", 30.0, 1.0, 300.0), 

273 pool=_bounded_timeout("INFERENCE_PROXY_POOL_TIMEOUT_SECONDS", 5.0, 0.1, 30.0), 

274 ) 

275 client = httpx.AsyncClient(timeout=timeout, follow_redirects=False, trust_env=False) 

276 try: 

277 upstream_request = client.build_request( 

278 request.method, 

279 upstream_url, 

280 params=list(request.query_params.multi_items()), 

281 headers=_request_headers(request), 

282 content=body, 

283 ) 

284 response = await client.send(upstream_request, stream=True) 

285 except httpx.TimeoutException as exc: 

286 await client.aclose() 

287 raise HTTPException(status_code=504, detail="Inference endpoint timed out") from exc 

288 except httpx.HTTPError as exc: 

289 await client.aclose() 

290 raise HTTPException(status_code=502, detail="Inference endpoint is unavailable") from exc 

291 

292 return StreamingResponse( 

293 _stream_response(response, client), 

294 status_code=response.status_code, 

295 headers=_response_headers(response), 

296 media_type=None, 

297 ) 

298 

299 

300async def proxy_inference_root(request: Request, endpoint_name: str) -> StreamingResponse: 

301 """Proxy an endpoint-root request after platform authentication.""" 

302 return await _proxy(request, endpoint_name) 

303 

304 

305async def proxy_inference_path( 

306 request: Request, 

307 endpoint_name: str, 

308 upstream_path: str, 

309) -> StreamingResponse: 

310 """Proxy an endpoint sub-path after platform authentication.""" 

311 return await _proxy(request, endpoint_name, upstream_path) 

312 

313 

314# Register one route per method rather than a single multi-method route. 

315# FastAPI derives an operation's ``operationId`` from ``generate_unique_id``, 

316# which appends ``list(route.methods)[0]`` — a single arbitrary member of an 

317# unordered set — and computes it once per route. A route carrying GET, HEAD, 

318# and POST therefore emits three OpenAPI operations sharing one operationId, 

319# which violates the spec's uniqueness requirement, makes generated clients 

320# collide, and raises a UserWarning on every schema build. One method per 

321# route keeps each generated operationId distinct while leaving request 

322# handling byte-for-byte identical. 

323for _path, _endpoint in ( 

324 ("/{endpoint_name}", proxy_inference_root), 

325 ("/{endpoint_name}/{upstream_path:path}", proxy_inference_path), 

326): 

327 for _method in _SUPPORTED_METHODS: 

328 router.add_api_route(_path, _endpoint, methods=[_method])