From 30a448927cb3785ea76735972b6e40bf7d2d174e Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 13 Aug 2026 10:42:01 +0200 Subject: [PATCH] Remove issues.md All tracked issues have been resolved and merged individually. Co-Authored-By: Claude Sonnet 5 --- issues.md | 1392 ----------------------------------------------------- 1 file changed, 1392 deletions(-) delete mode 100644 issues.md diff --git a/issues.md b/issues.md deleted file mode 100644 index 3c78d2f..0000000 --- a/issues.md +++ /dev/null @@ -1,1392 +0,0 @@ -# GIANT — architecture issues - -Software-engineering review of the `giant` codebase, conducted 2026-08-12 on branch -`v0.3.0-stage2-autoregressive` at commit `55332db`. - -This document is written to be read **without the context of the review conversation**. -Each issue states what the code does today, why that is a problem, how to verify the -claim independently, what the fix looks like, and — importantly — what the fix should -*not* touch. Line numbers are accurate as of `55332db`; if they have drifted, the -accompanying code excerpts and `grep` commands will still locate the code. - -## Baseline: what is already good - -Read this first, so the issues below are calibrated correctly. **This is a healthy -codebase.** The problems listed are about structural resilience as the v0.3.0 -architecture matrix grows, not about rot or breakage. - -- `pytest` — **725 tests pass in ~18 s**. Fast enough that there is no excuse for not - running it on every change. -- `uv run ruff check .`, `uv run ruff format --check .`, `uv run ty check .` — all clean, - and all three are enforced in CI (`.gitea/workflows/ci.yml`). -- Zero `TODO` / `FIXME` / `XXX` / `HACK` markers in `giant/` or `scripts/`. -- The internal import graph is **acyclic** with clean layering: - `constants → data → model → training → pipeline → cli`, and `analysis` almost fully - independent of the rest. -- Line coverage is **96–99.6 % on every core module** (`network.py` 98.6 %, - `trainers.py` 99.5 %, `rollout.py` 99.6 %, `transforms.py` 96.2 %, `pipeline.py` - 97.1 %, `analysis/catalog.py` 99.6 %). The single exception is `cli.py` at 35.8 % — - see Issue 4. -- Comment and docstring quality is unusually high. Docstrings routinely explain *why* - a decision was made, not just what the code does, and several call out their own - known limitations honestly. **Preserve this when refactoring.** A refactor that - deletes the reasoning in `giant/config.py`'s `DEFAULT_CONFIG` comments or - `giant/analysis/router_gating.py`'s module docstring is a net loss even if the code - gets shorter. -- `giant/analysis/` is the best-designed subsystem in the repo and should be treated as - the template the rest of the codebase moves toward. See Issue 11. - -## Issue index - -| # | Issue | Severity | Effort | Status | -|---|---|---|---|---| -| 1 | Config defaults are declared twice; `DEFAULT_CONFIG` and consumers already disagree | **High** | Medium | **Fixed** (`9bf5874`) | -| 2 | No unknown-key validation — a typo in `config.toml` silently trains the wrong model | **High** | Small | **Fixed** | -| 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 | **Fixed** | -| 7 | Positional tuple contracts (9-tuple, 7-tuple) between data, model and training layers | Medium | Small | **Fixed** | -| 8 | `network.py` is 1745 lines holding three distinct modules | Medium | Small | **Fixed** | -| 9 | `scripts` is published as a top-level distribution package | Medium | Small | **Fixed** | -| 10 | `torch.load(weights_only=False)` — checkpoints are arbitrary pickles | Low | Medium | Open | -| 11 | Minor: `echo=print` threading, `particles → data.loader` layering | Low | Small | Open | - ---- - -## Issue 1 — Config defaults are declared twice, and the two declarations already disagree - -> **Status: Fixed, commit `9bf5874` on `v0.3.0-stage2-autoregressive`.** `giant/config.py` now -> declares a hierarchy of frozen dataclasses (`GiantConfig` and its nested blocks, following -> `StageSpec`'s existing style) as the single source of truth; `DEFAULT_CONFIG` is generated from -> `GiantConfig().to_dict()` rather than hand-maintained, and `build_models`/`build_critics` -> (`network.py`) and `StageSpec.from_config` (`trainers.py`) convert their dict input into these -> dataclasses at the top of each function and read attributes instead of duplicating -> `.get(key, literal)` defaults. Both drifted keys (`decoder`, `particle_type.target`) now resolve -> to `DEFAULT_CONFIG`'s documented v0.3.0 values (`"autoregressive"`/`"onehot"`) unconditionally. -> Router/`n_sec` sub-blocks keep a dict-shaped `extra` catch-all for their genuinely dynamic keys -> (composed-router `axis{i}_*`, runtime-seeded `centers_init`, legacy `legacy_owner`) rather than -> being fully typed — see the "Recommended fix" section below, which this diverges from slightly -> on that one point. Fixing the fallback surfaced two existing callers that had been silently -> depending on the old (wrong) default — a `tests/test_train.py` fixture and -> `scripts/warm_setup_cache.py`'s hand-built minimal config (now merged against `DEFAULT_CONFIG` -> instead of hand-rolled) — both fixed in the same commit. Regression tests added in -> `tests/test_config.py`, `tests/test_network.py`, `tests/test_train.py`, including one pinning -> `GiantConfig().to_dict() == DEFAULT_CONFIG` so this class of drift can't recur silently. -> Everything below this point describes the pre-fix state and is kept for historical context. - -**Severity: High. Effort: Medium. Risk if unfixed: silent wrong-model training.** - -**Location:** `giant/config.py:32` (`DEFAULT_CONFIG`) versus `giant/model/network.py:1570-1693` -(`build_models`), `giant/training/trainers.py:130-175` (`StageSpec.from_config`), -`giant/pipeline.py:381`. - -### What the code does today - -`DEFAULT_CONFIG` in `giant/config.py:32` is a fully-specified nested dict. Every key has -a value and most have an explanatory comment. It is, on paper, the single source of -truth for configuration defaults. - -The consumers do not treat it that way. They re-declare the same defaults inline via -`dict.get(key, default)`: - -```python -# giant/model/network.py, inside build_models -cond_out_dim = conditioning.get("out_dim", 128) -... -hidden_dim=s1cfg.get("hidden_dim", 256), -n_res_blocks=s1cfg.get("n_res_blocks", 6), -dropout=s1cfg.get("dropout", 0.0), -time_dim=gen_sub.get("time_dim", 64), -``` - -```python -# giant/training/trainers.py, inside StageSpec.from_config -lambda_weight=stage_cfg.get("lambda", 1.0), -n_sec_lambda=cfg["stage2_model"].get("n_sec", {}).get("lambda", 0.1), -ddpm_n_steps=stage_cfg.get("ddpm", {}).get("n_steps", 1000), -n_critic=wgan_cfg.get("n_critic", 5), -gp_weight=wgan_cfg.get("gp_weight", 10.0), -``` - -There are **132 `.get("key", default)` call sites across `giant/`**, concentrated in -`network.py` (55) and `trainers.py` (25). Each one is an independent copy of a value -that `DEFAULT_CONFIG` also declares. - -### Why it's a problem - -Because `cfg` is typed `dict` (i.e. `dict[str, Any]`), neither `ty` nor the tests can -see that the two declarations are supposed to agree. When they drift, nothing fails — -the model is simply built differently than the config file says. - -**This has already happened.** Two defaults currently disagree: - -| Key | `DEFAULT_CONFIG` says | `build_models` fallback says | -|---|---|---| -| `stage2_model.decoder` | `"autoregressive"` | `"one_shot"` (`network.py:1632`) | -| `stage2_model.particle_type.target` | `"onehot"` | `"physical"` (`network.py:1644`, `trainers.py:143`, `pipeline.py:381`) | - -Both are *architecturally load-bearing*: `decoder` selects between `Stage2Autoregressive` -and `Stage2OneShot` — two different networks — and `particle_type.target` selects between -categorical class logits and continuous `(log-mass, charge)` regression, which is the -exact axis the v0.3.0 redesign was created to change (see `CLAUDE.md`, "v0.3.0 — -Stage-2 autoregressive redesign"). - -### Is it currently a live bug? - -**No — it is currently latent.** Two things save it today: - -1. The training path always passes a fully-merged config. `giant/pipeline.py:446` builds - `model_config` from `cfg["stage1_model"]` / `cfg["stage2_model"]`, and `cfg` always - originates from `merge_cli_overrides(DEFAULT_CONFIG, ...)`, so every key is present - and no fallback fires. -2. The legacy path also populates the keys explicitly. - `_migrate_legacy_model_config` (`network.py:1446`) hard-codes - `"decoder": "one_shot"` and `"particle_type": {"target": "physical"}` into the dict it - returns, so a v0.2 checkpoint does not rely on the fallbacks either. -3. The tests build their configs by deep-copying `DEFAULT_CONFIG` (e.g. - `tests/test_network.py:611` `_minimal_model_config`), so they never exercise the - fallback branch. - -That is precisely what makes this dangerous rather than harmless. The fallbacks are -**unreachable by every current caller, untested, and wrong.** The first caller that -hand-builds a partial `model_config` — a new test, a notebook, a debugging script, a -future `giant train --stage2-only` path — silently gets the v0.2 architecture while the -config documentation promises the v0.3 one. - -### How to verify - -```bash -uv run python - <<'EOF' -import copy -from giant import config as gc -from giant.model.network import build_models - -full = copy.deepcopy(gc.DEFAULT_CONFIG) -print("DEFAULT_CONFIG stage2 decoder :", full["stage2_model"]["decoder"]) -print("DEFAULT_CONFIG particle_type.target:", full["stage2_model"]["particle_type"]["target"]) - -partial = { - "pdg_vocab": 3, "mat_vocab": 2, - "conditioning": full["conditioning"], - "stage1_model": {"hidden_dim": 8, "n_res_blocks": 1}, - "stage2_model": {"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3}, -} -print("build_models fallback stage2 class :", type(build_models(partial)["stage2"]).__name__) -EOF -``` - -Observed output on `55332db`: - -``` -DEFAULT_CONFIG stage2 decoder : autoregressive -DEFAULT_CONFIG particle_type.target: onehot -build_models fallback stage2 class : Stage2OneShot -``` - -### Recommended fix - -Make the config **typed at the boundary**, so the defaults exist exactly once and the -type checker can see them. - -1. Define a frozen dataclass per config block in `giant/config.py`: - `ConditioningAxisConfig`, `ConditioningConfig`, `RouterConfig`, `WganConfig`, - `Stage1ModelConfig`, `Stage2ModelConfig`, `TrainConfig`, `GiantConfig`. The dataclass - field defaults become the *only* declaration of each default. Move the existing - `DEFAULT_CONFIG` comments onto the fields — they are the most valuable thing in that - file and must not be lost. -2. Give each a `from_dict(d)` classmethod that constructs from the merged TOML dict and - **rejects unknown keys** (that is Issue 2, and the same constructor solves both). -3. Change `build_models`, `build_critics` and `StageSpec.from_config` to take the - dataclasses instead of dicts. Every `.get(key, default)` becomes `cfg.key`. The 132 - duplicate defaults disappear by construction. -4. Keep `DEFAULT_CONFIG` as a derived artifact if anything still needs the dict shape - (`save_config`, `default_out_dir_name`, the setup-cache sidecar) — generate it from - the dataclasses via `dataclasses.asdict`, do not maintain it by hand. - -Precedent exists in this codebase and should be followed: -`@dataclass(frozen=True) StageSpec` (`trainers.py:79`) and `PlotSpec` -(`analysis/catalog.py:65`) already do exactly this, well. - -### Interim fix if the full change is too large - -If the dataclass migration is deferred, at minimum add a test that asserts the inline -fallbacks agree with `DEFAULT_CONFIG`, so the drift is caught: - -```python -def test_build_models_fallbacks_match_default_config(): - """Every `.get(key, default)` in build_models must use the same default - DEFAULT_CONFIG declares — see issues.md Issue 1.""" -``` - -and fix the two divergent values. Decide deliberately which way: if the fallbacks are -meant to encode *v0.2 legacy* behaviour, they are redundant (the legacy migration at -`network.py:1446` already sets both keys explicitly) and should be **deleted**, letting -a missing key raise `KeyError` instead of silently choosing an architecture. - -### Scope guard - -Do **not** change any actual default *value* as part of this refactor. The goal is to -make the two declarations agree and then have only one; changing what the model does is -a separate, physics-relevant decision that belongs in its own commit with its own -retraining. When resolving the two divergences above, the correct target is whatever -`DEFAULT_CONFIG` says (`autoregressive` / `onehot`), because that is the documented -v0.3.0 intent per `CLAUDE.md`. - ---- - -## Issue 2 — No unknown-key validation: a typo in `config.toml` silently trains the wrong model - -> **Status: Fixed.** `giant/config.py` now has `validate_config_keys(cfg)`, which walks a -> merged config dict recursively against `DEFAULT_CONFIG`'s tree (itself generated from -> the `GiantConfig` dataclasses added for Issue 1) and raises `ValueError` on any key not -> present there, with a `difflib`-based "did you mean" suggestion. `[meta]` is skipped -> unconditionally, and `stage{1,2}_model.router`'s `axis{i}_{field}` composed-router keys -> and runtime-seeded `centers_init` are allowed through explicitly. `merge_cli_overrides` -> calls it right before returning, so all three real entry points (`giant train`, -> `giant new-run`, `scripts/warm_setup_cache.py`) are covered automatically; checkpoint -> `model_config` loading (a separate code path, `network._migrate_legacy_model_config`) -> is untouched, so old checkpoints keep loading regardless of schema drift. This is the -> "standalone fallback" option from the recommended fix below rather than the -> dataclass-`from_dict`-rejects-unknown-keys option, specifically to keep the check scoped -> to the config.toml/CLI-overrides path without touching `validate_config`'s many direct -> unit-test callers or `run_train_job`'s redundant internal `validate_config` call. -> Regression tests added in `tests/test_config.py` cover top-level and nested typos, the -> axis/`centers_init` allowances, `[meta]` skipping, and both the file and CLI-overrides -> paths, plus the two real `configs/*.toml` fixtures that don't already fail -> `migrate_config` for unrelated reasons. Everything below this point describes the -> pre-fix state and is kept for historical context. - -**Severity: High. Effort: Small. Risk if unfixed: wasted GPU-days on a run that did not -use the setting you thought it did.** - -**Location:** `giant/config.py:456` (`_deep_merge`), `giant/config.py:628` -(`merge_cli_overrides`), `giant/config.py:661` (`validate_config`). - -### What the code does today - -`merge_cli_overrides` resolves configuration as `DEFAULT_CONFIG → TOML file → CLI -overrides`, deep-merging at each step. `_deep_merge` accepts any key: - -```python -def _deep_merge(base: dict, override: dict) -> dict: - result = dict(base) - for k, v in override.items(): - if isinstance(v, dict) and isinstance(result.get(k), dict): - result[k] = _deep_merge(result[k], v) - else: - result[k] = v # <-- unknown k is accepted and stored - return result -``` - -`migrate_config` explicitly preserves unrecognised sections -(`giant/config.py:612-615`: *"Anything else in the original dict (unrecognized top-level -sections) carries through untouched rather than being silently dropped"*). - -`validate_config` (`giant/config.py:661`) exists and is *good* — it catches cross-block -contradictions with genuinely excellent, actionable error messages (e.g. "`router.type = -'pdg'` builds its own training-vocab-scoped embedding, incompatible with -`conditioning.particle.type = 'physical'`"). But by design it only inspects keys it knows -about. It cannot detect a key that should not exist. - -### Why it's a problem - -A config file containing - -```toml -[stage1_model] -n_res_block = 12 # typo: should be n_res_blocks -``` - -merges cleanly, validates cleanly, trains to completion, and reports success — having -built a 6-block model. The stray `n_res_block` key is faithfully written back into the -run's saved `config.toml` by `save_config`, so the artifact of record also claims the -typo was a real setting. Nothing anywhere in the pipeline will ever say otherwise. - -The blast radius scales with the config surface. v0.3.0 introduced a large nested schema -(`conditioning` × `stage1_model` × `stage2_model` × per-generator sub-tables × router -sub-tables), which is a great deal of surface for typos, and the v0.3.0 workflow is -explicitly *"a sequence of architecture comparisons"* (`config.py`, the `wandb` default -comment). A comparison in which one arm silently ignored its distinguishing setting is -worse than no comparison — it produces a confident, wrong conclusion. - -There is a partial mitigation already: `warn_if_git_hash_mismatch` and -`warn_if_checkpoint_config_mismatch` warn about *provenance* drift. Neither looks at key -names. - -### How to verify - -```bash -uv run python - <<'EOF' -import tempfile, pathlib -from giant import config as gc - -p = pathlib.Path(tempfile.mkdtemp()) / "config.toml" -p.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n") -cfg = gc.merge_cli_overrides(gc.DEFAULT_CONFIG, p, {}) -gc.validate_config(cfg) # passes -print("typo key survived :", cfg["stage1_model"]["n_res_block"]) -print("real key unchanged:", cfg["stage1_model"]["n_res_blocks"]) -EOF -``` - -### Recommended fix - -Add strict key validation against the known schema. Two options, in order of preference: - -1. **Preferred — fold into Issue 1.** If each config block becomes a dataclass with a - `from_dict` that raises on unknown keys, this issue is solved for free and cannot - regress. Raise a `ValueError` in the established house style, naming the offending - key, its section, and the closest valid key by edit distance - (`difflib.get_close_matches`) — that last part matters, because "unknown key - `n_res_block`; did you mean `n_res_blocks`?" is the message that actually saves the - run. - -2. **Standalone fallback.** Add `validate_config_keys(cfg)` to `giant/config.py` that - walks `cfg` alongside `DEFAULT_CONFIG` and raises on any key not present in the - defaults tree. Call it from `merge_cli_overrides` right before returning, so every - entry point gets it. - -Either way, the following need explicit allowances, since they legitimately carry keys -not in `DEFAULT_CONFIG`: - -- **`meta`** — holds `config_version`, `git_hash`, and run metadata from - `build_run_meta`, none of which are in the defaults tree. -- **Composed-router axis keys** — `stage{1,2}_model.router.axis{i}_{field}` are - deliberately dynamic and deliberately absent from `DEFAULT_CONFIG` (see the comment at - the end of the `router` block in `config.py`, and `_parse_composed_axes` in - `network.py:499`). Validate these against the `axis(\d+)_(\w+)` pattern plus the known - per-axis field names rather than against a fixed key list. -- **`stage2_model.n_sec.legacy_owner`** — injected by `_migrate_legacy_model_config`, not - a user-facing key. It appears in checkpoint `model_config` dicts, not in `config.toml`, - so it should not reach this validator; confirm that before assuming. - -### Scope guard - -Strict validation must apply to the **`config.toml` path only**, not to loading old -checkpoints. A v0.2 checkpoint's `model_config` is migrated by a different function -(`network._migrate_legacy_model_config`, see Issue 6) and must keep loading. Adding -strictness that rejects historical checkpoints would break `giant predict` / `giant -rollout` against every model trained so far. - ---- - -## Issue 3 — `cli.py:train()` is a 58-parameter, 510-line fat controller - -> **Status: Fixed, commit `2bfb1ab` on `v0.3.0-stage2-autoregressive`.** `giant/config.py` -> now declares `FlagSpec` (a frozen dataclass: `name`, `paths` — one or more dotted config -> paths, `precedence`) and a `FLAG_SPECS` table covering every flag `train`/`new-run` map -> into the config tree, plus `overrides_from_flags(values: dict[str, object]) -> dict`, -> which applies specs in ascending precedence order (so a more-specific flag overwrites a -> shared/shorthand one written earlier at the same path) — one mechanism replacing the -> three different ad hoc "more specific wins" patterns identified below -> (`.update()`-call-order, `setdefault(...)[...] =` overwrite-order, and -> `{**shared, **specific}` merge). `train()`'s ~140-line override-building block -> (`cli.py:656-790` pre-fix) is now a single flat `flag_values` dict (mostly enum `.value` -> unwrapping) plus one call to `overrides_from_flags`; `new_run()`'s near-verbatim copy -> (`cli.py:915-983` pre-fix) collapsed the same way, reusing the identical table. Router -> overrides (`_router_cli_overrides`, unchanged, still shared by both commands) feed into -> the table as a single pre-aggregated `router_config` entry mapped only to -> `stage1_model.router` — the stage1-only asymmetry noted below is preserved exactly, with -> a regression test. No flag was added, removed, or renamed, and no precedence semantics -> changed: `giant train --help`/`giant new-run --help` are byte-identical before and after, -> verified by diffing both. Everything below this point describes the pre-fix state and is -> kept for historical context. - -**Severity: High. Effort: Medium.** - -**Location:** `giant/cli.py:329-839`. - -### What the code does today - -Measured on `55332db`: - -| Command | Parameters | Lines | -|---|---|---| -| `train` | **58** | **510** | -| `new_run` | 31 | 188 | -| `predict` | 9 | 387 | -| `rollout` | 14 | 237 | - -Of `train`'s 510 lines, roughly 430 are the Typer signature (one `Annotated[...]` block -per flag) and roughly 80 are hand-written translation from flags into the nested -overrides dict. That translation follows the same shape five times over — for `train`, -`stage1_model`, `stage2_model`, `conditioning`, and the WGAN sub-tables: - -```python -cli_stage1_model: dict[str, object] = { - k: v for k, v in { - "hidden_dim": hidden_dim, - "n_res_blocks": n_blocks, - "dropout": dropout, - }.items() if v is not None -} -cli_stage1_model.update( - {k: v for k, v in { - "hidden_dim": stage1_hidden_dim, - "n_res_blocks": stage1_n_res_blocks, - "dropout": stage1_dropout, - }.items() if v is not None} -) -``` - -On top of that sits genuinely intricate precedence logic, correctly implemented but -expressed imperatively: - -- `--mode` sets `generator` on **both** stages; `--stage1-generator` / `--stage2-generator` - then override a single stage (`cli.py:713-720`). -- `--hidden-dim` / `--n-blocks` / `--dropout` are stage-1-only backward-compatible - shorthands; `--stage1-*` wins when both are given (`cli.py:676-700`). -- `--n-critic` / `--gp-weight` / `--noise-dim` / `--critic-lr` fan out to - `stage{1,2}_model.wgan.*`, with `--stage{1,2}-*` variants overriding per stage — built - by merging a shared dict under a stage-specific one (`cli.py:722-753`). -- `--emb-dim` and `--conditioning` each set **two** places - (`conditioning.particle.*` and `conditioning.material.*`) (`cli.py:702-710`). - -### Why it's a problem - -1. **The mapping is data, written as code.** Flag → dotted config path → fan-out rule is - a table. Written as 80 lines of copy-adapted dict comprehensions, correctness has to - be re-verified by reading, every time a flag is added. `giant/config.py:446` already - provides `_set_path(d, "stage1_model.wgan.n_critic", v)` — the primitive a table-driven - version needs. -2. **Adding a flag touches three places at once** (signature, the right comprehension, - the right precedence rule) with no mechanism forcing them to stay consistent. v0.3.0's - design explicitly anticipates more per-stage flags. -3. **It is the least-tested code in the repo** (Issue 4), because unit-testing a - 58-parameter Typer command means going through `CliRunner` with argv strings. -4. It obscures the ~15 lines that actually do something (`cli.py:825-839`): resolve the - out-dir, echo two lines, call `run_train_job`. - -### Recommended fix - -Extract the mapping into a pure, table-driven function in `giant/config.py` (not -`cli.py`), so it is importable and directly testable without Typer: - -```python -@dataclass(frozen=True) -class FlagSpec: - """One CLI flag's mapping into the config tree.""" - name: str # "stage1_hidden_dim" - paths: tuple[str, ...] # ("stage1_model.hidden_dim",) — >1 means fan-out - precedence: int = 0 # higher wins; per-stage flags outrank shorthands - -FLAG_SPECS: tuple[FlagSpec, ...] = ( - FlagSpec("hidden_dim", ("stage1_model.hidden_dim",), precedence=0), - FlagSpec("stage1_hidden_dim", ("stage1_model.hidden_dim",), precedence=1), - FlagSpec("mode", ("stage1_model.generator", "stage2_model.generator")), - FlagSpec("n_critic", ("stage1_model.wgan.n_critic", "stage2_model.wgan.n_critic")), - FlagSpec("emb_dim", ("conditioning.particle.emb_dim", "conditioning.material.emb_dim")), - ... -) - -def overrides_from_flags(values: dict[str, object]) -> dict: - """Build the nested overrides dict from {flag_name: value}, dropping - None (= flag not given) and applying precedence.""" -``` - -`train()` then becomes: parse `--batch-size auto`, collect `locals()`-style flag values, -call `overrides_from_flags`, `merge_cli_overrides`, `validate_config`, resolve out-dir, -call `run_train_job`. The 430-line Typer signature stays — that is irreducible, it *is* -the user interface — but it becomes the only bulk in the function. - -The precedence table also becomes self-documenting, which today requires reading three -separate inline comment blocks to reconstruct. - -### Scope guard - -- **Do not remove or rename any flag.** The backward-compatible shorthands - (`--hidden-dim`, `--n-blocks`, `--dropout`, `--mode`) exist because they predate the - per-stage flags and are in people's shell history and job scripts. Their surprising - stage-1-only scoping is documented behaviour, not a bug to fix. -- **Do not change precedence semantics.** Reproduce the current rules exactly, then - test them (Issue 4). Any intentional change is a separate commit. -- Preserve the explanatory comments at `cli.py:676-680` and `cli.py:713-719` — they - record *why* the precedence is the way it is. - ---- - -## Issue 4 — `cli.py` sits at 35.8 % coverage and holds untested override-precedence logic - -> **Status: Fixed for the override-precedence logic (Issue 3's scope); the -> `predict`/`rollout` inference-bootstrap portion described below is still open — that is -> Issue 5, deliberately not attempted here.** Commit `2bfb1ab` on -> `v0.3.0-stage2-autoregressive` adds direct, `CliRunner`-free unit tests for -> `overrides_from_flags` in `tests/test_config.py` — one per precedence rule, including -> several with previously **zero** coverage: both legs of `--stage{1,2}-generator` -> overriding `--mode` (only the stage1 leg had a test before), all three stage1 -> shorthand-vs-`--stage1-*` pairs (previously only `hidden_dim`), all four WGAN knobs' -> shared-vs-per-stage precedence (previously only `n_critic`/`gp_weight`), the -> `--emb-dim`/`--conditioning` dual-axis fan-out (previously untested via `train` at all), -> and an explicit regression test that `router_config` only ever writes -> `stage1_model.router`. `tests/test_cli_train_overrides.py` keeps its original 4 -> `CliRunner` smoke tests unmodified, plus 4 new ones covering the `--batch-size auto` -> parse-error and success paths and all three `out_dir` resolution branches -> (`--out`/`--resume`/default) — previously entirely uncovered. The percentage barely moves -> (36 % → 35 %), because the extraction *deleted* more statements from `cli.py` (539→478) -> than the new tests cover elsewhere in the file, but covered statements rose in absolute -> terms (157→169) and every override-precedence line the original issue called out by -> number is now covered. Coverage of `cli.py`'s other listed gap — the `predict`/ -> `rollout` inference bootstrap (`cli.py:1104-1423`/`1528-1702` in the pre-fix numbering) — -> is unchanged, since fixing that requires the `load_for_inference` extraction described in -> Issue 5, which is explicitly out of scope for this fix (done separately, if at all). -> Everything below this point describes the pre-fix state and is kept for historical -> context. - -**Severity: High. Effort: Medium.** - -**Location:** `giant/cli.py` (1865 lines). - -### What the code does today - -Per-module line coverage from `coverage.xml`: - -| Module | Coverage | -|---|---| -| `rollout.py` | 99.6 % | -| `analysis/catalog.py` | 99.6 % | -| `training/trainers.py` | 99.5 % | -| `model/network.py` | 98.6 % | -| `pipeline.py` | 97.1 % | -| `data/transforms.py` | 96.2 % | -| `config.py` | 91.8 % | -| **`cli.py`** | **35.8 %** | - -Project total is 81.8 %. `cli.py` is not merely the lowest — it is a **60-point outlier** -against an otherwise uniformly high standard, and it is simultaneously the largest module -in the repo. - -Existing CLI tests (`tests/test_cli_train_overrides.py`, `tests/test_cli_new_run.py`, -`tests/test_cli_predict.py`) are the right idea and well-written; they simply cover a -fraction of the surface. - -### Why it's a problem - -The uncovered code is not boilerplate. It includes: - -- The entire flag → config override translation and its precedence rules (Issue 3). Every - "`--stage2-hidden-dim` beats `--hidden-dim`", "`--stage2-generator` beats `--mode`", - "`--n-critic` fans out to both stages unless `--stage1-n-critic` is given" rule is - currently **unverified by any test**. -- The complete inference bootstrap in `predict` and `rollout` — checkpoint loading, - normalizer reconstruction, conditioning-axis resolution, top-N map wiring (Issue 5). - This is the code where a divergence between the two commands produces silent - train/inference skew rather than a crash. -- The `--batch-size auto` parsing and estimation paths, and the six distinct - "retrain with the current code" checkpoint-compatibility guards. - -A failure here is expensive in a way a unit-test failure is not: a mis-scoped flag means -a multi-hour GPU run trains a model that differs from the one the experiment log claims. -Given `CLAUDE.md` describes v0.3.0 as a sequence of architecture comparisons, this is the -worst possible place for silent divergence. - -### How to verify - -```bash -uv run pytest -q --cov=giant --cov-report=term-missing:skip-covered 2>&1 | grep "cli.py" -``` - -### Recommended fix - -Coverage here is a *consequence* of Issue 3, not an independent goal. Sequence the work: - -1. **First, extract the testable logic** (Issue 3): `overrides_from_flags` in - `config.py`, and `load_for_inference` in a new `giant/checkpoint_io.py` (Issue 5). -2. **Then test the extracted functions directly** — no Typer, no `CliRunner`, no - filesystem. One test per precedence rule: - - ```python - def test_stage1_hidden_dim_beats_hidden_dim_shorthand(): ... - def test_mode_sets_both_stage_generators(): ... - def test_stage2_generator_overrides_mode_for_stage2_only(): ... - def test_n_critic_fans_out_to_both_wgan_subtables(): ... - def test_stage1_n_critic_overrides_shared_n_critic_for_stage1_only(): ... - def test_emb_dim_sets_both_conditioning_axes(): ... - ``` - -3. **Keep a thin layer of `CliRunner` smoke tests** over `train --help`, `predict`, - `rollout` to catch signature/wiring breakage, but do not try to reach high coverage - through the CLI surface — that is slow and brittle. - -Target: get `cli.py` to roughly the repo norm by *moving code out of it*, not by writing -elaborate CLI-invocation tests. A `cli.py` that is genuinely just argument declaration -plus delegation can sit at modest coverage without concern, because there will be -nothing in it left to get wrong. - -### Scope guard - -Do not add `# pragma: no cover` to close the gap. The gap is a real signal and the fix is -extraction. - ---- - -## Issue 5 — The inference bootstrap is duplicated verbatim between `predict` and `rollout` - -> **Status: Fixed.** `giant/checkpoint_io.py` now holds a single -> `load_for_inference(checkpoint, device, command_name, weights="raw", require_stage2=True)` -> plus an `InferenceContext` dataclass (stage1/stage2 models, all three normalizers, both -> vocab maps, both top-N maps, both conditioning axes, `k_max`, both stages' ddpm step -> counts, `other_policy`, the raw `model_config`, and `epoch`/`best_val_loss` for -> `rollout`'s YAML sidecar) — exactly the design this issue proposed, verified against the -> current (not `55332db`-era) code before writing it. `predict`/`rollout` in `cli.py` -> each shrink to one `try/except CheckpointCompatibilityError` call plus a block of -> `ctx.` unpacks; `cli.py` lost `_conditioning_axes`, `_stage_cfg`, `_ddpm_steps`, -> `_particle_type_other_policy`, `_load_pdg_topn_map`/`_load_mat_topn_map`, and -> `_load_model_weights` entirely (net ~180 lines off `cli.py`). The independent third copy -> in `giant/analysis/router_gating.py` was deleted in favour of a lazy -> `from giant.checkpoint_io import conditioning_axes` inside `load_router`'s existing -> lazy-import block, preserving that module's "polars/numpy only at module scope" -> contract (`checkpoint_io.py` imports torch eagerly, so it must never be imported at -> `router_gating.py` module scope). Two guard orderings from the two commands' drifted -> copies were consolidated into one (`warn_if_checkpoint_config_mismatch` now always -> runs right after the existence guards, and `predict`'s `--batch-size auto` estimate now -> reads `ctx.model_config` after the checkpoint loads rather than before) — both are -> console-output-order changes only, no error text or model behavior changed, confirmed by -> diffing `giant predict --help`/`giant rollout --help` byte-for-byte before and after (no -> flag touched) and re-reading both rewritten command bodies field-by-field against the -> original. `require_stage2` exists as a real parameter, exercised by a new test, even -> though both current callers pass the default `True`. New `tests/test_checkpoint_io.py` -> (17 tests: happy path, every guard individually with exact message-text assertions, the -> `require_stage2=False`/inactive-stage2 path, `conditioning_axes`/`stage_cfg` directly) -> plus thin `CliRunner` smoke tests in `tests/test_cli_predict.py` and the new -> `tests/test_cli_rollout.py` confirming `CheckpointCompatibilityError` actually surfaces -> as `typer.Exit(1)` through the CLI — previously this entire code path had zero test -> coverage. `uv run pytest -q` (803 passed, up from 784), `ruff check`, `ruff format -> --check`, and `ty check` all clean. Everything below this point describes the pre-fix -> state and is kept for historical context. - -**Severity: High. Effort: Small. Risk if unfixed: silent train/inference skew.** - -**Location:** `giant/cli.py:1121-1201` (`predict`), `giant/cli.py:1534-1593` (`rollout`), -plus a third partial copy at `giant/analysis/router_gating.py:86-120`. - -### What the code does today - -`predict` and `rollout` each contain ~65 lines that are near-identical line-for-line: - -1. `ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)` -2. Guard: no `model_config` → "retrain with the current code" → `Exit(1)` -3. Guard: no `sec_decoder` → same -4. Guard: no `normalizer.sec_phys` → same -5. `_conditioning_axes(model_cfg)` -6. `_load_pdg_topn_map(ckpt)` / `_load_mat_topn_map(ckpt)` -7. Guard: `particle_conditioning == "onehot"` but no `pdg_topn_map` → `Exit(1)` -8. Guard: `material_conditioning == "onehot"` but no `mat_topn_map` → `Exit(1)` -9. `pdg_map` / `mat_map` key-type coercion (`{int(k): v ...}` / `{str(k): v ...}`) -10. Three `Normalizer.from_dict(ckpt["normalizer"][...])` calls -11. `build_models(model_cfg)`, unpack `stage1` / `stage2` -12. Guard: either stage `None` → "needs both" → `Exit(1)` -13. `_load_model_weights(...)`, `.to(device).eval()` on both -14. `typer.echo(f"loaded checkpoint: ...")` -15. `gconfig.warn_if_checkpoint_config_mismatch(checkpoint)` - -Additionally, `_conditioning_axes` is **duplicated verbatim into a second module**. -`giant/analysis/router_gating.py:69` carries a byte-for-byte copy of -`giant/cli.py:74`, with the docstring openly acknowledging it: *"Mirrors -`giant.cli._conditioning_axes`."* - -```bash -grep -rn "def _conditioning_axes" giant/ -# giant/cli.py:74 -# giant/analysis/router_gating.py:69 -``` - -### Why it's a problem - -This is not cosmetic duplication. This code decides **how input features are assembled at -inference time** — which conditioning mode is used, which top-N vocabulary maps are -applied, which normalizer statistics are restored. If `predict` and `rollout` ever -diverge on any of those, the result is not an exception. It is two commands producing -subtly different physics from the same checkpoint, with no error and no warning. That -class of bug is found by noticing that a plot looks wrong, weeks later. - -The duplication is also **entirely untested** — it lives in the 35.8 %-covered region of -`cli.py` (Issue 4). And it is *already* growing: the third copy in `router_gating.py` -shows the pattern spreading into a package that is otherwise carefully isolated. - -The self-aware "Mirrors `giant.cli._conditioning_axes`" comment is the tell. When a -developer documents a copy rather than removing it, the missing abstraction has been -identified but not yet built. - -### Recommended fix - -Create `giant/checkpoint_io.py` (name it whatever fits; the point is that it is **not** -`cli.py`) exposing one function and one result object: - -```python -@dataclass(frozen=True) -class InferenceContext: - """Everything needed to run a trained checkpoint forward, resolved once.""" - stage1: nn.Module - stage2: nn.Module - cond_norm: Normalizer - tgt_norm: Normalizer - sec_phys_norm: Normalizer - pdg_map: dict[int, int] - mat_map: dict[str, int] - pdg_topn_map: TopNMap | None - mat_topn_map: TopNMap | None - particle_conditioning: str - material_conditioning: str - k_max: int - stage1_ddpm_steps: int - stage2_ddpm_steps: int - other_policy: str - model_config: dict - -def load_for_inference( - checkpoint: Path, - device: torch.device, - weights: str = "raw", - require_stage2: bool = True, -) -> InferenceContext: ... -``` - -Raise a dedicated `CheckpointCompatibilityError` (carrying the current, genuinely helpful -message text) instead of calling `typer.echo` + `typer.Exit`; the CLI catches it and does -the echo/exit. That keeps the module free of Typer and makes it unit-testable. - -`predict` and `rollout` each shrink by ~60 lines to a single call. `router_gating.py` -drops its `_conditioning_axes` copy and imports from the new module — note that -`router_gating` deliberately imports torch **lazily inside the function** -(`router_gating.py:88`) to keep the analysis workers' "polars/numpy only" contract; the -new module must be imported the same way there, and `checkpoint_io.py` must not be -imported at `analysis/` module scope. - -### Scope guard - -- Preserve every error message verbatim. They are specific and actionable - ("checkpoint's `conditioning.particle.type='onehot'` but has no `pdg_topn_map` — retrain - with the current code"), and users have seen them before. -- Preserve the `require_stage2` distinction: `giant predict` and `giant rollout` both need - both stages today, but `stage1_model.active = false` / `stage2_model.active = false` are - real config options (`DEFAULT_CONFIG`), so the parameter should exist rather than - hard-coding the requirement. -- Do not fold `giant/analysis/router_gating.py`'s router-specific loading - (`load_router`, which returns `None` for non-MoE checkpoints) into the shared function. - It has different semantics — absence is a normal outcome there, not an error. - ---- - -## 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 -`giant/model/network.py:1446` (`_migrate_legacy_model_config`). - -### What the code does today - -v0.3.0 broke the config format: the v0.2 single `[train]` + `[model]` layout became -`[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]`. That break has to be -absorbed in two different places, and today it is absorbed by two unrelated functions: - -- **`config.migrate_config`** translates a v0.2 **`config.toml`** on load. -- **`network._migrate_legacy_model_config`** translates a v0.2 **checkpoint's - `model_config` dict** on `build_models`. - -The split is acknowledged as deferred work in `migrate_config`'s own docstring: - -> *"Operates on the config.toml shape. A checkpoint's `model_config` dict (which -> additionally carries n_sec_head ownership and needs `network.build_models`'s -> cooperation) is a separate migration surface, deferred to the network.py refactor."* - -Both functions independently encode the same translation facts: - -| v0.2 concept | v0.3 destination | in `migrate_config` | in `_migrate_legacy_model_config` | -|---|---|---|---| -| `train.mode` / `model.mode` | `stage{1,2}_model.generator` | ✓ | ✓ | -| `model.emb_dim` | `conditioning.{particle,material}.emb_dim` | ✓ | ✓ | -| `model.conditioning` | `conditioning.{particle,material}.type` | ✓ | ✓ | -| `model.n_blocks` | `stage{1,2}_model.n_res_blocks` | ✓ | ✓ | -| `model.noise_dim` | `stage{1,2}_model.wgan.noise_dim` | ✓ | ✓ | -| conditioning MLP depth was always 2 | `conditioning.*.n_layers = 2` | ✓ | ✓ | -| `router.expert_hidden_dim` set → hard error | — | ✓ | ✓ (near-identical message) | -| n_sec head lived on stage 1 | `stage2_model.n_sec.legacy_owner` | ✗ | ✓ (only here) | - -The `expert_hidden_dim` rejection is the clearest symptom: the same policy, the same -reasoning, two hand-maintained copies of a ~10-line error message -(`config.py:581-591` and `network.py:1475-1489`). - -### Why it's a problem - -1. **Drift.** A future correction to the v0.2 interpretation must be applied to both. - Applying it to one produces a checkpoint that loads with different architecture than - its own config file describes. -2. **`legacy_owner` leaks into the builder.** Because the checkpoint migration is - downstream of the config migration rather than sharing it, `build_models` has to carry - legacy-specific branching at four sites (`network.py:1612`, `1642`, `1663`, `1688`): - - ```python - legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner") - n_sec_head_k_max = s2cfg.get("k_max", K_MAX) if legacy_owner == "stage1" else None - ... - build_n_sec_head=legacy_owner != "stage1", - ``` - - A builder for the current architecture should not need to know where v0.2 put its - `n_sec` head. -3. **The retention policy is undeclared.** `tests/legacy/network_v02_snapshot.py` is - **1009 lines** of frozen v0.2 network code kept purely so migration can be tested - against it. That is a substantial maintenance surface with no stated expiry. - -### Recommended fix - -1. **Extract the shared translation into one table.** A single - `_V02_TO_V03_TRANSLATION` mapping (old dotted key → tuple of new dotted keys) plus one - set of "v0.2 architectural facts with no config key" constants, consumed by both - functions. `config.py:475-511` already has the beginnings of this - (`_V02_TRAIN_PASSTHROUGH`, `_V02_MODEL_TO_BOTH_STAGES`, - `_V02_TRAIN_TO_BOTH_STAGES_WGAN`) — extend that pattern and share it. Put the shared - table in a neutral module (e.g. `giant/_migration.py`) so neither `config.py` nor - `network.py` has to import the other. -2. **Move `legacy_owner` handling out of `build_models`.** The migration function should - emit a `model_config` that `build_models` can consume without legacy branching — for - example by emitting an explicit `stage2_model.n_sec.owner` key that the *current* - schema also carries (with value `"stage2"` for new runs), so the builder reads one key - with two valid values rather than a nullable legacy sentinel. -3. **Write down the retention policy.** Add to `CLAUDE.md`: which v0.2 checkpoints must - remain loadable, until when, and what triggers dropping the shim and the 1009-line - snapshot. Without that, nobody will ever feel authorised to delete it. - -### Scope guard - -- **v0.2 checkpoints must keep loading** until the policy above says otherwise. There are - trained models on `/ceph` that predate v0.3.0 and analysis runs referencing them. -- Keep `tests/legacy/network_v02_snapshot.py` as a **frozen snapshot** — do not "clean it - up", reformat it, or make it share code with current `network.py`. Its entire value is - that it is an independent, unchanging record of what v0.2 did. If ruff/ty complain about - it, exclude it rather than edit it (it is already excluded from coverage via - `pyproject.toml`'s `omit = ["*/legacy/*"]`). - ---- - -## Issue 7 — Positional tuple contracts between the data, model and training layers - -> **Status: Fixed.** `giant/data/transforms.py` now declares `StepFeatures(NamedTuple)` -> (same 9 fields, same order) and `build_features` returns it instead of a bare tuple; -> `giant/data/dataset.py` declares `StepBatch(NamedTuple)` (same 7 fields, same order) and -> `StreamingStepsDataset` yields it. `_batch_to_device` (`trainers.py`) is retyped -> `(batch: StepBatch, device) -> StepBatch` and reconstructs via `type(batch)(*(...))` so -> the concrete `NamedTuple` type — not just a plain tuple — survives `.to(device)`, -> matching what torch's own `pin_memory` already does for `NamedTuple`s (verified against -> the installed torch 2.3.1: NamedTuple batches survive `DataLoader` worker-process -> pickling and are reconstructed as the same type, whereas a plain tuple is silently -> downgraded to a `list` by `pin_memory`). Investigation while fixing this corrected two -> details in the original issue text: `_batch_to_device` only ever has one real shape (the -> 7-tuple, used identically at both call sites) — the "5-tuple at a second -> `_batch_to_device` call site" was a misreading, that 5-element tuple is a **derived** -> bundle built from 5 of the same 7 already-`_batch_to_device`'d fields -> (`WGANStageTrainer._stage2_real_and_fake`), now its own `_Stage2RealFakeBatch(NamedTuple)` -> (trainers-internal); and `giant/validate.py:120` turned out to be a third real consumer -> of the same contract the issue didn't list, converted to `StepBatch` attribute access. -> The already-fully-named positional unpacks in `trainers.py`'s `_compute`/`step` were -> left as positional tuple-unpacks (still valid against a `NamedTuple`, per the issue's own -> "migration incremental" framing) — only the sites with heavy `_`-throwaways or bare -> unlabelled unpacks (`dataset.py`'s `build_features` consumption, `cli.py:980`, -> `pipeline.py:251`, `validate.py:120`) were converted to attribute access. Two test -> batch-construction helpers (`tests/test_train.py::_fake_batches`, -> `tests/test_validate.py::_loader`) now build real `StepBatch`s instead of plain tuples, -> which is what let `_batch_to_device` stay a simple, non-defensive one-liner. Verified live -> (then reverted) that inserting a field into `StepBatch` now makes `ty check` fail with -> "Too many values to unpack" at every positional-unpack call site and -> "No argument provided for required parameter" at every keyword-constructed one — exactly -> the protection this issue asked for, and precisely the failure mode that was previously -> invisible. No field was reordered anywhere. `uv run pytest -q` (804 passed), `ruff check`, -> `ruff format --check`, and `ty check` all clean. Everything below this point describes the -> pre-fix state and is kept for historical context. - -**Severity: Medium. Effort: Small.** - -**Location:** `giant/data/transforms.py:863-890` (`build_features`), -`giant/data/dataset.py:229-237` (`StreamingStepsDataset` yield), -`giant/training/trainers.py:574-582` (unpack), `giant/cli.py:1233` (unpack). - -### What the code does today - -`build_features` returns a **bare 9-tuple**: - -```python -) -> tuple[ - np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, - np.ndarray, np.ndarray, Normalizer | None, Normalizer | None, -]: - """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, - sec_type_idx) arrays. ...""" -``` - -The element *meanings* live only in the docstring. Call sites re-derive them positionally: - -```python -# giant/cli.py:1233 — predict -cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(...) -``` - -`StreamingStepsDataset` yields a **7-tuple** of tensors (`dataset.py:229`), unpacked -positionally in the trainer: - -```python -# giant/training/trainers.py:574 -( - cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx, sec_type_idx, -) = _batch_to_device(batch, device) -``` - -and again, with a *different arity*, in the WGAN path: - -```python -# giant/training/trainers.py:737 -cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors -``` - -`_batch_to_device` is typed `(batch: tuple, device) -> tuple` — no element information at -all. - -### Why it's a problem - -1. **Invisible to the type checker.** `ty check` passes on all of it. Reordering two - `np.ndarray` elements — or inserting a new one in the middle — type-checks cleanly and - produces a model trained on scrambled features. Several of these arrays are - shape-compatible with each other (`n_sec`, `proc_idx` and `sec_type_idx` are all - integer arrays), so a swap may not even produce a shape error at runtime. -2. **This is the hottest contract in the codebase.** It crosses three layer boundaries - (`data` → `training`, `data` → `cli`) and is the mechanism by which the model receives - physics. A silent corruption here is the most expensive possible failure: it does not - crash, it just trains something wrong. -3. **The nine-underscore unpack** at `cli.py:1233` is unreadable and does not survive any - change to the tuple. -4. **Two different batch arities** for the same conceptual batch (7 in `_compute`, 5 in - `_stage2_real_and_fake`) means the reader must track which slice is in play. - -### Recommended fix - -Convert both to `NamedTuple`. Zero runtime cost, full `ty` visibility, tuple-unpacking -still works so migration is incremental: - -```python -class StepFeatures(NamedTuple): - """Output of build_features. Field order is load-bearing for existing - positional unpacking — append only, never insert or reorder.""" - cond_cont: np.ndarray - cond_cat: np.ndarray - target_s1: np.ndarray - n_sec: np.ndarray - sec_cont: np.ndarray - proc_idx: np.ndarray - sec_type_idx: np.ndarray - cond_normalizer: Normalizer | None - target_normalizer: Normalizer | None - -class StepBatch(NamedTuple): - cond_cont: torch.Tensor - cond_cat: torch.Tensor - target_s1: torch.Tensor - n_sec: torch.Tensor - sec_cont: torch.Tensor - proc_idx: torch.Tensor - sec_type_idx: torch.Tensor -``` - -Then `cli.py:1233` becomes `feats = build_features(...)` / `feats.cond_cont`, and -`_batch_to_device` gets the real signature `(batch: StepBatch, device) -> StepBatch`. - -Precedent exists in this codebase: `RolloutSummary` (`rollout.py:276`, a `TypedDict`) and -`SetupStageResult` (`pipeline.py:39`, a dataclass) already do this correctly. - -Note that `StreamingStepsDataset` passes batches through a PyTorch `DataLoader` with -`batch_size=None` and `num_workers > 0`, so the batch type must survive worker-process -pickling and `pin_memory`. `NamedTuple` does — it is a plain tuple subclass, and -`pin_memory` recurses into tuples — but **verify this with a real multi-worker run**, not -just the test suite, since the tests may run single-process. - -### Scope guard - -Do not reorder any existing fields while converting. The whole point of the change is to -make future reordering safe; performing one during the conversion, when nothing yet -protects against it, is the single riskiest version of this change. - ---- - -## Issue 8 — `network.py` is 1745 lines holding three distinct modules - -> **Status: Fixed.** `giant/model/network.py` (1742 lines) is now split exactly along the -> seams this issue identified: `giant/model/layers.py` (`SinusoidalEmbedding`, `ResBlock`, -> `ContextAdapter`, `_make_axis_mlp`), `encoders.py` (`cat_col_layout`, `ConditionEncoder`), -> `routers.py` (`Router` + registry + all 4 router types + composed/config parsing — -> self-contained, no dependency on any other new submodule), `trunks.py` (`Trunk`, -> `MonolithicTrunk`, `RoutedTrunk`, `ExpertTrunk`, `build_trunk`), `history.py` -> (`HistoryEncoder`, `MarkovHistory`, `AttentionHistory`, `_CausalAttnBlock` — -> self-contained), `models.py` (`stage2_type_dim`, `stage2_trunk_sec_dim`, `Stage1Model`, -> `Stage2OneShot`, `Stage2Autoregressive`, `CriticModel`), `_legacy.py` -> (`_migrate_legacy_model_config`, `migrate_legacy_state_dict` — pure dict/state-dict -> translation, no dependency on the model classes), and `builders.py` (`build_models`, -> `build_critics`). `network.py` itself is now an 83-line re-export shim (`__all__` listing -> all 38 public + underscore-prefixed names any current caller reaches into), so none of -> the 22 external `from giant.model.network import ...` call sites (`sample.py`, -> `pipeline.py`, `checkpoint_io.py`, `training/trainers.py`, -> `analysis/router_gating.py`, and ~10 test files, several importing private names like -> `_build_router_from_cfg`/`_parse_composed_axes`) needed to change. Pure file move: no -> logic, signatures, or behavior changed — each new module carries only the imports it -> actually uses rather than a copy of the original file's full import block. `uv run -> pytest -q` (804 passed, unchanged), `uv run ruff check .`, `uv run ruff format --check -> .`, and `uv run ty check .` all clean. Everything below this point describes the -> pre-fix state and is kept for historical context. - -**Severity: Medium. Effort: Small (mechanical).** - -**Location:** `giant/model/network.py`. - -### What the code does today - -One file contains six unrelated concerns: - -| Lines (approx.) | Concern | -|---|---| -| 28-213 | Primitives: `SinusoidalEmbedding`, `cat_col_layout`, `_make_axis_mlp`, `ConditionEncoder`, `ContextAdapter`, `ResBlock` | -| 215-575 | **The entire router subsystem**: `Router` base, `register_router`/`build_router` registry, `EnergyRouter`, `PdgRouter`, `ProcessRouter`, `ComposedRouter`, `build_composed_router`, `_parse_composed_axes`, `_check_router_conditioning_compat`, `_build_router_from_cfg` | -| 577-727 | Trunks: `ExpertTrunk`, `_route_forward`, `Trunk`, `MonolithicTrunk`, `RoutedTrunk`, `build_trunk` | -| 729-881 | **History encoders**: `HistoryEncoder`, `MarkovHistory`, `_CausalAttnBlock`, `AttentionHistory` | -| 883-1443 | Top-level models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`, `CriticModel` | -| 1446-1568 | **Legacy migration**: `_migrate_legacy_model_config`, `migrate_legacy_state_dict` | -| 1570-1745 | Builders: `build_models`, `build_critics` | - -The v0.3.0 refactor described in `CLAUDE.md` — *"`network.py` is refactored from ten -permutation classes into composable parts (encoder × trunk × objective)"* — clearly -landed and is a genuine improvement. The composition is visible and clean. The file -simply was not split to match. - -### Why it's a problem - -- The seams are already there in the class layout; the file just does not honour them. A - reader looking for `AttentionHistory` has no reason to expect it 800 lines into a file - whose name suggests "the network". -- The router subsystem alone is ~370 lines with its own registry, its own config parsing, - and its own compatibility validation. It is a subsystem, not a section. -- Legacy migration (Issue 6) sitting in the same file as the current builders is exactly - what lets `legacy_owner` bleed into `build_models`. -- It is a merge-conflict magnet on a repo with parallel feature branches - (`condor-gpu-train-rollout`, `v0.3.0-stage2-autoregressive`). - -### Recommended fix - -Split along the existing seams into `giant/model/`: - -``` -giant/model/ - layers.py # SinusoidalEmbedding, ResBlock, ContextAdapter, _make_axis_mlp - encoders.py # ConditionEncoder, cat_col_layout - routers.py # Router base + registry + all 4 router types + config parsing - trunks.py # Trunk, MonolithicTrunk, RoutedTrunk, ExpertTrunk, build_trunk - history.py # HistoryEncoder, MarkovHistory, AttentionHistory, _CausalAttnBlock - models.py # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel - builders.py # build_models, build_critics - _legacy.py # _migrate_legacy_model_config, migrate_legacy_state_dict - network.py # re-export shim: `from giant.model.layers import *` etc. -``` - -Keep `network.py` as a **re-export shim** so no import site outside `giant/model/` has to -change in the same commit. `network.py` is imported by `cli.py`, `sample.py`, -`trainers.py`, `pipeline.py`, `analysis/router_gating.py`, and heavily by the test suite — -`grep -rn "from giant.model.network import" giant/ tests/ | wc -l` before starting, to -size the blast radius if you do decide to update call sites (21 sites on `55332db`). - -This is a pure file-move refactor: no logic changes, and the 725-test suite plus `ty -check` is a strong safety net for it. - -### Scope guard - -Do it as its own commit, containing **only** moves and imports. Do not combine with Issue 1 -or Issue 6, both of which change behaviour in this file — a diff mixing moves with logic -changes is effectively unreviewable. - ---- - -## Issue 9 — `scripts` is published as a top-level distribution package - -> **Status: Fixed.** `scripts/` moved to `giant/tools/` via `git mv` (preserving file -> history); `pyproject.toml`'s `[project.scripts]` entry point is now -> `dwarf = "giant.tools.dwarf:app"`, `[tool.hatch.build.targets.wheel] packages` is now -> `["giant"]` only, and `[tool.coverage.run] source` is now `["giant"]` only. The 7 -> `from scripts.X import Y` intra-package imports in `dwarf.py` became -> `from giant.tools.X import Y`; the one live (non-cosmetic) reference to the old module -> path — `steps_to_parquet_parallel.py`'s `_DWARF_CONVERT_CMD` subprocess invocation -> (`python -m scripts.dwarf convert`, used by every `dwarf convert --jobs N` fan-out -> worker) — is now `python -m giant.tools.dwarf`. The 5 test files that imported from -> `scripts` (`test_dwarf.py`, `test_steps_to_parquet.py`, -> `test_steps_to_parquet_parallel.py`, `test_bump_dataset_version.py`, -> `test_create_root_files.py`) now import from `giant.tools`. All remaining prose -> references to `scripts/`/`scripts.` — in `giant/cli.py`, `giant/config.py`, -> `giant/data/loader.py`, `giant/analysis/runtime_estimate.py`, `tests/test_dataset.py`, -> `CLAUDE.md`, `README.md`'s repo-tree block, and docstrings/comments inside the moved -> files themselves (`dwarf.py`, `profile_analysis_costs.py`, `check_migration_v02_v03.py`, -> `bump_dataset_version.py`, `create_root_files.py`) — were updated to `giant/tools/` / -> `giant.tools.` for accuracy. `giant/analysis/condor.py` and -> `.gitea/workflows/ci.yml` were checked per the scope guard below and needed no changes -> (neither hard-codes a `scripts/` path; CI's coverage source comes from `pyproject.toml`). -> The `dwarf` **command name is unchanged** — `giant train --help`/`dwarf --help` behavior -> is identical, only the Python import path and file location moved. `uv sync --extra cpu -> --extra dev` reinstalls cleanly with the new entry point; `uv run pytest -q` (804 -> passed, unchanged), `uv run ruff check .`, `uv run ruff format --check .`, and `uv run ty -> check .` all clean. Everything below this point describes the pre-fix state and is kept -> for historical context. - -**Severity: Medium. Effort: Small.** - -**Location:** `pyproject.toml`: - -```toml -[project.scripts] -giant = "giant.cli:app" -dwarf = "scripts.dwarf:app" - -[tool.hatch.build.targets.wheel] -packages = ["giant", "scripts"] -``` - -### What the code does today - -The repo's tooling CLI (`dwarf`) lives in a directory called `scripts/`, which is declared -as a wheel package and referenced by the console-script entry point as `scripts.dwarf:app`. -Installing `giant` therefore creates a top-level importable module named **`scripts`** in -`site-packages`. - -`scripts/` holds ten real modules: `dwarf.py`, `bump_dataset_version.py`, -`create_root_files.py`, `geometry_oracle.py`, `hparam_scan.py`, `migrate_geant_steps.py`, -`profile_analysis_costs.py`, `steps_to_parquet.py`, `steps_to_parquet_parallel.py`, -`warm_setup_cache.py`. - -### Why it's a problem - -`scripts` is one of the most generic names possible in the Python ecosystem. Consequences: - -1. **Collision.** Any other installed distribution that also ships a top-level `scripts` - package silently shadows or is shadowed by this one, depending on `sys.path` order. - The failure mode is an `ImportError` or — worse — importing someone else's `dwarf`-less - `scripts` and getting `AttributeError` at CLI startup. -2. **Environment-order fragility.** On the portal machines (`/work/lbogner`, shared with - other users, per `CLAUDE.md`), a stray `scripts/` directory in the CWD shadows the - installed package, because CWD precedes `site-packages` on `sys.path`. Running `dwarf` - from a directory that happens to contain a `scripts/` folder can break in a confusing - way. -3. **`scripts/` reads as "not part of the product"**, yet it is installed, has an entry - point, is covered by `[tool.coverage.run] source`, and has a full test suite - (`tests/test_dwarf.py` at 81.7 % coverage, plus `test_steps_to_parquet*.py`, - `test_create_root_files.py`, `test_bump_dataset_version.py`). Its name misrepresents - its status. - -### Recommended fix - -Move it under the `giant` namespace, where it cannot collide: - -``` -giant/tools/ # was scripts/ - dwarf.py - ... -``` - -```toml -[project.scripts] -giant = "giant.cli:app" -dwarf = "giant.tools.dwarf:app" - -[tool.hatch.build.targets.wheel] -packages = ["giant"] -``` - -Then update: - -- `[tool.coverage.run] source = ["giant"]` (drop the now-redundant `"scripts"`). -- All `from scripts.X import Y` in `tests/` (`grep -rn "scripts\." tests/`). -- Any `python scripts/foo.py` invocations in docs, `README.md`, `CLAUDE.md`, HTCondor - submit files, or shell history/job scripts on the portal machines. - -The `dwarf` **command name does not change**, so anything invoking `dwarf ...` keeps -working — only Python-level imports and direct `python scripts/...` paths move. - -### Scope guard - -Check `giant/analysis/condor.py` and any generated HTCondor submit descriptions for -hard-coded `scripts/` paths before moving. Jobs submitted to the cluster may reference the -path as it exists on `/work` or `/ceph`, and a rename that lands mid-flight breaks queued -jobs. Search for the literal string `scripts/` across the repo, not just Python imports. - ---- - -## Issue 10 — `torch.load(weights_only=False)`: checkpoints are arbitrary pickles - -**Severity: Low (given the threat model). Effort: Medium.** - -**Location:** `giant/cli.py:1122`, `giant/cli.py:1535`, -`giant/analysis/router_gating.py:93`, `giant/training/loop.py:153`. - -### What the code does today - -All four checkpoint loads pass `weights_only=False`: - -```python -ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) -``` - -This is *necessary* today, not careless: the checkpoint dict carries non-tensor objects -beyond raw state dicts — `model_config`, `pdg_map` / `mat_map`, `pdg_topn_map` / -`mat_topn_map`, and the `normalizer` sub-dict, assembled in -`giant/training/checkpoint.py:19` (`build_checkpoint`) from the `extras` argument. -`weights_only=True` would reject them. - -### Why it's worth recording - -`weights_only=False` means loading a checkpoint executes arbitrary pickle opcodes. The -practical threat model here is mild — checkpoints are produced by this codebase and live -on `/ceph/lbogner` — but the portal machines and `/ceph` are explicitly **shared with -other users** (`CLAUDE.md`, "Compute environment"), and `giant predict --checkpoint ` -will happily load a path someone else wrote. It is also the kind of thing that becomes a -blocker later, if a model is ever shared outside the group or published alongside a paper. - -Secondary practical cost: a pickled `model_config` cannot be inspected without importing -torch and unpickling. `dwarf status`-style tooling, or a human answering "what conditioning -mode was this trained with?", has to load the whole checkpoint. - -### Recommended fix - -Split the checkpoint into a tensor part and a JSON sidecar: - -- `ckpt.pt` — state dicts only, loadable with `weights_only=True`. -- `ckpt.meta.json` — `model_config`, vocab maps, top-N maps, normalizer statistics, - `epoch` / `global_step` / `best_val_loss`. - -This has real secondary benefits: the sidecar is greppable and diffable, so comparing two -runs' architectures becomes `diff` rather than a Python session, and it composes well -with the config-provenance machinery already in `config.py` -(`warn_if_checkpoint_config_mismatch`, `git_hash`). - -Note that `giant/data/setup_cache.py` already establishes the "JSON sidecar next to the -data" pattern for exactly this kind of non-tensor state — follow its conventions. - -### Scope guard - -This is a **format change**, so it needs a compatibility path: `load_for_inference` -(Issue 5) reads the sidecar if present and falls back to the pickled keys if not. Do not -attempt this before Issue 5 lands — with three separate copies of the loading code, a -format change means three separate compatibility paths. After Issue 5, it is one. - -Given the mild threat model, this is correctly the **lowest-priority** item here. It is -recorded so the decision is deliberate rather than accidental. - ---- - -## Issue 11 — Minor items - -### 11a — `echo=print` threaded through the pipeline instead of `logging` - -**Location:** `giant/pipeline.py:305` (`run_train_job(..., echo=print)`), -`giant/pipeline.py:86` (`run_setup_stage(..., echo=...)`); 52 `typer.echo` calls in -`cli.py`. - -Passing the output function as a parameter is *better* than hard-coding `typer.echo` deep -in the pipeline — it keeps `pipeline.py` free of Typer and makes output capturable in -tests. But at this scale it has costs: no severity levels (the `--num-workers` quota -warning at `pipeline.py:329` and routine progress output are indistinguishable), no -timestamps on long training runs, no way to route to a file without wrapping, and the -parameter has to be threaded through every function that might print. - -Suggested: standard `logging` with a Typer/rich handler configured once in `cli.py`. -Module-level `log = logging.getLogger(__name__)` replaces the threaded `echo`. Keep -`typer.echo` for genuine CLI output (the `device:` / `out_dir:` banners, error messages -before `Exit(1)`) — those are interface, not logs. - -Low priority; the current approach is defensible. Worth doing opportunistically if -`pipeline.py` is being touched anyway. - -### 11b — `giant.particles` imports from `giant.data.loader` - -**Location:** `giant/particles.py:30`. - -```python -if TYPE_CHECKING: - from giant.data.loader import TopNMap -``` - -`giant/particles.py` is a physics-domain module (PDG code decoding, mass/charge lookup). -`giant/data/loader.py` is data I/O. A domain module depending on an I/O module is a -layering inversion. - -Mitigating factors, which is why this is minor: the import is `TYPE_CHECKING`-guarded, so -there is **no runtime cycle**, and `ty check` passes. It is a design smell, not a defect. - -Suggested: move `TopNMap` to a neutral location — `giant/constants.py` or a new -`giant/types.py` — so both `particles.py` and `data/loader.py` depend on it rather than on -each other. Do this only if `TopNMap` is being touched for other reasons; it is not worth -a standalone commit. - ---- - -## Recommended sequence - -The issues are interdependent. This order minimises rework: - -1. **Issue 5** — extract `load_for_inference` into `giant/checkpoint_io.py`. Small, - self-contained, removes the highest-risk duplication, and unblocks Issue 10. -2. **Issue 3** — extract the flag→config mapping table into `config.py`. -3. **Issue 4** — test the two extracted units directly. Issues 3 and 5 make this cheap; - attempting it first means testing through `CliRunner`, which is slow and brittle. -4. **Issue 1 + Issue 2 together** — typed config dataclasses. One migration solves both, - and doing them separately means touching the same 132 call sites twice. -5. **Issue 6** — unify the migration surfaces (easier once config is typed). -6. **Issue 8** — split `network.py`. Pure file moves; do it as an isolated commit, after - the logic changes in Issues 1 and 6 have settled, to keep both diffs readable. -7. **Issue 7** — `NamedTuple` the feature/batch contracts. -8. **Issue 9** — move `scripts/` → `giant/tools/`. Independent of everything else; can be - done at any point. -9. **Issue 10** — checkpoint format split, if the threat model or a publication makes it - worthwhile. -10. **Issue 11** — opportunistically. - -## General guidance for whoever picks this up - -- **Run the full suite on every change.** `uv run pytest -q` takes 18 seconds. There is no - reason to batch up unverified changes. -- **Run all three checks before committing:** `uv run ruff check .`, - `uv run ruff format .`, `uv run ty check .`. CI enforces all three. -- **Do not change model behaviour while refactoring.** Several issues touch code that - determines what network gets built (Issues 1, 6) or what features it receives (Issues 5, - 7). A refactor that also changes a default silently invalidates every prior benchmark in - `/home/lars/knowledge-base/experiments/`. If a behaviour change is warranted, it is a - separate commit with a separate message saying so. -- **Preserve the comments.** This codebase's docstrings explain reasoning, trade-offs, and - known limitations to an unusually high standard. When code moves, the comments move with - it. When code is deleted, check whether the comment records a decision that still needs - recording elsewhere. -- **`giant/analysis/` is the template.** It is declarative where the rest of the codebase - is imperative (`PlotSpec` registry), explicit about its own contracts (the - `compute_partial` / `finalize` split, `chunkable=False`, the "polars/numpy only on - workers" boundary), and honest about its one exception (`router_gating.py`'s module - docstring explains precisely why it is allowed to import torch and why that is still - safe). When deciding how a refactored `config.py` or `cli.py` should look, look there.