src/fpgm/datagen/cache.py

"""Per-stage on-disk artifact cache: fingerprint + skip-if-done + resume.

Follows the repo's existing ``meta.json`` sidecar convention (see
``scripts/run_object_pipeline.py``'s per-stage ``meta.json`` and ``outputs/*/meta.json``'s
``limitations`` block) rather than inventing a new one: a stage directory holds
whatever artifacts it wants plus one ``meta.json`` recording the fingerprint of the
inputs that produced them, an arbitrary small ``payload`` (paths/stats worth
keeping without re-reading the artifacts), and a ``limitations`` list carried
straight through to the final export's own meta.json.

The one rule that matters more than the format: **a half-written artifact that
looks fresh is worse than no cache at all.** A resumed run that trusts a
crashed-mid-write ``meta.json`` will build every later stage on top of garbage
without ever knowing it. So every write goes through
:func:`fpgm.utils.io.atomic_write` (temp file + ``os.replace``), and every read
treats a missing *or* corrupt ``meta.json`` as "not fresh" rather than raising --
a crash must make the stage re-run, never silently reuse whatever bytes happened
to land on disk.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np

from fpgm.utils.io import atomic_write, ensure_dir
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

_META_FILENAME = "meta.json"


def _jsonable(value: Any) -> Any:
    """Recursively coerce ``value`` into something ``json.dumps`` accepts.

    Fingerprints are built out of whatever a caller had lying around -- numpy
    scalars from a config, ``Path`` objects, tuples -- and none of those survive
    a JSON round trip unchanged unless normalised first. Doing that normalisation
    once, here, on both the write and the compare side is what keeps
    :meth:`StageCache.is_fresh` from ever comparing "the same fingerprint" against
    itself and getting ``False`` back because a tuple turned into a list on one
    side but not the other.
    """
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, np.ndarray):
        return [_jsonable(v) for v in value.tolist()]
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, dict):
        return {str(k): _jsonable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_jsonable(v) for v in value]
    return value


class StageCache:
    """Artifact directory per stage, each with a fingerprinted ``meta.json``.

    Usage::

        cache = StageCache(outputs_root / episode_uuid)
        fp = {"config_hash": ..., "input_mtime": ..., "code_version": "s3.v1"}
        if not cache.is_fresh("dense_depth", fp):
            payload = run_dense_depth_stage(...)
            cache.write_meta("dense_depth", fp, payload, limitations=[...])
        meta = cache.read_meta("dense_depth")
    """

    def __init__(self, root: Path):
        self.root = ensure_dir(Path(root))

    def stage_dir(self, name: str) -> Path:
        """The directory a stage's artifacts (and its ``meta.json``) live in."""
        return ensure_dir(self.root / name)

    def _meta_path(self, name: str) -> Path:
        return self.stage_dir(name) / _META_FILENAME

    def is_fresh(self, name: str, fingerprint: dict[str, Any]) -> bool:
        """``True`` iff stage ``name`` has a valid ``meta.json`` matching ``fingerprint``.

        A fingerprint should capture everything the stage's output depends on:
        input file identities (path + mtime or hash), the relevant config values,
        and a code-version tag for the stage's own logic -- ``StageCache`` treats
        the fingerprint as an opaque dict and does not try to guess what belongs
        in it. Any mismatch (including keys present on only one side) counts as
        stale, by ordinary dict equality after normalisation.
        """
        meta = self.read_meta(name)
        if meta is None:
            return False
        return meta.get("fingerprint") == _jsonable(fingerprint)

    def read_meta(self, name: str) -> dict[str, Any] | None:
        """Parsed ``meta.json`` for stage ``name``, or ``None`` if missing/corrupt.

        Corruption (truncated JSON from a crash mid-write, before ``atomic_write``
        existed or if it was bypassed, or simple disk damage) is deliberately not
        an exception here: the whole point of caching is that a caller can do
        ``if not cache.is_fresh(...): recompute()`` without a try/except around
        every call, and "not fresh" already means "recompute", which is exactly
        the right response to corruption too.
        """
        path = self._meta_path(name)
        if not path.exists():
            return None
        try:
            raw = path.read_text()
        except OSError as exc:
            logger.warning("stage %r: could not read %s: %s", name, path, exc)
            return None
        try:
            meta = json.loads(raw)
        except json.JSONDecodeError as exc:
            logger.warning(
                "stage %r: %s is corrupt (%s); treating as not-fresh, not raising",
                name,
                path,
                exc,
            )
            return None
        if not isinstance(meta, dict):
            logger.warning("stage %r: %s did not contain a JSON object", name, path)
            return None
        return meta

    def write_meta(
        self,
        name: str,
        fingerprint: dict[str, Any],
        payload: dict[str, Any],
        limitations: list[str],
    ) -> Path:
        """Atomically write stage ``name``'s ``meta.json``.

        Writes through :func:`fpgm.utils.io.atomic_write`: the new content is
        written to a sibling temp file and ``os.replace``-d into place, so a
        process killed mid-write leaves either the previous ``meta.json``
        (correctly still "fresh" for the previous fingerprint, or stale and
        correctly re-run) or nothing at all -- never a truncated file that
        :meth:`read_meta` would have to guess about.
        """
        meta = {
            "fingerprint": _jsonable(fingerprint),
            "payload": _jsonable(payload),
            "limitations": [str(item) for item in limitations],
        }
        data = json.dumps(meta, indent=2, sort_keys=True).encode("utf-8")
        path = self._meta_path(name)
        atomic_write(path, data)
        return path