lars
e7478c36fb
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>
228 lines
7.5 KiB
Python
228 lines
7.5 KiB
Python
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
|