feat: add restartable sharded graph preparation

docs/perlmutter.md

@@ -80,7 +80,7 @@ allocation capacity:
 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`,
+The script defaults to 64 process workers 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

scripts/slurm/prepare_tth_cp_even_odd.sh

@@ -29,4 +29,5 @@ uv run gnn4colliders prepare-shard \
   --config-name "${CONFIG}" \
   --shard-index "${SLURM_ARRAY_TASK_ID}" \
   --shard-count "${SHARD_COUNT}" \
-  --output-dir "${SHARD_DIR}"
+  --output-dir "${SHARD_DIR}" \
+  data.num_workers="${SLURM_CPUS_PER_TASK:-1}"

src/gnn4colliders/config/application.py

@@ -178,6 +178,11 @@ def _write_graph_shard(args: tuple[Any, ...]) -> tuple[Path, int]:
     return shard_path, len(samples)
 
 
+def _build_graph_samples_from_args(args: tuple[Any, ...]) -> list[GraphSample]:
+    data, files, labels, start, stop = args
+    return _build_graph_samples(data, files, labels, start, stop)
+
+
 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")
@@ -237,7 +242,29 @@ def prepare_shard(
             logger.warning("discarding incomplete or incompatible shard %s", target)
 
     data = plain(config.data)
-    samples = _build_graph_samples(data, files, labels, start, stop)
+    workers = max(0, int(data.get("num_workers", 0)))
+    local_shards = max(1, workers)
+    local_ranges = [
+        (
+            start + (stop - start) * index // local_shards,
+            start + (stop - start) * (index + 1) // local_shards,
+        )
+        for index in range(local_shards)
+    ]
+    args = [
+        (data, files, labels, local_start, local_stop)
+        for local_start, local_stop in local_ranges
+        if local_start < local_stop
+    ]
+    if workers > 1:
+        with ProcessPoolExecutor(
+            max_workers=workers,
+            mp_context=multiprocessing.get_context("spawn"),
+        ) as pool:
+            local_results = list(pool.map(_build_graph_samples_from_args, args))
+    else:
+        local_results = [_build_graph_samples_from_args(item) for item in args]
+    samples = [sample for result in local_results for sample in result]
     if len(samples) != expected:
         raise RuntimeError(
             f"shard {shard_index} produced {len(samples)} samples, expected {expected}"

tests/integration/test_parallel_prepare.py

@@ -93,6 +93,7 @@ def test_shards_are_restartable_and_cover_source(tmp_path: Path):
         "global_features": [],
         "fold_var": "eventNumber",
         "weight_var": "weight",
+        "num_workers": 2,
     }
     config = OmegaConf.create({"data": data})
     output = tmp_path / "shards"