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>

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