From d25dfc03433cb53b682c4cdf9c259f2041624581 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 17 Aug 2026 08:54:14 +0200 Subject: [PATCH] Let dwarf warm-cache take --config so it can't under-warm a config's cache keys (gitea #59) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit warm-cache built its config from DEFAULT_CONFIG with only a handful of flags overridable, so it had no way to express settings like stage2_model.particle_type.n_classes. configs/baseline.toml sets that to 32; warm-cache always warmed the pdg top-N map under the emb_dim default (16) instead, so a `giant train --config configs/baseline.toml` run silently missed the cache and repaid the full parquet scan warm-cache exists to avoid. warm-cache now accepts the same --config a training run takes and resolves every value run_setup_stage needs (val_fraction/seed, conditioning types, both stages' router, particle_type.n_classes, ...) from one gconfig.merge_cli_overrides + validate_config pass, exactly like giant train's own pipeline does — so warming and training are guaranteed to agree. Per user decision, --config is mutually exclusive with the individual --val-fraction/--seed/--particle-conditioning/ --material-conditioning/--router*/flags (rejected outright rather than silently layered on top), since a hardcoded CLI default clobbering an unset config value is the same failure mode one level down. Also drops a hardcoded stage2_model.router/k_max override that was a no-op against today's defaults but would have clobbered a config setting either one away from its default — same bug class. Adding validate_config surfaced that the existing test_warm_cache_router_process_warms_proc_map test was warming a router.type="process" + conditioning.particle.type="physical" (the CLI's old hardcoded default) combination that giant train's own validate_config would already reject as incompatible — fixed by passing --particle-conditioning embedding, which is what a working --router-type process run actually requires. Co-Authored-By: Claude Sonnet 5 --- giant/tools/dwarf.py | 89 +++++++++++++++++++------- giant/tools/warm_setup_cache.py | 109 +++++++++++++++++++------------- tests/test_dwarf.py | 65 +++++++++++++++++++ 3 files changed, 197 insertions(+), 66 deletions(-) diff --git a/giant/tools/dwarf.py b/giant/tools/dwarf.py index 5983735..29b0661 100644 --- a/giant/tools/dwarf.py +++ b/giant/tools/dwarf.py @@ -442,43 +442,65 @@ def warm_cache( Path, typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"), ], + config: Annotated[ + Optional[Path], + typer.Option( + "--config", + "-c", + help="TOML config file to warm for — same file the `giant train` run(s) will use. " + "Mutually exclusive with the flags below (put val-fraction/seed/conditioning/router " + "settings in the file itself, so warming and training can't disagree on them)", + ), + ] = None, val_fraction: Annotated[ - float, + Optional[float], typer.Option( "--val-fraction", "-f", - help="Must match the `giant train` run(s) to warm for", + help="Must match the `giant train` run(s) to warm for. Not allowed together with --config", ), - ] = 0.1, + ] = None, seed: Annotated[ - int, - typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"), - ] = 0, + Optional[int], + typer.Option( + "--seed", + "-s", + help="Must match the `giant train` run(s) to warm for. Not allowed together with --config", + ), + ] = None, particle_conditioning: Annotated[ - Conditioning, + Optional[Conditioning], typer.Option( "--particle-conditioning", - help="Must match the `giant train` run(s)' conditioning.particle.type to warm for", + help="Must match the `giant train` run(s)' conditioning.particle.type to warm for. " + "Not allowed together with --config", ), - ] = Conditioning.physical, + ] = None, material_conditioning: Annotated[ - Conditioning, + Optional[Conditioning], typer.Option( "--material-conditioning", help="Must match the `giant train` run(s)' conditioning.material.type " "to warm for — independent of --particle-conditioning " - "(the two axes may differ)", + "(the two axes may differ). Not allowed together with --config", ), - ] = Conditioning.physical, + ] = None, router: Annotated[ - bool, + Optional[bool], typer.Option( "--router/--no-router", - help="Warm the process vocabulary too (only takes effect with --router-type process)", + help="Warm the process vocabulary too (only takes effect with --router-type process). " + "Not allowed together with --config", ), - ] = 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, + ] = None, + router_type: Annotated[ + Optional[str], + typer.Option("--router-type", help="Router implementation name. Not allowed together with --config"), + ] = None, + n_experts: Annotated[ + Optional[int], + typer.Option("--n-experts", help="Number of routed experts. Not allowed together with --config"), + ] = None, rebuild: Annotated[ bool, typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"), @@ -487,17 +509,38 @@ def warm_cache( """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/--particle-conditioning/ - --material-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. + either --config, or the given --val-fraction/--seed/ + --particle-conditioning/--material-conditioning/--router* flags, 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. """ + flag_overrides = { + "--val-fraction": val_fraction, + "--seed": seed, + "--particle-conditioning": particle_conditioning, + "--material-conditioning": material_conditioning, + "--router/--no-router": router, + "--router-type": router_type, + "--n-experts": n_experts, + } + if config is not None: + given = [name for name, value in flag_overrides.items() if value is not None] + if given: + typer.echo( + f"error: --config cannot be combined with {', '.join(given)} " + "— put these settings in the config file instead", + err=True, + ) + raise typer.Exit(1) + run_warm_setup_cache( data=str(data), + config_path=config, val_fraction=val_fraction, seed=seed, - particle_conditioning=particle_conditioning.value, - material_conditioning=material_conditioning.value, + particle_conditioning=particle_conditioning.value if particle_conditioning is not None else None, + material_conditioning=material_conditioning.value if material_conditioning is not None else None, router_enabled=router, router_type=router_type, n_experts=n_experts, diff --git a/giant/tools/warm_setup_cache.py b/giant/tools/warm_setup_cache.py index 1feed1c..9253311 100644 --- a/giant/tools/warm_setup_cache.py +++ b/giant/tools/warm_setup_cache.py @@ -10,63 +10,86 @@ for the sidecar itself. from pathlib import Path from giant import config as gconfig -from giant.constants import K_MAX from giant.pipeline import run_setup_stage def run_warm_setup_cache( data: str, - val_fraction: float = 0.1, - seed: int = 0, - particle_conditioning: str = "physical", - material_conditioning: str = "physical", - router_enabled: bool = False, - router_type: str = "energy", - n_experts: int = 4, + config_path: Path | None = None, + val_fraction: float | None = None, + seed: int | None = None, + particle_conditioning: str | None = None, + material_conditioning: str | None = None, + router_enabled: bool | None = None, + router_type: str | None = None, + n_experts: int | None = None, rebuild: bool = False, echo=print, ) -> None: """Populate (or refresh) the setup cache sidecar for `data`. - `val_fraction`/`seed`/`particle_conditioning`/`material_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. The two - conditioning axes are independent and may differ. - `router_enabled`/`router_type`/`n_experts` only matter for - `router_type == "process"` (warms that `n_experts`'s process map); the - energy-router quantile summary is always collected regardless, so a - later `--router-type energy` run never needs to rescan just to seed - centers. + Two mutually exclusive ways to select what to warm for (enforced by the + caller, `giant.tools.dwarf.warm_cache` — this function just trusts + whichever combination it's given): + + - `config_path`: the same TOML `giant train --config` takes. Every value + `run_setup_stage` needs (`train.val_fraction`/`seed`, + `conditioning.particle`/`material.type`, both stages' `router`, + `stage2_model.particle_type.n_classes`, ...) is read from the one + resulting merged `cfg`, so a later `giant train --config ` + run resolves to exactly the same cache keys — see gitea #59. + - The individual flags below: `val_fraction`/`seed`/ + `particle_conditioning`/`material_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. The two conditioning axes are independent and may + differ. `router_enabled`/`router_type`/`n_experts` only matter for + `router_type == "process"` (warms that `n_experts`'s process map); the + energy-router quantile summary is always collected regardless, so a + later `--router-type energy` run never needs to rescan just to seed + centers. + + Any flag left `None` is omitted from the merge, so it falls back to + `DEFAULT_CONFIG`'s own value (or the config file's, if `config_path` is + given) instead of silently overriding it — see gitea #59. """ - router_cfg = { - "enabled": router_enabled, - "type": router_type, - "n_experts": n_experts, - } - # Merged against DEFAULT_CONFIG (not a hand-rolled partial dict) so - # run_setup_stage always sees every key it might read (e.g. - # conditioning.particle.emb_dim, stage2_model.particle_type.target) at - # its real default, not silently missing/None — see issues.md Issue 1. + overrides: dict = {} + + conditioning_overrides: dict = {} + if particle_conditioning is not None: + conditioning_overrides["particle"] = {"type": particle_conditioning} + if material_conditioning is not None: + conditioning_overrides["material"] = {"type": material_conditioning} + if conditioning_overrides: + overrides["conditioning"] = conditioning_overrides + # This CLI only ever configures one router (matching today's single - # --router-type flag), so it's placed on stage1_model; stage2_model's - # stays disabled. - cfg = gconfig.merge_cli_overrides( - gconfig.DEFAULT_CONFIG, - None, - { - "conditioning": { - "particle": {"type": particle_conditioning}, - "material": {"type": material_conditioning}, - }, - "stage1_model": {"router": router_cfg}, - "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX}, - }, - ) + # --router-type flag), so it's placed on stage1_model; stage2_model's is + # left to DEFAULT_CONFIG/the config file rather than forced disabled. + router_overrides: dict = {} + if router_enabled is not None: + router_overrides["enabled"] = router_enabled + if router_type is not None: + router_overrides["type"] = router_type + if n_experts is not None: + router_overrides["n_experts"] = n_experts + if router_overrides: + overrides["stage1_model"] = {"router": router_overrides} + + train_overrides: dict = {} + if val_fraction is not None: + train_overrides["val_fraction"] = val_fraction + if seed is not None: + train_overrides["seed"] = seed + if train_overrides: + overrides["train"] = train_overrides + + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, overrides) + gconfig.validate_config(cfg) run_setup_stage( Path(data), - val_fraction=val_fraction, - seed=seed, + val_fraction=cfg["train"]["val_fraction"], + seed=cfg["train"]["seed"], cfg=cfg, cache_setup=True, rebuild_setup_cache=rebuild, diff --git a/tests/test_dwarf.py b/tests/test_dwarf.py index e312e0f..3b7d899 100644 --- a/tests/test_dwarf.py +++ b/tests/test_dwarf.py @@ -124,6 +124,11 @@ def test_warm_cache_router_process_warms_proc_map(tmp_path): [ "warm-cache", str(data), + # router.type="process" is incompatible with the default + # conditioning.particle.type="physical" (validate_config, now + # enforced by warm-cache too — see gitea #59). + "--particle-conditioning", + "embedding", "--router", "--router-type", "process", @@ -167,3 +172,63 @@ def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path): assert loaded is not None assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers assert "valfrac=0.3_seed=0_pcond=physical_mcond=physical" in loaded.normalizers + + +def test_warm_cache_config_warms_particle_type_n_classes(tmp_path): + """gitea #59: a config setting stage2_model.particle_type.n_classes away + from its 0 (= inherit conditioning.particle.emb_dim) default must warm + the pdg top-N map under that n_classes, not the emb_dim default, so a + later `giant train --config ` run hits it instead of quietly + re-scanning every parquet file.""" + data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20) + config_path = tmp_path / "config.toml" + config_path.write_text("[meta]\nconfig_version = 3\n\n[stage2_model.particle_type]\nn_classes = 32\n") + + runner.invoke(app, ["warm-cache", str(data), "--config", str(config_path)]) + result = runner.invoke(app, ["warm-cache", str(data), "--config", str(config_path)]) + + assert result.exit_code == 0, result.output + assert "pdg top-N map: cache hit" in result.output + assert "32 classes" in result.output + + +def test_warm_cache_config_rejects_val_fraction_flag(tmp_path): + data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20) + config_path = tmp_path / "config.toml" + config_path.write_text("[meta]\nconfig_version = 3\n") + + result = runner.invoke( + app, + ["warm-cache", str(data), "--config", str(config_path), "--val-fraction", "0.2"], + ) + + assert result.exit_code != 0 + assert "--config" in result.output + assert "--val-fraction" in result.output + + +def test_warm_cache_config_rejects_router_flags(tmp_path): + data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20) + config_path = tmp_path / "config.toml" + config_path.write_text("[meta]\nconfig_version = 3\n") + + result = runner.invoke( + app, + [ + "warm-cache", + str(data), + "--config", + str(config_path), + "--router", + "--router-type", + "process", + "--n-experts", + "3", + ], + ) + + assert result.exit_code != 0 + assert "--config" in result.output + assert "--router/--no-router" in result.output + assert "--router-type" in result.output + assert "--n-experts" in result.output