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,197 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Cut a new raw generation or processed schema version for the geant_steps
|
||||
dataset tree (see scripts/migrate_geant_steps.py for the layout):
|
||||
|
||||
raw/<kind>/<gen>/<detector>/shard-NNN.root
|
||||
processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet
|
||||
|
||||
`gen` bumps when the underlying ROOT changes (geometry/physics-list/macro).
|
||||
`schema` bumps when the parquet export (steps_to_parquet.py or similar)
|
||||
changes, and is scoped to its gen — a new gen always starts back at schema1.
|
||||
|
||||
Creates the new (empty) target directory and appends a dated, reasoned entry
|
||||
to VERSIONS.md. Defaults to a dry run; pass --execute to apply.
|
||||
|
||||
Usage:
|
||||
bump_dataset_version.py bump-gen --kind steps --reason "switched EM physics list"
|
||||
bump_dataset_version.py bump-schema --kind steps --gen gen1 --reason "added e_sec column"
|
||||
bump_dataset_version.py status
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import datetime as dt
|
||||
import re
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
GEN_RE = re.compile(r"^gen(\d+)$")
|
||||
SCHEMA_RE = re.compile(r"^schema(\d+)$")
|
||||
|
||||
|
||||
def _max_index(parent: Path, pattern: re.Pattern) -> int:
|
||||
"""Highest N across child dir names matching *pattern* (0 if none/missing)."""
|
||||
if not parent.is_dir():
|
||||
return 0
|
||||
best = 0
|
||||
for child in parent.iterdir():
|
||||
m = pattern.match(child.name)
|
||||
if m and child.is_dir():
|
||||
best = max(best, int(m.group(1)))
|
||||
return best
|
||||
|
||||
|
||||
def _git_user_name() -> str | None:
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["git", "config", "user.name"], capture_output=True, text=True, timeout=2
|
||||
)
|
||||
except OSError:
|
||||
return None
|
||||
name = out.stdout.strip()
|
||||
return name or None
|
||||
|
||||
|
||||
def plan_bump_gen(
|
||||
root: Path, kind: str, reason: str, by: str | None, date: str
|
||||
) -> tuple[list[Path], str]:
|
||||
"""New gen tag is one past the highest seen under raw/ or processed/ for *kind*
|
||||
— checking both, since a gen can exist in one tree before the other catches up."""
|
||||
next_gen = (
|
||||
max(
|
||||
_max_index(root / "raw" / kind, GEN_RE),
|
||||
_max_index(root / "processed" / kind, GEN_RE),
|
||||
)
|
||||
+ 1
|
||||
)
|
||||
gen_tag = f"gen{next_gen}"
|
||||
new_dirs = [
|
||||
root / "raw" / kind / gen_tag,
|
||||
root / "processed" / kind / gen_tag / "schema1",
|
||||
]
|
||||
by_suffix = f" ({by})" if by else ""
|
||||
log_line = f"- `{gen_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
|
||||
return new_dirs, log_line
|
||||
|
||||
|
||||
def plan_bump_schema(
|
||||
root: Path, kind: str, gen_tag: str, reason: str, by: str | None, date: str
|
||||
) -> tuple[list[Path], str]:
|
||||
if not GEN_RE.match(gen_tag):
|
||||
raise SystemExit(f"error: --gen must look like 'genN', got {gen_tag!r}")
|
||||
raw_gen_dir = root / "raw" / kind / gen_tag
|
||||
processed_gen_dir = root / "processed" / kind / gen_tag
|
||||
if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir():
|
||||
raise SystemExit(
|
||||
f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
|
||||
)
|
||||
next_schema = _max_index(processed_gen_dir, SCHEMA_RE) + 1
|
||||
schema_tag = f"schema{next_schema}"
|
||||
new_dirs = [processed_gen_dir / schema_tag]
|
||||
by_suffix = f" ({by})" if by else ""
|
||||
log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
|
||||
return new_dirs, log_line
|
||||
|
||||
|
||||
def apply_bump(root: Path, new_dirs: list[Path], log_line: str) -> None:
|
||||
for d in new_dirs:
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
versions_path = root / "VERSIONS.md"
|
||||
if not versions_path.exists():
|
||||
versions_path.write_text("# Dataset versions\n\n")
|
||||
with versions_path.open("a") as f:
|
||||
f.write(log_line + "\n")
|
||||
|
||||
|
||||
def print_status(root: Path) -> None:
|
||||
raw_root = root / "raw"
|
||||
if not raw_root.is_dir():
|
||||
print(f"no raw/ tree found under {root}")
|
||||
return
|
||||
for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
|
||||
kind = kind_dir.name
|
||||
gens = sorted(
|
||||
int(m.group(1))
|
||||
for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
|
||||
if m
|
||||
)
|
||||
print(f"kind={kind}")
|
||||
for gen in gens:
|
||||
gen_tag = f"gen{gen}"
|
||||
schema_dir = root / "processed" / kind / gen_tag
|
||||
schemas = sorted(
|
||||
int(m.group(1))
|
||||
for m in (
|
||||
SCHEMA_RE.match(p.name)
|
||||
for p in (schema_dir.iterdir() if schema_dir.is_dir() else [])
|
||||
if p.is_dir()
|
||||
)
|
||||
if m
|
||||
)
|
||||
schema_str = ", ".join(f"schema{s}" for s in schemas) or "(none)"
|
||||
print(f" {gen_tag}: {schema_str}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--root",
|
||||
default="/ceph/lbogner/geant_steps",
|
||||
help="Dataset root (default: /ceph/lbogner/geant_steps)",
|
||||
)
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
kind_help = "steps | hits | ... (default: steps)"
|
||||
|
||||
p_gen = sub.add_parser("bump-gen", help="Cut a new raw generation")
|
||||
p_gen.add_argument("--kind", default="steps", help=kind_help)
|
||||
p_gen.add_argument("--reason", required=True, help="Why this gen exists")
|
||||
p_gen.add_argument("--by", default=None, help="Attribution (default: git user.name)")
|
||||
p_gen.add_argument("--date", default=None, help="Override date (default: today, ISO)")
|
||||
p_gen.add_argument("--execute", action="store_true", help="Apply (default: dry run)")
|
||||
|
||||
p_schema = sub.add_parser("bump-schema", help="Cut a new schema within a gen")
|
||||
p_schema.add_argument("--kind", default="steps", help=kind_help)
|
||||
p_schema.add_argument("--gen", required=True, help="Existing gen tag, e.g. gen1")
|
||||
p_schema.add_argument("--reason", required=True, help="Why this schema exists")
|
||||
p_schema.add_argument("--by", default=None, help="Attribution (default: git user.name)")
|
||||
p_schema.add_argument("--date", default=None, help="Override date (default: today, ISO)")
|
||||
p_schema.add_argument("--execute", action="store_true", help="Apply (default: dry run)")
|
||||
|
||||
sub.add_parser("status", help="List existing gens/schemas per kind")
|
||||
|
||||
args = parser.parse_args()
|
||||
root = Path(args.root)
|
||||
if not root.is_dir():
|
||||
parser.error(f"{root} is not a directory")
|
||||
|
||||
if args.command == "status":
|
||||
print_status(root)
|
||||
return
|
||||
|
||||
date = args.date or dt.date.today().isoformat()
|
||||
by = args.by if args.by is not None else _git_user_name()
|
||||
|
||||
if args.command == "bump-gen":
|
||||
new_dirs, log_line = plan_bump_gen(root, args.kind, args.reason, by, date)
|
||||
else:
|
||||
new_dirs, log_line = plan_bump_schema(
|
||||
root, args.kind, args.gen, args.reason, by, date
|
||||
)
|
||||
|
||||
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
|
||||
print("new directories:")
|
||||
for d in new_dirs:
|
||||
print(f" {d}")
|
||||
print("VERSIONS.md entry:")
|
||||
print(f" {log_line}")
|
||||
|
||||
if not args.execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
apply_bump(root, new_dirs, log_line)
|
||||
print("\nDone.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,258 @@
|
||||
#!/usr/bin/env python3
|
||||
"""One-time migration of /ceph/lbogner/geant_steps into the versioned layout:
|
||||
|
||||
raw/<kind>/<gen>/<detector>/shard-NNN.root
|
||||
processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet
|
||||
pools/<detector>/{dev,full,holdout,unclassified}.txt (manifests, no copied bytes)
|
||||
derived/predictions/<bucket>/<detector>_shard-NNN_predicted[_local].parquet
|
||||
|
||||
Walks the whole tree (regardless of which legacy subfolder a file currently sits
|
||||
in — train/, sampling_train/, sampling_train/small/, or loose at the top level),
|
||||
classifies every .root/.parquet by detector+shard via filename pattern, and looks
|
||||
up pool membership from POOL_ASSIGNMENT — decoupling "where it sits today" from
|
||||
"which pool it belongs to".
|
||||
|
||||
Defaults to a dry run (prints the planned moves and manifest contents). Pass
|
||||
--execute to actually move files and write manifests.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
# Must match giant.data.loader.MANIFEST_SUFFIX.
|
||||
MANIFEST_SUFFIX = ".manifest"
|
||||
|
||||
# Shard membership is keyed by detector name, independent of current location,
|
||||
# since today's pool assignment is encoded only by *which folder a file's
|
||||
# parquet was copied into* — not by anything in the filename itself.
|
||||
POOL_ASSIGNMENT: dict[str, dict[str, range | list[int]]] = {
|
||||
"pbwo4": {"full": range(0, 6), "holdout": range(6, 10)},
|
||||
"sampling_fe_scint": {"dev": [0], "full": [1, 2], "holdout": [3]},
|
||||
"sampling_pb_lar": {"dev": [0], "full": [1, 2], "holdout": [3]},
|
||||
"sampling_pb_scint": {"dev": [0], "full": [1, 2], "holdout": [3]},
|
||||
"sampling_w_scint_ecal": {"dev": [0], "full": [1, 2], "holdout": [3]},
|
||||
}
|
||||
|
||||
# This is the migration baseline: everything that exists today is one
|
||||
# generation of raw data (gen1) exported with one parquet schema (schema1).
|
||||
GEN = "gen1"
|
||||
SCHEMA = "schema1"
|
||||
LEGACY_GEN = "gen0"
|
||||
LEGACY_SCHEMA = "schema0"
|
||||
|
||||
PREDICTED_RE = re.compile(
|
||||
r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)"
|
||||
r"_predicted(?P<local>_local)?\.parquet$"
|
||||
)
|
||||
SHARD_RE = re.compile(
|
||||
r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$"
|
||||
)
|
||||
LEGACY_PREDICTED_RE = re.compile(
|
||||
r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$"
|
||||
)
|
||||
LEGACY_RE = re.compile(r"^pbwo4_10000events_hits\.(?P<ext>root|parquet)$")
|
||||
|
||||
|
||||
def pool_for(detector: str, shard: int) -> str:
|
||||
rules = POOL_ASSIGNMENT.get(detector)
|
||||
if rules is None:
|
||||
return "unclassified"
|
||||
for pool, shards in rules.items():
|
||||
if shard in shards:
|
||||
return pool
|
||||
return "unclassified"
|
||||
|
||||
|
||||
def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
|
||||
"""Return (moves, unrecognized) where moves is [(src, dst), ...]."""
|
||||
moves: list[tuple[Path, Path]] = []
|
||||
unrecognized: list[Path] = []
|
||||
|
||||
for path in sorted(src_root.rglob("*")):
|
||||
if not path.is_file():
|
||||
continue
|
||||
name = path.name
|
||||
|
||||
m = PREDICTED_RE.match(name)
|
||||
if m:
|
||||
detector, shard = m["detector"], int(m["shard"])
|
||||
suffix = "_predicted_local" if m["local"] else "_predicted"
|
||||
dst = (
|
||||
src_root
|
||||
/ "derived"
|
||||
/ "predictions"
|
||||
/ "unknown-checkpoint"
|
||||
/ f"{detector}_shard-{shard:03d}{suffix}.parquet"
|
||||
)
|
||||
moves.append((path, dst))
|
||||
continue
|
||||
|
||||
m = LEGACY_PREDICTED_RE.match(name)
|
||||
if m:
|
||||
suffix = "_predicted_local" if m["local"] else "_predicted"
|
||||
dst = (
|
||||
src_root
|
||||
/ "derived"
|
||||
/ "predictions"
|
||||
/ "unknown-checkpoint-legacy"
|
||||
/ f"pbwo4_hits_shard-000{suffix}.parquet"
|
||||
)
|
||||
moves.append((path, dst))
|
||||
continue
|
||||
|
||||
m = SHARD_RE.match(name)
|
||||
if m:
|
||||
detector, shard, ext = m["detector"], int(m["shard"]), m["ext"]
|
||||
if ext == "root":
|
||||
dst = (
|
||||
src_root
|
||||
/ "raw"
|
||||
/ "steps"
|
||||
/ GEN
|
||||
/ detector
|
||||
/ f"shard-{shard:03d}.root"
|
||||
)
|
||||
else:
|
||||
dst = (
|
||||
src_root
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ GEN
|
||||
/ SCHEMA
|
||||
/ detector
|
||||
/ f"shard-{shard:03d}.parquet"
|
||||
)
|
||||
moves.append((path, dst))
|
||||
continue
|
||||
|
||||
m = LEGACY_RE.match(name)
|
||||
if m:
|
||||
ext = m["ext"]
|
||||
if ext == "root":
|
||||
dst = (
|
||||
src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
|
||||
)
|
||||
else:
|
||||
dst = (
|
||||
src_root
|
||||
/ "processed"
|
||||
/ "hits"
|
||||
/ LEGACY_GEN
|
||||
/ LEGACY_SCHEMA
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
moves.append((path, dst))
|
||||
continue
|
||||
|
||||
unrecognized.append(path)
|
||||
|
||||
return moves, unrecognized
|
||||
|
||||
|
||||
def plan_manifests(src_root: Path) -> dict[Path, list[str]]:
|
||||
"""Map manifest path -> sorted list of parquet paths, each relative to the
|
||||
manifest's own directory (so the manifest stays valid if the tree moves)."""
|
||||
manifests: dict[Path, list[tuple[int, Path]]] = defaultdict(list)
|
||||
for detector, rules in POOL_ASSIGNMENT.items():
|
||||
manifest_dir = src_root / "pools" / detector
|
||||
for pool, shards in rules.items():
|
||||
manifest_path = manifest_dir / f"{pool}{MANIFEST_SUFFIX}"
|
||||
for shard in shards:
|
||||
dst = (
|
||||
src_root
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ GEN
|
||||
/ SCHEMA
|
||||
/ detector
|
||||
/ f"shard-{shard:03d}.parquet"
|
||||
)
|
||||
manifests[manifest_path].append((shard, dst))
|
||||
return {
|
||||
k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)]
|
||||
for k, v in manifests.items()
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"root",
|
||||
nargs="?",
|
||||
default="/ceph/lbogner/geant_steps",
|
||||
help="Dataset root to migrate in place (default: /ceph/lbogner/geant_steps)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--execute",
|
||||
action="store_true",
|
||||
help="Actually move files and write manifests (default: dry run / print plan only)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
src_root = Path(args.root)
|
||||
if not src_root.is_dir():
|
||||
parser.error(f"{src_root} is not a directory")
|
||||
|
||||
moves, unrecognized = plan_moves(src_root)
|
||||
manifests = plan_manifests(src_root)
|
||||
|
||||
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'}: {src_root} ===\n")
|
||||
|
||||
print(f"-- {len(moves)} file move(s) --")
|
||||
for src, dst in moves:
|
||||
print(f" {src.relative_to(src_root)} -> {dst.relative_to(src_root)}")
|
||||
|
||||
print(f"\n-- {len(manifests)} manifest(s) --")
|
||||
for manifest_path, lines in sorted(manifests.items()):
|
||||
print(f" {manifest_path.relative_to(src_root)}:")
|
||||
for line in lines:
|
||||
print(f" {line}")
|
||||
|
||||
if unrecognized:
|
||||
print(f"\n-- {len(unrecognized)} UNRECOGNIZED file(s), left in place --")
|
||||
for path in unrecognized:
|
||||
print(f" {path.relative_to(src_root)}")
|
||||
|
||||
if not args.execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
for src, dst in moves:
|
||||
if dst.exists():
|
||||
raise FileExistsError(f"refusing to overwrite existing file: {dst}")
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.move(str(src), str(dst))
|
||||
|
||||
for manifest_path, lines in manifests.items():
|
||||
manifest_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
manifest_path.write_text("\n".join(lines) + "\n")
|
||||
|
||||
versions_path = src_root / "VERSIONS.md"
|
||||
if not versions_path.exists():
|
||||
versions_path.write_text(
|
||||
"# Dataset versions\n\n"
|
||||
f"- `{GEN}` / `{SCHEMA}` (steps) — migration baseline: all data that existed "
|
||||
"before the directory restructure, generation and export details unrecorded.\n"
|
||||
f"- `{LEGACY_GEN}` / `{LEGACY_SCHEMA}` (hits) — legacy single-shot "
|
||||
"`pbwo4_10000events_hits` dataset (Hits tree, predates the 10k-shard scheme).\n"
|
||||
)
|
||||
|
||||
# Legacy pool dirs (train/, sampling_train/, sampling_train/small/) are now
|
||||
# empty since their contents were classified by filename, not location —
|
||||
# remove them, but only if a move actually emptied them.
|
||||
for stale_dir in ("train", "sampling_train/small", "sampling_train"):
|
||||
d = src_root / stale_dir
|
||||
if d.is_dir() and not any(d.iterdir()):
|
||||
d.rmdir()
|
||||
print(f"removed now-empty directory: {d.relative_to(src_root)}")
|
||||
|
||||
print("\nDone.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,251 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Convert many ROOT files to Parquet by fanning out to steps_to_parquet.py.
|
||||
|
||||
steps_to_parquet.py itself converts a list of files one at a time; this wraps
|
||||
it to run up to --jobs conversions concurrently, each as its own subprocess
|
||||
(invoked with the same Python executable running this script, so it picks up
|
||||
the active venv/uv environment automatically).
|
||||
|
||||
Inputs must live under <dataset-root>/raw/<kind>/<gen>/<detector>/<file>.root
|
||||
(see scripts/migrate_geant_steps.py) — each is written to the matching
|
||||
processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet, where <schema>
|
||||
defaults to the highest schemaN already under processed/<kind>/<gen>/ (pass
|
||||
--schema to pick a specific one, e.g. one just created by
|
||||
bump_dataset_version.py bump-schema). A file that doesn't fit that layout is
|
||||
rejected up front, before any conversion runs.
|
||||
|
||||
Usage:
|
||||
steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/shard-000.root raw/steps/gen1/pbwo4/shard-001.root
|
||||
steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --jobs 8
|
||||
steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --schema schema2
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
|
||||
_STEPS_TO_PARQUET = Path(__file__).resolve().parent / "steps_to_parquet.py"
|
||||
|
||||
# Must match scripts/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
|
||||
GEN_RE = re.compile(r"^gen\d+$")
|
||||
SCHEMA_RE = re.compile(r"^schema(\d+)$")
|
||||
|
||||
|
||||
class DestinationError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
def latest_schema_tag(processed_gen_dir: Path) -> str | None:
|
||||
"""Highest schemaN dir directly under *processed_gen_dir*, or None if none exist."""
|
||||
if not processed_gen_dir.is_dir():
|
||||
return None
|
||||
best_tag, best_n = None, -1
|
||||
for child in processed_gen_dir.iterdir():
|
||||
m = SCHEMA_RE.match(child.name)
|
||||
if m and child.is_dir() and int(m.group(1)) > best_n:
|
||||
best_tag, best_n = child.name, int(m.group(1))
|
||||
return best_tag
|
||||
|
||||
|
||||
def resolve_destination(
|
||||
root_file: Path, dataset_root: Path, schema_override: str | None
|
||||
) -> Path:
|
||||
"""Map raw/<kind>/<gen>/<detector>/<file>.root (relative to *dataset_root*)
|
||||
to processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet.
|
||||
|
||||
Raises DestinationError if *root_file* doesn't fit that layout, or if no
|
||||
schema can be determined (no --schema and none exists yet for that gen).
|
||||
"""
|
||||
dataset_root = dataset_root.resolve()
|
||||
root_file = root_file.resolve()
|
||||
try:
|
||||
rel = root_file.relative_to(dataset_root)
|
||||
except ValueError:
|
||||
raise DestinationError(f"{root_file} is not under dataset root {dataset_root}")
|
||||
|
||||
parts = rel.parts
|
||||
if (
|
||||
len(parts) != 5
|
||||
or parts[0] != "raw"
|
||||
or not GEN_RE.match(parts[2])
|
||||
or not parts[4].endswith(".root")
|
||||
):
|
||||
raise DestinationError(
|
||||
f"{root_file} does not match raw/<kind>/<gen>/<detector>/<file>.root "
|
||||
f"under {dataset_root} (got relative path: {rel})"
|
||||
)
|
||||
_, kind, gen_tag, detector, filename = parts
|
||||
shard_stem = Path(filename).stem
|
||||
|
||||
processed_gen_dir = dataset_root / "processed" / kind / gen_tag
|
||||
schema_tag = schema_override or latest_schema_tag(processed_gen_dir)
|
||||
if schema_tag is None:
|
||||
raise DestinationError(
|
||||
f"no schema exists yet under {processed_gen_dir} — pass --schema or run "
|
||||
"bump_dataset_version.py bump-schema first"
|
||||
)
|
||||
return processed_gen_dir / schema_tag / detector / f"{shard_stem}.parquet"
|
||||
|
||||
|
||||
def _convert_one(
|
||||
root_file: str,
|
||||
batch_size: str,
|
||||
tree: str,
|
||||
compression: str,
|
||||
steps_to_parquet_path: Path,
|
||||
output_path: Path | None,
|
||||
) -> tuple[str, int, str, str]:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(steps_to_parquet_path),
|
||||
root_file,
|
||||
"--batch-size",
|
||||
batch_size,
|
||||
"--tree",
|
||||
tree,
|
||||
"--compression",
|
||||
compression,
|
||||
]
|
||||
if output_path is not None:
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
cmd += ["--output", str(output_path)]
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
return root_file, result.returncode, result.stdout, result.stderr
|
||||
|
||||
|
||||
def run_parallel(
|
||||
root_files: list[str],
|
||||
jobs: int,
|
||||
batch_size: str = "100 MB",
|
||||
tree: str = "Steps",
|
||||
compression: str = "snappy",
|
||||
steps_to_parquet_path: Path = _STEPS_TO_PARQUET,
|
||||
output_for: dict[str, Path] | None = None,
|
||||
) -> list[tuple[str, int, str, str]]:
|
||||
"""Run one steps_to_parquet.py subprocess per file, up to *jobs* at a time.
|
||||
|
||||
*output_for*, if given, maps each root_file to the parquet path it should
|
||||
be written to (passed through as steps_to_parquet.py's --output); files
|
||||
missing from the map fall back to steps_to_parquet.py's own default
|
||||
(parquet written next to the input .root).
|
||||
|
||||
Returns one (root_file, returncode, stdout, stderr) tuple per file, in
|
||||
completion order (not necessarily input order).
|
||||
"""
|
||||
output_for = output_for or {}
|
||||
results = []
|
||||
with ThreadPoolExecutor(max_workers=jobs) as pool:
|
||||
futures = {
|
||||
pool.submit(
|
||||
_convert_one,
|
||||
f,
|
||||
batch_size,
|
||||
tree,
|
||||
compression,
|
||||
steps_to_parquet_path,
|
||||
output_for.get(f),
|
||||
): f
|
||||
for f in root_files
|
||||
}
|
||||
for future in as_completed(futures):
|
||||
root_file, code, stdout, stderr = future.result()
|
||||
status = "ok" if code == 0 else f"FAILED (exit {code})"
|
||||
print(f"\n=== {root_file}: {status} ===")
|
||||
if stdout:
|
||||
print(stdout, end="")
|
||||
if stderr:
|
||||
print(stderr, end="", file=sys.stderr)
|
||||
results.append((root_file, code, stdout, stderr))
|
||||
return results
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert many ROOT files to Parquet in parallel via steps_to_parquet.py."
|
||||
)
|
||||
parser.add_argument("root_files", nargs="+", help="Input ROOT file(s)")
|
||||
parser.add_argument(
|
||||
"-j",
|
||||
"--jobs",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Number of conversions to run in parallel (default: 4)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-size",
|
||||
default="100 MB",
|
||||
help="Uproot read batch size (default: '100 MB'). E.g. '50 MB', '500000' (rows).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tree",
|
||||
default="Steps",
|
||||
help="Tree name inside the ROOT file (default: Steps)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--compression",
|
||||
default="snappy",
|
||||
choices=["snappy", "lz4", "zstd", "gzip", "none"],
|
||||
help="Parquet compression codec (default: snappy)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset-root",
|
||||
default="/ceph/lbogner/geant_steps",
|
||||
help="Dataset root containing raw/ and processed/ (default: /ceph/lbogner/geant_steps)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--schema",
|
||||
default=None,
|
||||
help="Schema tag to write parquets under, e.g. schema2 "
|
||||
"(default: highest schemaN already under processed/<kind>/<gen>/)",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.jobs < 1:
|
||||
parser.error("--jobs must be >= 1")
|
||||
|
||||
dataset_root = Path(args.dataset_root)
|
||||
output_for: dict[str, Path] = {}
|
||||
errors: list[str] = []
|
||||
for f in args.root_files:
|
||||
try:
|
||||
output_for[f] = resolve_destination(Path(f), dataset_root, args.schema)
|
||||
except DestinationError as exc:
|
||||
errors.append(str(exc))
|
||||
|
||||
if errors:
|
||||
for err in errors:
|
||||
print(f"error: {err}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
for root_file, dest in output_for.items():
|
||||
print(f"{root_file} -> {dest}")
|
||||
|
||||
results = run_parallel(
|
||||
args.root_files,
|
||||
jobs=args.jobs,
|
||||
batch_size=args.batch_size,
|
||||
tree=args.tree,
|
||||
compression=args.compression,
|
||||
output_for=output_for,
|
||||
)
|
||||
|
||||
failures = [root_file for root_file, code, _, _ in results if code != 0]
|
||||
if failures:
|
||||
print(f"\n{len(failures)} of {len(results)} conversion(s) failed:", file=sys.stderr)
|
||||
for root_file in failures:
|
||||
print(f" {root_file}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"\nAll {len(results)} conversion(s) completed.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user