Stream giant rollout output instead of buffering the whole run

_Recorder previously accumulated every generated step across all events/ tracks/steps in Python lists, materialised once at the end and written via a single pq.write_table — memory scaled with n_events * max_steps * avg_tracks_per_event. rollout() now takes an optional on_chunk callback that streams each non-empty batch immediately (fixed per-key dtypes via _RECORD_DTYPES keep every chunk's table schema identical, which pq.ParquetWriter requires across writes); giant rollout wires this to an incrementally-written ParquetWriter, mirroring the row-group streaming giant predict already does on its input side. Without on_chunk, rollout() keeps its old buffered return for existing callers/tests. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-15 11:10:04 +02:00
parent ac5fbd14b4
commit 550dc679c7
3 changed files with 218 additions and 31 deletions
@@ -804,7 +804,29 @@ def rollout(
    seeds = _seed_from_data(files, n_events)
    typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)")

-    records = run_rollout(
+    out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    # Written incrementally as each batch of steps is produced, rather than
+    # buffering the whole run (which scales with n_events * max_steps *
+    # avg_tracks_per_event) — mirrors the row-group streaming `giant predict`
+    # already does on its input side.
+    writer: pq.ParquetWriter | None = None
+
+    def _write_chunk(row: dict[str, np.ndarray]) -> None:
+        nonlocal writer
+        table = pa.table(row)
+        if writer is None:
+            table = table.replace_schema_metadata(
+                {
+                    PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE,
+                    PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
+                }
+            )
+            writer = pq.ParquetWriter(out, table.schema)
+        writer.write_table(table)
+
+    summary = run_rollout(
        model,
        sec_decoder,
        oracle,
@@ -820,18 +842,10 @@ def rollout(
        device=_device,
        max_tracks_per_event=max_tracks_per_event,
        escape_threshold=escape_threshold,
+        on_chunk=_write_chunk,
    )
-
-    out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
-    out.parent.mkdir(parents=True, exist_ok=True)
-
-    table = pa.table(records).replace_schema_metadata(
-        {
-            PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE,
-            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
-        }
-    )
-    pq.write_table(table, out)
+    if writer is not None:
+        writer.close()

    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path)
    ref = yaml.safe_load(ref_path.read_text())
@@ -848,10 +862,8 @@ def rollout(
    )
    ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))

-    n_rows = len(records["event_id"])
-    reasons = Counter(r for r in records["termination_reason"].tolist() if r)
-    typer.echo(f"wrote {n_rows:,} step rows → {out}")
-    typer.echo(f"terminations: {dict(reasons)}")
+    typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}")
+    typer.echo(f"terminations: {summary['termination_reason_counts']}")
    typer.echo(f"reference: {ref_path}")


@@ -16,6 +16,9 @@ treated as detector leakage and not deposited.

 from __future__ import annotations

+from collections import Counter
+from typing import Callable
+
 import numpy as np
 import torch

@@ -89,32 +92,95 @@ def _concat_frontiers(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarra
    return {k: np.concatenate([p[k] for p in parts], axis=0) for k in parts[0]}


-class _Recorder:
-    """Accumulates per-step rows into column lists, materialised at the end."""
+# Fixed per-key dtype, so every chunk table has an identical schema — needed
+# for `giant rollout --on_chunk` to stream chunks straight into one
+# pq.ParquetWriter (which requires matching schemas across writes), and a
+# side benefit even in the buffered path since np.concatenate would otherwise
+# silently upcast any stray int32/float32 chunk to the majority dtype.
+_RECORD_DTYPES: dict[str, type] = {
+    "event_id": np.int64,
+    "track_id": np.int64,
+    "parent_id": np.int64,
+    "generation": np.int64,
+    "step_no": np.int64,
+    "pdg": np.int64,
+    "pre_x": np.float64,
+    "pre_y": np.float64,
+    "pre_z": np.float64,
+    "pre_E": np.float64,
+    "pre_dx": np.float64,
+    "pre_dy": np.float64,
+    "pre_dz": np.float64,
+    "post_x": np.float64,
+    "post_y": np.float64,
+    "post_z": np.float64,
+    "post_E": np.float64,
+    "post_dx": np.float64,
+    "post_dy": np.float64,
+    "post_dz": np.float64,
+    "edep": np.float64,
+    "step_length": np.float64,
+    "material": object,
+    "layer_id": np.int64,
+    "n_sec_pred": np.int64,
+    "termination_reason": object,
+}

-    def __init__(self) -> None:
-        self._cols: dict[str, list] = {k: [] for k in _RECORD_KEYS}
+
+class _Recorder:
+    """Accumulates per-step rows into column lists, materialised at the end —
+    or, when `sink` is given, streams each non-empty chunk to it immediately
+    instead, keeping only row-count / termination-reason summaries in memory.
+
+    The streaming path is what lets `giant rollout` write output incrementally
+    (see `rollout`'s `on_chunk` parameter): without it, a whole run's steps —
+    scaling with `n_events * max_steps * avg_tracks_per_event` — would sit in
+    RAM until the very end.
+    """
+
+    def __init__(
+        self, sink: Callable[[dict[str, np.ndarray]], None] | None = None
+    ) -> None:
+        self._sink = sink
+        self._cols: dict[str, list] | None = (
+            None if sink is not None else {k: [] for k in _RECORD_KEYS}
+        )
+        self.n_rows = 0
+        self.termination_reason_counts: Counter[str] = Counter()

    def add(self, **cols) -> None:
        n = len(cols["event_id"])
        if n == 0:
            return
-        for k in _RECORD_KEYS:
-            v = cols[k]
-            self._cols[k].append(np.asarray(v).reshape(n))
+        row = {
+            k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n)
+            for k in _RECORD_KEYS
+        }
+        self.n_rows += n
+        reasons = row["termination_reason"]
+        nonempty = reasons[reasons != ""]
+        if len(nonempty):
+            for r, c in zip(*np.unique(nonempty, return_counts=True)):
+                self.termination_reason_counts[str(r)] += int(c)
+
+        if self._sink is not None:
+            self._sink(row)
+        else:
+            assert self._cols is not None
+            for k in _RECORD_KEYS:
+                self._cols[k].append(row[k])

    def to_dict(self) -> dict[str, np.ndarray]:
+        assert self._cols is not None, (
+            "to_dict() is unavailable when streaming to a sink — use "
+            "n_rows/termination_reason_counts instead"
+        )
        out = {}
        for k, chunks in self._cols.items():
            if chunks:
                out[k] = np.concatenate(chunks, axis=0)
            else:
-                out[k] = np.empty(
-                    0,
-                    dtype=object
-                    if k in ("material", "termination_reason")
-                    else np.float64,
-                )
+                out[k] = np.empty(0, dtype=_RECORD_DTYPES[k])
        return out


@@ -209,8 +275,20 @@ def rollout(
    device: torch.device | None = None,
    max_tracks_per_event: int | None = None,
    escape_threshold: float | None = None,
+    on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
 ) -> dict[str, np.ndarray]:
-    """Run showers to completion; return a step-record dict (see _RECORD_KEYS)."""
+    """Run showers to completion.
+
+    By default, returns a step-record dict (see _RECORD_KEYS) with the whole
+    run's rows materialised in memory.
+
+    If `on_chunk` is given, every non-empty batch of rows is streamed to it as
+    soon as it's produced instead — no per-run buffering — and this returns a
+    small summary dict instead: `{"n_rows": int, "termination_reason_counts":
+    dict[str, int]}`. Use this for large `--n-events`/`--max-steps` runs,
+    where the full record set would otherwise scale with
+    `n_events * max_steps * avg_tracks_per_event`.
+    """
    device = device or torch.device("cpu")
    stage1_model.eval()
    sec_decoder.eval()
@@ -227,7 +305,7 @@ def rollout(
        seeds["pre_E"],
        seeds["pre_dir"],
    )
-    rec = _Recorder()
+    rec = _Recorder(sink=on_chunk)

    while len(frontier["event_id"]) > 0:
        next_parts: list[dict[str, np.ndarray]] = []
@@ -257,6 +335,11 @@ def rollout(
            )
        frontier = _concat_frontiers(next_parts)

+    if on_chunk is not None:
+        return {
+            "n_rows": rec.n_rows,
+            "termination_reason_counts": dict(rec.termination_reason_counts),
+        }
    return rec.to_dict()


@@ -1,5 +1,6 @@
 """Tests for the autoregressive shower rollout driver."""

+from collections import Counter
 from pathlib import Path
 from unittest.mock import patch

@@ -161,3 +162,94 @@ def test_max_tracks_cap_conserves_energy():
        leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
        assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
        assert len(np.unique(rec["track_id"][m])) <= 3
+
+
+# ── Streaming output (on_chunk) ──────────────────────────────────────────────
+
+
+def _run_streaming(on_chunk, **kwargs):
+    torch.manual_seed(0)
+    np.random.seed(0)
+    s1, s2 = _models()
+    cond, tgt = _norms()
+    seeds = kwargs.pop("seeds", None) or _seeds()
+    return rollout(
+        s1,
+        s2,
+        _oracle(),
+        seeds,
+        cond,
+        tgt,
+        PDG_MAP,
+        MAT_MAP,
+        energy_cutoff=kwargs.pop("energy_cutoff", 1.0),
+        max_steps=kwargs.pop("max_steps", 30),
+        steps=4,
+        batch_size=128,
+        max_tracks_per_event=kwargs.pop("max_tracks_per_event", 300),
+        escape_threshold=kwargs.pop("escape_threshold", 1e9),
+        on_chunk=on_chunk,
+    )
+
+
+def test_on_chunk_receives_every_row_exactly_once():
+    """Concatenating the streamed chunks must reproduce the buffered result."""
+    from giant.rollout import _RECORD_KEYS
+
+    buffered = _run()
+
+    chunks: list[dict[str, np.ndarray]] = []
+    summary = _run_streaming(chunks.append)
+
+    streamed = {k: np.concatenate([c[k] for c in chunks]) for k in _RECORD_KEYS}
+    assert summary["n_rows"] == len(buffered["event_id"])
+    assert len(streamed["event_id"]) == len(buffered["event_id"])
+    for k in _RECORD_KEYS:
+        np.testing.assert_array_equal(streamed[k], buffered[k])
+
+
+def test_on_chunk_summary_termination_reason_counts_match_buffered():
+    buffered = _run()
+    summary = _run_streaming(lambda row: None)
+
+    expected = Counter(r for r in buffered["termination_reason"].tolist() if r)
+    assert summary["termination_reason_counts"] == dict(expected)
+
+
+def test_on_chunk_never_buffers_full_records():
+    """Streaming mode must not accumulate rows for later to_dict() retrieval."""
+    from giant.rollout import _Recorder
+
+    rec = _Recorder(sink=lambda row: None)
+    rec.add(
+        event_id=np.array([0]),
+        track_id=np.array([0]),
+        parent_id=np.array([-1]),
+        generation=np.array([0]),
+        step_no=np.array([0]),
+        pdg=np.array([11]),
+        pre_x=np.array([0.0]),
+        pre_y=np.array([0.0]),
+        pre_z=np.array([0.0]),
+        pre_E=np.array([1.0]),
+        pre_dx=np.array([0.0]),
+        pre_dy=np.array([0.0]),
+        pre_dz=np.array([1.0]),
+        post_x=np.array([0.0]),
+        post_y=np.array([0.0]),
+        post_z=np.array([1.0]),
+        post_E=np.array([0.0]),
+        post_dx=np.array([0.0]),
+        post_dy=np.array([0.0]),
+        post_dz=np.array([1.0]),
+        edep=np.array([1.0]),
+        step_length=np.array([1.0]),
+        material=np.array(["G4_AIR"], dtype=object),
+        layer_id=np.array([0]),
+        n_sec_pred=np.array([0]),
+        termination_reason=np.array(["natural_end"], dtype=object),
+    )
+    assert rec.n_rows == 1
+    assert rec.termination_reason_counts == {"natural_end": 1}
+    with pytest.raises(AssertionError):
+        rec.to_dict()