lars
24445b7427
CI / Lint (ruff check) (push) Successful in 28s
CI / Format (ruff format) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 34s
CI / Type check (ty) (push) Successful in 35s
CI / Format (ruff format) (pull_request) Successful in 33s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 34s
CI / Tests (pull_request) Failing after 3m23s
CI / Tests (push) Failing after 3m32s
Closes the highest-value coverage gaps found via pytest-cov: pipeline.py's EnergyRouter quantile-seeding (the roadmap's flagged fix for the failed MoE rollout benchmark) had zero coverage, geometry.py's real parquet-batch reader was always mocked, the material top-N-map cache-hit branch was untested (only pdg's was), and setup_cache.py was missing malformed-cache-body and unknown-axis error paths. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
413 lines
16 KiB
Python
413 lines
16 KiB
Python
import copy
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
import torch
|
|
|
|
from giant import config as gconfig
|
|
from giant.constants import COND_DIM
|
|
from giant.data import setup_cache
|
|
from giant.data.transforms import Normalizer
|
|
from giant.pipeline import _seed_energy_router, run_train_job
|
|
|
|
|
|
def _unit(v):
|
|
v = np.asarray(v, dtype=np.float64)
|
|
n = np.linalg.norm(v)
|
|
return v / n if n > 1e-9 else np.array([0.0, 0.0, 1.0])
|
|
|
|
|
|
def _make_synthetic_steps(path, n_events=20, seed=0):
|
|
"""A tiny but schema-complete synthetic steps parquet for run_train_job.
|
|
|
|
pdg/material/process are assigned deterministically by row index (not
|
|
random) so tests that assert on the resulting vocab/proc maps aren't
|
|
flaky; only continuous quantities (positions/energies/directions) are
|
|
drawn from `rng`.
|
|
"""
|
|
rng = np.random.default_rng(seed)
|
|
materials = ["G4_AIR", "G4_Fe"]
|
|
pdgs = [11, 22]
|
|
processes = ["eIoni", "phot", "compt"]
|
|
rows = []
|
|
row_idx = 0
|
|
for event_id in range(n_events):
|
|
n_steps = int(rng.integers(2, 4))
|
|
for s in range(n_steps):
|
|
pre_E = float(rng.uniform(50.0, 500.0))
|
|
n_sec = int(rng.integers(0, 3))
|
|
frac_dep = float(rng.uniform(0.05, 0.3))
|
|
frac_sec = float(rng.uniform(0.05, 0.2)) if n_sec > 0 else 0.0
|
|
frac_post = 1.0 - frac_dep - frac_sec
|
|
edep = pre_E * frac_dep
|
|
e_sec = pre_E * frac_sec
|
|
post_E = pre_E * frac_post
|
|
pre_pos = rng.uniform(-10, 10, size=3)
|
|
step_length = float(rng.uniform(0.1, 5.0))
|
|
pre_dir = np.array([0.0, 0.0, 1.0])
|
|
post_dir = _unit(rng.normal(size=3))
|
|
post_pos = pre_pos + step_length * pre_dir
|
|
sec_energies = (
|
|
list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
|
|
)
|
|
sec_pdgs = [pdgs[(row_idx + j) % 2] for j in range(n_sec)]
|
|
sec_dirs = [_unit(rng.normal(size=3)) for _ in range(n_sec)]
|
|
rows.append(
|
|
{
|
|
"event_id": event_id,
|
|
"pdg": pdgs[row_idx % 2],
|
|
"pre_x": pre_pos[0],
|
|
"pre_y": pre_pos[1],
|
|
"pre_z": pre_pos[2],
|
|
"pre_E": pre_E,
|
|
"pre_dx": pre_dir[0],
|
|
"pre_dy": pre_dir[1],
|
|
"pre_dz": pre_dir[2],
|
|
"material": materials[row_idx % 2],
|
|
"layer_id": s,
|
|
"child_track_ids": list(range(n_sec)),
|
|
"e_sec": e_sec,
|
|
"process": processes[row_idx % 3],
|
|
"step_length": step_length,
|
|
"post_E": post_E,
|
|
"edep": edep,
|
|
"post_dx": post_dir[0],
|
|
"post_dy": post_dir[1],
|
|
"post_dz": post_dir[2],
|
|
"post_x": post_pos[0],
|
|
"post_y": post_pos[1],
|
|
"post_z": post_pos[2],
|
|
"sec_E_list": sec_energies,
|
|
"sec_pdg_list": sec_pdgs,
|
|
"sec_dx_list": [d[0] for d in sec_dirs],
|
|
"sec_dy_list": [d[1] for d in sec_dirs],
|
|
"sec_dz_list": [d[2] for d in sec_dirs],
|
|
}
|
|
)
|
|
row_idx += 1
|
|
pd.DataFrame(rows).to_parquet(path)
|
|
return path
|
|
|
|
|
|
def _tiny_cfg(**train_overrides):
|
|
cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
|
|
cfg["train"].update(
|
|
{
|
|
"epochs": 1,
|
|
"batch_size": 8,
|
|
"val_fraction": 0.2,
|
|
"seed": 0,
|
|
"warmup_epochs": 0,
|
|
"validate_every": 0,
|
|
"max_val_batches": 1,
|
|
"wandb": False,
|
|
}
|
|
)
|
|
cfg["train"].update(train_overrides)
|
|
cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0})
|
|
cfg["stage2_model"].update(
|
|
# decoder="autoregressive" is DEFAULT_CONFIG's default (v0.3.0 step 5)
|
|
# and left as-is here on purpose, so this pipeline-level fixture
|
|
# exercises the real default end-to-end against actual data.
|
|
{"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0}
|
|
)
|
|
cfg["conditioning"]["particle"]["emb_dim"] = 4
|
|
cfg["conditioning"]["material"]["emb_dim"] = 4
|
|
return cfg
|
|
|
|
|
|
def _run(data, out_dir, cfg=None, **kwargs):
|
|
echoed: list[str] = []
|
|
kwargs.setdefault("num_workers", 0)
|
|
run_train_job(
|
|
data=data,
|
|
cfg=cfg or _tiny_cfg(),
|
|
out_dir=out_dir,
|
|
device=torch.device("cpu"),
|
|
shuffle_buffer=64,
|
|
echo=echoed.append,
|
|
**kwargs,
|
|
)
|
|
return echoed
|
|
|
|
|
|
@pytest.fixture
|
|
def data(tmp_path):
|
|
return _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
|
|
|
|
|
|
def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
|
|
echo1 = _run(data, tmp_path / "out1")
|
|
assert any("fitting normalizer (streaming)" in m for m in echo1)
|
|
|
|
def _forbidden(*a, **k):
|
|
raise AssertionError("should be served from cache, not recomputed")
|
|
|
|
monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
|
|
monkeypatch.setattr("giant.pipeline.iter_file_chunks", _forbidden)
|
|
|
|
echo2 = _run(data, tmp_path / "out2")
|
|
joined = "\n".join(echo2)
|
|
assert "event index: cache hit" in joined
|
|
assert "vocabulary maps: cache hit" in joined
|
|
assert "normalizer: cache hit" in joined
|
|
|
|
|
|
def test_run_train_job_builds_caches_and_persists_pdg_topn_map(tmp_path, data):
|
|
"""DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
|
|
"onehot" — a plain _tiny_cfg() run must
|
|
build the shared pdg top-N map, cache it in the setup-cache sidecar, and
|
|
persist it into the checkpoint, with no extra config needed."""
|
|
echo1 = _run(data, tmp_path / "out1")
|
|
assert any("building pdg top-N map" in m for m in echo1)
|
|
|
|
loaded = setup_cache.load(data, [data])
|
|
assert loaded is not None
|
|
key = setup_cache.topn_key("pdg", 4) # conditioning.particle.emb_dim = 4
|
|
assert key in loaded.topn_maps
|
|
assert set(loaded.topn_maps[key].class_map.keys()) >= {11, 22}
|
|
|
|
ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False)
|
|
assert "pdg_topn_map" in ckpt
|
|
assert set(ckpt["pdg_topn_map"]["class_map"].keys()) >= {"11", "22"}
|
|
|
|
echo2 = _run(data, tmp_path / "out2")
|
|
assert any("pdg top-N map: cache hit" in m for m in echo2)
|
|
|
|
|
|
def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, data):
|
|
"""conditioning.material.type="onehot" is an independent axis from the
|
|
pdg one above, with its own build/cache-hit branch in run_setup_stage —
|
|
exercise both here the same way the pdg test above does."""
|
|
cfg = _tiny_cfg()
|
|
cfg["conditioning"]["material"]["type"] = "onehot"
|
|
echo1 = _run(data, tmp_path / "out1", cfg=cfg)
|
|
assert any("building material top-N map" in m for m in echo1)
|
|
|
|
loaded = setup_cache.load(data, [data])
|
|
assert loaded is not None
|
|
key = setup_cache.topn_key("material", 4) # conditioning.material.emb_dim = 4
|
|
assert key in loaded.topn_maps
|
|
assert set(loaded.topn_maps[key].class_map.keys()) >= {"G4_AIR", "G4_Fe"}
|
|
|
|
ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False)
|
|
assert "mat_topn_map" in ckpt
|
|
assert set(ckpt["mat_topn_map"]["class_map"].keys()) >= {"G4_AIR", "G4_Fe"}
|
|
|
|
echo2 = _run(data, tmp_path / "out2", cfg=cfg)
|
|
assert any("material top-N map: cache hit" in m for m in echo2)
|
|
|
|
|
|
def test_run_train_job_no_topn_map_for_physical_target(tmp_path, data):
|
|
cfg = _tiny_cfg()
|
|
cfg["stage2_model"]["particle_type"] = {"target": "physical", "lambda": 1.0}
|
|
echo = _run(data, tmp_path / "out", cfg=cfg)
|
|
assert not any("top-N map" in m for m in echo)
|
|
|
|
loaded = setup_cache.load(data, [data])
|
|
assert loaded is not None
|
|
assert loaded.topn_maps == {}
|
|
|
|
|
|
@pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
|
|
def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
|
|
tmp_path, data, monkeypatch
|
|
):
|
|
# num_workers>0 makes DataLoader actually fork worker subprocesses
|
|
# (unlike every other test here, which runs with num_workers=0) — pytest
|
|
# itself is multi-threaded, hence Python's fork-safety warning below.
|
|
monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8) # quota = 2
|
|
echo = _run(data, tmp_path / "out", num_workers=3)
|
|
assert any("num-workers=3" in m and "exceeds" in m for m in echo)
|
|
|
|
|
|
@pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
|
|
def test_run_train_job_no_warning_when_num_workers_within_shared_quota(
|
|
tmp_path, data, monkeypatch
|
|
):
|
|
monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8) # quota = 2
|
|
echo = _run(data, tmp_path / "out", num_workers=2)
|
|
assert not any("exceeds" in m for m in echo)
|
|
|
|
|
|
def test_run_train_job_no_cache_setup_never_writes_sidecar(tmp_path, data):
|
|
_run(data, tmp_path / "out", cache_setup=False)
|
|
assert not setup_cache.sidecar_path(data).exists()
|
|
|
|
|
|
def test_run_train_job_rebuild_setup_cache_ignores_existing(tmp_path, data):
|
|
files = [data]
|
|
stale = setup_cache.SetupCache.empty(files)
|
|
stale.vocab = ({999999: 0}, {"G4_AIR": 0}) # deliberately wrong
|
|
setup_cache.save(data, files, stale)
|
|
|
|
_run(data, tmp_path / "out", rebuild_setup_cache=True)
|
|
|
|
loaded = setup_cache.load(data, files)
|
|
assert loaded is not None
|
|
assert loaded.vocab is not None
|
|
assert set(loaded.vocab[0].keys()) == {11, 22}
|
|
assert set(loaded.vocab[1].keys()) == {"G4_AIR", "G4_Fe"}
|
|
|
|
|
|
def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypatch):
|
|
_run(data, tmp_path / "out1", cfg=_tiny_cfg(val_fraction=0.1))
|
|
|
|
def _forbidden(*a, **k):
|
|
raise AssertionError("vocab should be served from cache")
|
|
|
|
monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
|
|
|
|
echo2 = _run(data, tmp_path / "out2", cfg=_tiny_cfg(val_fraction=0.3))
|
|
joined = "\n".join(echo2)
|
|
assert "vocabulary maps: cache hit" in joined
|
|
assert "fitting normalizer (streaming)" in joined
|
|
|
|
|
|
def test_run_train_job_custom_k_max_end_to_end(tmp_path, data):
|
|
"""Regression: stage2_model.k_max other
|
|
than the K_MAX module constant's default (15) must not produce a shape
|
|
mismatch between the data pipeline (loader.py/transforms.py padding) and
|
|
the model (network.py's trunks, sized from this same config value)."""
|
|
cfg = _tiny_cfg()
|
|
cfg["stage2_model"]["k_max"] = 3
|
|
_run(data, tmp_path / "out", cfg=cfg)
|
|
ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
|
|
assert ckpt["model_config"]["stage2_model"]["k_max"] == 3
|
|
|
|
|
|
def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, data):
|
|
"""Regression: conditioning.particle.type
|
|
and conditioning.material.type are configured independently and may mix
|
|
freely — e.g. particle "embedding" with
|
|
material "physical" — end-to-end through the real data pipeline, not
|
|
just accepted by validate_config."""
|
|
cfg = _tiny_cfg()
|
|
cfg["conditioning"]["particle"]["type"] = "embedding"
|
|
cfg["conditioning"]["material"]["type"] = "physical"
|
|
_run(data, tmp_path / "out", cfg=cfg)
|
|
ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
|
|
cond_cfg = ckpt["model_config"]["conditioning"]
|
|
assert cond_cfg["particle"]["type"] == "embedding"
|
|
assert cond_cfg["material"]["type"] == "physical"
|
|
|
|
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
|
|
# Particle block ([COND_DIM_BASE:COND_DIM_BASE+PARTICLE_PHYS_DIM]) stays
|
|
# unfitted (mean=0/std=1) since "embedding" never computes real values
|
|
# for it; the material block is fit for real under "physical".
|
|
from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM
|
|
|
|
assert cond_norm.mean is not None and cond_norm.std is not None
|
|
np.testing.assert_allclose(
|
|
cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0
|
|
)
|
|
np.testing.assert_allclose(
|
|
cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0
|
|
)
|
|
material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :]
|
|
assert np.all(material_std > 0) and not np.allclose(material_std, 1.0)
|
|
|
|
|
|
def test_run_train_job_share_stages_end_to_end(tmp_path, data):
|
|
"""Regression: conditioning.share_stages
|
|
= true must actually train (not raise NotImplementedError), and the
|
|
resulting checkpoint's two stages must reload into a single shared
|
|
ConditionEncoder instance rather than two independent ones."""
|
|
from giant.model.network import build_models
|
|
|
|
cfg = _tiny_cfg()
|
|
cfg["conditioning"]["share_stages"] = True
|
|
_run(data, tmp_path / "out", cfg=cfg)
|
|
ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
|
|
assert ckpt["model_config"]["conditioning"]["share_stages"] is True
|
|
|
|
built = build_models(ckpt["model_config"])
|
|
stage1, stage2 = built["stage1"], built["stage2"]
|
|
assert stage1 is not None and stage2 is not None
|
|
assert stage1.cond_enc is stage2.cond_enc
|
|
|
|
stage1.load_state_dict(ckpt["model"])
|
|
stage2.load_state_dict(ckpt["sec_decoder"])
|
|
for p1, p2 in zip(stage1.cond_enc.parameters(), stage2.cond_enc.parameters()):
|
|
assert torch.equal(p1, p2)
|
|
|
|
|
|
def test_run_train_job_matches_uncached_output(tmp_path, data):
|
|
_run(data, tmp_path / "uncached", cache_setup=False)
|
|
_run(data, tmp_path / "cached1", cache_setup=True)
|
|
_run(data, tmp_path / "cached2", cache_setup=True) # second is a cache hit
|
|
|
|
uncached = torch.load(tmp_path / "uncached" / "last.pt", weights_only=False)
|
|
cached = torch.load(tmp_path / "cached2" / "last.pt", weights_only=False)
|
|
|
|
for key in ("cond", "target", "sec_phys"):
|
|
np.testing.assert_allclose(
|
|
uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"]
|
|
)
|
|
np.testing.assert_allclose(
|
|
uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"]
|
|
)
|
|
assert uncached["pdg_map"] == cached["pdg_map"]
|
|
assert uncached["mat_map"] == cached["mat_map"]
|
|
|
|
|
|
def _fitted_cond_norm(seed=0):
|
|
rng = np.random.default_rng(seed)
|
|
return Normalizer().fit(rng.normal(size=(64, COND_DIM)).astype(np.float32))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"router_cfg",
|
|
[
|
|
{"enabled": False, "type": "energy", "n_experts": 4},
|
|
{"enabled": True, "type": "pdg", "n_experts": 4},
|
|
],
|
|
)
|
|
def test_seed_energy_router_noop_when_not_an_enabled_energy_router(router_cfg):
|
|
cond_norm = _fitted_cond_norm()
|
|
echoed = []
|
|
_seed_energy_router(router_cfg, cond_norm, np.array([1.0, 2.0]), 3, echoed.append)
|
|
assert "centers_init" not in router_cfg
|
|
assert echoed == []
|
|
|
|
|
|
def test_seed_energy_router_falls_back_to_default_and_warns_when_no_samples():
|
|
router_cfg = {"enabled": True, "type": "energy", "n_experts": 4}
|
|
cond_norm = _fitted_cond_norm()
|
|
echoed = []
|
|
_seed_energy_router(
|
|
router_cfg, cond_norm, np.empty(0), energy_idx=3, echo=echoed.append
|
|
)
|
|
assert "centers_init" not in router_cfg
|
|
assert len(echoed) == 1
|
|
assert "falls back to default centers" in echoed[0]
|
|
|
|
|
|
def test_seed_energy_router_seeds_centers_from_data_quantiles():
|
|
router_cfg = {"enabled": True, "type": "energy", "n_experts": 4}
|
|
cond_norm = _fitted_cond_norm()
|
|
energy_idx = 3
|
|
# A grid of "raw" quantile values as setup_cache.energy_quantiles_from_sample
|
|
# would produce them: monotonically increasing, in the same (log-energy)
|
|
# units as the conditioning column being normalized against.
|
|
energy_quantiles = np.linspace(1.0, 10.0, 33).astype(np.float32)
|
|
echoed = []
|
|
_seed_energy_router(
|
|
router_cfg, cond_norm, energy_quantiles, energy_idx, echoed.append
|
|
)
|
|
|
|
assert "centers_init" in router_cfg
|
|
centers = np.asarray(router_cfg["centers_init"], dtype=np.float32)
|
|
assert centers.shape == (router_cfg["n_experts"],)
|
|
|
|
levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
|
|
raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
|
|
expected = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx]
|
|
np.testing.assert_allclose(centers, expected, rtol=1e-5)
|
|
# Quantile levels are increasing, and the normalizer's std is positive, so
|
|
# the seeded centers must preserve that order rather than e.g. reversing it.
|
|
assert np.all(np.diff(centers) > 0)
|
|
assert len(echoed) == 1
|
|
assert "seeded EnergyRouter centers" in echoed[0]
|