feat: add restartable sharded graph preparation

docs/perlmutter.md

@@ -56,3 +56,32 @@ An unavailable DGL wheel, incompatible driver, missing cache, or mismatched
 cache schema should be fixed in the environment/input rather than hidden with
 package-level path changes. GPU kernels and distributed execution are not
 promised bitwise deterministic.
+
+## Large ROOT graph preparation
+
+Large samples should be prepared as persistent shards. Each Slurm array task
+reads one contiguous global event range and writes one `shard-NNNNN.pt` file;
+completed shards are validated and reused when an interrupted array is
+resubmitted. No full-cache in-memory merge is performed.
+
+Create the range manifest once:
+
+```bash
+uv run gnn4colliders prepare-manifest \
+  --config-name config_tth_cp_even_odd \
+  --shard-count 256 \
+  --output-dir outputs/ttH_cp_even_odd/shards
+```
+
+Submit the CPU array, limiting concurrent tasks to match the filesystem and
+allocation capacity:
+
+```bash
+sbatch --array=0-255%32 scripts/slurm/prepare_tth_cp_even_odd.sh
+```
+
+The script defaults to 64 CPUs per task and 256 shards. Set `SHARD_DIR`,
+`SHARD_COUNT`, or `CONFIG` at submission time to use another location or
+configuration. The current training loader still expects a single cache, so
+the next step after extraction is a streaming sharded loader; do not merge the
+full benchmark into one `.pt` file.

scripts/slurm/prepare_tth_cp_even_odd.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+# Prepare one independent graph-cache shard for the ttH CP benchmark.
+# Submit with: sbatch --array=0-255%32 scripts/slurm/prepare_tth_cp_even_odd.sh
+# Override SHARD_DIR, SHARD_COUNT, and CONFIG when needed.
+#SBATCH -N 1
+#SBATCH -C cpu
+#SBATCH -t 04:00:00
+#SBATCH -n 1
+#SBATCH -c 64
+#SBATCH -o logs/tth_cp_even_odd_%A_%a.out
+#SBATCH -e logs/tth_cp_even_odd_%A_%a.err
+
+set -euo pipefail
+cd "${SLURM_SUBMIT_DIR:-$(pwd)}"
+
+: "${SLURM_ARRAY_TASK_ID:?submit this script as a Slurm array}"
+SHARD_COUNT="${SHARD_COUNT:-256}"
+SHARD_DIR="${SHARD_DIR:-outputs/ttH_cp_even_odd/shards}"
+CONFIG="${CONFIG:-config_tth_cp_even_odd}"
+
+# Prevent native math libraries from multiplying the Slurm CPU allocation.
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export NUMEXPR_NUM_THREADS=1
+
+mkdir -p logs "${SHARD_DIR}"
+uv run gnn4colliders prepare-shard \
+  --config-name "${CONFIG}" \
+  --shard-index "${SLURM_ARRAY_TASK_ID}" \
+  --shard-count "${SHARD_COUNT}" \
+  --output-dir "${SHARD_DIR}"

src/gnn4colliders/cli/__init__.py

@@ -14,10 +14,20 @@ from gnn4colliders.config.application import (
     export_onnx,
     predict_or_evaluate,
     prepare,
+    prepare_shard,
     train,
+    write_shard_manifest,
 )
 
-_COMMANDS = {"prepare", "train", "evaluate", "predict", "export"}
+_COMMANDS = {
+    "prepare",
+    "prepare-manifest",
+    "prepare-shard",
+    "train",
+    "evaluate",
+    "predict",
+    "export",
+}
 
 
 def _is_main_process() -> bool:
@@ -30,7 +40,8 @@ def _help(command: str | None = None) -> None:
         print("Compose Hydra configuration with semantic overrides.")
     else:
         print(
-            "Usage: gnn4colliders <prepare|train|evaluate|predict|export> "
+            "Usage: gnn4colliders <prepare|prepare-manifest|prepare-shard|train|"
+            "evaluate|predict|export> "
             "[key=value ...]"
         )
         print("Use gnn4colliders <command> --help for command help.")
@@ -68,6 +79,39 @@ def _config_name(args: list[str]) -> tuple[str, list[str]]:
     return config_name, remaining
 
 
+def _shard_options(
+    args: list[str],
+) -> tuple[int | None, int | None, str | None, list[str]]:
+    values: dict[str, int | str] = {}
+    remaining: list[str] = []
+    names = {
+        "--shard-index": "index",
+        "--shard-count": "count",
+        "--output-dir": "output",
+    }
+    index = 0
+    while index < len(args):
+        argument = args[index]
+        matched = next((name for name in names if argument == name), None)
+        if matched:
+            if index + 1 >= len(args):
+                raise ValueError(f"{matched} requires a value")
+            values[names[matched]] = args[index + 1]
+            index += 2
+            continue
+        prefix = next((name for name in names if argument.startswith(name + "=")), None)
+        if prefix:
+            values[names[prefix]] = argument.split("=", 1)[1]
+            index += 1
+            continue
+        remaining.append(argument)
+        index += 1
+    shard_index = int(values["index"]) if "index" in values else None
+    shard_count = int(values["count"]) if "count" in values else None
+    output_dir = str(values["output"]) if "output" in values else None
+    return shard_index, shard_count, output_dir, remaining
+
+
 def main(argv: list[str] | None = None) -> int:
     args = list(sys.argv[1:] if argv is None else argv)
     if not args or args[0] in {"-h", "--help"}:
@@ -85,12 +129,35 @@ def main(argv: list[str] | None = None) -> int:
         return 0
     try:
         config_name, overrides = _config_name(args)
+        shard_index, shard_count, output_dir, overrides = _shard_options(overrides)
         config = _config(overrides, config_name=config_name)
         logging.basicConfig(
             level=getattr(logging, str(config.logging.level).upper(), logging.INFO)
         )
         if command == "prepare":
             print(f"prepared cache: {prepare(config)}")
+        elif command == "prepare-manifest":
+            if shard_count is None or output_dir is None:
+                raise ValueError(
+                    "prepare-manifest requires --shard-count and --output-dir"
+                )
+            manifest = write_shard_manifest(
+                config, shard_count=shard_count, output_dir=output_dir
+            )
+            print(f"wrote shard manifest: {manifest}")
+        elif command == "prepare-shard":
+            if shard_index is None or shard_count is None or output_dir is None:
+                raise ValueError(
+                    "prepare-shard requires --shard-index, --shard-count, and "
+                    "--output-dir"
+                )
+            shard = prepare_shard(
+                config,
+                shard_index=shard_index,
+                shard_count=shard_count,
+                output_dir=output_dir,
+            )
+            print(f"prepared graph shard: {shard}")
         elif command == "train":
             if _is_main_process():
                 print(f"saved checkpoint: {train(config)}")

src/gnn4colliders/config/application.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 import logging
 import multiprocessing
 import os
@@ -177,6 +178,108 @@ def _write_graph_shard(args: tuple[Any, ...]) -> tuple[Path, int]:
     return shard_path, len(samples)
 
 
+def _shard_range(total: int, shard_index: int, shard_count: int) -> tuple[int, int]:
+    if shard_count < 1:
+        raise ValueError("shard_count must be positive")
+    if shard_index < 0 or shard_index >= shard_count:
+        raise ValueError("shard_index must be in [0, shard_count)")
+    return (
+        total * shard_index // shard_count,
+        total * (shard_index + 1) // shard_count,
+    )
+
+
+def _prepare_source(config: DictConfig):
+    data = config.data
+    files, labels = resolve_data_files(data)
+    labels_per_file = (
+        list(labels) if isinstance(labels, (list, tuple)) else [labels] * len(files)
+    )
+    if not isinstance(labels, (list, tuple)) and isinstance(labels, Sequence):
+        labels_per_file = list(labels)
+    source = RootEventDataset(
+        files,
+        tree_name=str(data.tree_name),
+        label=labels_per_file,
+        feature_branches=plain(data.get("feature_branches")),
+        global_features=plain(data.get("global_features", [])),
+        fold_var=str(data.get("fold_var", "eventNumber")),
+        weight_var=data.get("weight_var"),
+    )
+    return source, files, labels_per_file
+
+
+def prepare_shard(
+    config: DictConfig,
+    *,
+    shard_index: int,
+    shard_count: int,
+    output_dir: str | Path,
+) -> Path:
+    """Prepare one persistent, restartable graph-cache shard.
+
+    Shards use contiguous global event ranges and are independent of one
+    another, making this function suitable for a Slurm job array. Existing
+    valid shards are reused so interrupted arrays can be resubmitted safely.
+    """
+    source, files, labels = _prepare_source(config)
+    start, stop = _shard_range(len(source), shard_index, shard_count)
+    target_dir = Path(output_dir)
+    target_dir.mkdir(parents=True, exist_ok=True)
+    target = target_dir / f"shard-{shard_index:05d}.pt"
+    expected = stop - start
+    if target.is_file():
+        try:
+            if len(GraphSampleCache(target).load()) == expected:
+                logger.info("reusing completed shard %s", target)
+                return target
+        except (OSError, RuntimeError, ValueError, EOFError):
+            logger.warning("discarding incomplete or incompatible shard %s", target)
+
+    data = plain(config.data)
+    samples = _build_graph_samples(data, files, labels, start, stop)
+    if len(samples) != expected:
+        raise RuntimeError(
+            f"shard {shard_index} produced {len(samples)} samples, expected {expected}"
+        )
+    temporary = target.with_suffix(target.suffix + f".tmp-{os.getpid()}")
+    GraphSampleCache(temporary).save(samples)
+    os.replace(temporary, target)
+    return target
+
+
+def write_shard_manifest(
+    config: DictConfig, *, shard_count: int, output_dir: str | Path
+) -> Path:
+    """Write an immutable description of the event ranges in a shard set."""
+    source, files, labels = _prepare_source(config)
+    target_dir = Path(output_dir)
+    target_dir.mkdir(parents=True, exist_ok=True)
+    manifest = {
+        "format": 1,
+        "total_events": len(source),
+        "shard_count": int(shard_count),
+        "tree_name": str(config.data.tree_name),
+        "files": [str(path) for path in files],
+        "labels": labels,
+        "shards": [
+            {
+                "index": index,
+                "start": start,
+                "stop": stop,
+                "path": f"shard-{index:05d}.pt",
+            }
+            for index in range(shard_count)
+            for start, stop in [_shard_range(len(source), index, shard_count)]
+        ],
+    }
+    target = target_dir / "manifest.json"
+    temporary = target.with_suffix(".json.tmp")
+    temporary.write_text(json.dumps(manifest, indent=2) + "\n")
+    os.replace(temporary, target)
+    return target
+
+
 def prepare(config: DictConfig) -> Path:
     if bool(config.distributed.enabled) or int(os.environ.get("WORLD_SIZE", "1")) > 1:
         raise ValueError(

tests/integration/test_parallel_prepare.py

@@ -6,7 +6,11 @@ import pytest
 import uproot
 from omegaconf import OmegaConf
 
-from gnn4colliders.config.application import prepare
+from gnn4colliders.config.application import (
+    prepare,
+    prepare_shard,
+    write_shard_manifest,
+)
 from gnn4colliders.data import GraphSampleCache
 
 
@@ -65,3 +69,49 @@ def test_parallel_prepare_matches_serial_cache(tmp_path: Path):
         assert left.label.equal(right.label)
         assert left.graph.ndata["features"].equal(right.graph.ndata["features"])
         assert left.graph.edata["features"].equal(right.graph.edata["features"])
+
+
+@pytest.mark.integration
+def test_shards_are_restartable_and_cover_source(tmp_path: Path):
+    pytest.importorskip("dgl")
+    root_file = tmp_path / "events.root"
+    _write_fixture(root_file, events=10)
+    data = {
+        "files": [str(root_file)],
+        "tree_name": "Events",
+        "feature_branches": [
+            ["jet_pt"],
+            ["jet_eta"],
+            ["jet_phi"],
+            "CALC_E",
+            [1.0],
+            [0.0],
+            "NODE_TYPE",
+        ],
+        "object_types": ["vector"],
+        "scales": [1, 1, 1, 1, 1, 1, 1],
+        "global_features": [],
+        "fold_var": "eventNumber",
+        "weight_var": "weight",
+    }
+    config = OmegaConf.create({"data": data})
+    output = tmp_path / "shards"
+
+    manifest = write_shard_manifest(config, shard_count=3, output_dir=output)
+    assert manifest.is_file()
+    paths = [
+        prepare_shard(
+            config,
+            shard_index=index,
+            shard_count=3,
+            output_dir=output,
+        )
+        for index in range(3)
+    ]
+    assert [len(GraphSampleCache(path).load()) for path in paths] == [3, 3, 4]
+    assert prepare_shard(
+        config,
+        shard_index=1,
+        shard_count=3,
+        output_dir=output,
+    ) == paths[1]