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

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

2 

3from __future__ import annotations 

4 

5import logging 

6import os 

7from typing import Any 

8 

9from fastapi import FastAPI, Request 

10from fastapi.responses import JSONResponse 

11 

12from gco.services.auth_middleware import AuthenticationMiddleware 

13from gco.services.request_size_middleware import ( 

14 DEFAULT_MAX_REQUEST_BODY_BYTES, 

15 RequestSizeLimitMiddleware, 

16) 

17from gco.services.service_metrics import mount_metrics 

18 

19logging.basicConfig( 

20 level=logging.INFO, 

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

22) 

23logger = logging.getLogger(__name__) 

24 

25DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS = 900 

26 

27app = FastAPI( 

28 title="GCO Inference Proxy API", 

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

30 version="1.0.0", 

31) 

32 

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

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

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

36app.add_middleware(AuthenticationMiddleware) 

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

38app.add_middleware(RequestSizeLimitMiddleware, max_body_bytes=_max_body_bytes) 

39 

40mount_metrics(app, "inference-proxy") 

41 

42from gco.services.api_routes.inference_proxy import router as inference_proxy_router # noqa: E402 

43 

44app.include_router(inference_proxy_router) 

45 

46 

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

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

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

50 return { 

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

52 "version": "1.0.0", 

53 "status": "running", 

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

55 } 

56 

57 

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

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

60 """Kubernetes and ALB liveness probe.""" 

61 return {"status": "ok"} 

62 

63 

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

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

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

67 return {"status": "ready"} 

68 

69 

70@app.exception_handler(Exception) 

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

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

73 logger.exception( 

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

75 ) 

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

77 

78 

79def create_app() -> FastAPI: 

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

81 return app 

82 

83 

84def _run_server() -> None: 

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

86 import uvicorn 

87 

88 host = os.getenv("HOST", "0.0.0.0") # nosec B104 — container listener 

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

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

91 graceful_shutdown_seconds = int( 

92 os.getenv( 

93 "GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS", 

94 str(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS), 

95 ) 

96 ) 

97 logger.info("Starting Inference Proxy API on %s:%d", host, port) 

98 uvicorn.run( 

99 "gco.services.inference_api:app", 

100 host=host, 

101 port=port, 

102 log_level=log_level, 

103 reload=False, 

104 timeout_graceful_shutdown=graceful_shutdown_seconds, 

105 ) 

106 

107 

108if __name__ == "__main__": 

109 _run_server()