From e7478c36fb2679eca087a6c59eb3942e5aa54111 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 30 Jul 2026 10:28:37 +0200 Subject: [PATCH] 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 --- giant/cli.py | 22 +++ giant/data/setup_cache.py | 280 ++++++++++++++++++++++++++++++++++++++ giant/pipeline.py | 178 ++++++++++++++++-------- tests/test_pipeline.py | 196 ++++++++++++++++++++++++++ tests/test_setup_cache.py | 227 ++++++++++++++++++++++++++++++ 5 files changed, 847 insertions(+), 56 deletions(-) create mode 100644 giant/data/setup_cache.py create mode 100644 tests/test_pipeline.py create mode 100644 tests/test_setup_cache.py diff --git a/giant/cli.py b/giant/cli.py index fc4c743..462355b 100644 --- a/giant/cli.py +++ b/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, ) diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py new file mode 100644 index 0000000..cf210dd --- /dev/null +++ b/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()) diff --git a/giant/pipeline.py b/giant/pipeline.py index 15452ed..f24fcd4 100644 --- a/giant/pipeline.py +++ b/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, diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..f0711ee --- /dev/null +++ b/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"] diff --git a/tests/test_setup_cache.py b/tests/test_setup_cache.py new file mode 100644 index 0000000..c0b1e48 --- /dev/null +++ b/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