From d5853d5a75454d5eb8669ff458b52270e1f52fb2 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 2 Jul 2026 09:52:37 +0200 Subject: [PATCH] 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 --- CLAUDE.md | 3 + README.md | 19 +- pyproject.toml | 6 +- scripts/bump_dataset_version.py | 302 +++++++++------------ scripts/create_root_files.py | 111 +++----- scripts/dwarf.py | 334 ++++++++++++++++++++++++ scripts/hparam_scan.py | 39 +-- scripts/migrate_geant_steps.py | 53 +--- scripts/steps_to_parquet.py | 55 +--- scripts/steps_to_parquet_parallel.py | 133 ++++------ tests/test_create_root_files.py | 21 -- tests/test_dwarf.py | 68 +++++ tests/test_steps_to_parquet_parallel.py | 23 +- 13 files changed, 662 insertions(+), 505 deletions(-) create mode 100644 scripts/dwarf.py create mode 100644 tests/test_dwarf.py diff --git a/CLAUDE.md b/CLAUDE.md index 71e91fa..4214e06 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -11,6 +11,9 @@ uv sync --extra cpu --extra dev # add dev extras (pytest, etc.) pytest # run tests giant train path/to/steps.parquet --mode flow # train (flow matching) giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline) +dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen, + # bump-schema, status, update-manifest, create-manifest, + # make-root, hparam-scan (see scripts/dwarf.py) ``` `cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support. diff --git a/README.md b/README.md index f30b903..b5c2124 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ The model is developed in two phases: ## Data -Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `uv run steps-to-parquet`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle) to avoid leaking correlated steps from the same shower. +Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `uv run dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle) to avoid leaking correlated steps from the same shower. ## Project structure @@ -55,13 +55,16 @@ giant/ │ ├── validate.py # step-level marginal + KL-divergence validation │ ├── analysis.py # notebook diagnostics: marginals, correlations, constraint checks │ └── cli.py # `giant train` / `giant predict` Typer app -├── scripts/ # also exposed as uv entry points, e.g. `uv run steps-to-parquet` -│ ├── steps_to_parquet.py # ROOT → parquet conversion (uproot/awkward/polars) -│ ├── steps_to_parquet_parallel.py # fan out steps_to_parquet.py over several ROOT files -│ ├── migrate_geant_steps.py # one-time move into the raw/processed/pools/derived layout -│ ├── bump_dataset_version.py # cut a new raw gen or parquet schema, with a logged reason -│ ├── create_root_files.py # generate new ROOT shards via a minicalosim executable -│ └── hparam_scan.py # hyperparameter grid scan over `giant train` runs +├── scripts/ # dataset/tooling logic, unified under the `dwarf` CLI (`uv run dwarf --help`) +│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status, +│ │ # update-manifest, create-manifest, make-root, hparam-scan +│ ├── steps_to_parquet.py # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert` +│ ├── steps_to_parquet_parallel.py # fan out conversion over several ROOT files — `dwarf convert --jobs N` +│ ├── migrate_geant_steps.py # one-time move into the raw/processed/pools/derived layout — `dwarf migrate` +│ ├── bump_dataset_version.py # cut a new raw gen or parquet schema, with a logged reason — +│ │ # `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest` +│ ├── create_root_files.py # generate new ROOT shards via a minicalosim executable — `dwarf make-root` +│ └── hparam_scan.py # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan` └── tests/ ``` diff --git a/pyproject.toml b/pyproject.toml index 4e93ee0..997b936 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,11 +38,7 @@ analysis = [ [project.scripts] giant = "giant.cli:app" -steps-to-parquet = "scripts.steps_to_parquet:main" -steps-to-parquet-parallel = "scripts.steps_to_parquet_parallel:main" -migrate-geant-steps = "scripts.migrate_geant_steps:main" -bump-dataset-version = "scripts.bump_dataset_version:main" -create-root-files = "scripts.create_root_files:main" +dwarf = "scripts.dwarf:app" [build-system] requires = ["hatchling"] diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py index 5bc5b2a..837d72d 100644 --- a/scripts/bump_dataset_version.py +++ b/scripts/bump_dataset_version.py @@ -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//.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 /pools//.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.") diff --git a/scripts/create_root_files.py b/scripts/create_root_files.py index 8ca23ad..257b4eb 100644 --- a/scripts/create_root_files.py +++ b/scripts/create_root_files.py @@ -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// — 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//, 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() diff --git a/scripts/dwarf.py b/scripts/dwarf.py new file mode 100644 index 0000000..aaa3eb0 --- /dev/null +++ b/scripts/dwarf.py @@ -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: .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///)", + ), + ] = 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 " + "/pools//.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//, 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() diff --git a/scripts/hparam_scan.py b/scripts/hparam_scan.py index ac0cb1f..dfb5ca2 100644 --- a/scripts/hparam_scan.py +++ b/scripts/hparam_scan.py @@ -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() diff --git a/scripts/migrate_geant_steps.py b/scripts/migrate_geant_steps.py index c3c0181..27e96c9 100644 --- a/scripts/migrate_geant_steps.py +++ b/scripts/migrate_geant_steps.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 """One-time migration of /ceph/lbogner/geant_steps into the versioned layout: raw////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() diff --git a/scripts/steps_to_parquet.py b/scripts/steps_to_parquet.py index 3b58645..e245a8b 100644 --- a/scripts/steps_to_parquet.py +++ b/scripts/steps_to_parquet.py @@ -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: .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() diff --git a/scripts/steps_to_parquet_parallel.py b/scripts/steps_to_parquet_parallel.py index 7344b04..d18db42 100644 --- a/scripts/steps_to_parquet_parallel.py +++ b/scripts/steps_to_parquet_parallel.py @@ -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 /raw////.root (see scripts/migrate_geant_steps.py) — each is written to the matching processed/////.parquet, where defaults to the highest schemaN already under processed/// (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///)", - ) - 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() diff --git a/tests/test_create_root_files.py b/tests/test_create_root_files.py index ccd6047..5c86908 100644 --- a/tests/test_create_root_files.py +++ b/tests/test_create_root_files.py @@ -244,24 +244,3 @@ def test_run_all_caps_concurrency(tmp_path): concurrent += delta peak = max(peak, concurrent) assert peak <= 2 - - -def test_build_parser_defaults(): - args = create_root_files.build_parser().parse_args( - [ - "--executable", - "fake", - "--detector", - "pbwo4", - "--num-files", - "2", - "--events-per-file", - "10000", - "--gen", - "gen1", - ] - ) - assert args.jobs == 4 - assert args.kind == "steps" - assert args.dataset_root == "/ceph/lbogner/geant_steps" - assert args.execute is False diff --git a/tests/test_dwarf.py b/tests/test_dwarf.py new file mode 100644 index 0000000..5ad4367 --- /dev/null +++ b/tests/test_dwarf.py @@ -0,0 +1,68 @@ +from typer.testing import CliRunner + +from scripts.dwarf import app + +runner = CliRunner() + + +def test_convert_rejects_jobs_below_one(tmp_path): + root_file = tmp_path / "shard.root" + root_file.touch() + result = runner.invoke(app, ["convert", str(root_file), "--jobs", "0"]) + assert result.exit_code != 0 + assert "--jobs must be >= 1" in result.output + + +def test_convert_rejects_output_with_multiple_files(tmp_path): + a = tmp_path / "a.root" + b = tmp_path / "b.root" + a.touch() + b.touch() + result = runner.invoke(app, ["convert", str(a), str(b), "--output", "out.parquet"]) + assert result.exit_code != 0 + assert "--output can only be used with a single input file" in result.output + + +def test_convert_rejects_output_with_parallel_jobs(tmp_path): + root_file = tmp_path / "shard.root" + root_file.touch() + result = runner.invoke( + app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"] + ) + assert result.exit_code != 0 + assert "--output cannot be combined with --jobs > 1" in result.output + + +def test_convert_default_jobs_is_one(): + result = runner.invoke(app, ["convert", "--help"]) + assert result.exit_code == 0 + assert "default: 1" in result.output + + +def test_bump_gen_requires_reason(): + result = runner.invoke(app, ["bump-gen"]) + assert result.exit_code != 0 + assert "reason" in result.output.lower() + + +def test_create_manifest_requires_exactly_one_of_output_or_pool(tmp_path): + f = tmp_path / "a.parquet" + f.touch() + result = runner.invoke(app, ["create-manifest", str(f)]) + assert result.exit_code != 0 + assert "exactly one of --output or --pool is required" in result.output + + +def test_create_manifest_requires_type_with_pool(tmp_path): + f = tmp_path / "a.parquet" + f.touch() + result = runner.invoke(app, ["create-manifest", "--pool", "pbwo4", str(f)]) + assert result.exit_code != 0 + assert "--type is required when --pool is given" in result.output + + +def test_status_reports_missing_root(tmp_path): + missing = tmp_path / "does-not-exist" + result = runner.invoke(app, ["status", "--root", str(missing)]) + assert result.exit_code != 0 + assert "is not a directory" in result.output diff --git a/tests/test_steps_to_parquet_parallel.py b/tests/test_steps_to_parquet_parallel.py index b5b2bb7..101d36f 100644 --- a/tests/test_steps_to_parquet_parallel.py +++ b/tests/test_steps_to_parquet_parallel.py @@ -1,4 +1,5 @@ import json +import sys from pathlib import Path from scripts import steps_to_parquet_parallel @@ -40,7 +41,7 @@ def test_runs_one_job_per_file_and_reports_success(tmp_path): fake = _write_fake_executable(tmp_path, marker_dir) files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(3)] - results = run_parallel(files, jobs=4, steps_to_parquet_path=fake) + results = run_parallel(files, jobs=4, cmd_prefix=[sys.executable, str(fake)]) assert {r[0] for r in results} == set(files) assert all(code == 0 for _, code, _, _ in results) @@ -53,7 +54,7 @@ def test_failures_are_reported_with_nonzero_exit_code(tmp_path): fake = _write_fake_executable(tmp_path, marker_dir) files = [str(tmp_path / "shard-000.root"), str(tmp_path / "shard-fail.root")] - results = run_parallel(files, jobs=4, steps_to_parquet_path=fake) + results = run_parallel(files, jobs=4, cmd_prefix=[sys.executable, str(fake)]) codes = {Path(f).stem: code for f, code, _, _ in results} assert codes["shard-000"] == 0 @@ -66,7 +67,7 @@ def test_jobs_caps_concurrency(tmp_path): fake = _write_fake_executable(tmp_path, marker_dir) files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(6)] - run_parallel(files, jobs=2, steps_to_parquet_path=fake) + run_parallel(files, jobs=2, cmd_prefix=[sys.executable, str(fake)]) intervals = [] for f in files: @@ -83,11 +84,6 @@ def test_jobs_caps_concurrency(tmp_path): assert peak <= 2 -def test_default_jobs_is_four(): - args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"]) - assert args.jobs == 4 - - def test_output_for_is_passed_through_as_output_flag(tmp_path): marker_dir = tmp_path / "markers" marker_dir.mkdir() @@ -108,7 +104,10 @@ sys.exit(0) root_file = str(tmp_path / "shard-000.root") dest = tmp_path / "processed" / "shard-000.parquet" run_parallel( - [root_file], jobs=1, steps_to_parquet_path=fake, output_for={root_file: dest} + [root_file], + jobs=1, + cmd_prefix=[sys.executable, str(fake)], + output_for={root_file: dest}, ) assert (marker_dir / "shard-000.txt").read_text() == str(dest) @@ -171,9 +170,3 @@ def test_resolve_destination_errors_on_wrong_shape(tmp_path): def test_latest_schema_tag_returns_none_when_missing(tmp_path): assert latest_schema_tag(tmp_path / "does" / "not" / "exist") is None - - -def test_dataset_root_and_schema_flags_default(tmp_path): - args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"]) - assert args.dataset_root == "/ceph/lbogner/geant_steps" - assert args.schema is None