code/tt_diffusion_planner/ttaw/server/app.py

# SPDX-License-Identifier: Apache-2.0
"""C08 server: the FastAPI app factory implementing the HTTP contract of all 13 bundles (BUNDLE_CONVENTIONS.md 7).

A bundle's ``server/app.py`` is a few lines::

    from ..api import CenterPoint
    from ..ttaw.server.app import ServerSpec, create_app, parse_mesh_shape  # noqa: F401 (verify: line)

    SPEC = ServerSpec(model_name=CenterPoint.MODEL_NAME, env_prefix="CENTERPOINT", model_cls=CenterPoint,
                      task="LiDAR 3D object detection", default_weights=CenterPoint.DEFAULT_REPO,
                      io="a LiDAR point cloud in, 3D boxes out",
                      autoware={"package": "autoware_lidar_centerpoint",
                                "path": "perception/autoware_lidar_centerpoint",
                                "autoware_universe": "9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd"},
                      source={"repo": "https://huggingface.co/changh95/centerpoint-p150", "license": "Apache-2.0"},
                      calib_dir=Path(__file__).resolve().parents[1] / "calib")
    app = create_app(SPEC)

Served by tt-model-manager as ``kind: tt-dit-server`` (``python -m uvicorn --lifespan on <pkg>.server.app:app``).
Everything device-touching happens in the lifespan: weights -> device -> graph -> trace capture of every warm-up
variant; uvicorn logs ``Application startup complete`` only after that, so READY means warm. Shutdown (SIGTERM,
120 s budget) closes the model under the lock. Importing the module has no side effects beyond importing
fastapi / pydantic (no device, no download, no environment read).

Tests replace the model with a stub: ``app.state.ttaw.model_factory = StubModel`` before ``TestClient(app)``.
``app.state.ttaw.predict(request)`` is the route handler itself (for API == server checks).

Environment (read in the lifespan): ``HF_MODEL``, ``TT_MODEL_WEIGHTS_REVISION`` (then ``TT_WEIGHTS_REVISION``),
``<ENV>_WEIGHTS_DIR``, ``TT_MESH_SHAPE`` (only 1x1), ``TT_DEVICE_ID``, ``<ENV>_DISPATCH``, ``<ENV>_NUM_CQS``,
``<ENV>_VARIANT``, ``<ENV>_WARMUP`` (JSON list, ``default`` or ``none``), ``<ENV>_MAX_BODY_MB`` (default 256),
``<ENV>_TORCH_THREADS``.
"""
import json
import logging
import math
import os
import re
import threading
import time
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, List, Literal, Mapping, Optional, Type

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel, ConfigDict, Field

from .. import io as tio

__all__ = [
    "HARDWARE", "PointsSpec", "SweepSpec", "ImageSpec", "StreamSpec", "ArraysSpec", "PredictRequest", "ServerSpec",
    "ServerState", "parse_mesh_shape", "config_from_env", "default_model_factory", "decode_request", "create_app",
]

HARDWARE = "Tenstorrent Blackhole p150 (single chip) via tt-nn"
log = logging.getLogger("ttaw.server")


# ------------------------------------------------------------------------------------------- request schema

class _Strict(BaseModel):
    model_config = ConfigDict(extra="forbid")  # a misspelled field is a 422, not a silently ignored option


class PointsSpec(_Strict):
    format: Literal["bin", "npy", "npz", "pcd", "list"] = "bin"
    data: Optional[str] = None                   # base64 of the file bytes (all formats except "list")
    values: Optional[List[List[float]]] = None   # "list" only
    fields: Optional[List[str]] = None           # column names (required meaning for "bin" / "list")
    dtype: Literal["float32", "float64"] = "float32"
    key: Optional[str] = None                    # "npz" only
    frame_id: Optional[str] = None


class SweepSpec(_Strict):
    points: PointsSpec
    time_lag_s: float = 0.0                      # age of this sweep relative to the current one
    T_current_from_sweep: Optional[Any] = None   # 4x4 / {translation, rotation_wxyz} / {x, y, z, roll, pitch, yaw}


class ImageSpec(_Strict):
    camera: str
    data: str                                    # base64 PNG / JPEG (or .npy uint8 HxWx3)
    format: Literal["auto", "png", "jpeg", "npy"] = "auto"
    intrinsics: Optional[Any] = None             # 3x3 / 3x4 / {fx, fy, cx, cy}
    T_ref_from_camera: Optional[Any] = None
    distortion: Optional[Dict[str, Any]] = None
    timestamp_s: Optional[float] = None


class StreamSpec(_Strict):
    id: str = "default"
    reset: bool = False
    timestamp_s: Optional[float] = None
    T_world_from_ego: Optional[Any] = None


class ArraysSpec(_Strict):
    format: Literal["npz", "json"] = "npz"
    data: Optional[str] = None
    arrays: Optional[Dict[str, Any]] = None


class PredictRequest(_Strict):
    """The common ``POST /predict`` envelope. Bundles with extra inputs subclass it (e.g. PointPainting ``rois``)
    and decode the extra fields in ``ServerSpec.decode_extra``."""

    points: Optional[PointsSpec] = None
    sweeps: Optional[List[SweepSpec]] = None
    images: Optional[List[ImageSpec]] = None
    calibration: Optional[Dict[str, Any]] = None
    stream: Optional[StreamSpec] = None
    inputs: Optional[ArraysSpec] = None
    params: Dict[str, Any] = Field(default_factory=dict)
    output_format: Literal["json", "npz"] = "json"


# ----------------------------------------------------------------------------------------------- config

@dataclass
class ServerSpec:
    """Identity and hooks of one bundle's server. ``model_cls`` is the bundle's ``ModelBase`` subclass (used by the
    default factory and for ``CAMERA_ORDER`` / calibration rules)."""

    model_name: str
    env_prefix: str
    model_cls: Any = None
    task: str = ""
    default_weights: str = ""
    owner: str = "changh95"
    hardware: str = HARDWARE
    io: str = ""
    autoware: Mapping[str, Any] = field(default_factory=dict)
    source: Mapping[str, Any] = field(default_factory=dict)
    calib_dir: Optional[Path] = None
    default_variant: Optional[str] = None
    version: str = "0.1.0"
    description: str = ""
    request_model: Type[PredictRequest] = PredictRequest
    # decode_extra(request, kwargs, max_bytes): add model-specific call kwargs from extra request fields
    decode_extra: Optional[Callable[[Any, Dict[str, Any], int], None]] = None
    # info_extra(model) -> dict merged into /info
    info_extra: Optional[Callable[[Any], Dict[str, Any]]] = None


_SHAPE_RE = re.compile(r"^\s*[\(\[]?\s*(\d+)\s*[xX,]\s*(\d+)\s*[\)\]]?\s*$")


def parse_mesh_shape(value: Optional[str]) -> tuple:
    """``"1x1"`` / ``"(1, 1)"`` / ``"1,1"`` -> ``(1, 1)``; unset -> ``(1, 1)``."""
    if value is None or not value.strip():
        return (1, 1)
    m = _SHAPE_RE.match(value)
    if not m:
        raise RuntimeError(f"TT_MESH_SHAPE={value!r} is not a mesh shape; expected '1x1', '(1, 1)' or '1,1'")
    return (int(m.group(1)), int(m.group(2)))


def config_from_env(spec: ServerSpec, env: Optional[Mapping[str, str]] = None) -> Dict[str, Any]:
    """The server configuration from the environment (read once, in the lifespan)."""
    env = os.environ if env is None else env
    p = spec.env_prefix
    warm = (env.get(f"{p}_WARMUP") or "default").strip() or "default"
    default_variant = spec.default_variant or getattr(spec.model_cls, "DEFAULT_VARIANT", None) or "default"
    return {
        "model_id": env.get("HF_MODEL") or spec.default_weights,
        "revision": env.get("TT_MODEL_WEIGHTS_REVISION") or env.get("TT_WEIGHTS_REVISION") or None,
        "weights_dir": env.get(f"{p}_WEIGHTS_DIR") or None,
        "mesh_shape": parse_mesh_shape(env.get("TT_MESH_SHAPE")),
        "device_id": int(env.get("TT_DEVICE_ID") or 0),
        "dispatch": (env.get(f"{p}_DISPATCH") or "eth").strip().lower(),
        "num_command_queues": int(env.get(f"{p}_NUM_CQS") or 0) or None,
        "variant": env.get(f"{p}_VARIANT") or default_variant,
        "warmup_variants": warm if warm in ("default", "none") else json.loads(warm),
        "max_body_bytes": int(float(env.get(f"{p}_MAX_BODY_MB") or 256) * (1 << 20)),
        "torch_threads": env.get(f"{p}_TORCH_THREADS") or None,
    }


def default_model_factory(spec: ServerSpec) -> Callable[[Dict[str, Any]], Any]:
    """``cfg -> spec.model_cls.from_pretrained(...)``."""

    def factory(cfg: Dict[str, Any]) -> Any:
        if spec.model_cls is None:
            raise RuntimeError(f"{spec.model_name}: ServerSpec.model_cls is not set and no model_factory was given")
        return spec.model_cls.from_pretrained(cfg["model_id"] or None, revision=cfg["revision"],
                                              weights_dir=cfg["weights_dir"], variant=cfg["variant"],
                                              device_id=cfg["device_id"], dispatch=cfg["dispatch"],
                                              num_command_queues=cfg["num_command_queues"],
                                              warmup_variants=cfg["warmup_variants"])

    return factory


@dataclass
class ServerState:
    """Per-app state (``app.state.ttaw``)."""

    spec: ServerSpec
    model_factory: Callable[[Dict[str, Any]], Any]
    ready: bool = False
    error: Optional[str] = None
    model: Any = None
    config: Dict[str, Any] = field(default_factory=dict)
    boot_ms: Optional[float] = None
    lock: Any = field(default_factory=threading.Lock)   # one chip, batch 1: device calls are serialized
    predict: Optional[Callable[[Any], Dict[str, Any]]] = None

    def status(self) -> str:
        return "ok" if self.ready else ("error" if self.error else "starting")

    def device_desc(self) -> Optional[Dict[str, Any]]:
        return getattr(self.model, "device_info", None) if self.model is not None else None


def _non_finite_path(obj: Any, path: str = "body") -> Optional[str]:
    """Location of the first NaN / infinite float in a JSON-able body (None if there is none)."""
    if isinstance(obj, float):
        return None if math.isfinite(obj) else path
    if isinstance(obj, Mapping):
        for k, v in obj.items():
            found = _non_finite_path(v, f"{path}.{k}")
            if found is not None:
                return found
    elif isinstance(obj, (list, tuple)):
        for i, v in enumerate(obj):
            found = _non_finite_path(v, f"{path}[{i}]")
            if found is not None:
                return found
    return None


def decode_request(req: PredictRequest, spec: ServerSpec, max_bytes: int) -> Dict[str, Any]:
    """JSON envelope -> the Python-API call kwargs (all client errors raise ``io.InputError``)."""
    cls = spec.model_cls
    fields = tuple(getattr(cls, "POINT_FIELDS", tio.DEFAULT_POINT_FIELDS))
    kw: Dict[str, Any] = {}
    if req.points is not None:
        kw["points"] = tio.decode_points(req.points.model_dump(exclude_none=True), max_bytes=max_bytes,
                                         default_fields=fields)
    if req.sweeps:
        kw["sweeps"] = [{"points": tio.decode_points(s.points.model_dump(exclude_none=True), max_bytes=max_bytes,
                                                     default_fields=fields),
                         "time_lag_s": s.time_lag_s,
                         "T_current_from_sweep": None if s.T_current_from_sweep is None else
                         tio.parse_transform(s.T_current_from_sweep, field="sweeps[].T_current_from_sweep")}
                        for s in req.sweeps]
    calib = tio.resolve_calibration(req.calibration, spec.calib_dir)
    if req.images:
        order = tuple(getattr(cls, "CAMERA_ORDER", ()) or ()) or None
        require = cls.requires_calibration() if hasattr(cls, "requires_calibration") else False
        kw["images"] = tio.decode_cameras([i.model_dump(exclude_none=True) for i in req.images], calib, order=order,
                                          require_calibration=require, max_bytes=max_bytes)
    if calib is not None:
        kw["calibration"] = calib
    if req.stream is not None:
        s = req.stream.model_dump(exclude_none=True)
        if "T_world_from_ego" in s:
            s["T_world_from_ego"] = tio.parse_transform(s["T_world_from_ego"], field="stream.T_world_from_ego")
        kw["stream"] = s
    if req.inputs is not None:
        kw["inputs"] = tio.decode_named_arrays(req.inputs.model_dump(exclude_none=True), max_bytes=max_bytes)
    if spec.decode_extra is not None:
        spec.decode_extra(req, kw, max_bytes)
    return kw


def create_app(spec: ServerSpec, *, model_factory: Optional[Callable[[Dict[str, Any]], Any]] = None) -> FastAPI:
    """The ASGI app of one bundle (see the module docstring)."""
    state = ServerState(spec, model_factory or default_model_factory(spec))
    RequestModel = spec.request_model

    @asynccontextmanager
    async def lifespan(_app: FastAPI):
        logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
        cfg = config_from_env(spec)
        if cfg["mesh_shape"] != (1, 1):
            raise RuntimeError(f"TT_MESH_SHAPE={cfg['mesh_shape']} is a multi-chip mesh; {spec.model_name} runs on "
                               "exactly one Blackhole chip (mesh_device P150, shape 1x1)")
        if cfg["torch_threads"]:
            import torch

            torch.set_num_threads(int(cfg["torch_threads"]))
        state.config, state.ready, state.error = cfg, False, None
        t0 = time.perf_counter()
        try:
            state.model = state.model_factory(cfg)  # logs: Loading weights / Opening device / Warming up / complete
        except Exception as e:  # startup failures exit uvicorn non-zero: no silent CPU fallback
            state.error = f"{type(e).__name__}: {e}"
            log.exception("startup failed")
            raise
        state.boot_ms = (time.perf_counter() - t0) * 1e3
        state.ready = True
        log.info("%s ready in %.1f s (%s)", spec.model_name, state.boot_ms / 1e3, state.device_desc())
        try:
            yield
        finally:
            state.ready = False
            model, state.model = state.model, None
            if model is not None:
                log.info("Closing the model and the device")
                with state.lock:
                    model.close()

    app = FastAPI(title=spec.model_name, version=spec.version, lifespan=lifespan,
                  description=spec.description or f"{spec.task} on one Tenstorrent Blackhole p150. "
                                                  "Not an OpenAI-compatible API.")
    app.state.ttaw = state

    @app.middleware("http")
    async def _body_limit(request: Request, call_next):
        limit = (state.config or {}).get("max_body_bytes")
        length = request.headers.get("content-length")
        if limit and length and length.isdigit() and int(length) > limit * 4 // 3 + 4096:
            return JSONResponse(status_code=400, content={"detail": f"request body of {length} bytes exceeds "
                                                                    f"{spec.env_prefix}_MAX_BODY_MB"})
        return await call_next(request)

    @app.get("/")
    def root() -> dict:
        return {"model": spec.model_name, "routes": ["/health", "/v1/health", "/info", "/v1/models", "/predict"],
                "docs": "/docs"}

    @app.get("/health")
    def health() -> dict:
        """Always 200: ``ok`` after warm-up, ``starting`` before, ``error`` if the boot failed."""
        return {"status": state.status(), "model": spec.model_name, "device": state.device_desc(),
                "error": state.error}

    @app.get("/v1/health")
    def v1_health() -> dict:  # tt-model's hint for non-chat packages points here
        return health()

    @app.get("/v1/models")
    def v1_models() -> dict:
        """OpenAI-shaped stub so tt-model's ready card / probes do not 404. This is NOT a chat API."""
        return {"object": "list", "data": [{"id": state.config.get("model_id") or spec.default_weights,
                                            "object": "model", "owned_by": spec.owner}]}

    @app.get("/info")
    def info() -> dict:
        m = state.model
        minfo = dict(getattr(m, "info", {}) or {}) if m is not None else {}
        calib = sorted(p.stem for p in Path(spec.calib_dir).glob("*.json")) if spec.calib_dir else []
        body = {
            "model": spec.model_name, "task": spec.task, "status": state.status(), "hardware": spec.hardware,
            "io": spec.io, "autoware": dict(spec.autoware),
            "weights": minfo.get("weights") or {"repo": state.config.get("model_id") or spec.default_weights},
            "device": state.device_desc(), "mesh_shape": os.environ.get("TT_MESH_SHAPE", "1x1"),
            "input": {"kind": minfo.get("input_kind"), "point_fields": minfo.get("point_fields"),
                      "camera_order": minfo.get("camera_order"), "batch_size": 1,
                      "point_formats": list(tio.POINT_FORMATS), "max_body_bytes": state.config.get("max_body_bytes"),
                      "calibration_presets": calib, "extra_inputs": minfo.get("extra_inputs"),
                      "input_schema": minfo.get("input_schema")},
            "output": {"labels": minfo.get("labels")},
            "variant": minfo.get("variant"), "warm_variants": minfo.get("warm_variants"),
            "runtime_params": minfo.get("runtime_params"), "compile_params": minfo.get("compile_params"),
            "warmup_ms": minfo.get("warmup_ms"), "boot_ms": state.boot_ms, "source": dict(spec.source),
        }
        if spec.info_extra is not None and m is not None:
            body.update(spec.info_extra(m))
        return body

    @app.post("/predict")
    def predict(req: RequestModel) -> dict:
        """One frame -> the model output's ``to_dict`` + ``timing_ms``. Sync handler: uvicorn runs it in a worker
        thread; the lock serializes the device."""
        if not state.ready:
            raise HTTPException(status_code=503, detail="model is still starting" if not state.error
                                else f"model failed to start: {state.error}")
        model = state.model
        t0 = time.perf_counter()
        try:
            kw = decode_request(req, spec, state.config.get("max_body_bytes") or (256 << 20))
            params = model.validate_params(req.params)
        except tio.InputError as e:
            raise HTTPException(status_code=400, detail=str(e)) from None
        t1 = time.perf_counter()
        try:
            with state.lock:
                out = model(**kw, **params)
        except tio.InputError as e:  # raised by the model's own pre-processing (shape / range checks)
            raise HTTPException(status_code=400, detail=str(e)) from None
        except Exception as e:  # noqa: BLE001 -- a device / runtime failure
            log.exception("inference failed")
            raise HTTPException(status_code=500, detail=f"inference failed: {type(e).__name__}: {e}") from None
        t2 = time.perf_counter()
        body = out.to_dict(req.output_format)
        bad = _non_finite_path(body)
        if bad is not None:  # strict JSON has no NaN / Infinity: say where, instead of a bare 500 from the encoder
            log.error("inference produced a non-finite value at %s", bad)
            raise HTTPException(status_code=500, detail=f"inference failed: non-finite value at {bad}")
        timing = dict(body.get("timing_ms") or {})
        timing.update({"decode": round((t1 - t0) * 1e3, 3), "model_call": round((t2 - t1) * 1e3, 3),
                       "total": round((time.perf_counter() - t0) * 1e3, 3)})
        body["timing_ms"] = timing
        return body

    state.predict = predict
    return app