Extract giant train/new-run's CLI override mapping into a table-driven function (issues.md Issues 3 & 4)
train()'s ~140-line hand-written flag->config translation (three different ad hoc "more specific flag wins" patterns) and new_run()'s near-verbatim copy are replaced by a shared FlagSpec/FLAG_SPECS table and overrides_from_flags() in config.py, reused by both commands. This makes the override/precedence logic directly unit-testable without CliRunner, closing coverage gaps that had zero tests (e.g. --emb-dim/--conditioning dual-axis fan-out, three of four WGAN knob legs, --stage2-generator overriding --mode, router's stage1-only asymmetry). No CLI flags, help text, or precedence semantics changed -- `giant train --help`/`giant new-run --help` are byte-identical before and after, and all previously-passing CliRunner tests still pass unmodified. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
+74
-201
@@ -653,141 +653,57 @@ def train(
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
cli_train = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"epochs": epochs,
|
||||
"batch_size": batch_size_value,
|
||||
"lr": lr,
|
||||
"weight_decay": weight_decay,
|
||||
"ema_decay": ema_decay,
|
||||
"warmup_epochs": warmup_epochs,
|
||||
"val_fraction": val_fraction,
|
||||
"num_workers": num_workers,
|
||||
"seed": seed,
|
||||
"validate_every": validate_every,
|
||||
"validate_steps": validate_steps,
|
||||
"max_val_batches": max_val_batches,
|
||||
"wandb": wandb,
|
||||
"wandb_project": wandb_project,
|
||||
"wandb_run_name": wandb_run_name,
|
||||
"wandb_log_every": wandb_log_every,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
# --hidden-dim/--n-blocks/--dropout are stage-1-only shorthands kept for
|
||||
# backward compatibility (they predate stage2_model having its own
|
||||
# flags); --stage1-*/--stage2-* below are the explicit, discoverable
|
||||
# per-stage flags, and take precedence when both are given.
|
||||
cli_stage1_model: dict[str, object] = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"hidden_dim": hidden_dim,
|
||||
"n_res_blocks": n_blocks,
|
||||
"dropout": dropout,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
cli_stage1_model.update(
|
||||
{
|
||||
k: v
|
||||
for k, v in {
|
||||
"hidden_dim": stage1_hidden_dim,
|
||||
"n_res_blocks": stage1_n_res_blocks,
|
||||
"dropout": stage1_dropout,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
)
|
||||
cli_stage2_model: dict[str, object] = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"hidden_dim": stage2_hidden_dim,
|
||||
"n_res_blocks": stage2_n_res_blocks,
|
||||
"dropout": stage2_dropout,
|
||||
"decoder": stage2_decoder.value if stage2_decoder is not None else None,
|
||||
"k_max": stage2_k_max,
|
||||
"context_dim": stage2_context_dim,
|
||||
"stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
|
||||
if cli_router:
|
||||
cli_stage1_model["router"] = cli_router
|
||||
|
||||
# --emb-dim/--conditioning set both conditioning axes (v0.2 had one
|
||||
# shared value for particle+material).
|
||||
cli_conditioning: dict[str, dict] = {}
|
||||
if emb_dim is not None:
|
||||
cli_conditioning["particle"] = {"emb_dim": emb_dim}
|
||||
cli_conditioning["material"] = {"emb_dim": emb_dim}
|
||||
if conditioning is not None:
|
||||
cli_conditioning.setdefault("particle", {})["type"] = conditioning.value
|
||||
cli_conditioning.setdefault("material", {})["type"] = conditioning.value
|
||||
|
||||
overrides: dict[str, dict] = {}
|
||||
if cli_train:
|
||||
overrides["train"] = cli_train
|
||||
if cli_stage1_model:
|
||||
overrides["stage1_model"] = cli_stage1_model
|
||||
if cli_stage2_model:
|
||||
overrides["stage2_model"] = cli_stage2_model
|
||||
if cli_conditioning:
|
||||
overrides["conditioning"] = cli_conditioning
|
||||
|
||||
# --mode/--n-critic/--gp-weight/--critic-lr/--noise-dim apply to BOTH
|
||||
# stages by default (v0.2 had one global mode/wgan config shared by both
|
||||
# — see giant.config.migrate_config's train.mode /
|
||||
# train.{n_critic,gp_weight,critic_lr} precedent); the --stage1-*/
|
||||
# --stage2-* variants below override a single stage independently (decision
|
||||
# 7), which is what actually enables e.g. `--stage1-generator flow
|
||||
# --stage2-generator wgan`.
|
||||
if mode is not None:
|
||||
overrides.setdefault("stage1_model", {})["generator"] = mode.value
|
||||
overrides.setdefault("stage2_model", {})["generator"] = mode.value
|
||||
if stage1_generator is not None:
|
||||
overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
|
||||
if stage2_generator is not None:
|
||||
overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
|
||||
|
||||
shared_wgan_overrides = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"n_critic": n_critic,
|
||||
"gp_weight": gp_weight,
|
||||
"noise_dim": noise_dim,
|
||||
"critic_lr": critic_lr,
|
||||
}.items()
|
||||
if v is not None
|
||||
flag_values: dict[str, object] = {
|
||||
"epochs": epochs,
|
||||
"batch_size": batch_size_value,
|
||||
"lr": lr,
|
||||
"weight_decay": weight_decay,
|
||||
"ema_decay": ema_decay,
|
||||
"warmup_epochs": warmup_epochs,
|
||||
"val_fraction": val_fraction,
|
||||
"num_workers": num_workers,
|
||||
"seed": seed,
|
||||
"validate_every": validate_every,
|
||||
"validate_steps": validate_steps,
|
||||
"max_val_batches": max_val_batches,
|
||||
"wandb": wandb,
|
||||
"wandb_project": wandb_project,
|
||||
"wandb_run_name": wandb_run_name,
|
||||
"wandb_log_every": wandb_log_every,
|
||||
"hidden_dim": hidden_dim,
|
||||
"n_blocks": n_blocks,
|
||||
"dropout": dropout,
|
||||
"stage1_hidden_dim": stage1_hidden_dim,
|
||||
"stage1_n_res_blocks": stage1_n_res_blocks,
|
||||
"stage1_dropout": stage1_dropout,
|
||||
"stage2_hidden_dim": stage2_hidden_dim,
|
||||
"stage2_n_res_blocks": stage2_n_res_blocks,
|
||||
"stage2_dropout": stage2_dropout,
|
||||
"stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None,
|
||||
"stage2_k_max": stage2_k_max,
|
||||
"stage2_context_dim": stage2_context_dim,
|
||||
"stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
|
||||
"mode": mode.value if mode is not None else None,
|
||||
"stage1_generator": stage1_generator.value if stage1_generator is not None else None,
|
||||
"stage2_generator": stage2_generator.value if stage2_generator is not None else None,
|
||||
"conditioning": conditioning.value if conditioning is not None else None,
|
||||
"emb_dim": emb_dim,
|
||||
"router_config": cli_router or None,
|
||||
"n_critic": n_critic,
|
||||
"gp_weight": gp_weight,
|
||||
"noise_dim": noise_dim,
|
||||
"critic_lr": critic_lr,
|
||||
"stage1_n_critic": stage1_n_critic,
|
||||
"stage1_gp_weight": stage1_gp_weight,
|
||||
"stage1_noise_dim": stage1_noise_dim,
|
||||
"stage1_critic_lr": stage1_critic_lr,
|
||||
"stage2_n_critic": stage2_n_critic,
|
||||
"stage2_gp_weight": stage2_gp_weight,
|
||||
"stage2_noise_dim": stage2_noise_dim,
|
||||
"stage2_critic_lr": stage2_critic_lr,
|
||||
}
|
||||
stage1_wgan_overrides = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"n_critic": stage1_n_critic,
|
||||
"gp_weight": stage1_gp_weight,
|
||||
"noise_dim": stage1_noise_dim,
|
||||
"critic_lr": stage1_critic_lr,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
stage2_wgan_overrides = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"n_critic": stage2_n_critic,
|
||||
"gp_weight": stage2_gp_weight,
|
||||
"noise_dim": stage2_noise_dim,
|
||||
"critic_lr": stage2_critic_lr,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
for stage_name, stage_specific in (
|
||||
("stage1_model", stage1_wgan_overrides),
|
||||
("stage2_model", stage2_wgan_overrides),
|
||||
):
|
||||
stage_wgan = {**shared_wgan_overrides, **stage_specific}
|
||||
if stage_wgan:
|
||||
overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update(stage_wgan)
|
||||
overrides = gconfig.overrides_from_flags(flag_values)
|
||||
|
||||
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
|
||||
gconfig.validate_config(cfg)
|
||||
@@ -912,75 +828,32 @@ def new_run(
|
||||
(with the full dataset-derived meta section), so this scaffold's meta
|
||||
section is just a placeholder recording what was asked for and when.
|
||||
"""
|
||||
cli_train = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"epochs": epochs,
|
||||
"batch_size": batch_size,
|
||||
"lr": lr,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
cli_stage1_model: dict[str, object] = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"hidden_dim": hidden_dim,
|
||||
"n_res_blocks": n_blocks,
|
||||
"dropout": dropout,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
cli_stage1_model.update(
|
||||
{
|
||||
k: v
|
||||
for k, v in {
|
||||
"hidden_dim": stage1_hidden_dim,
|
||||
"n_res_blocks": stage1_n_res_blocks,
|
||||
"dropout": stage1_dropout,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
)
|
||||
cli_stage2_model: dict[str, object] = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"hidden_dim": stage2_hidden_dim,
|
||||
"n_res_blocks": stage2_n_res_blocks,
|
||||
"dropout": stage2_dropout,
|
||||
"decoder": stage2_decoder.value if stage2_decoder is not None else None,
|
||||
"k_max": stage2_k_max,
|
||||
"context_dim": stage2_context_dim,
|
||||
"stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
|
||||
if cli_router:
|
||||
cli_stage1_model["router"] = cli_router
|
||||
cli_conditioning: dict[str, dict] = {}
|
||||
if emb_dim is not None:
|
||||
cli_conditioning["particle"] = {"emb_dim": emb_dim}
|
||||
cli_conditioning["material"] = {"emb_dim": emb_dim}
|
||||
if conditioning is not None:
|
||||
cli_conditioning.setdefault("particle", {})["type"] = conditioning.value
|
||||
cli_conditioning.setdefault("material", {})["type"] = conditioning.value
|
||||
|
||||
overrides: dict[str, dict] = {}
|
||||
if cli_train:
|
||||
overrides["train"] = cli_train
|
||||
if cli_stage1_model:
|
||||
overrides["stage1_model"] = cli_stage1_model
|
||||
if cli_stage2_model:
|
||||
overrides["stage2_model"] = cli_stage2_model
|
||||
if cli_conditioning:
|
||||
overrides["conditioning"] = cli_conditioning
|
||||
if mode is not None:
|
||||
overrides.setdefault("stage1_model", {})["generator"] = mode.value
|
||||
overrides.setdefault("stage2_model", {})["generator"] = mode.value
|
||||
if stage1_generator is not None:
|
||||
overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
|
||||
if stage2_generator is not None:
|
||||
overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
|
||||
flag_values: dict[str, object] = {
|
||||
"epochs": epochs,
|
||||
"batch_size": batch_size,
|
||||
"lr": lr,
|
||||
"hidden_dim": hidden_dim,
|
||||
"n_blocks": n_blocks,
|
||||
"dropout": dropout,
|
||||
"stage1_hidden_dim": stage1_hidden_dim,
|
||||
"stage1_n_res_blocks": stage1_n_res_blocks,
|
||||
"stage1_dropout": stage1_dropout,
|
||||
"stage2_hidden_dim": stage2_hidden_dim,
|
||||
"stage2_n_res_blocks": stage2_n_res_blocks,
|
||||
"stage2_dropout": stage2_dropout,
|
||||
"stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None,
|
||||
"stage2_k_max": stage2_k_max,
|
||||
"stage2_context_dim": stage2_context_dim,
|
||||
"stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
|
||||
"mode": mode.value if mode is not None else None,
|
||||
"stage1_generator": stage1_generator.value if stage1_generator is not None else None,
|
||||
"stage2_generator": stage2_generator.value if stage2_generator is not None else None,
|
||||
"conditioning": conditioning.value if conditioning is not None else None,
|
||||
"emb_dim": emb_dim,
|
||||
"router_config": cli_router or None,
|
||||
}
|
||||
overrides = gconfig.overrides_from_flags(flag_values)
|
||||
|
||||
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
|
||||
gconfig.validate_config(cfg)
|
||||
|
||||
+108
@@ -893,6 +893,114 @@ def _deep_merge(base: dict, override: dict) -> dict:
|
||||
return result
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FlagSpec:
|
||||
"""One CLI flag's mapping into the config-overrides tree.
|
||||
|
||||
`paths` lists every dotted config path this flag writes (>1 means fan-out
|
||||
to multiple stages/axes, e.g. `--mode` -> both stages' `generator`).
|
||||
`precedence` controls write order when two flags target the same path:
|
||||
specs are applied in ascending precedence, so a higher-precedence (more
|
||||
specific) flag overwrites a lower-precedence (shared/shorthand) one —
|
||||
this is the "build a shared dict, then let a more specific dict win"
|
||||
pattern `giant train`/`giant new-run` need (e.g. `--hidden-dim` vs
|
||||
`--stage1-hidden-dim`, or `--n-critic` vs `--stage1-n-critic`),
|
||||
generalized to one mechanism instead of three different ad hoc ones.
|
||||
"""
|
||||
|
||||
name: str
|
||||
paths: tuple[str, ...]
|
||||
precedence: int = 0
|
||||
|
||||
|
||||
# Flag -> config-path table shared by `giant train`/`giant new-run`
|
||||
# (giant/cli.py) so both commands resolve CLI overrides identically. See
|
||||
# issues.md Issue 3: this replaces ~140 lines of hand-written, imperative
|
||||
# dict-building in cli.py with one declarative table plus
|
||||
# `overrides_from_flags` below.
|
||||
FLAG_SPECS: tuple[FlagSpec, ...] = (
|
||||
# train block -- flat pass-through, unique paths, precedence irrelevant.
|
||||
FlagSpec("epochs", ("train.epochs",)),
|
||||
FlagSpec("batch_size", ("train.batch_size",)),
|
||||
FlagSpec("lr", ("train.lr",)),
|
||||
FlagSpec("weight_decay", ("train.weight_decay",)),
|
||||
FlagSpec("ema_decay", ("train.ema_decay",)),
|
||||
FlagSpec("warmup_epochs", ("train.warmup_epochs",)),
|
||||
FlagSpec("val_fraction", ("train.val_fraction",)),
|
||||
FlagSpec("num_workers", ("train.num_workers",)),
|
||||
FlagSpec("seed", ("train.seed",)),
|
||||
FlagSpec("validate_every", ("train.validate_every",)),
|
||||
FlagSpec("validate_steps", ("train.validate_steps",)),
|
||||
FlagSpec("max_val_batches", ("train.max_val_batches",)),
|
||||
FlagSpec("wandb", ("train.wandb",)),
|
||||
FlagSpec("wandb_project", ("train.wandb_project",)),
|
||||
FlagSpec("wandb_run_name", ("train.wandb_run_name",)),
|
||||
FlagSpec("wandb_log_every", ("train.wandb_log_every",)),
|
||||
# --hidden-dim/--n-blocks/--dropout are stage-1-only backward-compat
|
||||
# shorthands (they predate stage2_model having its own flags);
|
||||
# --stage1-* wins when both are given.
|
||||
FlagSpec("hidden_dim", ("stage1_model.hidden_dim",), precedence=0),
|
||||
FlagSpec("stage1_hidden_dim", ("stage1_model.hidden_dim",), precedence=1),
|
||||
FlagSpec("n_blocks", ("stage1_model.n_res_blocks",), precedence=0),
|
||||
FlagSpec("stage1_n_res_blocks", ("stage1_model.n_res_blocks",), precedence=1),
|
||||
FlagSpec("dropout", ("stage1_model.dropout",), precedence=0),
|
||||
FlagSpec("stage1_dropout", ("stage1_model.dropout",), precedence=1),
|
||||
# stage2-only knobs.
|
||||
FlagSpec("stage2_hidden_dim", ("stage2_model.hidden_dim",)),
|
||||
FlagSpec("stage2_n_res_blocks", ("stage2_model.n_res_blocks",)),
|
||||
FlagSpec("stage2_dropout", ("stage2_model.dropout",)),
|
||||
FlagSpec("stage2_decoder", ("stage2_model.decoder",)),
|
||||
FlagSpec("stage2_k_max", ("stage2_model.k_max",)),
|
||||
FlagSpec("stage2_context_dim", ("stage2_model.context_dim",)),
|
||||
FlagSpec("stage2_stage1_context", ("stage2_model.stage1_context",)),
|
||||
# --mode applies to both stages by default (v0.2 had one shared
|
||||
# mode/wgan config); --stage{1,2}-generator override a single stage.
|
||||
FlagSpec("mode", ("stage1_model.generator", "stage2_model.generator"), precedence=0),
|
||||
FlagSpec("stage1_generator", ("stage1_model.generator",), precedence=1),
|
||||
FlagSpec("stage2_generator", ("stage2_model.generator",), precedence=1),
|
||||
# --emb-dim/--conditioning set both conditioning axes (v0.2 had one
|
||||
# shared value for particle+material).
|
||||
FlagSpec("conditioning", ("conditioning.particle.type", "conditioning.material.type")),
|
||||
FlagSpec("emb_dim", ("conditioning.particle.emb_dim", "conditioning.material.emb_dim")),
|
||||
# Pre-aggregated router override dict (built by `_router_cli_overrides`
|
||||
# in cli.py from --router/--router-type/--n-experts/--router-axis).
|
||||
# Router overrides only ever land on stage1_model -- this asymmetry is
|
||||
# deliberate (see cli.py) and must not be "fixed" into a fan-out here.
|
||||
FlagSpec("router_config", ("stage1_model.router",)),
|
||||
# WGAN: shared knobs apply to both stages by default (v0.2 had one
|
||||
# shared wgan config); --stage{1,2}-* override a single stage.
|
||||
FlagSpec("n_critic", ("stage1_model.wgan.n_critic", "stage2_model.wgan.n_critic"), precedence=0),
|
||||
FlagSpec("stage1_n_critic", ("stage1_model.wgan.n_critic",), precedence=1),
|
||||
FlagSpec("stage2_n_critic", ("stage2_model.wgan.n_critic",), precedence=1),
|
||||
FlagSpec("gp_weight", ("stage1_model.wgan.gp_weight", "stage2_model.wgan.gp_weight"), precedence=0),
|
||||
FlagSpec("stage1_gp_weight", ("stage1_model.wgan.gp_weight",), precedence=1),
|
||||
FlagSpec("stage2_gp_weight", ("stage2_model.wgan.gp_weight",), precedence=1),
|
||||
FlagSpec("noise_dim", ("stage1_model.wgan.noise_dim", "stage2_model.wgan.noise_dim"), precedence=0),
|
||||
FlagSpec("stage1_noise_dim", ("stage1_model.wgan.noise_dim",), precedence=1),
|
||||
FlagSpec("stage2_noise_dim", ("stage2_model.wgan.noise_dim",), precedence=1),
|
||||
FlagSpec("critic_lr", ("stage1_model.wgan.critic_lr", "stage2_model.wgan.critic_lr"), precedence=0),
|
||||
FlagSpec("stage1_critic_lr", ("stage1_model.wgan.critic_lr",), precedence=1),
|
||||
FlagSpec("stage2_critic_lr", ("stage2_model.wgan.critic_lr",), precedence=1),
|
||||
)
|
||||
|
||||
|
||||
def overrides_from_flags(values: dict[str, object]) -> dict:
|
||||
"""Build the nested, section-keyed config-overrides dict
|
||||
`merge_cli_overrides` expects, from `{flag_name: value}`.
|
||||
|
||||
Flags absent from `values`, or mapped to `None` (= not given on the
|
||||
CLI), are skipped. See `FlagSpec`/`FLAG_SPECS` above for the precedence
|
||||
rule applied when two flags target the same path.
|
||||
"""
|
||||
overrides: dict = {}
|
||||
for spec in sorted(FLAG_SPECS, key=lambda s: s.precedence):
|
||||
if spec.name not in values or values[spec.name] is None:
|
||||
continue
|
||||
for path in spec.paths:
|
||||
_set_path(overrides, path, values[spec.name])
|
||||
return overrides
|
||||
|
||||
|
||||
# v0.2 [train] keys that pass through to v0.3 [train] unchanged (same name,
|
||||
# same meaning) when present in the loaded file — everything model-shaped
|
||||
# moved to the stage/conditioning blocks instead (see the rest of
|
||||
|
||||
@@ -94,3 +94,83 @@ def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
|
||||
assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
|
||||
assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
|
||||
assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5
|
||||
|
||||
|
||||
def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
|
||||
result = runner.invoke(
|
||||
cli.app,
|
||||
["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
|
||||
)
|
||||
assert result.exit_code == 1
|
||||
assert "--batch-size must be an integer or 'auto'" in result.output
|
||||
|
||||
|
||||
def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_path):
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
|
||||
captured["out_dir"] = out_dir
|
||||
|
||||
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
||||
|
||||
resume_dir = tmp_path / "resumed_run"
|
||||
resume_dir.mkdir()
|
||||
(resume_dir / "last.pt").touch()
|
||||
explicit_out = tmp_path / "explicit_run"
|
||||
|
||||
result = runner.invoke(
|
||||
cli.app,
|
||||
["train", "dummy.parquet", "--out", str(explicit_out), "--resume", str(resume_dir / "last.pt")],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["out_dir"] == explicit_out
|
||||
|
||||
|
||||
def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
|
||||
captured["out_dir"] = out_dir
|
||||
|
||||
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
||||
|
||||
resume_dir = tmp_path / "resumed_run"
|
||||
resume_dir.mkdir()
|
||||
(resume_dir / "last.pt").touch()
|
||||
|
||||
result = runner.invoke(cli.app, ["train", "dummy.parquet", "--resume", str(resume_dir / "last.pt")])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["out_dir"] == resume_dir
|
||||
|
||||
|
||||
def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypatch, tmp_path):
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
|
||||
captured["out_dir"] = out_dir
|
||||
|
||||
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
result = runner.invoke(cli.app, ["train", "dummy.parquet"])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["out_dir"] == Path("checkpoints") / cli.gconfig.default_out_dir_name(cli.gconfig.DEFAULT_CONFIG)
|
||||
|
||||
|
||||
def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
|
||||
captured["batch_size"] = cfg["train"]["batch_size"]
|
||||
|
||||
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
||||
monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)
|
||||
|
||||
result = runner.invoke(
|
||||
cli.app,
|
||||
["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "auto"],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["batch_size"] == 123
|
||||
assert "batch_size: 123 (auto-estimated from free GPU memory)" in result.output
|
||||
|
||||
@@ -889,6 +889,125 @@ def test_merge_cli_overrides_real_config_fixtures_pass_key_validation(fixture_na
|
||||
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / fixture_name, {}) # must not raise
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# overrides_from_flags (issues.md Issue 3): the flag -> config-path table
|
||||
# shared by `giant train`/`giant new-run`. Each test below pins one
|
||||
# precedence rule directly, without CliRunner — see also
|
||||
# tests/test_cli_train_overrides.py for the thin end-to-end smoke coverage.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_overrides_from_flags_empty_values_yield_empty_overrides():
|
||||
assert gconfig.overrides_from_flags({}) == {}
|
||||
assert gconfig.overrides_from_flags({"epochs": None, "hidden_dim": None}) == {}
|
||||
|
||||
|
||||
def test_overrides_from_flags_train_block_passthrough():
|
||||
overrides = gconfig.overrides_from_flags({"epochs": 5, "lr": 1e-3, "hidden_dim": None})
|
||||
assert overrides == {"train": {"epochs": 5, "lr": 1e-3}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("shorthand", "explicit", "path_key"),
|
||||
[
|
||||
("hidden_dim", "stage1_hidden_dim", "hidden_dim"),
|
||||
("n_blocks", "stage1_n_res_blocks", "n_res_blocks"),
|
||||
("dropout", "stage1_dropout", "dropout"),
|
||||
],
|
||||
)
|
||||
def test_overrides_from_flags_stage1_explicit_overrides_shorthand(shorthand, explicit, path_key):
|
||||
overrides = gconfig.overrides_from_flags({shorthand: 1, explicit: 2})
|
||||
assert overrides["stage1_model"][path_key] == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("shorthand", "path_key"),
|
||||
[("hidden_dim", "hidden_dim"), ("n_blocks", "n_res_blocks"), ("dropout", "dropout")],
|
||||
)
|
||||
def test_overrides_from_flags_stage1_shorthand_alone(shorthand, path_key):
|
||||
overrides = gconfig.overrides_from_flags({shorthand: 7})
|
||||
assert overrides["stage1_model"][path_key] == 7
|
||||
|
||||
|
||||
def test_overrides_from_flags_stage2_only_knobs():
|
||||
overrides = gconfig.overrides_from_flags(
|
||||
{
|
||||
"stage2_hidden_dim": 32,
|
||||
"stage2_n_res_blocks": 4,
|
||||
"stage2_dropout": 0.1,
|
||||
"stage2_decoder": "one_shot",
|
||||
"stage2_k_max": 8,
|
||||
"stage2_context_dim": 16,
|
||||
"stage2_stage1_context": "sampled",
|
||||
}
|
||||
)
|
||||
assert overrides["stage2_model"] == {
|
||||
"hidden_dim": 32,
|
||||
"n_res_blocks": 4,
|
||||
"dropout": 0.1,
|
||||
"decoder": "one_shot",
|
||||
"k_max": 8,
|
||||
"context_dim": 16,
|
||||
"stage1_context": "sampled",
|
||||
}
|
||||
assert "stage1_model" not in overrides
|
||||
|
||||
|
||||
def test_overrides_from_flags_mode_fans_to_both_stages():
|
||||
overrides = gconfig.overrides_from_flags({"mode": "wgan"})
|
||||
assert overrides["stage1_model"]["generator"] == "wgan"
|
||||
assert overrides["stage2_model"]["generator"] == "wgan"
|
||||
|
||||
|
||||
def test_overrides_from_flags_stage1_generator_overrides_mode_for_stage1_only():
|
||||
overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage1_generator": "flow"})
|
||||
assert overrides["stage1_model"]["generator"] == "flow"
|
||||
assert overrides["stage2_model"]["generator"] == "wgan"
|
||||
|
||||
|
||||
def test_overrides_from_flags_stage2_generator_overrides_mode_for_stage2_only():
|
||||
overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage2_generator": "flow"})
|
||||
assert overrides["stage1_model"]["generator"] == "wgan"
|
||||
assert overrides["stage2_model"]["generator"] == "flow"
|
||||
|
||||
|
||||
def test_overrides_from_flags_emb_dim_sets_both_conditioning_axes():
|
||||
overrides = gconfig.overrides_from_flags({"emb_dim": 24})
|
||||
assert overrides["conditioning"]["particle"]["emb_dim"] == 24
|
||||
assert overrides["conditioning"]["material"]["emb_dim"] == 24
|
||||
|
||||
|
||||
def test_overrides_from_flags_conditioning_sets_both_axes_type():
|
||||
overrides = gconfig.overrides_from_flags({"conditioning": "onehot"})
|
||||
assert overrides["conditioning"]["particle"]["type"] == "onehot"
|
||||
assert overrides["conditioning"]["material"]["type"] == "onehot"
|
||||
|
||||
|
||||
def test_overrides_from_flags_router_config_only_touches_stage1():
|
||||
overrides = gconfig.overrides_from_flags({"router_config": {"enabled": True, "type": "energy"}})
|
||||
assert overrides["stage1_model"]["router"] == {"enabled": True, "type": "energy"}
|
||||
assert "stage2_model" not in overrides
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("shared", "stage1_specific", "stage2_specific", "path_key"),
|
||||
[
|
||||
("n_critic", "stage1_n_critic", "stage2_n_critic", "n_critic"),
|
||||
("gp_weight", "stage1_gp_weight", "stage2_gp_weight", "gp_weight"),
|
||||
("noise_dim", "stage1_noise_dim", "stage2_noise_dim", "noise_dim"),
|
||||
("critic_lr", "stage1_critic_lr", "stage2_critic_lr", "critic_lr"),
|
||||
],
|
||||
)
|
||||
def test_overrides_from_flags_wgan_knobs_split_per_stage(shared, stage1_specific, stage2_specific, path_key):
|
||||
overrides = gconfig.overrides_from_flags({shared: 5.0, stage1_specific: 3.0})
|
||||
assert overrides["stage1_model"]["wgan"][path_key] == 3.0
|
||||
assert overrides["stage2_model"]["wgan"][path_key] == 5.0
|
||||
|
||||
overrides = gconfig.overrides_from_flags({shared: 5.0, stage2_specific: 2.5})
|
||||
assert overrides["stage1_model"]["wgan"][path_key] == 5.0
|
||||
assert overrides["stage2_model"]["wgan"][path_key] == 2.5
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# checkpoint config-mismatch warnings (unchanged surface, still exercised)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user