Fold --to/--gen dataset-versioning flags into the dwarf CLI

origin/energy-conservation-poc grew bump-gen/bump-schema --to and
update-manifest --gen flags (091b23a) plus a train output-dir date
prefix (305e436) after the dwarf unification was written locally.
Reconcile: bring plan_bump_gen/plan_bump_schema/plan_update_manifest's
target/target_gen support into the plain-function (argparse-free) form,
thread --to/--gen through scripts/dwarf.py's bump-gen/bump-schema/
update-manifest commands, and take giant/cli.py's date-prefix change
and the associated tests as-is.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-02 09:57:17 +02:00
parent d5853d5a75
commit 917835e182
4 changed files with 197 additions and 55 deletions
+3 -2
View File
@@ -1,5 +1,5 @@
from collections import Counter
from datetime import datetime, timezone
from datetime import date, datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Optional
@@ -245,7 +245,8 @@ def train(
)
out_dir = out or Path(
f"checkpoints/{t['mode']}"
f"checkpoints/{date.today().strftime('%Y%m%d')}"
f"_{t['mode']}"
f"_h{m['hidden_dim']}"
f"_b{m['n_blocks']}"
f"_e{m['emb_dim']}"
+93 -49
View File
@@ -26,6 +26,14 @@ 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():
@@ -50,18 +58,24 @@ def _git_user_name() -> str | None:
def plan_bump_gen(
root: Path, kind: str, reason: str, by: str | None, date: str
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*
— checking both, since a gen can exist in one tree before the other catches up."""
next_gen = (
max(
_max_index(root / "raw" / kind, GEN_RE),
_max_index(root / "processed" / kind, GEN_RE),
"""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
)
+ 1
)
gen_tag = f"gen{next_gen}"
gen_tag = f"gen{next_gen}"
new_dirs = [
root / "raw" / kind / gen_tag,
root / "processed" / kind / gen_tag / "schema1",
@@ -72,7 +86,8 @@ def plan_bump_gen(
def plan_bump_schema(
root: Path, kind: str, gen_tag: str, reason: str, by: str | None, date: str
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}")
@@ -82,8 +97,13 @@ def plan_bump_schema(
raise SystemExit(
f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
)
next_schema = _max_index(processed_gen_dir, SCHEMA_RE) + 1
schema_tag = f"schema{next_schema}"
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}"
@@ -133,18 +153,12 @@ def print_status(root: Path) -> None:
# 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
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 schema replacements for each data line.
"""Parse a manifest and plan gen/schema replacements for each data line.
Returns:
lines: list of (original_line, new_relative_path_or_None)
@@ -168,30 +182,39 @@ def plan_update_manifest(
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])
changed = False
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]
# 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
if new_schema == old_schema:
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
parts[schema_idx] = new_schema
new_abs = Path(*parts)
if not new_abs.exists():
missing.append(new_abs)
@@ -299,6 +322,7 @@ def _run_bump(
execute: bool,
root: str,
gen: str | None,
to: str | None,
) -> None:
root_path = Path(root)
if not root_path.is_dir():
@@ -307,9 +331,9 @@ def _run_bump(
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)
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)
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:")
@@ -326,27 +350,47 @@ def _run_bump(
def run_bump_gen(
kind: str, reason: str, by: str | None, date: str | None, execute: bool, root: str
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)
_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
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)
_run_bump(kind, reason, by, date, execute, root, gen=gen, to=to)
def run_update_manifest(manifests: list[str], schema: str | None, execute: bool) -> None:
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)
plan, missing = plan_update_manifest(mp, schema, gen)
all_plans.append((mp.resolve(), plan))
all_missing.extend(missing)
+33 -4
View File
@@ -175,11 +175,19 @@ def bump_gen(
date: Annotated[
Optional[str], typer.Option("--date", help="Override date (default: today, ISO)")
] = None,
to: Annotated[
Optional[str],
typer.Option(
"--to", metavar="genN", help="Target gen tag (default: one past the current highest)"
),
] = None,
execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
"""Cut a new raw generation."""
run_bump_gen(kind=kind, reason=reason, by=by, date=date, execute=execute, root=str(root))
run_bump_gen(
kind=kind, reason=reason, by=by, date=date, execute=execute, root=str(root), to=to
)
@app.command("bump-schema")
@@ -193,12 +201,27 @@ def bump_schema(
date: Annotated[
Optional[str], typer.Option("--date", help="Override date (default: today, ISO)")
] = None,
to: Annotated[
Optional[str],
typer.Option(
"--to",
metavar="schemaN",
help="Target schema tag (default: one past the current highest)",
),
] = None,
execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
"""Cut a new schema within a gen."""
run_bump_schema(
kind=kind, gen=gen, reason=reason, by=by, date=date, execute=execute, root=str(root)
kind=kind,
gen=gen,
reason=reason,
by=by,
date=date,
execute=execute,
root=str(root),
to=to,
)
@@ -223,12 +246,18 @@ def update_manifest(
help="Target schema tag (default: highest schema found in the same gen dir)",
),
] = None,
gen: Annotated[
Optional[str],
typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
] = None,
execute: Annotated[
bool, typer.Option("--execute", help="Write updated manifests (default: dry run)")
] = False,
) -> None:
"""Repoint manifest(s) to a new schema, verifying all target files exist."""
run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute)
"""Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
run_update_manifest(
[str(m) for m in manifests], schema=schema, execute=execute, gen=gen
)
@app.command("create-manifest")
+68
View File
@@ -74,6 +74,40 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path):
pass
def test_bump_gen_to_specific_tag(tmp_path):
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
dirs, log_line = plan_bump_gen(tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5")
assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
assert "`gen5`" in log_line
def test_bump_gen_rejects_invalid_to_tag(tmp_path):
try:
plan_bump_gen(tmp_path, "steps", "bad tag", None, "2026-01-01", target="v5")
assert False, "expected SystemExit"
except SystemExit:
pass
def test_bump_schema_to_specific_tag(tmp_path):
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
(tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
dirs, log_line = plan_bump_schema(
tmp_path, "steps", "gen1", "jump to schema5", None, "2026-01-01", target="schema5"
)
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema5"]
assert "`schema5`" in log_line
def test_bump_schema_rejects_invalid_to_tag(tmp_path):
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
try:
plan_bump_schema(tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3")
assert False, "expected SystemExit"
except SystemExit:
pass
def test_apply_bump_creates_dirs_and_appends_log(tmp_path):
dirs, log_line = plan_bump_gen(tmp_path, "steps", "reason A", "alice", "2026-01-01")
apply_bump(tmp_path, dirs, log_line)
@@ -185,6 +219,40 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path):
assert lines[2][1] is not None # the data line was updated
def test_update_manifest_bumps_gen(tmp_path):
parquet = tmp_path / "processed" / "steps" / "gen2" / "schema1" / "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")
lines, missing = plan_update_manifest(manifest, None, target_gen="gen2")
assert missing == []
changed = [(old, new) for old, new in lines if new is not None]
assert len(changed) == 1
assert "gen2" in changed[0][1]
assert "gen1" not in changed[0][1]
def test_update_manifest_bumps_gen_and_schema(tmp_path):
parquet = tmp_path / "processed" / "steps" / "gen2" / "schema3" / "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")
lines, missing = plan_update_manifest(manifest, "schema3", target_gen="gen2")
assert missing == []
changed = [(old, new) for old, new in lines if new is not None]
assert len(changed) == 1
assert "gen2" in changed[0][1]
assert "schema3" in changed[0][1]
def test_apply_update_manifest_writes_file(tmp_path):
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
_make_parquet(parquet)