SPB Git forge

spb/llm-api

Public
0commits 0branches 0releases
0 Bsize
maindefault branch
—last push
6.9 KB · 167 lines python
Raw Blame History
1"""LLM API application: startup sequence, routers, shutdown."""23from __future__ import annotations45import logging6import logging.handlers7import os8import sys9import time10from contextlib import asynccontextmanager1112import httpx13from fastapi import FastAPI, Request14from fastapi.exceptions import RequestValidationError15from fastapi.responses import JSONResponse1617from . import __version__18from .api import admin_routes, openai_routes19from .api.admin_routes import public_health20from .auth import Auth21from .config import Settings, get_settings22from .db import Database23from .downloads import Downloader24from .errors import APIError, api_error_handler25from .events import bus26from .hardware import detect_hardware27from .harvester import Harvester28from .jobs import JobRunner29from .manager import ModelManager30from .metrics import MetricsCollector31from .models.registry import Registry32from .proxy_ui import router as ui_router3334log = logging.getLogger("llm_api")353637def setup_logging(settings: Settings) -> None:38    settings.logs_path.mkdir(parents=True, exist_ok=True)39    fmt = logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")40    root = logging.getLogger()41    root.setLevel(logging.INFO)42    if not any(isinstance(h, logging.StreamHandler) for h in root.handlers):43        sh = logging.StreamHandler(sys.stdout)44        sh.setFormatter(fmt)45        root.addHandler(sh)46    fh = logging.handlers.RotatingFileHandler(settings.logs_path / "llm-api.log", maxBytes=20_000_000, backupCount=5)47    fh.setFormatter(fmt)48    root.addHandler(fh)49    logging.getLogger("httpx").setLevel(logging.WARNING)50    logging.getLogger("httpcore").setLevel(logging.WARNING)51    logging.getLogger("uvicorn.access").setLevel(logging.WARNING)525354@asynccontextmanager55async def lifespan(app: FastAPI):56    settings: Settings = app.state.settings57    t0 = time.time()58    # 1. config + dirs59    settings.ensure_dirs()60    # 2. database61    db = Database(settings.db_path)62    await db.connect()63    app.state.db = db64    # 3. hardware65    hw = detect_hardware(settings.models_dir)66    log.info("hardware: %s, %.0f GB, %s CPU cores, %s GPU cores, %s %s", hw.chip, hw.memory_gb, hw.cpu_cores, hw.gpu_cores, hw.os, hw.os_version)67    if settings.max_model_memory_gb > hw.memory_gb - settings.macos_reserve_gb:68        log.warning("MAX_MODEL_MEMORY_GB=%s is above memory minus the macOS reserve (%.0f GB)", settings.max_model_memory_gb,69                    hw.memory_gb - settings.macos_reserve_gb)70    # 4. auth71    auth = Auth(db, settings)72    app.state.auth = auth73    if await auth.user_count() == 0 and settings.admin_email and settings.admin_password:74        try:75            await auth.create_user(settings.admin_email, settings.admin_password)76            log.info("seeded admin user %s", settings.admin_email)77        except Exception as e:78            log.warning("could not seed admin user: %s", e)79    # 5. registry + manager (clears stale state)80    registry = Registry(db, settings)81    manager = ModelManager(settings, db, registry)82    app.state.manager = manager83    app.state.registry = registry84    jobs = JobRunner(db)85    await jobs.start()86    app.state.jobs = jobs87    downloader = Downloader(settings, db, registry, jobs)88    app.state.downloader = downloader89    app.state.harvester = Harvester(settings, db, registry, jobs, downloader)90    metrics = MetricsCollector(settings, db, manager)91    app.state.metrics = metrics92    app.state.ui_client = httpx.AsyncClient(timeout=httpx.Timeout(30.0, read=120.0), follow_redirects=False)93    # 6. scan registry94    try:95        summary = await registry.rescan()96        log.info("registry: %s", summary)97    except Exception:98        log.exception("initial rescan failed")99    # 7. start manager loops (kills stale workers, preload)100    await manager.start()101    await metrics.start()102    await db.audit("server.start", actor="system", detail={"version": __version__, "seconds": round(time.time() - t0, 2)})103    log.info("LLM API %s ready on %s:%s (%.1fs)", __version__, settings.host, settings.port, time.time() - t0)104    bus.publish("server", {"event": "started", "version": __version__})105    try:106        yield107    finally:108        log.info("shutting down: draining requests, unloading models")109        await metrics.stop()110        await manager.stop()111        await app.state.ui_client.aclose()112        await db.audit("server.stop", actor="system")113        await db.close()114115116def create_app(settings: Settings | None = None) -> FastAPI:117    settings = settings or get_settings()118    setup_logging(settings)119    app = FastAPI(title="LLM API", version=__version__, docs_url="/openapi", redoc_url=None, openapi_url="/openapi.json",120                  lifespan=lifespan)121    app.state.settings = settings122    app.add_exception_handler(APIError, api_error_handler)  # type: ignore[arg-type]123124    @app.exception_handler(RequestValidationError)125    async def _validation(_: Request, exc: RequestValidationError):126        errs = exc.errors()127        msg = "; ".join(f"{'.'.join(str(x) for x in e.get('loc', []))}: {e.get('msg')}" for e in errs[:3])128        return JSONResponse(status_code=400, content={"error": {"message": msg or "invalid request", "type": "invalid_request_error",129                                                                  "code": "INVALID_REQUEST", "param": None}})130131    @app.exception_handler(Exception)132    async def _unhandled(_: Request, exc: Exception):133        log.exception("unhandled error: %s", exc)134        return JSONResponse(status_code=500, content={"error": {"message": "Internal server error.", "type": "server_error",135                                                                  "code": "INTERNAL", "param": None}})136137    @app.middleware("http")138    async def _security_headers(request: Request, call_next):139        # Body size guard for JSON APIs140        cl = request.headers.get("content-length")141        if cl and cl.isdigit() and int(cl) > settings.max_body_bytes and request.url.path.startswith(("/v1", "/api")):142            return JSONResponse(status_code=413, content={"error": {"message": "Request body too large.", "type": "invalid_request_error",143                                                                      "code": "BODY_TOO_LARGE", "param": None}})144        resp = await call_next(request)145        resp.headers.setdefault("X-Content-Type-Options", "nosniff")146        resp.headers.setdefault("Referrer-Policy", "same-origin")147        resp.headers.setdefault("X-Frame-Options", "DENY")148        if request.url.path.startswith(("/v1", "/api")):149            resp.headers.setdefault("Cache-Control", "no-store")150        return resp151152    @app.get("/health", include_in_schema=False)153    async def health(request: Request):154        return await public_health(request)155156    @app.get("/api/status", include_in_schema=False)157    async def status(request: Request):158        return await public_health(request)159160    app.include_router(openai_routes.router)161    app.include_router(admin_routes.router)162    app.include_router(ui_router)  # must be last (catch-all)163    return app164165166app = create_app()167