From 320365606ae4e7dd12834562b459d738ad5feab7 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 25 Jun 2026 16:46:48 +0200 Subject: [PATCH] Add tooling for a versioned geant_steps dataset layout Introduces raw////shard-NNN.root and processed/////shard-NNN.parquet as the dataset convention, plus scripts to operate on it: migrate_geant_steps.py for the one-time move into this layout, bump_dataset_version.py to cut new gen/schema versions with a logged reason, steps_to_parquet_parallel.py to convert ROOT shards to parquet in parallel and place them correctly, and create_root_files.py to generate new ROOT shards via a minicalosim executable. The loader gains .manifest file support so pools/ (train/dev/ holdout shard lists) can be passed straight to `giant train`. Co-Authored-By: Claude Sonnet 4.6 --- giant/data/loader.py | 24 ++ scripts/bump_dataset_version.py | 197 +++++++++++++++++ scripts/create_root_files.py | 282 ++++++++++++++++++++++++ scripts/migrate_geant_steps.py | 258 ++++++++++++++++++++++ scripts/steps_to_parquet_parallel.py | 251 +++++++++++++++++++++ tests/test_bump_dataset_version.py | 95 ++++++++ tests/test_create_root_files.py | 274 +++++++++++++++++++++++ tests/test_loader.py | 60 +++++ tests/test_steps_to_parquet_parallel.py | 186 ++++++++++++++++ 9 files changed, 1627 insertions(+) create mode 100644 scripts/bump_dataset_version.py create mode 100644 scripts/create_root_files.py create mode 100644 scripts/migrate_geant_steps.py create mode 100644 scripts/steps_to_parquet_parallel.py create mode 100644 tests/test_bump_dataset_version.py create mode 100644 tests/test_create_root_files.py create mode 100644 tests/test_loader.py create mode 100644 tests/test_steps_to_parquet_parallel.py diff --git a/giant/data/loader.py b/giant/data/loader.py index 72f3bbc..ceeb195 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -5,9 +5,33 @@ import numpy as np import pandas as pd import pyarrow.parquet as pq +# A manifest is a plain text file listing one parquet path per line, used to +# name a curated subset of files (e.g. a train/holdout pool) without copying +# or symlinking the underlying parquet files. Lines are resolved relative to +# the manifest's own directory, so the manifest stays valid if the whole +# dataset tree is moved or copied elsewhere intact. +MANIFEST_SUFFIX = ".manifest" + + +def _read_manifest(path: Path) -> list[Path]: + files = [] + for line in path.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + resolved = (path.parent / line).resolve() + if not resolved.is_file(): + raise FileNotFoundError(f"{path} lists missing file: {resolved}") + files.append(resolved) + if not files: + raise FileNotFoundError(f"manifest {path} lists no files") + return files + def find_parquet_files(path: str | Path) -> list[Path]: p = Path(path) + if p.suffix == MANIFEST_SUFFIX: + return _read_manifest(p) if p.is_dir(): files = sorted(p.glob("*.parquet")) if not files: diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py new file mode 100644 index 0000000..958728b --- /dev/null +++ b/scripts/bump_dataset_version.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +"""Cut a new raw generation or processed schema version for the geant_steps +dataset tree (see scripts/migrate_geant_steps.py for the layout): + + raw////shard-NNN.root + processed/////shard-NNN.parquet + +`gen` bumps when the underlying ROOT changes (geometry/physics-list/macro). +`schema` bumps when the parquet export (steps_to_parquet.py or similar) +changes, and is scoped to its gen — a new gen always starts back at schema1. + +Creates the new (empty) target directory and appends a dated, reasoned entry +to VERSIONS.md. Defaults to a dry run; pass --execute to apply. + +Usage: + bump_dataset_version.py bump-gen --kind steps --reason "switched EM physics list" + bump_dataset_version.py bump-schema --kind steps --gen gen1 --reason "added e_sec column" + bump_dataset_version.py status +""" + +import argparse +import datetime as dt +import re +import subprocess +from pathlib import Path + +GEN_RE = re.compile(r"^gen(\d+)$") +SCHEMA_RE = re.compile(r"^schema(\d+)$") + + +def _max_index(parent: Path, pattern: re.Pattern) -> int: + """Highest N across child dir names matching *pattern* (0 if none/missing).""" + if not parent.is_dir(): + return 0 + best = 0 + for child in parent.iterdir(): + m = pattern.match(child.name) + if m and child.is_dir(): + best = max(best, int(m.group(1))) + return best + + +def _git_user_name() -> str | None: + try: + out = subprocess.run( + ["git", "config", "user.name"], capture_output=True, text=True, timeout=2 + ) + except OSError: + return None + name = out.stdout.strip() + return name or None + + +def plan_bump_gen( + root: Path, kind: str, reason: str, by: str | None, date: str +) -> tuple[list[Path], str]: + """New gen tag is one past the highest seen under raw/ or processed/ for *kind* + — checking both, since a gen can exist in one tree before the other catches up.""" + next_gen = ( + max( + _max_index(root / "raw" / kind, GEN_RE), + _max_index(root / "processed" / kind, GEN_RE), + ) + + 1 + ) + gen_tag = f"gen{next_gen}" + new_dirs = [ + root / "raw" / kind / gen_tag, + root / "processed" / kind / gen_tag / "schema1", + ] + by_suffix = f" ({by})" if by else "" + log_line = f"- `{gen_tag}` (kind={kind}) — {date} — {reason}{by_suffix}" + return new_dirs, log_line + + +def plan_bump_schema( + root: Path, kind: str, gen_tag: str, reason: str, by: str | None, date: str +) -> tuple[list[Path], str]: + if not GEN_RE.match(gen_tag): + raise SystemExit(f"error: --gen must look like 'genN', got {gen_tag!r}") + raw_gen_dir = root / "raw" / kind / gen_tag + processed_gen_dir = root / "processed" / kind / gen_tag + if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir(): + raise SystemExit( + f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first" + ) + next_schema = _max_index(processed_gen_dir, SCHEMA_RE) + 1 + schema_tag = f"schema{next_schema}" + new_dirs = [processed_gen_dir / schema_tag] + by_suffix = f" ({by})" if by else "" + log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}" + return new_dirs, log_line + + +def apply_bump(root: Path, new_dirs: list[Path], log_line: str) -> None: + for d in new_dirs: + d.mkdir(parents=True, exist_ok=True) + versions_path = root / "VERSIONS.md" + if not versions_path.exists(): + versions_path.write_text("# Dataset versions\n\n") + with versions_path.open("a") as f: + f.write(log_line + "\n") + + +def print_status(root: Path) -> None: + raw_root = root / "raw" + if not raw_root.is_dir(): + print(f"no raw/ tree found under {root}") + return + for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()): + kind = kind_dir.name + gens = sorted( + int(m.group(1)) + for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir()) + if m + ) + print(f"kind={kind}") + for gen in gens: + gen_tag = f"gen{gen}" + schema_dir = root / "processed" / kind / gen_tag + schemas = sorted( + int(m.group(1)) + for m in ( + SCHEMA_RE.match(p.name) + for p in (schema_dir.iterdir() if schema_dir.is_dir() else []) + if p.is_dir() + ) + if m + ) + schema_str = ", ".join(f"schema{s}" for s in schemas) or "(none)" + print(f" {gen_tag}: {schema_str}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--root", + default="/ceph/lbogner/geant_steps", + help="Dataset root (default: /ceph/lbogner/geant_steps)", + ) + sub = parser.add_subparsers(dest="command", required=True) + + kind_help = "steps | hits | ... (default: steps)" + + p_gen = sub.add_parser("bump-gen", help="Cut a new raw generation") + p_gen.add_argument("--kind", default="steps", help=kind_help) + p_gen.add_argument("--reason", required=True, help="Why this gen exists") + p_gen.add_argument("--by", default=None, help="Attribution (default: git user.name)") + p_gen.add_argument("--date", default=None, help="Override date (default: today, ISO)") + p_gen.add_argument("--execute", action="store_true", help="Apply (default: dry run)") + + p_schema = sub.add_parser("bump-schema", help="Cut a new schema within a gen") + p_schema.add_argument("--kind", default="steps", help=kind_help) + p_schema.add_argument("--gen", required=True, help="Existing gen tag, e.g. gen1") + p_schema.add_argument("--reason", required=True, help="Why this schema exists") + p_schema.add_argument("--by", default=None, help="Attribution (default: git user.name)") + p_schema.add_argument("--date", default=None, help="Override date (default: today, ISO)") + p_schema.add_argument("--execute", action="store_true", help="Apply (default: dry run)") + + sub.add_parser("status", help="List existing gens/schemas per kind") + + args = parser.parse_args() + root = Path(args.root) + if not root.is_dir(): + parser.error(f"{root} is not a directory") + + if args.command == "status": + print_status(root) + return + + date = args.date or dt.date.today().isoformat() + by = args.by if args.by is not None else _git_user_name() + + if args.command == "bump-gen": + new_dirs, log_line = plan_bump_gen(root, args.kind, args.reason, by, date) + else: + new_dirs, log_line = plan_bump_schema( + root, args.kind, args.gen, args.reason, by, date + ) + + print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===") + print("new directories:") + for d in new_dirs: + print(f" {d}") + print("VERSIONS.md entry:") + print(f" {log_line}") + + if not args.execute: + print("\nDry run only — pass --execute to apply.") + return + + apply_bump(root, new_dirs, log_line) + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/scripts/create_root_files.py b/scripts/create_root_files.py new file mode 100644 index 0000000..8ca23ad --- /dev/null +++ b/scripts/create_root_files.py @@ -0,0 +1,282 @@ +#!/usr/bin/env python3 +"""Generate new ROOT shards by running a minicalosim executable (e.g. +run_pbwo4, run_sampling) and filing the output into the dataset's raw/ tree: + + raw////shard-NNN.root + +These executables take `[configName] nEvents` and always write a fixed-name +*.root file into the current directory — so running several in parallel +needs separate working directories, and the output filename has to be +discovered rather than assumed (it differs per executable: run_pbwo4 writes +pbwo4_events_hits.root, run_sampling writes sampling__events_hits.root, +others may differ again). This script gives each run its own scratch +directory under /.sim-tmp/, requires exactly one *.root to +appear there, and moves it to the next free shard index for that detector +(existing shards are never overwritten). + +--gen must already exist under raw// — create one first with +bump_dataset_version.py bump-gen. + +Usage: + create_root_files.py --executable build/run_pbwo4 --detector pbwo4 \\ + --gen gen1 --num-files 4 --events-per-file 10000 --execute + + create_root_files.py --executable build/run_sampling \\ + --detector sampling_pb_scint:pb_scint \\ + --detector sampling_fe_scint:fe_scint \\ + --gen gen1 --num-files 4 --events-per-file 10000 --jobs 8 --execute +""" + +import argparse +import os +import re +import shutil +import subprocess +import sys +import uuid +from dataclasses import dataclass +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +# Must match scripts/bump_dataset_version.py's GEN_RE. +GEN_RE = re.compile(r"^gen\d+$") +SHARD_RE = re.compile(r"^shard-(\d+)\.root$") + + +class PlanError(ValueError): + pass + + +@dataclass(frozen=True) +class SimJob: + detector: str + config: str | None + shard_index: int + + +@dataclass(frozen=True) +class JobResult: + job: SimJob + ok: bool + dest: Path | None + message: str + stdout: str + stderr: str + + +def parse_detector_spec(spec: str) -> tuple[str, str | None]: + """'sampling_pb_scint:pb_scint' -> ('sampling_pb_scint', 'pb_scint'); + 'pbwo4' -> ('pbwo4', None).""" + if ":" in spec: + label, config = spec.split(":", 1) + if not label or not config: + raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG") + return label, config + return spec, None + + +def next_shard_index(detector_dir: Path) -> int: + """One past the highest existing shard-NNN.root in *detector_dir* (0 if none/missing).""" + if not detector_dir.is_dir(): + return 0 + best = -1 + for child in detector_dir.iterdir(): + m = SHARD_RE.match(child.name) + if m and child.is_file(): + best = max(best, int(m.group(1))) + return best + 1 + + +def plan_jobs( + detector_specs: list[str], + num_files: int, + dataset_root: Path, + kind: str, + gen: str, +) -> list[SimJob]: + if not GEN_RE.match(gen): + raise PlanError(f"--gen must look like 'genN', got {gen!r}") + gen_dir = dataset_root / "raw" / kind / gen + if not gen_dir.is_dir(): + raise PlanError( + f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first" + ) + + jobs = [] + for spec in detector_specs: + label, config = parse_detector_spec(spec) + start = next_shard_index(gen_dir / label) + for i in range(num_files): + jobs.append(SimJob(detector=label, config=config, shard_index=start + i)) + return jobs + + +def run_job( + job: SimJob, + executable: Path, + events_per_file: int, + dataset_root: Path, + kind: str, + gen: str, + tmp_root: Path, +) -> JobResult: + workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}" + workdir.mkdir(parents=True) + + cmd = [str(executable)] + if job.config: + cmd.append(job.config) + cmd.append(str(events_per_file)) + + result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True) + + if result.returncode != 0: + return JobResult( + job, False, None, + f"executable exited {result.returncode}", + result.stdout, result.stderr, + ) + + produced = sorted(workdir.glob("*.root")) + if len(produced) != 1: + return JobResult( + job, False, None, + f"expected exactly one .root output in {workdir}, found {len(produced)}: " + f"{[p.name for p in produced]}", + result.stdout, result.stderr, + ) + + dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root" + if dest.exists(): + return JobResult( + job, False, None, f"refusing to overwrite existing {dest}", + result.stdout, result.stderr, + ) + + dest.parent.mkdir(parents=True, exist_ok=True) + shutil.move(str(produced[0]), str(dest)) + # rmtree, not rmdir: the executable may leave other side-effect files + # (logs, seed state, ...) behind in its scratch workdir besides the .root. + shutil.rmtree(workdir) + return JobResult(job, True, dest, "ok", result.stdout, result.stderr) + + +def run_all( + jobs: list[SimJob], + executable: Path, + events_per_file: int, + dataset_root: Path, + kind: str, + gen: str, + max_workers: int, + tmp_root: Path, +) -> list[JobResult]: + results = [] + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = { + pool.submit( + run_job, job, executable, events_per_file, dataset_root, kind, gen, tmp_root + ): job + for job in jobs + } + for future in as_completed(futures): + result = future.result() + label = f"{result.job.detector} shard-{result.job.shard_index:03d}" + status = "ok" if result.ok else f"FAILED: {result.message}" + print(f"\n=== {label}: {status} ===") + if result.stdout: + print(result.stdout, end="") + if result.stderr: + print(result.stderr, end="", file=sys.stderr) + results.append(result) + return results + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--executable", required=True, type=Path, help="Built minicalosim run_* executable" + ) + parser.add_argument( + "--detector", + action="append", + required=True, + metavar="NAME[:CONFIG]", + help="Dataset detector label, optionally with ':CONFIG' to pass as the " + "executable's config-name argument (e.g. sampling_pb_scint:pb_scint). " + "Omit ':CONFIG' for executables that take no config selector (e.g. run_pbwo4). " + "Repeatable.", + ) + parser.add_argument( + "--num-files", type=int, required=True, help="New shards to create per detector" + ) + parser.add_argument( + "--events-per-file", type=int, required=True, help="nEvents passed to the executable" + ) + parser.add_argument("--kind", default="steps", help="steps | hits | ... (default: steps)") + parser.add_argument("--gen", required=True, help="Existing gen tag under raw//, e.g. gen1") + parser.add_argument( + "--dataset-root", + default="/ceph/lbogner/geant_steps", + help="Dataset root (default: /ceph/lbogner/geant_steps)", + ) + parser.add_argument( + "-j", "--jobs", type=int, default=4, help="Parallel simulation runs (default: 4)" + ) + parser.add_argument( + "--execute", action="store_true", help="Actually run jobs (default: dry run / print plan)" + ) + return parser + + +def main() -> None: + parser = build_parser() + args = parser.parse_args() + + if args.jobs < 1: + parser.error("--jobs must be >= 1") + if args.num_files < 1: + parser.error("--num-files must be >= 1") + if args.events_per_file < 1: + parser.error("--events-per-file must be >= 1") + if not args.executable.is_file() or not os.access(args.executable, os.X_OK): + parser.error(f"{args.executable} is not an executable file") + + dataset_root = Path(args.dataset_root) + try: + jobs = plan_jobs(args.detector, args.num_files, dataset_root, args.kind, args.gen) + except PlanError as exc: + parser.error(str(exc)) + + print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===") + print(f"executable: {args.executable}") + for job in jobs: + cmd = [str(args.executable)] + ([job.config] if job.config else []) + [str(args.events_per_file)] + dest = dataset_root / "raw" / args.kind / args.gen / job.detector / f"shard-{job.shard_index:03d}.root" + print(f" {' '.join(cmd)} -> {dest}") + + if not args.execute: + print("\nDry run only — pass --execute to apply.") + return + + tmp_root = dataset_root / ".sim-tmp" + tmp_root.mkdir(parents=True, exist_ok=True) + results = run_all( + jobs, args.executable, args.events_per_file, dataset_root, args.kind, args.gen, + max_workers=args.jobs, tmp_root=tmp_root, + ) + if tmp_root.is_dir() and not any(tmp_root.iterdir()): + tmp_root.rmdir() + + failures = [r for r in results if not r.ok] + if failures: + print(f"\n{len(failures)} of {len(results)} job(s) failed:", file=sys.stderr) + for r in failures: + print(f" {r.job.detector} shard-{r.job.shard_index:03d}: {r.message}", file=sys.stderr) + sys.exit(1) + + print(f"\nAll {len(results)} job(s) completed.") + + +if __name__ == "__main__": + main() diff --git a/scripts/migrate_geant_steps.py b/scripts/migrate_geant_steps.py new file mode 100644 index 0000000..ec49718 --- /dev/null +++ b/scripts/migrate_geant_steps.py @@ -0,0 +1,258 @@ +#!/usr/bin/env python3 +"""One-time migration of /ceph/lbogner/geant_steps into the versioned layout: + + raw////shard-NNN.root + processed/////shard-NNN.parquet + pools//{dev,full,holdout,unclassified}.txt (manifests, no copied bytes) + derived/predictions//_shard-NNN_predicted[_local].parquet + +Walks the whole tree (regardless of which legacy subfolder a file currently sits +in — train/, sampling_train/, sampling_train/small/, or loose at the top level), +classifies every .root/.parquet by detector+shard via filename pattern, and looks +up pool membership from POOL_ASSIGNMENT — decoupling "where it sits today" from +"which pool it belongs to". + +Defaults to a dry run (prints the planned moves and manifest contents). Pass +--execute to actually move files and write manifests. +""" + +import argparse +import os +import re +import shutil +from collections import defaultdict +from pathlib import Path + +# Must match giant.data.loader.MANIFEST_SUFFIX. +MANIFEST_SUFFIX = ".manifest" + +# Shard membership is keyed by detector name, independent of current location, +# since today's pool assignment is encoded only by *which folder a file's +# parquet was copied into* — not by anything in the filename itself. +POOL_ASSIGNMENT: dict[str, dict[str, range | list[int]]] = { + "pbwo4": {"full": range(0, 6), "holdout": range(6, 10)}, + "sampling_fe_scint": {"dev": [0], "full": [1, 2], "holdout": [3]}, + "sampling_pb_lar": {"dev": [0], "full": [1, 2], "holdout": [3]}, + "sampling_pb_scint": {"dev": [0], "full": [1, 2], "holdout": [3]}, + "sampling_w_scint_ecal": {"dev": [0], "full": [1, 2], "holdout": [3]}, +} + +# This is the migration baseline: everything that exists today is one +# generation of raw data (gen1) exported with one parquet schema (schema1). +GEN = "gen1" +SCHEMA = "schema1" +LEGACY_GEN = "gen0" +LEGACY_SCHEMA = "schema0" + +PREDICTED_RE = re.compile( + r"^(?P[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P\d+)" + r"_predicted(?P_local)?\.parquet$" +) +SHARD_RE = re.compile( + r"^(?P[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P\d+)\.(?Proot|parquet)$" +) +LEGACY_PREDICTED_RE = re.compile( + r"^pbwo4_10000events_hits_predicted(?P_local)?\.parquet$" +) +LEGACY_RE = re.compile(r"^pbwo4_10000events_hits\.(?Proot|parquet)$") + + +def pool_for(detector: str, shard: int) -> str: + rules = POOL_ASSIGNMENT.get(detector) + if rules is None: + return "unclassified" + for pool, shards in rules.items(): + if shard in shards: + return pool + return "unclassified" + + +def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]: + """Return (moves, unrecognized) where moves is [(src, dst), ...].""" + moves: list[tuple[Path, Path]] = [] + unrecognized: list[Path] = [] + + for path in sorted(src_root.rglob("*")): + if not path.is_file(): + continue + name = path.name + + m = PREDICTED_RE.match(name) + if m: + detector, shard = m["detector"], int(m["shard"]) + suffix = "_predicted_local" if m["local"] else "_predicted" + dst = ( + src_root + / "derived" + / "predictions" + / "unknown-checkpoint" + / f"{detector}_shard-{shard:03d}{suffix}.parquet" + ) + moves.append((path, dst)) + continue + + m = LEGACY_PREDICTED_RE.match(name) + if m: + suffix = "_predicted_local" if m["local"] else "_predicted" + dst = ( + src_root + / "derived" + / "predictions" + / "unknown-checkpoint-legacy" + / f"pbwo4_hits_shard-000{suffix}.parquet" + ) + moves.append((path, dst)) + continue + + m = SHARD_RE.match(name) + if m: + detector, shard, ext = m["detector"], int(m["shard"]), m["ext"] + if ext == "root": + dst = ( + src_root + / "raw" + / "steps" + / GEN + / detector + / f"shard-{shard:03d}.root" + ) + else: + dst = ( + src_root + / "processed" + / "steps" + / GEN + / SCHEMA + / detector + / f"shard-{shard:03d}.parquet" + ) + moves.append((path, dst)) + continue + + m = LEGACY_RE.match(name) + if m: + ext = m["ext"] + if ext == "root": + dst = ( + src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root" + ) + else: + dst = ( + src_root + / "processed" + / "hits" + / LEGACY_GEN + / LEGACY_SCHEMA + / "pbwo4" + / "shard-000.parquet" + ) + moves.append((path, dst)) + continue + + unrecognized.append(path) + + return moves, unrecognized + + +def plan_manifests(src_root: Path) -> dict[Path, list[str]]: + """Map manifest path -> sorted list of parquet paths, each relative to the + manifest's own directory (so the manifest stays valid if the tree moves).""" + manifests: dict[Path, list[tuple[int, Path]]] = defaultdict(list) + for detector, rules in POOL_ASSIGNMENT.items(): + manifest_dir = src_root / "pools" / detector + for pool, shards in rules.items(): + manifest_path = manifest_dir / f"{pool}{MANIFEST_SUFFIX}" + for shard in shards: + dst = ( + src_root + / "processed" + / "steps" + / GEN + / SCHEMA + / detector + / f"shard-{shard:03d}.parquet" + ) + manifests[manifest_path].append((shard, dst)) + return { + k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)] + for k, v in manifests.items() + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "root", + nargs="?", + default="/ceph/lbogner/geant_steps", + help="Dataset root to migrate in place (default: /ceph/lbogner/geant_steps)", + ) + parser.add_argument( + "--execute", + action="store_true", + help="Actually move files and write manifests (default: dry run / print plan only)", + ) + args = parser.parse_args() + + src_root = Path(args.root) + if not src_root.is_dir(): + parser.error(f"{src_root} is not a directory") + + moves, unrecognized = plan_moves(src_root) + manifests = plan_manifests(src_root) + + print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'}: {src_root} ===\n") + + print(f"-- {len(moves)} file move(s) --") + for src, dst in moves: + print(f" {src.relative_to(src_root)} -> {dst.relative_to(src_root)}") + + print(f"\n-- {len(manifests)} manifest(s) --") + for manifest_path, lines in sorted(manifests.items()): + print(f" {manifest_path.relative_to(src_root)}:") + for line in lines: + print(f" {line}") + + if unrecognized: + print(f"\n-- {len(unrecognized)} UNRECOGNIZED file(s), left in place --") + for path in unrecognized: + print(f" {path.relative_to(src_root)}") + + if not args.execute: + print("\nDry run only — pass --execute to apply.") + return + + for src, dst in moves: + if dst.exists(): + raise FileExistsError(f"refusing to overwrite existing file: {dst}") + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.move(str(src), str(dst)) + + for manifest_path, lines in manifests.items(): + manifest_path.parent.mkdir(parents=True, exist_ok=True) + manifest_path.write_text("\n".join(lines) + "\n") + + versions_path = src_root / "VERSIONS.md" + if not versions_path.exists(): + versions_path.write_text( + "# Dataset versions\n\n" + f"- `{GEN}` / `{SCHEMA}` (steps) — migration baseline: all data that existed " + "before the directory restructure, generation and export details unrecorded.\n" + f"- `{LEGACY_GEN}` / `{LEGACY_SCHEMA}` (hits) — legacy single-shot " + "`pbwo4_10000events_hits` dataset (Hits tree, predates the 10k-shard scheme).\n" + ) + + # Legacy pool dirs (train/, sampling_train/, sampling_train/small/) are now + # empty since their contents were classified by filename, not location — + # remove them, but only if a move actually emptied them. + for stale_dir in ("train", "sampling_train/small", "sampling_train"): + d = src_root / stale_dir + if d.is_dir() and not any(d.iterdir()): + d.rmdir() + print(f"removed now-empty directory: {d.relative_to(src_root)}") + + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/scripts/steps_to_parquet_parallel.py b/scripts/steps_to_parquet_parallel.py new file mode 100644 index 0000000..7344b04 --- /dev/null +++ b/scripts/steps_to_parquet_parallel.py @@ -0,0 +1,251 @@ +#!/usr/bin/env python3 +"""Convert many ROOT files to Parquet by fanning out to steps_to_parquet.py. + +steps_to_parquet.py itself converts a list of files one at a time; this wraps +it to run up to --jobs conversions concurrently, each as its own subprocess +(invoked with the same Python executable running this script, so it picks up +the active venv/uv environment automatically). + +Inputs must live under /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. + +Usage: + steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/shard-000.root raw/steps/gen1/pbwo4/shard-001.root + steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --jobs 8 + steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --schema schema2 +""" + +import argparse +import re +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +_STEPS_TO_PARQUET = Path(__file__).resolve().parent / "steps_to_parquet.py" + +# Must match scripts/bump_dataset_version.py's GEN_RE / SCHEMA_RE. +GEN_RE = re.compile(r"^gen\d+$") +SCHEMA_RE = re.compile(r"^schema(\d+)$") + + +class DestinationError(ValueError): + pass + + +def latest_schema_tag(processed_gen_dir: Path) -> str | None: + """Highest schemaN dir directly under *processed_gen_dir*, or None if none exist.""" + if not processed_gen_dir.is_dir(): + return None + best_tag, best_n = None, -1 + for child in processed_gen_dir.iterdir(): + m = SCHEMA_RE.match(child.name) + if m and child.is_dir() and int(m.group(1)) > best_n: + best_tag, best_n = child.name, int(m.group(1)) + return best_tag + + +def resolve_destination( + root_file: Path, dataset_root: Path, schema_override: str | None +) -> Path: + """Map raw////.root (relative to *dataset_root*) + to processed/////.parquet. + + Raises DestinationError if *root_file* doesn't fit that layout, or if no + schema can be determined (no --schema and none exists yet for that gen). + """ + dataset_root = dataset_root.resolve() + root_file = root_file.resolve() + try: + rel = root_file.relative_to(dataset_root) + except ValueError: + raise DestinationError(f"{root_file} is not under dataset root {dataset_root}") + + parts = rel.parts + if ( + len(parts) != 5 + or parts[0] != "raw" + or not GEN_RE.match(parts[2]) + or not parts[4].endswith(".root") + ): + raise DestinationError( + f"{root_file} does not match raw////.root " + f"under {dataset_root} (got relative path: {rel})" + ) + _, kind, gen_tag, detector, filename = parts + shard_stem = Path(filename).stem + + processed_gen_dir = dataset_root / "processed" / kind / gen_tag + schema_tag = schema_override or latest_schema_tag(processed_gen_dir) + if schema_tag is None: + raise DestinationError( + f"no schema exists yet under {processed_gen_dir} — pass --schema or run " + "bump_dataset_version.py bump-schema first" + ) + return processed_gen_dir / schema_tag / detector / f"{shard_stem}.parquet" + + +def _convert_one( + root_file: str, + batch_size: str, + tree: str, + compression: str, + steps_to_parquet_path: Path, + output_path: Path | None, +) -> tuple[str, int, str, str]: + cmd = [ + sys.executable, + str(steps_to_parquet_path), + root_file, + "--batch-size", + batch_size, + "--tree", + tree, + "--compression", + compression, + ] + if output_path is not None: + output_path.parent.mkdir(parents=True, exist_ok=True) + cmd += ["--output", str(output_path)] + result = subprocess.run(cmd, capture_output=True, text=True) + return root_file, result.returncode, result.stdout, result.stderr + + +def run_parallel( + root_files: list[str], + jobs: int, + batch_size: str = "100 MB", + tree: str = "Steps", + compression: str = "snappy", + steps_to_parquet_path: Path = _STEPS_TO_PARQUET, + output_for: dict[str, Path] | None = None, +) -> list[tuple[str, int, str, str]]: + """Run one steps_to_parquet.py subprocess per file, up to *jobs* at a time. + + *output_for*, if given, maps each root_file to the parquet path it should + be written to (passed through as steps_to_parquet.py's --output); files + missing from the map fall back to steps_to_parquet.py's own default + (parquet written next to the input .root). + + Returns one (root_file, returncode, stdout, stderr) tuple per file, in + completion order (not necessarily input order). + """ + output_for = output_for or {} + results = [] + with ThreadPoolExecutor(max_workers=jobs) as pool: + futures = { + pool.submit( + _convert_one, + f, + batch_size, + tree, + compression, + steps_to_parquet_path, + output_for.get(f), + ): f + for f in root_files + } + for future in as_completed(futures): + root_file, code, stdout, stderr = future.result() + status = "ok" if code == 0 else f"FAILED (exit {code})" + print(f"\n=== {root_file}: {status} ===") + if stdout: + print(stdout, end="") + if stderr: + print(stderr, end="", file=sys.stderr) + results.append((root_file, code, stdout, stderr)) + return results + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Convert many ROOT files to Parquet in parallel via steps_to_parquet.py." + ) + parser.add_argument("root_files", nargs="+", help="Input ROOT file(s)") + parser.add_argument( + "-j", + "--jobs", + type=int, + default=4, + help="Number of conversions to run in parallel (default: 4)", + ) + parser.add_argument( + "--batch-size", + default="100 MB", + help="Uproot read batch size (default: '100 MB'). E.g. '50 MB', '500000' (rows).", + ) + parser.add_argument( + "--tree", + default="Steps", + help="Tree name inside the ROOT file (default: Steps)", + ) + parser.add_argument( + "--compression", + default="snappy", + choices=["snappy", "lz4", "zstd", "gzip", "none"], + help="Parquet compression codec (default: snappy)", + ) + parser.add_argument( + "--dataset-root", + default="/ceph/lbogner/geant_steps", + help="Dataset root containing raw/ and processed/ (default: /ceph/lbogner/geant_steps)", + ) + parser.add_argument( + "--schema", + default=None, + help="Schema tag to write parquets under, e.g. schema2 " + "(default: highest schemaN already under processed///)", + ) + return parser + + +def main() -> None: + parser = build_parser() + args = parser.parse_args() + + if args.jobs < 1: + parser.error("--jobs must be >= 1") + + dataset_root = Path(args.dataset_root) + output_for: dict[str, Path] = {} + errors: list[str] = [] + for f in args.root_files: + try: + output_for[f] = resolve_destination(Path(f), dataset_root, args.schema) + except DestinationError as exc: + errors.append(str(exc)) + + if errors: + for err in errors: + print(f"error: {err}", file=sys.stderr) + sys.exit(1) + + for root_file, dest in output_for.items(): + print(f"{root_file} -> {dest}") + + results = run_parallel( + args.root_files, + jobs=args.jobs, + batch_size=args.batch_size, + tree=args.tree, + compression=args.compression, + output_for=output_for, + ) + + failures = [root_file for root_file, code, _, _ in results if code != 0] + if failures: + print(f"\n{len(failures)} of {len(results)} conversion(s) failed:", file=sys.stderr) + for root_file in failures: + print(f" {root_file}", file=sys.stderr) + sys.exit(1) + + print(f"\nAll {len(results)} conversion(s) completed.") + + +if __name__ == "__main__": + main() diff --git a/tests/test_bump_dataset_version.py b/tests/test_bump_dataset_version.py new file mode 100644 index 0000000..e1de496 --- /dev/null +++ b/tests/test_bump_dataset_version.py @@ -0,0 +1,95 @@ +import importlib.util +from pathlib import Path + +# scripts/ is not an installed package — load the module straight from its path. +_SPEC = importlib.util.spec_from_file_location( + "bump_dataset_version", + Path(__file__).resolve().parents[1] / "scripts" / "bump_dataset_version.py", +) +bump_dataset_version = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(bump_dataset_version) + +plan_bump_gen = bump_dataset_version.plan_bump_gen +plan_bump_schema = bump_dataset_version.plan_bump_schema +apply_bump = bump_dataset_version.apply_bump + + +def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path): + dirs, log_line = plan_bump_gen(tmp_path, "steps", "first generation", None, "2026-01-01") + assert dirs == [ + tmp_path / "raw" / "steps" / "gen1", + tmp_path / "processed" / "steps" / "gen1" / "schema1", + ] + assert "`gen1`" in log_line + assert "first generation" in log_line + + +def test_bump_gen_increments_past_existing(tmp_path): + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + (tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True) + dirs, _ = plan_bump_gen(tmp_path, "steps", "next gen", None, "2026-01-01") + assert dirs[0] == tmp_path / "raw" / "steps" / "gen3" + + +def test_bump_gen_checks_both_raw_and_processed_trees(tmp_path): + # processed/ is ahead of raw/ — next gen must still be past the max of both. + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + (tmp_path / "processed" / "steps" / "gen4" / "schema1").mkdir(parents=True) + dirs, _ = plan_bump_gen(tmp_path, "steps", "next gen", None, "2026-01-01") + assert dirs[0] == tmp_path / "raw" / "steps" / "gen5" + + +def test_bump_gen_kinds_are_independent(tmp_path): + (tmp_path / "raw" / "steps" / "gen5").mkdir(parents=True) + dirs, _ = plan_bump_gen(tmp_path, "hits", "first hits gen", None, "2026-01-01") + assert dirs[0] == tmp_path / "raw" / "hits" / "gen1" + + +def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path): + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + dirs, log_line = plan_bump_schema( + tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01" + ) + assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema1"] + assert "`gen1`/`schema1`" in log_line + + +def test_bump_schema_increments_within_its_gen(tmp_path): + (tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True) + (tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True) + dirs, _ = plan_bump_schema(tmp_path, "steps", "gen1", "next schema", None, "2026-01-01") + assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema3"] + + +def test_bump_schema_does_not_see_other_gens_schemas(tmp_path): + (tmp_path / "processed" / "steps" / "gen1" / "schema5").mkdir(parents=True) + (tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True) + dirs, _ = plan_bump_schema(tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01") + assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"] + + +def test_bump_schema_rejects_nonexistent_gen(tmp_path): + try: + plan_bump_schema(tmp_path, "steps", "gen9", "oops", None, "2026-01-01") + assert False, "expected SystemExit" + except SystemExit: + pass + + +def test_apply_bump_creates_dirs_and_appends_log(tmp_path): + dirs, log_line = plan_bump_gen(tmp_path, "steps", "reason A", "alice", "2026-01-01") + apply_bump(tmp_path, dirs, log_line) + for d in dirs: + assert d.is_dir() + text = (tmp_path / "VERSIONS.md").read_text() + assert "reason A" in text + assert "alice" in text + + +def test_apply_bump_appends_without_clobbering_existing_log(tmp_path): + (tmp_path / "VERSIONS.md").write_text("# Dataset versions\n\n- existing entry\n") + dirs, log_line = plan_bump_gen(tmp_path, "steps", "reason B", None, "2026-01-02") + apply_bump(tmp_path, dirs, log_line) + text = (tmp_path / "VERSIONS.md").read_text() + assert "existing entry" in text + assert "reason B" in text diff --git a/tests/test_create_root_files.py b/tests/test_create_root_files.py new file mode 100644 index 0000000..18d374f --- /dev/null +++ b/tests/test_create_root_files.py @@ -0,0 +1,274 @@ +import importlib.util +import json +import stat +from pathlib import Path + +import pytest + +# scripts/ is not an installed package — load the module straight from its path. +_SPEC = importlib.util.spec_from_file_location( + "create_root_files", + Path(__file__).resolve().parents[1] / "scripts" / "create_root_files.py", +) +create_root_files = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(create_root_files) + +parse_detector_spec = create_root_files.parse_detector_spec +next_shard_index = create_root_files.next_shard_index +plan_jobs = create_root_files.plan_jobs +run_job = create_root_files.run_job +run_all = create_root_files.run_all +SimJob = create_root_files.SimJob +PlanError = create_root_files.PlanError + + +def _write_fake_executable( + path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0 +) -> Path: + """Stand-in for run_pbwo4/run_sampling: writes *output_count* .root files + into its own cwd (so callers can verify each job gets an isolated workdir + and that the workdir ends up holding *only* the .root output, matching + real executables that may also drop other side-effect files). Each .root + file's content is a JSON record of argv/cwd/timing, so a test can recover + that info from the moved file after the workdir is gone. Also writes a + non-.root side file unconditionally, to catch any cleanup code that + assumes the workdir is empty after the .root is moved out. + """ + path.write_text( + f"""#!/usr/bin/env python3 +import json, os, sys, time + +start = time.time() +time.sleep({sleep}) +end = time.time() +payload = json.dumps( + {{"argv": sys.argv[1:], "cwd": os.getcwd(), "start": start, "end": end}} +) +for i in range({output_count}): + with open(f"out_{{i}}.root", "w") as f: + f.write(payload) +with open("side_effect.log", "w") as f: + f.write("not a root file") +sys.exit({exit_code}) +""" + ) + path.chmod(path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH) + return path + + +def test_parse_detector_spec_with_config(): + assert parse_detector_spec("sampling_pb_scint:pb_scint") == ("sampling_pb_scint", "pb_scint") + + +def test_parse_detector_spec_without_config(): + assert parse_detector_spec("pbwo4") == ("pbwo4", None) + + +def test_parse_detector_spec_rejects_empty_parts(): + with pytest.raises(PlanError): + parse_detector_spec(":pb_scint") + with pytest.raises(PlanError): + parse_detector_spec("sampling_pb_scint:") + + +def test_next_shard_index_missing_dir(tmp_path): + assert next_shard_index(tmp_path / "nope") == 0 + + +def test_next_shard_index_continues_past_existing(tmp_path): + d = tmp_path / "pbwo4" + d.mkdir() + (d / "shard-000.root").touch() + (d / "shard-005.root").touch() + (d / "not-a-shard.root").touch() + assert next_shard_index(d) == 6 + + +def test_plan_jobs_rejects_missing_gen(tmp_path): + with pytest.raises(PlanError): + plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1") + + +def test_plan_jobs_rejects_malformed_gen(tmp_path): + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + with pytest.raises(PlanError): + plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen") + + +def test_plan_jobs_continues_from_existing_shards(tmp_path): + gen_dir = tmp_path / "raw" / "steps" / "gen1" + (gen_dir / "pbwo4").mkdir(parents=True) + (gen_dir / "pbwo4" / "shard-000.root").touch() + (gen_dir / "pbwo4" / "shard-001.root").touch() + + jobs = plan_jobs(["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1") + + assert [j.shard_index for j in jobs] == [2, 3, 4] + assert all(j.detector == "pbwo4" and j.config is None for j in jobs) + + +def test_plan_jobs_multiple_detectors_each_start_independently(tmp_path): + gen_dir = tmp_path / "raw" / "steps" / "gen1" + (gen_dir / "sampling_pb_scint").mkdir(parents=True) + (gen_dir / "sampling_pb_scint" / "shard-003.root").touch() + (gen_dir / "sampling_fe_scint").mkdir(parents=True) + + jobs = plan_jobs( + ["sampling_pb_scint:pb_scint", "sampling_fe_scint:fe_scint"], + num_files=2, + dataset_root=tmp_path, + kind="steps", + gen="gen1", + ) + + by_detector = {} + for j in jobs: + by_detector.setdefault(j.detector, []).append(j.shard_index) + assert by_detector["sampling_pb_scint"] == [4, 5] + assert by_detector["sampling_fe_scint"] == [0, 1] + + +def test_run_job_moves_output_to_correct_shard_path(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py") + gen_dir = tmp_path / "raw" / "steps" / "gen1" + gen_dir.mkdir(parents=True) + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + job = SimJob(detector="pbwo4", config=None, shard_index=7) + result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + + assert result.ok + assert result.dest == gen_dir / "pbwo4" / "shard-007.root" + assert result.dest.is_file() + assert not any(tmp_root.iterdir()) # workdir cleaned up + + +def test_run_job_passes_config_arg_and_isolates_cwd(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py") + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + job = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0) + result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + + assert result.ok + payload = json.loads(result.dest.read_text()) + assert payload["argv"] == ["pb_scint", "10000"] + # ran in its own scratch workdir under .sim-tmp, not directly in dataset_root + assert payload["cwd"] != str(tmp_path) + assert str(tmp_root) in payload["cwd"] + + +def test_run_job_omits_config_arg_when_none(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py") + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + job = SimJob(detector="pbwo4", config=None, shard_index=0) + result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + + payload = json.loads(result.dest.read_text()) + assert payload["argv"] == ["10000"] + + +def test_run_job_fails_when_executable_errors(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py", exit_code=1) + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + job = SimJob(detector="pbwo4", config=None, shard_index=0) + result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + + assert not result.ok + assert "exited 1" in result.message + + +def test_run_job_fails_when_no_root_file_produced(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py", output_count=0) + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + job = SimJob(detector="pbwo4", config=None, shard_index=0) + result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + + assert not result.ok + assert "found 0" in result.message + + +def test_run_job_fails_when_multiple_root_files_produced(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py", output_count=2) + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + job = SimJob(detector="pbwo4", config=None, shard_index=0) + result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + + assert not result.ok + assert "found 2" in result.message + + +def test_run_job_refuses_to_overwrite_existing_shard(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py") + detector_dir = tmp_path / "raw" / "steps" / "gen1" / "pbwo4" + detector_dir.mkdir(parents=True) + (detector_dir / "shard-000.root").write_text("already here") + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + job = SimJob(detector="pbwo4", config=None, shard_index=0) + result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + + assert not result.ok + assert "overwrite" in result.message + assert (detector_dir / "shard-000.root").read_text() == "already here" + + +def test_run_all_caps_concurrency(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py", sleep=0.2) + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + jobs = [SimJob(detector="pbwo4", config=None, shard_index=i) for i in range(6)] + results = run_all(jobs, fake, 10000, tmp_path, "steps", "gen1", max_workers=2, tmp_root=tmp_root) + + assert all(r.ok for r in results) + assert {r.dest.name for r in results} == {f"shard-{i:03d}.root" for i in range(6)} + + intervals = [json.loads(r.dest.read_text()) for r in results] + events = sorted( + [(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals] + ) + concurrent = 0 + peak = 0 + for _, delta in events: + 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_loader.py b/tests/test_loader.py new file mode 100644 index 0000000..6c562f3 --- /dev/null +++ b/tests/test_loader.py @@ -0,0 +1,60 @@ +import pytest + +from giant.data.loader import find_parquet_files + + +def _touch(path): + path.parent.mkdir(parents=True, exist_ok=True) + path.touch() + return path + + +def test_find_parquet_files_single_file(tmp_path): + f = _touch(tmp_path / "shard-000.parquet") + assert find_parquet_files(f) == [f] + + +def test_find_parquet_files_directory_glob(tmp_path): + a = _touch(tmp_path / "shard-000.parquet") + b = _touch(tmp_path / "shard-001.parquet") + _touch(tmp_path / "not_a_parquet.root") + assert find_parquet_files(tmp_path) == sorted([a, b]) + + +def test_find_parquet_files_empty_directory_raises(tmp_path): + with pytest.raises(FileNotFoundError): + find_parquet_files(tmp_path) + + +def test_manifest_resolves_relative_to_its_own_directory(tmp_path): + target = _touch(tmp_path / "processed" / "pbwo4" / "shard-000.parquet") + manifest_dir = tmp_path / "pools" / "pbwo4" + manifest_dir.mkdir(parents=True) + manifest = manifest_dir / "full.manifest" + manifest.write_text("../../processed/pbwo4/shard-000.parquet\n") + + assert find_parquet_files(manifest) == [target.resolve()] + + +def test_manifest_skips_blank_lines_and_comments(tmp_path): + target = _touch(tmp_path / "shard-000.parquet") + manifest = tmp_path / "full.manifest" + manifest.write_text("\n# a comment\nshard-000.parquet\n\n") + + assert find_parquet_files(manifest) == [target.resolve()] + + +def test_manifest_missing_file_raises(tmp_path): + manifest = tmp_path / "full.manifest" + manifest.write_text("does_not_exist.parquet\n") + + with pytest.raises(FileNotFoundError): + find_parquet_files(manifest) + + +def test_manifest_with_no_entries_raises(tmp_path): + manifest = tmp_path / "full.manifest" + manifest.write_text("# only comments\n") + + with pytest.raises(FileNotFoundError): + find_parquet_files(manifest) diff --git a/tests/test_steps_to_parquet_parallel.py b/tests/test_steps_to_parquet_parallel.py new file mode 100644 index 0000000..2daedf8 --- /dev/null +++ b/tests/test_steps_to_parquet_parallel.py @@ -0,0 +1,186 @@ +import importlib.util +import json +from pathlib import Path + +# scripts/ is not an installed package — load the module straight from its path. +_SPEC = importlib.util.spec_from_file_location( + "steps_to_parquet_parallel", + Path(__file__).resolve().parents[1] / "scripts" / "steps_to_parquet_parallel.py", +) +steps_to_parquet_parallel = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(steps_to_parquet_parallel) + +run_parallel = steps_to_parquet_parallel.run_parallel +resolve_destination = steps_to_parquet_parallel.resolve_destination +latest_schema_tag = steps_to_parquet_parallel.latest_schema_tag +DestinationError = steps_to_parquet_parallel.DestinationError + + +def _write_fake_executable(tmp_path: Path, marker_dir: Path) -> Path: + """A stand-in for steps_to_parquet.py: records its own start/end time per + input file (so tests can check overlap), then exits 1 if "fail" is in the + filename, else 0. Accepts (and ignores) the same flags as the real script. + """ + fake = tmp_path / "fake_steps_to_parquet.py" + fake.write_text( + f""" +import json, sys, time +from pathlib import Path + +marker_dir = Path({str(marker_dir)!r}) +root_file = sys.argv[1] +name = Path(root_file).stem +start = time.time() +time.sleep(0.2) +end = time.time() +(marker_dir / f"{{name}}.json").write_text(json.dumps({{"start": start, "end": end}})) +print(f"converted {{root_file}}") +sys.exit(1 if "fail" in name else 0) +""" + ) + return fake + + +def test_runs_one_job_per_file_and_reports_success(tmp_path): + marker_dir = tmp_path / "markers" + marker_dir.mkdir() + 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) + + assert {r[0] for r in results} == set(files) + assert all(code == 0 for _, code, _, _ in results) + assert all((marker_dir / f"{Path(f).stem}.json").exists() for f in files) + + +def test_failures_are_reported_with_nonzero_exit_code(tmp_path): + marker_dir = tmp_path / "markers" + marker_dir.mkdir() + 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) + + codes = {Path(f).stem: code for f, code, _, _ in results} + assert codes["shard-000"] == 0 + assert codes["shard-fail"] == 1 + + +def test_jobs_caps_concurrency(tmp_path): + marker_dir = tmp_path / "markers" + marker_dir.mkdir() + 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) + + intervals = [] + for f in files: + data = json.loads((marker_dir / f"{Path(f).stem}.json").read_text()) + intervals.append((data["start"], data["end"])) + + # max number of intervals overlapping any given instant must not exceed jobs + events = sorted([(s, 1) for s, _ in intervals] + [(e, -1) for _, e in intervals]) + concurrent = 0 + peak = 0 + for _, delta in events: + concurrent += delta + peak = max(peak, concurrent) + 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() + fake = tmp_path / "fake_steps_to_parquet.py" + fake.write_text( + f""" +import sys +from pathlib import Path + +marker_dir = Path({str(marker_dir)!r}) +root_file = sys.argv[1] +name = Path(root_file).stem +output = sys.argv[sys.argv.index("--output") + 1] if "--output" in sys.argv else "" +(marker_dir / f"{{name}}.txt").write_text(output) +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} + ) + assert (marker_dir / "shard-000.txt").read_text() == str(dest) + + +def _make_dataset(tmp_path: Path, schemas: list[str] | None = None) -> Path: + root_file = tmp_path / "raw" / "steps" / "gen1" / "pbwo4" / "shard-000.root" + root_file.parent.mkdir(parents=True) + root_file.touch() + for schema in schemas or []: + (tmp_path / "processed" / "steps" / "gen1" / schema).mkdir(parents=True) + return root_file + + +def test_resolve_destination_uses_latest_schema(tmp_path): + root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3", "schema2"]) + dest = resolve_destination(root_file, tmp_path, schema_override=None) + assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet" + + +def test_resolve_destination_schema_override_wins(tmp_path): + root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3"]) + dest = resolve_destination(root_file, tmp_path, schema_override="schema9") + assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema9" / "pbwo4" / "shard-000.parquet" + + +def test_resolve_destination_errors_without_any_schema(tmp_path): + root_file = _make_dataset(tmp_path, schemas=[]) + try: + resolve_destination(root_file, tmp_path, schema_override=None) + assert False, "expected DestinationError" + except DestinationError: + pass + + +def test_resolve_destination_errors_outside_dataset_root(tmp_path): + other_root = tmp_path / "other" + other_root.mkdir() + stray_file = other_root / "shard-000.root" + stray_file.touch() + dataset_root = tmp_path / "dataset" + dataset_root.mkdir() + try: + resolve_destination(stray_file, dataset_root, schema_override=None) + assert False, "expected DestinationError" + except DestinationError: + pass + + +def test_resolve_destination_errors_on_wrong_shape(tmp_path): + # missing the segment + bad_file = tmp_path / "raw" / "steps" / "pbwo4" / "shard-000.root" + bad_file.parent.mkdir(parents=True) + bad_file.touch() + try: + resolve_destination(bad_file, tmp_path, schema_override=None) + assert False, "expected DestinationError" + except DestinationError: + pass + + +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