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_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