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/cli.py

@@ -385,6 +385,26 @@ def train(
             "--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"
         ),
     ] = 65536,
+    cache_setup: Annotated[
+        bool,
+        typer.Option(
+            "--cache-setup/--no-cache-setup",
+            help="Cache the training setup stage's expensive per-file "
+            "precomputation (vocab maps, event split index, normalizer stats) "
+            "in a JSON sidecar next to the data, so a repeat `giant train` "
+            "against the same dataset (e.g. a hyperparameter sweep) can skip "
+            "re-deriving it",
+        ),
+    ] = True,
+    rebuild_setup_cache: Annotated[
+        bool,
+        typer.Option(
+            "--rebuild-setup-cache/--no-rebuild-setup-cache",
+            help="Ignore any existing setup cache sidecar and recompute every "
+            "section fresh for this run (still writes the refreshed sections "
+            "back to the sidecar for later runs; no effect if --no-cache-setup)",
+        ),
+    ] = False,
     out: Annotated[
         Optional[Path],
         typer.Option(
@@ -550,6 +570,8 @@ def train(
         shuffle_buffer=shuffle_buffer,
         num_workers=t["num_workers"],
         resume=resume,
+        cache_setup=cache_setup,
+        rebuild_setup_cache=rebuild_setup_cache,
         echo=typer.echo,
     )
 

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())

giant/pipeline.py

@@ -13,9 +13,9 @@ from giant.constants import (
     SEC_SLOT_DIM,
     X_DIM,
 )
+from giant.data import setup_cache
 from giant.data.loader import (
     find_parquet_files,
-    load_event_ids,
     iter_file_chunks,
     build_index_maps_from_files,
     build_process_map_from_files,
@@ -39,6 +39,8 @@ def run_train_job(
     shuffle_buffer: int,
     num_workers: int,
     resume: Path | None = None,
+    cache_setup: bool = True,
+    rebuild_setup_cache: bool = False,
     echo=print,
 ) -> None:
     t, m = cfg["train"], cfg["model"]
@@ -49,21 +51,47 @@ def run_train_job(
     files = find_parquet_files(data)
     echo(f"found {len(files)} parquet file(s)")
 
-    echo("scanning event IDs …")
-    all_event_ids = np.concatenate([load_event_ids(f) for f in files])
+    cache: setup_cache.SetupCache | None = None
+    if cache_setup:
+        if rebuild_setup_cache:
+            echo("setup cache: --rebuild-setup-cache given, recomputing all sections")
+            loaded = None
+        else:
+            loaded = setup_cache.load(data, files, echo=echo)
+        cache = loaded if loaded is not None else setup_cache.SetupCache.empty(files)
+
+    if cache is not None and cache.event_index is not None:
+        unique_ids, counts = cache.event_index
+        echo(f"event index: cache hit ({len(unique_ids):,} unique events)")
+    else:
+        echo("scanning event IDs …")
+        unique_ids, counts = setup_cache.compute_event_index_from_files(files)
+        if cache is not None:
+            cache.event_index = (unique_ids, counts)
+
     train_events, val_events = make_event_split(
-        all_event_ids, val_fraction=t["val_fraction"]
+        unique_ids, val_fraction=t["val_fraction"]
     )
     events_arr = np.array(sorted(train_events))
+    n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
     echo(
-        f"  {len(all_event_ids):,} steps  |  "
+        f"  {int(counts.sum()):,} steps  |  "
         f"{len(train_events)} train events  |  "
         f"{len(val_events)} val events"
     )
 
-    echo("building vocabulary maps …")
-    pdg_map, mat_map = build_index_maps_from_files(files)
-    echo(f"  {len(pdg_map)} PDG codes  |  {len(mat_map)} materials")
+    if cache is not None and cache.vocab is not None:
+        pdg_map, mat_map = cache.vocab
+        echo(
+            f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, "
+            f"{len(mat_map)} materials)"
+        )
+    else:
+        echo("building vocabulary maps …")
+        pdg_map, mat_map = build_index_maps_from_files(files)
+        echo(f"  {len(pdg_map)} PDG codes  |  {len(mat_map)} materials")
+        if cache is not None:
+            cache.vocab = (pdg_map, mat_map)
 
     router_cfg = m["router"]
     if t["mode"] == "wgan" and router_cfg.get("enabled"):
@@ -73,66 +101,96 @@ def run_train_job(
         )
     proc_map: dict[str, int] | None = None
     if router_cfg.get("enabled") and router_cfg.get("type") == "process":
-        echo("building process vocabulary …")
-        proc_map = build_process_map_from_files(
-            files, n_experts=router_cfg["n_experts"]
-        )
-        echo(
-            f"  {len(proc_map)} process labels mapped to {router_cfg['n_experts']} experts"
-        )
+        n_experts = router_cfg["n_experts"]
+        cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
+        if cached_proc_map is not None:
+            proc_map = cached_proc_map
+            echo(
+                f"process vocabulary: cache hit ({len(proc_map)} labels, "
+                f"{n_experts} experts)"
+            )
+        else:
+            echo("building process vocabulary …")
+            proc_map = build_process_map_from_files(files, n_experts=n_experts)
+            echo(f"  {len(proc_map)} process labels mapped to {n_experts} experts")
+            if cache is not None:
+                cache.proc_maps[n_experts] = proc_map
 
-    echo("fitting normalizer (streaming) …")
     conditioning = m["conditioning"]
-    cond_acc = _WelfordAccumulator(COND_DIM)
-    tgt_acc = _WelfordAccumulator(X_DIM)
-    sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
-    # EnergyRouter's default center spread (linspace over [-2, 2]) assumes
-    # the z-normalized energy column is roughly uniform, which real energy
-    # spectra rarely are — collect a reservoir sample here (reusing this
-    # same pass, not a second scan) so centers can instead be seeded from
-    # actual data quantiles below.
     energy_router_active = (
         router_cfg.get("enabled") and router_cfg.get("type") == "energy"
     )
     energy_idx = router_cfg.get("energy_idx", 3)
-    energy_sampler = (
-        _ReservoirSampler(capacity=100_000) if energy_router_active else None
-    )
-    n_train_steps = 0
-    for path in files:
-        for chunk in iter_file_chunks(path):
-            mask = sorted_membership(chunk["event_id"], events_arr)
-            if not mask.any():
-                continue
-            n_train_steps += int(mask.sum())
-            chunk_tr = {k: v[mask] for k, v in chunk.items()}
-            cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features(
-                chunk_tr,
-                pdg_map,
-                mat_map,
-                proc_map=proc_map,
-                require_secondaries=True,
-                conditioning=conditioning,
-                sec_phys_only=True,
+    norm_key = setup_cache.normalizer_key(t["val_fraction"], t["seed"], conditioning)
+    entry = cache.normalizers.get(norm_key) if cache is not None else None
+
+    if entry is not None:
+        echo(f"normalizer: cache hit (key={norm_key!r})")
+        cond_norm = entry.cond_norm
+        tgt_norm = entry.tgt_norm
+        sec_phys_norm = entry.sec_phys_norm
+        energy_sample = entry.energy_reservoir_sample
+    else:
+        echo("fitting normalizer (streaming) …")
+        cond_acc = _WelfordAccumulator(COND_DIM)
+        tgt_acc = _WelfordAccumulator(X_DIM)
+        sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
+        # EnergyRouter's default center spread (linspace over [-2, 2]) assumes
+        # the z-normalized energy column is roughly uniform, which real energy
+        # spectra rarely are — collect a reservoir sample here (reusing this
+        # same pass, not a second scan) so centers can instead be seeded from
+        # actual data quantiles below. Collected whenever the setup cache is
+        # being populated, not only when *this* run's router is
+        # energy-typed, so a later run enabling --router-type energy against
+        # this same (val_fraction, seed, conditioning) key never needs to
+        # rescan just to seed centers.
+        collect_energy_sample = energy_router_active or cache is not None
+        energy_sampler = (
+            _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
+        )
+        for path in files:
+            for chunk in iter_file_chunks(path):
+                mask = sorted_membership(chunk["event_id"], events_arr)
+                if not mask.any():
+                    continue
+                chunk_tr = {k: v[mask] for k, v in chunk.items()}
+                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features(
+                    chunk_tr,
+                    pdg_map,
+                    mat_map,
+                    proc_map=proc_map,
+                    require_secondaries=True,
+                    conditioning=conditioning,
+                    sec_phys_only=True,
+                )
+                cond_acc.update(cond_cont)
+                tgt_acc.update(target_s1)
+                if energy_sampler is not None:
+                    energy_sampler.update(cond_cont[:, energy_idx])
+                sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
+                sec_phys = sec_cont[:, :, 4:6][sec_valid]
+                if len(sec_phys) > 0:
+                    sec_phys_acc.update(sec_phys)
+        cond_norm = cond_acc.to_normalizer()
+        tgt_norm = tgt_acc.to_normalizer()
+        sec_phys_norm = sec_phys_acc.to_normalizer()
+        energy_sample = (
+            energy_sampler.sample
+            if energy_sampler is not None
+            else np.empty(0, dtype=np.float32)
+        )
+        if cache is not None:
+            cache.normalizers[norm_key] = setup_cache.NormalizerEntry(
+                cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_sample
             )
-            cond_acc.update(cond_cont)
-            tgt_acc.update(target_s1)
-            if energy_sampler is not None:
-                energy_sampler.update(cond_cont[:, energy_idx])
-            sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
-            sec_phys = sec_cont[:, :, 4:6][sec_valid]
-            if len(sec_phys) > 0:
-                sec_phys_acc.update(sec_phys)
-    cond_norm = cond_acc.to_normalizer()
-    tgt_norm = tgt_acc.to_normalizer()
-    sec_phys_norm = sec_phys_acc.to_normalizer()
+
     total_train_batches = n_train_steps // t["batch_size"]
     echo(f"  ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches")
 
-    if energy_sampler is not None and energy_sampler.n_seen > 0:
+    if energy_router_active and energy_sample.size > 0:
         assert cond_norm.mean is not None and cond_norm.std is not None
         normalized_sample = (
-            energy_sampler.sample - cond_norm.mean[energy_idx]
+            energy_sample - cond_norm.mean[energy_idx]
         ) / cond_norm.std[energy_idx]
         quantiles = np.linspace(0.0, 1.0, router_cfg["n_experts"])
         centers_init = np.quantile(normalized_sample, quantiles).astype(np.float32)
@@ -140,6 +198,14 @@ def run_train_job(
         echo(
             f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
         )
+    elif energy_router_active:
+        echo(
+            "  warning: no energy samples collected — EnergyRouter falls back to "
+            "default centers"
+        )
+
+    if cache is not None:
+        setup_cache.save(data, files, cache, echo=echo)
 
     train_ds = StreamingStepsDataset(
         files=files,

tests/test_pipeline.py

@@ -0,0 +1,196 @@
+import copy
+
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+
+from giant import config as gconfig
+from giant.data import setup_cache
+from giant.pipeline import run_train_job
+
+
+def _unit(v):
+    v = np.asarray(v, dtype=np.float64)
+    n = np.linalg.norm(v)
+    return v / n if n > 1e-9 else np.array([0.0, 0.0, 1.0])
+
+
+def _make_synthetic_steps(path, n_events=20, seed=0):
+    """A tiny but schema-complete synthetic steps parquet for run_train_job.
+
+    pdg/material/process are assigned deterministically by row index (not
+    random) so tests that assert on the resulting vocab/proc maps aren't
+    flaky; only continuous quantities (positions/energies/directions) are
+    drawn from `rng`.
+    """
+    rng = np.random.default_rng(seed)
+    materials = ["G4_AIR", "G4_Fe"]
+    pdgs = [11, 22]
+    processes = ["eIoni", "phot", "compt"]
+    rows = []
+    row_idx = 0
+    for event_id in range(n_events):
+        n_steps = int(rng.integers(2, 4))
+        for s in range(n_steps):
+            pre_E = float(rng.uniform(50.0, 500.0))
+            n_sec = int(rng.integers(0, 3))
+            frac_dep = float(rng.uniform(0.05, 0.3))
+            frac_sec = float(rng.uniform(0.05, 0.2)) if n_sec > 0 else 0.0
+            frac_post = 1.0 - frac_dep - frac_sec
+            edep = pre_E * frac_dep
+            e_sec = pre_E * frac_sec
+            post_E = pre_E * frac_post
+            pre_pos = rng.uniform(-10, 10, size=3)
+            step_length = float(rng.uniform(0.1, 5.0))
+            pre_dir = np.array([0.0, 0.0, 1.0])
+            post_dir = _unit(rng.normal(size=3))
+            post_pos = pre_pos + step_length * pre_dir
+            sec_energies = (
+                list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
+            )
+            sec_pdgs = [pdgs[(row_idx + j) % 2] for j in range(n_sec)]
+            sec_dirs = [_unit(rng.normal(size=3)) for _ in range(n_sec)]
+            rows.append(
+                {
+                    "event_id": event_id,
+                    "pdg": pdgs[row_idx % 2],
+                    "pre_x": pre_pos[0],
+                    "pre_y": pre_pos[1],
+                    "pre_z": pre_pos[2],
+                    "pre_E": pre_E,
+                    "pre_dx": pre_dir[0],
+                    "pre_dy": pre_dir[1],
+                    "pre_dz": pre_dir[2],
+                    "material": materials[row_idx % 2],
+                    "layer_id": s,
+                    "child_track_ids": list(range(n_sec)),
+                    "e_sec": e_sec,
+                    "process": processes[row_idx % 3],
+                    "step_length": step_length,
+                    "post_E": post_E,
+                    "edep": edep,
+                    "post_dx": post_dir[0],
+                    "post_dy": post_dir[1],
+                    "post_dz": post_dir[2],
+                    "post_x": post_pos[0],
+                    "post_y": post_pos[1],
+                    "post_z": post_pos[2],
+                    "sec_E_list": sec_energies,
+                    "sec_pdg_list": sec_pdgs,
+                    "sec_dx_list": [d[0] for d in sec_dirs],
+                    "sec_dy_list": [d[1] for d in sec_dirs],
+                    "sec_dz_list": [d[2] for d in sec_dirs],
+                }
+            )
+            row_idx += 1
+    pd.DataFrame(rows).to_parquet(path)
+    return path
+
+
+def _tiny_cfg(**train_overrides):
+    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
+    cfg["train"].update(
+        {
+            "epochs": 1,
+            "batch_size": 8,
+            "val_fraction": 0.2,
+            "seed": 0,
+            "warmup_epochs": 0,
+            "validate_every": 0,
+            "max_val_batches": 1,
+        }
+    )
+    cfg["train"].update(train_overrides)
+    cfg["model"].update({"hidden_dim": 8, "n_blocks": 1, "emb_dim": 4, "dropout": 0.0})
+    return cfg
+
+
+def _run(data, out_dir, cfg=None, **kwargs):
+    echoed: list[str] = []
+    run_train_job(
+        data=data,
+        cfg=cfg or _tiny_cfg(),
+        out_dir=out_dir,
+        device=torch.device("cpu"),
+        shuffle_buffer=64,
+        num_workers=0,
+        echo=echoed.append,
+        **kwargs,
+    )
+    return echoed
+
+
+@pytest.fixture
+def data(tmp_path):
+    return _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+
+
+def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
+    echo1 = _run(data, tmp_path / "out1")
+    assert any("fitting normalizer (streaming)" in m for m in echo1)
+
+    def _forbidden(*a, **k):
+        raise AssertionError("should be served from cache, not recomputed")
+
+    monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
+    monkeypatch.setattr("giant.pipeline.iter_file_chunks", _forbidden)
+
+    echo2 = _run(data, tmp_path / "out2")
+    joined = "\n".join(echo2)
+    assert "event index: cache hit" in joined
+    assert "vocabulary maps: cache hit" in joined
+    assert "normalizer: cache hit" in joined
+
+
+def test_run_train_job_no_cache_setup_never_writes_sidecar(tmp_path, data):
+    _run(data, tmp_path / "out", cache_setup=False)
+    assert not setup_cache.sidecar_path(data).exists()
+
+
+def test_run_train_job_rebuild_setup_cache_ignores_existing(tmp_path, data):
+    files = [data]
+    stale = setup_cache.SetupCache.empty(files)
+    stale.vocab = ({999999: 0}, {"G4_AIR": 0})  # deliberately wrong
+    setup_cache.save(data, files, stale)
+
+    _run(data, tmp_path / "out", rebuild_setup_cache=True)
+
+    loaded = setup_cache.load(data, files)
+    assert loaded is not None
+    assert loaded.vocab is not None
+    assert set(loaded.vocab[0].keys()) == {11, 22}
+    assert set(loaded.vocab[1].keys()) == {"G4_AIR", "G4_Fe"}
+
+
+def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypatch):
+    _run(data, tmp_path / "out1", cfg=_tiny_cfg(val_fraction=0.1))
+
+    def _forbidden(*a, **k):
+        raise AssertionError("vocab should be served from cache")
+
+    monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
+
+    echo2 = _run(data, tmp_path / "out2", cfg=_tiny_cfg(val_fraction=0.3))
+    joined = "\n".join(echo2)
+    assert "vocabulary maps: cache hit" in joined
+    assert "fitting normalizer (streaming)" in joined
+
+
+def test_run_train_job_matches_uncached_output(tmp_path, data):
+    _run(data, tmp_path / "uncached", cache_setup=False)
+    _run(data, tmp_path / "cached1", cache_setup=True)
+    _run(data, tmp_path / "cached2", cache_setup=True)  # second is a cache hit
+
+    uncached = torch.load(tmp_path / "uncached" / "last.pt", weights_only=False)
+    cached = torch.load(tmp_path / "cached2" / "last.pt", weights_only=False)
+
+    for key in ("cond", "target", "sec_phys"):
+        np.testing.assert_allclose(
+            uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"]
+        )
+        np.testing.assert_allclose(
+            uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"]
+        )
+    assert uncached["pdg_map"] == cached["pdg_map"]
+    assert uncached["mat_map"] == cached["mat_map"]

tests/test_setup_cache.py

@@ -0,0 +1,227 @@
+import json
+import os
+import stat
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from giant.data import setup_cache
+from giant.data.setup_cache import NormalizerEntry, SetupCache
+from giant.data.transforms import Normalizer
+
+
+def _touch_parquet(path, n=1):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    pd.DataFrame(
+        {"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}
+    ).to_parquet(path)
+    return path
+
+
+def _normalizer(width=3):
+    norm = Normalizer()
+    norm.mean = np.zeros(width, dtype=np.float32)
+    norm.std = np.ones(width, dtype=np.float32)
+    return norm
+
+
+def _entry(n_train_steps=100, sample=None):
+    sample = np.array([1.0, 2.0, 3.0], dtype=np.float32) if sample is None else sample
+    return NormalizerEntry(
+        _normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample
+    )
+
+
+# ── sidecar_path ─────────────────────────────────────────────────────────
+
+
+def test_sidecar_path_single_file(tmp_path):
+    f = tmp_path / "shard.parquet"
+    assert (
+        setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
+    )
+
+
+def test_sidecar_path_directory(tmp_path):
+    d = tmp_path / "pbwo4"
+    assert setup_cache.sidecar_path(d) == tmp_path / "pbwo4.giant_train_cache.json"
+
+
+def test_sidecar_path_manifest(tmp_path):
+    m = tmp_path / "pools" / "full.manifest"
+    assert (
+        setup_cache.sidecar_path(m)
+        == tmp_path / "pools" / "full.manifest.giant_train_cache.json"
+    )
+
+
+# ── fingerprint_files ────────────────────────────────────────────────────
+
+
+def test_fingerprint_files_order_preserving(tmp_path):
+    a = _touch_parquet(tmp_path / "a.parquet")
+    b = _touch_parquet(tmp_path / "b.parquet")
+
+    forward = setup_cache.fingerprint_files([a, b])
+    backward = setup_cache.fingerprint_files([b, a])
+
+    assert forward[0][0] == str(a.resolve())
+    assert forward[1][0] == str(b.resolve())
+    assert backward[0][0] == str(b.resolve())
+    assert forward != backward
+
+
+# ── save / load round trip ──────────────────────────────────────────────
+
+
+def test_save_load_round_trip(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+
+    cache = SetupCache.empty(files)
+    cache.vocab = ({11: 0, 22: 1}, {"G4_AIR": 0})
+    cache.event_index = (np.array([1, 2, 3]), np.array([10, 20, 30]))
+    cache.proc_maps[4] = {"eIoni": 0, "phot": 1}
+    cache.normalizers["valfrac=0.1_seed=0_cond=physical"] = _entry()
+
+    setup_cache.save(data, files, cache)
+    loaded = setup_cache.load(data, files)
+
+    assert loaded is not None
+    assert loaded.vocab == ({11: 0, 22: 1}, {"G4_AIR": 0})
+    assert loaded.event_index is not None
+    np.testing.assert_array_equal(loaded.event_index[0], [1, 2, 3])
+    np.testing.assert_array_equal(loaded.event_index[1], [10, 20, 30])
+    assert loaded.proc_maps == {4: {"eIoni": 0, "phot": 1}}
+    entry = loaded.normalizers["valfrac=0.1_seed=0_cond=physical"]
+    assert entry.cond_norm.mean is not None
+    np.testing.assert_allclose(entry.cond_norm.mean, np.zeros(3, dtype=np.float32))
+    assert entry.n_train_steps == 100
+    np.testing.assert_allclose(entry.energy_reservoir_sample, [1.0, 2.0, 3.0])
+
+
+def test_load_missing_sidecar_returns_none(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    assert setup_cache.load(data, [data]) is None
+
+
+def test_load_corrupt_json_returns_none(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    setup_cache.sidecar_path(data).write_text("not valid json {{{")
+    assert setup_cache.load(data, [data]) is None
+
+
+def test_load_invalidates_on_dims_mismatch(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    path = setup_cache.sidecar_path(data)
+    raw = json.loads(path.read_text())
+    raw["dims"]["K_MAX"] = raw["dims"]["K_MAX"] + 1
+    path.write_text(json.dumps(raw))
+
+    assert setup_cache.load(data, files) is None
+
+
+def test_load_invalidates_on_format_version_mismatch(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    path = setup_cache.sidecar_path(data)
+    raw = json.loads(path.read_text())
+    raw["format_version"] = raw["format_version"] + 1
+    path.write_text(json.dumps(raw))
+
+    assert setup_cache.load(data, files) is None
+
+
+def test_load_invalidates_on_file_content_change(tmp_path):
+    data = tmp_path / "shard.parquet"
+    _touch_parquet(data, n=1)
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    _touch_parquet(data, n=50)  # different size -> fingerprint changes
+
+    assert setup_cache.load(data, files) is None
+
+
+def test_load_soft_warns_on_git_hash_mismatch_but_still_hits(tmp_path, capsys):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    cache = SetupCache.empty(files)
+    cache.git_hash = "not-a-real-git-hash"
+    setup_cache.save(data, files, cache)
+
+    loaded = setup_cache.load(data, files)
+
+    assert loaded is not None
+    err = capsys.readouterr().err
+    assert "not-a-real-git-hash" in err
+
+
+# ── save: atomicity / robustness ────────────────────────────────────────
+
+
+def test_save_is_atomic_no_stray_tmp_file(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    leftovers = [p for p in tmp_path.iterdir() if ".tmp." in p.name]
+    assert leftovers == []
+
+
+def test_save_degrades_gracefully_on_permission_error(tmp_path):
+    if os.geteuid() == 0:
+        pytest.skip("root bypasses directory permission bits")
+    data_dir = tmp_path / "ro"
+    data_dir.mkdir()
+    data = _touch_parquet(data_dir / "shard.parquet")
+    files = [data]
+
+    warnings = []
+    mode = data_dir.stat().st_mode
+    data_dir.chmod(stat.S_IREAD | stat.S_IEXEC)
+    try:
+        setup_cache.save(data, files, SetupCache.empty(files), echo=warnings.append)
+    finally:
+        data_dir.chmod(mode)
+
+    assert any("could not write" in w for w in warnings)
+    assert not setup_cache.sidecar_path(data).exists()
+
+
+def test_save_merges_non_colliding_normalizer_keys(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+
+    cache1 = SetupCache.empty(files)
+    cache1.normalizers["k1"] = _entry(n_train_steps=1)
+    setup_cache.save(data, files, cache1)
+
+    cache2 = SetupCache.empty(files)
+    cache2.normalizers["k2"] = _entry(n_train_steps=2)
+    setup_cache.save(data, files, cache2)
+
+    loaded = setup_cache.load(data, files)
+    assert loaded is not None
+    assert set(loaded.normalizers.keys()) == {"k1", "k2"}
+    assert loaded.normalizers["k1"].n_train_steps == 1
+    assert loaded.normalizers["k2"].n_train_steps == 2
+
+
+# ── n_train_steps_for_split ──────────────────────────────────────────────
+
+
+def test_n_train_steps_for_split_matches_full_scan():
+    unique_ids = np.array([1, 2, 3, 4, 5])
+    counts = np.array([10, 20, 30, 40, 50])
+    train_events_arr = np.array([2, 4, 5])
+
+    result = setup_cache.n_train_steps_for_split(unique_ids, counts, train_events_arr)
+
+    assert result == 20 + 40 + 50