lars
c0746c40e0
Shows per-schema, per-gen (raw/processed split), per-kind, and grand total sizes so the dataset tree's footprint is visible at a glance.
527 lines
17 KiB
Python
527 lines
17 KiB
Python
"""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 (`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
|
|
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 _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 _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"
|
|
|
|
|
|
_ROW_WIDTH = 48
|
|
|
|
|
|
def _row(label: str, size_bytes: int, indent: int = 0) -> str:
|
|
text = " " * indent + label
|
|
size_str = _human_size(size_bytes)
|
|
pad = max(1, _ROW_WIDTH - len(text) - len(size_str))
|
|
return text + " " * pad + size_str
|
|
|
|
|
|
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
|
|
grand_total = 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(f"{kind}/")
|
|
kind_total = 0
|
|
for gen in gens:
|
|
gen_tag = f"gen{gen}"
|
|
raw_size = _du(raw_root / kind / gen_tag)
|
|
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_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas}
|
|
processed_size = sum(schema_sizes.values())
|
|
gen_total = raw_size + processed_size
|
|
kind_total += gen_total
|
|
|
|
print(_row(gen_tag, gen_total, indent=1))
|
|
print(_row("raw", raw_size, indent=2))
|
|
print(_row("processed", processed_size, indent=2))
|
|
if schemas:
|
|
for s in schemas:
|
|
print(_row(f"schema{s}", schema_sizes[s], indent=3))
|
|
else:
|
|
print(" (none)")
|
|
print(_row(f"{kind} total", kind_total, indent=1))
|
|
print()
|
|
grand_total += kind_total
|
|
|
|
derived_size = _du(root / "derived")
|
|
pools_size = _du(root / "pools")
|
|
grand_total += derived_size + pools_size
|
|
print(_row("derived/", derived_size))
|
|
print(_row("pools/", pools_size))
|
|
print("-" * _ROW_WIDTH)
|
|
print(_row("grand total", grand_total))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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 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.")
|