Spaces:
Running
Running
File size: 3,857 Bytes
2322c22 846dbfc 677e796 846dbfc 2322c22 aeebd04 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | import asyncio
import importlib
import logging
import os
from contextlib import asynccontextmanager
from typing import Callable
from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from registry import PROJECT_REGISTRY
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("hf-backend")
loaded_projects: dict[str, str] = {}
failed_projects: dict[str, str] = {}
max_body_size = int(os.getenv("MAX_BODY_SIZE_BYTES", str(10 * 1024 * 1024)))
request_timeout = float(os.getenv("REQUEST_TIMEOUT_SECONDS", "60.0"))
class RequestLimitMiddleware(BaseHTTPMiddleware):
def __init__(self, app, max_size: int, timeout: float):
super().__init__(app)
self.max_size = max_size
self.timeout = timeout
async def dispatch(self, request: Request, call_next: Callable) -> Response:
content_length = request.headers.get("content-length")
if content_length and int(content_length) > self.max_size:
return Response(content="Payload Too Large", status_code=413)
bytes_received = 0
original_receive = request._receive
async def limited_receive():
nonlocal bytes_received
message = await original_receive()
if message["type"] == "http.request":
bytes_received += len(message.get("body", b""))
if bytes_received > self.max_size:
raise HTTPException(status_code=413, detail="Payload Too Large")
return message
request._receive = limited_receive
try:
return await asyncio.wait_for(call_next(request), timeout=self.timeout)
except asyncio.TimeoutError:
return Response(content="Request Timeout", status_code=504)
app = FastAPI(
title="HF Backend",
docs_url=None,
redoc_url=None,
openapi_url=None,
)
app.add_middleware(RequestLimitMiddleware, max_size=max_body_size, timeout=request_timeout)
raw_origins = [o.strip() for o in os.getenv("CORS_ALLOWED_ORIGINS", "").split(",") if o.strip()]
allowed_origins = [o for o in raw_origins if o != "*"]
if allowed_origins:
app.add_middleware(
CORSMiddleware,
allow_origins=allowed_origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
def load_routers():
disabled_raw = [p.strip().strip("/") for p in os.getenv("DISABLE_PROJECTS", "").split(",") if p.strip()]
disabled_set = {p for p in disabled_raw}
for prefix, target in PROJECT_REGISTRY.items():
clean_prefix = prefix.strip("/")
if clean_prefix in disabled_set:
logger.warning(f"[{prefix}] Skipped: disabled by DISABLE_PROJECTS")
continue
try:
logger.info(f"[{prefix}] Importing {target}")
module_name, attr_name = target.split(":")
module = importlib.import_module(module_name)
router = getattr(module, attr_name)
app.include_router(router, prefix=prefix)
loaded_projects[prefix] = target
logger.info(f"[{prefix}] Successfully mounted at {prefix}")
except Exception as e:
failed_projects[prefix] = f"{type(e).__name__}: {str(e)}"
logger.error(f"[{prefix}] Failed to mount router: {e}", exc_info=True)
load_routers()
@app.get("/health")
def health():
if not failed_projects:
overall_status = "healthy"
elif loaded_projects:
overall_status = "degraded"
else:
overall_status = "unhealthy"
return {
"status": overall_status,
"loaded": loaded_projects,
"failed": failed_projects,
}
@app.get("/")
def root():
return health()
|