Unify dataset/tooling scripts into a single dwarf Typer CLI
Replace the five separately-hyphenated uv entry points (steps-to-parquet, steps-to-parquet-parallel, migrate-geant-steps, bump-dataset-version, create-root-files) plus the unregistered hparam_scan.py with one `dwarf` command exposing convert/migrate/bump-gen/bump-schema/status/ update-manifest/create-manifest/make-root/hparam-scan as subcommands. Each scripts/*.py module now only holds argparse-free business logic; scripts/dwarf.py wires it up with Typer, matching giant/cli.py's style. `dwarf convert` merges the old serial/parallel conversion scripts behind a --jobs flag (default 1: sequential with plain -o; >1: dataset-layout fan-out via subprocess). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
+128
-174
@@ -1,4 +1,3 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Cut a new raw generation or processed schema version for the geant_steps
|
||||
dataset tree (see scripts/migrate_geant_steps.py for the layout):
|
||||
|
||||
@@ -7,21 +6,16 @@ dataset tree (see scripts/migrate_geant_steps.py for the layout):
|
||||
pools/<detector>/<pool>.manifest
|
||||
|
||||
`gen` bumps when the underlying ROOT changes (geometry/physics-list/macro).
|
||||
`schema` bumps when the parquet export (steps_to_parquet.py or similar)
|
||||
changes, and is scoped to its gen — a new gen always starts back at schema1.
|
||||
`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 to apply.
|
||||
to VERSIONS.md. Defaults to a dry run; pass execute=True to apply.
|
||||
|
||||
Usage:
|
||||
bump_dataset_version.py bump-gen --kind steps --reason "switched EM physics list"
|
||||
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 create-manifest --output pools/pbwo4/train.manifest a.parquet b.parquet
|
||||
bump_dataset_version.py status
|
||||
See `uv run dwarf bump-gen/bump-schema/update-manifest/create-manifest/status
|
||||
--help` for the CLI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import datetime as dt
|
||||
import os
|
||||
import re
|
||||
@@ -287,183 +281,143 @@ def apply_create_manifest(output_path: Path, lines: list[str]) -> None:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# CLI entry points (called from scripts/dwarf.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--root",
|
||||
default="/ceph/lbogner/geant_steps",
|
||||
help="Dataset root for bump-gen/bump-schema/status (default: /ceph/lbogner/geant_steps)",
|
||||
)
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
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)
|
||||
|
||||
kind_help = "steps | hits | ... (default: steps)"
|
||||
|
||||
p_gen = sub.add_parser("bump-gen", help="Cut a new raw generation")
|
||||
p_gen.add_argument("--kind", default="steps", help=kind_help)
|
||||
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("--execute", action="store_true", help="Apply (default: dry run)")
|
||||
def _run_bump(
|
||||
kind: str,
|
||||
reason: str,
|
||||
by: str | None,
|
||||
date: str | None,
|
||||
execute: bool,
|
||||
root: str,
|
||||
gen: str | None,
|
||||
) -> None:
|
||||
root_path = Path(root)
|
||||
if not root_path.is_dir():
|
||||
raise SystemExit(f"error: {root_path} is not a directory")
|
||||
|
||||
p_schema = sub.add_parser("bump-schema", help="Cut a new schema within a gen")
|
||||
p_schema.add_argument("--kind", default="steps", help=kind_help)
|
||||
p_schema.add_argument("--gen", required=True, help="Existing gen tag, e.g. gen1")
|
||||
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("--execute", action="store_true", help="Apply (default: dry run)")
|
||||
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)
|
||||
else:
|
||||
new_dirs, log_line = plan_bump_schema(root_path, kind, gen, reason, by, date)
|
||||
|
||||
p_update = sub.add_parser(
|
||||
"update-manifest",
|
||||
help="Repoint manifest(s) to a new schema, verifying all target files exist",
|
||||
)
|
||||
p_update.add_argument(
|
||||
"manifests", nargs="+", metavar="MANIFEST", help="One or more .manifest files to update"
|
||||
)
|
||||
p_update.add_argument(
|
||||
"--schema",
|
||||
default=None,
|
||||
metavar="schemaN",
|
||||
help="Target schema tag (default: highest schema found in the same gen dir)",
|
||||
)
|
||||
p_update.add_argument("--execute", action="store_true", help="Write updated manifests (default: dry run)")
|
||||
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}")
|
||||
|
||||
p_create = sub.add_parser(
|
||||
"create-manifest",
|
||||
help="Create a new manifest from a list of parquet files",
|
||||
)
|
||||
dest_group = p_create.add_mutually_exclusive_group(required=True)
|
||||
dest_group.add_argument(
|
||||
"--output", "-o", metavar="PATH", help="Explicit path for the new .manifest file"
|
||||
)
|
||||
dest_group.add_argument(
|
||||
"--pool", metavar="DETECTOR",
|
||||
help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
|
||||
)
|
||||
p_create.add_argument(
|
||||
"--type", choices=["full", "holdout", "dev"],
|
||||
help="Pool type — full, holdout, or dev (required with --pool)",
|
||||
)
|
||||
p_create.add_argument(
|
||||
"files", nargs="+", metavar="FILE", help="Parquet files to include"
|
||||
)
|
||||
p_create.add_argument("--execute", action="store_true", help="Write the manifest (default: dry run)")
|
||||
if not execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
apply_bump(root_path, new_dirs, log_line)
|
||||
print("\nDone.")
|
||||
|
||||
sub.add_parser("status", help="List existing gens/schemas per kind")
|
||||
|
||||
args = parser.parse_args()
|
||||
def run_bump_gen(
|
||||
kind: str, reason: str, by: str | None, date: str | None, execute: bool, root: str
|
||||
) -> None:
|
||||
_run_bump(kind, reason, by, date, execute, root, gen=None)
|
||||
|
||||
# Commands that need --root
|
||||
if args.command in ("bump-gen", "bump-schema", "status"):
|
||||
root = Path(args.root)
|
||||
if not root.is_dir():
|
||||
parser.error(f"{root} is not a directory")
|
||||
|
||||
if args.command == "status":
|
||||
print_status(root)
|
||||
def run_bump_schema(
|
||||
kind: str, gen: str, reason: str, by: str | None, date: str | None, execute: bool, root: str
|
||||
) -> None:
|
||||
_run_bump(kind, reason, by, date, execute, root, gen=gen)
|
||||
|
||||
|
||||
def run_update_manifest(manifests: list[str], schema: str | None, execute: bool) -> None:
|
||||
if schema and not SCHEMA_RE.match(schema):
|
||||
raise SystemExit(f"error: --schema must look like 'schemaN', got {schema!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)
|
||||
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
|
||||
|
||||
if args.command in ("bump-gen", "bump-schema"):
|
||||
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)
|
||||
else:
|
||||
new_dirs, log_line = plan_bump_schema(root, args.kind, args.gen, args.reason, by, date)
|
||||
for mp, plan in all_plans:
|
||||
apply_update_manifest(mp, plan)
|
||||
print("\nDone.")
|
||||
|
||||
print(f"=== {'EXECUTING' if args.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 args.execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
apply_bump(root, new_dirs, log_line)
|
||||
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
|
||||
|
||||
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}")
|
||||
|
||||
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)
|
||||
all_plans.append((mp.resolve(), plan))
|
||||
all_missing.extend(missing)
|
||||
|
||||
print(f"=== {'EXECUTING' if args.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 args.execute:
|
||||
raise SystemExit("error: refusing to write manifests with missing targets")
|
||||
|
||||
if not args.execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
for mp, plan in all_plans:
|
||||
apply_update_manifest(mp, plan)
|
||||
print("\nDone.")
|
||||
return
|
||||
|
||||
if args.command == "create-manifest":
|
||||
if args.pool is not None and args.type is None:
|
||||
parser.error("--type is required when --pool is given")
|
||||
|
||||
if args.pool is not None:
|
||||
root = Path(args.root)
|
||||
output_path = root / "pools" / args.pool / f"{args.type}.manifest"
|
||||
else:
|
||||
output_path = Path(args.output)
|
||||
|
||||
parquet_files = [Path(f) for f in args.files]
|
||||
lines, missing, resolved = plan_create_manifest(output_path, parquet_files)
|
||||
overlaps = check_holdout_overlap(output_path, resolved)
|
||||
|
||||
print(f"=== {'EXECUTING' if args.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 args.execute:
|
||||
raise SystemExit("error: refusing to write manifest (see above)")
|
||||
|
||||
if not args.execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
apply_create_manifest(output_path, lines)
|
||||
print("\nDone.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
apply_create_manifest(output_path, lines)
|
||||
print("\nDone.")
|
||||
|
||||
Reference in New Issue
Block a user