lars
9112e845e0
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 / Type check (ty) (push) Successful in 23s
CI / Tests (push) Successful in 1m1s
CI / Lint (ruff check) (pull_request) Successful in 26s
CI / Format (ruff format) (pull_request) Successful in 27s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 27s
CI / Tests (pull_request) Successful in 1m2s
Replaces train.py's single global training loop with a StageTrainer hierarchy (FlowDDPMStageTrainer, WGANStageTrainer) — one per active stage, each owning its own optimizer/LR schedule/EMA and reading only the shared batch tuple (stage 2 always teacher-forces on the ground-truth x1_s1, so stages never need each other's output at train time). Supports every stage1/stage2 generator combination, including the design doc's headline mixed case (stage1=flow + stage2=wgan) and its reverse, plus stage1-only/stage2-only ablation runs, routed+gumbel stages, and checkpoint save/resume. metrics.csv/wandb logging are stage-prefixed. validate_marginals calls are guarded with a one-time warning and a Wasserstein-magnitude fallback for wgan best-checkpoint selection, since giant/sample.py still assumes stage1 always owns n_sec_head (decision 1 moved it to stage 2 by default) — deferred to design doc step 6, not silently papered over. pipeline.py's run_setup_stage/run_train_job now read the new nested config directly; the dangling resolve_expert_dims call and the --mode wgan --router rejection are both gone (routed WGAN works). cli.py's train/new-run build correctly-shaped config overrides (architecture flags -> stage1_model only per the approved decision; --mode/--n-critic/--gp-weight/--critic-lr broadcast to both stages, matching migrate_config's own precedent and avoiding a regression on the common --mode case); predict/rollout's dangling build_models tuple-unpack is fixed; new-run now tags config_version, fixing a bug where a re-loaded v0.3 config.toml would have been silently corrupted by migrate_config mistaking it for v0.2. config.py's validate_config rejects mixed particle/material conditioning types for now (ConditionEncoder supports it, the data pipeline in giant/data/transforms.py doesn't yet). analysis/render.py and router_gating.py handle both the new nested model_config shape and legacy flat checkpoints. scripts/warm_setup_cache.py updated for run_setup_stage's new signature. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
686 lines
23 KiB
Python
686 lines
23 KiB
Python
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
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_not_supported():
|
|
"""The data pipeline doesn't support mixed conditioning types yet, even
|
|
though ConditionEncoder itself already can (docs/v0.3.0-design.md §3.1
|
|
vs. giant/data/transforms.py's still-single conditioning param)."""
|
|
cfg = _cfg_with(
|
|
**{
|
|
"conditioning.particle.type": "physical",
|
|
"conditioning.material.type": "embedding",
|
|
}
|
|
)
|
|
try:
|
|
gconfig.validate_config(cfg)
|
|
assert False, "expected ValueError"
|
|
except ValueError as e:
|
|
assert "mixed" in str(e) or "conditioning.material.type" in str(e)
|
|
|
|
|
|
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)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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 == ""
|