Extract predict/rollout's duplicated inference bootstrap into giant.checkpoint_io (issues.md Issue 5)

giant predict and giant rollout each carried a ~65-line, independently
drifting copy of "load checkpoint -> validate -> resolve conditioning axes
-> restore normalizers/vocab maps -> build models -> load weights", plus a
third partial copy of _conditioning_axes in analysis/router_gating.py. A
silent divergence there doesn't crash, it makes the two commands run
different physics from the same checkpoint with no test coverage anywhere
along that path.

giant/checkpoint_io.py now holds the single implementation:
load_for_inference() + an InferenceContext dataclass, raising
CheckpointCompatibilityError (verbatim message text preserved) instead of
calling typer directly, so it can be unit-tested and imported from
non-Typer code. router_gating.py's load_router imports conditioning_axes
from it lazily, keeping its "no torch at module scope" contract intact.

Adds 17 direct unit tests for load_for_inference/conditioning_axes/stage_cfg
plus CLI smoke tests confirming the error surfaces as typer.Exit(1) through
predict and rollout — previously zero coverage on this path.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-08-12 15:31:14 +02:00
parent 6a21c3b908
commit 818c380fd0
6 changed files with 539 additions and 224 deletions
+22
View File
@@ -1,13 +1,18 @@
import uuid
import torch
import yaml
from typer.testing import CliRunner
from giant.cli import (
_CEPH_PREDICTIONS,
_resolve_prediction_output,
_write_prediction_ref,
app,
)
runner = CliRunner()
# ---------------------------------------------------------------------------
# _resolve_prediction_output
@@ -150,3 +155,20 @@ def test_ref_checkpoint_path_is_absolute(tmp_path):
data = yaml.safe_load(ref_path.read_text())
assert data["checkpoint"].startswith("/")
# ---------------------------------------------------------------------------
# Bootstrap failure surfaces via the CLI (issues.md Issue 5 — confirms
# CheckpointCompatibilityError -> typer.Exit(1) actually wires up end-to-end,
# not just at the giant.checkpoint_io unit level).
# ---------------------------------------------------------------------------
def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
checkpoint = tmp_path / "bad.pt"
torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
result = runner.invoke(app, ["predict", "dummy.parquet", "--checkpoint", str(checkpoint)])
assert result.exit_code == 1
assert "checkpoint has no model_config" in result.output