diff --git a/CLAUDE.md b/CLAUDE.md index 7c89dab..8b73e3c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -91,6 +91,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`). -**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. +**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. This config break is why v0.2-shaped configs/checkpoints need migrating at all (`config.migrate_config`, `model.network._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`) — v0.2 checkpoint-loading support has **no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first. **Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet. diff --git a/giant/_migration.py b/giant/_migration.py new file mode 100644 index 0000000..007a023 --- /dev/null +++ b/giant/_migration.py @@ -0,0 +1,70 @@ +"""Shared v0.2 -> v0.3 migration knowledge. + +v0.3.0 broke the config format (single `[train]` + `[model]` -> `[conditioning]`/ +`[stage1_model]`/`[stage2_model]`/`[train]`), and that break has to be absorbed by two +independent migration surfaces: `giant.config.migrate_config` (a v0.2 `config.toml`) and +`giant.model.network._migrate_legacy_model_config` (a v0.2 checkpoint's flat +`model_config` dict). Both translate the same v0.2 facts into the same v0.3 shape, so +the facts live here once rather than as two hand-maintained copies — see issues.md +Issue 6. + +A dependency-free leaf module so neither `config.py` nor `network.py` has to import the +other to share this. +""" + +# v0.2 model-shaped keys (config.toml's [model] table, or a checkpoint's flat +# model_config dict — same key names in both) applied identically to both v0.3 stage +# blocks, because v0.2 had only one trunk shape shared by both stages. +V02_MODEL_KEY_TO_STAGES: tuple[tuple[str, str], ...] = ( + ("hidden_dim", "hidden_dim"), + ("n_blocks", "n_res_blocks"), + ("dropout", "dropout"), +) + +# v0.2 architectural facts that had no corresponding config key at all — always true of +# a v0.2 model, so both migration surfaces inject them unconditionally. Keyed by dotted +# path relative to the migrated dict's root. NOTE: conditioning.*.n_layers (2) differs +# from the v0.3 *default* (1) — not a typo, v0.2's conditioning MLP was always 2 layers +# deep. +V02_FIXED_FACTS: dict[str, object] = { + "conditioning.out_dim": 128, + "conditioning.particle.n_layers": 2, + "conditioning.material.n_layers": 2, + "stage1_model.active": True, + "stage1_model.flow.time_dim": 64, + "stage1_model.ddpm.time_dim": 64, + "stage2_model.active": True, + "stage2_model.flow.time_dim": 64, + "stage2_model.ddpm.time_dim": 64, + "stage2_model.context_dim": 64, + "stage2_model.decoder": "one_shot", + "stage2_model.particle_type.target": "physical", +} + + +def reject_legacy_router_expert_sizing(router_cfg: dict, *, source: str) -> None: + """Pop and validate v0.2's per-expert width/depth override, in place. + + v0.3.0 removed per-expert sizing — experts always inherit the stage's + hidden_dim/n_res_blocks — so a v0.2 router config/checkpoint that set a non-default + `expert_hidden_dim`/`expert_n_blocks` describes experts with a different width/depth + than the monolith, and can only be reproduced by v0.2 code. Silently dropping these + keys (a router builder's kwarg filtering would do this for free) would resize the + experts instead of refusing, so this raises loudly. + + Always pops both keys, whether or not they were non-default, so callers can go on + to use the (now-cleaned) `router_cfg` unconditionally. `source` names what's being + migrated (e.g. "v0.2 config's model.router" or "this checkpoint's + model_config.router") for the error message. + """ + expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0) + expert_n_blocks = router_cfg.pop("expert_n_blocks", 0) + if not (expert_hidden_dim or expert_n_blocks): + return + raise ValueError( + f"{source} sets expert_hidden_dim/expert_n_blocks to a non-default value " + f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed per-expert " + "sizing (experts always inherit the stage's hidden_dim/n_res_blocks), so " + "this router's experts have a different width/depth than the monolith. " + "This checkpoint/config can only be loaded by v0.2 code." + ) diff --git a/giant/config.py b/giant/config.py index c7d10ec..0f0ae77 100644 --- a/giant/config.py +++ b/giant/config.py @@ -14,6 +14,8 @@ from pathlib import Path import numpy as np import torch +from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing + class Conditioning(str, Enum): """`conditioning.particle.type` / `conditioning.material.type` choices — @@ -47,13 +49,11 @@ CONFIG_VERSION = 3 # `lambda` is a Python keyword, so dict key "lambda" is always exposed as the # field `lambda_weight`. # -# Two sub-blocks — router and n_sec — carry genuinely dynamic keys that don't -# fit a fixed schema: composed-router `axis{i}_{field}` flags (see -# giant.model.network._parse_composed_axes) and pipeline.py's runtime-seeded -# `centers_init`, plus n_sec's `legacy_owner` (injected only by -# _migrate_legacy_model_config for v0.2 checkpoints). Both dataclasses carry -# an `extra: dict` catch-all so these keys round-trip losslessly without -# becoming named fields that would leak into every new run's config.toml. +# The router sub-block carries genuinely dynamic keys that don't fit a fixed schema: +# composed-router `axis{i}_{field}` flags (see giant.model.network._parse_composed_axes) +# and pipeline.py's runtime-seeded `centers_init`. It carries an `extra: dict` catch-all +# so these keys round-trip losslessly without becoming named fields that would leak into +# every new run's config.toml. @dataclass(frozen=True) @@ -384,26 +384,24 @@ class NSecConfig: # only, never for rollout. mode: str = "head" lambda_weight: float = 0.1 # dict key "lambda" — cross-entropy weight for the head - # Holds "legacy_owner" when injected by _migrate_legacy_model_config - # (v0.2 checkpoints only) — not a user-facing config.toml key. - extra: dict = field(default_factory=dict) - - @property - def legacy_owner(self) -> str | None: - return self.extra.get("legacy_owner") + # Which stage's module physically owns the n_sec_head weights: "stage2" (default, + # fresh v0.3.0 runs — Stage2OneShot/Stage2Autoregressive builds it) or "stage1" + # (a migrated v0.2 checkpoint — see network._migrate_legacy_model_config, whose + # n_sec head was trained against Stage 1's own ConditionEncoder output and so has + # to stay attached there, not just be labeled as such). + owner: str = "stage2" @classmethod def from_dict(cls, d: dict | None) -> "NSecConfig": d = d or {} - known = {"mode", "lambda"} return cls( mode=d.get("mode", "head"), lambda_weight=d.get("lambda", 0.1), - extra={k: v for k, v in d.items() if k not in known}, + owner=d.get("owner", "stage2"), ) def to_dict(self) -> dict: - return {"mode": self.mode, "lambda": self.lambda_weight, **self.extra} + return {"mode": self.mode, "lambda": self.lambda_weight, "owner": self.owner} @dataclass(frozen=True) @@ -1024,14 +1022,6 @@ _V02_TRAIN_PASSTHROUGH = ( "wandb_log_every", ) -# v0.2 model.hidden_dim/n_blocks/dropout applied identically to both stages -# (there was only ever one trunk shape) -> copied to both stage{1,2}_model. -_V02_MODEL_TO_BOTH_STAGES = ( - ("hidden_dim", "hidden_dim"), - ("n_blocks", "n_res_blocks"), - ("dropout", "dropout"), -) - # v0.2 train.{n_critic,gp_weight,critic_lr} applied identically to both # stages' wgan sub-table (there was only ever one wgan objective, shared). _V02_TRAIN_TO_BOTH_STAGES_WGAN = ( @@ -1091,7 +1081,7 @@ def migrate_config(cfg: dict) -> dict: _set_path(new, f"stage1_model.wgan.{new_key}", old_train[old_key]) _set_path(new, f"stage2_model.wgan.{new_key}", old_train[old_key]) - for old_key, new_key in _V02_MODEL_TO_BOTH_STAGES: + for old_key, new_key in V02_MODEL_KEY_TO_STAGES: if old_key in old_model: _set_path(new, f"stage1_model.{new_key}", old_model[old_key]) _set_path(new, f"stage2_model.{new_key}", old_model[old_key]) @@ -1108,18 +1098,7 @@ def migrate_config(cfg: dict) -> dict: _set_path(new, "stage2_model.k_max", old_model["k_max"]) if old_router: - expert_hidden_dim = old_router.pop("expert_hidden_dim", 0) - expert_n_blocks = old_router.pop("expert_n_blocks", 0) - if expert_hidden_dim or expert_n_blocks: - raise ValueError( - "v0.2 config sets model.router.expert_hidden_dim/" - f"expert_n_blocks to a non-default value " - f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed " - "per-expert sizing (experts always inherit the stage's " - "hidden_dim/n_res_blocks), so this config's routed experts " - "have a different width/depth than the monolith and its " - "checkpoint can only be loaded by v0.2 code." - ) + reject_legacy_router_expert_sizing(old_router, source="v0.2 config's model.router") _set_path(new, "stage1_model.router", dict(old_router)) stage2_router = dict(old_router) stage2_router["tie_to_stage1"] = False @@ -1127,21 +1106,9 @@ def migrate_config(cfg: dict) -> dict: # v0.2 architectural facts with no corresponding config key at all — # always set once we've determined we're migrating a v0.2 dict, - # independent of what the file did/didn't specify. NOTE: n_layers here - # (2) differs from the v0.3 *default* (1) — this is not a typo, see the - # docstring above. - _set_path(new, "conditioning.out_dim", 128) - _set_path(new, "conditioning.particle.n_layers", 2) - _set_path(new, "conditioning.material.n_layers", 2) - _set_path(new, "stage1_model.active", True) - _set_path(new, "stage1_model.flow.time_dim", 64) - _set_path(new, "stage1_model.ddpm.time_dim", 64) - _set_path(new, "stage2_model.active", True) - _set_path(new, "stage2_model.flow.time_dim", 64) - _set_path(new, "stage2_model.ddpm.time_dim", 64) - _set_path(new, "stage2_model.context_dim", 64) - _set_path(new, "stage2_model.decoder", "one_shot") - _set_path(new, "stage2_model.particle_type.target", "physical") + # independent of what the file did/didn't specify (see giant._migration). + for path, value in V02_FIXED_FACTS.items(): + _set_path(new, path, value) new_meta = dict(cfg.pop("meta", {})) new_meta["config_version"] = CONFIG_VERSION diff --git a/giant/model/network.py b/giant/model/network.py index e946e20..4289529 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from giant._migration import V02_FIXED_FACTS, reject_legacy_router_expert_sizing from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig from giant.constants import ( COND_DIM, @@ -983,7 +984,7 @@ class Stage1Model(nn.Module): raise RuntimeError( "this Stage1Model has no n_sec_head — n_sec now lives on " "stage 2 by default; this method only exists " - "for a migrated v0.2 checkpoint (legacy_owner='stage1')" + "for a migrated v0.2 checkpoint (n_sec.owner='stage1')" ) c_emb = self.cond_enc(cond_cont, cond_cat) return self.n_sec_head(c_emb) @@ -1103,7 +1104,7 @@ class Stage2OneShot(nn.Module): if self.n_sec_head is None: raise RuntimeError( "this Stage2OneShot has no n_sec_head — it belongs to a " - "migrated v0.2 checkpoint (legacy_owner='stage1'); call " + "migrated v0.2 checkpoint (n_sec.owner='stage1'); call " "stage1.predict_n_sec(cond_cont, cond_cat) instead" ) c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) @@ -1345,7 +1346,7 @@ class Stage2Autoregressive(nn.Module): if self.n_sec_head is None: raise RuntimeError( "this Stage2Autoregressive has no n_sec_head — it belongs to " - "a migrated v0.2 checkpoint (legacy_owner='stage1'); call " + "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call " "stage1.predict_n_sec(cond_cont, cond_cat) instead" ) return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out)) @@ -1451,12 +1452,11 @@ def _migrate_legacy_model_config(model_config: dict) -> dict: `{"pdg_vocab", "mat_vocab", "conditioning", "stage1_model", "stage2_model"}` shape `build_models` expects. - Sets `stage2_model.n_sec.legacy_owner = "stage1"` so the n_sec_head - weights a v0.2 checkpoint carries on its Stage-1 module keep loading - there instead of the new default location - (`Stage2OneShot`) — the n_sec head was trained against Stage 1's own - `ConditionEncoder` output, so it has to stay attached to Stage 1's - module, not just be labeled as such. + Sets `stage2_model.n_sec.owner = "stage1"` so the n_sec_head weights a v0.2 + checkpoint carries on its Stage-1 module keep loading there instead of the new + default location (`Stage2OneShot`) — the n_sec head was trained against Stage 1's + own `ConditionEncoder` output, so it has to stay attached to Stage 1's module, not + just be labeled as such. Only the monolithic (non-routed) trunk shape is exercised by the step-2 migration test; a routed v0.2 checkpoint still builds correctly here @@ -1473,55 +1473,44 @@ def _migrate_legacy_model_config(model_config: dict) -> dict: k_max = m.get("k_max", K_MAX) noise_dim = m.get("noise_dim", 64) router_cfg = dict(m.get("router") or {}) - expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0) - expert_n_blocks = router_cfg.pop("expert_n_blocks", 0) - if expert_hidden_dim or expert_n_blocks: - raise ValueError( - "this checkpoint's model_config.router sets expert_hidden_dim/" - f"expert_n_blocks to a non-default value " - f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed " - "per-expert sizing (experts always inherit the stage's " - "hidden_dim/n_res_blocks), so this checkpoint's routed experts " - "have a different width/depth than the monolith — silently " - "dropping these keys would resize the experts instead of " - "refusing. This checkpoint can " - "only be loaded by v0.2 code." - ) + reject_legacy_router_expert_sizing(router_cfg, source="this checkpoint's model_config.router") router_cfg.setdefault("enabled", False) + F = V02_FIXED_FACTS + cond_n_layers = F["conditioning.particle.n_layers"] # same fact for both axes return { "pdg_vocab": m["pdg_vocab"], "mat_vocab": m["mat_vocab"], "conditioning": { - "out_dim": 128, + "out_dim": F["conditioning.out_dim"], "share_stages": False, - "particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": 2}, - "material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": 2}, + "particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers}, + "material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers}, }, "stage1_model": { - "active": True, + "active": F["stage1_model.active"], "generator": generator, "hidden_dim": hidden_dim, "n_res_blocks": n_blocks, "dropout": dropout, - "flow": {"time_dim": 64}, - "ddpm": {"time_dim": 64}, + "flow": {"time_dim": F["stage1_model.flow.time_dim"]}, + "ddpm": {"time_dim": F["stage1_model.ddpm.time_dim"]}, "wgan": {"noise_dim": noise_dim}, "router": dict(router_cfg), }, "stage2_model": { - "active": True, - "decoder": "one_shot", + "active": F["stage2_model.active"], + "decoder": F["stage2_model.decoder"], "generator": generator, "hidden_dim": hidden_dim, "n_res_blocks": n_blocks, "dropout": dropout, "k_max": k_max, - "context_dim": 64, - "n_sec": {"mode": "head", "legacy_owner": "stage1"}, - "particle_type": {"target": "physical"}, - "flow": {"time_dim": 64}, - "ddpm": {"time_dim": 64}, + "context_dim": F["stage2_model.context_dim"], + "n_sec": {"mode": "head", "owner": "stage1"}, + "particle_type": {"target": F["stage2_model.particle_type.target"]}, + "flow": {"time_dim": F["stage2_model.flow.time_dim"]}, + "ddpm": {"time_dim": F["stage2_model.ddpm.time_dim"]}, "wgan": {"noise_dim": noise_dim}, "router": {**router_cfg, "tie_to_stage1": False}, }, @@ -1545,7 +1534,7 @@ def migrate_legacy_state_dict(old_stage1_sd: dict, old_stage2_sd: dict) -> tuple new_stage1 = {} for k, v in old_stage1_sd.items(): if k.startswith("n_sec_head."): - new_stage1[k] = v # stays top-level (legacy_owner="stage1") + new_stage1[k] = v # stays top-level (n_sec.owner="stage1") else: new_stage1[_trunk_prefix(k)] = v @@ -1614,8 +1603,8 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: # matches the pre-dataclass .get("time_dim", 64) fallback, which # always hit its default for a wgan sub-block too. time_dim = getattr(s1_spec, generator).time_dim if generator != "wgan" else 64 - legacy_owner = s2_spec.n_sec.legacy_owner - n_sec_head_k_max = s2_spec.k_max if legacy_owner == "stage1" else None + n_sec_owner = s2_spec.n_sec.owner + n_sec_head_k_max = s2_spec.k_max if n_sec_owner == "stage1" else None result["stage1"] = Stage1Model( pdg_vocab=pdg_vocab, mat_vocab=mat_vocab, @@ -1646,7 +1635,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: # wgan has no time_dim concept — see the matching comment in stage 1 # above. time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64 - legacy_owner = s2_spec.n_sec.legacy_owner + n_sec_owner = s2_spec.n_sec.owner k_max = s2_spec.k_max particle_type_cfg = s2_spec.particle_type.to_dict() @@ -1667,7 +1656,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: noise_dim=s2_spec.wgan.noise_dim, k_max=k_max, router=stage2_router, - build_n_sec_head=legacy_owner != "stage1", + build_n_sec_head=n_sec_owner != "stage1", particle_type_cfg=particle_type_cfg, history=ar_cfg.history, attn_n_heads=ar_cfg.attn_n_heads, @@ -1692,7 +1681,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: noise_dim=s2_spec.wgan.noise_dim, k_max=k_max, router=stage2_router, - build_n_sec_head=legacy_owner != "stage1", + build_n_sec_head=n_sec_owner != "stage1", particle_type_cfg=particle_type_cfg, cond_enc=shared_cond_enc, ) diff --git a/issues.md b/issues.md index 7d5a085..0c8f62a 100644 --- a/issues.md +++ b/issues.md @@ -45,7 +45,7 @@ architecture matrix grows, not about rot or breakage. | 3 | `cli.py:train()` is a 58-parameter, 510-line fat controller | **High** | Medium | **Fixed** (`2bfb1ab`) | | 4 | `cli.py` is at 35.8 % coverage and holds untested override-precedence logic | **High** | Medium | **Fixed** (`2bfb1ab`, partial — see status note) | | 5 | Inference bootstrap is duplicated verbatim between `predict` and `rollout` | **High** | Small | **Fixed** | -| 6 | Two independent v0.2→v0.3 migration surfaces encode the same knowledge | Medium | Medium | Open | +| 6 | Two independent v0.2→v0.3 migration surfaces encode the same knowledge | Medium | Medium | **Fixed** | | 7 | Positional tuple contracts (9-tuple, 7-tuple) between data, model and training layers | Medium | Small | Open | | 8 | `network.py` is 1745 lines holding three distinct modules | Medium | Small | Open | | 9 | `scripts` is published as a top-level distribution package | Medium | Small | Open | @@ -770,6 +770,34 @@ imported at `analysis/` module scope. ## Issue 6 — Two independent v0.2→v0.3 migration surfaces encode the same knowledge +> **Status: Fixed.** A new leaf module, `giant/_migration.py`, now holds the v0.2 facts +> both `config.migrate_config` and `network._migrate_legacy_model_config` need: +> `V02_MODEL_KEY_TO_STAGES` (the `hidden_dim`/`n_blocks`/`dropout` → both-stages +> mapping), `V02_FIXED_FACTS` (the ~12 "architectural facts with no config key" — +> conditioning `out_dim`/`n_layers`, both stages' `active`/`flow.time_dim`/ +> `ddpm.time_dim`, stage 2's `context_dim`/`decoder`/`particle_type.target`), and +> `reject_legacy_router_expert_sizing` (one canonical error message for a non-default +> `router.expert_hidden_dim`/`expert_n_blocks`, replacing the two hand-written copies). +> Both migration functions now consume this module instead of hardcoding their own +> copies, so a future correction needs one edit, not two. Separately (fix item 2), the +> checkpoint-only `n_sec.legacy_owner` nullable sentinel (living in `NSecConfig.extra`, +> `None` unless a v0.2 checkpoint was migrated) is replaced by a normal, always-set +> `NSecConfig.owner: str = "stage2"` field — the current schema now carries +> `stage2_model.n_sec.owner` for every run (`"stage2"` by default, `"stage1"` only for a +> migrated v0.2 checkpoint), so `build_models` reads one concrete two-valued key instead +> of branching on a nullable legacy marker; `NSecConfig.extra` was dropped entirely since +> `legacy_owner` was its only consumer. Fix item 3 (retention policy) is recorded in +> `CLAUDE.md`'s v0.3.0 roadmap paragraph as **explicitly undecided** rather than a +> concrete expiry — `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs +> referencing them, and no criterion for dropping the shim (or +> `tests/legacy/network_v02_snapshot.py`) has been agreed yet. `tests/legacy/ +> network_v02_snapshot.py` itself is untouched, per the scope guard below. Existing +> regression tests updated for the rename +> (`tests/test_config.py`, `tests/test_migration_v02_v03.py`); the four bit-identical +> `test_migration_*` cases and the full suite (804 tests) pass unchanged, confirming no +> migrated value changed. Everything below this point describes the pre-fix state and is +> kept for historical context. + **Severity: Medium. Effort: Medium.** **Location:** `giant/config.py:514` (`migrate_config`) and diff --git a/tests/test_config.py b/tests/test_config.py index 58d4d1c..0f1411c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -95,10 +95,15 @@ def test_stage1_router_config_has_no_tie_to_stage1_key(): 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"} +def test_n_sec_config_owner_defaults_to_stage2(): + n_sec = gconfig.NSecConfig() + assert n_sec.owner == "stage2" + + +def test_n_sec_config_owner_round_trips(): + n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"}) + assert n_sec.owner == "stage1" + assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1"} # --------------------------------------------------------------------------- diff --git a/tests/test_migration_v02_v03.py b/tests/test_migration_v02_v03.py index 35f4db5..126e237 100644 --- a/tests/test_migration_v02_v03.py +++ b/tests/test_migration_v02_v03.py @@ -130,7 +130,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None: new_stage1, new_stage2 = new_models["stage1"], new_models["stage2"] assert isinstance(new_stage1, net.Stage1Model) assert isinstance(new_stage2, net.Stage2OneShot) - # legacy_owner="stage1": n_sec lives on stage1, not stage2, for a + # n_sec.owner="stage1": n_sec lives on stage1, not stage2, for a # migrated v0.2 checkpoint. assert new_stage1.n_sec_head is not None assert new_stage2.n_sec_head is None @@ -177,7 +177,7 @@ def test_migration_wgan_physical(): def test_migrate_legacy_model_config_shape(): """_migrate_legacy_model_config produces the nested shape build_models - expects, with the legacy_owner marker set so build_models routes the + expects, with the n_sec.owner marker set so build_models routes the n_sec head back onto stage 1.""" legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical") migrated = net._migrate_legacy_model_config(legacy_cfg) @@ -187,7 +187,7 @@ def test_migrate_legacy_model_config_shape(): assert migrated["conditioning"]["particle"]["n_layers"] == 2 assert migrated["conditioning"]["material"]["n_layers"] == 2 assert migrated["stage1_model"]["hidden_dim"] == HIDDEN_DIM - assert migrated["stage2_model"]["n_sec"]["legacy_owner"] == "stage1" + assert migrated["stage2_model"]["n_sec"]["owner"] == "stage1" assert migrated["stage2_model"]["decoder"] == "one_shot" @@ -278,6 +278,6 @@ def test_build_models_accepts_new_nested_shape_unchanged(): models = net.build_models(cfg) assert isinstance(models["stage1"], net.Stage1Model) assert isinstance(models["stage2"], net.Stage2OneShot) - # Fresh v0.3.0 config, no legacy_owner: n_sec lives on stage 2. + # Fresh v0.3.0 config, n_sec.owner defaults to "stage2": n_sec lives on stage 2. assert models["stage1"].n_sec_head is None assert models["stage2"].n_sec_head is not None