Spaces:
Sleeping
Sleeping
File size: 4,260 Bytes
ec855e6 04c779b ec855e6 04c779b ec855e6 | 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 | # app/main.py
"""
Application entrypoint.
- Configures FastAPI app and CORS middleware.
- Includes embed router.
- Starts uvicorn when run as __main__.
"""
import torch
from contextlib import asynccontextmanager
import uvicorn
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from .config import settings
from .logger import logger
from .controllers.embed_controller import router as embed_router, service as embedding_service
@asynccontextmanager
async def lifespan(app: FastAPI):
# Startup logic
gpu_available = torch.cuda.is_available()
device_to_use = settings.DEVICE or ("cuda" if gpu_available else "cpu")
logger.info("┌────────────────────────────────────────┐")
logger.info("│ Embedding Server Start │")
logger.info("└────────────────────────────────────────┘")
if gpu_available:
logger.info(f"GPU Detected: [bold]{torch.cuda.get_device_name(0)}[/bold]")
else:
logger.warning("GPU not detected, falling back to CPU.")
logger.info(f"Active Device: {device_to_use}")
logger.info(f"Serving on: {settings.HOST}:{settings.PORT}")
# Eagerly load the model to RAM/VRAM on startup
try:
await embedding_service.load_model()
except Exception as e:
logger.error("CRITICAL: Failed to load model during startup.")
logger.exception(e)
# We don't exit here to allow the health check to remain accessible,
# but the model status will be 'unloaded'.
yield
# Shutdown logic
logger.info("Shutting down server...")
app = FastAPI(
title="BGE-M3 Multilingual Embedding Server",
summary="High-performance API for BAAI/bge-m3 embeddings.",
description="""
## Features
- **Multilingual Support**: Embed text in 100+ languages.
- **GPU Acceleration**: Automatically uses CUDA if available.
- **Batch Processing**: Efficiently handle lists of text.
- **Security**: Strict API Key protection.
## Usage
- **Required Header**: `x-api-key` must be provided for all `POST` / `GET` requests.
- **Interactive Docs**: Use the **Try it out** button to test endpoints directly from the browser.
""",
version="1.0.0",
lifespan=lifespan,
docs_url="/docs",
redoc_url="/redoc"
)
# Global Exception Handler for unhandled errors
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
logger.error(f"Unhandled Exception: {type(exc).__name__}")
logger.error(f"Path: {request.url.path}")
logger.exception(exc) # This will log the full traceback in red
return JSONResponse(
status_code=500,
content={"detail": "An internal server error occurred.", "type": type(exc).__name__},
)
# CORS middleware - allow all (security handled by API Key guard)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["GET", "POST"],
allow_headers=["*"],
)
import time
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
start_time = time.perf_counter()
response = await call_next(request)
process_time = (time.perf_counter() - start_time) * 1000
logger.info(f"Request: {request.method} {request.url.path} - Completed in {process_time:.2f}ms")
return response
@app.get("/", tags=["status"])
async def root():
"""Root endpoint providing server status and documentation links."""
return {
"status": "online",
"message": "BGE-M3 Multilingual Embedding Server is running.",
"version": app.version,
"docs_url": "/docs",
"health_check": "/embed/health"
}
app.include_router(embed_router)
if __name__ == "__main__":
uvicorn.run("app.main:app", host=settings.HOST, port=settings.PORT, workers=settings.WORKERS)
|