From 74343d3e4847ff19e93e4e80d86c61d8fbe013cf Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 3 Aug 2026 13:47:29 +0200 Subject: [PATCH] Add data-integrity guards against silent NaN/Inf propagation and races MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - log_transform / _validate_unit_pre_dir now raise on non-finite input instead of letting a NaN row silently poison the persisted normalizer cache (norm < 1e-6 was always False for NaN, so the existing guard never caught it). - encode_secondaries warns when a row's secondary energies cumulatively exceed e_sec, instead of silently saturating the overflowing slot's stick-breaking logit via the _EPS floor. - EVENT_ID_FILE_STRIDE overflow now raises instead of silently colliding two files' event ids together (reintroducing train/val leakage). - make_event_split(val_fraction=0.0) now actually holds out nothing, instead of always forcing at least 1 validation event. - setup_cache.save() is now serialized with a flock, since two concurrent writers (a real scenario on this repo's shared portal/condor machines) could otherwise race and silently drop one writer's freshly-computed cache section. - Documented (no behavior change) the pre_dir ≈ -ẑ antipodal rotation singularity in _rodrigues_axis, which is real but inherent to any single-valued local-frame convention. Each fix has a regression test. Co-Authored-By: Claude Sonnet 5 --- giant/data/dataset.py | 5 ++- giant/data/loader.py | 26 +++++++++++-- giant/data/setup_cache.py | 28 ++++++++++--- giant/data/transforms.py | 67 +++++++++++++++++++++++++++++--- tests/test_dataset.py | 10 +++++ tests/test_loader.py | 9 +++++ tests/test_setup_cache.py | 35 +++++++++++++++++ tests/test_transforms.py | 82 ++++++++++++++++++++++++++++++++++++++- 8 files changed, 247 insertions(+), 15 deletions(-) diff --git a/giant/data/dataset.py b/giant/data/dataset.py index 94eff71..3b56a79 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -19,7 +19,10 @@ def make_event_split( rng = np.random.default_rng(seed) unique = np.unique(all_event_ids) rng.shuffle(unique) - n_val = max(1, int(len(unique) * val_fraction)) + # max(1, ...) only applies when a validation split was actually + # requested — val_fraction=0.0 is an explicit "train on everything" + # request and must not be silently overridden into holding out 1 event. + n_val = max(1, int(len(unique) * val_fraction)) if val_fraction > 0 else 0 val_set = set(unique[:n_val].tolist()) train_set = set(unique[n_val:].tolist()) return train_set, val_set diff --git a/giant/data/loader.py b/giant/data/loader.py index febf885..e76af49 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -27,6 +27,26 @@ def event_id_offset(file_index: int) -> int: return file_index * EVENT_ID_FILE_STRIDE +def _offset_event_id(raw_ids: np.ndarray, offset: int) -> np.ndarray: + """Add this file's `event_id_offset`, after checking the raw ids fit in one stride block. + + Without this check, a file whose own raw event_id numbering reaches + `EVENT_ID_FILE_STRIDE` (an unusually large job, or non-contiguous + numbering) would silently collide into the next file's offset block, + merging unrelated events across files — reintroducing exactly the + train/val event leakage this offset scheme exists to prevent. + """ + raw_ids = np.asarray(raw_ids, dtype=np.int64) + if raw_ids.size and int(raw_ids.max()) >= EVENT_ID_FILE_STRIDE: + raise ValueError( + f"event_id {int(raw_ids.max())} >= EVENT_ID_FILE_STRIDE " + f"({EVENT_ID_FILE_STRIDE}) — this file has a larger event_id " + "than the per-file offset scheme can support without colliding " + "with the next file's id block." + ) + return raw_ids + offset + + def _read_manifest(path: Path) -> list[Path]: files = [] for line in path.read_text().splitlines(): @@ -98,7 +118,7 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]: has_sec_lists = "sec_E_list" in df.columns d: dict[str, np.ndarray] = { - "event_id": df["event_id"].to_numpy().astype(np.int64) + offset, + "event_id": _offset_event_id(df["event_id"].to_numpy(), offset), "pdg": df["pdg"].to_numpy(dtype=np.int32), "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32), "pre_E": df["pre_E"].to_numpy(dtype=np.float32), @@ -142,7 +162,7 @@ def load_steps(path: str | Path, offset: int = 0) -> dict[str, np.ndarray]: def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray: """Read only the event_id column — cheap scan for split assignment.""" ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy() - return ids.astype(np.int64) + offset + return _offset_event_id(ids, offset) def iter_file_chunks( @@ -173,7 +193,7 @@ _COND_COLS = [ def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]: return { - "event_id": df["event_id"].to_numpy().astype(np.int64) + offset, + "event_id": _offset_event_id(df["event_id"].to_numpy(), offset), "pdg": df["pdg"].to_numpy(dtype=np.int32), "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32), "pre_E": df["pre_E"].to_numpy(dtype=np.float32), diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py index 229a99f..b1c69d4 100644 --- a/giant/data/setup_cache.py +++ b/giant/data/setup_cache.py @@ -13,6 +13,7 @@ before reuse — see `load`/`save`. from __future__ import annotations +import fcntl import json import os from dataclasses import dataclass, field @@ -270,15 +271,32 @@ def save( Best-effort: any OSError (permission denied on a read-only mount, disk full, ...) is caught, echoed as a warning, and swallowed — a failure to cache must never fail training. + + The load-merge-write is serialized with an exclusive flock on a sidecar + lockfile: `os.replace` alone only guarantees the *file* is never + corrupt, not that concurrent writers don't race. Without the lock, two + concurrent `giant train`/condor jobs against the same `data` path (this + repo's shared-portal/condor usage makes that a real scenario, not just + theoretical) could both `load()` the same base state, merge their own + `sections` in independently, and whichever `os.replace()` lands last + silently discards the other's freshly-computed section. """ path = sidecar_path(data) + lock_path = path.parent / f".{path.name}.lock" tmp = path.parent / f".{path.name}.tmp.{os.getpid()}" try: - base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(files) - merged = base.merge(sections) - payload = json.dumps(merged.to_json(), separators=(",", ":")) - tmp.write_text(payload) - os.replace(tmp, path) + with open(lock_path, "a") as lock_file: + fcntl.flock(lock_file, fcntl.LOCK_EX) + try: + base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty( + files + ) + merged = base.merge(sections) + payload = json.dumps(merged.to_json(), separators=(",", ":")) + tmp.write_text(payload) + os.replace(tmp, path) + finally: + fcntl.flock(lock_file, fcntl.LOCK_UN) except OSError as exc: echo( f"setup cache: could not write {path} ({exc}) — continuing without caching" diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 8189e87..50bd1c8 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -12,7 +12,17 @@ _SIMPLEX_FLOOR = 1e-5 def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray: - return np.log(np.asarray(x, dtype=np.float32) + eps) + x = np.asarray(x, dtype=np.float32) + y = np.log(x + eps) + if not np.all(np.isfinite(y)): + bad = int(np.sum(~np.isfinite(y))) + raise ValueError( + f"log_transform: {bad} value(s) produced non-finite output (input " + f"< -eps={eps:g}, or already NaN/Inf); every quantity this is " + "applied to should be non-negative, so this indicates upstream " + "data corruption rather than expected float noise." + ) + return y def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray: @@ -110,8 +120,22 @@ def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray: [pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1 ) axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1) - # Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis - # are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant). + # axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and + # pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to + # ẑ) vanishes at both. The "choice is irrelevant" claim below only holds + # at +ẑ, where sin_t~0 AND (1-cos_t)~0 so every axis-dependent Rodrigues + # term vanishes. At -ẑ, sin_t~0 but (1-cos_t)~2 — not negligible — so + # snapping to a fixed x̂ there is a genuine (if physically rare) + # modeling choice, not a no-op: it picks one representative out of an + # inherently ambiguous family of 180°-about-any-transverse-axis + # rotations (no single-valued frame convention can be continuous through + # this antipode — same obstruction as a sphere's tangent frame having no + # continuous choice at a pole). x̂ is still fine to use — it's a fixed, + # self-consistent convention that `local_frame_rotation`/ + # `inv_local_frame_rotation` (same threshold) round-trip correctly + # through — but steps whose pre_dir falls in this tiny near-backscatter + # cone get a discontinuous "roll" relative to their non-degenerate + # neighbors, injecting a small amount of label noise there. safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm) return np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm) @@ -134,8 +158,20 @@ def _validate_unit_pre_dir(pre_dir: np.ndarray) -> np.ndarray: drift is corrected silently; a near-zero-norm row has no well-defined direction, so it's raised loudly instead of producing a meaningless rotation (previously it fell through to an arbitrary axis with no error). + + NaN/Inf rows are also raised on explicitly: `norm < 1e-6` is False for a + NaN norm, so without this check a non-finite row would silently pass + through and poison everything downstream (e.g. the persisted normalizer + stats in `setup_cache`, if the row is swept into a Welford accumulator). """ pre_dir = np.asarray(pre_dir, dtype=np.float32) + if not np.all(np.isfinite(pre_dir)): + bad = int(np.sum(~np.all(np.isfinite(pre_dir), axis=1))) + raise ValueError( + f"pre_dir has {bad} row(s) with non-finite (NaN/Inf) components; " + "local/inv_local_frame_rotation require a well-defined incoming " + "direction for every row." + ) norm = np.linalg.norm(pre_dir, axis=1, keepdims=True) if np.any(norm < 1e-6): raise ValueError( @@ -423,11 +459,21 @@ def encode_secondaries( else: cumsum = np.cumsum(sec_E_list.astype(np.float64), axis=1) stick_logits = np.zeros((N, K), dtype=np.float32) + # A valid slot whose cumulative secondary energy so far exceeds + # e_sec by more than float noise means sec_E_list sums to more than + # e_sec — a real upstream data mismatch, not something to paper + # over. Flagged once after the loop rather than let `remaining`'s + # np.maximum(..., _EPS) floor silently absorb it by saturating that + # slot's stick-breaking logit with no signal that anything was off. + _SHORTFALL_TOL = 1e-3 + shortfall_flagged = np.zeros(N, dtype=bool) for i in range(K): if i == 0: - remaining = np.maximum(e_sec, _EPS) + remaining_raw = e_sec else: - remaining = np.maximum(e_sec - cumsum[:, i - 1], _EPS) + remaining_raw = e_sec - cumsum[:, i - 1] + shortfall_flagged |= sec_valid[:, i] & (remaining_raw < -_SHORTFALL_TOL) + remaining = np.maximum(remaining_raw, _EPS) f = np.clip( sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS ) @@ -444,6 +490,17 @@ def encode_secondaries( ) stick_logits[:, i] = logit.astype(np.float32) + if shortfall_flagged.any(): + n = int(shortfall_flagged.sum()) + warnings.warn( + f"encode_secondaries: {n}/{N} row(s) have sec_E_list summing " + "to more than e_sec (beyond float noise) — the overflowing " + "slot(s)' stick-breaking logit was saturated instead of " + "reflecting a real fraction; check upstream secondary " + "energy accounting for these rows.", + stacklevel=2, + ) + # Rotate each slot's direction into the local frame of the primary. # pre_dir is broadcast across all K slots. dir_local = np.zeros((N, K, 3), dtype=np.float32) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 23bf20f..e311dae 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -30,6 +30,16 @@ def test_make_event_split_no_empty_sets(): assert len(val_set) > 0 +def test_make_event_split_val_fraction_zero_holds_out_nothing(): + """val_fraction=0.0 is an explicit "train on everything" request and + must not be silently overridden into holding out 1 event.""" + rng = np.random.default_rng(3) + event_ids = rng.integers(0, 50, size=1000) + train_set, val_set = make_event_split(event_ids, val_fraction=0.0) + assert val_set == set() + assert train_set == set(np.unique(event_ids).tolist()) + + def test_make_event_split_reproducible(): event_ids = np.arange(100) a_tr, a_val = make_event_split(event_ids, val_fraction=0.1, seed=42) diff --git a/tests/test_loader.py b/tests/test_loader.py index 80fab28..f0900c7 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -344,6 +344,15 @@ def test_load_event_ids_applies_offset(tmp_path): ) +def test_load_event_ids_raises_when_event_id_reaches_stride(tmp_path): + """A raw event_id >= EVENT_ID_FILE_STRIDE would collide into the next + file's offset block if silently allowed through — must raise instead.""" + path = tmp_path / "a.parquet" + pd.DataFrame({"event_id": [0, 1, EVENT_ID_FILE_STRIDE]}).to_parquet(path) + with pytest.raises(ValueError, match="EVENT_ID_FILE_STRIDE"): + load_event_ids(path) + + def test_load_steps_applies_offset_to_event_id(tmp_path): path = tmp_path / "a.parquet" _steps_df([0, 1]).to_parquet(path) diff --git a/tests/test_setup_cache.py b/tests/test_setup_cache.py index 5e1c7bd..a243a53 100644 --- a/tests/test_setup_cache.py +++ b/tests/test_setup_cache.py @@ -214,6 +214,41 @@ def test_save_merges_non_colliding_normalizer_keys(tmp_path): assert loaded.normalizers["k2"].n_train_steps == 2 +def test_save_is_serialized_against_concurrent_writers(tmp_path): + """Without the flock in setup_cache.save(), two concurrent writers can + both load() the same base state and merge their own section in + independently, so whichever os.replace() lands last silently drops the + other's key — a lost-update race, not a corrupt file. Each of these + threads writes a distinct normalizer key many times over; if the + load-merge-write critical section isn't actually serialized, at least + one thread's key is likely to go missing from the final merged cache.""" + import threading + + data = _touch_parquet(tmp_path / "shard.parquet") + files = [data] + setup_cache.save(data, files, SetupCache.empty(files)) + + n_writers, n_rounds = 6, 15 + + def _writer(idx: int) -> None: + for r in range(n_rounds): + cache = SetupCache.empty(files) + cache.normalizers[f"k{idx}"] = _entry(n_train_steps=r) + setup_cache.save(data, files, cache) + + threads = [threading.Thread(target=_writer, args=(i,)) for i in range(n_writers)] + for t in threads: + t.start() + for t in threads: + t.join() + + loaded = setup_cache.load(data, files) + assert loaded is not None + assert set(loaded.normalizers.keys()) == {f"k{i}" for i in range(n_writers)} + for i in range(n_writers): + assert loaded.normalizers[f"k{i}"].n_train_steps == n_rounds - 1 + + # ── energy_quantiles_from_sample / energy_quantile_at ─────────────────── diff --git a/tests/test_transforms.py b/tests/test_transforms.py index d598938..899e976 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1,9 +1,12 @@ +import warnings + import numpy as np import pytest from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX from giant.data.transforms import ( build_cond_features, build_features, + encode_secondaries, energy_simplex_decode, energy_simplex_encode, inv_local_frame_rotation, @@ -24,6 +27,50 @@ def test_log_transform_invertible(): np.testing.assert_allclose(inv_log_transform(log_transform(x)), x, rtol=1e-5) +def test_log_transform_raises_on_input_below_negative_eps(): + """A meaningfully negative input (upstream data corruption, not float + noise near 0) must raise instead of silently returning NaN.""" + x = np.array([1.0, -5.0], dtype=np.float32) + with np.errstate(invalid="ignore"), pytest.raises(ValueError, match="non-finite"): + log_transform(x) + + +def test_log_transform_raises_on_nan_input(): + x = np.array([1.0, np.nan], dtype=np.float32) + with pytest.raises(ValueError, match="non-finite"): + log_transform(x) + + +def test_encode_secondaries_warns_when_sec_energies_exceed_e_sec(): + """sec_E_list summing to more than e_sec (before the last slot is even + reached) is a real upstream data mismatch — must warn instead of + silently saturating the overflowing slot's stick-breaking logit via the + _EPS floor. (A single slot alone exceeding what's left of the budget is + the normal, expected "last slot takes the remainder" case and must NOT + warn — the mismatch here is the *cumulative* sum through an earlier + slot already exceeding e_sec.)""" + sec_E_list = np.array([[5.0, 4.0, 1.0]], dtype=np.float32) # sums to 10 + sec_dir_list = np.tile([0.0, 0.0, 1.0], (1, 3, 1)).astype(np.float32) + sec_valid = np.array([[True, True, True]]) + e_sec = np.array([6.0], dtype=np.float32) # cumsum already 9 by slot 2 + pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32) + + with pytest.warns(UserWarning, match="sec_E_list summing to more than e_sec"): + encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir) + + +def test_encode_secondaries_no_warning_when_energies_are_consistent(): + sec_E_list = np.array([[3.0, 2.0]], dtype=np.float32) # sums to 5 + sec_dir_list = np.tile([0.0, 0.0, 1.0], (1, 2, 1)).astype(np.float32) + sec_valid = np.array([[True, True]]) + e_sec = np.array([6.0], dtype=np.float32) # >= 5, no shortfall + pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir) + + def test_local_frame_rotation_noop_when_aligned(): N = 8 pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) @@ -68,10 +115,43 @@ def test_local_frame_rotation_rejects_near_zero_pre_dir(): post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32) with pytest.raises(ValueError, match="near-zero norm"): local_frame_rotation(pre_dir, post_dir) - with pytest.raises(ValueError, match="near-zero norm"): + + +def test_local_frame_rotation_rejects_nan_pre_dir(): + """A NaN pre_dir must raise loudly — `norm < 1e-6` is False for NaN, so + without an explicit isfinite check this would silently poison the + rotation (and any normalizer stats it feeds) instead of erroring.""" + pre_dir = np.array([[np.nan, 0.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32) + post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32) + with pytest.raises(ValueError, match="non-finite"): + local_frame_rotation(pre_dir, post_dir) + with pytest.raises(ValueError, match="non-finite"): inv_local_frame_rotation(pre_dir, post_dir) +def test_local_frame_rotation_antipodal_pre_dir_uses_x_axis_convention(): + """pre_dir ~ -ẑ (near-exact backscatter) is a second axis_norm~0 + degeneracy besides pre_dir ~ +ẑ; unlike the forward case, the Rodrigues + axis-dependent terms are NOT negligible there ((1-cos_t)~2), so the x̂ + fallback is a real (if arbitrary and physically rare) convention choice + rather than a no-op. Pin it explicitly — angle-preservation and the + round-trip property must still hold even though the "roll" is degenerate. + """ + pre_dir = np.array([[0.0, 0.0, -1.0]], dtype=np.float32) + post_dir = np.array([[0.3, 0.4, 0.5]], dtype=np.float32) + post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True) + + rotated = local_frame_rotation(pre_dir, post_dir) + + cos_before = (pre_dir * post_dir).sum(axis=1) + cos_after = rotated[:, 2] + np.testing.assert_allclose(cos_after, cos_before, atol=1e-5) + np.testing.assert_allclose(np.linalg.norm(rotated, axis=1), 1.0, atol=1e-5) + + recovered = inv_local_frame_rotation(pre_dir, rotated) + np.testing.assert_allclose(recovered, post_dir, atol=1e-5) + + def test_local_frame_rotation_normalizes_non_unit_pre_dir(): """A pre_dir with float32-drift norm (not exactly 1) must still produce the same result as its exactly-normalized counterpart, not a skewed frame."""