Add dwarf warm-cache to precompute the setup-stage sidecar

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 <noreply@anthropic.com>

CLAUDE.md

@@ -20,7 +20,7 @@ giant analyze submit rollout.yaml --accounting-group cms   # parallel rollout-vs
 giant analyze render <run_dir> --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)
 ```
 

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,

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,

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.")

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