src/fpgm/physics/priors.py

"""Orchestrate the VLM material worker from the ``fpgm`` env, with caching.

--- Where this sits ------------------------------------------------------------
:mod:`fpgm.physics.materials` turns a :class:`~fpgm.physics.types.MaterialVerdict`
into a :class:`~fpgm.physics.types.GaussianPrior`; it never touches a GPU or a
subprocess. This module is the other half: getting that verdict in the first
place, by shelling out to ``scripts/_vlm_material_worker.py`` in the isolated
``vla457`` env (Qwen3-VL, transformers) exactly the way
``scripts/settle_after_release.py`` shells out to
``scripts/_mujoco_settle_worker.py`` in the isolated ``mujoco`` env -- JSON
files across a subprocess boundary, never an in-process import, because
``transformers`` must not be installed into ``fpgm`` (see that module's own
docstring on why: it would risk a numpy/torch downgrade cascade that breaks
the working SAM3/TAPNext/pyrender stack there).

--- Why classification is cached per (scene_id, label), not per episode -------
All 62 local AUTOLab episodes this repo currently has share one scene
(``scene_id`` ``8756300955``) and the same physical objects appear across many
of that scene's episodes -- the mesh-reconstruction stages upstream already
key their own caches this way for exactly that reason. Asking a VLM "what
material is this brick" once per episode instead of once per (scene, object)
would be ~62x wasted GPU-VLM calls for an answer that cannot change between
episodes of the same scene: the object is the same object. The fingerprint
:meth:`VlmPriorProposer.propose` builds is therefore keyed on
``(scene_id, label)``, via :class:`~fpgm.datagen.cache.StageCache`'s existing
fingerprint-match convention -- a cache miss re-runs the worker, a hit reads
the previous verdict straight off disk with no subprocess at all.

--- Failure discipline: raise, never silently substitute -----------------------
:meth:`VlmPriorProposer.propose` raises :class:`~fpgm.physics.types.PhysicsError`
(with the worker's stderr tail attached) on any worker failure -- model load
OOM, a crashed CUDA context, a malformed JSON reply. It does **not** catch
that and quietly return :meth:`VlmPriorProposer.fallback_verdict` instead.
The reason is the same one ``fpgm.datagen.cache.StageCache``'s own docstring
gives for corrupt caches: a half-covered failure that resembles success is
worse than a loud one. If a caller's policy is "fall back to the unknown-
material prior when the VLM is unavailable," that is a legitimate policy --
but it must be a decision the caller makes explicitly (catch
``PhysicsError``, call ``fallback_verdict`` itself), not a behaviour buried
inside this class where a report could never tell "the VLM said unknown with
high confidence" apart from "the VLM never ran." ``fallback_verdict``'s own
``source="table:default"`` is exactly the marker that makes the two
distinguishable downstream (see its docstring).
"""

from __future__ import annotations

import json
import os
import subprocess
import tempfile
import time
from contextlib import contextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Any

from fpgm.physics.types import MaterialVerdict, PhysicsError
from fpgm.utils.logging import get_logger

if TYPE_CHECKING:
    from fpgm.datagen.cache import StageCache
    from fpgm.utils.timing import StepTimer

logger = get_logger(__name__)

__all__ = ["VlmPriorProposer"]

_REPO_ROOT = Path(__file__).resolve().parents[3]
_WORKER_SCRIPT = _REPO_ROOT / "scripts" / "_vlm_material_worker.py"
_DEFAULT_VLM_PYTHON = Path("/home/quang/miniconda3/envs/vla457/bin/python")
#: Overridable per the task spec; also lets a test stub in a fake interpreter
#: without monkeypatching module state.
_VLM_PYTHON_ENV_VAR = "FPGM_VLM_PYTHON"

#: Chars of subprocess stderr kept in a raised PhysicsError -- enough to see
#: a Python traceback's final frame/exception line without dragging an entire
#: CUDA/transformers import-warning wall of text into the caller's error.
_STDERR_TAIL_CHARS = 4000

_STAGE_NAME = "vlm_material"
_CACHE_KEY = "verdicts"  # single sub-key inside the stage's payload


def _vlm_python() -> Path:
    override = os.environ.get(_VLM_PYTHON_ENV_VAR)
    return Path(override) if override else _DEFAULT_VLM_PYTHON


class VlmPriorProposer:
    """Runs the VLM material worker (subprocess) and returns parsed verdicts.

    Args:
        vlm_python: Interpreter to invoke the worker with. Defaults to the
            ``vla457`` env, overridable via the ``FPGM_VLM_PYTHON`` env var
            (checked at call time, not construction time, so a test or a
            caller can flip it per-call by setting the env var).
        scene_id: Identifies the physical scene the crops were taken from,
            for the cache key (see module docstring on why caching is keyed
            on scene, not episode).
    """

    def __init__(self, *, vlm_python: Path | None = None, scene_id: str | int = "default") -> None:
        self._vlm_python_override = vlm_python
        self.scene_id = str(scene_id)

    def _resolve_vlm_python(self) -> Path:
        return self._vlm_python_override or _vlm_python()

    # -- public API ---------------------------------------------------------- #

    def propose(
        self,
        crops: dict[str, Path],
        *,
        cache: StageCache | None = None,
        timer: StepTimer | None = None,
    ) -> dict[str, MaterialVerdict]:
        """Classify each ``label -> crop_path`` entry, one worker call for all misses.

        Args:
            crops: ``{label: path_to_rgba_png}`` -- typically the output of
                ``fpgm.objects.crop.save_debug`` for each object in a scene.
            cache: If given, a hit for ``(scene_id, label)`` is read straight
                from disk and never sent to the worker; every miss across
                ``crops`` is batched into a *single* worker invocation (model
                load dominates -- see the worker's own docstring), and the
                results are written back to the cache one-by-one so a later
                call with a different, overlapping ``crops`` dict only pays
                for its own new misses.
            timer: Optional :class:`~fpgm.utils.timing.StepTimer`; wraps the
                subprocess call as ``propose`` and folds in the worker's own
                self-reported model-load/inference split via ``timer.mark``.

        Returns:
            ``{label: MaterialVerdict}``, one entry per key in ``crops``.

        Raises:
            PhysicsError: Any crop path does not exist, or the worker
                subprocess fails or returns unparseable output. Never
                silently substituted with :meth:`fallback_verdict` -- see
                module docstring.
        """
        if not crops:
            return {}
        for label, path in crops.items():
            if not Path(path).exists():
                raise PhysicsError(f"VlmPriorProposer: crop for {label!r} does not exist: {path}")

        cm = _null_timer_step if timer is None else timer.step

        with cm("propose", n=len(crops)):
            results: dict[str, MaterialVerdict] = {}
            misses: dict[str, Path] = {}

            if cache is not None:
                for label, path in crops.items():
                    cached = self._read_cached(cache, label)
                    if cached is not None:
                        results[label] = cached
                    else:
                        misses[label] = path
            else:
                misses = dict(crops)

            if misses:
                verdicts, worker_timing = self._run_worker(misses, timer=timer)
                for label, verdict in verdicts.items():
                    results[label] = verdict
                    if cache is not None:
                        self._write_cached(cache, label, verdict)
                if timer is not None:
                    ml = worker_timing.get("model_load_seconds")
                    if ml is not None:
                        timer.mark("propose/vlm_model_load", float(ml))
                    ti = worker_timing.get("total_inference_seconds")
                    if ti is not None:
                        timer.mark(
                            "propose/vlm_inference", float(ti), n=len(misses)
                        )

        missing = set(crops) - set(results)
        if missing:
            raise PhysicsError(
                f"VlmPriorProposer: worker did not return verdicts for {sorted(missing)}"
            )
        return {label: results[label] for label in crops}

    @staticmethod
    def fallback_verdict(label: str) -> MaterialVerdict:
        """``unknown`` at probability 1, for when the VLM is unavailable.

        Explicit-fallback marker: ``source="table:default"`` is the one
        field a downstream report can grep for to tell "the VLM never ran
        for this object, ``fpgm.physics.materials``' widest, most
        uninformative prior was used" apart from a genuine VLM verdict that
        happened to land on ``unknown`` (which would carry
        ``source="vlm:Qwen3-VL-2B-Instruct"`` instead, from the worker's own
        stamp) -- see module docstring on why this class never calls this
        method on a worker failure itself.
        """
        return MaterialVerdict(
            label=label,
            classes=(("unknown", 1.0),),
            source="table:default",
            raw={"reason": "vlm_unavailable"},
        )

    # -- cache plumbing -------------------------------------------------------- #

    def _cache_stage(self, label: str) -> str:
        # One StageCache "stage name" per label so a hit/miss on one object
        # never touches another's meta.json -- matches the granularity a
        # caller actually wants to invalidate at (re-classify one object,
        # not the whole scene).
        return f"{_STAGE_NAME}/{self.scene_id}/{label}"

    def _fingerprint(self, label: str) -> dict[str, str]:
        return {"scene_id": self.scene_id, "label": label, "code_version": "s9.materials.v1"}

    def _read_cached(self, cache: StageCache, label: str) -> MaterialVerdict | None:
        stage = self._cache_stage(label)
        fp = self._fingerprint(label)
        if not cache.is_fresh(stage, fp):
            return None
        meta = cache.read_meta(stage)
        if meta is None:
            return None
        payload = meta.get("payload", {})
        verdict_dict = payload.get(_CACHE_KEY)
        if verdict_dict is None:
            return None
        try:
            return MaterialVerdict.from_dict(verdict_dict)
        except (KeyError, PhysicsError) as exc:
            logger.warning(
                "VlmPriorProposer: cached verdict for %r is malformed (%s); treating as miss",
                label,
                exc,
            )
            return None

    def _write_cached(self, cache: StageCache, label: str, verdict: MaterialVerdict) -> None:
        stage = self._cache_stage(label)
        cache.write_meta(
            stage,
            fingerprint=self._fingerprint(label),
            payload={_CACHE_KEY: verdict.as_dict()},
            limitations=[],
        )

    # -- worker subprocess ------------------------------------------------------ #

    def _run_worker(
        self, misses: dict[str, Path], *, timer: StepTimer | None = None
    ) -> tuple[dict[str, MaterialVerdict], dict]:
        vlm_python = self._resolve_vlm_python()
        labels = list(misses.keys())
        images = [misses[label] for label in labels]

        with tempfile.TemporaryDirectory(prefix="fpgm_vlm_worker_") as tmp:
            out_json = Path(tmp) / "verdicts_out.json"
            cmd = [
                str(vlm_python),
                str(_WORKER_SCRIPT),
                "--images",
                *[str(p) for p in images],
                "--labels",
                *labels,
                "--out",
                str(out_json),
            ]
            logger.info("VlmPriorProposer: invoking vlm worker for %d object(s)", len(labels))
            t0 = time.perf_counter()
            proc = subprocess.run(cmd, capture_output=True, text=True)
            wall = time.perf_counter() - t0
            logger.info(
                "VlmPriorProposer: worker finished in %.2fs, returncode=%d", wall, proc.returncode
            )
            if proc.returncode != 0:
                tail = (proc.stderr or "")[-_STDERR_TAIL_CHARS:]
                raise PhysicsError(
                    f"vlm material worker failed (code {proc.returncode}) for labels "
                    f"{labels}: ...{tail}"
                )
            if not out_json.exists():
                tail = (proc.stderr or "")[-_STDERR_TAIL_CHARS:]
                raise PhysicsError(
                    f"vlm material worker exited 0 but wrote no output ({out_json}); "
                    f"stderr tail: ...{tail}"
                )
            return self._parse_worker_output(out_json.read_text(), labels)

    @staticmethod
    def _parse_worker_output(
        raw_text: str, expected_labels: list[str]
    ) -> tuple[dict[str, MaterialVerdict], dict]:
        """Shared by :meth:`_run_worker` and tests that stub the subprocess.

        Raises :class:`PhysicsError` on anything short of a fully well-formed
        list of verdict dicts -- a malformed blob must never yield a
        half-built verdict (e.g. skipping the entries that fail to parse and
        returning the rest), because a caller that only checks "did I get a
        MaterialVerdict for every label" would not notice.
        """
        try:
            parsed = json.loads(raw_text)
        except json.JSONDecodeError as exc:
            raise PhysicsError(f"vlm material worker: output is not valid JSON: {exc}") from exc
        if not isinstance(parsed, list):
            raise PhysicsError(
                "vlm material worker: expected a JSON list of verdicts, got "
                f"{type(parsed).__name__}"
            )

        verdicts: dict[str, MaterialVerdict] = {}
        for i, item in enumerate(parsed):
            if not isinstance(item, dict):
                raise PhysicsError(
                    f"vlm material worker: verdict #{i} is not a JSON object: {item!r}"
                )
            try:
                verdict = MaterialVerdict.from_dict(item)
            except (KeyError, TypeError, ValueError, PhysicsError) as exc:
                raise PhysicsError(
                    f"vlm material worker: verdict #{i} is malformed: {exc}"
                ) from exc
            verdicts[verdict.label] = verdict

        missing = set(expected_labels) - set(verdicts)
        if missing:
            raise PhysicsError(
                f"vlm material worker: output missing verdicts for {sorted(missing)}"
            )
        return verdicts, {}


# --------------------------------------------------------------------------- #
# No-op context manager for timer=None, matching StepTimer.step's signature
# closely enough for this module's single use (label + n=).
# --------------------------------------------------------------------------- #


@contextmanager
def _null_timer_step(label: str, *, n: int | None = None, **fields: Any):
    yield None