| """HTTP service around one Bauxi model, built to stay inside a 100m / 80 MiB pod.""" |
|
|
| from collections.abc import AsyncIterator |
| from contextlib import asynccontextmanager |
| from pathlib import Path |
| from typing import Annotated, Any |
|
|
| import anyio |
| from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request |
| from pydantic import BaseModel, ConfigDict, Field |
|
|
| from bauxi import MAX_CHARS, Bauxi, Label, __version__, release_unused_memory |
|
|
|
|
| class DecideRequest(BaseModel): |
| model_config = ConfigDict(str_strip_whitespace=True, extra="forbid") |
| decision: str = Field(min_length=1, max_length=MAX_CHARS) |
| evidence: str = Field(min_length=1, max_length=MAX_CHARS) |
|
|
|
|
| class DecideResponse(BaseModel): |
| prediction: Label |
| probabilities: dict[Label, float] |
| calibrated: bool |
|
|
|
|
| @asynccontextmanager |
| async def lifespan(_: FastAPI) -> AsyncIterator[None]: |
| |
| anyio.to_thread.current_default_thread_limiter().total_tokens = 1 |
| yield |
|
|
|
|
| app = FastAPI(title="Bauxi", version=__version__, lifespan=lifespan) |
| router = APIRouter() |
|
|
|
|
| def get_model(request: Request) -> Bauxi: |
| model: Bauxi = request.app.state.model |
| return model |
|
|
|
|
| @router.post("/decide") |
| def decide(body: DecideRequest, model: Annotated[Bauxi, Depends(get_model)]) -> DecideResponse: |
| try: |
| result = model.predict(body.decision, body.evidence) |
| except ValueError as error: |
| raise HTTPException(status_code=422, detail=str(error)) from error |
| finally: |
| release_unused_memory() |
| return DecideResponse( |
| prediction=result.label, |
| probabilities=result.probabilities, |
| calibrated=result.calibrated, |
| ) |
|
|
|
|
| @router.get("/health") |
| def health() -> dict[str, str]: |
| return {"status": "ok", "device": "cpu"} |
|
|
|
|
| @router.get("/runtime") |
| def runtime(model: Annotated[Bauxi, Depends(get_model)]) -> dict[str, Any]: |
| """What is loaded and what the cgroup measures, for the verifier and benchmarks.""" |
| root = Path("/sys/fs/cgroup") |
| names = ("memory.current", "memory.peak", "memory.max", "memory.events", "cpu.max", "cpu.stat") |
| return { |
| "version": model.metadata["version"], |
| "artifact": model.artifact, |
| "artifact_sha256": model.metadata["artifacts"][model.artifact]["sha256"], |
| "providers": model.session.get_providers(), |
| "cgroup": { |
| name: (root / name).read_text().strip() for name in names if (root / name).exists() |
| }, |
| } |
|
|
|
|
| app.include_router(router) |
|
|
