lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Unify dataset/tooling scripts into a single dwarf Typer CLI

Browse Source 
Replace the five separately-hyphenated uv entry points (steps-to-parquet,
steps-to-parquet-parallel, migrate-geant-steps, bump-dataset-version,
create-root-files) plus the unregistered hparam_scan.py with one `dwarf`
command exposing convert/migrate/bump-gen/bump-schema/status/
update-manifest/create-manifest/make-root/hparam-scan as subcommands.

Each scripts/*.py module now only holds argparse-free business logic;
scripts/dwarf.py wires it up with Typer, matching giant/cli.py's style.
`dwarf convert` merges the old serial/parallel conversion scripts behind
a --jobs flag (default 1: sequential with plain -o; >1: dataset-layout
fan-out via subprocess).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-02 09:52:37 +02:00
parent 7b37b284f8
commit d5853d5a75
13 changed files with 662 additions and 505 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/steps_to_parquet_parallel.py
+46 -87
View File
@@ -1,34 +1,27 @@
#!/usr/bin/env python3
"""Convert many ROOT files to Parquet by fanning out to steps_to_parquet.py.
"""Convert many ROOT files to Parquet by fanning out to `dwarf convert`.
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
A single `dwarf convert` call converts a list of files one at a time; this
module runs up to --jobs conversions concurrently, each as its own `dwarf
convert` subprocess (invoked via `python -m scripts.dwarf`, 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.
--schema to pick a specific one, e.g. one just created by `dwarf 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
See `uv run dwarf convert --help` for the CLI.
"""
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+)$")
@@ -90,17 +83,19 @@ def resolve_destination(
return processed_gen_dir / schema_tag / detector / f"{shard_stem}.parquet"
_DWARF_CONVERT_CMD = [sys.executable, "-m", "scripts.dwarf", "convert"]
def _convert_one(
root_file: str,
batch_size: str,
tree: str,
compression: str,
steps_to_parquet_path: Path,
output_path: Path | None,
cmd_prefix: list[str],
) -> tuple[str, int, str, str]:
cmd = [
sys.executable,
str(steps_to_parquet_path),
*cmd_prefix,
root_file,
"--batch-size",
batch_size,
@@ -122,20 +117,25 @@ def run_parallel(
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,
cmd_prefix: list[str] | None = None,
) -> list[tuple[str, int, str, str]]:
"""Run one steps_to_parquet.py subprocess per file, up to *jobs* at a time.
"""Run one `dwarf convert` 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).
be written to (passed through as `dwarf convert`'s --output); files
missing from the map fall back to `dwarf convert`'s own default (parquet
written next to the input .root).
*cmd_prefix* overrides the subprocess command run per file (defaults to
`python -m scripts.dwarf convert`) — used by tests to substitute a fake
conversion script.
Returns one (root_file, returncode, stdout, stderr) tuple per file, in
completion order (not necessarily input order).
"""
output_for = output_for or {}
cmd_prefix = cmd_prefix if cmd_prefix is not None else _DWARF_CONVERT_CMD
results = []
with ThreadPoolExecutor(max_workers=jobs) as pool:
futures = {
@@ -145,8 +145,8 @@ def run_parallel(
batch_size,
tree,
compression,
steps_to_parquet_path,
output_for.get(f),
cmd_prefix,
): f
for f in root_files
}
@@ -162,78 +162,41 @@ def run_parallel(
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)
def run_parallel_job(
root_files: list[str],
jobs: int,
dataset_root: Path,
schema: str | None,
batch_size: str,
tree: str,
compression: str,
) -> None:
"""Resolve each file's dataset-layout destination, convert in parallel, and
report results. Exits the process (via SystemExit) on destination or
conversion failure — this is the top-level entry point `dwarf convert`
delegates to when --jobs > 1."""
output_for: dict[str, Path] = {}
errors: list[str] = []
for f in args.root_files:
for f in root_files:
try:
output_for[f] = resolve_destination(Path(f), dataset_root, args.schema)
output_for[f] = resolve_destination(Path(f), dataset_root, 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)
raise SystemExit(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,
root_files,
jobs=jobs,
batch_size=batch_size,
tree=tree,
compression=compression,
output_for=output_for,
)
@@ -242,10 +205,6 @@ def main() -> None:
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)
raise SystemExit(1)
print(f"\nAll {len(results)} conversion(s) completed.")
if __name__ == "__main__":
main()