Add tooling for a versioned geant_steps dataset layout
Introduces raw/<kind>/<gen>/<detector>/shard-NNN.root and processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet as the dataset convention, plus scripts to operate on it: migrate_geant_steps.py for the one-time move into this layout, bump_dataset_version.py to cut new gen/schema versions with a logged reason, steps_to_parquet_parallel.py to convert ROOT shards to parquet in parallel and place them correctly, and create_root_files.py to generate new ROOT shards via a minicalosim executable. The loader gains .manifest file support so pools/ (train/dev/ holdout shard lists) can be passed straight to `giant train`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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/<kind>/<gen>/<detector>/shard-NNN.root
|
||||
processed/<kind>/<gen>/<schema>/<detector>/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()
|
||||
@@ -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/<kind>/<gen>/<detector>/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_<n>events_hits.root, run_sampling writes sampling_<config>_<n>events_hits.root,
|
||||
others may differ again). This script gives each run its own scratch
|
||||
directory under <dataset-root>/.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/<kind>/ — 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/<kind>/, e.g. gen1")
|
||||
parser.add_argument(
|
||||
"--dataset-root",
|
||||
default="/ceph/lbogner/geant_steps",
|
||||
help="Dataset root (default: /ceph/lbogner/geant_steps)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-j", "--jobs", type=int, default=4, help="Parallel simulation runs (default: 4)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--execute", action="store_true", help="Actually run jobs (default: dry run / print plan)"
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def 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()
|
||||
@@ -0,0 +1,258 @@
|
||||
#!/usr/bin/env python3
|
||||
"""One-time migration of /ceph/lbogner/geant_steps into the versioned layout:
|
||||
|
||||
raw/<kind>/<gen>/<detector>/shard-NNN.root
|
||||
processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet
|
||||
pools/<detector>/{dev,full,holdout,unclassified}.txt (manifests, no copied bytes)
|
||||
derived/predictions/<bucket>/<detector>_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<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)"
|
||||
r"_predicted(?P<local>_local)?\.parquet$"
|
||||
)
|
||||
SHARD_RE = re.compile(
|
||||
r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$"
|
||||
)
|
||||
LEGACY_PREDICTED_RE = re.compile(
|
||||
r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$"
|
||||
)
|
||||
LEGACY_RE = re.compile(r"^pbwo4_10000events_hits\.(?P<ext>root|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()
|
||||
@@ -0,0 +1,251 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Convert many ROOT files to Parquet by fanning out to steps_to_parquet.py.
|
||||
|
||||
steps_to_parquet.py itself converts a list of files one at a time; this wraps
|
||||
it to run up to --jobs conversions concurrently, each as its own subprocess
|
||||
(invoked with the same Python executable running this script, so it picks up
|
||||
the active venv/uv environment automatically).
|
||||
|
||||
Inputs must live under <dataset-root>/raw/<kind>/<gen>/<detector>/<file>.root
|
||||
(see scripts/migrate_geant_steps.py) — each is written to the matching
|
||||
processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet, where <schema>
|
||||
defaults to the highest schemaN already under processed/<kind>/<gen>/ (pass
|
||||
--schema to pick a specific one, e.g. one just created by
|
||||
bump_dataset_version.py bump-schema). A file that doesn't fit that layout is
|
||||
rejected up front, before any conversion runs.
|
||||
|
||||
Usage:
|
||||
steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/shard-000.root raw/steps/gen1/pbwo4/shard-001.root
|
||||
steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --jobs 8
|
||||
steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --schema schema2
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
|
||||
_STEPS_TO_PARQUET = Path(__file__).resolve().parent / "steps_to_parquet.py"
|
||||
|
||||
# Must match scripts/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
|
||||
GEN_RE = re.compile(r"^gen\d+$")
|
||||
SCHEMA_RE = re.compile(r"^schema(\d+)$")
|
||||
|
||||
|
||||
class DestinationError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
def latest_schema_tag(processed_gen_dir: Path) -> str | None:
|
||||
"""Highest schemaN dir directly under *processed_gen_dir*, or None if none exist."""
|
||||
if not processed_gen_dir.is_dir():
|
||||
return None
|
||||
best_tag, best_n = None, -1
|
||||
for child in processed_gen_dir.iterdir():
|
||||
m = SCHEMA_RE.match(child.name)
|
||||
if m and child.is_dir() and int(m.group(1)) > best_n:
|
||||
best_tag, best_n = child.name, int(m.group(1))
|
||||
return best_tag
|
||||
|
||||
|
||||
def resolve_destination(
|
||||
root_file: Path, dataset_root: Path, schema_override: str | None
|
||||
) -> Path:
|
||||
"""Map raw/<kind>/<gen>/<detector>/<file>.root (relative to *dataset_root*)
|
||||
to processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet.
|
||||
|
||||
Raises DestinationError if *root_file* doesn't fit that layout, or if no
|
||||
schema can be determined (no --schema and none exists yet for that gen).
|
||||
"""
|
||||
dataset_root = dataset_root.resolve()
|
||||
root_file = root_file.resolve()
|
||||
try:
|
||||
rel = root_file.relative_to(dataset_root)
|
||||
except ValueError:
|
||||
raise DestinationError(f"{root_file} is not under dataset root {dataset_root}")
|
||||
|
||||
parts = rel.parts
|
||||
if (
|
||||
len(parts) != 5
|
||||
or parts[0] != "raw"
|
||||
or not GEN_RE.match(parts[2])
|
||||
or not parts[4].endswith(".root")
|
||||
):
|
||||
raise DestinationError(
|
||||
f"{root_file} does not match raw/<kind>/<gen>/<detector>/<file>.root "
|
||||
f"under {dataset_root} (got relative path: {rel})"
|
||||
)
|
||||
_, kind, gen_tag, detector, filename = parts
|
||||
shard_stem = Path(filename).stem
|
||||
|
||||
processed_gen_dir = dataset_root / "processed" / kind / gen_tag
|
||||
schema_tag = schema_override or latest_schema_tag(processed_gen_dir)
|
||||
if schema_tag is None:
|
||||
raise DestinationError(
|
||||
f"no schema exists yet under {processed_gen_dir} — pass --schema or run "
|
||||
"bump_dataset_version.py bump-schema first"
|
||||
)
|
||||
return processed_gen_dir / schema_tag / detector / f"{shard_stem}.parquet"
|
||||
|
||||
|
||||
def _convert_one(
|
||||
root_file: str,
|
||||
batch_size: str,
|
||||
tree: str,
|
||||
compression: str,
|
||||
steps_to_parquet_path: Path,
|
||||
output_path: Path | None,
|
||||
) -> tuple[str, int, str, str]:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(steps_to_parquet_path),
|
||||
root_file,
|
||||
"--batch-size",
|
||||
batch_size,
|
||||
"--tree",
|
||||
tree,
|
||||
"--compression",
|
||||
compression,
|
||||
]
|
||||
if output_path is not None:
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
cmd += ["--output", str(output_path)]
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
return root_file, result.returncode, result.stdout, result.stderr
|
||||
|
||||
|
||||
def run_parallel(
|
||||
root_files: list[str],
|
||||
jobs: int,
|
||||
batch_size: str = "100 MB",
|
||||
tree: str = "Steps",
|
||||
compression: str = "snappy",
|
||||
steps_to_parquet_path: Path = _STEPS_TO_PARQUET,
|
||||
output_for: dict[str, Path] | None = None,
|
||||
) -> list[tuple[str, int, str, str]]:
|
||||
"""Run one steps_to_parquet.py subprocess per file, up to *jobs* at a time.
|
||||
|
||||
*output_for*, if given, maps each root_file to the parquet path it should
|
||||
be written to (passed through as steps_to_parquet.py's --output); files
|
||||
missing from the map fall back to steps_to_parquet.py's own default
|
||||
(parquet written next to the input .root).
|
||||
|
||||
Returns one (root_file, returncode, stdout, stderr) tuple per file, in
|
||||
completion order (not necessarily input order).
|
||||
"""
|
||||
output_for = output_for or {}
|
||||
results = []
|
||||
with ThreadPoolExecutor(max_workers=jobs) as pool:
|
||||
futures = {
|
||||
pool.submit(
|
||||
_convert_one,
|
||||
f,
|
||||
batch_size,
|
||||
tree,
|
||||
compression,
|
||||
steps_to_parquet_path,
|
||||
output_for.get(f),
|
||||
): f
|
||||
for f in root_files
|
||||
}
|
||||
for future in as_completed(futures):
|
||||
root_file, code, stdout, stderr = future.result()
|
||||
status = "ok" if code == 0 else f"FAILED (exit {code})"
|
||||
print(f"\n=== {root_file}: {status} ===")
|
||||
if stdout:
|
||||
print(stdout, end="")
|
||||
if stderr:
|
||||
print(stderr, end="", file=sys.stderr)
|
||||
results.append((root_file, code, stdout, stderr))
|
||||
return results
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert many ROOT files to Parquet in parallel via steps_to_parquet.py."
|
||||
)
|
||||
parser.add_argument("root_files", nargs="+", help="Input ROOT file(s)")
|
||||
parser.add_argument(
|
||||
"-j",
|
||||
"--jobs",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Number of conversions to run in parallel (default: 4)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-size",
|
||||
default="100 MB",
|
||||
help="Uproot read batch size (default: '100 MB'). E.g. '50 MB', '500000' (rows).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tree",
|
||||
default="Steps",
|
||||
help="Tree name inside the ROOT file (default: Steps)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--compression",
|
||||
default="snappy",
|
||||
choices=["snappy", "lz4", "zstd", "gzip", "none"],
|
||||
help="Parquet compression codec (default: snappy)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset-root",
|
||||
default="/ceph/lbogner/geant_steps",
|
||||
help="Dataset root containing raw/ and processed/ (default: /ceph/lbogner/geant_steps)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--schema",
|
||||
default=None,
|
||||
help="Schema tag to write parquets under, e.g. schema2 "
|
||||
"(default: highest schemaN already under processed/<kind>/<gen>/)",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.jobs < 1:
|
||||
parser.error("--jobs must be >= 1")
|
||||
|
||||
dataset_root = Path(args.dataset_root)
|
||||
output_for: dict[str, Path] = {}
|
||||
errors: list[str] = []
|
||||
for f in args.root_files:
|
||||
try:
|
||||
output_for[f] = resolve_destination(Path(f), dataset_root, args.schema)
|
||||
except DestinationError as exc:
|
||||
errors.append(str(exc))
|
||||
|
||||
if errors:
|
||||
for err in errors:
|
||||
print(f"error: {err}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
for root_file, dest in output_for.items():
|
||||
print(f"{root_file} -> {dest}")
|
||||
|
||||
results = run_parallel(
|
||||
args.root_files,
|
||||
jobs=args.jobs,
|
||||
batch_size=args.batch_size,
|
||||
tree=args.tree,
|
||||
compression=args.compression,
|
||||
output_for=output_for,
|
||||
)
|
||||
|
||||
failures = [root_file for root_file, code, _, _ in results if code != 0]
|
||||
if failures:
|
||||
print(f"\n{len(failures)} of {len(results)} conversion(s) failed:", file=sys.stderr)
|
||||
for root_file in failures:
|
||||
print(f" {root_file}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"\nAll {len(results)} conversion(s) completed.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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 <gen> 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
|
||||
Reference in New Issue
Block a user