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>

tests/test_checkpoint_io.py

@@ -0,0 +1,236 @@
+"""Tests for giant.checkpoint_io.load_for_inference (issues.md Issue 5) —
+the shared bootstrap `giant predict`/`giant rollout` use to go from a
+checkpoint path to ready-to-run models."""
+
+from __future__ import annotations
+
+import copy
+
+import numpy as np
+import pytest
+import torch
+
+from giant import config as gconfig
+from giant.checkpoint_io import (
+    CheckpointCompatibilityError,
+    InferenceContext,
+    conditioning_axes,
+    load_for_inference,
+    stage_cfg,
+)
+from giant.data.loader import TopNMap
+from giant.data.setup_cache import topnmap_to_json
+from giant.data.transforms import Normalizer
+from giant.model.network import build_models
+
+PDG_MAP = {11: 0, 22: 1, -11: 2}
+MAT_MAP = {"G4_PbWO4": 0, "G4_AIR": 1}
+
+
+def _model_cfg(stage2_active: bool = True) -> dict:
+    """DEFAULT_CONFIG-derived, shrunk for speed — same pattern as
+    tests/test_network.py::_minimal_model_config. Default `conditioning`
+    (both axes "physical") needs no top-N vocab map, so this is a cheap,
+    fully self-contained happy-path config."""
+    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
+    cfg["conditioning"]["particle"]["emb_dim"] = 4
+    cfg["conditioning"]["material"]["emb_dim"] = 4
+    cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
+    cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
+    cfg["stage2_model"]["active"] = stage2_active
+    return {
+        "pdg_vocab": len(PDG_MAP),
+        "mat_vocab": len(MAT_MAP),
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+    }
+
+
+def _norms() -> tuple[Normalizer, Normalizer, Normalizer]:
+    rng = np.random.default_rng(0)
+    cond = Normalizer().fit(rng.standard_normal((100, 15)).astype(np.float32))
+    tgt = Normalizer().fit(rng.standard_normal((100, 9)).astype(np.float32))
+    sec_phys = Normalizer().fit(rng.standard_normal((100, 2)).astype(np.float32))
+    return cond, tgt, sec_phys
+
+
+def _write_checkpoint(tmp_path, model_cfg=None, ema: bool = False, **ckpt_overrides):
+    cfg = model_cfg if model_cfg is not None else _model_cfg()
+    built = build_models(cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    cond, tgt, sec_phys = _norms()
+    ckpt: dict = {
+        "model_config": cfg,
+        "model": stage1.state_dict() if stage1 is not None else {},
+        "sec_decoder": stage2.state_dict() if stage2 is not None else {},
+        "pdg_map": PDG_MAP,
+        "mat_map": MAT_MAP,
+        "normalizer": {"cond": cond.to_dict(), "target": tgt.to_dict(), "sec_phys": sec_phys.to_dict()},
+        "epoch": 3,
+        "best_val_loss": 0.5,
+    }
+    if ema:
+        ckpt["model_ema"] = stage1.state_dict() if stage1 is not None else {}
+        ckpt["sec_decoder_ema"] = stage2.state_dict() if stage2 is not None else {}
+    ckpt.update(ckpt_overrides)
+    path = tmp_path / "ckpt.pt"
+    torch.save(ckpt, path)
+    return path
+
+
+def _onehot_model_cfg() -> dict:
+    cfg = _model_cfg()
+    cfg["conditioning"]["particle"]["type"] = "onehot"
+    return cfg
+
+
+# ---------------------------------------------------------------------------
+# Happy path
+# ---------------------------------------------------------------------------
+
+
+def test_happy_path_returns_populated_context(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+    assert isinstance(ctx, InferenceContext)
+    assert ctx.stage1 is not None and ctx.stage2 is not None
+    assert not ctx.stage1.training
+    assert not ctx.stage2.training
+    assert next(ctx.stage1.parameters()).device == torch.device("cpu")
+    assert ctx.pdg_map == PDG_MAP
+    assert ctx.mat_map == MAT_MAP
+    assert all(isinstance(k, int) for k in ctx.pdg_map)
+    assert all(isinstance(k, str) for k in ctx.mat_map)
+    assert ctx.particle_conditioning == "physical"
+    assert ctx.material_conditioning == "physical"
+    assert ctx.k_max == 3
+    assert ctx.epoch == 3
+    assert ctx.best_val_loss == 0.5
+    assert ctx.model_config["stage1_model"]["hidden_dim"] == 8
+
+
+def test_happy_path_normalizer_values_round_trip(tmp_path):
+    cond, tgt, sec_phys = _norms()
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+    assert ctx.cond_norm.mean is not None and cond.mean is not None
+    assert ctx.tgt_norm.mean is not None and tgt.mean is not None
+    assert ctx.sec_phys_norm.mean is not None and sec_phys.mean is not None
+    np.testing.assert_allclose(ctx.cond_norm.mean, cond.mean)
+    np.testing.assert_allclose(ctx.tgt_norm.mean, tgt.mean)
+    np.testing.assert_allclose(ctx.sec_phys_norm.mean, sec_phys.mean)
+
+
+# ---------------------------------------------------------------------------
+# Guards
+# ---------------------------------------------------------------------------
+
+
+def test_missing_model_config_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["model_config"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no model_config"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_missing_sec_decoder_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_decoder"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no sec_decoder"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_missing_sec_phys_normalizer_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["normalizer"]["sec_phys"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no normalizer.sec_phys"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_onehot_particle_conditioning_without_topn_map_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_onehot_model_cfg())
+
+    with pytest.raises(CheckpointCompatibilityError, match="pdg_topn_map"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_onehot_particle_conditioning_with_topn_map_succeeds(tmp_path):
+    topn = TopNMap(class_map={11: 0, 22: 1}, other_members={})
+    checkpoint = _write_checkpoint(
+        tmp_path,
+        model_cfg=_onehot_model_cfg(),
+        pdg_topn_map=topnmap_to_json(topn),
+    )
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.particle_conditioning == "onehot"
+    assert ctx.pdg_topn_map is not None
+    assert ctx.pdg_topn_map.class_map == {11: 0, 22: 1}
+
+
+def test_ema_weights_requested_but_missing_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, ema=False)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no EMA weights"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
+
+
+def test_ema_weights_requested_and_present_succeeds(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, ema=True)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
+    assert ctx.stage1 is not None and ctx.stage2 is not None
+
+
+@pytest.mark.parametrize("command_name", ["predict", "rollout"])
+def test_inactive_stage_with_require_stage2_raises_with_command_name(tmp_path, command_name):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
+
+    with pytest.raises(CheckpointCompatibilityError, match=f"{command_name} needs both"):
+        load_for_inference(checkpoint, torch.device("cpu"), command_name)
+
+
+def test_inactive_stage_with_require_stage2_false_succeeds_with_stage2_none(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
+
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", require_stage2=False)
+    assert ctx.stage1 is not None
+    assert ctx.stage2 is None
+
+
+# ---------------------------------------------------------------------------
+# conditioning_axes / stage_cfg
+# ---------------------------------------------------------------------------
+
+
+def test_conditioning_axes_v02_flat_string_applies_to_both_axes():
+    assert conditioning_axes({"conditioning": "embedding"}) == ("embedding", "embedding")
+
+
+def test_conditioning_axes_v03_nested_dict_independent_per_axis():
+    model_cfg = {"conditioning": {"particle": {"type": "onehot"}, "material": {"type": "physical"}}}
+    assert conditioning_axes(model_cfg) == ("onehot", "physical")
+
+
+def test_conditioning_axes_missing_key_uses_default():
+    assert conditioning_axes({}, default="embedding") == ("embedding", "embedding")
+
+
+def test_stage_cfg_new_shape_returns_subdict():
+    model_cfg = {"stage2_model": {"k_max": 7}}
+    assert stage_cfg(model_cfg, "stage2") == {"k_max": 7}
+
+
+def test_stage_cfg_v02_flat_shape_returns_empty_dict():
+    model_cfg = {"hidden_dim": 32, "n_blocks": 4}
+    assert stage_cfg(model_cfg, "stage2") == {}

tests/test_cli_predict.py

@@ -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

tests/test_cli_rollout.py

@@ -0,0 +1,33 @@
+"""Thin CLI smoke coverage for `giant rollout` (issues.md Issue 5) — confirms
+the CheckpointCompatibilityError raised by giant.checkpoint_io.load_for_inference
+surfaces as a clean typer.Exit(1) with the expected message, end-to-end
+through the CLI, not just at the giant.checkpoint_io unit level."""
+
+from __future__ import annotations
+
+import torch
+from typer.testing import CliRunner
+
+from giant.cli import app
+
+runner = CliRunner()
+
+
+def test_rollout_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,
+        [
+            "rollout",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            "dummy_geometry.pkl",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "checkpoint has no model_config" in result.output