Add tooling for a versioned geant_steps dataset layout
Introduces raw/<kind>/<gen>/<detector>/shard-NNN.root and processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet as the dataset convention, plus scripts to operate on it: migrate_geant_steps.py for the one-time move into this layout, bump_dataset_version.py to cut new gen/schema versions with a logged reason, steps_to_parquet_parallel.py to convert ROOT shards to parquet in parallel and place them correctly, and create_root_files.py to generate new ROOT shards via a minicalosim executable. The loader gains .manifest file support so pools/ (train/dev/ holdout shard lists) can be passed straight to `giant train`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,282 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate new ROOT shards by running a minicalosim executable (e.g.
|
||||
run_pbwo4, run_sampling) and filing the output into the dataset's raw/ tree:
|
||||
|
||||
raw/<kind>/<gen>/<detector>/shard-NNN.root
|
||||
|
||||
These executables take `[configName] nEvents` and always write a fixed-name
|
||||
*.root file into the current directory — so running several in parallel
|
||||
needs separate working directories, and the output filename has to be
|
||||
discovered rather than assumed (it differs per executable: run_pbwo4 writes
|
||||
pbwo4_<n>events_hits.root, run_sampling writes sampling_<config>_<n>events_hits.root,
|
||||
others may differ again). This script gives each run its own scratch
|
||||
directory under <dataset-root>/.sim-tmp/, requires exactly one *.root to
|
||||
appear there, and moves it to the next free shard index for that detector
|
||||
(existing shards are never overwritten).
|
||||
|
||||
--gen must already exist under raw/<kind>/ — create one first with
|
||||
bump_dataset_version.py bump-gen.
|
||||
|
||||
Usage:
|
||||
create_root_files.py --executable build/run_pbwo4 --detector pbwo4 \\
|
||||
--gen gen1 --num-files 4 --events-per-file 10000 --execute
|
||||
|
||||
create_root_files.py --executable build/run_sampling \\
|
||||
--detector sampling_pb_scint:pb_scint \\
|
||||
--detector sampling_fe_scint:fe_scint \\
|
||||
--gen gen1 --num-files 4 --events-per-file 10000 --jobs 8 --execute
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
|
||||
# Must match scripts/bump_dataset_version.py's GEN_RE.
|
||||
GEN_RE = re.compile(r"^gen\d+$")
|
||||
SHARD_RE = re.compile(r"^shard-(\d+)\.root$")
|
||||
|
||||
|
||||
class PlanError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SimJob:
|
||||
detector: str
|
||||
config: str | None
|
||||
shard_index: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class JobResult:
|
||||
job: SimJob
|
||||
ok: bool
|
||||
dest: Path | None
|
||||
message: str
|
||||
stdout: str
|
||||
stderr: str
|
||||
|
||||
|
||||
def parse_detector_spec(spec: str) -> tuple[str, str | None]:
|
||||
"""'sampling_pb_scint:pb_scint' -> ('sampling_pb_scint', 'pb_scint');
|
||||
'pbwo4' -> ('pbwo4', None)."""
|
||||
if ":" in spec:
|
||||
label, config = spec.split(":", 1)
|
||||
if not label or not config:
|
||||
raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG")
|
||||
return label, config
|
||||
return spec, None
|
||||
|
||||
|
||||
def next_shard_index(detector_dir: Path) -> int:
|
||||
"""One past the highest existing shard-NNN.root in *detector_dir* (0 if none/missing)."""
|
||||
if not detector_dir.is_dir():
|
||||
return 0
|
||||
best = -1
|
||||
for child in detector_dir.iterdir():
|
||||
m = SHARD_RE.match(child.name)
|
||||
if m and child.is_file():
|
||||
best = max(best, int(m.group(1)))
|
||||
return best + 1
|
||||
|
||||
|
||||
def plan_jobs(
|
||||
detector_specs: list[str],
|
||||
num_files: int,
|
||||
dataset_root: Path,
|
||||
kind: str,
|
||||
gen: str,
|
||||
) -> list[SimJob]:
|
||||
if not GEN_RE.match(gen):
|
||||
raise PlanError(f"--gen must look like 'genN', got {gen!r}")
|
||||
gen_dir = dataset_root / "raw" / kind / gen
|
||||
if not gen_dir.is_dir():
|
||||
raise PlanError(
|
||||
f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first"
|
||||
)
|
||||
|
||||
jobs = []
|
||||
for spec in detector_specs:
|
||||
label, config = parse_detector_spec(spec)
|
||||
start = next_shard_index(gen_dir / label)
|
||||
for i in range(num_files):
|
||||
jobs.append(SimJob(detector=label, config=config, shard_index=start + i))
|
||||
return jobs
|
||||
|
||||
|
||||
def run_job(
|
||||
job: SimJob,
|
||||
executable: Path,
|
||||
events_per_file: int,
|
||||
dataset_root: Path,
|
||||
kind: str,
|
||||
gen: str,
|
||||
tmp_root: Path,
|
||||
) -> JobResult:
|
||||
workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
|
||||
workdir.mkdir(parents=True)
|
||||
|
||||
cmd = [str(executable)]
|
||||
if job.config:
|
||||
cmd.append(job.config)
|
||||
cmd.append(str(events_per_file))
|
||||
|
||||
result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True)
|
||||
|
||||
if result.returncode != 0:
|
||||
return JobResult(
|
||||
job, False, None,
|
||||
f"executable exited {result.returncode}",
|
||||
result.stdout, result.stderr,
|
||||
)
|
||||
|
||||
produced = sorted(workdir.glob("*.root"))
|
||||
if len(produced) != 1:
|
||||
return JobResult(
|
||||
job, False, None,
|
||||
f"expected exactly one .root output in {workdir}, found {len(produced)}: "
|
||||
f"{[p.name for p in produced]}",
|
||||
result.stdout, result.stderr,
|
||||
)
|
||||
|
||||
dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
|
||||
if dest.exists():
|
||||
return JobResult(
|
||||
job, False, None, f"refusing to overwrite existing {dest}",
|
||||
result.stdout, result.stderr,
|
||||
)
|
||||
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.move(str(produced[0]), str(dest))
|
||||
# rmtree, not rmdir: the executable may leave other side-effect files
|
||||
# (logs, seed state, ...) behind in its scratch workdir besides the .root.
|
||||
shutil.rmtree(workdir)
|
||||
return JobResult(job, True, dest, "ok", result.stdout, result.stderr)
|
||||
|
||||
|
||||
def run_all(
|
||||
jobs: list[SimJob],
|
||||
executable: Path,
|
||||
events_per_file: int,
|
||||
dataset_root: Path,
|
||||
kind: str,
|
||||
gen: str,
|
||||
max_workers: int,
|
||||
tmp_root: Path,
|
||||
) -> list[JobResult]:
|
||||
results = []
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as pool:
|
||||
futures = {
|
||||
pool.submit(
|
||||
run_job, job, executable, events_per_file, dataset_root, kind, gen, tmp_root
|
||||
): job
|
||||
for job in jobs
|
||||
}
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
label = f"{result.job.detector} shard-{result.job.shard_index:03d}"
|
||||
status = "ok" if result.ok else f"FAILED: {result.message}"
|
||||
print(f"\n=== {label}: {status} ===")
|
||||
if result.stdout:
|
||||
print(result.stdout, end="")
|
||||
if result.stderr:
|
||||
print(result.stderr, end="", file=sys.stderr)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--executable", required=True, type=Path, help="Built minicalosim run_* executable"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--detector",
|
||||
action="append",
|
||||
required=True,
|
||||
metavar="NAME[:CONFIG]",
|
||||
help="Dataset detector label, optionally with ':CONFIG' to pass as the "
|
||||
"executable's config-name argument (e.g. sampling_pb_scint:pb_scint). "
|
||||
"Omit ':CONFIG' for executables that take no config selector (e.g. run_pbwo4). "
|
||||
"Repeatable.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-files", type=int, required=True, help="New shards to create per detector"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--events-per-file", type=int, required=True, help="nEvents passed to the executable"
|
||||
)
|
||||
parser.add_argument("--kind", default="steps", help="steps | hits | ... (default: steps)")
|
||||
parser.add_argument("--gen", required=True, help="Existing gen tag under raw/<kind>/, e.g. gen1")
|
||||
parser.add_argument(
|
||||
"--dataset-root",
|
||||
default="/ceph/lbogner/geant_steps",
|
||||
help="Dataset root (default: /ceph/lbogner/geant_steps)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-j", "--jobs", type=int, default=4, help="Parallel simulation runs (default: 4)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--execute", action="store_true", help="Actually run jobs (default: dry run / print plan)"
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.jobs < 1:
|
||||
parser.error("--jobs must be >= 1")
|
||||
if args.num_files < 1:
|
||||
parser.error("--num-files must be >= 1")
|
||||
if args.events_per_file < 1:
|
||||
parser.error("--events-per-file must be >= 1")
|
||||
if not args.executable.is_file() or not os.access(args.executable, os.X_OK):
|
||||
parser.error(f"{args.executable} is not an executable file")
|
||||
|
||||
dataset_root = Path(args.dataset_root)
|
||||
try:
|
||||
jobs = plan_jobs(args.detector, args.num_files, dataset_root, args.kind, args.gen)
|
||||
except PlanError as exc:
|
||||
parser.error(str(exc))
|
||||
|
||||
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
|
||||
print(f"executable: {args.executable}")
|
||||
for job in jobs:
|
||||
cmd = [str(args.executable)] + ([job.config] if job.config else []) + [str(args.events_per_file)]
|
||||
dest = dataset_root / "raw" / args.kind / args.gen / job.detector / f"shard-{job.shard_index:03d}.root"
|
||||
print(f" {' '.join(cmd)} -> {dest}")
|
||||
|
||||
if not args.execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
tmp_root = dataset_root / ".sim-tmp"
|
||||
tmp_root.mkdir(parents=True, exist_ok=True)
|
||||
results = run_all(
|
||||
jobs, args.executable, args.events_per_file, dataset_root, args.kind, args.gen,
|
||||
max_workers=args.jobs, tmp_root=tmp_root,
|
||||
)
|
||||
if tmp_root.is_dir() and not any(tmp_root.iterdir()):
|
||||
tmp_root.rmdir()
|
||||
|
||||
failures = [r for r in results if not r.ok]
|
||||
if failures:
|
||||
print(f"\n{len(failures)} of {len(results)} job(s) failed:", file=sys.stderr)
|
||||
for r in failures:
|
||||
print(f" {r.job.detector} shard-{r.job.shard_index:03d}: {r.message}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"\nAll {len(results)} job(s) completed.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user