lars
ca3a2a3462
- run_create_manifest gains --force; it previously overwrote an existing manifest (including holdout.manifest, which check_holdout_overlap exists specifically to protect) with no warning or backup on a second run. - _git_user_name only caught OSError, not subprocess.TimeoutExpired (a SubprocessError, not an OSError) — a slow/loaded shared portal machine could crash `dwarf bump-gen`/`bump-schema` instead of degrading to by=None as intended. - `dwarf convert --jobs`/`make-root --jobs` now warn (never block) when the requested count exceeds ~1/4 of the machine's CPUs, matching the same shared-machine etiquette check added to giant train in the previous commit. - The Conditioning enum was independently redefined in both giant/cli.py and scripts/dwarf.py; moved to a single giant.config.Conditioning both now import, removing the drift risk of a third conditioning mode being added to one but not the other. Each fix has a regression test. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
585 lines
18 KiB
Python
585 lines
18 KiB
Python
"""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.
|
|
"""
|
|
|
|
import os
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import typer
|
|
from typing_extensions import Annotated
|
|
|
|
from giant.config import Conditioning
|
|
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.geometry_oracle import run_build_geometry_oracle
|
|
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
|
|
from scripts.warm_setup_cache import run_warm_setup_cache
|
|
|
|
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)."""
|
|
|
|
|
|
def _warn_if_exceeds_shared_quota(n: int, flag: str) -> None:
|
|
"""Soft warning (never blocks) when a worker/job count looks likely to
|
|
grab more than this repo's documented shared-portal-machine etiquette
|
|
(CLAUDE.md's Compute environment: stay within ~1/4 of CPU/RAM and a
|
|
single GPU, since portal1/deepthought{,2}/bms{1..3} are shared with
|
|
other users). Not a hard cap — a legitimate big machine or a
|
|
deliberately aggressive run is still the caller's call.
|
|
"""
|
|
cpu_count = os.cpu_count() or 1
|
|
quota = max(1, cpu_count // 4)
|
|
if n > quota:
|
|
typer.echo(
|
|
f"warning: {flag}={n} exceeds ~1/4 of this machine's "
|
|
f"{cpu_count} CPU(s) ({quota}) — portal machines are shared "
|
|
"with other users (see CLAUDE.md's Compute environment section)",
|
|
err=True,
|
|
)
|
|
|
|
|
|
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)
|
|
_warn_if_exceeds_shared_quota(jobs, "--jobs")
|
|
|
|
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)
|
|
total_orphaned = 0
|
|
for root_file in root_files:
|
|
_, n_orphaned = convert_steps_to_parquet(
|
|
root_file,
|
|
output_path=output,
|
|
batch_size=batch_size,
|
|
tree_name=tree,
|
|
compression=compression_value,
|
|
)
|
|
total_orphaned += n_orphaned
|
|
if total_orphaned:
|
|
typer.echo(
|
|
f"\n{total_orphaned} orphaned child track(s) dropped across "
|
|
f"{len(root_files)} file(s)."
|
|
)
|
|
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,
|
|
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),
|
|
to=to,
|
|
)
|
|
|
|
|
|
@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,
|
|
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),
|
|
to=to,
|
|
)
|
|
|
|
|
|
@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,
|
|
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 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")
|
|
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,
|
|
force: Annotated[
|
|
bool,
|
|
typer.Option("--force", help="Overwrite the manifest if it already exists"),
|
|
] = 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),
|
|
force=force,
|
|
)
|
|
|
|
|
|
@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")
|
|
],
|
|
energy_gev: Annotated[
|
|
float | None,
|
|
typer.Option(
|
|
"--energy-gev",
|
|
help="energy_GeV passed to the executable (default: executable's own "
|
|
"default, currently 1.0). Note the dataset detector label is not "
|
|
"derived from this — e.g. use '--detector pbwo4_10gev --energy-gev 10' "
|
|
"to name the dataset accordingly.",
|
|
),
|
|
] = None,
|
|
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."""
|
|
_warn_if_exceeds_shared_quota(jobs, "--jobs")
|
|
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,
|
|
energy_gev=energy_gev,
|
|
)
|
|
|
|
|
|
class OracleMethod(str, Enum):
|
|
slab = "slab"
|
|
knn = "knn"
|
|
svm = "svm"
|
|
|
|
|
|
@app.command("build-geometry-oracle")
|
|
def build_geometry_oracle(
|
|
data: Annotated[
|
|
Path, typer.Argument(help="Steps parquet file or directory of steps files")
|
|
],
|
|
out: Annotated[Path, typer.Option("--out", "-o", help="Output oracle .pkl path")],
|
|
method: Annotated[
|
|
OracleMethod,
|
|
typer.Option(
|
|
"--method",
|
|
help=(
|
|
"Lookup strategy: slab (default; exact O(log #segments) fast "
|
|
"path for the layered-slab detector geometry), knn, or svm "
|
|
"(generic fallbacks for non-slab geometries)"
|
|
),
|
|
),
|
|
] = OracleMethod.slab,
|
|
k: Annotated[
|
|
int, typer.Option("--k", help="Neighbours for the knn classifier")
|
|
] = 1,
|
|
subsample: Annotated[
|
|
int,
|
|
typer.Option("--subsample", help="Max reference points sampled from the data"),
|
|
] = 500_000,
|
|
escape_factor: Annotated[
|
|
float,
|
|
typer.Option(
|
|
"--escape-factor",
|
|
help="escape_threshold = this x median spacing of reference points",
|
|
),
|
|
] = 5.0,
|
|
seed: Annotated[int, typer.Option("--seed", help="Sampling seed")] = 0,
|
|
depth_axis: Annotated[
|
|
int,
|
|
typer.Option(
|
|
"--depth-axis",
|
|
help="0/1/2 -> x/y/z axis the layers stack along (method=slab only)",
|
|
),
|
|
] = 2,
|
|
n_bins: Annotated[
|
|
int,
|
|
typer.Option(
|
|
"--n-bins",
|
|
help="Depth-axis resolution, finer than the thinnest layer (method=slab only)",
|
|
),
|
|
] = 2000,
|
|
) -> None:
|
|
"""Fit a position -> (material, layer_id) oracle for `giant rollout`."""
|
|
run_build_geometry_oracle(
|
|
data=data,
|
|
out=out,
|
|
method=method.value,
|
|
k=k,
|
|
subsample=subsample,
|
|
escape_factor=escape_factor,
|
|
seed=seed,
|
|
depth_axis=depth_axis,
|
|
n_bins=n_bins,
|
|
)
|
|
|
|
|
|
@app.command("warm-cache")
|
|
def warm_cache(
|
|
data: Annotated[
|
|
Path,
|
|
typer.Argument(
|
|
help="Parquet file, directory, or .manifest — same as `giant train`'s"
|
|
),
|
|
],
|
|
val_fraction: Annotated[
|
|
float,
|
|
typer.Option(
|
|
"--val-fraction",
|
|
"-f",
|
|
help="Must match the `giant train` run(s) to warm for",
|
|
),
|
|
] = 0.1,
|
|
seed: Annotated[
|
|
int,
|
|
typer.Option(
|
|
"--seed", "-s", help="Must match the `giant train` run(s) to warm for"
|
|
),
|
|
] = 0,
|
|
conditioning: Annotated[
|
|
Conditioning,
|
|
typer.Option(
|
|
"--conditioning", help="Must match the `giant train` run(s) to warm for"
|
|
),
|
|
] = Conditioning.physical,
|
|
router: Annotated[
|
|
bool,
|
|
typer.Option(
|
|
"--router/--no-router",
|
|
help="Warm the process vocabulary too (only takes effect with "
|
|
"--router-type process)",
|
|
),
|
|
] = False,
|
|
router_type: Annotated[
|
|
str, typer.Option("--router-type", help="Router implementation name")
|
|
] = "energy",
|
|
n_experts: Annotated[
|
|
int, typer.Option("--n-experts", help="Number of routed experts")
|
|
] = 4,
|
|
rebuild: Annotated[
|
|
bool,
|
|
typer.Option(
|
|
"--rebuild", help="Ignore any existing sidecar and recompute every section"
|
|
),
|
|
] = False,
|
|
) -> None:
|
|
"""Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.
|
|
|
|
Warms the vocab maps, event-id split index, and the normalizer entry for
|
|
the given --val-fraction/--seed/--conditioning, so a later `giant train`
|
|
run (or a `dwarf hparam-scan` sweep, which shares one such entry across
|
|
every run) skips straight to training. See giant/data/setup_cache.py.
|
|
"""
|
|
run_warm_setup_cache(
|
|
data=str(data),
|
|
val_fraction=val_fraction,
|
|
seed=seed,
|
|
conditioning=conditioning.value,
|
|
router_enabled=router,
|
|
router_type=router_type,
|
|
n_experts=n_experts,
|
|
rebuild=rebuild,
|
|
echo=typer.echo,
|
|
)
|
|
|
|
|
|
@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()
|