"""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=True to actually move files and write manifests, and copy=True to copy instead of move, leaving the original files in place — e.g. if another process is still reading them from their current location. See `uv run dwarf migrate --help` for the CLI. """ import 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 run_migration(root: str, execute: bool, copy: bool) -> None: src_root = Path(root) if not src_root.is_dir(): raise SystemExit(f"error: {src_root} is not a directory") moves, unrecognized = plan_moves(src_root) manifests = plan_manifests(src_root) verb = "COPY" if copy else "MOVE" print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ({verb}): {src_root} ===\n") print(f"-- {len(moves)} file {'copy' if copy else '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 execute: print("\nDry run only — pass --execute to apply.") return transfer = shutil.copy2 if copy else shutil.move for src, dst in moves: if dst.exists(): raise FileExistsError(f"refusing to overwrite existing file: {dst}") dst.parent.mkdir(parents=True, exist_ok=True) transfer(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. In --copy mode the # originals are still there by design, so leave these alone entirely. if not copy: for stale_dir in ("train", "sampling_train/small", "sampling_train"): d = src_root / stale_dir if d.is_dir() and not any(d.iterdir()): d.rmdir() print(f"removed now-empty directory: {d.relative_to(src_root)}") print("\nDone.")