From 2bfb1ab0564c2e3b8ba8de617ad99cfe9d72c02f Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 12 Aug 2026 15:04:49 +0200 Subject: [PATCH] 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 --- giant/cli.py | 275 ++++++++---------------------- giant/config.py | 108 ++++++++++++ tests/test_cli_train_overrides.py | 80 +++++++++ tests/test_config.py | 119 +++++++++++++ 4 files changed, 381 insertions(+), 201 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index ee49c1d..533034c 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -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) diff --git a/giant/config.py b/giant/config.py index 59ec1e0..c7d10ec 100644 --- a/giant/config.py +++ b/giant/config.py @@ -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 diff --git a/tests/test_cli_train_overrides.py b/tests/test_cli_train_overrides.py index 486ea4b..7df258f 100644 --- a/tests/test_cli_train_overrides.py +++ b/tests/test_cli_train_overrides.py @@ -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 diff --git a/tests/test_config.py b/tests/test_config.py index ceb5e73..58d4d1c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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) # ---------------------------------------------------------------------------