Records what commit 2bfb1ab actually changed and its scope, matching the
status-blockquote convention already used for Issues 1 and 2.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
60 KiB
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/HACKmarkers ingiant/orscripts/. - The internal import graph is acyclic with clean layering:
constants → data → model → training → pipeline → cli, andanalysisalmost fully independent of the rest. - Line coverage is 96–99.6 % on every core module (
network.py98.6 %,trainers.py99.5 %,rollout.py99.6 %,transforms.py96.2 %,pipeline.py97.1 %,analysis/catalog.py99.6 %). The single exception iscli.pyat 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'sDEFAULT_CONFIGcomments orgiant/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 | Open |
| 6 | Two independent v0.2→v0.3 migration surfaces encode the same knowledge | Medium | Medium | Open |
| 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 |
| 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
9bf5874onv0.3.0-stage2-autoregressive.giant/config.pynow declares a hierarchy of frozen dataclasses (GiantConfigand its nested blocks, followingStageSpec's existing style) as the single source of truth;DEFAULT_CONFIGis generated fromGiantConfig().to_dict()rather than hand-maintained, andbuild_models/build_critics(network.py) andStageSpec.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 toDEFAULT_CONFIG's documented v0.3.0 values ("autoregressive"/"onehot") unconditionally. Router/n_secsub-blocks keep a dict-shapedextracatch-all for their genuinely dynamic keys (composed-routeraxis{i}_*, runtime-seededcenters_init, legacylegacy_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 — atests/test_train.pyfixture andscripts/warm_setup_cache.py's hand-built minimal config (now merged againstDEFAULT_CONFIGinstead of hand-rolled) — both fixed in the same commit. Regression tests added intests/test_config.py,tests/test_network.py,tests/test_train.py, including one pinningGiantConfig().to_dict() == DEFAULT_CONFIGso 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):
# 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),
# 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:
- The training path always passes a fully-merged config.
giant/pipeline.py:446buildsmodel_configfromcfg["stage1_model"]/cfg["stage2_model"], andcfgalways originates frommerge_cli_overrides(DEFAULT_CONFIG, ...), so every key is present and no fallback fires. - 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. - 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
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.
- 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 existingDEFAULT_CONFIGcomments onto the fields — they are the most valuable thing in that file and must not be lost. - 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). - Change
build_models,build_criticsandStageSpec.from_configto take the dataclasses instead of dicts. Every.get(key, default)becomescfg.key. The 132 duplicate defaults disappear by construction. - Keep
DEFAULT_CONFIGas 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 viadataclasses.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:
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.pynow hasvalidate_config_keys(cfg), which walks a merged config dict recursively againstDEFAULT_CONFIG's tree (itself generated from theGiantConfigdataclasses added for Issue 1) and raisesValueErroron any key not present there, with adifflib-based "did you mean" suggestion.[meta]is skipped unconditionally, andstage{1,2}_model.router'saxis{i}_{field}composed-router keys and runtime-seededcenters_initare allowed through explicitly.merge_cli_overridescalls it right before returning, so all three real entry points (giant train,giant new-run,scripts/warm_setup_cache.py) are covered automatically; checkpointmodel_configloading (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 touchingvalidate_config's many direct unit-test callers orrun_train_job's redundant internalvalidate_configcall. Regression tests added intests/test_config.pycover top-level and nested typos, the axis/centers_initallowances,[meta]skipping, and both the file and CLI-overrides paths, plus the two realconfigs/*.tomlfixtures that don't already failmigrate_configfor 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:
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
[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
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:
-
Preferred — fold into Issue 1. If each config block becomes a dataclass with a
from_dictthat raises on unknown keys, this issue is solved for free and cannot regress. Raise aValueErrorin 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 keyn_res_block; did you meann_res_blocks?" is the message that actually saves the run. -
Standalone fallback. Add
validate_config_keys(cfg)togiant/config.pythat walkscfgalongsideDEFAULT_CONFIGand raises on any key not present in the defaults tree. Call it frommerge_cli_overridesright 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— holdsconfig_version,git_hash, and run metadata frombuild_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 fromDEFAULT_CONFIG(see the comment at the end of therouterblock inconfig.py, and_parse_composed_axesinnetwork.py:499). Validate these against theaxis(\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 checkpointmodel_configdicts, not inconfig.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
2bfb1abonv0.3.0-stage2-autoregressive.giant/config.pynow declaresFlagSpec(a frozen dataclass:name,paths— one or more dotted config paths,precedence) and aFLAG_SPECStable covering every flagtrain/new-runmap into the config tree, plusoverrides_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-790pre-fix) is now a single flatflag_valuesdict (mostly enum.valueunwrapping) plus one call tooverrides_from_flags;new_run()'s near-verbatim copy (cli.py:915-983pre-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-aggregatedrouter_configentry mapped only tostage1_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 --helpare 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:
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:
--modesetsgeneratoron both stages;--stage1-generator/--stage2-generatorthen override a single stage (cli.py:713-720).--hidden-dim/--n-blocks/--dropoutare stage-1-only backward-compatible shorthands;--stage1-*wins when both are given (cli.py:676-700).--n-critic/--gp-weight/--noise-dim/--critic-lrfan out tostage{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-dimand--conditioningeach set two places (conditioning.particle.*andconditioning.material.*) (cli.py:702-710).
Why it's a problem
- 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:446already provides_set_path(d, "stage1_model.wgan.n_critic", v)— the primitive a table-driven version needs. - 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.
- It is the least-tested code in the repo (Issue 4), because unit-testing a
58-parameter Typer command means going through
CliRunnerwith argv strings. - It obscures the ~15 lines that actually do something (
cli.py:825-839): resolve the out-dir, echo two lines, callrun_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:
@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-680andcli.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/rolloutinference-bootstrap portion described below is still open — that is Issue 5, deliberately not attempted here. Commit2bfb1abonv0.3.0-stage2-autoregressiveadds direct,CliRunner-free unit tests foroverrides_from_flagsintests/test_config.py— one per precedence rule, including several with previously zero coverage: both legs of--stage{1,2}-generatoroverriding--mode(only the stage1 leg had a test before), all three stage1 shorthand-vs---stage1-*pairs (previously onlyhidden_dim), all four WGAN knobs' shared-vs-per-stage precedence (previously onlyn_critic/gp_weight), the--emb-dim/--conditioningdual-axis fan-out (previously untested viatrainat all), and an explicit regression test thatrouter_configonly ever writesstage1_model.router.tests/test_cli_train_overrides.pykeeps its original 4CliRunnersmoke tests unmodified, plus 4 new ones covering the--batch-size autoparse-error and success paths and all threeout_dirresolution branches (--out/--resume/default) — previously entirely uncovered. The percentage barely moves (36 % → 35 %), because the extraction deleted more statements fromcli.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 ofcli.py's other listed gap — thepredict/rolloutinference bootstrap (cli.py:1104-1423/1528-1702in the pre-fix numbering) — is unchanged, since fixing that requires theload_for_inferenceextraction 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-dimbeats--hidden-dim", "--stage2-generatorbeats--mode", "--n-criticfans out to both stages unless--stage1-n-criticis given" rule is currently unverified by any test. - The complete inference bootstrap in
predictandrollout— 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 autoparsing 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
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:
-
First, extract the testable logic (Issue 3):
overrides_from_flagsinconfig.py, andload_for_inferencein a newgiant/checkpoint_io.py(Issue 5). -
Then test the extracted functions directly — no Typer, no
CliRunner, no filesystem. One test per precedence rule: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(): ... -
Keep a thin layer of
CliRunnersmoke tests overtrain --help,predict,rolloutto 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
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:
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)- Guard: no
model_config→ "retrain with the current code" →Exit(1) - Guard: no
sec_decoder→ same - Guard: no
normalizer.sec_phys→ same _conditioning_axes(model_cfg)_load_pdg_topn_map(ckpt)/_load_mat_topn_map(ckpt)- Guard:
particle_conditioning == "onehot"but nopdg_topn_map→Exit(1) - Guard:
material_conditioning == "onehot"but nomat_topn_map→Exit(1) pdg_map/mat_mapkey-type coercion ({int(k): v ...}/{str(k): v ...})- Three
Normalizer.from_dict(ckpt["normalizer"][...])calls build_models(model_cfg), unpackstage1/stage2- Guard: either stage
None→ "needs both" →Exit(1) _load_model_weights(...),.to(device).eval()on bothtyper.echo(f"loaded checkpoint: ...")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."
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:
@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 nopdg_topn_map— retrain with the current code"), and users have seen them before. - Preserve the
require_stage2distinction:giant predictandgiant rolloutboth need both stages today, butstage1_model.active = false/stage2_model.active = falseare 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 returnsNonefor 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
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_configtranslates a v0.2config.tomlon load.network._migrate_legacy_model_configtranslates a v0.2 checkpoint'smodel_configdict onbuild_models.
The split is acknowledged as deferred work in migrate_config's own docstring:
"Operates on the config.toml shape. A checkpoint's
model_configdict (which additionally carries n_sec_head ownership and needsnetwork.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
-
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.
-
legacy_ownerleaks into the builder. Because the checkpoint migration is downstream of the config migration rather than sharing it,build_modelshas to carry legacy-specific branching at four sites (network.py:1612,1642,1663,1688):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_sechead. -
The retention policy is undeclared.
tests/legacy/network_v02_snapshot.pyis 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
- Extract the shared translation into one table. A single
_V02_TO_V03_TRANSLATIONmapping (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-511already 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 neitherconfig.pynornetwork.pyhas to import the other. - Move
legacy_ownerhandling out ofbuild_models. The migration function should emit amodel_configthatbuild_modelscan consume without legacy branching — for example by emitting an explicitstage2_model.n_sec.ownerkey 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. - 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
/cephthat predate v0.3.0 and analysis runs referencing them. - Keep
tests/legacy/network_v02_snapshot.pyas a frozen snapshot — do not "clean it up", reformat it, or make it share code with currentnetwork.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 viapyproject.toml'somit = ["*/legacy/*"]).
Issue 7 — Positional tuple contracts between the data, model and training layers
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:
) -> 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:
# 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:
# 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:
# 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
- Invisible to the type checker.
ty checkpasses on all of it. Reordering twonp.ndarrayelements — 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_idxandsec_type_idxare all integer arrays), so a swap may not even produce a shape error at runtime. - 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. - The nine-underscore unpack at
cli.py:1233is unreadable and does not survive any change to the tuple. - 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:
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
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
AttentionHistoryhas 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_ownerbleed intobuild_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
Severity: Medium. Effort: Small.
Location: pyproject.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:
- Collision. Any other installed distribution that also ships a top-level
scriptspackage silently shadows or is shadowed by this one, depending onsys.pathorder. The failure mode is anImportErroror — worse — importing someone else'sdwarf-lessscriptsand gettingAttributeErrorat CLI startup. - Environment-order fragility. On the portal machines (
/work/lbogner, shared with other users, perCLAUDE.md), a strayscripts/directory in the CWD shadows the installed package, because CWD precedessite-packagesonsys.path. Runningdwarffrom a directory that happens to contain ascripts/folder can break in a confusing way. 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.pyat 81.7 % coverage, plustest_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
...
[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 Yintests/(grep -rn "scripts\." tests/). - Any
python scripts/foo.pyinvocations 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:
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 <path>
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 withweights_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.
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:
- Issue 5 — extract
load_for_inferenceintogiant/checkpoint_io.py. Small, self-contained, removes the highest-risk duplication, and unblocks Issue 10. - Issue 3 — extract the flag→config mapping table into
config.py. - 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. - Issue 1 + Issue 2 together — typed config dataclasses. One migration solves both, and doing them separately means touching the same 132 call sites twice.
- Issue 6 — unify the migration surfaces (easier once config is typed).
- 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. - Issue 7 —
NamedTuplethe feature/batch contracts. - Issue 9 — move
scripts/→giant/tools/. Independent of everything else; can be done at any point. - Issue 10 — checkpoint format split, if the threat model or a publication makes it worthwhile.
- Issue 11 — opportunistically.
General guidance for whoever picks this up
- Run the full suite on every change.
uv run pytest -qtakes 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 (PlotSpecregistry), explicit about its own contracts (thecompute_partial/finalizesplit,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 refactoredconfig.pyorcli.pyshould look, look there.