server.py

"""
Nori multi-policy cloud-inference server (INFERENCE_ENDPOINT_PLAN step 5).

ONE image, model ADAPTERS: `MODEL_KIND` selects which policy this container
serves — each Inference Endpoint runs exactly one kind (separate endpoints per
model = independent scaling + failure domains, plan §3.4).

    MODEL_KIND=molmoact2   today's predict_action path, UNCHANGED semantics
    MODEL_KIND=pi05        LeRobot policies.pi05 checkpoint (a Nori finetune)
    MODEL_KIND=groot       GR00T N1.7 (stub — lands with the finetune lane)

The /act JSON contract is FROZEN (laptop client in the field) and adapters
extend, never break:

    POST /act { images:[b64...], state:[floats], instruction:str, num_steps? }
      -> { actions:[[...DOF...] x horizon], compute_ms, meta? }

Per-policy chunk semantics (chunk_hz / horizon / joint + camera contract) are
NOT constants — pi05 must not silently inherit MolmoAct2's 30 Hz x 30. Each
adapter reports its `meta` (surfaced in /health and every /act response) and
the step-6 client wiring reads it via /load policy_kind.

Auth: X-Nori-Token primary (constant-time), Bearer kept for transition — same
as the step-3 single-model server. NO pickle on the wire (LeRobot PolicyServer
CVE class); plain JSON only.

Torch/model imports live INSIDE adapters' load() so this module imports light
(tests run the HTTP surface with MODEL_KIND=fake and no GPU stack installed).
"""

from __future__ import annotations

import base64
import io
import os
import secrets
import threading
import time
from typing import Any, Optional

import numpy as np
from fastapi import FastAPI, Header, HTTPException
from PIL import Image
from pydantic import BaseModel

MODEL_KIND = os.environ.get("MODEL_KIND", "molmoact2")
AUTH_TOKEN = os.environ.get("NORI_INFER_TOKEN")  # REQUIRED
# Inference Endpoints mount the endpoint's model repo here (no boot download);
# adapters fall back to their Hub default when absent (Space / bare GPU box).
MODEL_PATH = os.environ.get("MODEL_PATH", "/repository")

app = FastAPI(title=f"nori-serve-{MODEL_KIND}")
_adapter = None                       # set by the background load
_load_error: Optional[str] = None
_lock = threading.Lock()              # single GPU: serialize inference


def _make_adapter():
    """Resolve MODEL_KIND -> adapter instance. Import inside the function so an
    unused kind's heavy deps are never touched."""
    if MODEL_KIND == "molmoact2":
        from adapters.molmoact2 import MolmoAct2Adapter
        return MolmoAct2Adapter(MODEL_PATH)
    if MODEL_KIND == "pi05":
        from adapters.pi05 import Pi05Adapter
        return Pi05Adapter(MODEL_PATH)
    if MODEL_KIND == "groot":
        from adapters.groot import GrootAdapter
        return GrootAdapter(MODEL_PATH)
    if MODEL_KIND == "fake":
        from adapters.fake import FakeAdapter  # tests/CI only — no GPU deps
        return FakeAdapter(MODEL_PATH)
    raise RuntimeError(f"unknown MODEL_KIND {MODEL_KIND!r}")


def _load() -> None:
    """Background load (port must be up immediately; see step-3 notes)."""
    global _adapter, _load_error
    try:
        a = _make_adapter()
        a.load()
        _adapter = a
        print(f"[serve] {MODEL_KIND} ready: {a.meta()}", flush=True)
    except Exception as exc:  # surfaced via /health + /ready, not a dead port
        _load_error = f"{type(exc).__name__}: {exc}"
        print(f"[serve] LOAD FAILED — {_load_error}", flush=True)


@app.on_event("startup")
def _startup() -> None:
    if not AUTH_TOKEN:
        raise RuntimeError("NORI_INFER_TOKEN must be set")
    threading.Thread(target=_load, name=f"{MODEL_KIND}-load", daemon=True).start()


def _require_auth(x_nori_token: Optional[str], authorization: Optional[str]) -> None:
    """X-Nori-Token primary (a protected endpoint's edge consumes Authorization;
    custom headers pass through); Bearer kept for the transition client. Each
    compare constant-time; checked before any model work."""
    if x_nori_token and secrets.compare_digest(x_nori_token, AUTH_TOKEN):
        return
    if authorization and secrets.compare_digest(authorization, f"Bearer {AUTH_TOKEN}"):
        return
    raise HTTPException(status_code=401, detail="bad or missing auth token")


def _decode(b64: str) -> np.ndarray:
    if b64.lstrip().startswith("data:") and "," in b64[:64]:
        b64 = b64.split(",", 1)[1]
    img = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
    return np.asarray(img)


def _status() -> dict:
    ready = _adapter is not None
    return {
        "ok": ready,
        "status": "ready" if ready else ("error" if _load_error else "loading"),
        "error": _load_error,
        "kind": MODEL_KIND,
        "meta": _adapter.meta() if ready else None,
    }


@app.get("/")
def root() -> dict:
    # Spaces route external traffic only after a 2xx on "/" — keep it 200 always.
    return _status()


@app.get("/health")
def health() -> dict:
    return _status()


@app.get("/ready")
def ready() -> dict:
    """503-until-loaded — set as the Inference Endpoint's health_route."""
    if _adapter is None:
        detail = f"model load failed: {_load_error}" if _load_error else "model loading"
        raise HTTPException(status_code=503, detail=detail)
    return _status()


class ActRequest(BaseModel):
    images: list[str]
    state: list[float]
    instruction: str
    num_steps: int = 10
    # MolmoAct2-only knobs — other adapters ignore/report them (see extras).
    enable_cuda_graph: bool = True
    enable_grad: bool = False
    rtc: Optional[dict] = None


class ActResponse(BaseModel):
    actions: list[list[float]]
    compute_ms: Optional[float] = None
    rtc: Optional[dict] = None
    # Per-policy chunk semantics so the client NEVER assumes another model's
    # constants (chunk_hz/horizon/dof/cameras). Additive — old clients that
    # read only `actions` are unaffected.
    meta: Optional[dict] = None


@app.post("/act", response_model=ActResponse)
def act(req: ActRequest, authorization: Optional[str] = Header(None),
        x_nori_token: Optional[str] = Header(None)) -> ActResponse:
    _require_auth(x_nori_token, authorization)
    if _adapter is None:
        detail = f"model load failed: {_load_error}" if _load_error else "model not loaded yet"
        raise HTTPException(status_code=503, detail=detail)
    limits = _adapter.meta()
    max_images = int(limits.get("max_images", 6))
    if not 1 <= len(req.images) <= max_images:
        raise HTTPException(status_code=422, detail=f"need 1..{max_images} camera images")
    images = [_decode(b) for b in req.images]
    state = np.asarray(req.state, dtype=np.float32)
    t0 = time.perf_counter()
    with _lock:
        try:
            actions, extra = _adapter.act(
                images=images, state=state, instruction=req.instruction,
                num_steps=req.num_steps,
                extras={"enable_cuda_graph": req.enable_cuda_graph,
                        "enable_grad": req.enable_grad, "rtc": req.rtc},
            )
        except HTTPException:
            raise
        except MemoryError as e:
            raise HTTPException(status_code=507, detail=str(e))
        except Exception as e:
            # Adapters raise HTTPException for contract errors; anything else is
            # a real inference failure worth a 500 with the class name.
            raise HTTPException(status_code=500,
                                detail=f"inference failed: {type(e).__name__}: {e}")
    compute_ms = (time.perf_counter() - t0) * 1000.0
    acts = np.asarray(actions, dtype=np.float32)
    if acts.ndim == 3 and acts.shape[0] == 1:
        acts = acts[0]
    return ActResponse(actions=acts.tolist(), compute_ms=round(compute_ms, 1),
                       rtc=(extra or {}).get("rtc"), meta=_adapter.meta())


class PointRequest(BaseModel):
    image: str
    query: str = "the red cup"
    max_new_tokens: int = 96


@app.post("/point")
def point(req: PointRequest, authorization: Optional[str] = Header(None),
          x_nori_token: Optional[str] = Header(None)) -> dict:
    """Perception probe — only meaningful for adapters with a pointing-capable
    VLM backbone (MolmoAct2). Others 501 so the client can feature-detect."""
    _require_auth(x_nori_token, authorization)
    if _adapter is None:
        detail = f"model load failed: {_load_error}" if _load_error else "model not loaded yet"
        raise HTTPException(status_code=503, detail=detail)
    if not hasattr(_adapter, "point"):
        raise HTTPException(status_code=501,
                            detail=f"{MODEL_KIND} has no pointing backbone")
    with _lock:
        return _adapter.point(image=_decode(req.image), query=req.query,
                              max_new_tokens=req.max_new_tokens)