app.py

#!/usr/bin/env python
"""Hugging Face Space entry point for the hallucination-signal correlation demo.

Hosts `interactive_corr_v2.html` — the live comparison of u_r / u_f / u_s and WAV
against the true rollout error — on ZeroGPU.

Why it is shaped like this
--------------------------
The demo is a stateful, keyboard-driven rollout: each frame depends on the
latents and env produced by the previous one. Three constraints decide the
architecture, and all three point the same way.

1. **ZeroGPU allocates the GPU per call, for a bounded duration.** So the frame
   loop lives inside a `@spaces.GPU(duration=...)` *generator*: one allocation
   streams many frames instead of paying the allocation cost per frame. Models
   are loaded once at module scope and stay resident.

2. **MuJoCo pins its EGL context to one thread.** `_run_blocking`'s docstring in
   `interactive_uncertainty.py` spells this out: anything that may touch the env
   — session creation, `env.reset` inside a render step, a task switch — has to
   happen on a single thread. So each session owns exactly one worker thread and
   *everything* env-touching happens there. Client messages that would create an
   env (task switches) are queued for that thread rather than applied inline.

3. **The page accumulates its statistics client-side** and clears them whenever
   the step counter goes backwards. So a lapsed GPU window must NOT reset the
   session — otherwise a viewer's correlation history is wiped every time the
   allocation renews, which is the entire point of the demo. The session state
   (small CPU-resident latent history plus a cached env) therefore outlives the
   GPU window: the generator returns *before* its deadline and the worker simply
   requests another, with `step` continuing to climb.

The browser client is unchanged: it speaks raw WebSocket with JSON status frames
and binary JPEG, which FastAPI serves just as aiohttp did.
"""
# `spaces` must be imported before torch so ZeroGPU can patch CUDA
# initialisation; module-level `.to("cuda")` is only legal because of that.
import spaces  # noqa: F401  (import order is load-bearing)

import argparse
import asyncio
import shlex
import json
import multiprocessing
import os
import queue
import sys
import threading
import time
import traceback
import uuid
from contextvars import copy_context
from dataclasses import dataclass, field, fields as dataclass_fields
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import gradio as gr
import torch
from fastapi import WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, PlainTextResponse
from gradio.context import LocalContext

APP_DIR = Path(__file__).resolve().parent
SRC_DIR = APP_DIR / "src"
sys.path.insert(0, str(SRC_DIR))
# The repo's modules assume cwd == src/ (flat imports; `--tasks_json ../tasks.json`).
os.chdir(SRC_DIR)

# Headless rendering. Set here rather than as Space variables so the defaults
# travel with the code — and set before any env module is imported, because
# MuJoCo picks its GL backend at import time and pygame at display init.
os.environ.setdefault("MUJOCO_GL", "egl")
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
os.environ.setdefault("SDL_VIDEODRIVER", "dummy")


def _env_int(name: str, default: int) -> int:
    try:
        return int(os.environ.get(name, default))
    except (TypeError, ValueError):
        return default


GPU_DURATION = _env_int("GPU_DURATION", 90)   # seconds per ZeroGPU allocation
GPU_MARGIN = _env_int("GPU_MARGIN", 8)        # leave the window cleanly, before it is revoked
IDLE_TIMEOUT = _env_int("IDLE_TIMEOUT", 120)  # drop a session this long after the socket goes away
VARIANT = os.environ.get("WM_VARIANT", "base")

# ---------------------------------------------------------------------------
# Weights, then models — all at module scope so ZeroGPU keeps them resident.
# ---------------------------------------------------------------------------
from huggingface_hub import snapshot_download  # noqa: E402

CKPT_DIR = Path(os.environ.get("CKPT_DIR", str(APP_DIR / "checkpoints")))
print(f"[startup] fetching '{VARIANT}' world-model weights -> {CKPT_DIR}", flush=True)
snapshot_download(
    repo_id="nicklashansen/mmbench2-models",
    repo_type="model",
    local_dir=str(CKPT_DIR),
    allow_patterns=[f"{VARIANT}/*"],
)
print("[startup] weights ready.", flush=True)

from interactive_four_corr import FourCorrServer, build_parser  # noqa: E402
from interactive_uncertainty import ACTION_KEYS, SessionState  # noqa: E402
import envs as _envs  # noqa: E402
from task_set import task_to_domain  # noqa: E402

# EXTRA_ARGS lets a launcher tune the sampler without editing this file --
# run_local.sh uses it for its speed presets. Parsed through the real parser so
# an unknown flag fails loudly instead of being silently ignored.
_extra = shlex.split(os.environ.get("EXTRA_ARGS", ""))
if _extra:
    print(f"[startup] extra args: {' '.join(_extra)}", flush=True)
_args = build_parser().parse_args(_extra)
_args.tokenizer_ckpt = str(CKPT_DIR / VARIANT / "tokenizer.pt")
_args.dynamics_ckpt = str(CKPT_DIR / VARIANT / "dynamics.pt")
_args.idm_ckpt = os.environ.get("IDM_CKPT", str(APP_DIR / "assets" / "idm_tokenizer.pt"))
_args.html = str(SRC_DIR / "interactive_corr_v2.html")
_args.tasks_json = str(APP_DIR / "tasks.json")
_args.host, _args.port = "0.0.0.0", _env_int("PORT", 7860)
_args.uncertainty_overlay = True
# WAV costs 2 rollouts x wav_horizon_long sampler calls each time it fires, which
# competes with the frame loop inside a bounded GPU window. But the panels stay
# empty until three samples exist, so too sparse and the demo looks broken on
# arrival: at 4 the third sample lands ~8 steps sooner than at 8.
_args.wav_every = _env_int("WAV_EVERY", 4)
_args.wav_horizon = _env_int("WAV_HORIZON", 4)
_args.wav_horizon_long = _env_int("WAV_HORIZON_LONG", 8)

print("[startup] building server (tokenizer + dynamics + IDM) ...", flush=True)
SERVER = FourCorrServer(_args)
print(f"[startup] server ready on device={SERVER.device}", flush=True)
if SERVER.device.type != "cuda":
    print("[startup] WARNING: models landed on CPU. On ZeroGPU this means CUDA was not "
          "patched in — check that `import spaces` precedes torch.", flush=True)

# Advertise only tasks whose domain actually imported. `envs/__init__.py` binds a
# raising stub for any domain whose dependency is missing, so without this filter
# a viewer could pick a task that can never start.
_ALL_TASKS = list(SERVER.tasks)
TASKS = [t for t in _ALL_TASKS if task_to_domain(t) in _envs.AVAILABLE_DOMAINS]
_DROPPED = sorted({task_to_domain(t) for t in _ALL_TASKS} - _envs.AVAILABLE_DOMAINS)
if _DROPPED:
    print(f"[startup] domains unavailable, their tasks are hidden: {_DROPPED}", flush=True)
    for mod, err in _envs.unavailable_domains().items():
        print(f"[startup]   {mod}: {type(err).__name__}: {err}", flush=True)
if not TASKS:
    raise RuntimeError("no task domain imported successfully — nothing to serve")
INITIAL_TASK = SERVER.initial_task if SERVER.initial_task in TASKS else TASKS[0]
print(f"[startup] serving {len(TASKS)}/{len(_ALL_TASKS)} tasks, initial={INITIAL_TASK}", flush=True)

# Bound the simulator cache: a public demo has viewers trying many tasks, and
# every one used to stay resident for the life of the process.
SERVER.env_cache_max = _env_int("ENV_CACHE_MAX", 4)
print(f"[startup] env cache capped at {SERVER.env_cache_max} simulators", flush=True)

# One GPU, so one render at a time across sessions. A threading.Lock (not the
# server's asyncio one) because rendering happens on worker threads.
_RENDER_LOCK = threading.Lock()


# ZeroGPU refuses to start a Space that declares no GPU function.
spaces.GPU(lambda: None)


# ---------------------------------------------------------------------------
# Sessions
#
# ZeroGPU runs every `@spaces.GPU` call in a *forked child process* and pickles
# the arguments across. That rules out handing the child a Session: it holds
# locks, threading events and queues, none of which pickle -- and even a
# hand-rolled picklable copy would be a snapshot, so nothing the parent wrote
# afterwards (a keypress, a task switch) would ever reach the running rollout.
#
# So the split is:
#   parent -> child   one multiprocessing.Queue per session, created *before*
#                     the first fork, therefore inherited rather than pickled.
#                     Every interaction is a message on it, including keys.
#   child  -> parent  whatever the generator yields; ZeroGPU pickles that back.
#                     Frames go this way, and so does the end-of-window state
#                     snapshot that lets the next window resume where this one
#                     stopped instead of restarting the episode.
# ---------------------------------------------------------------------------
# Match the context `spaces` itself uses (spaces/zero/wrappers.py) so the queue
# and the worker that reads it agree on how the child is started.
_MP = multiprocessing.get_context("fork")


@dataclass
class Session:
    sid: str
    # Parent-side only: the worker thread hands frames to the socket coroutine.
    frame_queue: "queue.Queue[Tuple[Optional[bytes], Dict[str, Any]]]" = field(
        default_factory=lambda: queue.Queue(maxsize=3))
    # Parent -> child. Inherited through the fork; never pickled as an argument.
    cmd_q: Any = field(default_factory=lambda: _MP.Queue(maxsize=256))
    stop_event: threading.Event = field(default_factory=threading.Event)
    ready: threading.Event = field(default_factory=threading.Event)
    # Picklable SessionState carried from one GPU window to the next. The live
    # state lives in the child; this is the only part of it the parent sees.
    snapshot: Optional[Dict[str, Any]] = None
    task: Optional[str] = None
    thread: Optional[threading.Thread] = None
    error: Optional[str] = None
    last_seen: float = field(default_factory=time.monotonic)

    def touch(self) -> None:
        self.last_seen = time.monotonic()

    def send(self, kind: str, payload: Any = None) -> None:
        """Queue one command for the rollout. Dropping is better than blocking
        the socket coroutine: a lost keyup self-corrects on the next event."""
        try:
            self.cmd_q.put_nowait((kind, payload))
        except queue.Full:
            pass


_sessions: Dict[str, Session] = {}
_sessions_lock = threading.Lock()


def _tasks_in_use() -> set:
    """Tasks a live session is stepping — the env cache must not evict these."""
    with _sessions_lock:
        return {s.task for s in _sessions.values() if s.task}


# Defined after the registry it reads, then handed to the server so eviction
# can never close a simulator a live session is still stepping.
SERVER.env_in_use = _tasks_in_use


def _get_session(sid: str) -> Optional[Session]:
    with _sessions_lock:
        return _sessions.get(sid)


def _drop_session(sid: str) -> Optional[Session]:
    with _sessions_lock:
        return _sessions.pop(sid, None)


def _reap_idle() -> None:
    """Free envs and threads for sockets that went away and never came back."""
    now = time.monotonic()
    with _sessions_lock:
        stale = [s for s in _sessions.values() if now - s.last_seen > IDLE_TIMEOUT]
        for s in stale:
            _sessions.pop(s.sid, None)
    for s in stale:
        s.stop_event.set()
        s.send("stop")
        print(f"[session {s.sid[:8]}] reaped after {IDLE_TIMEOUT}s idle", flush=True)


# ---------------------------------------------------------------------------
# GPU streaming
# ---------------------------------------------------------------------------
def _to_cpu(v: Any) -> Any:
    if isinstance(v, torch.Tensor):
        return v.detach().to("cpu")
    if isinstance(v, list):
        return [_to_cpu(x) for x in v]
    return v


def _to_device(v: Any, dev: Any) -> Any:
    if isinstance(v, torch.Tensor):
        return v.to(dev)
    if isinstance(v, list):
        return [_to_device(x, dev) for x in v]
    return v


# Rebuilt from the env on the far side; carrying raw frames across the fork
# would dwarf everything else in the snapshot.
_SNAPSHOT_SKIP = {"recorded_frames"}


def _snapshot(st: SessionState) -> Dict[str, Any]:
    """SessionState in a form ZeroGPU can pickle back to the parent.

    Tensors move to host memory first: a CUDA tensor pickles as an IPC handle
    into a process that is about to exit, leaving the parent holding a dangling
    reference to freed device memory.
    """
    return {f.name: _to_cpu(getattr(st, f.name))
            for f in dataclass_fields(st) if f.name not in _SNAPSHOT_SKIP}


def _restore(snap: Dict[str, Any]) -> SessionState:
    """Rehydrate in the child. The simulator is not part of this — the render
    path calls `_get_or_make_env(st.task)`, so it is rebuilt on first use."""
    return SessionState(**{k: _to_device(v, SERVER.device) for k, v in snap.items()})


def _apply_commands(sess: Session, st: SessionState) -> bool:
    """Drain queued client commands onto the live state, in the child.

    Everything the viewer does arrives here, keys included. Before ZeroGPU the
    socket coroutine wrote `st.keys_down` directly, which cannot work once the
    rollout lives in another process.

    Returns False when the parent has asked this window to end early.
    """
    keep = True
    while True:
        try:
            kind, payload = sess.cmd_q.get_nowait()
        except queue.Empty:
            return keep
        if kind == "keydown":
            if payload in ACTION_KEYS:
                st.keys_down.add(payload)
        elif kind == "keyup":
            st.keys_down.discard(payload)
        elif kind == "toggle_pause":
            st.paused = not st.paused
        elif kind == "reset":
            st.reset_requested = True
        elif kind == "set_task":
            if payload in TASKS and payload != st.task:
                try:
                    SERVER._switch_task_sync(st, payload)
                except Exception as e:
                    print(f"[session {sess.sid[:8]}] task switch to {payload!r} failed: {e}",
                          flush=True)
        elif kind == "stop":
            keep = False


@spaces.GPU(duration=GPU_DURATION)
def _gpu_stream(sid: str):
    """Stream frames for one GPU allocation, then return so another can be taken.

    Runs in a process ZeroGPU forks for this call, which is why the argument is
    the session id and not the Session (see the note above the dataclass). The
    child reaches the registry through the fork; everything travelling the other
    way is yielded.
    """
    sess = _sessions.get(sid)
    if sess is None:                    # reaped while we waited for the GPU
        return

    # Nothing in this process trains, and grad must be off for a second reason:
    # ZeroGPU reconstructs the packed model weights in the child as *inference*
    # tensors, and autograd refuses to save those for backward. With grad live
    # the first Linear in the reward head dies with
    #   RuntimeError: Inference tensors cannot be saved for backward
    # even though the call is pure inference. Off HF the weights are ordinary
    # tensors, so the same code merely builds a graph nobody uses -- which is
    # why this only ever showed up on the Space.
    with torch.no_grad():
        if sess.snapshot is None:
            st = SERVER.new_session()
            print(f"[session {sid[:8]}] opened task={st.task}", flush=True)
        else:
            st = _restore(sess.snapshot)
            # One line per allocation, so "did the episode survive the window?"
            # is answerable from the Space logs alone. A second `opened task=`
            # for the same session means it did not.
            print(f"[session {sid[:8]}] window resumed at step={st.step} task={st.task}",
                  flush=True)

        dt = 1.0 / max(1e-6, float(st.fps))
        next_t = time.monotonic()
        # Return under our own power, before ZeroGPU revokes the allocation: a
        # clean return keeps the session alive for the next window.
        deadline = time.monotonic() + max(1, GPU_DURATION - GPU_MARGIN)
        running = True

        while running and time.monotonic() < deadline:
            now = time.monotonic()
            if now < next_t:
                time.sleep(min(next_t - now, 0.05))
                continue
            next_t += dt
            if now - next_t > 1.0:  # resync after a stall rather than bursting
                next_t = now

            running = _apply_commands(sess, st)
            if not running:
                break
            t0 = time.monotonic()
            with _RENDER_LOCK:
                t_lock = time.monotonic()
                jpeg, status = SERVER._render_step_sync(st)
            now = time.monotonic()
            status["ms"] = round((now - t_lock) * 1000.0, 1)
            # Sessions no longer share this process, so the lock above only
            # orders this one. Contention for the actual device is ZeroGPU's to
            # schedule, and shows up as time waiting for the allocation.
            queue_ms = (t_lock - t0) * 1000.0
            n_sess = len(_sessions)
            status["queue_ms"] = round(queue_ms, 1)
            status["sessions"] = n_sess
            if n_sess > 1 or queue_ms > 50:
                status["text"] = status.get("text", "") + (
                    f" | {n_sess} session(s), queued {queue_ms:.0f} ms")
            yield "frame", (jpeg, status)

        # Hand the state back before the allocation lapses. Without this the
        # next window would call new_session() and restart the episode -- and
        # the page drops every accumulated correlation sample when the step
        # counter goes backwards, which the demo cannot afford.
        print(f"[session {sid[:8]}] window ended at step={st.step}", flush=True)
        yield "state", _snapshot(st)


def _worker(sess: Session, gradio_request) -> None:
    """Own one session: keep taking GPU windows until the client goes away."""
    if gradio_request is not None:
        # Carry the visitor's ZeroGPU token into this thread so their quota is
        # billed, not the Space owner's.
        try:
            LocalContext.request.set(gradio_request)
        except Exception:
            pass

    try:
        while not sess.stop_event.is_set():
            produced = False
            for kind, payload in _gpu_stream(sess.sid):
                if sess.stop_event.is_set():
                    break
                if kind == "frame":
                    produced = True
                    sess.ready.set()
                    try:
                        sess.frame_queue.put(payload, timeout=1.0)
                    except queue.Full:
                        pass    # reader is behind; drop the frame, keep stepping
                elif kind == "state":
                    sess.snapshot = payload
                    sess.task = payload.get("task")
            if not produced and sess.snapshot is None:
                # A window that yielded nothing and left no state to resume from
                # means the child died on arrival. Looping would spin against
                # whatever broke; fail the session so the page can say so.
                raise RuntimeError("GPU window produced no frames")
    except Exception as e:
        sess.error = f"{type(e).__name__}: {e}"
        sess.ready.set()        # unblock a connect that is still waiting
        print(f"[session {sess.sid[:8]}] worker died: {sess.error}\n{traceback.format_exc()}",
              flush=True)
    finally:
        step = (sess.snapshot or {}).get("step", 0)
        print(f"[session {sess.sid[:8]}] closed steps={step}", flush=True)


def _start_session(gradio_request) -> Session:
    _reap_idle()
    sess = Session(sid=uuid.uuid4().hex)
    with _sessions_lock:
        _sessions[sess.sid] = sess

    def _entry():
        _worker(sess, gradio_request)

    ctx = copy_context()
    sess.thread = threading.Thread(target=ctx.run, args=(_entry,), daemon=True,
                                   name=f"wm-{sess.sid[:8]}")
    sess.thread.start()
    return sess


# ---------------------------------------------------------------------------
# HTTP + WebSocket
# ---------------------------------------------------------------------------
app = gr.Server()


def _page_build() -> str:
    """Short hash of the page on disk — changes whenever the file does."""
    try:
        import hashlib
        return hashlib.sha1(Path(_args.html).read_bytes()).hexdigest()[:8]
    except OSError:
        return "unknown"


@app.get("/")
def index() -> HTMLResponse:
    # Re-read per request rather than serving the copy loaded at startup. The
    # page is 35 KB, so the read is free, and it means editing the HTML no
    # longer needs a server restart -- a stale in-memory copy sent us chasing
    # bugs that had already been fixed on disk.
    try:
        html = Path(_args.html).read_text()
    except OSError:
        html = SERVER.html
    html = html.replace("__TASK_SET__", json.dumps(TASKS))
    html = html.replace("__INITIAL_TASK__", INITIAL_TASK)
    # A visible build id ends the "is the browser showing me a cached page?"
    # question: compare what the header shows against /healthz.
    html = html.replace("__BUILD__", _page_build())
    return HTMLResponse(html, headers={"Cache-Control": "no-store"})


@app.get("/healthz")
def healthz() -> PlainTextResponse:
    with _sessions_lock:
        n = len(_sessions)
    return PlainTextResponse(
        f"ok sessions={n} tasks={len(TASKS)} device={SERVER.device} build={_page_build()}\n")


@app.websocket("/ws")
async def ws_endpoint(websocket: WebSocket) -> None:
    await websocket.accept()
    loop = asyncio.get_running_loop()
    sess = _start_session(LocalContext.request.get(None))

    # The first GPU window has to be granted before there is any state to talk
    # about; tell the client rather than leaving it staring at a blank frame.
    await websocket.send_text(json.dumps({"type": "status", "text": "waiting for GPU…"}))
    # `ready` is now set by the worker thread when the child's first frame comes
    # back, rather than by the rollout itself: the state it used to signal from
    # lives in another process.
    ready = await loop.run_in_executor(None, sess.ready.wait, 180.0)
    if not ready or sess.error:
        await websocket.send_text(json.dumps(
            {"type": "status", "text": f"failed to start: {sess.error or 'timeout'}"}))
        sess.stop_event.set()
        _drop_session(sess.sid)
        await websocket.close()
        return

    async def send_loop() -> None:
        while not sess.stop_event.is_set():
            try:
                jpeg, status = await loop.run_in_executor(
                    None, lambda: sess.frame_queue.get(timeout=0.5))
            except queue.Empty:
                continue
            sess.touch()
            await websocket.send_text(json.dumps(status))
            if jpeg is not None:
                await websocket.send_bytes(jpeg)

    async def recv_loop() -> None:
        # Every branch here queues rather than mutating state directly: the
        # rollout runs in the GPU child, so an attribute write on this side
        # would land on a copy nothing reads.
        last_reset = 0.0
        while not sess.stop_event.is_set():
            data = json.loads(await websocket.receive_text())
            if not isinstance(data, dict):
                continue
            sess.touch()
            t = str(data.get("type", ""))
            if t == "keydown":
                k = str(data.get("key", ""))
                if k == "Space":
                    sess.send("toggle_pause")
                elif k in ("r", "R"):
                    now = time.monotonic()
                    if now - last_reset >= 0.3:
                        last_reset = now
                        sess.send("reset")
                elif k in ("q", "Q", "Escape"):
                    return
                elif k in ACTION_KEYS:
                    sess.send("keydown", k)
            elif t == "keyup":
                sess.send("keyup", str(data.get("key", "")))
            elif t == "toggle_pause":
                sess.send("toggle_pause")
            elif t == "reset":
                now = time.monotonic()
                if now - last_reset >= 0.3:
                    last_reset = now
                    sess.send("reset")
            elif t == "set_task":
                sess.send("set_task", str(data.get("task", ""))[:128])
            elif t == "disconnect":
                return

    try:
        await asyncio.gather(send_loop(), recv_loop())
    except (WebSocketDisconnect, RuntimeError, json.JSONDecodeError):
        pass
    except Exception as e:
        print(f"[session {sess.sid[:8]}] ws error: {type(e).__name__}: {e}", flush=True)
    finally:
        sess.stop_event.set()       # stops the parent's worker loop
        sess.send("stop")           # ...and cuts the child's window short
        _drop_session(sess.sid)
        try:
            await websocket.close()
        except Exception:
            pass


if __name__ == "__main__":
    # ssr_mode must be off. Spaces sets GRADIO_SSR_MODE=True, which puts a Node
    # proxy in front of uvicorn ("Node proxy -> Python :7861" in the startup
    # log). It relays plain HTTP fine -- /healthz and the page both work -- but
    # it does not carry a WebSocket upgrade through: the request reaches our
    # route and a session opens, while the client never sees the 101 and gives
    # up. Locally GRADIO_SSR_MODE defaults to False, so the whole failure is
    # invisible off-Space. Client-side rendering costs us nothing: Gradio's own
    # UI is a stub here, the page is served by our own route.
    app.launch(server_name="0.0.0.0", server_port=_env_int("PORT", 7860),
               ssr_mode=False)