Files
giant/tests/test_config_consumed_keys.py
T
lars 48faaee79d
CI / Lint (ruff check) (push) Successful in 27s
CI / Format (ruff format) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 31s
CI / Tests (push) Successful in 2m26s
CI / Lint (ruff check) (pull_request) Successful in 27s
CI / Format (ruff format) (pull_request) Successful in 27s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 27s
CI / Tests (pull_request) Successful in 2m23s
Implement stage2_model.stage1_context = "sampled" (gitea #41)
Stage 2 was trained on ground-truth stage-1 outcomes but deployed on
sampled ones, and in a rollout that gap compounds over every step of
every track — the same train/inference gap teacher_forcing="scheduled"
already closes within stage 2, just never applied at the stage
boundary. "sampled" was declared in the schema but rejected loudly by
validate_config as unimplemented; this lands the real implementation.

Mirrors the existing scheduled-sampling precedent rather than a hard
switch: new stage2_model.ctx_p_start/ctx_p_end (defaults 1.0 -> 0.0)
linearly ramp P(condition on ground truth) from epoch 0 to the final
epoch, so stage 2 doesn't chase a wildly moving stage-1 target early in
training. Per the plan discussed with the user: the sample is drawn
from stage 1's sampling_model() (EMA weights when present, matching
what inference actually deploys), mixed per example via a Bernoulli
draw (never blended within a row), and validation always uses the
ground truth regardless of the schedule. Fixes a latent bug the same
pattern would otherwise have hit: every sampler in giant/sample.py
flips its model to .eval() with no restore, so sampling from the raw
(non-EMA) stage-1 model mid-step now explicitly restores its .training
flag afterward to avoid silently corrupting stage 1's own training mode
for the rest of the epoch.

validate_config now enforces stage1_context in {"truth", "sampled"},
requires both stages active for "sampled" (nothing to sample from
otherwise), range-checks ctx_p_start/ctx_p_end, and rejects the
ctx_p_start = ctx_p_end = 1.0 configuration as an unadvertised no-op
identical to "truth".

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-17 12:27:13 +02:00

138 lines
5.8 KiB
Python

"""Consumed-keys audit (issues.md Issue 5).
`validate_config_keys` (`giant/config.py`) only checks that a config key is
*declared* — present somewhere in `DEFAULT_CONFIG`, which is generated from
the frozen dataclasses. It says nothing about whether anything actually
*reads* the value once parsed. Issues 1, 2 and 4 are three keys that slipped
through exactly that gap: declared, round-tripped, silently ignored. This
module walks every leaf path in `DEFAULT_CONFIG` and asserts each is either
genuinely consumed by the model-building/training/rollout code, or explicitly
recorded in `_KNOWN_UNUSED` with a reason.
"Consumed" is approximated by static analysis rather than true call-graph
reachability: for each leaf path's field name, does it appear anywhere in a
fixed whitelist of source files as a real attribute access, a dict-key-shaped
string constant, or a function/constructor parameter name (the last of these
because `Router` subclasses receive their config via `**kwargs` filtered by
signature — see `giant.model.routers.build_router`)? Docstrings are excluded
from the string-constant scan so prose mentioning a dotted config path in
passing can't masquerade as a read of it. This whitelist-based approach is
deliberately narrower than "anywhere in `giant/`": scanning the whole package
produces false negatives from unrelated identifier collisions (e.g.
`giant/analysis/router_gating.py`'s `_top1_shares(..., order: list, ...)`
parameter would otherwise make `stage2_model.autoregressive.order` read as
"consumed").
"""
import ast
from pathlib import Path
from giant.config import DEFAULT_CONFIG
from giant.config import leaf_paths as _leaf_paths
_REPO_ROOT = Path(__file__).resolve().parents[1]
# Files that legitimately consume model_config / training config at
# build/train/rollout time. Not `giant/cli.py` (a CLI flag existing is not
# consumption — that's precisely how Issue 1 slipped through), not
# `giant/config.py` itself (declaring/parsing a field is not reading it), and
# not `giant/model/_legacy.py` (the protected v0.2 migration surface, which
# intentionally re-derives old flat keys under old names).
_CONSUMER_ROOTS = ("giant/model", "giant/training")
_CONSUMER_FILES = (
"giant/sample.py",
"giant/pipeline.py",
"giant/rollout.py",
"giant/checkpoint_io.py",
"giant/particles.py",
"giant/materials.py",
)
_EXCLUDED_FILES = ("giant/model/_legacy.py",)
# Leaf DEFAULT_CONFIG paths that are declared but not (yet) read anywhere in
# the consumer whitelist above. Each entry must name the issue that tracks
# it. If a key here starts showing up as consumed, the fix landed and this
# entry is stale — see test_known_unused_allow_list_has_no_stale_entries.
_KNOWN_UNUSED = {
"stage2_model.autoregressive.order": (
"gitea #30 — validate_config now checks order is 'energy_desc', but "
"nothing in the build/train/rollout consumer whitelist reads the "
"value itself since it's still single-valued"
),
}
# "lambda" is a Python keyword, so the dataclasses expose the dict key
# "lambda" as the field `lambda_weight` (giant/config.py:49-50).
_FIELD_NAME_OVERRIDES = {"lambda": "lambda_weight"}
def _field_name(leaf_path: str) -> str:
name = leaf_path.rsplit(".", 1)[-1]
return _FIELD_NAME_OVERRIDES.get(name, name)
def _is_docstring_expr(expr: ast.Expr) -> bool:
return isinstance(expr.value, ast.Constant) and isinstance(expr.value.value, str)
def _collect_names(source: str, filename: str) -> set[str]:
tree = ast.parse(source, filename=filename)
docstring_ids = set()
for node in ast.walk(tree):
if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
body = getattr(node, "body", [])
if body and isinstance(body[0], ast.Expr) and _is_docstring_expr(body[0]):
docstring_ids.add(id(body[0].value))
names: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Attribute):
names.add(node.attr)
elif isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstring_ids:
names.add(node.value)
elif isinstance(node, ast.arg):
names.add(node.arg)
elif isinstance(node, ast.keyword) and node.arg is not None:
names.add(node.arg)
return names
def _consumer_files() -> list[Path]:
files: set[Path] = {_REPO_ROOT / f for f in _CONSUMER_FILES}
for root in _CONSUMER_ROOTS:
files |= set((_REPO_ROOT / root).rglob("*.py"))
files -= {_REPO_ROOT / f for f in _EXCLUDED_FILES}
return sorted(files)
def _consumed_names() -> set[str]:
names: set[str] = set()
for path in _consumer_files():
names |= _collect_names(path.read_text(), str(path))
return names
def test_every_config_key_is_consumed_or_allow_listed():
consumed = _consumed_names()
unconsumed = {p for p in _leaf_paths(DEFAULT_CONFIG) if _field_name(p) not in consumed}
unexplained = unconsumed - _KNOWN_UNUSED.keys()
assert not unexplained, (
f"config key(s) {sorted(unexplained)} are declared in DEFAULT_CONFIG "
"but not read anywhere in the build/train/rollout consumer files "
f"({[str(f.relative_to(_REPO_ROOT)) for f in _consumer_files()]}) — "
"either wire the key up, or add it to _KNOWN_UNUSED with a reason "
"(see issues.md Issue 5)"
)
def test_known_unused_allow_list_has_no_stale_entries():
consumed = _consumed_names()
all_paths = set(_leaf_paths(DEFAULT_CONFIG))
stale = {p for p in _KNOWN_UNUSED if p not in all_paths or _field_name(p) in consumed}
assert not stale, (
f"_KNOWN_UNUSED entry/entries {sorted(stale)} no longer belong on the "
"allow-list — either the key was removed from DEFAULT_CONFIG, or it "
"is now consumed (the underlying issue was fixed). Remove the stale "
"entry/entries."
)