lars
da7cde3ef9
CI / Lint (ruff check) (push) Successful in 27s
CI / Format (ruff format) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 36s
CI / Type check (ty) (push) Successful in 39s
CI / Format (ruff format) (pull_request) Successful in 30s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 30s
CI / Tests (pull_request) Successful in 2m50s
CI / Tests (push) Successful in 2m58s
Works through docs/v0.3.0-followups.md item by item, closing the gap between the design doc and the shipped v0.3.0-stage2-autoregressive code: 1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2 dispatch, stage-2 particle-type-class marginal. 2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run. 3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/ train instead of the hardcoded K_MAX constant. 4. Mixed conditioning.particle.type / conditioning.material.type support end-to-end (data pipeline + dwarf warm-cache). 5. conditioning.share_stages = true: one shared ConditionEncoder instance across both stages. 6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2 (was silently unimplemented). 7. giant predict/rollout: implement conditioning.*.type = "onehot" via the checkpoint's saved pdg_topn_map/mat_topn_map. 8. network.py's checkpoint-path model_config migration now fails loudly on non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's TOML-load path (§4.2). 9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a rollout-capable checkpoint (§9). Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics), mostly a test-helper dict-unpack pattern that made every unrelated constructor keyword look like a type error. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
802 lines
27 KiB
Python
802 lines
27 KiB
Python
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",
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _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():
|
|
"""docs/v0.3.0-design.md §3.1: 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():
|
|
"""docs/v0.3.0-design.md §9: '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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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 == ""
|