diff --git a/giant/cli.py b/giant/cli.py index 6636d40..de53608 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -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']}" diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py index 837d72d..9def3ef 100644 --- a/scripts/bump_dataset_version.py +++ b/scripts/bump_dataset_version.py @@ -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) diff --git a/scripts/dwarf.py b/scripts/dwarf.py index aaa3eb0..a81bf39 100644 --- a/scripts/dwarf.py +++ b/scripts/dwarf.py @@ -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") diff --git a/tests/test_bump_dataset_version.py b/tests/test_bump_dataset_version.py index 6f3bc30..ccdf3ce 100644 --- a/tests/test_bump_dataset_version.py +++ b/tests/test_bump_dataset_version.py @@ -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)