From 471a81b5e77084c4d1864c3c796a03c31d4dba0c Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 30 Jul 2026 10:53:31 +0200 Subject: [PATCH] Add dwarf warm-cache to precompute the setup-stage sidecar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lets the vocab maps, event-id split index, and normalizer stats be warmed once for a dataset (right after `dwarf convert`, or before a `dwarf hparam-scan` sweep) without needing to also start training. Extracts the setup-stage logic out of giant/pipeline.py:run_train_job into a standalone run_setup_stage() (returning a SetupStageResult), reused by both run_train_job and the new dwarf command's scripts/warm_setup_cache.py — a behavior-preserving refactor, covered by the existing test_pipeline.py suite. Co-Authored-By: Claude Sonnet 5 --- CLAUDE.md | 2 +- giant/pipeline.py | 131 +++++++++++++++++++++++++++++------- scripts/dwarf.py | 75 +++++++++++++++++++++ scripts/warm_setup_cache.py | 52 ++++++++++++++ tests/test_dwarf.py | 80 ++++++++++++++++++++++ 5 files changed, 314 insertions(+), 26 deletions(-) create mode 100644 scripts/warm_setup_cache.py diff --git a/CLAUDE.md b/CLAUDE.md index 27c19c2..dc1a500 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -20,7 +20,7 @@ giant analyze submit rollout.yaml --accounting-group cms # parallel rollout-vs giant analyze render --gallery # render PDFs + HTML gallery (run_dir from prep/submit) dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen, # bump-schema, status, update-manifest, create-manifest, - # make-root, build-geometry-oracle, hparam-scan + # make-root, build-geometry-oracle, warm-cache, hparam-scan # (see scripts/dwarf.py) ``` diff --git a/giant/pipeline.py b/giant/pipeline.py index b265dbe..b6bdfa5 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from pathlib import Path import numpy as np @@ -21,6 +22,7 @@ from giant.data.loader import ( build_process_map_from_files, ) from giant.data.transforms import ( + Normalizer, build_features, _WelfordAccumulator, _ReservoirSampler, @@ -31,23 +33,47 @@ from giant.model.network import build_models, build_critics from giant.train import train as run_training -def run_train_job( - data: Path, - cfg: dict, - out_dir: Path, - device: torch.device, - shuffle_buffer: int, - num_workers: int, - resume: Path | None = None, +@dataclass +class SetupStageResult: + """Everything `run_train_job`'s pre-epoch setup stage derives from `data`. + + Also returned standalone by `run_setup_stage` for callers (e.g. `dwarf + warm-cache`) that only want to populate/refresh the setup cache sidecar + without actually training a model. + """ + + files: list[Path] + pdg_map: dict[int, int] + mat_map: dict[str, int] + proc_map: dict[str, int] | None + cond_norm: Normalizer + tgt_norm: Normalizer + sec_phys_norm: Normalizer + train_events: set + val_events: set + n_train_steps: int + + +def run_setup_stage( + data: str | Path, + val_fraction: float, + seed: int, + conditioning: str, + router_cfg: dict, cache_setup: bool = True, rebuild_setup_cache: bool = False, echo=print, -) -> None: - t, m = cfg["train"], cfg["model"] - config.seed_everything(t["seed"]) - - out_dir = Path(out_dir) +) -> SetupStageResult: + """Scan `data` for everything training needs before the epoch loop: the + train/val event split, pdg/material vocab maps, an optional process map + (`router_cfg.type == "process"`), and the Stage-1/Stage-2 normalizers. + Reads from and writes to the `giant.data.setup_cache` sidecar when + `cache_setup` is set (`rebuild_setup_cache` ignores — but still + refreshes — any existing sidecar content). `router_cfg` may be mutated + in place: an active `EnergyRouter` (`router_cfg["type"] == "energy"`) + gets its `centers_init` seeded from real data quantiles here. + """ files = find_parquet_files(data) echo(f"found {len(files)} parquet file(s)") @@ -70,7 +96,7 @@ def run_train_job( cache.event_index = (unique_ids, counts) train_events, val_events = make_event_split( - unique_ids, val_fraction=t["val_fraction"], seed=t["seed"] + unique_ids, val_fraction=val_fraction, seed=seed ) events_arr = np.array(sorted(train_events)) n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr) @@ -93,12 +119,6 @@ def run_train_job( if cache is not None: cache.vocab = (pdg_map, mat_map) - router_cfg = m["router"] - if t["mode"] == "wgan" and router_cfg.get("enabled"): - raise ValueError( - "--mode wgan does not support --router (no routed WGAN generator/" - "critic exists) — disable one or the other" - ) proc_map: dict[str, int] | None = None if router_cfg.get("enabled") and router_cfg.get("type") == "process": n_experts = router_cfg["n_experts"] @@ -116,12 +136,11 @@ def run_train_job( if cache is not None: cache.proc_maps[n_experts] = proc_map - conditioning = m["conditioning"] energy_router_active = ( router_cfg.get("enabled") and router_cfg.get("type") == "energy" ) energy_idx = router_cfg.get("energy_idx", 3) - norm_key = setup_cache.normalizer_key(t["val_fraction"], t["seed"], conditioning) + norm_key = setup_cache.normalizer_key(val_fraction, seed, conditioning) entry = cache.normalizers.get(norm_key) if cache is not None else None if entry is not None: @@ -184,9 +203,6 @@ def run_train_job( cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_sample ) - total_train_batches = n_train_steps // t["batch_size"] - echo(f" ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches") - 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 = ( @@ -207,6 +223,71 @@ def run_train_job( if cache is not None: setup_cache.save(data, files, cache, echo=echo) + return SetupStageResult( + files=files, + pdg_map=pdg_map, + mat_map=mat_map, + proc_map=proc_map, + cond_norm=cond_norm, + tgt_norm=tgt_norm, + sec_phys_norm=sec_phys_norm, + train_events=train_events, + val_events=val_events, + n_train_steps=n_train_steps, + ) + + +def run_train_job( + data: Path, + cfg: dict, + out_dir: Path, + device: torch.device, + 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"] + config.seed_everything(t["seed"]) + + out_dir = Path(out_dir) + + router_cfg = m["router"] + if t["mode"] == "wgan" and router_cfg.get("enabled"): + raise ValueError( + "--mode wgan does not support --router (no routed WGAN generator/" + "critic exists) — disable one or the other" + ) + + conditioning = m["conditioning"] + setup = run_setup_stage( + data, + val_fraction=t["val_fraction"], + seed=t["seed"], + conditioning=conditioning, + router_cfg=router_cfg, + cache_setup=cache_setup, + rebuild_setup_cache=rebuild_setup_cache, + echo=echo, + ) + files = setup.files + pdg_map, mat_map, proc_map = setup.pdg_map, setup.mat_map, setup.proc_map + cond_norm, tgt_norm, sec_phys_norm = ( + setup.cond_norm, + setup.tgt_norm, + setup.sec_phys_norm, + ) + train_events, val_events, n_train_steps = ( + setup.train_events, + setup.val_events, + setup.n_train_steps, + ) + + total_train_batches = n_train_steps // t["batch_size"] + echo(f" ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches") + train_ds = StreamingStepsDataset( files=files, split_events=train_events, diff --git a/scripts/dwarf.py b/scripts/dwarf.py index de519f0..beacc89 100644 --- a/scripts/dwarf.py +++ b/scripts/dwarf.py @@ -25,6 +25,7 @@ from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan from scripts.migrate_geant_steps import run_migration from scripts.steps_to_parquet import convert_steps_to_parquet from scripts.steps_to_parquet_parallel import run_parallel_job +from scripts.warm_setup_cache import run_warm_setup_cache app = typer.Typer(no_args_is_help=True) @@ -471,6 +472,80 @@ def build_geometry_oracle( ) +class Conditioning(str, Enum): + physical = "physical" + embedding = "embedding" + + +@app.command("warm-cache") +def warm_cache( + data: Annotated[ + Path, + typer.Argument( + help="Parquet file, directory, or .manifest — same as `giant train`'s" + ), + ], + val_fraction: Annotated[ + float, + typer.Option( + "--val-fraction", + "-f", + help="Must match the `giant train` run(s) to warm for", + ), + ] = 0.1, + seed: Annotated[ + int, + typer.Option( + "--seed", "-s", help="Must match the `giant train` run(s) to warm for" + ), + ] = 0, + conditioning: Annotated[ + Conditioning, + typer.Option( + "--conditioning", help="Must match the `giant train` run(s) to warm for" + ), + ] = Conditioning.physical, + router: Annotated[ + bool, + typer.Option( + "--router/--no-router", + help="Warm the process vocabulary too (only takes effect with " + "--router-type process)", + ), + ] = False, + router_type: Annotated[ + str, typer.Option("--router-type", help="Router implementation name") + ] = "energy", + n_experts: Annotated[ + int, typer.Option("--n-experts", help="Number of routed experts") + ] = 4, + rebuild: Annotated[ + bool, + typer.Option( + "--rebuild", help="Ignore any existing sidecar and recompute every section" + ), + ] = False, +) -> None: + """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time. + + Warms the vocab maps, event-id split index, and the normalizer entry for + the given --val-fraction/--seed/--conditioning, so a later `giant train` + run (or a `dwarf hparam-scan` sweep, which shares one such entry across + every run) skips straight to training. See giant/data/setup_cache.py. + """ + run_warm_setup_cache( + data=str(data), + val_fraction=val_fraction, + seed=seed, + conditioning=conditioning.value, + router_enabled=router, + router_type=router_type, + n_experts=n_experts, + rebuild=rebuild, + echo=typer.echo, + ) + + @app.command("hparam-scan") def hparam_scan( data: Annotated[str, typer.Option("--data")] = DATA_DEFAULT, diff --git a/scripts/warm_setup_cache.py b/scripts/warm_setup_cache.py new file mode 100644 index 0000000..5885069 --- /dev/null +++ b/scripts/warm_setup_cache.py @@ -0,0 +1,52 @@ +"""dwarf warm-cache — precompute `giant train`'s setup-stage sidecar ahead of time. + +Thin wrapper around `giant.pipeline.run_setup_stage` so a dataset's vocab +maps, event-id split index, and normalizer stats can be warmed once — e.g. +right after `dwarf convert`, or before kicking off a `dwarf hparam-scan` +sweep — without needing to also start training. See giant/data/setup_cache.py +for the sidecar itself. +""" + +from pathlib import Path + +from giant.pipeline import run_setup_stage + + +def run_warm_setup_cache( + data: str, + val_fraction: float = 0.1, + seed: int = 0, + conditioning: str = "physical", + router_enabled: bool = False, + router_type: str = "energy", + n_experts: int = 4, + rebuild: bool = False, + echo=print, +) -> None: + """Populate (or refresh) the setup cache sidecar for `data`. + + `val_fraction`/`seed`/`conditioning` select the normalizer cache entry + (`giant.data.setup_cache.normalizer_key`) — pass the same values a later + `giant train` invocation will use so it hits this warmed entry. + `router_enabled`/`router_type`/`n_experts` only matter for + `router_type == "process"` (warms that `n_experts`'s process map); the + energy-router reservoir sample is always collected regardless, so a + later `--router-type energy` run never needs to rescan just to seed + centers. + """ + router_cfg = { + "enabled": router_enabled, + "type": router_type, + "n_experts": n_experts, + } + run_setup_stage( + Path(data), + val_fraction=val_fraction, + seed=seed, + conditioning=conditioning, + router_cfg=router_cfg, + cache_setup=True, + rebuild_setup_cache=rebuild, + echo=echo, + ) + echo("setup cache warmed.") diff --git a/tests/test_dwarf.py b/tests/test_dwarf.py index 5ad4367..9e29958 100644 --- a/tests/test_dwarf.py +++ b/tests/test_dwarf.py @@ -1,6 +1,8 @@ from typer.testing import CliRunner +from giant.data import setup_cache from scripts.dwarf import app +from test_pipeline import _make_synthetic_steps runner = CliRunner() @@ -66,3 +68,81 @@ def test_status_reports_missing_root(tmp_path): result = runner.invoke(app, ["status", "--root", str(missing)]) assert result.exit_code != 0 assert "is not a directory" in result.output + + +def test_warm_cache_writes_sidecar(tmp_path): + data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20) + + result = runner.invoke(app, ["warm-cache", str(data)]) + + assert result.exit_code == 0, result.output + loaded = setup_cache.load(data, [data]) + assert loaded is not None + assert loaded.vocab is not None + assert loaded.event_index is not None + assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers + + +def test_warm_cache_second_run_hits_cache(tmp_path): + data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20) + runner.invoke(app, ["warm-cache", str(data)]) + + result = runner.invoke(app, ["warm-cache", str(data)]) + + assert result.exit_code == 0, result.output + assert "event index: cache hit" in result.output + assert "vocabulary maps: cache hit" in result.output + assert "normalizer: cache hit" in result.output + + +def test_warm_cache_router_process_warms_proc_map(tmp_path): + data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20) + + result = runner.invoke( + app, + [ + "warm-cache", + str(data), + "--router", + "--router-type", + "process", + "--n-experts", + "3", + ], + ) + + assert result.exit_code == 0, result.output + loaded = setup_cache.load(data, [data]) + assert loaded is not None + assert 3 in loaded.proc_maps + + +def test_warm_cache_rebuild_ignores_existing(tmp_path): + data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20) + files = [data] + stale = setup_cache.SetupCache.empty(files) + stale.vocab = ({999999: 0}, {"G4_AIR": 0}) # deliberately wrong + setup_cache.save(data, files, stale) + + result = runner.invoke(app, ["warm-cache", str(data), "--rebuild"]) + + assert result.exit_code == 0, result.output + 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} + + +def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path): + data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20) + runner.invoke(app, ["warm-cache", str(data), "--val-fraction", "0.1"]) + + result = runner.invoke(app, ["warm-cache", str(data), "--val-fraction", "0.3"]) + + assert result.exit_code == 0, result.output + assert "vocabulary maps: cache hit" in result.output + assert "fitting normalizer (streaming)" in result.output + loaded = setup_cache.load(data, [data]) + assert loaded is not None + assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers + assert "valfrac=0.3_seed=0_cond=physical" in loaded.normalizers