Add dwarf warm-cache to precompute the setup-stage sidecar
CI / Format (ruff format) (push) Successful in 26s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 37s
CI / Lint (ruff check) (pull_request) Successful in 36s
CI / Format (ruff format) (pull_request) Successful in 31s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 33s
CI / Tests (push) Successful in 1m43s
CI / Tests (pull_request) Successful in 1m39s

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>
This commit is contained in:
2026-07-30 10:53:31 +02:00
parent 26aa9d3fde
commit 471a81b5e7
5 changed files with 314 additions and 26 deletions
+80
View File
@@ -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