lars
320365606a
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>
283 lines
9.3 KiB
Python
283 lines
9.3 KiB
Python
#!/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()
|