diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py index 958728b..befc90a 100644 --- a/scripts/bump_dataset_version.py +++ b/scripts/bump_dataset_version.py @@ -4,6 +4,7 @@ dataset tree (see scripts/migrate_geant_steps.py for the layout): raw////shard-NNN.root processed/////shard-NNN.parquet + pools//.manifest `gen` bumps when the underlying ROOT changes (geometry/physics-list/macro). `schema` bumps when the parquet export (steps_to_parquet.py or similar) @@ -15,11 +16,14 @@ 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 update-manifest pools/pbwo4/full.manifest [--schema schema2] + 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 @@ -131,12 +135,161 @@ def print_status(root: Path) -> None: print(f" {gen_tag}: {schema_str}") +# --------------------------------------------------------------------------- +# update-manifest +# --------------------------------------------------------------------------- + +def _find_schema_idx(abs_path: Path) -> int | None: + """Return the index of the first schemaN component in abs_path.parts, or None.""" + for i, part in enumerate(abs_path.parts): + if SCHEMA_RE.match(part): + return i + return None + + +def plan_update_manifest( + manifest_path: Path, target_schema: str | None +) -> tuple[list[tuple[str, str | None]], list[Path]]: + """Parse a manifest and plan 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() + schema_idx = _find_schema_idx(old_abs) + if schema_idx is None: + result.append((raw, None)) + continue + + parts = list(old_abs.parts) + 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: + result.append((raw, None)) + continue + + parts[schema_idx] = new_schema + 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]] = [] + for existing in sorted(manifest_dir.glob("*.manifest")): + 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 (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) @@ -157,40 +310,153 @@ def main() -> None: 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)") + p_update = sub.add_parser( + "update-manifest", + help="Repoint manifest(s) to a new 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("--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 /pools//.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() - root = Path(args.root) - if not root.is_dir(): - parser.error(f"{root} is not a directory") + + # 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 - date = args.date or dt.date.today().isoformat() - by = args.by if args.by is not None else _git_user_name() + 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) + else: + new_dirs, log_line = plan_bump_schema(root, args.kind, args.gen, args.reason, by, date) - 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}") - 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.") + if not args.execute: + print("\nDry run only — pass --execute to apply.") + return + apply_bump(root, new_dirs, log_line) + print("\nDone.") return - apply_bump(root, new_dirs, log_line) - print("\nDone.") + 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}") + + 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) + 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__": diff --git a/tests/test_bump_dataset_version.py b/tests/test_bump_dataset_version.py index 62c0f36..948e5f8 100644 --- a/tests/test_bump_dataset_version.py +++ b/tests/test_bump_dataset_version.py @@ -1,8 +1,15 @@ +import os +import pytest from scripts import 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 +plan_update_manifest = bump_dataset_version.plan_update_manifest +apply_update_manifest = bump_dataset_version.apply_update_manifest +plan_create_manifest = bump_dataset_version.plan_create_manifest +apply_create_manifest = bump_dataset_version.apply_create_manifest +check_holdout_overlap = bump_dataset_version.check_holdout_overlap def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path): @@ -84,3 +91,219 @@ def test_apply_bump_appends_without_clobbering_existing_log(tmp_path): text = (tmp_path / "VERSIONS.md").read_text() assert "existing entry" in text assert "reason B" in text + + +# --------------------------------------------------------------------------- +# update-manifest +# --------------------------------------------------------------------------- + +def _make_parquet(path): + """Create a zero-byte stand-in for a parquet file.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.touch() + + +def test_update_manifest_bumps_to_specified_schema(tmp_path): + parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet" + _make_parquet(parquet) + + manifest_dir = tmp_path / "pools" / "pbwo4" + manifest_dir.mkdir(parents=True) + manifest = manifest_dir / "full.manifest" + old_rel = "../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet" + manifest.write_text(old_rel + "\n") + + lines, missing = plan_update_manifest(manifest, "schema2") + assert missing == [] + # exactly one data line was updated + changed = [(old, new) for old, new in lines if new is not None] + assert len(changed) == 1 + assert "schema2" in changed[0][1] + assert "schema1" not in changed[0][1] + + +def test_update_manifest_auto_detects_highest_schema(tmp_path): + for schema in ("schema1", "schema2", "schema3"): + d = tmp_path / "processed" / "steps" / "gen1" / schema / "pbwo4" + d.mkdir(parents=True) + parquet = tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet" + parquet.touch() + + manifest_dir = tmp_path / "pools" / "pbwo4" + manifest_dir.mkdir(parents=True) + manifest = manifest_dir / "full.manifest" + manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n") + + lines, missing = plan_update_manifest(manifest, None) + assert missing == [] + changed = [(old, new) for old, new in lines if new is not None] + assert "schema3" in changed[0][1] + + +def test_update_manifest_reports_missing_targets(tmp_path): + manifest_dir = tmp_path / "pools" / "pbwo4" + manifest_dir.mkdir(parents=True) + manifest = manifest_dir / "full.manifest" + manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n") + + # schema2 dir exists but the parquet file does not + (tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True) + + lines, missing = plan_update_manifest(manifest, "schema2") + assert len(missing) == 1 + assert "schema2" in str(missing[0]) + + +def test_update_manifest_skips_already_at_target(tmp_path): + parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet" + _make_parquet(parquet) + + manifest_dir = tmp_path / "pools" / "pbwo4" + manifest_dir.mkdir(parents=True) + manifest = manifest_dir / "full.manifest" + manifest.write_text("../../processed/steps/gen1/schema2/pbwo4/shard-000.parquet\n") + + lines, missing = plan_update_manifest(manifest, "schema2") + assert missing == [] + # line is unchanged — new is None + assert all(new is None for _, new in lines) + + +def test_update_manifest_preserves_comments_and_blanks(tmp_path): + parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet" + _make_parquet(parquet) + + manifest_dir = tmp_path / "pools" / "pbwo4" + manifest_dir.mkdir(parents=True) + manifest = manifest_dir / "full.manifest" + content = "# header\n\n../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n" + manifest.write_text(content) + + lines, _ = plan_update_manifest(manifest, "schema2") + assert lines[0] == ("# header", None) + assert lines[1] == ("", None) + assert lines[2][1] is not None # the data line was updated + + +def test_apply_update_manifest_writes_file(tmp_path): + parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet" + _make_parquet(parquet) + + manifest_dir = tmp_path / "pools" / "pbwo4" + manifest_dir.mkdir(parents=True) + manifest = manifest_dir / "full.manifest" + manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n") + + plan, _ = plan_update_manifest(manifest, "schema2") + apply_update_manifest(manifest, plan) + + written = manifest.read_text() + assert "schema2" in written + assert "schema1" not in written + + +# --------------------------------------------------------------------------- +# create-manifest +# --------------------------------------------------------------------------- + +def test_create_manifest_writes_relative_paths(tmp_path): + pq1 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet" + pq2 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-001.parquet" + _make_parquet(pq1) + _make_parquet(pq2) + + output = tmp_path / "pools" / "pbwo4" / "train.manifest" + lines, missing, resolved = plan_create_manifest(output, [pq1, pq2]) + + assert missing == [] + assert len(lines) == 2 + assert all("schema2" in l for l in lines) + assert all(not l.startswith("/") for l in lines) + assert resolved == [pq1.resolve(), pq2.resolve()] + + apply_create_manifest(output, lines) + assert output.exists() + written = output.read_text().strip().splitlines() + assert len(written) == 2 + + +def test_create_manifest_reports_missing_files(tmp_path): + ghost = tmp_path / "processed" / "gen1" / "schema2" / "shard-000.parquet" + output = tmp_path / "pools" / "full.manifest" + lines, missing, _ = plan_create_manifest(output, [ghost]) + assert len(missing) == 1 + assert missing[0] == ghost.resolve() + + +def test_create_manifest_creates_parent_dirs(tmp_path): + pq = tmp_path / "a.parquet" + pq.touch() + output = tmp_path / "deep" / "nested" / "pool.manifest" + lines, _, _ = plan_create_manifest(output, [pq]) + apply_create_manifest(output, lines) + assert output.exists() + + +# --------------------------------------------------------------------------- +# check_holdout_overlap +# --------------------------------------------------------------------------- + +def test_no_overlap_check_when_no_holdout_involved(tmp_path): + pool_dir = tmp_path / "pools" / "pbwo4" + pool_dir.mkdir(parents=True) + pq = tmp_path / "a.parquet" + pq.touch() + # writing full.manifest, no holdout.manifest exists + output = pool_dir / "full.manifest" + overlaps = check_holdout_overlap(output, [pq.resolve()]) + assert overlaps == [] + + +def test_overlap_detected_when_creating_holdout(tmp_path): + pool_dir = tmp_path / "pools" / "pbwo4" + pool_dir.mkdir(parents=True) + pq = tmp_path / "a.parquet" + pq.touch() + + # full.manifest already lists the same file + full = pool_dir / "full.manifest" + full.write_text(os.path.relpath(pq.resolve(), start=pool_dir) + "\n") + + output = pool_dir / "holdout.manifest" + overlaps = check_holdout_overlap(output, [pq.resolve()]) + assert len(overlaps) == 1 + assert overlaps[0][0] == "full.manifest" + assert overlaps[0][1] == pq.resolve() + + +def test_overlap_detected_when_holdout_already_exists(tmp_path): + pool_dir = tmp_path / "pools" / "pbwo4" + pool_dir.mkdir(parents=True) + pq = tmp_path / "a.parquet" + pq.touch() + + # holdout.manifest already lists the file + holdout = pool_dir / "holdout.manifest" + holdout.write_text(os.path.relpath(pq.resolve(), start=pool_dir) + "\n") + + # now creating full.manifest with the same file + output = pool_dir / "full.manifest" + overlaps = check_holdout_overlap(output, [pq.resolve()]) + assert len(overlaps) == 1 + assert overlaps[0][0] == "holdout.manifest" + + +def test_no_overlap_when_files_are_disjoint(tmp_path): + pool_dir = tmp_path / "pools" / "pbwo4" + pool_dir.mkdir(parents=True) + pq_holdout = tmp_path / "holdout.parquet" + pq_full = tmp_path / "full.parquet" + pq_holdout.touch() + pq_full.touch() + + holdout = pool_dir / "holdout.manifest" + holdout.write_text(os.path.relpath(pq_holdout.resolve(), start=pool_dir) + "\n") + + output = pool_dir / "full.manifest" + overlaps = check_holdout_overlap(output, [pq_full.resolve()]) + assert overlaps == []