lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
2bfb1ab0564c2e3b8ba8de617ad99cfe9d72c02f
giant/tests/test_config.py
T
lars 2bfb1ab056 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>
2026-08-12 15:04:49 +02:00

1047 lines
38 KiB
Python
Raw Blame History

from datetime import datetime
from pathlib import Path
import pytest
from giant import config as gconfig
_CONFIGS_DIR = Path(__file__).resolve().parents[1] / "configs"
def _write_toml(path, git_hash=None, extra=""):
meta = f'\n[meta]\ngit_hash = "{git_hash}"\n' if git_hash is not None else ""
path.write_text(extra + meta)
# ---------------------------------------------------------------------------
# Conditioning enum
# ---------------------------------------------------------------------------
def test_conditioning_enum_has_onehot():
assert gconfig.Conditioning.onehot == "onehot"
assert {c.value for c in gconfig.Conditioning} == {
"physical",
"embedding",
"onehot",
}
# ---------------------------------------------------------------------------
# Config dataclasses (issues.md Issue 1)
# ---------------------------------------------------------------------------
def test_giant_config_to_dict_matches_default_config():
"""DEFAULT_CONFIG is generated from GiantConfig().to_dict() (not
hand-maintained), so the two cannot structurally drift apart — but this
pins the *equality* too, catching e.g. a stray in-place mutation of
DEFAULT_CONFIG added elsewhere after import."""
assert gconfig.GiantConfig().to_dict() == gconfig.DEFAULT_CONFIG
@pytest.mark.parametrize(
"cls",
[
gconfig.ConditioningAxisConfig,
gconfig.ConditioningConfig,
gconfig.FlowConfig,
gconfig.DdpmConfig,
gconfig.Stage1WganConfig,
gconfig.Stage2WganConfig,
gconfig.RouterConfig,
gconfig.Stage2RouterConfig,
gconfig.NSecConfig,
gconfig.ParticleTypeConfig,
gconfig.AutoregressiveConfig,
gconfig.Stage1ModelConfig,
gconfig.Stage2ModelConfig,
gconfig.TrainConfig,
gconfig.GiantConfig,
],
)
def test_config_dataclass_from_dict_round_trips_through_to_dict(cls):
assert cls.from_dict(cls().to_dict()) == cls()
assert cls.from_dict(None) == cls()
def test_stage2_model_config_defaults_match_documented_v030_intent():
"""The two keys issues.md Issue 1 found drifted between DEFAULT_CONFIG
and build_models/StageSpec.from_config's own .get(key, default)
fallbacks — pinned directly against the dataclass that is now their
shared single source of truth."""
spec = gconfig.Stage2ModelConfig()
assert spec.decoder == "autoregressive"
assert spec.particle_type.target == "onehot"
def test_router_config_extra_round_trips_composed_axis_keys():
d = {"enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4}
router = gconfig.RouterConfig.from_dict(d)
assert router.enabled is True
assert router.extra == {"axis0_type": "energy", "axis0_n_experts": 4}
assert router.to_dict()["axis0_type"] == "energy"
def test_stage2_router_config_tie_to_stage1_not_leaked_into_extra():
router = gconfig.Stage2RouterConfig.from_dict({"tie_to_stage1": True})
assert router.tie_to_stage1 is True
assert "tie_to_stage1" not in router.extra
def test_stage1_router_config_has_no_tie_to_stage1_key():
"""Stage 1's router schema must not gain stage 2's tie_to_stage1 key —
that would change every future run's saved config.toml shape."""
assert "tie_to_stage1" not in gconfig.RouterConfig().to_dict()
def test_n_sec_config_extra_round_trips_legacy_owner():
n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "legacy_owner": "stage1"})
assert n_sec.legacy_owner == "stage1"
assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "legacy_owner": "stage1"}
# ---------------------------------------------------------------------------
# _deep_merge
# ---------------------------------------------------------------------------
def test_deep_merge_leaf_override_keeps_untouched_siblings():
base = {"a": 1, "b": {"c": 2, "d": 3}}
result = gconfig._deep_merge(base, {"b": {"c": 99}})
assert result == {"a": 1, "b": {"c": 99, "d": 3}}
def test_deep_merge_recurses_at_multiple_levels():
base = {
"stage1_model": {
"hidden_dim": 256,
"router": {"enabled": False, "type": "energy", "n_experts": 4},
}
}
result = gconfig._deep_merge(base, {"stage1_model": {"router": {"enabled": True}}})
assert result["stage1_model"]["hidden_dim"] == 256
assert result["stage1_model"]["router"] == {
"enabled": True,
"type": "energy",
"n_experts": 4,
}
def test_deep_merge_does_not_mutate_base():
base = {"a": {"b": 1}}
gconfig._deep_merge(base, {"a": {"b": 2}})
assert base == {"a": {"b": 1}}
def test_deep_merge_non_dict_override_replaces_wholesale():
base = {"a": {"b": 1}}
result = gconfig._deep_merge(base, {"a": 5})
assert result == {"a": 5}
# ---------------------------------------------------------------------------
# migrate_config
# ---------------------------------------------------------------------------
def test_migrate_config_already_v3_returned_unchanged():
cfg = {"meta": {"config_version": 3}, "stage1_model": {"generator": "flow"}}
result = gconfig.migrate_config(cfg)
assert result == cfg
result["stage1_model"]["generator"] = "wgan"
assert cfg["stage1_model"]["generator"] == "flow" # deep-copied, not aliased
def test_migrate_config_empty_dict_still_injects_hardcoded_v02_facts():
# No [train]/[model] at all still counts as "v0.2" (config_version
# absent) — the hardcoded architectural facts fire unconditionally.
new = gconfig.migrate_config({})
assert "train" not in new
assert new["conditioning"]["out_dim"] == 128
assert new["conditioning"]["particle"]["n_layers"] == 2
assert new["conditioning"]["material"]["n_layers"] == 2
assert new["stage1_model"]["active"] is True
assert new["stage1_model"]["flow"]["time_dim"] == 64
assert new["stage1_model"]["ddpm"]["time_dim"] == 64
assert new["stage2_model"]["active"] is True
assert new["stage2_model"]["flow"]["time_dim"] == 64
assert new["stage2_model"]["ddpm"]["time_dim"] == 64
assert new["stage2_model"]["context_dim"] == 64
assert new["stage2_model"]["decoder"] == "one_shot"
assert new["stage2_model"]["particle_type"]["target"] == "physical"
assert new["meta"] == {"config_version": 3}
def test_migrate_config_mode_maps_to_both_stage_generators():
new = gconfig.migrate_config({"train": {"mode": "wgan"}})
assert new["stage1_model"]["generator"] == "wgan"
assert new["stage2_model"]["generator"] == "wgan"
def test_migrate_config_lambda_nsec_and_lambda_s2():
new = gconfig.migrate_config({"train": {"lambda_nsec": 0.2, "lambda_s2": 2.0}})
assert new["stage2_model"]["n_sec"]["lambda"] == 0.2
assert new["stage2_model"]["lambda"] == 2.0
def test_migrate_config_wgan_knobs_map_to_both_stages():
new = gconfig.migrate_config({"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}})
for stage in ("stage1_model", "stage2_model"):
assert new[stage]["wgan"]["n_critic"] == 3
assert new[stage]["wgan"]["gp_weight"] == 5.0
assert new[stage]["wgan"]["critic_lr"] == 1e-4
def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():
new = gconfig.migrate_config({"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}})
for stage in ("stage1_model", "stage2_model"):
assert new[stage]["hidden_dim"] == 128
assert new[stage]["n_res_blocks"] == 4
assert new[stage]["dropout"] == 0.2
def test_migrate_config_emb_dim_and_conditioning_map_to_both_axes():
new = gconfig.migrate_config({"model": {"emb_dim": 32, "conditioning": "embedding"}})
for axis in ("particle", "material"):
assert new["conditioning"][axis]["emb_dim"] == 32
assert new["conditioning"][axis]["type"] == "embedding"
def test_migrate_config_noise_dim_maps_to_both_stages_wgan():
new = gconfig.migrate_config({"model": {"noise_dim": 128}})
assert new["stage1_model"]["wgan"]["noise_dim"] == 128
assert new["stage2_model"]["wgan"]["noise_dim"] == 128
def test_migrate_config_k_max_maps_to_stage2_only():
new = gconfig.migrate_config({"model": {"k_max": 20}})
assert new["stage2_model"]["k_max"] == 20
assert "k_max" not in new.get("stage1_model", {})
def test_migrate_config_router_copied_to_both_stages_with_tie_to_stage1_false():
new = gconfig.migrate_config(
{
"model": {
"router": {
"enabled": True,
"type": "energy",
"n_experts": 10,
"temperature": 0.05,
}
}
}
)
assert new["stage1_model"]["router"] == {
"enabled": True,
"type": "energy",
"n_experts": 10,
"temperature": 0.05,
}
assert new["stage2_model"]["router"] == {
"enabled": True,
"type": "energy",
"n_experts": 10,
"temperature": 0.05,
"tie_to_stage1": False,
}
def test_migrate_config_router_nonzero_expert_dims_raises():
cfg = {"model": {"router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}}}
try:
gconfig.migrate_config(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "expert_hidden_dim" in str(e)
def test_migrate_config_router_zero_expert_dims_dropped_silently():
new = gconfig.migrate_config(
{
"model": {
"router": {
"enabled": True,
"type": "energy",
"n_experts": 4,
"expert_hidden_dim": 0,
"expert_n_blocks": 0,
}
}
}
)
assert "expert_hidden_dim" not in new["stage1_model"]["router"]
assert "expert_n_blocks" not in new["stage1_model"]["router"]
def test_migrate_config_train_passthrough_is_exact():
new = gconfig.migrate_config({"train": {"epochs": 7, "batch_size": 999, "seed": 3}})
assert new["train"] == {"epochs": 7, "batch_size": 999, "seed": 3}
def test_migrate_config_preserves_meta_git_hash():
new = gconfig.migrate_config({"meta": {"git_hash": "abc123"}})
assert new["meta"] == {"git_hash": "abc123", "config_version": 3}
def test_migrate_config_real_router_fixture_raises_on_nonzero_expert_dims():
cfg = gconfig.load_toml(_CONFIGS_DIR / "router_energy_n10_embedding.toml")
try:
gconfig.migrate_config(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "expert_hidden_dim" in str(e)
def test_migrate_config_real_wgan_fixture():
cfg = gconfig.load_toml(_CONFIGS_DIR / "wgan_h128_b4_physical.toml")
new = gconfig.migrate_config(cfg)
assert new["stage1_model"]["generator"] == "wgan"
assert new["stage2_model"]["generator"] == "wgan"
for stage in ("stage1_model", "stage2_model"):
assert new[stage]["hidden_dim"] == 128
assert new[stage]["n_res_blocks"] == 4
assert new[stage]["dropout"] == 0.0
assert new["conditioning"]["particle"]["type"] == "physical"
assert new["conditioning"]["material"]["type"] == "physical"
assert new["train"]["epochs"] == 30
assert new["train"]["warmup_epochs"] == 3
# ---------------------------------------------------------------------------
# merge_cli_overrides
# ---------------------------------------------------------------------------
def test_merge_cli_overrides_defaults_only_matches_default_config():
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {})
assert cfg == gconfig.DEFAULT_CONFIG
assert cfg is not gconfig.DEFAULT_CONFIG
def test_merge_cli_overrides_nested_override_keeps_siblings():
cfg = gconfig.merge_cli_overrides(
gconfig.DEFAULT_CONFIG,
None,
{"stage1_model": {"router": {"enabled": True}}},
)
assert cfg["stage1_model"]["router"]["enabled"] is True
assert cfg["stage1_model"]["router"]["type"] == "energy" # default preserved
assert cfg["stage1_model"]["hidden_dim"] == 256 # untouched sibling section
def test_merge_cli_overrides_file_then_explicit_override_precedence(tmp_path, monkeypatch):
monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
path = tmp_path / "config.toml"
_write_toml(
path,
git_hash="abc123",
extra="[train]\nepochs = 5\n\n[model]\nhidden_dim = 64\n",
)
cfg = gconfig.merge_cli_overrides(
gconfig.DEFAULT_CONFIG,
path,
{"stage1_model": {"hidden_dim": 128}},
)
assert cfg["train"]["epochs"] == 5 # from file
assert cfg["stage1_model"]["hidden_dim"] == 128 # explicit override wins over file
assert cfg["stage2_model"]["hidden_dim"] == 64 # migrated from file, not overridden
def test_merge_cli_overrides_migrates_v2_file_transparently(tmp_path, monkeypatch):
monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
path = tmp_path / "config.toml"
_write_toml(
path,
git_hash="abc123",
extra='[train]\nmode = "wgan"\n\n[model]\nconditioning = "embedding"\n',
)
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
assert cfg["stage1_model"]["generator"] == "wgan"
assert cfg["stage2_model"]["generator"] == "wgan"
assert cfg["conditioning"]["particle"]["type"] == "embedding"
# hardcoded v0.2 fact still applied even though it's not a CLI-settable key
assert cfg["conditioning"]["particle"]["n_layers"] == 2
def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
path = tmp_path / "config.toml"
_write_toml(path, git_hash="old111", extra="[train]\nepochs = 5\n")
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
captured = capsys.readouterr()
assert "warning" in captured.err
assert "old111" in captured.err
assert "current999" in captured.err
def test_merge_cli_overrides_no_warning_on_matching_git_hash(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
path = tmp_path / "config.toml"
_write_toml(path, git_hash="same123", extra="[train]\nepochs = 5\n")
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
assert capsys.readouterr().err == ""
def test_merge_cli_overrides_no_warning_when_git_hash_unknown(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown")
path = tmp_path / "config.toml"
_write_toml(path, git_hash="abc123", extra="[train]\nepochs = 5\n")
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
assert capsys.readouterr().err == ""
def test_merge_cli_overrides_no_warning_when_meta_section_absent(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
path = tmp_path / "config.toml"
path.write_text("[train]\nepochs = 5\n")
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
assert capsys.readouterr().err == ""
def test_merge_cli_overrides_real_default_toml_fixture(monkeypatch):
monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243")
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {})
assert cfg["stage1_model"]["generator"] == "flow"
assert cfg["stage1_model"]["hidden_dim"] == 256
assert cfg["stage2_model"]["hidden_dim"] == 256
assert cfg["conditioning"]["particle"]["emb_dim"] == 16
assert cfg["conditioning"]["particle"]["n_layers"] == 2 # migrated hardcoded fact
assert cfg["train"]["epochs"] == 100
assert cfg["stage2_model"]["decoder"] == "one_shot"
# ---------------------------------------------------------------------------
# save_config
# ---------------------------------------------------------------------------
def test_save_config_round_trips_multi_level_nesting(tmp_path):
cfg = {
"stage1_model": {
"hidden_dim": 256,
"router": {
"enabled": True,
"type": "energy",
"n_experts": 4,
},
},
"train": {"epochs": 100},
}
meta = {"config_version": 3, "git_hash": "abc123"}
gconfig.save_config(cfg, tmp_path, meta)
loaded = gconfig.load_toml(tmp_path / "config.toml")
assert loaded["stage1_model"]["hidden_dim"] == 256
assert loaded["stage1_model"]["router"] == {
"enabled": True,
"type": "energy",
"n_experts": 4,
}
assert loaded["train"] == {"epochs": 100}
assert loaded["meta"] == meta
def test_save_config_round_trips_three_level_nesting(tmp_path):
cfg = {
"stage2_model": {
"decoder": "autoregressive",
"n_sec": {"mode": "head", "lambda": 0.1},
"router": {"tie_to_stage1": True},
}
}
gconfig.save_config(cfg, tmp_path, {"config_version": 3})
loaded = gconfig.load_toml(tmp_path / "config.toml")
assert loaded["stage2_model"]["decoder"] == "autoregressive"
assert loaded["stage2_model"]["n_sec"] == {"mode": "head", "lambda": 0.1}
assert loaded["stage2_model"]["router"] == {"tie_to_stage1": True}
# ---------------------------------------------------------------------------
# default_out_dir_name
# ---------------------------------------------------------------------------
_NOW = datetime(2026, 7, 29, 14, 30)
def _cfg_with(**dotted_overrides):
"""Build a full DEFAULT_CONFIG-shaped dict with dotted-path overrides
applied via _deep_merge, e.g. _cfg_with(**{"stage1_model.hidden_dim": 512})."""
overrides: dict = {}
for dotted, value in dotted_overrides.items():
gconfig._set_path(overrides, dotted, value)
return gconfig._deep_merge(gconfig.DEFAULT_CONFIG, overrides)
def test_default_out_dir_name_all_defaults_is_just_the_timestamp():
assert gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW) == "20260729_1430"
def test_default_out_dir_name_stage1_generator_shown_bare_no_prefix():
cfg = _cfg_with(**{"stage1_model.generator": "wgan"})
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan"
def test_default_out_dir_name_stage2_decoder_shown():
cfg = _cfg_with(**{"stage2_model.decoder": "one_shot"})
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_dec-one_shot"
def test_default_out_dir_name_particle_type_target_shown():
cfg = _cfg_with(**{"stage2_model.particle_type.target": "physical"})
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_pt-physical"
def test_default_out_dir_name_particle_conditioning_embedding_abbreviated():
cfg = _cfg_with(**{"conditioning.particle.type": "embedding"})
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb"
def test_default_out_dir_name_stage1_router_shown_as_unit():
cfg = _cfg_with(
**{
"stage1_model.router.enabled": True,
"stage1_model.router.type": "energy",
"stage1_model.router.n_experts": 8,
}
)
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8"
def test_default_out_dir_name_stage2_router_shown_as_unit_distinct_from_stage1():
cfg = _cfg_with(
**{
"stage2_model.router.enabled": True,
"stage2_model.router.type": "pdg",
"stage2_model.router.n_experts": 3,
}
)
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s2r-pdg3"
def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault():
cfg = _cfg_with(
**{
"stage1_model.router.enabled": False,
"stage1_model.router.type": "pdg",
"stage1_model.router.n_experts": 8,
}
)
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
def test_default_out_dir_name_router_gumbel_shown_when_enabled():
cfg = _cfg_with(
**{
"stage1_model.router.enabled": True,
"stage1_model.router.type": "energy",
"stage1_model.router.n_experts": 8,
"stage1_model.router.gumbel": True,
}
)
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"
def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
cfg = _cfg_with(
**{
"stage1_model.generator": "wgan",
"stage2_model.generator": "flow",
"stage2_model.decoder": "one_shot",
"stage2_model.autoregressive.history": "attention",
"stage2_model.particle_type.target": "physical",
"stage1_model.router.enabled": True,
"stage1_model.router.type": "energy",
"stage1_model.router.n_experts": 8,
"stage2_model.router.enabled": True,
"stage2_model.router.type": "pdg",
"stage2_model.router.n_experts": 3,
}
)
name = gconfig.default_out_dir_name(cfg, now=_NOW)
# First 6 by priority: stage1_generator, stage2_generator, stage2_decoder,
# stage2_history, particle_type_target, stage1_router — stage2_router
# overflows into the hash suffix.
assert name.startswith("20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+")
def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive():
overrides = {
"stage1_model.generator": "wgan",
"stage2_model.generator": "flow",
"stage2_model.decoder": "one_shot",
"stage2_model.autoregressive.history": "attention",
"stage2_model.particle_type.target": "physical",
"stage1_model.router.enabled": True,
"stage1_model.router.type": "energy",
"stage1_model.router.n_experts": 8,
"train.seed": 3,
}
name_a = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW)
name_b = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW)
assert name_a == name_b
changed = dict(overrides, **{"train.seed": 99})
name_c = gconfig.default_out_dir_name(_cfg_with(**changed), now=_NOW)
assert name_c != name_a
# ---------------------------------------------------------------------------
# validate_config
# ---------------------------------------------------------------------------
def test_validate_config_default_config_passes():
gconfig.validate_config(gconfig.DEFAULT_CONFIG) # must not raise
def test_validate_config_embedding_target_requires_embedding_conditioning():
cfg = _cfg_with(
**{
"stage2_model.particle_type.target": "embedding",
"conditioning.particle.type": "physical",
}
)
try:
gconfig.validate_config(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "embedding" in str(e)
def test_validate_config_embedding_target_passes_with_embedding_conditioning():
cfg = _cfg_with(
**{
"stage2_model.particle_type.target": "embedding",
"conditioning.particle.type": "embedding",
"conditioning.material.type": "embedding",
}
)
gconfig.validate_config(cfg) # must not raise
def test_validate_config_mixed_particle_material_conditioning_is_valid():
"""The particle and material conditioning axes are configured
independently and may mix freely — e.g. material
"physical" with particle "embedding" — and the data pipeline
(giant/data/transforms.py) now implements that end-to-end, so
validate_config must not reject it."""
cfg = _cfg_with(
**{
"conditioning.particle.type": "physical",
"conditioning.material.type": "embedding",
}
)
gconfig.validate_config(cfg) # must not raise
def test_validate_config_pdg_router_incompatible_with_physical_conditioning():
cfg = _cfg_with(
**{
"stage1_model.router.enabled": True,
"stage1_model.router.type": "pdg",
"conditioning.particle.type": "physical",
}
)
try:
gconfig.validate_config(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "pdg" in str(e)
def test_validate_config_tie_to_stage1_requires_stage1_active():
cfg = _cfg_with(
**{
"stage2_model.router.tie_to_stage1": True,
"stage1_model.active": False,
}
)
try:
gconfig.validate_config(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "tie_to_stage1" in str(e)
def test_validate_config_stop_token_not_implemented():
cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"})
try:
gconfig.validate_config(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "stop_token" in str(e)
def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():
"""'n_sec.mode = "truth" is invalid for a rollout-capable checkpoint' —
both stages active means giant rollout
could load this checkpoint, but 'truth' has no ground truth to draw
n_sec from at rollout time."""
cfg = _cfg_with(
**{
"stage2_model.n_sec.mode": "truth",
"stage1_model.active": True,
"stage2_model.active": True,
}
)
try:
gconfig.validate_config(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "n_sec.mode" in str(e) and "truth" in str(e)
def test_validate_config_n_sec_truth_allowed_for_stage2_only_checkpoint():
"""'truth' is exactly the standalone stage-2 evaluation mode the design
doc carves out — stage1_model.active = false must still pass."""
cfg = _cfg_with(
**{
"stage2_model.n_sec.mode": "truth",
"stage1_model.active": False,
"stage2_model.active": True,
}
)
gconfig.validate_config(cfg) # must not raise
def test_validate_config_n_sec_truth_allowed_when_stage2_inactive():
cfg = _cfg_with(
**{
"stage2_model.n_sec.mode": "truth",
"stage1_model.active": True,
"stage2_model.active": False,
}
)
gconfig.validate_config(cfg) # must not raise
def test_validate_config_ar_default_markov_always_passes():
"""DEFAULT_CONFIG already has decoder='autoregressive',
history='markov', teacher_forcing='always' — must not raise (v0.3.0
step 5; see also test_validate_config_default_config_passes)."""
cfg = _cfg_with(**{"stage2_model.decoder": "autoregressive"})
gconfig.validate_config(cfg) # must not raise
def test_validate_config_ar_history_attention_passes():
"""v0.3.0 step 7 implements history='attention' — must not raise."""
cfg = _cfg_with(
**{
"stage2_model.decoder": "autoregressive",
"stage2_model.autoregressive.history": "attention",
}
)
gconfig.validate_config(cfg) # must not raise
@pytest.mark.parametrize("teacher_forcing", ["scheduled", "never"])
def test_validate_config_ar_teacher_forcing_scheduled_or_never_passes(teacher_forcing):
"""v0.3.0 step 7 implements teacher_forcing in {'scheduled', 'never'} —
must not raise."""
cfg = _cfg_with(
**{
"stage2_model.decoder": "autoregressive",
"stage2_model.autoregressive.teacher_forcing": teacher_forcing,
}
)
gconfig.validate_config(cfg) # must not raise
def test_validate_config_ar_history_invalid_value_rejected():
cfg = _cfg_with(
**{
"stage2_model.decoder": "autoregressive",
"stage2_model.autoregressive.history": "bogus",
}
)
try:
gconfig.validate_config(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "history" in str(e)
def test_validate_config_ar_teacher_forcing_invalid_value_rejected():
cfg = _cfg_with(
**{
"stage2_model.decoder": "autoregressive",
"stage2_model.autoregressive.teacher_forcing": "bogus",
}
)
try:
gconfig.validate_config(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "teacher_forcing" in str(e)
def test_validate_config_ar_checks_skipped_under_one_shot():
"""history/teacher_forcing values that would fail under AR are irrelevant
(and unchecked) when decoder='one_shot'."""
cfg = _cfg_with(
**{
"stage2_model.decoder": "one_shot",
"stage2_model.autoregressive.history": "attention",
"stage2_model.autoregressive.teacher_forcing": "scheduled",
}
)
gconfig.validate_config(cfg) # must not raise
# ---------------------------------------------------------------------------
# validate_config_keys / merge_cli_overrides unknown-key rejection
# ---------------------------------------------------------------------------
def test_validate_config_keys_default_config_passes():
gconfig.validate_config_keys(gconfig.DEFAULT_CONFIG) # must not raise
def test_validate_config_keys_rejects_unknown_top_level_key():
cfg = _cfg_with(**{"bogus_section.foo": 1})
try:
gconfig.validate_config_keys(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "bogus_section" in str(e)
def test_validate_config_keys_rejects_unknown_nested_key_with_close_match_hint():
cfg = _cfg_with(**{"stage1_model.n_res_block": 12})
try:
gconfig.validate_config_keys(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "stage1_model.n_res_block" in str(e)
assert "n_res_blocks" in str(e)
def test_validate_config_keys_allows_composed_router_axis_keys():
cfg = _cfg_with(
**{
"stage1_model.router.enabled": True,
"stage1_model.router.type": "composed",
"stage1_model.router.axis0_type": "energy",
"stage1_model.router.axis0_n_experts": 4,
"stage1_model.router.axis1_type": "pdg",
"stage1_model.router.axis1_emb_dim": 8,
}
)
gconfig.validate_config_keys(cfg) # must not raise
def test_validate_config_keys_allows_centers_init():
cfg = _cfg_with(**{"stage1_model.router.centers_init": [-1.0, 0.0, 1.0]})
gconfig.validate_config_keys(cfg) # must not raise
def test_validate_config_keys_rejects_unrelated_unknown_router_key():
cfg = _cfg_with(**{"stage1_model.router.n_expert": 4}) # typo for n_experts
try:
gconfig.validate_config_keys(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "stage1_model.router.n_expert" in str(e)
assert "n_experts" in str(e)
def test_validate_config_keys_skips_meta_section():
cfg = _cfg_with()
cfg["meta"] = {"config_version": 3, "git_hash": "abc123"}
gconfig.validate_config_keys(cfg) # must not raise
def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path):
path = tmp_path / "config.toml"
path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n")
try:
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
assert False, "expected ValueError"
except ValueError as e:
assert "n_res_block" in str(e)
def test_merge_cli_overrides_rejects_typo_in_cli_overrides():
try:
gconfig.merge_cli_overrides(
gconfig.DEFAULT_CONFIG,
None,
{"stage1_model": {"n_res_block": 12}},
)
assert False, "expected ValueError"
except ValueError as e:
assert "n_res_block" in str(e)
@pytest.mark.parametrize("fixture_name", ["default.toml", "wgan_h128_b4_physical.toml"])
def test_merge_cli_overrides_real_config_fixtures_pass_key_validation(fixture_name, monkeypatch):
monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243")
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)
# ---------------------------------------------------------------------------
def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
ckpt_path = tmp_path / "best.pt"
ckpt_path.write_bytes(b"")
_write_toml(tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n")
gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
captured = capsys.readouterr()
assert "warning" in captured.err
assert "old111" in captured.err
assert "current999" in captured.err
def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
ckpt_path = tmp_path / "best.pt"
ckpt_path.write_bytes(b"")
gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
assert capsys.readouterr().err == ""
def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(tmp_path, monkeypatch, capsys):
monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
ckpt_path = tmp_path / "best.pt"
ckpt_path.write_bytes(b"")
_write_toml(tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n")
gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
assert capsys.readouterr().err == ""
Reference in New Issue View Git Blame Copy Permalink