lars
091b23a60f
bump-gen and bump-schema now accept --to genN/schemaN to target a specific version instead of always auto-incrementing. update-manifest gains --gen genN to repoint the gen component of manifest paths (combinable with --schema). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
511 lines
19 KiB
Python
511 lines
19 KiB
Python
#!/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
|
|
pools/<detector>/<pool>.manifest
|
|
|
|
`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-gen --kind steps --to gen5 --reason "skip to gen5"
|
|
bump_dataset_version.py bump-schema --kind steps --gen gen1 --reason "added e_sec column"
|
|
bump_dataset_version.py bump-schema --kind steps --gen gen1 --to schema3 --reason "skip to schema3"
|
|
bump_dataset_version.py update-manifest pools/pbwo4/full.manifest [--schema schema2] [--gen gen2]
|
|
bump_dataset_version.py create-manifest --output pools/pbwo4/train.manifest a.parquet b.parquet
|
|
bump_dataset_version.py status
|
|
"""
|
|
|
|
import argparse
|
|
import datetime as dt
|
|
import os
|
|
import re
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
GEN_RE = re.compile(r"^gen(\d+)$")
|
|
SCHEMA_RE = re.compile(r"^schema(\d+)$")
|
|
|
|
|
|
def _find_component_idx(abs_path: Path, pattern: re.Pattern) -> int | None:
|
|
"""Return the index of the first path component matching *pattern*, or None."""
|
|
for i, part in enumerate(abs_path.parts):
|
|
if pattern.match(part):
|
|
return i
|
|
return None
|
|
|
|
|
|
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,
|
|
target: str | None = None,
|
|
) -> tuple[list[Path], str]:
|
|
"""New gen tag is one past the highest seen under raw/ or processed/ for *kind*,
|
|
or *target* if explicitly provided."""
|
|
if target is not None:
|
|
if not GEN_RE.match(target):
|
|
raise SystemExit(f"error: --to must look like 'genN', got {target!r}")
|
|
gen_tag = target
|
|
else:
|
|
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,
|
|
target: str | None = None,
|
|
) -> 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"
|
|
)
|
|
if target is not None:
|
|
if not SCHEMA_RE.match(target):
|
|
raise SystemExit(f"error: --to must look like 'schemaN', got {target!r}")
|
|
schema_tag = target
|
|
else:
|
|
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}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# update-manifest
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def plan_update_manifest(
|
|
manifest_path: Path,
|
|
target_schema: str | None,
|
|
target_gen: str | None = None,
|
|
) -> tuple[list[tuple[str, str | None]], list[Path]]:
|
|
"""Parse a manifest and plan gen/schema replacements for each data line.
|
|
|
|
Returns:
|
|
lines: list of (original_line, new_relative_path_or_None)
|
|
None means the line is unchanged (comment, blank, or already at target).
|
|
missing: resolved absolute paths that don't exist on disk.
|
|
"""
|
|
manifest_path = manifest_path.resolve()
|
|
if not manifest_path.exists():
|
|
raise SystemExit(f"error: manifest not found: {manifest_path}")
|
|
manifest_dir = manifest_path.parent
|
|
raw_lines = manifest_path.read_text().splitlines()
|
|
|
|
result: list[tuple[str, str | None]] = []
|
|
missing: list[Path] = []
|
|
schema_cache: dict[Path, str] = {}
|
|
|
|
for raw in raw_lines:
|
|
stripped = raw.strip()
|
|
if not stripped or stripped.startswith("#"):
|
|
result.append((raw, None))
|
|
continue
|
|
|
|
old_abs = (manifest_dir / stripped).resolve()
|
|
parts = list(old_abs.parts)
|
|
changed = False
|
|
|
|
# Apply gen replacement first (shifts subsequent indices).
|
|
if target_gen is not None:
|
|
gen_idx = _find_component_idx(Path(*parts), GEN_RE)
|
|
if gen_idx is not None and parts[gen_idx] != target_gen:
|
|
parts[gen_idx] = target_gen
|
|
changed = True
|
|
|
|
schema_idx = _find_component_idx(Path(*parts), SCHEMA_RE)
|
|
if schema_idx is not None:
|
|
old_schema = parts[schema_idx]
|
|
gen_dir = Path(*parts[:schema_idx])
|
|
|
|
if target_schema is not None:
|
|
new_schema = target_schema
|
|
else:
|
|
if gen_dir not in schema_cache:
|
|
n = _max_index(gen_dir, SCHEMA_RE)
|
|
if n == 0:
|
|
raise SystemExit(f"error: no schema dirs found under {gen_dir}")
|
|
schema_cache[gen_dir] = f"schema{n}"
|
|
new_schema = schema_cache[gen_dir]
|
|
|
|
if new_schema != old_schema:
|
|
parts[schema_idx] = new_schema
|
|
changed = True
|
|
|
|
if not changed:
|
|
result.append((raw, None))
|
|
continue
|
|
|
|
new_abs = Path(*parts)
|
|
if not new_abs.exists():
|
|
missing.append(new_abs)
|
|
|
|
new_rel = os.path.relpath(new_abs, start=manifest_dir)
|
|
result.append((raw, new_rel))
|
|
|
|
return result, missing
|
|
|
|
|
|
def apply_update_manifest(manifest_path: Path, lines: list[tuple[str, str | None]]) -> None:
|
|
out = [replacement if replacement is not None else original for original, replacement in lines]
|
|
manifest_path.write_text("\n".join(out) + "\n")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# create-manifest
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _resolve_manifest_files(manifest_path: Path) -> list[Path]:
|
|
"""Read a manifest and return its entries as resolved absolute paths."""
|
|
files = []
|
|
for line in manifest_path.read_text().splitlines():
|
|
line = line.strip()
|
|
if not line or line.startswith("#"):
|
|
continue
|
|
files.append((manifest_path.parent / line).resolve())
|
|
return files
|
|
|
|
|
|
def plan_create_manifest(
|
|
output_path: Path, parquet_files: list[Path]
|
|
) -> tuple[list[str], list[Path], list[Path]]:
|
|
"""Return (relative_lines, missing_files, resolved_abs_paths)."""
|
|
manifest_dir = output_path.resolve().parent
|
|
lines: list[str] = []
|
|
missing: list[Path] = []
|
|
resolved: list[Path] = []
|
|
for p in parquet_files:
|
|
abs_p = p.resolve()
|
|
resolved.append(abs_p)
|
|
if not abs_p.exists():
|
|
missing.append(abs_p)
|
|
lines.append(os.path.relpath(abs_p, start=manifest_dir))
|
|
return lines, missing, resolved
|
|
|
|
|
|
def check_holdout_overlap(
|
|
output_path: Path, resolved_new_files: list[Path]
|
|
) -> list[tuple[str, Path]]:
|
|
"""Return (other_manifest_name, file) pairs where new files clash with existing manifests.
|
|
|
|
The check is triggered when output_path is (or will be) holdout.manifest, or when a
|
|
holdout.manifest already exists in the same directory — in either case holdout data
|
|
must be strictly isolated from all other pools.
|
|
"""
|
|
output_resolved = output_path.resolve()
|
|
manifest_dir = output_resolved.parent
|
|
holdout_path = manifest_dir / "holdout.manifest"
|
|
|
|
if output_resolved.name != "holdout.manifest" and not holdout_path.exists():
|
|
return []
|
|
|
|
new_set = set(resolved_new_files)
|
|
overlaps: list[tuple[str, Path]] = []
|
|
# When creating holdout, check against all other manifests (dev, full, …).
|
|
# When creating dev/full, only check against holdout — dev vs full overlap is allowed.
|
|
if output_resolved.name == "holdout.manifest":
|
|
candidates = sorted(manifest_dir.glob("*.manifest"))
|
|
else:
|
|
candidates = [holdout_path]
|
|
for existing in candidates:
|
|
if existing.resolve() == output_resolved:
|
|
continue
|
|
try:
|
|
existing_files = set(_resolve_manifest_files(existing))
|
|
except OSError:
|
|
continue
|
|
for f in sorted(new_set & existing_files):
|
|
overlaps.append((existing.name, f))
|
|
return overlaps
|
|
|
|
|
|
def apply_create_manifest(output_path: Path, lines: list[str]) -> None:
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
output_path.write_text("\n".join(lines) + "\n")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--root",
|
|
default="/ceph/lbogner/geant_steps",
|
|
help="Dataset root for bump-gen/bump-schema/status (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(
|
|
"--to", default=None, metavar="genN",
|
|
help="Target gen tag (default: one past the current highest)",
|
|
)
|
|
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(
|
|
"--to", default=None, metavar="schemaN",
|
|
help="Target schema tag (default: one past the current highest)",
|
|
)
|
|
p_schema.add_argument("--execute", action="store_true", help="Apply (default: dry run)")
|
|
|
|
p_update = sub.add_parser(
|
|
"update-manifest",
|
|
help="Repoint manifest(s) to a new gen and/or schema, verifying all target files exist",
|
|
)
|
|
p_update.add_argument(
|
|
"manifests", nargs="+", metavar="MANIFEST", help="One or more .manifest files to update"
|
|
)
|
|
p_update.add_argument(
|
|
"--schema",
|
|
default=None,
|
|
metavar="schemaN",
|
|
help="Target schema tag (default: highest schema found in the same gen dir)",
|
|
)
|
|
p_update.add_argument(
|
|
"--gen",
|
|
default=None,
|
|
metavar="genN",
|
|
help="Target gen tag (default: keep existing gen)",
|
|
)
|
|
p_update.add_argument("--execute", action="store_true", help="Write updated manifests (default: dry run)")
|
|
|
|
p_create = sub.add_parser(
|
|
"create-manifest",
|
|
help="Create a new manifest from a list of parquet files",
|
|
)
|
|
dest_group = p_create.add_mutually_exclusive_group(required=True)
|
|
dest_group.add_argument(
|
|
"--output", "-o", metavar="PATH", help="Explicit path for the new .manifest file"
|
|
)
|
|
dest_group.add_argument(
|
|
"--pool", metavar="DETECTOR",
|
|
help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
|
|
)
|
|
p_create.add_argument(
|
|
"--type", choices=["full", "holdout", "dev"],
|
|
help="Pool type — full, holdout, or dev (required with --pool)",
|
|
)
|
|
p_create.add_argument(
|
|
"files", nargs="+", metavar="FILE", help="Parquet files to include"
|
|
)
|
|
p_create.add_argument("--execute", action="store_true", help="Write the manifest (default: dry run)")
|
|
|
|
sub.add_parser("status", help="List existing gens/schemas per kind")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Commands that need --root
|
|
if args.command in ("bump-gen", "bump-schema", "status"):
|
|
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
|
|
|
|
if args.command in ("bump-gen", "bump-schema"):
|
|
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, args.to)
|
|
else:
|
|
new_dirs, log_line = plan_bump_schema(root, args.kind, args.gen, args.reason, by, date, args.to)
|
|
|
|
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.")
|
|
return
|
|
|
|
if args.command == "update-manifest":
|
|
if args.schema and not SCHEMA_RE.match(args.schema):
|
|
parser.error(f"--schema must look like 'schemaN', got {args.schema!r}")
|
|
if args.gen and not GEN_RE.match(args.gen):
|
|
parser.error(f"--gen must look like 'genN', got {args.gen!r}")
|
|
|
|
all_plans: list[tuple[Path, list[tuple[str, str | None]]]] = []
|
|
all_missing: list[Path] = []
|
|
|
|
for raw in args.manifests:
|
|
mp = Path(raw)
|
|
plan, missing = plan_update_manifest(mp, args.schema, args.gen)
|
|
all_plans.append((mp.resolve(), plan))
|
|
all_missing.extend(missing)
|
|
|
|
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
|
|
for mp, plan in all_plans:
|
|
changes = [(old, new) for old, new in plan if new is not None]
|
|
print(f"\n{mp} ({len(changes)} path(s) to update)")
|
|
for old, new in changes:
|
|
print(f" - {old.strip()}")
|
|
print(f" + {new}")
|
|
|
|
if all_missing:
|
|
print(f"\nMISSING ({len(all_missing)} file(s) — target paths do not exist):")
|
|
for p in all_missing:
|
|
print(f" {p}")
|
|
if args.execute:
|
|
raise SystemExit("error: refusing to write manifests with missing targets")
|
|
|
|
if not args.execute:
|
|
print("\nDry run only — pass --execute to apply.")
|
|
return
|
|
|
|
for mp, plan in all_plans:
|
|
apply_update_manifest(mp, plan)
|
|
print("\nDone.")
|
|
return
|
|
|
|
if args.command == "create-manifest":
|
|
if args.pool is not None and args.type is None:
|
|
parser.error("--type is required when --pool is given")
|
|
|
|
if args.pool is not None:
|
|
root = Path(args.root)
|
|
output_path = root / "pools" / args.pool / f"{args.type}.manifest"
|
|
else:
|
|
output_path = Path(args.output)
|
|
|
|
parquet_files = [Path(f) for f in args.files]
|
|
lines, missing, resolved = plan_create_manifest(output_path, parquet_files)
|
|
overlaps = check_holdout_overlap(output_path, resolved)
|
|
|
|
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
|
|
print(f"manifest: {output_path.resolve()}")
|
|
for line in lines:
|
|
print(f" {line}")
|
|
|
|
if missing:
|
|
print(f"\nMISSING ({len(missing)} file(s) do not exist):")
|
|
for p in missing:
|
|
print(f" {p}")
|
|
|
|
if overlaps:
|
|
print(f"\nHOLDOUT OVERLAP ({len(overlaps)} file(s) appear in other manifests):")
|
|
for name, f in overlaps:
|
|
print(f" {f} (also in {name})")
|
|
|
|
if (missing or overlaps) and args.execute:
|
|
raise SystemExit("error: refusing to write manifest (see above)")
|
|
|
|
if not args.execute:
|
|
print("\nDry run only — pass --execute to apply.")
|
|
return
|
|
|
|
apply_create_manifest(output_path, lines)
|
|
print("\nDone.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|