lars
e5bf7c51cb
- energy_simplex_encode: warn when clipping post_E to pre_E discards recorded edep/e_sec instead of silently zeroing them - local/inv_local_frame_rotation: validate and normalize pre_dir instead of silently assuming unit norm; raise on near-zero-norm rows - train(): make --lr authoritative on resume instead of being silently overwritten by the checkpoint's optimizer/scheduler state; print and exit cleanly instead of silently training zero epochs when the checkpoint already meets --epochs; truncate metrics.csv on a fresh run instead of always appending - dwarf update-manifest: check file existence for every manifest line, not just ones whose gen/schema actually changed - pyproject.toml: dev extra now pulls in convert+analysis so the documented `uv sync --extra cpu --extra dev` + `pytest` actually passes collection Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
732 lines
25 KiB
Python
732 lines
25 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
|
|
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"^- `(?P<gen>gen\d+)`/`(?P<schema>schema\d+)` \(kind=(?P<kind>[\w-]+)\)"
|
|
r" — \d{4}-\d{2}-\d{2} — (?P<reason>.+)$"
|
|
)
|
|
VERSIONS_GEN_LINE_RE = re.compile(
|
|
r"^- `(?P<gen>gen\d+)` \(kind=(?P<kind>[\w-]+)\) — \d{4}-\d{2}-\d{2} — (?P<reason>.+)$"
|
|
)
|
|
|
|
# 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.")
|