diff --git a/giant/cli.py b/giant/cli.py index ea40946..070a00c 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -211,9 +211,10 @@ class Mode(str, Enum): wgan = "wgan" -class Conditioning(str, Enum): - physical = "physical" - embedding = "embedding" +# Conditioning itself lives in giant.config (imported below as gconfig) — +# shared with scripts/dwarf.py's Typer commands so the two CLIs can't +# silently drift apart on the option's valid values. +Conditioning = gconfig.Conditioning class Coord(str, Enum): diff --git a/giant/config.py b/giant/config.py index d021693..e4945e5 100644 --- a/giant/config.py +++ b/giant/config.py @@ -4,11 +4,23 @@ import subprocess import sys import tomllib from datetime import datetime, timezone +from enum import Enum from pathlib import Path import numpy as np import torch + +class Conditioning(str, Enum): + """`model.conditioning` choices — shared by `giant.cli` and `scripts.dwarf`'s + Typer commands so the two CLIs can't silently drift apart on the option's + valid values (see DEFAULT_CONFIG["model"]["conditioning"] for what each + value means).""" + + physical = "physical" + embedding = "embedding" + + DEFAULT_CONFIG: dict = { "train": { "mode": "flow", diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py index dd603c8..216748a 100644 --- a/scripts/bump_dataset_version.py +++ b/scripts/bump_dataset_version.py @@ -85,7 +85,12 @@ def _git_user_name() -> str | None: out = subprocess.run( ["git", "config", "user.name"], capture_output=True, text=True, timeout=2 ) - except OSError: + except (OSError, subprocess.SubprocessError): + # OSError (e.g. git not on PATH) and subprocess.SubprocessError + # (e.g. TimeoutExpired) are unrelated hierarchies — TimeoutExpired + # is not an OSError, so catching only OSError (as before) let a + # slow/loaded NFS-backed portal machine crash this instead of + # degrading to by=None as intended. return None name = out.stdout.strip() return name or None @@ -730,6 +735,7 @@ def run_create_manifest( pool: str | None = None, type_: str | None = None, root: str = "/ceph/lbogner/geant_steps", + force: bool = False, ) -> None: if (output is None) == (pool is None): raise SystemExit("error: exactly one of --output or --pool is required") @@ -745,6 +751,12 @@ def run_create_manifest( 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) + # Unlike missing/overlaps this is a hard stop even without --execute + # reaching the write, since create_manifest has no in-place "update" mode + # (unlike update_manifest) — a second run against the same output_path + # (e.g. holdout.manifest, the file check_holdout_overlap exists to + # protect) would otherwise silently clobber it with no diff/backup. + already_exists = output_path.exists() and not force print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===") print(f"manifest: {output_path.resolve()}") @@ -761,7 +773,10 @@ def run_create_manifest( for name, f in overlaps: print(f" {f} (also in {name})") - if (missing or overlaps) and execute: + if already_exists: + print(f"\n{output_path} already exists — pass --force to overwrite it.") + + if (missing or overlaps or already_exists) and execute: raise SystemExit("error: refusing to write manifest (see above)") if not execute: diff --git a/scripts/dwarf.py b/scripts/dwarf.py index beacc89..200b1ec 100644 --- a/scripts/dwarf.py +++ b/scripts/dwarf.py @@ -5,6 +5,7 @@ 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 @@ -12,6 +13,7 @@ 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, @@ -37,6 +39,25 @@ 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" @@ -106,6 +127,7 @@ def convert( 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 @@ -327,6 +349,10 @@ def create_manifest( 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( @@ -336,6 +362,7 @@ def create_manifest( pool=pool, type_=type_.value if type_ is not None else None, root=str(root), + force=force, ) @@ -391,6 +418,7 @@ def make_root( ] = False, ) -> None: """Generate new ROOT shards via a minicalosim executable.""" + _warn_if_exceeds_shared_quota(jobs, "--jobs") run_make_root( executable=executable, detector=detector, @@ -472,11 +500,6 @@ def build_geometry_oracle( ) -class Conditioning(str, Enum): - physical = "physical" - embedding = "embedding" - - @app.command("warm-cache") def warm_cache( data: Annotated[ diff --git a/tests/test_bump_dataset_version.py b/tests/test_bump_dataset_version.py index c2835f0..3f169a5 100644 --- a/tests/test_bump_dataset_version.py +++ b/tests/test_bump_dataset_version.py @@ -1,4 +1,6 @@ import os +import subprocess + from scripts import bump_dataset_version plan_bump_gen = bump_dataset_version.plan_bump_gen @@ -11,6 +13,14 @@ apply_create_manifest = bump_dataset_version.apply_create_manifest check_holdout_overlap = bump_dataset_version.check_holdout_overlap +def test_git_user_name_returns_none_on_timeout(monkeypatch): + def _raise_timeout(*args, **kwargs): + raise subprocess.TimeoutExpired(cmd=["git"], timeout=2) + + monkeypatch.setattr(subprocess, "run", _raise_timeout) + assert bump_dataset_version._git_user_name() is None + + def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path): dirs, log_line = plan_bump_gen( tmp_path, "steps", "first generation", None, "2026-01-01" @@ -401,6 +411,36 @@ def test_create_manifest_creates_parent_dirs(tmp_path): assert output.exists() +def test_run_create_manifest_refuses_to_overwrite_existing_output(tmp_path): + pq = tmp_path / "a.parquet" + pq.touch() + output = tmp_path / "pools" / "pbwo4" / "holdout.manifest" + output.parent.mkdir(parents=True) + output.write_text("original contents\n") + + try: + bump_dataset_version.run_create_manifest( + [str(pq)], execute=True, output=str(output) + ) + assert False, "expected SystemExit" + except SystemExit: + pass + assert output.read_text() == "original contents\n" + + +def test_run_create_manifest_force_overwrites_existing_output(tmp_path): + pq = tmp_path / "a.parquet" + pq.touch() + output = tmp_path / "pools" / "pbwo4" / "holdout.manifest" + output.parent.mkdir(parents=True) + output.write_text("original contents\n") + + bump_dataset_version.run_create_manifest( + [str(pq)], execute=True, output=str(output), force=True + ) + assert output.read_text() != "original contents\n" + + # --------------------------------------------------------------------------- # check_holdout_overlap # --------------------------------------------------------------------------- diff --git a/tests/test_dwarf.py b/tests/test_dwarf.py index 9e29958..14b2673 100644 --- a/tests/test_dwarf.py +++ b/tests/test_dwarf.py @@ -1,12 +1,35 @@ from typer.testing import CliRunner +from giant import cli as giant_cli +from giant.config import Conditioning from giant.data import setup_cache +from scripts import dwarf from scripts.dwarf import app from test_pipeline import _make_synthetic_steps runner = CliRunner() +def test_conditioning_enum_shared_across_both_clis(): + """giant.cli and scripts.dwarf must use the one giant.config.Conditioning + enum, not independently redefined copies that could silently drift apart + on valid --conditioning values.""" + assert dwarf.Conditioning is Conditioning + assert giant_cli.Conditioning is Conditioning + + +def test_warn_if_exceeds_shared_quota_warns_over_quarter_cpu(monkeypatch, capsys): + monkeypatch.setattr(dwarf.os, "cpu_count", lambda: 8) # quota = 2 + dwarf._warn_if_exceeds_shared_quota(3, "--jobs") + assert "warning: --jobs=3 exceeds" in capsys.readouterr().err + + +def test_warn_if_exceeds_shared_quota_silent_within_quota(monkeypatch, capsys): + monkeypatch.setattr(dwarf.os, "cpu_count", lambda: 8) # quota = 2 + dwarf._warn_if_exceeds_shared_quota(2, "--jobs") + assert capsys.readouterr().err == "" + + def test_convert_rejects_jobs_below_one(tmp_path): root_file = tmp_path / "shard.root" root_file.touch()