code/tt_diffusion_planner/tests/test_e2e_device.py

# SPDX-License-Identifier: Apache-2.0
"""End-to-end output agreement vs the CPU reference on the shipped and the public sample scenes (needs one p150).

    bin/devrun -t 1800 -- python -m pytest -q -s code/tt_diffusion_planner/tests/test_e2e_device.py

No dataset-level metric: the gate is agreement with the fp32 CPU reference of the same Autoware network on the same
inputs (PLAN.md 2.12, SPEC 6.3) on every scene with stored reference outputs: the shipped samples and research scenes
(``research/diffusion-planner/goldens/<scene>.npz``) and the nuScenes-derived public-data instants
(``research/diffusion-planner/goldens/public/<dataset>/<id>.npz``; inputs in ``public_data/inputs``, CC BY-NC-SA, local
only), all written by ``code/scripts/ref_golden.py``. Without the research directory (an installed package) the two
shipped samples are run through the CPU reference on the fly. Plus API == server bit identity.

- ego trajectory (post-processed x, y): max displacement <= 1.0 m and mean displacement <= 0.3 m per scene;
- turn-indicator command identical on every scene;
- neighbours (multi-modal futures, SPEC 10.2): the median over valid neighbours of the per-agent max displacement
  <= 1.5 m.

``GATES`` freeze like the PCC gates (``test_e2e_device.gates.json``); keep ``server/smoke_test.py``
``REFERENCE_GATES`` in line with them. ``DIFFUSION_PLANNER_E2E_REPORT=<path>`` writes the per-scene numbers as JSON.
"""
from __future__ import annotations

import base64
import json
import os
from pathlib import Path

import numpy as np
import pytest

from tt_diffusion_planner.reference import config as C
from tt_diffusion_planner.ttaw.golden import GateRegistry

pytestmark = pytest.mark.device

PKG = Path(__file__).resolve().parents[1]
RESEARCH = PKG.parents[3] / "research" / "diffusion-planner"
GOLDENS = RESEARCH / "goldens"
SAMPLE = PKG / "samples" / "kashiwanoha_dense.npz"
GATES = GateRegistry.for_test(__file__, {
    "ego.max_err_m": (1.0, "max", "max_abs_err"),
    "ego.mean_err_m": (0.3, "max", "ade"),
    "turn.command_agreement": (1.0, "min", "agreement"),
    "neighbors.median_max_err_m": (1.5, "max", "center_err"),
})


def _raw(path: Path, prefix: str = ""):
    with np.load(path, allow_pickle=False) as z:
        return {k: np.array(z[prefix + k]) for k in C.INPUT_NAMES}


def _expected(golden: Path):
    with np.load(golden, allow_pickle=False) as z:
        return z["out.trajectory"], z["out.predicted_agents"], int(z["out.turn_command"])


def _cases():
    """``[(name, raw inputs, (trajectory, predicted_agents, turn_command) or None)]``."""
    cases = []
    for p in sorted((PKG / "samples").glob("*.npz")):
        g = GOLDENS / f"{p.stem}.npz"
        cases.append((p.stem, _raw(p), _expected(g) if g.is_file() else None))
    for p in sorted((RESEARCH / "ort").glob("golden_*.npz")):
        g = GOLDENS / f"{p.stem[len('golden_'):]}.npz"
        if g.is_file():
            cases.append((g.stem, _raw(p, "raw/"), _expected(g)))
    for g in sorted(GOLDENS.glob("public/*/*.npz")):
        src = RESEARCH / "public_data" / "inputs" / g.parent.name / g.name
        if src.is_file():
            cases.append((f"{g.parent.name}/{g.stem}", _raw(src, "raw/"), _expected(g)))
    return cases


@pytest.fixture(scope="module")
def model(device):
    from tt_diffusion_planner import DiffusionPlanner
    from tt_diffusion_planner.reference.weights import find_weights_dir

    wd = find_weights_dir()        # $DIFFUSION_PLANNER_WEIGHTS_DIR > workspace assets > the HF cache (offline)
    if wd is None:
        pytest.skip("v5.0 weights not found (DIFFUSION_PLANNER_WEIGHTS_DIR, workspace assets or the HF cache)")
    with DiffusionPlanner.from_pretrained(device=device, weights_dir=str(wd)) as m:
        yield m


def test_scene_agreement(model):
    ego_max, ego_mean, turn_equal, nb_median = [], [], [], []
    report = {}
    reference = None
    for name, raw, want in _cases():
        if want is None:  # no stored goldens (installed package): the CPU reference of the shipped sample
            if reference is None:
                from tt_diffusion_planner.reference.pipeline import ReferencePlanner

                reference = ReferencePlanner(threads=4)
            out = reference(inputs=raw)
            want = (out.poses, out.predicted_agents, out.turn_indicator["command"])
        dev = model(inputs=raw)
        d = np.hypot(*(dev.poses[:, :2] - want[0][:, :2]).T)
        ego_max.append(float(d.max()))
        ego_mean.append(float(d.mean()))
        turn_equal.append(dev.turn_indicator["command"] == want[2])
        if want[1].shape[0]:
            per_agent = np.hypot(*(dev.predicted_agents[..., :2] - want[1][..., :2]).transpose(2, 0, 1))
            nb_median.append(float(np.median(per_agent.max(axis=1))))
        report[name] = {"ego_max_m": round(ego_max[-1], 4), "ego_mean_m": round(ego_mean[-1], 4),
                        "turn_equal": bool(turn_equal[-1]),
                        "nb_median_max_m": round(nb_median[-1], 4) if want[1].shape[0] else None,
                        "device_ms": round(float(dev.timing_ms.get("device", float("nan"))), 3)}
        print(f"{name}: ego max {ego_max[-1]:.3f} m mean {ego_mean[-1]:.3f} m turn {turn_equal[-1]} "
              f"nb median-max {nb_median[-1] if want[1].shape[0] else float('nan'):.3f} m")
    print(f"{len(ego_max)} scenes")
    out = os.environ.get("DIFFUSION_PLANNER_E2E_REPORT")
    if out:
        Path(out).parent.mkdir(parents=True, exist_ok=True)
        Path(out).write_text(json.dumps({"scenes": report, "info": model.info}, indent=1, default=str) + "\n")
    GATES.require("ego.max_err_m", max(ego_max))
    GATES.require("ego.mean_err_m", max(ego_mean))
    GATES.require("turn.command_agreement", float(np.mean(turn_equal)))
    GATES.require("neighbors.median_max_err_m", max(nb_median))


def test_api_equals_server(model, monkeypatch):
    """The server's ``/predict`` handler and ``model(...)`` give identical outputs (same decoders, same graph)."""
    pytest.importorskip("fastapi")
    from tt_diffusion_planner.server.app import PredictRequest, app

    state = app.state.ttaw
    monkeypatch.setattr(state, "model", model)
    monkeypatch.setattr(state, "ready", True)
    req = PredictRequest(inputs={"format": "npz", "data": base64.b64encode(SAMPLE.read_bytes()).decode()})
    body = state.predict(req)
    want = model(inputs=str(SAMPLE)).to_dict()
    assert body["trajectory"] == want["trajectory"] and body["turn_indicator"] == want["turn_indicator"]
    assert body["predicted_agents"] == want["predicted_agents"]