Unify dataset/tooling scripts into a single dwarf Typer CLI
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:
+128
-174
@@ -1,4 +1,3 @@
|
||||
#!/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):
|
||||
|
||||
@@ -7,21 +6,16 @@ dataset tree (see scripts/migrate_geant_steps.py for the layout):
|
||||
pools/<detector>/<pool>.manifest
|
||||
|
||||
`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.
|
||||
`schema` bumps when the parquet export (`dwarf convert` 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.
|
||||
to VERSIONS.md. Defaults to a dry run; pass execute=True 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 update-manifest pools/pbwo4/full.manifest [--schema schema2]
|
||||
bump_dataset_version.py create-manifest --output pools/pbwo4/train.manifest a.parquet b.parquet
|
||||
bump_dataset_version.py status
|
||||
See `uv run dwarf bump-gen/bump-schema/update-manifest/create-manifest/status
|
||||
--help` for the CLI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import datetime as dt
|
||||
import os
|
||||
import re
|
||||
@@ -287,183 +281,143 @@ def apply_create_manifest(output_path: Path, lines: list[str]) -> None:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# CLI entry points (called from scripts/dwarf.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--root",
|
||||
default="/ceph/lbogner/geant_steps",
|
||||
help="Dataset root for bump-gen/bump-schema/status (default: /ceph/lbogner/geant_steps)",
|
||||
)
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
def run_status(root: str) -> None:
|
||||
root_path = Path(root)
|
||||
if not root_path.is_dir():
|
||||
raise SystemExit(f"error: {root_path} is not a directory")
|
||||
print_status(root_path)
|
||||
|
||||
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)")
|
||||
def _run_bump(
|
||||
kind: str,
|
||||
reason: str,
|
||||
by: str | None,
|
||||
date: str | None,
|
||||
execute: bool,
|
||||
root: str,
|
||||
gen: str | None,
|
||||
) -> None:
|
||||
root_path = Path(root)
|
||||
if not root_path.is_dir():
|
||||
raise SystemExit(f"error: {root_path} is not a directory")
|
||||
|
||||
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)")
|
||||
date = date or dt.date.today().isoformat()
|
||||
by = by if by is not None else _git_user_name()
|
||||
if gen is None:
|
||||
new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date)
|
||||
else:
|
||||
new_dirs, log_line = plan_bump_schema(root_path, kind, gen, reason, by, date)
|
||||
|
||||
p_update = sub.add_parser(
|
||||
"update-manifest",
|
||||
help="Repoint manifest(s) to a new schema, verifying all target files exist",
|
||||
)
|
||||
p_update.add_argument(
|
||||
"manifests", nargs="+", metavar="MANIFEST", help="One or more .manifest files to update"
|
||||
)
|
||||
p_update.add_argument(
|
||||
"--schema",
|
||||
default=None,
|
||||
metavar="schemaN",
|
||||
help="Target schema tag (default: highest schema found in the same gen dir)",
|
||||
)
|
||||
p_update.add_argument("--execute", action="store_true", help="Write updated manifests (default: dry run)")
|
||||
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
|
||||
print("new directories:")
|
||||
for d in new_dirs:
|
||||
print(f" {d}")
|
||||
print("VERSIONS.md entry:")
|
||||
print(f" {log_line}")
|
||||
|
||||
p_create = sub.add_parser(
|
||||
"create-manifest",
|
||||
help="Create a new manifest from a list of parquet files",
|
||||
)
|
||||
dest_group = p_create.add_mutually_exclusive_group(required=True)
|
||||
dest_group.add_argument(
|
||||
"--output", "-o", metavar="PATH", help="Explicit path for the new .manifest file"
|
||||
)
|
||||
dest_group.add_argument(
|
||||
"--pool", metavar="DETECTOR",
|
||||
help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
|
||||
)
|
||||
p_create.add_argument(
|
||||
"--type", choices=["full", "holdout", "dev"],
|
||||
help="Pool type — full, holdout, or dev (required with --pool)",
|
||||
)
|
||||
p_create.add_argument(
|
||||
"files", nargs="+", metavar="FILE", help="Parquet files to include"
|
||||
)
|
||||
p_create.add_argument("--execute", action="store_true", help="Write the manifest (default: dry run)")
|
||||
if not execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
apply_bump(root_path, new_dirs, log_line)
|
||||
print("\nDone.")
|
||||
|
||||
sub.add_parser("status", help="List existing gens/schemas per kind")
|
||||
|
||||
args = parser.parse_args()
|
||||
def run_bump_gen(
|
||||
kind: str, reason: str, by: str | None, date: str | None, execute: bool, root: str
|
||||
) -> None:
|
||||
_run_bump(kind, reason, by, date, execute, root, gen=None)
|
||||
|
||||
# Commands that need --root
|
||||
if args.command in ("bump-gen", "bump-schema", "status"):
|
||||
root = Path(args.root)
|
||||
if not root.is_dir():
|
||||
parser.error(f"{root} is not a directory")
|
||||
|
||||
if args.command == "status":
|
||||
print_status(root)
|
||||
def run_bump_schema(
|
||||
kind: str, gen: str, reason: str, by: str | None, date: str | None, execute: bool, root: str
|
||||
) -> None:
|
||||
_run_bump(kind, reason, by, date, execute, root, gen=gen)
|
||||
|
||||
|
||||
def run_update_manifest(manifests: list[str], schema: str | None, execute: bool) -> None:
|
||||
if schema and not SCHEMA_RE.match(schema):
|
||||
raise SystemExit(f"error: --schema must look like 'schemaN', got {schema!r}")
|
||||
|
||||
all_plans: list[tuple[Path, list[tuple[str, str | None]]]] = []
|
||||
all_missing: list[Path] = []
|
||||
|
||||
for raw in manifests:
|
||||
mp = Path(raw)
|
||||
plan, missing = plan_update_manifest(mp, schema)
|
||||
all_plans.append((mp.resolve(), plan))
|
||||
all_missing.extend(missing)
|
||||
|
||||
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
|
||||
for mp, plan in all_plans:
|
||||
changes = [(old, new) for old, new in plan if new is not None]
|
||||
print(f"\n{mp} ({len(changes)} path(s) to update)")
|
||||
for old, new in changes:
|
||||
print(f" - {old.strip()}")
|
||||
print(f" + {new}")
|
||||
|
||||
if all_missing:
|
||||
print(f"\nMISSING ({len(all_missing)} file(s) — target paths do not exist):")
|
||||
for p in all_missing:
|
||||
print(f" {p}")
|
||||
if execute:
|
||||
raise SystemExit("error: refusing to write manifests with missing targets")
|
||||
|
||||
if not execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
if args.command in ("bump-gen", "bump-schema"):
|
||||
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)
|
||||
for mp, plan in all_plans:
|
||||
apply_update_manifest(mp, plan)
|
||||
print("\nDone.")
|
||||
|
||||
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.")
|
||||
def run_create_manifest(
|
||||
files: list[str],
|
||||
execute: bool,
|
||||
output: str | None = None,
|
||||
pool: str | None = None,
|
||||
type_: str | None = None,
|
||||
root: str = "/ceph/lbogner/geant_steps",
|
||||
) -> None:
|
||||
if (output is None) == (pool is None):
|
||||
raise SystemExit("error: exactly one of --output or --pool is required")
|
||||
if pool is not None and type_ is None:
|
||||
raise SystemExit("error: --type is required when --pool is given")
|
||||
|
||||
if pool is not None:
|
||||
output_path = Path(root) / "pools" / pool / f"{type_}.manifest"
|
||||
else:
|
||||
assert output is not None # guaranteed by the exclusivity check above
|
||||
output_path = Path(output)
|
||||
|
||||
parquet_files = [Path(f) for f in files]
|
||||
lines, missing, resolved = plan_create_manifest(output_path, parquet_files)
|
||||
overlaps = check_holdout_overlap(output_path, resolved)
|
||||
|
||||
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
|
||||
print(f"manifest: {output_path.resolve()}")
|
||||
for line in lines:
|
||||
print(f" {line}")
|
||||
|
||||
if missing:
|
||||
print(f"\nMISSING ({len(missing)} file(s) do not exist):")
|
||||
for p in missing:
|
||||
print(f" {p}")
|
||||
|
||||
if overlaps:
|
||||
print(f"\nHOLDOUT OVERLAP ({len(overlaps)} file(s) appear in other manifests):")
|
||||
for name, f in overlaps:
|
||||
print(f" {f} (also in {name})")
|
||||
|
||||
if (missing or overlaps) and execute:
|
||||
raise SystemExit("error: refusing to write manifest (see above)")
|
||||
|
||||
if not execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
if args.command == "update-manifest":
|
||||
if args.schema and not SCHEMA_RE.match(args.schema):
|
||||
parser.error(f"--schema must look like 'schemaN', got {args.schema!r}")
|
||||
|
||||
all_plans: list[tuple[Path, list[tuple[str, str | None]]]] = []
|
||||
all_missing: list[Path] = []
|
||||
|
||||
for raw in args.manifests:
|
||||
mp = Path(raw)
|
||||
plan, missing = plan_update_manifest(mp, args.schema)
|
||||
all_plans.append((mp.resolve(), plan))
|
||||
all_missing.extend(missing)
|
||||
|
||||
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
|
||||
for mp, plan in all_plans:
|
||||
changes = [(old, new) for old, new in plan if new is not None]
|
||||
print(f"\n{mp} ({len(changes)} path(s) to update)")
|
||||
for old, new in changes:
|
||||
print(f" - {old.strip()}")
|
||||
print(f" + {new}")
|
||||
|
||||
if all_missing:
|
||||
print(f"\nMISSING ({len(all_missing)} file(s) — target paths do not exist):")
|
||||
for p in all_missing:
|
||||
print(f" {p}")
|
||||
if args.execute:
|
||||
raise SystemExit("error: refusing to write manifests with missing targets")
|
||||
|
||||
if not args.execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
for mp, plan in all_plans:
|
||||
apply_update_manifest(mp, plan)
|
||||
print("\nDone.")
|
||||
return
|
||||
|
||||
if args.command == "create-manifest":
|
||||
if args.pool is not None and args.type is None:
|
||||
parser.error("--type is required when --pool is given")
|
||||
|
||||
if args.pool is not None:
|
||||
root = Path(args.root)
|
||||
output_path = root / "pools" / args.pool / f"{args.type}.manifest"
|
||||
else:
|
||||
output_path = Path(args.output)
|
||||
|
||||
parquet_files = [Path(f) for f in args.files]
|
||||
lines, missing, resolved = plan_create_manifest(output_path, parquet_files)
|
||||
overlaps = check_holdout_overlap(output_path, resolved)
|
||||
|
||||
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
|
||||
print(f"manifest: {output_path.resolve()}")
|
||||
for line in lines:
|
||||
print(f" {line}")
|
||||
|
||||
if missing:
|
||||
print(f"\nMISSING ({len(missing)} file(s) do not exist):")
|
||||
for p in missing:
|
||||
print(f" {p}")
|
||||
|
||||
if overlaps:
|
||||
print(f"\nHOLDOUT OVERLAP ({len(overlaps)} file(s) appear in other manifests):")
|
||||
for name, f in overlaps:
|
||||
print(f" {f} (also in {name})")
|
||||
|
||||
if (missing or overlaps) and args.execute:
|
||||
raise SystemExit("error: refusing to write manifest (see above)")
|
||||
|
||||
if not args.execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
apply_create_manifest(output_path, lines)
|
||||
print("\nDone.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
apply_create_manifest(output_path, lines)
|
||||
print("\nDone.")
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
#!/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:
|
||||
|
||||
@@ -15,19 +14,11 @@ 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.
|
||||
`dwarf 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
|
||||
See `uv run dwarf make-root --help` for the CLI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
@@ -192,78 +183,48 @@ def run_all(
|
||||
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 run_make_root(
|
||||
executable: Path,
|
||||
detector: list[str],
|
||||
num_files: int,
|
||||
events_per_file: int,
|
||||
kind: str,
|
||||
gen: str,
|
||||
dataset_root: str,
|
||||
jobs: int,
|
||||
execute: bool,
|
||||
) -> None:
|
||||
if jobs < 1:
|
||||
raise SystemExit("error: --jobs must be >= 1")
|
||||
if num_files < 1:
|
||||
raise SystemExit("error: --num-files must be >= 1")
|
||||
if events_per_file < 1:
|
||||
raise SystemExit("error: --events-per-file must be >= 1")
|
||||
if not executable.is_file() or not os.access(executable, os.X_OK):
|
||||
raise SystemExit(f"error: {executable} is not an executable file")
|
||||
|
||||
|
||||
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)
|
||||
dataset_root_path = Path(dataset_root)
|
||||
try:
|
||||
jobs = plan_jobs(args.detector, args.num_files, dataset_root, args.kind, args.gen)
|
||||
planned_jobs = plan_jobs(detector, num_files, dataset_root_path, kind, gen)
|
||||
except PlanError as exc:
|
||||
parser.error(str(exc))
|
||||
raise SystemExit(f"error: {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"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
|
||||
print(f"executable: {executable}")
|
||||
for job in planned_jobs:
|
||||
cmd = [str(executable)] + ([job.config] if job.config else []) + [str(events_per_file)]
|
||||
dest = dataset_root_path / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
|
||||
print(f" {' '.join(cmd)} -> {dest}")
|
||||
|
||||
if not args.execute:
|
||||
if not execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
tmp_root = dataset_root / ".sim-tmp"
|
||||
tmp_root = dataset_root_path / ".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,
|
||||
planned_jobs, executable, events_per_file, dataset_root_path, kind, gen,
|
||||
max_workers=jobs, tmp_root=tmp_root,
|
||||
)
|
||||
if tmp_root.is_dir() and not any(tmp_root.iterdir()):
|
||||
tmp_root.rmdir()
|
||||
@@ -273,10 +234,6 @@ def main() -> None:
|
||||
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)
|
||||
raise SystemExit(1)
|
||||
|
||||
print(f"\nAll {len(results)} job(s) completed.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -0,0 +1,334 @@
|
||||
"""dwarf — little helper to `giant`: dataset/tooling CLI for the geant_steps pipeline.
|
||||
|
||||
Unifies the standalone scripts/*.py conversion, migration, versioning, and
|
||||
simulation-fanout tools into one Typer app so there's a single command name
|
||||
(and `--help`) to remember instead of five differently-hyphenated ones.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
from typing_extensions import Annotated
|
||||
|
||||
from scripts.bump_dataset_version import (
|
||||
run_bump_gen,
|
||||
run_bump_schema,
|
||||
run_create_manifest,
|
||||
run_status,
|
||||
run_update_manifest,
|
||||
)
|
||||
from scripts.create_root_files import run_make_root
|
||||
from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
|
||||
from scripts.migrate_geant_steps import run_migration
|
||||
from scripts.steps_to_parquet import convert_steps_to_parquet
|
||||
from scripts.steps_to_parquet_parallel import run_parallel_job
|
||||
|
||||
app = typer.Typer(no_args_is_help=True)
|
||||
|
||||
_DATASET_ROOT_DEFAULT = Path("/ceph/lbogner/geant_steps")
|
||||
|
||||
|
||||
@app.callback()
|
||||
def _main() -> None:
|
||||
"""dwarf — dataset/tooling CLI (ROOT<->parquet conversion, dataset versioning, sim fanout)."""
|
||||
|
||||
|
||||
class Compression(str, Enum):
|
||||
snappy = "snappy"
|
||||
lz4 = "lz4"
|
||||
zstd = "zstd"
|
||||
gzip = "gzip"
|
||||
none = "none"
|
||||
|
||||
|
||||
class PoolType(str, Enum):
|
||||
full = "full"
|
||||
holdout = "holdout"
|
||||
dev = "dev"
|
||||
|
||||
|
||||
@app.command()
|
||||
def convert(
|
||||
root_files: Annotated[
|
||||
list[Path], typer.Argument(help="Input ROOT file(s)")
|
||||
],
|
||||
output: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option(
|
||||
"--output",
|
||||
"-o",
|
||||
help="Output Parquet file (default: <input>.parquet). Only valid "
|
||||
"with a single input file and --jobs 1.",
|
||||
),
|
||||
] = None,
|
||||
batch_size: Annotated[
|
||||
str,
|
||||
typer.Option(
|
||||
"--batch-size",
|
||||
help="Uproot read batch size, e.g. '100 MB' or '500000' (rows)",
|
||||
),
|
||||
] = "100 MB",
|
||||
tree: Annotated[
|
||||
str, typer.Option("--tree", help="Tree name inside the ROOT file")
|
||||
] = "Steps",
|
||||
compression: Annotated[
|
||||
Compression, typer.Option("--compression", help="Parquet compression codec")
|
||||
] = Compression.snappy,
|
||||
jobs: Annotated[
|
||||
int,
|
||||
typer.Option(
|
||||
"--jobs",
|
||||
"-j",
|
||||
help="Convert N files in parallel, resolving each destination from "
|
||||
"--dataset-root/--schema (default: 1, sequential, any file layout)",
|
||||
),
|
||||
] = 1,
|
||||
dataset_root: Annotated[
|
||||
Path,
|
||||
typer.Option(
|
||||
"--dataset-root",
|
||||
help="Dataset root containing raw/ and processed/ (only used with --jobs > 1)",
|
||||
),
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
schema: Annotated[
|
||||
Optional[str],
|
||||
typer.Option(
|
||||
"--schema",
|
||||
help="Schema tag to write parquets under, e.g. schema2 (only used "
|
||||
"with --jobs > 1; default: highest schemaN already under "
|
||||
"processed/<kind>/<gen>/)",
|
||||
),
|
||||
] = None,
|
||||
) -> None:
|
||||
"""Convert ROOT Steps tree(s) to Parquet."""
|
||||
if jobs < 1:
|
||||
typer.echo("error: --jobs must be >= 1", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
compression_value = "uncompressed" if compression is Compression.none else compression.value
|
||||
|
||||
if jobs == 1:
|
||||
if output is not None and len(root_files) > 1:
|
||||
typer.echo("error: --output can only be used with a single input file", err=True)
|
||||
raise typer.Exit(1)
|
||||
for root_file in root_files:
|
||||
convert_steps_to_parquet(
|
||||
root_file,
|
||||
output_path=output,
|
||||
batch_size=batch_size,
|
||||
tree_name=tree,
|
||||
compression=compression_value,
|
||||
)
|
||||
return
|
||||
|
||||
if output is not None:
|
||||
typer.echo(
|
||||
"error: --output cannot be combined with --jobs > 1 "
|
||||
"(destinations are derived from --dataset-root/--schema)",
|
||||
err=True,
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
run_parallel_job(
|
||||
[str(f) for f in root_files],
|
||||
jobs=jobs,
|
||||
dataset_root=dataset_root,
|
||||
schema=schema,
|
||||
batch_size=batch_size,
|
||||
tree=tree,
|
||||
compression=compression_value,
|
||||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def migrate(
|
||||
root: Annotated[
|
||||
Path, typer.Argument(help="Dataset root to migrate in place")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
execute: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
"--execute", help="Actually move/copy files and write manifests (default: dry run)"
|
||||
),
|
||||
] = False,
|
||||
copy: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
"--copy",
|
||||
help="Copy instead of move, leaving the originals in place "
|
||||
"(e.g. if another process is still reading them)",
|
||||
),
|
||||
] = False,
|
||||
) -> None:
|
||||
"""One-time migration into the versioned raw/processed/pools/derived layout."""
|
||||
run_migration(str(root), execute=execute, copy=copy)
|
||||
|
||||
|
||||
@app.command("bump-gen")
|
||||
def bump_gen(
|
||||
reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
|
||||
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
|
||||
by: Annotated[
|
||||
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
|
||||
] = None,
|
||||
date: Annotated[
|
||||
Optional[str], typer.Option("--date", help="Override date (default: today, ISO)")
|
||||
] = None,
|
||||
execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
|
||||
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
|
||||
) -> None:
|
||||
"""Cut a new raw generation."""
|
||||
run_bump_gen(kind=kind, reason=reason, by=by, date=date, execute=execute, root=str(root))
|
||||
|
||||
|
||||
@app.command("bump-schema")
|
||||
def bump_schema(
|
||||
gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
|
||||
reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
|
||||
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
|
||||
by: Annotated[
|
||||
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
|
||||
] = None,
|
||||
date: Annotated[
|
||||
Optional[str], typer.Option("--date", help="Override date (default: today, ISO)")
|
||||
] = None,
|
||||
execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
|
||||
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
|
||||
) -> None:
|
||||
"""Cut a new schema within a gen."""
|
||||
run_bump_schema(
|
||||
kind=kind, gen=gen, reason=reason, by=by, date=date, execute=execute, root=str(root)
|
||||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def status(
|
||||
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
|
||||
) -> None:
|
||||
"""List existing gens/schemas per kind."""
|
||||
run_status(str(root))
|
||||
|
||||
|
||||
@app.command("update-manifest")
|
||||
def update_manifest(
|
||||
manifests: Annotated[
|
||||
list[Path], typer.Argument(help="One or more .manifest files to update")
|
||||
],
|
||||
schema: Annotated[
|
||||
Optional[str],
|
||||
typer.Option(
|
||||
"--schema",
|
||||
metavar="schemaN",
|
||||
help="Target schema tag (default: highest schema found in the same gen dir)",
|
||||
),
|
||||
] = None,
|
||||
execute: Annotated[
|
||||
bool, typer.Option("--execute", help="Write updated manifests (default: dry run)")
|
||||
] = False,
|
||||
) -> None:
|
||||
"""Repoint manifest(s) to a new schema, verifying all target files exist."""
|
||||
run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute)
|
||||
|
||||
|
||||
@app.command("create-manifest")
|
||||
def create_manifest(
|
||||
files: Annotated[list[Path], typer.Argument(help="Parquet files to include")],
|
||||
output: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option("--output", "-o", help="Explicit path for the new .manifest file"),
|
||||
] = None,
|
||||
pool: Annotated[
|
||||
Optional[str],
|
||||
typer.Option(
|
||||
"--pool",
|
||||
metavar="DETECTOR",
|
||||
help="Detector name; combined with --type and --root to form "
|
||||
"<root>/pools/<detector>/<type>.manifest",
|
||||
),
|
||||
] = None,
|
||||
type_: Annotated[
|
||||
Optional[PoolType],
|
||||
typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
|
||||
] = None,
|
||||
root: Annotated[
|
||||
Path, typer.Option("--root", help="Dataset root (used with --pool)")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
execute: Annotated[
|
||||
bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
|
||||
] = False,
|
||||
) -> None:
|
||||
"""Create a new manifest from a list of parquet files."""
|
||||
run_create_manifest(
|
||||
[str(f) for f in files],
|
||||
execute=execute,
|
||||
output=str(output) if output is not None else None,
|
||||
pool=pool,
|
||||
type_=type_.value if type_ is not None else None,
|
||||
root=str(root),
|
||||
)
|
||||
|
||||
|
||||
@app.command("make-root")
|
||||
def make_root(
|
||||
executable: Annotated[
|
||||
Path, typer.Option("--executable", help="Built minicalosim run_* executable")
|
||||
],
|
||||
detector: Annotated[
|
||||
list[str],
|
||||
typer.Option(
|
||||
"--detector",
|
||||
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.",
|
||||
),
|
||||
],
|
||||
num_files: Annotated[
|
||||
int, typer.Option("--num-files", help="New shards to create per detector")
|
||||
],
|
||||
events_per_file: Annotated[
|
||||
int, typer.Option("--events-per-file", help="nEvents passed to the executable")
|
||||
],
|
||||
gen: Annotated[
|
||||
str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
|
||||
],
|
||||
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
|
||||
dataset_root: Annotated[
|
||||
Path, typer.Option("--dataset-root", help="Dataset root")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
jobs: Annotated[
|
||||
int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")
|
||||
] = 4,
|
||||
execute: Annotated[
|
||||
bool, typer.Option("--execute", help="Actually run jobs (default: dry run / print plan)")
|
||||
] = False,
|
||||
) -> None:
|
||||
"""Generate new ROOT shards via a minicalosim executable."""
|
||||
run_make_root(
|
||||
executable=executable,
|
||||
detector=detector,
|
||||
num_files=num_files,
|
||||
events_per_file=events_per_file,
|
||||
kind=kind,
|
||||
gen=gen,
|
||||
dataset_root=str(dataset_root),
|
||||
jobs=jobs,
|
||||
execute=execute,
|
||||
)
|
||||
|
||||
|
||||
@app.command("hparam-scan")
|
||||
def hparam_scan(
|
||||
data: Annotated[str, typer.Option("--data")] = DATA_DEFAULT,
|
||||
scan_dir: Annotated[str, typer.Option("--scan-dir")] = SCAN_DIR_DEFAULT,
|
||||
seed: Annotated[int, typer.Option("--seed")] = 0,
|
||||
dry_run: Annotated[bool, typer.Option("--dry-run")] = False,
|
||||
) -> None:
|
||||
"""Grid-scan dropout x n_blocks x hidden_dim via sequential `giant train` runs."""
|
||||
run_hparam_scan(data=data, scan_dir=scan_dir, seed=seed, dry_run=dry_run)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
+14
-25
@@ -1,4 +1,3 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Hyperparameter scan over dropout x n_blocks x hidden_dim.
|
||||
|
||||
Runs `giant train` sequentially (this machine has a single GPU) for every
|
||||
@@ -6,13 +5,9 @@ combination, plus one extra run at the default architecture with a higher
|
||||
learning rate. Runs are shuffled so the parameter space gets coarse coverage
|
||||
early rather than exhausting one corner of the grid first.
|
||||
|
||||
Usage:
|
||||
uv run python scripts/hparam_scan.py
|
||||
uv run python scripts/hparam_scan.py --dry-run
|
||||
uv run python scripts/hparam_scan.py --seed 1 --data /path/to/parquet
|
||||
See `uv run dwarf hparam-scan --help` for the CLI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import itertools
|
||||
import os
|
||||
@@ -87,31 +82,29 @@ def append_summary(summary_path: Path, row: dict) -> None:
|
||||
writer.writerow(row)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--data", default=DATA_DEFAULT)
|
||||
parser.add_argument("--scan-dir", default=SCAN_DIR_DEFAULT)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--dry-run", action="store_true")
|
||||
args = parser.parse_args()
|
||||
def run_hparam_scan(
|
||||
data: str = DATA_DEFAULT,
|
||||
scan_dir: str = SCAN_DIR_DEFAULT,
|
||||
seed: int = 0,
|
||||
dry_run: bool = False,
|
||||
) -> None:
|
||||
runs = build_runs(seed)
|
||||
scan_dir_path = Path(scan_dir)
|
||||
|
||||
runs = build_runs(args.seed)
|
||||
scan_dir = Path(args.scan_dir)
|
||||
|
||||
if args.dry_run:
|
||||
if dry_run:
|
||||
for i, run in enumerate(runs, 1):
|
||||
print(f"[{i}/{len(runs)}] {run_name(run)}")
|
||||
return
|
||||
|
||||
scan_dir.mkdir(parents=True, exist_ok=True)
|
||||
summary_path = scan_dir / "scan_summary.csv"
|
||||
scan_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
summary_path = scan_dir_path / "scan_summary.csv"
|
||||
|
||||
env = os.environ.copy()
|
||||
env["TQDM_DISABLE"] = "1"
|
||||
|
||||
for i, run in enumerate(runs, 1):
|
||||
name = run_name(run)
|
||||
out_dir = scan_dir / name
|
||||
out_dir = scan_dir_path / name
|
||||
metrics_path = out_dir / "metrics.csv"
|
||||
last_ckpt = out_dir / "last.pt"
|
||||
|
||||
@@ -124,7 +117,7 @@ def main() -> None:
|
||||
cmd = [
|
||||
"giant",
|
||||
"train",
|
||||
args.data,
|
||||
data,
|
||||
"--mode",
|
||||
"flow",
|
||||
"--epochs",
|
||||
@@ -186,7 +179,3 @@ def main() -> None:
|
||||
print(
|
||||
f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
#!/usr/bin/env python3
|
||||
"""One-time migration of /ceph/lbogner/geant_steps into the versioned layout:
|
||||
|
||||
raw/<kind>/<gen>/<detector>/shard-NNN.root
|
||||
@@ -13,12 +12,13 @@ up pool membership from POOL_ASSIGNMENT — decoupling "where it sits today" fro
|
||||
"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. Pass --copy as well to
|
||||
copy instead of move, leaving the original files in place — e.g. if another
|
||||
process is still reading them from their current location.
|
||||
execute=True to actually move files and write manifests, and copy=True to copy
|
||||
instead of move, leaving the original files in place — e.g. if another process
|
||||
is still reading them from their current location.
|
||||
|
||||
See `uv run dwarf migrate --help` for the CLI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
@@ -181,39 +181,18 @@ def plan_manifests(src_root: Path) -> dict[Path, list[str]]:
|
||||
}
|
||||
|
||||
|
||||
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)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--copy",
|
||||
action="store_true",
|
||||
help="Copy instead of move, leaving the originals in place "
|
||||
"(e.g. if another process is still reading them). Implies the legacy "
|
||||
"train/ etc. directories are left as-is too, since they won't be empty.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
src_root = Path(args.root)
|
||||
def run_migration(root: str, execute: bool, copy: bool) -> None:
|
||||
src_root = Path(root)
|
||||
if not src_root.is_dir():
|
||||
parser.error(f"{src_root} is not a directory")
|
||||
raise SystemExit(f"error: {src_root} is not a directory")
|
||||
|
||||
moves, unrecognized = plan_moves(src_root)
|
||||
manifests = plan_manifests(src_root)
|
||||
|
||||
verb = "COPY" if args.copy else "MOVE"
|
||||
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ({verb}): {src_root} ===\n")
|
||||
verb = "COPY" if copy else "MOVE"
|
||||
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ({verb}): {src_root} ===\n")
|
||||
|
||||
print(f"-- {len(moves)} file {'copy' if args.copy else 'move'}(s) --")
|
||||
print(f"-- {len(moves)} file {'copy' if copy else 'move'}(s) --")
|
||||
for src, dst in moves:
|
||||
print(f" {src.relative_to(src_root)} -> {dst.relative_to(src_root)}")
|
||||
|
||||
@@ -228,11 +207,11 @@ def main() -> None:
|
||||
for path in unrecognized:
|
||||
print(f" {path.relative_to(src_root)}")
|
||||
|
||||
if not args.execute:
|
||||
if not execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
transfer = shutil.copy2 if args.copy else shutil.move
|
||||
transfer = shutil.copy2 if copy else shutil.move
|
||||
for src, dst in moves:
|
||||
if dst.exists():
|
||||
raise FileExistsError(f"refusing to overwrite existing file: {dst}")
|
||||
@@ -257,7 +236,7 @@ def main() -> None:
|
||||
# empty since their contents were classified by filename, not location —
|
||||
# remove them, but only if a move actually emptied them. In --copy mode the
|
||||
# originals are still there by design, so leave these alone entirely.
|
||||
if not args.copy:
|
||||
if not copy:
|
||||
for stale_dir in ("train", "sampling_train/small", "sampling_train"):
|
||||
d = src_root / stale_dir
|
||||
if d.is_dir() and not any(d.iterdir()):
|
||||
@@ -265,7 +244,3 @@ def main() -> None:
|
||||
print(f"removed now-empty directory: {d.relative_to(src_root)}")
|
||||
|
||||
print("\nDone.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -1,14 +1,8 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Convert the Steps tree from a ROOT file to Parquet.
|
||||
|
||||
Usage:
|
||||
uv run python steps_to_parquet.py input.root
|
||||
uv run python steps_to_parquet.py input.root -o output.parquet
|
||||
uv run python steps_to_parquet.py input.root --batch-size "200 MB" --tree Hits
|
||||
uv run python steps_to_parquet.py input1.root input2.root input3.root
|
||||
See `uv run dwarf convert --help` for the CLI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
@@ -125,50 +119,3 @@ def convert_steps_to_parquet(
|
||||
df.write_parquet(output_path, compression=compression)
|
||||
print(f"done ({output_path.stat().st_size / 1e6:.1f} MB)")
|
||||
return output_path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert a Steps (or any flat+jagged) tree in a ROOT file to Parquet."
|
||||
)
|
||||
parser.add_argument("root_files", nargs="+", help="Input ROOT file(s)")
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--output",
|
||||
help="Output Parquet file (default: <input>.parquet). "
|
||||
"Only valid with a single input file.",
|
||||
)
|
||||
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)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.output is not None and len(args.root_files) > 1:
|
||||
parser.error("--output can only be used with a single input file")
|
||||
|
||||
compression = "uncompressed" if args.compression == "none" else args.compression
|
||||
for root_file in args.root_files:
|
||||
convert_steps_to_parquet(
|
||||
root_file,
|
||||
output_path=args.output,
|
||||
batch_size=args.batch_size,
|
||||
tree_name=args.tree,
|
||||
compression=compression,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user