Coverage for gco / services / inference_api.py: 100.00%

1"""Dedicated authenticated streaming proxy API for managed inference endpoints.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6import os 

7from collections.abc import AsyncIterator 

8from contextlib import asynccontextmanager 

9from typing import Any 

10 

11from fastapi import FastAPI, Request 

12from fastapi.responses import JSONResponse 

13 

14from gco.services import tracing 

15from gco.services.api_routes.inference_proxy import router as inference_proxy_router 

16from gco.services.auth_middleware import AuthenticationMiddleware 

17from gco.services.request_size_middleware import ( 

18 DEFAULT_MAX_REQUEST_BODY_BYTES, 

19 RequestSizeLimitMiddleware, 

20) 

21from gco.services.service_metrics import mount_metrics 

22 

23logging.basicConfig( 

24 level=logging.INFO, 

25 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", 

26) 

27logger = logging.getLogger(__name__) 

28 

29DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS = 900 

30 

31 

32@asynccontextmanager 

33async def lifespan(app: FastAPI) -> AsyncIterator[None]: 

34 """Flush queued spans at shutdown; startup has nothing to prepare. 

35 

36 The proxy keeps no process-wide upstream state: every request builds its 

37 model-endpoint client and releases it when its stream ends, so shutdown 

38 has no client to close. Uvicorn runs this shutdown only after it has 

39 drained (or, past the graceful budget, cancelled) the in-flight streams, 

40 so the spans of the last requests are already queued when the provider 

41 flushes. 

42 """ 

43 try: 

44 yield 

45 finally: 

46 tracing.shutdown_tracing() 

47 

48 

49app = FastAPI( 

50 title="GCO Inference Proxy API", 

51 description="Authenticated streaming reverse proxy for managed GCO inference endpoints", 

52 version="1.0.0", 

53 lifespan=lifespan, 

54) 

55 

56# Server spans for every route except the probe and scrape paths, with the 

57# model-endpoint client spans beneath them; inert unless GCO_TRACING_ENABLED=true. 

58tracing.configure_tracing("inference-proxy") 

59tracing.instrument_fastapi_app(app) 

60 

61# Starlette executes middleware in reverse registration order. The size limit 

62# runs first so an oversized request is rejected before secret retrieval or 

63# HMAC work; AuthenticationMiddleware still validates the exact cached body. 

64app.add_middleware(AuthenticationMiddleware) 

65_max_body_bytes = int(os.getenv("MAX_REQUEST_BODY_BYTES", str(DEFAULT_MAX_REQUEST_BODY_BYTES))) 

66app.add_middleware(RequestSizeLimitMiddleware, max_body_bytes=_max_body_bytes) 

67 

68mount_metrics(app, "inference-proxy") 

69 

70app.include_router(inference_proxy_router) 

71 

72 

73@app.get("/", tags=["Info"]) 

74async def root() -> dict[str, Any]: 

75 """Return a small authenticated service descriptor.""" 

76 return { 

77 "service": "GCO Inference Proxy API", 

78 "version": "1.0.0", 

79 "status": "running", 

80 "region": os.getenv("REGION", "unknown"), 

81 } 

82 

83 

84@app.get("/healthz", tags=["Health"]) 

85async def kubernetes_health_check() -> dict[str, str]: 

86 """Kubernetes and ALB liveness probe.""" 

87 return {"status": "ok"} 

88 

89 

90@app.get("/readyz", tags=["Health"]) 

91async def kubernetes_readiness_check() -> dict[str, str]: 

92 """Readiness probe; external dependencies are resolved per request.""" 

93 return {"status": "ready"} 

94 

95 

96@app.exception_handler(Exception) 

97async def global_exception_handler(request: Request, exc: Exception) -> JSONResponse: 

98 """Log unexpected failures without exposing internal routing details. 

99 

100 ``exc`` is intentionally unused: the traceback comes from the active 

101 exception, and the response deliberately says nothing about it. 

102 """ 

103 # LOG004 is suppressed because this IS an exception handler. Starlette 

104 # awaits it from inside its own `except` block, so the exception is still 

105 # active and `.exception()` records the traceback; ruff only sees that the 

106 # call is not lexically inside an `except`. Downgrading to `.error()` as it 

107 # suggests would silently drop the traceback that makes a 500 debuggable. 

108 logger.exception( # noqa: LOG004 

109 "Unhandled inference proxy exception for %s %s", request.method, request.url.path 

110 ) 

111 return JSONResponse(status_code=500, content={"error": "Internal server error"}) 

112 

113 

114def create_app() -> FastAPI: 

115 """Return the process-wide FastAPI application.""" 

116 return app 

117 

118 

119def _run_server() -> None: 

120 """Run Uvicorn with the same drain budget declared by the pod manifest.""" 

121 import uvicorn 

122 

123 from gco.services.uvicorn_runtime import describe_uvicorn_runtime 

124 

125 host = os.getenv("HOST", "0.0.0.0") # container listener 

126 port = int(os.getenv("PORT", "8080")) 

127 log_level = os.getenv("LOG_LEVEL", "info").lower() 

128 graceful_shutdown_seconds = int( 

129 os.getenv( 

130 "GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS", 

131 str(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS), 

132 ) 

133 ) 

134 loop_impl, http_impl = describe_uvicorn_runtime() 

135 logger.info( 

136 "Starting Inference Proxy API on %s:%d (uvicorn loop=%s http=%s)", 

137 host, 

138 port, 

139 loop_impl, 

140 http_impl, 

141 ) 

142 uvicorn.run( 

143 "gco.services.inference_api:app", 

144 host=host, 

145 port=port, 

146 log_level=log_level, 

147 reload=False, 

148 timeout_graceful_shutdown=graceful_shutdown_seconds, 

149 ) 

150 

151 

152if __name__ == "__main__": 

153 _run_server()