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:
Claude Sonnet 4.6
@@ -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