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:
@@ -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()
|
||||
Reference in New Issue
Block a user