code/tt_diffusion_planner/server/app.py

# SPDX-License-Identifier: Apache-2.0
"""ASGI serving app of diffusion-planner-p150 on one Tenstorrent Blackhole p150: the HTTP contract of the Autoware
collection (vendored ``ttaw.server.app``, BUNDLE_CONVENTIONS.md section 7) bound to :class:`DiffusionPlanner`.

Served by tt-model-manager as ``kind: tt-dit-server``::

    python -m uvicorn --host 0.0.0.0 --port <p> --lifespan on tt_diffusion_planner.server.app:app

Routes: ``GET /``, ``/health`` and ``/v1/health`` (always 200: ``ok`` / ``starting`` / ``error``), ``/info``,
``/v1/models`` (stub), ``POST /predict``; errors 400 / 422 / 503 / 500 (SERVING.md section 3). Everything that
touches the device happens in the lifespan: weights -> device (ETH dispatch, 12x10) -> graph -> trace capture of
every warm-up variant, so uvicorn's ``Application startup complete`` (the line ``tt-model serve`` waits for) means
warm; SIGTERM (``tt-model stop``, 120 s) closes the model under the lock. ``/predict`` calls the Python API, so it
returns exactly what ``model(...)`` returns. Importing this module has no side effects beyond importing fastapi and
pydantic (the image's ``verify:`` imports it without a device); the environment is read in the lifespan only.

Model-specific request fields (e.g. PointPainting ``rois``): subclass ``PredictRequest``, pass ``request_model=`` and
``decode_extra=`` (which adds the decoded field to the call kwargs) to :class:`ServerSpec`, and list the keyword in
``EXTRA_INPUTS`` of the model class. Host tests swap the model before starting the app:
``app.state.ttaw.model_factory = Stub``; ``app.state.ttaw.predict(request)`` is the route handler itself.
"""
from pathlib import Path

from .. import __version__
from ..api import DiffusionPlanner
from ..ttaw.server.app import PredictRequest, ServerSpec, create_app, parse_mesh_shape

__all__ = ["SPEC", "app", "PredictRequest", "parse_mesh_shape"]

SPEC = ServerSpec(
    model_name=DiffusionPlanner.MODEL_NAME,
    env_prefix=DiffusionPlanner.ENV_PREFIX,
    model_cls=DiffusionPlanner,
    task="ego trajectory planning with a diffusion model (DPM-Solver++ 10 steps), neighbour prediction and a "
         "turn-indicator command",
    default_weights=DiffusionPlanner.DEFAULT_REPO,
    owner="changh95",
    io="the Autoware planner tensors (ego and neighbour histories, lanes, route, polygons, line strings, goal, ego "
       "shape, turn-indicator history) in, an 8 s ego trajectory, predicted paths of the valid neighbours and a "
       "turn-indicator command out",
    autoware={"package": "autoware_diffusion_planner",
              "path": "planning/autoware_diffusion_planner",
              "autoware_universe": "9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd"},
    source={"repo": "https://huggingface.co/changh95/diffusion-planner-p150", "license": "Apache-2.0"},
    calib_dir=Path(__file__).resolve().parents[1] / "calib",
    version=__version__,
    description="Diffusion Planner v5.0 (Autoware diffusion_planner) on one Tenstorrent Blackhole p150. "
                "Not an OpenAI-compatible API.",
)

app = create_app(SPEC)