Let dwarf warm-cache take --config so it can't under-warm a config's cache keys (gitea #59)
CI / Format (ruff format) (push) Successful in 40s
CI / Lint (ruff check) (push) Successful in 40s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 42s
CI / Type check (ty) (push) Successful in 44s
CI / Format (ruff format) (pull_request) Successful in 36s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 37s
CI / Tests (push) Successful in 3m48s
CI / Tests (pull_request) Successful in 3m44s
CI / Format (ruff format) (push) Successful in 40s
CI / Lint (ruff check) (push) Successful in 40s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 42s
CI / Type check (ty) (push) Successful in 44s
CI / Format (ruff format) (pull_request) Successful in 36s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 37s
CI / Tests (push) Successful in 3m48s
CI / Tests (pull_request) Successful in 3m44s
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 <noreply@anthropic.com>
This commit is contained in:
+66
-23
@@ -442,43 +442,65 @@ def warm_cache(
|
|||||||
Path,
|
Path,
|
||||||
typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
|
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[
|
val_fraction: Annotated[
|
||||||
float,
|
Optional[float],
|
||||||
typer.Option(
|
typer.Option(
|
||||||
"--val-fraction",
|
"--val-fraction",
|
||||||
"-f",
|
"-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[
|
seed: Annotated[
|
||||||
int,
|
Optional[int],
|
||||||
typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"),
|
typer.Option(
|
||||||
] = 0,
|
"--seed",
|
||||||
|
"-s",
|
||||||
|
help="Must match the `giant train` run(s) to warm for. Not allowed together with --config",
|
||||||
|
),
|
||||||
|
] = None,
|
||||||
particle_conditioning: Annotated[
|
particle_conditioning: Annotated[
|
||||||
Conditioning,
|
Optional[Conditioning],
|
||||||
typer.Option(
|
typer.Option(
|
||||||
"--particle-conditioning",
|
"--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[
|
material_conditioning: Annotated[
|
||||||
Conditioning,
|
Optional[Conditioning],
|
||||||
typer.Option(
|
typer.Option(
|
||||||
"--material-conditioning",
|
"--material-conditioning",
|
||||||
help="Must match the `giant train` run(s)' conditioning.material.type "
|
help="Must match the `giant train` run(s)' conditioning.material.type "
|
||||||
"to warm for — independent of --particle-conditioning "
|
"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[
|
router: Annotated[
|
||||||
bool,
|
Optional[bool],
|
||||||
typer.Option(
|
typer.Option(
|
||||||
"--router/--no-router",
|
"--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,
|
] = None,
|
||||||
router_type: Annotated[str, typer.Option("--router-type", help="Router implementation name")] = "energy",
|
router_type: Annotated[
|
||||||
n_experts: Annotated[int, typer.Option("--n-experts", help="Number of routed experts")] = 4,
|
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[
|
rebuild: Annotated[
|
||||||
bool,
|
bool,
|
||||||
typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"),
|
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.
|
"""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
|
Warms the vocab maps, event-id split index, and the normalizer entry for
|
||||||
the given --val-fraction/--seed/--particle-conditioning/
|
either --config, or the given --val-fraction/--seed/
|
||||||
--material-conditioning, so a later `giant train` run (or a `dwarf
|
--particle-conditioning/--material-conditioning/--router* flags, so a
|
||||||
hparam-scan` sweep, which shares one such entry across every run) skips
|
later `giant train` run (or a `dwarf hparam-scan` sweep, which shares one
|
||||||
straight to training. See giant/data/setup_cache.py.
|
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(
|
run_warm_setup_cache(
|
||||||
data=str(data),
|
data=str(data),
|
||||||
|
config_path=config,
|
||||||
val_fraction=val_fraction,
|
val_fraction=val_fraction,
|
||||||
seed=seed,
|
seed=seed,
|
||||||
particle_conditioning=particle_conditioning.value,
|
particle_conditioning=particle_conditioning.value if particle_conditioning is not None else None,
|
||||||
material_conditioning=material_conditioning.value,
|
material_conditioning=material_conditioning.value if material_conditioning is not None else None,
|
||||||
router_enabled=router,
|
router_enabled=router,
|
||||||
router_type=router_type,
|
router_type=router_type,
|
||||||
n_experts=n_experts,
|
n_experts=n_experts,
|
||||||
|
|||||||
@@ -10,63 +10,86 @@ for the sidecar itself.
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from giant import config as gconfig
|
from giant import config as gconfig
|
||||||
from giant.constants import K_MAX
|
|
||||||
from giant.pipeline import run_setup_stage
|
from giant.pipeline import run_setup_stage
|
||||||
|
|
||||||
|
|
||||||
def run_warm_setup_cache(
|
def run_warm_setup_cache(
|
||||||
data: str,
|
data: str,
|
||||||
val_fraction: float = 0.1,
|
config_path: Path | None = None,
|
||||||
seed: int = 0,
|
val_fraction: float | None = None,
|
||||||
particle_conditioning: str = "physical",
|
seed: int | None = None,
|
||||||
material_conditioning: str = "physical",
|
particle_conditioning: str | None = None,
|
||||||
router_enabled: bool = False,
|
material_conditioning: str | None = None,
|
||||||
router_type: str = "energy",
|
router_enabled: bool | None = None,
|
||||||
n_experts: int = 4,
|
router_type: str | None = None,
|
||||||
|
n_experts: int | None = None,
|
||||||
rebuild: bool = False,
|
rebuild: bool = False,
|
||||||
echo=print,
|
echo=print,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Populate (or refresh) the setup cache sidecar for `data`.
|
"""Populate (or refresh) the setup cache sidecar for `data`.
|
||||||
|
|
||||||
`val_fraction`/`seed`/`particle_conditioning`/`material_conditioning`
|
Two mutually exclusive ways to select what to warm for (enforced by the
|
||||||
select the normalizer cache entry
|
caller, `giant.tools.dwarf.warm_cache` — this function just trusts
|
||||||
(`giant.data.setup_cache.normalizer_key`) — pass the same values a later
|
whichever combination it's given):
|
||||||
`giant train` invocation will use so it hits this warmed entry. The two
|
|
||||||
conditioning axes are independent and may differ.
|
- `config_path`: the same TOML `giant train --config` takes. Every value
|
||||||
`router_enabled`/`router_type`/`n_experts` only matter for
|
`run_setup_stage` needs (`train.val_fraction`/`seed`,
|
||||||
`router_type == "process"` (warms that `n_experts`'s process map); the
|
`conditioning.particle`/`material.type`, both stages' `router`,
|
||||||
energy-router quantile summary is always collected regardless, so a
|
`stage2_model.particle_type.n_classes`, ...) is read from the one
|
||||||
later `--router-type energy` run never needs to rescan just to seed
|
resulting merged `cfg`, so a later `giant train --config <same file>`
|
||||||
centers.
|
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 = {
|
overrides: dict = {}
|
||||||
"enabled": router_enabled,
|
|
||||||
"type": router_type,
|
conditioning_overrides: dict = {}
|
||||||
"n_experts": n_experts,
|
if particle_conditioning is not None:
|
||||||
}
|
conditioning_overrides["particle"] = {"type": particle_conditioning}
|
||||||
# Merged against DEFAULT_CONFIG (not a hand-rolled partial dict) so
|
if material_conditioning is not None:
|
||||||
# run_setup_stage always sees every key it might read (e.g.
|
conditioning_overrides["material"] = {"type": material_conditioning}
|
||||||
# conditioning.particle.emb_dim, stage2_model.particle_type.target) at
|
if conditioning_overrides:
|
||||||
# its real default, not silently missing/None — see issues.md Issue 1.
|
overrides["conditioning"] = conditioning_overrides
|
||||||
|
|
||||||
# This CLI only ever configures one router (matching today's single
|
# This CLI only ever configures one router (matching today's single
|
||||||
# --router-type flag), so it's placed on stage1_model; stage2_model's
|
# --router-type flag), so it's placed on stage1_model; stage2_model's is
|
||||||
# stays disabled.
|
# left to DEFAULT_CONFIG/the config file rather than forced disabled.
|
||||||
cfg = gconfig.merge_cli_overrides(
|
router_overrides: dict = {}
|
||||||
gconfig.DEFAULT_CONFIG,
|
if router_enabled is not None:
|
||||||
None,
|
router_overrides["enabled"] = router_enabled
|
||||||
{
|
if router_type is not None:
|
||||||
"conditioning": {
|
router_overrides["type"] = router_type
|
||||||
"particle": {"type": particle_conditioning},
|
if n_experts is not None:
|
||||||
"material": {"type": material_conditioning},
|
router_overrides["n_experts"] = n_experts
|
||||||
},
|
if router_overrides:
|
||||||
"stage1_model": {"router": router_cfg},
|
overrides["stage1_model"] = {"router": router_overrides}
|
||||||
"stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
|
|
||||||
},
|
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(
|
run_setup_stage(
|
||||||
Path(data),
|
Path(data),
|
||||||
val_fraction=val_fraction,
|
val_fraction=cfg["train"]["val_fraction"],
|
||||||
seed=seed,
|
seed=cfg["train"]["seed"],
|
||||||
cfg=cfg,
|
cfg=cfg,
|
||||||
cache_setup=True,
|
cache_setup=True,
|
||||||
rebuild_setup_cache=rebuild,
|
rebuild_setup_cache=rebuild,
|
||||||
|
|||||||
@@ -124,6 +124,11 @@ def test_warm_cache_router_process_warms_proc_map(tmp_path):
|
|||||||
[
|
[
|
||||||
"warm-cache",
|
"warm-cache",
|
||||||
str(data),
|
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",
|
||||||
"--router-type",
|
"--router-type",
|
||||||
"process",
|
"process",
|
||||||
@@ -167,3 +172,63 @@ def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path):
|
|||||||
assert loaded is not None
|
assert loaded is not None
|
||||||
assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
|
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
|
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 <same file>` 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
|
||||||
|
|||||||
Reference in New Issue
Block a user