"""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////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 (`dwarf convert` 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=True to apply. See `uv run dwarf bump-gen/bump-schema/update-manifest/create-manifest/status --help` for the CLI. """ import datetime as dt import os import re import subprocess import sys from pathlib import Path GEN_RE = re.compile(r"^gen(\d+)$") SCHEMA_RE = re.compile(r"^schema(\d+)$") # Match the log lines written by apply_bump()/plan_bump_gen()/plan_bump_schema(): # - `gen2` (kind=steps) — 2026-01-01 — reason text (by) # - `gen2`/`schema2` (kind=steps) — 2026-01-01 — reason text (by) VERSIONS_SCHEMA_LINE_RE = re.compile( r"^- `(?Pgen\d+)`/`(?Pschema\d+)` \(kind=(?P[\w-]+)\)" r" — \d{4}-\d{2}-\d{2} — (?P.+)$" ) VERSIONS_GEN_LINE_RE = re.compile( r"^- `(?Pgen\d+)` \(kind=(?P[\w-]+)\) — \d{4}-\d{2}-\d{2} — (?P.+)$" ) # One color per tree level in `dwarf status` output, so the eye can jump # straight to e.g. "all the schema rows" or "all the totals". _LEVEL_COLORS = { "kind": "\033[1;36m", # bold cyan — top-level kind/ header "gen": "\033[1;33m", # bold yellow — genN row + kind total "bucket": "\033[34m", # blue — raw/processed subtotals "schema": "\033[32m", # green — schemaN rows "root": "\033[1;35m", # bold magenta — derived/, pools/, grand total "reason": "\033[2m", # dim — VERSIONS.md reason extract under gen/schema rows } _RESET = "\033[0m" def _use_color() -> bool: return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None def _colorize(text: str, level: str) -> str: if not _use_color(): return text return f"{_LEVEL_COLORS[level]}{text}{_RESET}" 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 _du(path: Path) -> int: """Total size in bytes of all regular files under *path* (0 if missing).""" if not path.is_dir(): return 0 total = 0 for dirpath, _dirnames, filenames in os.walk(path): for name in filenames: fp = Path(dirpath) / name try: total += fp.stat().st_size except OSError: pass return total def _count_files(path: Path) -> int: """Total number of regular files under *path*, recursively (0 if missing).""" if not path.is_dir(): return 0 total = 0 for _dirpath, _dirnames, filenames in os.walk(path): total += len(filenames) return total # Must match giant.data.loader.MANIFEST_SUFFIX. MANIFEST_SUFFIX = ".manifest" def _manifest_referenced_files(pools_root: Path) -> set[Path]: """Resolved absolute paths of every file listed in any *.manifest under *pools_root*.""" referenced: set[Path] = set() if not pools_root.is_dir(): return referenced for manifest_path in pools_root.rglob(f"*{MANIFEST_SUFFIX}"): try: referenced.update(_resolve_manifest_files(manifest_path)) except OSError: continue return referenced def _referenced_root_count( raw_gen_dir: Path, processed_gen_dir: Path ) -> tuple[int, int]: """(total .root files, count with a same-named .parquet under any schema) for one gen.""" if not raw_gen_dir.is_dir(): return 0, 0 parquet_stems: set[tuple[str, str]] = set() if processed_gen_dir.is_dir(): for schema_dir in processed_gen_dir.iterdir(): if not schema_dir.is_dir(): continue for detector_dir in schema_dir.iterdir(): if not detector_dir.is_dir(): continue for f in detector_dir.iterdir(): if f.is_file() and f.suffix == ".parquet": parquet_stems.add((detector_dir.name, f.stem)) total = 0 referenced = 0 for detector_dir in raw_gen_dir.iterdir(): if not detector_dir.is_dir(): continue for f in detector_dir.iterdir(): if f.is_file() and f.suffix == ".root": total += 1 if (detector_dir.name, f.stem) in parquet_stems: referenced += 1 return total, referenced def _referenced_parquet_count( schema_dir: Path, manifest_referenced: set[Path] ) -> tuple[int, int]: """(total .parquet files, count listed in at least one manifest) for one schema dir.""" if not schema_dir.is_dir(): return 0, 0 total = 0 referenced = 0 for detector_dir in schema_dir.iterdir(): if not detector_dir.is_dir(): continue for f in detector_dir.iterdir(): if f.is_file() and f.suffix == ".parquet": total += 1 if f.resolve() in manifest_referenced: referenced += 1 return total, referenced def _parse_versions( versions_path: Path, ) -> tuple[dict[tuple[str, str], str], dict[tuple[str, str, str], str]]: """Read VERSIONS.md and return (gen_reasons, schema_reasons) keyed by (kind, gen_tag) and (kind, gen_tag, schema_tag) respectively. Later entries for the same key win, since VERSIONS.md is append-only and chronological.""" gen_reasons: dict[tuple[str, str], str] = {} schema_reasons: dict[tuple[str, str, str], str] = {} if not versions_path.is_file(): return gen_reasons, schema_reasons for line in versions_path.read_text().splitlines(): line = line.strip() m = VERSIONS_SCHEMA_LINE_RE.match(line) if m: schema_reasons[(m["kind"], m["gen"], m["schema"])] = m["reason"] continue m = VERSIONS_GEN_LINE_RE.match(line) if m: gen_reasons[(m["kind"], m["gen"])] = m["reason"] return gen_reasons, schema_reasons def _truncate(text: str, width: int = 72) -> str: text = text.strip() if len(text) <= width: return text return text[: width - 1].rstrip() + "…" def _human_size(n: int) -> str: size = float(n) for unit in ("B", "KB", "MB", "GB", "TB"): if size < 1024 or unit == "TB": return f"{size:.1f} {unit}" if unit != "B" else f"{int(size)} {unit}" size /= 1024 return f"{size:.1f} TB" _LABEL_WIDTH = 26 _COUNT_WIDTH = 20 _ROW_WIDTH = _LABEL_WIDTH + _COUNT_WIDTH + 10 def _count_str(count: int, referenced: int | None = None) -> str: files = "file" if count == 1 else "files" if referenced is not None: return f"{count} {files} ({referenced} ref)" return f"{count} {files}" def _row( label: str, size_bytes: int, indent: int = 0, level: str | None = None, count: int | None = None, referenced: int | None = None, ) -> str: text = " " * indent + label count_str = _count_str(count, referenced) if count is not None else "" size_str = _human_size(size_bytes) row = f"{text:<{_LABEL_WIDTH}}{count_str:<{_COUNT_WIDTH}}{size_str:>10}" return _colorize(row, level) if level else row def _reason_line(reason: str, indent: int) -> str: return _colorize(" " * indent + "↳ " + _truncate(reason), "reason") 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 manifest_referenced = _manifest_referenced_files(root / "pools") gen_reasons, schema_reasons = _parse_versions(root / "VERSIONS.md") grand_total = 0 grand_files = 0 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(_colorize(f"{kind}/", "kind")) kind_total = 0 kind_files = 0 for gen in gens: gen_tag = f"gen{gen}" raw_gen_dir = raw_root / kind / gen_tag raw_size = _du(raw_gen_dir) processed_gen_dir = root / "processed" / kind / gen_tag schema_dir = processed_gen_dir 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_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas} schema_counts = { s: _referenced_parquet_count( schema_dir / f"schema{s}", manifest_referenced ) for s in schemas } processed_size = sum(schema_sizes.values()) processed_files = sum(c[0] for c in schema_counts.values()) processed_referenced = sum(c[1] for c in schema_counts.values()) raw_files, raw_referenced = _referenced_root_count( raw_gen_dir, processed_gen_dir ) gen_total = raw_size + processed_size gen_files = raw_files + processed_files kind_total += gen_total kind_files += gen_files print(_row(gen_tag, gen_total, indent=1, level="gen", count=gen_files)) gen_reason = gen_reasons.get((kind, gen_tag)) if gen_reason: print(_reason_line(gen_reason, indent=2)) print( _row( "raw", raw_size, indent=2, level="bucket", count=raw_files, referenced=raw_referenced, ) ) print( _row( "processed", processed_size, indent=2, level="bucket", count=processed_files, referenced=processed_referenced, ) ) if schemas: for s in schemas: s_total, s_referenced = schema_counts[s] print( _row( f"schema{s}", schema_sizes[s], indent=3, level="schema", count=s_total, referenced=s_referenced, ) ) schema_reason = schema_reasons.get((kind, gen_tag, f"schema{s}")) if schema_reason: print(_reason_line(schema_reason, indent=4)) else: print(_colorize(" (none)", "schema")) print( _row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files) ) print() grand_total += kind_total grand_files += kind_files derived_dir = root / "derived" pools_dir = root / "pools" derived_size = _du(derived_dir) pools_size = _du(pools_dir) derived_files = _count_files(derived_dir) pools_files = _count_files(pools_dir) grand_total += derived_size + pools_size grand_files += derived_files + pools_files print(_row("derived/", derived_size, level="root", count=derived_files)) print(_row("pools/", pools_size, level="root", count=pools_files)) print("-" * _ROW_WIDTH) print(_row("grand total", grand_total, level="root", count=grand_files)) # --------------------------------------------------------------------------- # 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 # Check existence for every data line, not just ones whose gen/schema # actually changed — an already-correct-looking line can still point # at a file that was deleted or moved out-of-band. new_abs = Path(*parts) if not new_abs.exists(): missing.append(new_abs) if not changed: result.append((raw, None)) continue 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 entry points (called from scripts/dwarf.py) # --------------------------------------------------------------------------- def run_status(root: str) -> None: root_path = Path(root) if not root_path.is_dir(): raise SystemExit(f"error: {root_path} is not a directory") print_status(root_path) def _run_bump( kind: str, reason: str, by: str | None, date: str | None, execute: bool, root: str, gen: str | None, to: str | None, ) -> None: root_path = Path(root) if not root_path.is_dir(): raise SystemExit(f"error: {root_path} is not a directory") date = date or dt.date.today().isoformat() by = by if by is not None else _git_user_name() if gen is None: new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to) else: new_dirs, log_line = plan_bump_schema( root_path, kind, gen, reason, by, date, to ) print(f"=== {'EXECUTING' if 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 execute: print("\nDry run only — pass --execute to apply.") return apply_bump(root_path, new_dirs, log_line) print("\nDone.") def run_bump_gen( kind: str, reason: str, by: str | None, date: str | None, execute: bool, root: str, to: str | None = None, ) -> None: _run_bump(kind, reason, by, date, execute, root, gen=None, to=to) def run_bump_schema( kind: str, gen: str, reason: str, by: str | None, date: str | None, execute: bool, root: str, to: str | None = None, ) -> None: _run_bump(kind, reason, by, date, execute, root, gen=gen, to=to) def run_update_manifest( manifests: list[str], schema: str | None, execute: bool, gen: str | None = None, ) -> None: if schema and not SCHEMA_RE.match(schema): raise SystemExit(f"error: --schema must look like 'schemaN', got {schema!r}") if gen and not GEN_RE.match(gen): raise SystemExit(f"error: --gen must look like 'genN', got {gen!r}") all_plans: list[tuple[Path, list[tuple[str, str | None]]]] = [] all_missing: list[Path] = [] for raw in manifests: mp = Path(raw) plan, missing = plan_update_manifest(mp, schema, gen) all_plans.append((mp.resolve(), plan)) all_missing.extend(missing) print(f"=== {'EXECUTING' if 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 execute: raise SystemExit("error: refusing to write manifests with missing targets") if not execute: print("\nDry run only — pass --execute to apply.") return for mp, plan in all_plans: apply_update_manifest(mp, plan) print("\nDone.") def run_create_manifest( files: list[str], execute: bool, output: str | None = None, pool: str | None = None, type_: str | None = None, root: str = "/ceph/lbogner/geant_steps", ) -> None: if (output is None) == (pool is None): raise SystemExit("error: exactly one of --output or --pool is required") if pool is not None and type_ is None: raise SystemExit("error: --type is required when --pool is given") if pool is not None: output_path = Path(root) / "pools" / pool / f"{type_}.manifest" else: assert output is not None # guaranteed by the exclusivity check above output_path = Path(output) parquet_files = [Path(f) for f in files] lines, missing, resolved = plan_create_manifest(output_path, parquet_files) overlaps = check_holdout_overlap(output_path, resolved) print(f"=== {'EXECUTING' if 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 execute: raise SystemExit("error: refusing to write manifest (see above)") if not execute: print("\nDry run only — pass --execute to apply.") return apply_create_manifest(output_path, lines) print("\nDone.")