51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import Depends, FastAPI, Header, HTTPException, status
|
|
|
|
from app.config import Settings, get_settings
|
|
from app.engine import InferenceEngine
|
|
from app.schemas import GenerateRequest, GenerateResponse, HealthResponse
|
|
|
|
engine: InferenceEngine | None = None
|
|
|
|
|
|
def verify_api_key(
|
|
settings: Settings = Depends(get_settings), x_api_key: str | None = Header(default=None)
|
|
) -> None:
|
|
if settings.api_key and x_api_key != settings.api_key:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Invalid API key",
|
|
)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_: FastAPI):
|
|
global engine
|
|
settings = get_settings()
|
|
engine = InferenceEngine(settings)
|
|
yield
|
|
if engine is not None:
|
|
engine.close()
|
|
engine = None
|
|
|
|
|
|
app = FastAPI(title="ROCm vLLM Inference API", version="1.0.0", lifespan=lifespan)
|
|
|
|
|
|
@app.get("/health", response_model=HealthResponse)
|
|
def health(settings: Settings = Depends(get_settings)) -> HealthResponse:
|
|
return HealthResponse(status="ok", model=settings.served_model_name or settings.model_name)
|
|
|
|
|
|
@app.post("/v1/generate", response_model=GenerateResponse, dependencies=[Depends(verify_api_key)])
|
|
def generate(req: GenerateRequest, settings: Settings = Depends(get_settings)) -> GenerateResponse:
|
|
if engine is None:
|
|
raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Engine not ready")
|
|
if req.max_tokens > settings.max_tokens:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
|
detail=f"max_tokens must be <= {settings.max_tokens}",
|
|
)
|
|
return engine.generate(req)
|