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:
@@ -0,0 +1,334 @@
|
||||
"""dwarf — little helper to `giant`: dataset/tooling CLI for the geant_steps pipeline.
|
||||
|
||||
Unifies the standalone scripts/*.py conversion, migration, versioning, and
|
||||
simulation-fanout tools into one Typer app so there's a single command name
|
||||
(and `--help`) to remember instead of five differently-hyphenated ones.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
from typing_extensions import Annotated
|
||||
|
||||
from scripts.bump_dataset_version import (
|
||||
run_bump_gen,
|
||||
run_bump_schema,
|
||||
run_create_manifest,
|
||||
run_status,
|
||||
run_update_manifest,
|
||||
)
|
||||
from scripts.create_root_files import run_make_root
|
||||
from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
|
||||
from scripts.migrate_geant_steps import run_migration
|
||||
from scripts.steps_to_parquet import convert_steps_to_parquet
|
||||
from scripts.steps_to_parquet_parallel import run_parallel_job
|
||||
|
||||
app = typer.Typer(no_args_is_help=True)
|
||||
|
||||
_DATASET_ROOT_DEFAULT = Path("/ceph/lbogner/geant_steps")
|
||||
|
||||
|
||||
@app.callback()
|
||||
def _main() -> None:
|
||||
"""dwarf — dataset/tooling CLI (ROOT<->parquet conversion, dataset versioning, sim fanout)."""
|
||||
|
||||
|
||||
class Compression(str, Enum):
|
||||
snappy = "snappy"
|
||||
lz4 = "lz4"
|
||||
zstd = "zstd"
|
||||
gzip = "gzip"
|
||||
none = "none"
|
||||
|
||||
|
||||
class PoolType(str, Enum):
|
||||
full = "full"
|
||||
holdout = "holdout"
|
||||
dev = "dev"
|
||||
|
||||
|
||||
@app.command()
|
||||
def convert(
|
||||
root_files: Annotated[
|
||||
list[Path], typer.Argument(help="Input ROOT file(s)")
|
||||
],
|
||||
output: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option(
|
||||
"--output",
|
||||
"-o",
|
||||
help="Output Parquet file (default: <input>.parquet). Only valid "
|
||||
"with a single input file and --jobs 1.",
|
||||
),
|
||||
] = None,
|
||||
batch_size: Annotated[
|
||||
str,
|
||||
typer.Option(
|
||||
"--batch-size",
|
||||
help="Uproot read batch size, e.g. '100 MB' or '500000' (rows)",
|
||||
),
|
||||
] = "100 MB",
|
||||
tree: Annotated[
|
||||
str, typer.Option("--tree", help="Tree name inside the ROOT file")
|
||||
] = "Steps",
|
||||
compression: Annotated[
|
||||
Compression, typer.Option("--compression", help="Parquet compression codec")
|
||||
] = Compression.snappy,
|
||||
jobs: Annotated[
|
||||
int,
|
||||
typer.Option(
|
||||
"--jobs",
|
||||
"-j",
|
||||
help="Convert N files in parallel, resolving each destination from "
|
||||
"--dataset-root/--schema (default: 1, sequential, any file layout)",
|
||||
),
|
||||
] = 1,
|
||||
dataset_root: Annotated[
|
||||
Path,
|
||||
typer.Option(
|
||||
"--dataset-root",
|
||||
help="Dataset root containing raw/ and processed/ (only used with --jobs > 1)",
|
||||
),
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
schema: Annotated[
|
||||
Optional[str],
|
||||
typer.Option(
|
||||
"--schema",
|
||||
help="Schema tag to write parquets under, e.g. schema2 (only used "
|
||||
"with --jobs > 1; default: highest schemaN already under "
|
||||
"processed/<kind>/<gen>/)",
|
||||
),
|
||||
] = None,
|
||||
) -> None:
|
||||
"""Convert ROOT Steps tree(s) to Parquet."""
|
||||
if jobs < 1:
|
||||
typer.echo("error: --jobs must be >= 1", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
compression_value = "uncompressed" if compression is Compression.none else compression.value
|
||||
|
||||
if jobs == 1:
|
||||
if output is not None and len(root_files) > 1:
|
||||
typer.echo("error: --output can only be used with a single input file", err=True)
|
||||
raise typer.Exit(1)
|
||||
for root_file in root_files:
|
||||
convert_steps_to_parquet(
|
||||
root_file,
|
||||
output_path=output,
|
||||
batch_size=batch_size,
|
||||
tree_name=tree,
|
||||
compression=compression_value,
|
||||
)
|
||||
return
|
||||
|
||||
if output is not None:
|
||||
typer.echo(
|
||||
"error: --output cannot be combined with --jobs > 1 "
|
||||
"(destinations are derived from --dataset-root/--schema)",
|
||||
err=True,
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
run_parallel_job(
|
||||
[str(f) for f in root_files],
|
||||
jobs=jobs,
|
||||
dataset_root=dataset_root,
|
||||
schema=schema,
|
||||
batch_size=batch_size,
|
||||
tree=tree,
|
||||
compression=compression_value,
|
||||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def migrate(
|
||||
root: Annotated[
|
||||
Path, typer.Argument(help="Dataset root to migrate in place")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
execute: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
"--execute", help="Actually move/copy files and write manifests (default: dry run)"
|
||||
),
|
||||
] = False,
|
||||
copy: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
"--copy",
|
||||
help="Copy instead of move, leaving the originals in place "
|
||||
"(e.g. if another process is still reading them)",
|
||||
),
|
||||
] = False,
|
||||
) -> None:
|
||||
"""One-time migration into the versioned raw/processed/pools/derived layout."""
|
||||
run_migration(str(root), execute=execute, copy=copy)
|
||||
|
||||
|
||||
@app.command("bump-gen")
|
||||
def bump_gen(
|
||||
reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
|
||||
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
|
||||
by: Annotated[
|
||||
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
|
||||
] = None,
|
||||
date: Annotated[
|
||||
Optional[str], typer.Option("--date", help="Override date (default: today, ISO)")
|
||||
] = 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))
|
||||
|
||||
|
||||
@app.command("bump-schema")
|
||||
def bump_schema(
|
||||
gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
|
||||
reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
|
||||
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
|
||||
by: Annotated[
|
||||
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
|
||||
] = None,
|
||||
date: Annotated[
|
||||
Optional[str], typer.Option("--date", help="Override date (default: today, ISO)")
|
||||
] = 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)
|
||||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def status(
|
||||
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
|
||||
) -> None:
|
||||
"""List existing gens/schemas per kind."""
|
||||
run_status(str(root))
|
||||
|
||||
|
||||
@app.command("update-manifest")
|
||||
def update_manifest(
|
||||
manifests: Annotated[
|
||||
list[Path], typer.Argument(help="One or more .manifest files to update")
|
||||
],
|
||||
schema: Annotated[
|
||||
Optional[str],
|
||||
typer.Option(
|
||||
"--schema",
|
||||
metavar="schemaN",
|
||||
help="Target schema tag (default: highest schema found in the same gen dir)",
|
||||
),
|
||||
] = 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)
|
||||
|
||||
|
||||
@app.command("create-manifest")
|
||||
def create_manifest(
|
||||
files: Annotated[list[Path], typer.Argument(help="Parquet files to include")],
|
||||
output: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option("--output", "-o", help="Explicit path for the new .manifest file"),
|
||||
] = None,
|
||||
pool: Annotated[
|
||||
Optional[str],
|
||||
typer.Option(
|
||||
"--pool",
|
||||
metavar="DETECTOR",
|
||||
help="Detector name; combined with --type and --root to form "
|
||||
"<root>/pools/<detector>/<type>.manifest",
|
||||
),
|
||||
] = None,
|
||||
type_: Annotated[
|
||||
Optional[PoolType],
|
||||
typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
|
||||
] = None,
|
||||
root: Annotated[
|
||||
Path, typer.Option("--root", help="Dataset root (used with --pool)")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
execute: Annotated[
|
||||
bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
|
||||
] = False,
|
||||
) -> None:
|
||||
"""Create a new manifest from a list of parquet files."""
|
||||
run_create_manifest(
|
||||
[str(f) for f in files],
|
||||
execute=execute,
|
||||
output=str(output) if output is not None else None,
|
||||
pool=pool,
|
||||
type_=type_.value if type_ is not None else None,
|
||||
root=str(root),
|
||||
)
|
||||
|
||||
|
||||
@app.command("make-root")
|
||||
def make_root(
|
||||
executable: Annotated[
|
||||
Path, typer.Option("--executable", help="Built minicalosim run_* executable")
|
||||
],
|
||||
detector: Annotated[
|
||||
list[str],
|
||||
typer.Option(
|
||||
"--detector",
|
||||
metavar="NAME[:CONFIG]",
|
||||
help="Dataset detector label, optionally with ':CONFIG' to pass as the "
|
||||
"executable's config-name argument (e.g. sampling_pb_scint:pb_scint). "
|
||||
"Omit ':CONFIG' for executables that take no config selector (e.g. run_pbwo4). "
|
||||
"Repeatable.",
|
||||
),
|
||||
],
|
||||
num_files: Annotated[
|
||||
int, typer.Option("--num-files", help="New shards to create per detector")
|
||||
],
|
||||
events_per_file: Annotated[
|
||||
int, typer.Option("--events-per-file", help="nEvents passed to the executable")
|
||||
],
|
||||
gen: Annotated[
|
||||
str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
|
||||
],
|
||||
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
|
||||
dataset_root: Annotated[
|
||||
Path, typer.Option("--dataset-root", help="Dataset root")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
jobs: Annotated[
|
||||
int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")
|
||||
] = 4,
|
||||
execute: Annotated[
|
||||
bool, typer.Option("--execute", help="Actually run jobs (default: dry run / print plan)")
|
||||
] = False,
|
||||
) -> None:
|
||||
"""Generate new ROOT shards via a minicalosim executable."""
|
||||
run_make_root(
|
||||
executable=executable,
|
||||
detector=detector,
|
||||
num_files=num_files,
|
||||
events_per_file=events_per_file,
|
||||
kind=kind,
|
||||
gen=gen,
|
||||
dataset_root=str(dataset_root),
|
||||
jobs=jobs,
|
||||
execute=execute,
|
||||
)
|
||||
|
||||
|
||||
@app.command("hparam-scan")
|
||||
def hparam_scan(
|
||||
data: Annotated[str, typer.Option("--data")] = DATA_DEFAULT,
|
||||
scan_dir: Annotated[str, typer.Option("--scan-dir")] = SCAN_DIR_DEFAULT,
|
||||
seed: Annotated[int, typer.Option("--seed")] = 0,
|
||||
dry_run: Annotated[bool, typer.Option("--dry-run")] = False,
|
||||
) -> None:
|
||||
"""Grid-scan dropout x n_blocks x hidden_dim via sequential `giant train` runs."""
|
||||
run_hparam_scan(data=data, scan_dir=scan_dir, seed=seed, dry_run=dry_run)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
Reference in New Issue
Block a user