giant/scripts/migrate_geant_steps.py

"""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=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<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 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.")