lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Add --to flag for bump-gen/bump-schema and --gen flag for update-manifest

Browse Source 
bump-gen and bump-schema now accept --to genN/schemaN to target a specific
version instead of always auto-incrementing. update-manifest gains --gen genN
to repoint the gen component of manifest paths (combinable with --schema).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-01 15:58:26 +02:00
parent 7b37b284f8
commit 091b23a60f
2 changed files with 155 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/bump_dataset_version.py
+87 -46
View File
@@ -15,8 +15,10 @@ to VERSIONS.md. Defaults to a dry run; pass --execute to apply.
Usage:
bump_dataset_version.py bump-gen --kind steps --reason "switched EM physics list"
bump_dataset_version.py bump-gen --kind steps --to gen5 --reason "skip to gen5"
bump_dataset_version.py bump-schema --kind steps --gen gen1 --reason "added e_sec column"
bump_dataset_version.py update-manifest pools/pbwo4/full.manifest [--schema schema2]
bump_dataset_version.py bump-schema --kind steps --gen gen1 --to schema3 --reason "skip to schema3"
bump_dataset_version.py update-manifest pools/pbwo4/full.manifest [--schema schema2] [--gen gen2]
bump_dataset_version.py create-manifest --output pools/pbwo4/train.manifest a.parquet b.parquet
bump_dataset_version.py status
"""
@@ -32,6 +34,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():
@@ -56,18 +66,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",
@@ -78,7 +94,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}")
@@ -88,8 +105,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}"
@@ -139,18 +161,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)
@@ -174,30 +190,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)
@@ -306,6 +331,10 @@ def main() -> None:
p_gen.add_argument("--reason", required=True, help="Why this gen exists")
p_gen.add_argument("--by", default=None, help="Attribution (default: git user.name)")
p_gen.add_argument("--date", default=None, help="Override date (default: today, ISO)")
p_gen.add_argument(
"--to", default=None, metavar="genN",
help="Target gen tag (default: one past the current highest)",
)
p_gen.add_argument("--execute", action="store_true", help="Apply (default: dry run)")
p_schema = sub.add_parser("bump-schema", help="Cut a new schema within a gen")
@@ -314,11 +343,15 @@ def main() -> None:
p_schema.add_argument("--reason", required=True, help="Why this schema exists")
p_schema.add_argument("--by", default=None, help="Attribution (default: git user.name)")
p_schema.add_argument("--date", default=None, help="Override date (default: today, ISO)")
p_schema.add_argument(
"--to", default=None, metavar="schemaN",
help="Target schema tag (default: one past the current highest)",
)
p_schema.add_argument("--execute", action="store_true", help="Apply (default: dry run)")
p_update = sub.add_parser(
"update-manifest",
help="Repoint manifest(s) to a new schema, verifying all target files exist",
help="Repoint manifest(s) to a new gen and/or schema, verifying all target files exist",
)
p_update.add_argument(
"manifests", nargs="+", metavar="MANIFEST", help="One or more .manifest files to update"
@@ -329,6 +362,12 @@ def main() -> None:
metavar="schemaN",
help="Target schema tag (default: highest schema found in the same gen dir)",
)
p_update.add_argument(
"--gen",
default=None,
metavar="genN",
help="Target gen tag (default: keep existing gen)",
)
p_update.add_argument("--execute", action="store_true", help="Write updated manifests (default: dry run)")
p_create = sub.add_parser(
@@ -370,9 +409,9 @@ def main() -> None:
date = args.date or dt.date.today().isoformat()
by = args.by if args.by is not None else _git_user_name()
if args.command == "bump-gen":
new_dirs, log_line = plan_bump_gen(root, args.kind, args.reason, by, date)
new_dirs, log_line = plan_bump_gen(root, args.kind, args.reason, by, date, args.to)
else:
new_dirs, log_line = plan_bump_schema(root, args.kind, args.gen, args.reason, by, date)
new_dirs, log_line = plan_bump_schema(root, args.kind, args.gen, args.reason, by, date, args.to)
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
print("new directories:")
@@ -391,13 +430,15 @@ def main() -> None:
if args.command == "update-manifest":
if args.schema and not SCHEMA_RE.match(args.schema):
parser.error(f"--schema must look like 'schemaN', got {args.schema!r}")
if args.gen and not GEN_RE.match(args.gen):
parser.error(f"--gen must look like 'genN', got {args.gen!r}")
all_plans: list[tuple[Path, list[tuple[str, str | None]]]] = []
all_missing: list[Path] = []
for raw in args.manifests:
mp = Path(raw)
plan, missing = plan_update_manifest(mp, args.schema)
plan, missing = plan_update_manifest(mp, args.schema, args.gen)
all_plans.append((mp.resolve(), plan))
all_missing.extend(missing)
tests/test_bump_dataset_version.py
+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)