Cache giant train's setup stage in a sidecar file

Building the pdg/material vocab maps, the process map, and fitting the
Stage-1/Stage-2 normalizers all require scanning the training dataset
before a single epoch runs, which is wasted work whenever the same
data path is reused across runs (hyperparameter sweeps via `dwarf
hparam-scan`, repeated manual training attempts, ...). Persist those
setup-stage outputs to a JSON sidecar next to the input data
(giant/data/setup_cache.py), validated by a file fingerprint plus
fixed dimension constants and a manually-bumped format version before
reuse, with a soft warning (not a hard invalidation) on a git-hash
mismatch alone.

Also derives n_train_steps instantly from cached per-event row counts
instead of accumulating it during the normalizer scan, and always
collects the energy-router reservoir sample while the cache is being
populated (not only when the current run's router is energy-typed) so
a later run enabling --router-type energy never needs to rescan just
to seed expert centers.

New --cache-setup/--no-cache-setup (default on) and
--rebuild-setup-cache/--no-rebuild-setup-cache flags on `giant train`.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/data/setup_cache.py

@@ -0,0 +1,280 @@
+"""Sidecar cache for `giant train`'s setup stage (vocab maps, event-id split
+index, process maps, normalizer stats).
+
+The setup stage scans the full training dataset before a single epoch runs
+(see giant/pipeline.py:run_train_job); on multi-hundred-million-row datasets
+that scan is itself expensive, and it's pure waste to repeat when the same
+`data` path is reused across runs (hyperparameter sweeps via `dwarf
+hparam-scan`, repeated manual training attempts, ...). This module persists
+those scan outputs to a JSON file next to `data`, validated by a file
+fingerprint + fixed dimension constants + a manually-bumped format version
+before reuse — see `load`/`save`.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+from giant import config
+from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
+from giant.data.loader import load_event_ids
+from giant.data.transforms import Normalizer, sorted_membership
+
+# Bump manually on a change to the data-encoding semantics (e.g. a future
+# energy_simplex_encode bugfix) that doesn't also move one of _DIMS below —
+# a dims change already hard-invalidates on its own.
+_CACHE_FORMAT_VERSION = 1
+
+_DIMS = {
+    "COND_DIM": COND_DIM,
+    "X_DIM": X_DIM,
+    "K_MAX": K_MAX,
+    "PARTICLE_PHYS_DIM": PARTICLE_PHYS_DIM,
+    "SEC_SLOT_DIM": SEC_SLOT_DIM,
+}
+
+
+def sidecar_path(data: str | Path) -> Path:
+    """The cache sidecar for `data`, always a sibling of `data` itself.
+
+    A directory `data` gets a sidecar *next to* it (not inside), since the
+    directory may be a shared/read-only dataset mount, and other code globs
+    `*.parquet` directly inside it.
+    """
+    p = Path(data)
+    return p.parent / f"{p.name}.giant_train_cache.json"
+
+
+def fingerprint_files(files: list[Path]) -> list[list]:
+    """`[[resolved_path_str, size, mtime_ns], ...]`, in `files` order (not sorted).
+
+    Order must be preserved rather than normalized (e.g. sorted): file scan
+    order affects `build_process_map_from_files`'s tie-breaking (see
+    tests/test_loader.py), so the cached fingerprint has to reflect the same
+    order `find_parquet_files` produced.
+    """
+    out = []
+    for f in files:
+        resolved = Path(f).resolve()
+        st = resolved.stat()
+        out.append([str(resolved), st.st_size, st.st_mtime_ns])
+    return out
+
+
+def normalizer_key(val_fraction: float, seed: int, conditioning: str) -> str:
+    # .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing
+    # spurious cache misses between runs with the "same" val_fraction.
+    return f"valfrac={val_fraction:.6g}_seed={seed}_cond={conditioning}"
+
+
+@dataclass
+class NormalizerEntry:
+    cond_norm: Normalizer
+    tgt_norm: Normalizer
+    sec_phys_norm: Normalizer
+    n_train_steps: int
+    energy_reservoir_sample: np.ndarray
+
+    def to_json(self) -> dict:
+        return {
+            "cond_norm": self.cond_norm.to_dict(),
+            "tgt_norm": self.tgt_norm.to_dict(),
+            "sec_phys_norm": self.sec_phys_norm.to_dict(),
+            "n_train_steps": self.n_train_steps,
+            "energy_reservoir_sample": np.asarray(
+                self.energy_reservoir_sample, dtype=np.float32
+            ).tolist(),
+        }
+
+    @classmethod
+    def from_json(cls, d: dict) -> "NormalizerEntry":
+        return cls(
+            cond_norm=Normalizer.from_dict(d["cond_norm"]),
+            tgt_norm=Normalizer.from_dict(d["tgt_norm"]),
+            sec_phys_norm=Normalizer.from_dict(d["sec_phys_norm"]),
+            n_train_steps=int(d["n_train_steps"]),
+            energy_reservoir_sample=np.array(
+                d["energy_reservoir_sample"], dtype=np.float32
+            ),
+        )
+
+
+@dataclass
+class SetupCache:
+    fingerprint: list
+    git_hash: str = field(default_factory=config.git_hash)
+    vocab: tuple[dict[int, int], dict[str, int]] | None = None
+    event_index: tuple[np.ndarray, np.ndarray] | None = None
+    proc_maps: dict[int, dict[str, int]] = field(default_factory=dict)
+    normalizers: dict[str, NormalizerEntry] = field(default_factory=dict)
+
+    @classmethod
+    def empty(cls, files: list[Path]) -> "SetupCache":
+        return cls(fingerprint=fingerprint_files(files))
+
+    def to_json(self) -> dict:
+        d: dict = {
+            "format_version": _CACHE_FORMAT_VERSION,
+            "dims": dict(_DIMS),
+            "git_hash": self.git_hash,
+            "fingerprint": self.fingerprint,
+            "proc_maps": {str(k): v for k, v in self.proc_maps.items()},
+            "normalizers": {k: v.to_json() for k, v in self.normalizers.items()},
+        }
+        if self.vocab is not None:
+            pdg_map, mat_map = self.vocab
+            d["vocab"] = {
+                "pdg_map": {str(k): v for k, v in pdg_map.items()},
+                "mat_map": dict(mat_map),
+            }
+        if self.event_index is not None:
+            unique_ids, counts = self.event_index
+            d["event_index"] = {
+                "event_ids": np.asarray(unique_ids).tolist(),
+                "counts": np.asarray(counts).tolist(),
+            }
+        return d
+
+    @classmethod
+    def from_json(cls, d: dict) -> "SetupCache":
+        vocab = None
+        if "vocab" in d:
+            pdg_map = {int(k): v for k, v in d["vocab"]["pdg_map"].items()}
+            mat_map = dict(d["vocab"]["mat_map"])
+            vocab = (pdg_map, mat_map)
+        event_index = None
+        if "event_index" in d:
+            event_index = (
+                np.array(d["event_index"]["event_ids"], dtype=np.int64),
+                np.array(d["event_index"]["counts"], dtype=np.int64),
+            )
+        proc_maps = {int(k): v for k, v in d.get("proc_maps", {}).items()}
+        normalizers = {
+            k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()
+        }
+        return cls(
+            fingerprint=d["fingerprint"],
+            git_hash=d.get("git_hash", "unknown"),
+            vocab=vocab,
+            event_index=event_index,
+            proc_maps=proc_maps,
+            normalizers=normalizers,
+        )
+
+    def merge(self, other: "SetupCache") -> "SetupCache":
+        """Union of both caches; `other`'s populated fields win on a shared key.
+
+        Used by `save` to combine freshly-computed sections with whatever a
+        concurrent writer already persisted, so two runs against the same
+        dataset with different (e.g.) val_fraction don't clobber each
+        other's normalizer entries.
+        """
+        return SetupCache(
+            fingerprint=other.fingerprint,
+            git_hash=other.git_hash,
+            vocab=other.vocab if other.vocab is not None else self.vocab,
+            event_index=(
+                other.event_index if other.event_index is not None else self.event_index
+            ),
+            proc_maps={**self.proc_maps, **other.proc_maps},
+            normalizers={**self.normalizers, **other.normalizers},
+        )
+
+
+def load(
+    data: str | Path, files: list[Path], echo=lambda *a, **k: None
+) -> SetupCache | None:
+    """Load and validate the sidecar for `data`; `None` on any miss (never raises).
+
+    A missing file, corrupt JSON, format-version mismatch, dimension-constant
+    mismatch, or file-fingerprint mismatch are all clean misses. A git-hash
+    mismatch alone is a soft warning only (see
+    `config.warn_if_git_hash_mismatch`) — most commits to this repo don't
+    touch data-encoding semantics, so hard-invalidating on every one would
+    defeat the cache.
+    """
+    path = sidecar_path(data)
+    if not path.exists():
+        return None
+    try:
+        raw = json.loads(path.read_text())
+    except (json.JSONDecodeError, OSError) as exc:
+        echo(f"setup cache: {path} is corrupt ({exc}) — ignoring")
+        return None
+
+    try:
+        if raw.get("format_version") != _CACHE_FORMAT_VERSION:
+            echo("setup cache: format version changed — ignoring stale cache")
+            return None
+        if raw.get("dims") != _DIMS:
+            echo(
+                "setup cache: model dimension constants changed — ignoring stale cache"
+            )
+            return None
+        fp = fingerprint_files(files)
+        if raw.get("fingerprint") != fp:
+            echo("setup cache: input files changed — ignoring stale cache")
+            return None
+        cache = SetupCache.from_json(raw)
+    except (KeyError, TypeError, ValueError) as exc:
+        echo(f"setup cache: {path} is malformed ({exc}) — ignoring")
+        return None
+
+    config.warn_if_git_hash_mismatch({"meta": {"git_hash": cache.git_hash}}, path)
+    return cache
+
+
+def save(
+    data: str | Path,
+    files: list[Path],
+    sections: SetupCache,
+    echo=lambda *a, **k: None,
+) -> None:
+    """Merge `sections` into the on-disk sidecar and write it atomically.
+
+    Best-effort: any OSError (permission denied on a read-only mount, disk
+    full, ...) is caught, echoed as a warning, and swallowed — a failure to
+    cache must never fail training.
+    """
+    path = sidecar_path(data)
+    tmp = path.parent / f".{path.name}.tmp.{os.getpid()}"
+    try:
+        base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(files)
+        merged = base.merge(sections)
+        payload = json.dumps(merged.to_json(), separators=(",", ":"))
+        tmp.write_text(payload)
+        os.replace(tmp, path)
+    except OSError as exc:
+        echo(
+            f"setup cache: could not write {path} ({exc}) — continuing without caching"
+        )
+        try:
+            tmp.unlink(missing_ok=True)
+        except OSError:
+            pass
+
+
+def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.ndarray]:
+    """Unique event ids + per-event row (step) counts, across all `files`."""
+    if not files:
+        return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
+    all_ids = np.concatenate([load_event_ids(f) for f in files])
+    unique_ids, counts = np.unique(all_ids, return_counts=True)
+    return unique_ids, counts
+
+
+def n_train_steps_for_split(
+    unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray
+) -> int:
+    """Row (step) count summed over whichever `unique_ids` fall in `train_events_arr`.
+
+    `train_events_arr` must be ascending and duplicate-free (as produced by
+    `np.array(sorted(train_events))` in giant/pipeline.py).
+    """
+    mask = sorted_membership(unique_ids, train_events_arr)
+    return int(counts[mask].sum())