Spaces:
Sleeping
Sleeping
File size: 6,326 Bytes
a4508e8 | 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 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | """
GPU Perf Prophet — FastAPI application.
Routes
------
GET /health liveness probe
GET /gpus list all GPU ids in the spec DB
GET /models list all supported LLM model names
POST /predict predict throughput for one (GPU, workload) pair
POST /predict/batch vectorised prediction for up to 50 pairs
POST /recommend Pareto GPU recommendation for a workload
"""
from __future__ import annotations
import logging
from contextlib import asynccontextmanager
from typing import Annotated
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from pydantic import Field
from src.api.schemas import (
PredictRequest,
PredictResponse,
RecommendRequest,
RecommendResponse,
)
from src.data.gpu_spec_db import load_specs
from src.models.predictor import GpuPredictor, VALID_MODELS
from src.recommend.recommender import GpuRecommender
# Hard cap on raw request body size — prevents parsing-based DoS before any
# Pydantic validation runs. 50 PredictRequest objects at ~200 bytes each ≈
# 10 KB; 1 MB is a generous ceiling that still blocks runaway bodies.
_MAX_BODY_BYTES = 1 * 1024 * 1024 # 1 MB
log = logging.getLogger(__name__)
_predictor: GpuPredictor | None = None
_recommender: GpuRecommender | None = None
_gpu_list: list[dict] | None = None
_model_list: list[str] = sorted(VALID_MODELS)
@asynccontextmanager
async def lifespan(app: FastAPI):
global _predictor, _recommender, _gpu_list
log.info("Loading model …")
_predictor = GpuPredictor()
_recommender = GpuRecommender(_predictor)
specs = load_specs()
_gpu_list = [
{
"id": s["id"],
"name": s.get("name", s["id"]),
"vendor": s.get("vendor"),
"architecture": s.get("architecture"),
"vram_gb": s.get("vram_gb"),
"in_model_scope": s.get("in_model_scope", False),
}
for s in specs
]
log.info("Ready.")
yield
_predictor = None
_recommender = None
_gpu_list = None
app = FastAPI(
title="GPU Perf Prophet",
description=(
"Cross-vendor LLM inference performance forecasting "
"and workload-aware GPU recommendation."
),
version="1.0.0",
lifespan=lifespan,
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["GET", "POST"],
allow_headers=["Content-Type"],
)
@app.middleware("http")
async def limit_body_size(request: Request, call_next) -> Response:
# Fast path: reject if Content-Length header exceeds the limit.
# A malformed or missing Content-Length header is not an error here —
# it just means we must check the actual body below.
cl = request.headers.get("content-length")
if cl is not None:
try:
if int(cl) > _MAX_BODY_BYTES:
return Response(status_code=413, content="Request body too large")
except ValueError:
return Response(status_code=400, content="Invalid Content-Length header")
# GET / HEAD / OPTIONS carry no request body — no route handler reads one,
# so the size limit is irrelevant and buffering is wasted work.
if request.method in {"GET", "HEAD", "OPTIONS"}:
return await call_next(request)
# Slow path: consume the actual body so chunked-transfer or lying clients
# cannot bypass the Content-Length check above. Without this step, a
# client sending without Content-Length (chunked encoding) would bypass
# the limit entirely because uvicorn imposes no default body size cap.
body = await request.body()
if len(body) > _MAX_BODY_BYTES:
return Response(status_code=413, content="Request body too large")
# Re-inject the consumed body so downstream route handlers can read it.
async def _receive() -> dict:
return {"type": "http.request", "body": body, "more_body": False}
request._receive = _receive
return await call_next(request)
def _get_predictor() -> GpuPredictor:
if _predictor is None:
raise HTTPException(status_code=503, detail="Model not loaded")
return _predictor
def _get_recommender() -> GpuRecommender:
if _recommender is None:
raise HTTPException(status_code=503, detail="Recommender not loaded")
return _recommender
# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------
@app.get("/health")
def health() -> dict:
return {"status": "ok", "model_loaded": _predictor is not None}
@app.get("/gpus")
def list_gpus() -> dict:
return {"gpus": _gpu_list or []}
@app.get("/models")
def list_models() -> dict:
return {"models": _model_list}
@app.post("/predict", response_model=PredictResponse)
def predict(req: PredictRequest) -> dict:
predictor = _get_predictor()
try:
return predictor.predict(
gpu_id=req.gpu_id,
model_name=req.model_name,
scenario=req.scenario,
accuracy_tier=req.accuracy_tier,
framework=req.framework,
)
except ValueError as exc:
raise HTTPException(status_code=422, detail=str(exc)) from exc
@app.post("/predict/batch", response_model=list[PredictResponse])
def predict_batch(
requests: Annotated[list[PredictRequest], Field(max_length=50)],
) -> list[dict]:
predictor = _get_predictor()
try:
return predictor.predict_batch([r.model_dump() for r in requests])
except ValueError as exc:
raise HTTPException(status_code=422, detail=str(exc)) from exc
@app.post("/recommend", response_model=RecommendResponse)
def recommend(req: RecommendRequest) -> dict:
recommender = _get_recommender()
try:
return recommender.recommend(
model_name=req.model_name,
scenario=req.scenario,
accuracy_tier=req.accuracy_tier,
framework=req.framework,
budget_per_gpu_hr=req.budget_per_gpu_hr,
min_throughput_tok_per_sec=req.min_throughput_tok_per_sec,
)
except ValueError as exc:
raise HTTPException(status_code=422, detail=str(exc)) from exc
|