diff --git a/CLAUDE.md b/CLAUDE.md index 81fcdae..0025238 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -10,10 +10,13 @@ uv sync --extra cuda # install dependencies with CUDA 11.8 torch uv sync --extra cpu --extra dev # add dev extras (pytest, etc.) uv sync --extra cpu --extra geometry # add scikit-learn for the geometry oracle (giant rollout) pytest # run tests +giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir for a new training run giant train path/to/steps.parquet --mode flow # train (flow matching) giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline) +giant train-submit path/to/steps.parquet --config run/config.toml --accounting-group cms # train as a remote-GPU HTCondor job (TOpAS/NEMO2) giant predict path/to/steps.parquet --checkpoint ckpt/best.pt # per-step predictions giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl # full showers +giant rollout-submit path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl --accounting-group cms # rollout as a remote-GPU HTCondor job giant analyze submit rollout.yaml --accounting-group cms # parallel rollout-vs-reference analysis on HTCondor giant analyze render --gallery # render PDFs + HTML gallery (run_dir from prep/submit) dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen, @@ -62,6 +65,8 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep **Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. +**GPU HTCondor submission** (`giant/condor.py`, `giant train-submit`/`giant rollout-submit`): submits a single `giant train`/`giant rollout` invocation as a remote-GPU HTCondor job (ETP's TOpAS V100/A100 or NEMO2 L40S workers — see the ETP HTCondor wiki's "GPU Jobs" section). GPU workers are remote-only, so submit descriptions always carry `+RemoteJob = True`/`RequestGPUs`, and reach `/ceph` via `requirements = TARGET.ProvidesEtpCeph =?= True` rather than the local-only `TARGET.ProvidesETPResources` `giant analyze submit`'s CPU jobs use — this means the `giant` checkout submitting these jobs (and its `uv sync --extra cuda` venv) needs to live under `/ceph`, not `/work`/`/home`. Each submission writes a self-contained `condor/` (train, nested in `--out`) or `condor_/` (rollout, sibling to the output parquet) directory: `run.sh` (the wrapper actually executed — for training this re-checks `out_dir/last.pt` on every invocation and adds `--resume`, so a preempted/retried job resumes rather than restarting), `job.sub`, and `submission.json` (a `CondorJobMeta`: accounting group, GPU request, docker image, the exact command, and the assigned cluster id once known) — enough on its own to `condor_history ` a run later. `train-submit` requires `--config` (a TOML) rather than mirroring `train`'s hyperparameter flags, since the config file is already the reproducible source of truth (`giant/config.py:save_config`/`merge_cli_overrides`) and passing the same one on every retry is what keeps a resumed run's architecture from drifting. `giant new-run` scaffolds that TOML: resolves a base `--config` (or built-in defaults) plus a handful of override flags into a fresh run dir via `gconfig.default_out_dir` (hyperparam-named, uuid-tagged so repeat/same-day runs never collide), and prints the matching `giant train`/`giant train-submit` invocations — `giant train` overwrites this scaffolded `config.toml` in place once it actually runs, filling in the real dataset-derived meta section. + ## Roadmap **Phase 1 (done):** number of secondaries and their total energy were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC). diff --git a/giant/cli.py b/giant/cli.py index e371120..211d612 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -1,9 +1,10 @@ from collections import Counter -from datetime import date, datetime, timezone +from datetime import datetime, timezone from enum import Enum import math from pathlib import Path import re +import shlex from typing import Optional import uuid as uuid_mod @@ -95,6 +96,30 @@ def _batch_size_estimate_dims(model_cfg: dict, training: bool) -> tuple[int, int return model_cfg["hidden_dim"], model_cfg["n_blocks"] +def _router_cli_overrides( + router: bool | None, + router_type: str | None, + n_experts: int | None, + router_axis: list[str] | None, +) -> dict[str, object]: + """Build the `model.router` override dict from `--router`/`--router-type`/ + `--n-experts`/`--router-axis` flags (empty if none were given). Shared by + `train` and `new-run` so both resolve router overrides identically. + """ + cli_router: dict[str, object] = { + k: v + for k, v in { + "enabled": router, + "type": router_type, + "n_experts": n_experts, + }.items() + if v is not None + } + if router_axis: + cli_router.update(_parse_router_axis_flags(router_axis)) + return cli_router + + def _coerce_scalar(value: str) -> object: """Best-effort str -> bool/int/float, else leave as str. @@ -135,15 +160,21 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]: _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions") -def _resolve_prediction_output(data: Path, out: Path | None) -> tuple[Path, Path, str]: +def _resolve_prediction_output( + data: Path, out: Path | None, pred_uuid: str | None = None +) -> tuple[Path, Path, str]: """Return (out_path, resolved_dataset_path, pred_uuid). When *out* is None the output path is derived from *data*: - under /ceph/ → fixed central store with a UUID filename - elsewhere → sibling of *data* with a UUID filename + + *pred_uuid* can be pinned by the caller (e.g. `rollout-submit`, which + needs the same id for its condor run-dir name, the output filename, and + the YAML sidecar) — default is a freshly generated one, as before. """ dataset_path = data.resolve() - pred_uuid = str(uuid_mod.uuid4()) + pred_uuid = pred_uuid or str(uuid_mod.uuid4()) if out is None: if str(dataset_path).startswith("/ceph/"): out = _CEPH_PREDICTIONS / f"{pred_uuid}.parquet" @@ -451,17 +482,7 @@ def train( }.items() if v is not None } - cli_router: dict[str, object] = { - k: v - for k, v in { - "enabled": router, - "type": router_type, - "n_experts": n_experts, - }.items() - if v is not None - } - if router_axis: - cli_router.update(_parse_router_axis_flags(router_axis)) + cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis) if cli_router: cli_model["router"] = cli_router cfg = gconfig.merge_cli_overrides( @@ -484,16 +505,7 @@ def train( f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)" ) - out_dir = out or Path( - f"checkpoints/{date.today().strftime('%Y%m%d')}" - f"_{t['mode']}" - f"_h{m['hidden_dim']}" - f"_b{m['n_blocks']}" - f"_e{m['emb_dim']}" - f"_c{m['conditioning']}" - f"_lr{t['lr']}" - f"_bs{t['batch_size']}" - ) + out_dir = out or gconfig.default_out_dir(cfg) typer.echo(f"device: {_device}") typer.echo(f"out_dir: {out_dir}") @@ -510,6 +522,151 @@ def train( ) +@app.command("new-run") +def new_run( + config: Annotated[ + Optional[Path], + typer.Option( + "--config", + "-c", + help="Base TOML to start from (default: built-in defaults)", + ), + ] = None, + mode: Annotated[Optional[Mode], typer.Option("--mode", "-m")] = None, + epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None, + batch_size: Annotated[Optional[int], typer.Option("--batch-size", "-b")] = None, + lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None, + hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None, + n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None, + emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None, + dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None, + conditioning: Annotated[ + Optional[Conditioning], typer.Option("--conditioning") + ] = None, + router: Annotated[ + Optional[bool], typer.Option("--router/--no-router") + ] = None, + router_type: Annotated[Optional[str], typer.Option("--router-type")] = None, + n_experts: Annotated[Optional[int], typer.Option("--n-experts")] = None, + router_axis: Annotated[ + Optional[list[str]], typer.Option("--router-axis") + ] = None, + out: Annotated[ + Optional[Path], + typer.Option("--out", "-o", help="Run dir (default: auto from hyperparams)"), + ] = None, + comment: Annotated[ + Optional[str], + typer.Option( + "--comment", help="Free-text note recorded in config.toml's meta section" + ), + ] = None, + data: Annotated[ + Optional[Path], + typer.Option( + "--data", + help="Dataset path to fill in the printed next-step commands " + "(not stored in the config)", + ), + ] = None, + force: Annotated[ + bool, + typer.Option( + "--force", + help="Overwrite config.toml even if --out already has checkpoints", + ), + ] = False, + dry_run: Annotated[ + bool, + typer.Option( + "--dry-run", help="Print the resolved config without writing anything" + ), + ] = False, +) -> None: + """Scaffold a new training run: resolve hyperparams to a config.toml and lay out its run dir. + + This is the config-file-first counterpart to hand-editing a TOML: start + from a base --config (or built-in defaults), override a few hyperparams + inline, and this resolves+writes the full `config.toml` into a fresh (or + explicit --out) run dir — the same file `giant train --config ...` reads + and `giant train-submit --config ...` requires. `giant train` itself + overwrites this file in place once it actually runs (with the full + dataset-derived meta section), so this scaffold's meta section is just a + placeholder recording what was asked for and when. + """ + cli_train = { + k: v + for k, v in { + "mode": mode.value if mode is not None else None, + "epochs": epochs, + "batch_size": batch_size, + "lr": lr, + }.items() + if v is not None + } + cli_model: dict[str, object] = { + k: v + for k, v in { + "hidden_dim": hidden_dim, + "n_blocks": n_blocks, + "emb_dim": emb_dim, + "dropout": dropout, + "conditioning": conditioning.value if conditioning is not None else None, + }.items() + if v is not None + } + cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis) + if cli_router: + cli_model["router"] = cli_router + + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, cli_train, cli_model) + run_dir = (out or gconfig.default_out_dir(cfg)).resolve() + + if not force: + existing = [n for n in ("last.pt", "best.pt") if (run_dir / n).exists()] + if existing: + typer.echo( + f"error: {run_dir} already has {', '.join(existing)} — pass " + "--force to overwrite its config.toml anyway", + err=True, + ) + raise typer.Exit(1) + + typer.echo(f"run dir: {run_dir}") + + if dry_run: + typer.echo("dry-run: not writing anything. Resolved config:") + for section in ("train", "model"): + typer.echo(f"[{section}]") + for k, v in cfg[section].items(): + if k == "router": + continue + typer.echo(f" {k} = {v}") + return + + meta = { + "git_hash": gconfig.git_hash(), + "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "created_by": "giant new-run", + } + if comment: + meta["comment"] = comment + + run_dir.mkdir(parents=True, exist_ok=True) + gconfig.save_config(cfg, run_dir, meta) + config_path = run_dir / "config.toml" + typer.echo(f"wrote {config_path}") + + data_arg = str(data) if data is not None else "" + typer.echo("") + typer.echo("next:") + typer.echo(f" giant train {data_arg} --config {config_path} --out {run_dir}") + typer.echo( + f" giant train-submit {data_arg} --config {config_path} --out {run_dir} " + "--accounting-group " + ) + + @app.command() def predict( data: Annotated[ @@ -998,6 +1155,15 @@ def rollout( Optional[int], typer.Option("--seed", help="Torch/numpy seed for reproducibility"), ] = None, + prediction_id: Annotated[ + Optional[str], + typer.Option( + "--prediction-id", + help="Pin the prediction uuid (output filename, YAML sidecar name) " + "instead of generating a fresh one — used by `rollout-submit` so " + "its condor run-dir shares an id with the run it submitted", + ), + ] = None, ) -> None: """Roll the surrogate forward into full showers (autoregressive).""" if seed is not None: @@ -1048,7 +1214,9 @@ def rollout( seeds = _seed_from_data(files, n_events) typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)") - out, dataset_path, pred_uuid = _resolve_prediction_output(data, out) + out, dataset_path, pred_uuid = _resolve_prediction_output( + data, out, pred_uuid=prediction_id + ) out.parent.mkdir(parents=True, exist_ok=True) # Written incrementally as each batch of steps is produced, rather than @@ -1251,5 +1419,362 @@ def analyze_submit( subprocess.run(["condor_submit", str(sub)], check=True) +def _warn_if_not_ceph(path: Path) -> None: + """Remote GPU condor workers only reliably reach /ceph (see giant/condor.py).""" + if not str(path).startswith("/ceph/"): + typer.echo( + f"warning: {path} is not under /ceph/ — remote GPU condor workers " + "may not be able to reach it (see TARGET.ProvidesEtpCeph)", + err=True, + ) + + +@app.command("train-submit") +def train_submit( + data: Annotated[ + Path, + typer.Argument( + help="Parquet file or directory (must resolve under /ceph so a " + "remote GPU worker can reach it)" + ), + ], + config: Annotated[ + Path, + typer.Option( + "--config", + "-c", + help="TOML config — the full, reproducible spec for this run. " + "Hyperparams aren't exposed as flags here; edit the TOML instead " + "(this is also what makes resume-after-preemption safe: the same " + "file is passed on every condor retry, so the architecture never " + "drifts from the checkpoint being resumed).", + ), + ], + accounting_group: Annotated[str, typer.Option("--accounting-group")], + out: Annotated[ + Optional[Path], + typer.Option( + "--out", "-o", help="Checkpoint dir (default: auto from hyperparams)" + ), + ] = None, + request_gpus: Annotated[int, typer.Option("--request-gpus")] = 1, + gpu_type: Annotated[ + Optional[str], + typer.Option( + "--gpu-type", help='Pin GPU model, e.g. "Tesla V100-PCIE-32GB"' + ), + ] = None, + gpu_memory_mb: Annotated[Optional[int], typer.Option("--gpu-memory-mb")] = None, + request_memory: Annotated[ + int, typer.Option("--request-memory", help="MB") + ] = 16384, + request_walltime: Annotated[ + int, typer.Option("--request-walltime", help="Seconds (default: 2 days)") + ] = 172800, + docker_image: Annotated[ + str, typer.Option("--docker-image") + ] = "mschnepf/slc7-condocker", + repo_dir: Annotated[ + Optional[Path], + typer.Option( + "--repo-dir", + help="The /ceph checkout `uv run` executes from (default: cwd)", + ), + ] = None, + dry_run: Annotated[ + bool, typer.Option("--dry-run", help="Write files but don't condor_submit") + ] = False, +) -> None: + """Submit `giant train` as a single remote-GPU HTCondor job. + + GPU workers (TOpAS/NEMO2) are remote-only, so this always sets + `+RemoteJob = True` / `RequestGPUs` and reaches data via + `TARGET.ProvidesEtpCeph` rather than the local-only + `TARGET.ProvidesETPResources` `giant analyze submit` uses — see + giant/condor.py and the ETP HTCondor wiki's "GPU Jobs" section. + """ + from giant.condor import ( + CondorJobMeta, + GpuSubmitConfig, + parse_cluster_id, + write_gpu_submit, + ) + + data_path = data.resolve() + config_path = config.resolve() + _warn_if_not_ceph(data_path) + + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, {}, {}) + out_dir = (out or gconfig.default_out_dir(cfg)).resolve() + out_dir.mkdir(parents=True, exist_ok=True) + _warn_if_not_ceph(out_dir) + + repo_path = (repo_dir or Path.cwd()).resolve() + run_dir = out_dir / "condor" + + q = shlex.quote + last_ckpt = out_dir / "last.pt" + command = ( + f"uv run giant train {q(str(data_path))} --config {q(str(config_path))} " + f"--out {q(str(out_dir))} --device cuda" + ) + wrapper_body = ( + "#!/bin/bash\n" + "set -euo pipefail\n" + f"cd {q(str(repo_path))}\n" + 'RESUME=""\n' + f'[ -f {q(str(last_ckpt))} ] && RESUME="--resume {last_ckpt}"\n' + f"exec {command} $RESUME\n" + ) + + submit_cfg = GpuSubmitConfig( + run_dir=run_dir, + accounting_group=accounting_group, + repo_dir=repo_path, + command=command, + docker_image=docker_image, + request_memory_mb=request_memory, + request_gpus=request_gpus, + gpu_type=gpu_type, + gpu_memory_mb=gpu_memory_mb, + request_walltime_s=request_walltime, + ) + sub = write_gpu_submit(submit_cfg, wrapper_body) + + meta = CondorJobMeta( + accounting_group=accounting_group, + request_gpus=request_gpus, + gpu_type=gpu_type, + gpu_memory_mb=gpu_memory_mb, + request_memory_mb=request_memory, + request_walltime_s=request_walltime, + docker_image=docker_image, + repo_dir=str(repo_path), + command=command, + submitted_at=datetime.now(timezone.utc).isoformat(), + ) + meta_path = run_dir / "submission.json" + meta.save(meta_path) + + typer.echo(f"out dir: {out_dir}") + typer.echo(f"wrote submit description: {sub}") + if dry_run: + typer.echo("dry-run: not submitting") + return + + import subprocess + + result = subprocess.run( + ["condor_submit", str(sub)], check=True, capture_output=True, text=True + ) + typer.echo(result.stdout) + meta.cluster_id = parse_cluster_id(result.stdout) + meta.save(meta_path) + if meta.cluster_id is not None: + typer.echo(f"cluster id: {meta.cluster_id}") + + +@app.command("rollout-submit") +def rollout_submit( + data: Annotated[ + Path, + typer.Argument( + help="Parquet file/dir to seed showers from (must resolve under " + "/ceph so a remote GPU worker can reach it)" + ), + ], + checkpoint: Annotated[ + Path, + typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)"), + ], + geometry: Annotated[ + Path, + typer.Option( + "--geometry", + "-g", + help="Geometry oracle .pkl (dwarf build-geometry-oracle)", + ), + ], + accounting_group: Annotated[str, typer.Option("--accounting-group")], + energy_cutoff: Annotated[ + float, + typer.Option( + "--energy-cutoff", + help="Stop a track when its energy drops below this [MeV]", + ), + ] = 0.1, + max_steps: Annotated[ + int, typer.Option("--max-steps", help="Max steps per individual track") + ] = 1000, + steps: Annotated[ + int, + typer.Option( + "--steps", + "-s", + help="Flow matching ODE steps per model call (ignored for a wgan checkpoint)", + ), + ] = 10, + weights: Annotated[ + Weights, typer.Option("--weights", help="raw or ema (see `giant rollout --help`)") + ] = Weights.raw, + batch_size: Annotated[ + int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward") + ] = 4096, + max_tracks_per_event: Annotated[ + Optional[int], + typer.Option( + "--max-tracks-per-event", + help="Safety cap on tracks per shower (sub-cap secondaries deposit in place)", + ), + ] = None, + escape_threshold: Annotated[ + Optional[float], + typer.Option( + "--escape-threshold", + help="Override the oracle's NN-distance escape threshold [mm]", + ), + ] = None, + n_events: Annotated[ + Optional[int], typer.Option("--n-events", help="Cap number of seed events") + ] = None, + seed: Annotated[ + Optional[int], + typer.Option("--seed", help="Torch/numpy seed for reproducibility"), + ] = None, + out: Annotated[ + Optional[Path], typer.Option("--out", "-o", help="Output steps parquet") + ] = None, + request_gpus: Annotated[int, typer.Option("--request-gpus")] = 1, + gpu_type: Annotated[ + Optional[str], + typer.Option( + "--gpu-type", help='Pin GPU model, e.g. "Tesla V100-PCIE-32GB"' + ), + ] = None, + gpu_memory_mb: Annotated[Optional[int], typer.Option("--gpu-memory-mb")] = None, + request_memory: Annotated[ + int, typer.Option("--request-memory", help="MB") + ] = 16384, + request_walltime: Annotated[ + int, typer.Option("--request-walltime", help="Seconds (default: 2 days)") + ] = 172800, + docker_image: Annotated[ + str, typer.Option("--docker-image") + ] = "mschnepf/slc7-condocker", + repo_dir: Annotated[ + Optional[Path], + typer.Option( + "--repo-dir", + help="The /ceph checkout `uv run` executes from (default: cwd)", + ), + ] = None, + dry_run: Annotated[ + bool, typer.Option("--dry-run", help="Write files but don't condor_submit") + ] = False, +) -> None: + """Submit `giant rollout` as a single remote-GPU HTCondor job. + + No resume logic (unlike `train-submit`) — a retried rollout just starts + over, which is fine since it's deterministic given `--seed` and has no + epoch-loop state to preserve. + """ + from giant.condor import ( + CondorJobMeta, + GpuSubmitConfig, + parse_cluster_id, + write_gpu_submit, + ) + + data_path = data.resolve() + checkpoint_path = checkpoint.resolve() + geometry_path = geometry.resolve() + _warn_if_not_ceph(data_path) + _warn_if_not_ceph(checkpoint_path) + + out_path, _dataset_path, pred_uuid = _resolve_prediction_output(data_path, out) + out_path.parent.mkdir(parents=True, exist_ok=True) + _warn_if_not_ceph(out_path) + + repo_path = (repo_dir or Path.cwd()).resolve() + run_dir = out_path.parent / f"condor_{pred_uuid[:8]}" + + q = shlex.quote + flags = [ + f"--checkpoint {q(str(checkpoint_path))}", + f"--geometry {q(str(geometry_path))}", + f"--energy-cutoff {energy_cutoff}", + f"--max-steps {max_steps}", + f"--steps {steps}", + f"--weights {weights.value}", + f"--batch-size {batch_size}", + f"--out {q(str(out_path))}", + f"--prediction-id {pred_uuid}", + "--device cuda", + ] + if max_tracks_per_event is not None: + flags.append(f"--max-tracks-per-event {max_tracks_per_event}") + if escape_threshold is not None: + flags.append(f"--escape-threshold {escape_threshold}") + if n_events is not None: + flags.append(f"--n-events {n_events}") + if seed is not None: + flags.append(f"--seed {seed}") + + command = f"uv run giant rollout {q(str(data_path))} " + " ".join(flags) + wrapper_body = ( + "#!/bin/bash\n" + "set -euo pipefail\n" + f"cd {q(str(repo_path))}\n" + f"exec {command}\n" + ) + + submit_cfg = GpuSubmitConfig( + run_dir=run_dir, + accounting_group=accounting_group, + repo_dir=repo_path, + command=command, + docker_image=docker_image, + request_memory_mb=request_memory, + request_gpus=request_gpus, + gpu_type=gpu_type, + gpu_memory_mb=gpu_memory_mb, + request_walltime_s=request_walltime, + ) + sub = write_gpu_submit(submit_cfg, wrapper_body) + + meta = CondorJobMeta( + accounting_group=accounting_group, + request_gpus=request_gpus, + gpu_type=gpu_type, + gpu_memory_mb=gpu_memory_mb, + request_memory_mb=request_memory, + request_walltime_s=request_walltime, + docker_image=docker_image, + repo_dir=str(repo_path), + command=command, + submitted_at=datetime.now(timezone.utc).isoformat(), + ) + meta_path = run_dir / "submission.json" + meta.save(meta_path) + + typer.echo(f"prediction id: {pred_uuid}") + typer.echo(f"output: {out_path}") + typer.echo(f"wrote submit description: {sub}") + if dry_run: + typer.echo("dry-run: not submitting") + return + + import subprocess + + result = subprocess.run( + ["condor_submit", str(sub)], check=True, capture_output=True, text=True + ) + typer.echo(result.stdout) + meta.cluster_id = parse_cluster_id(result.stdout) + meta.save(meta_path) + if meta.cluster_id is not None: + typer.echo(f"cluster id: {meta.cluster_id}") + + if __name__ == "__main__": app() diff --git a/giant/condor.py b/giant/condor.py new file mode 100644 index 0000000..16c69a6 --- /dev/null +++ b/giant/condor.py @@ -0,0 +1,129 @@ +"""HTCondor GPU-job submission for `giant train` / `giant rollout`. + +Trains and rollouts run as single condor jobs (no plot-style fan-out) on the +ETP cluster's remote GPU resources (TOpAS V100/A100, NEMO2 L40S — see the ETP +HTCondor wiki's "GPU Jobs" section). GPU workers are remote-only, so these +jobs always carry ``+RemoteJob = True`` and ``RequestGPUs``; the local-only +``TARGET.ProvidesETPResources`` requirement ``giant.analysis.condor``'s CPU +jobs use is replaced by ``TARGET.ProvidesEtpCeph =?= True``, the +wiki-documented way for a remote job to reach ``/ceph`` without HTCondor +file-transferring multi-GB checkpoints/datasets. + +`giant/cli.py`'s `train-submit`/`rollout-submit` commands build a +`GpuSubmitConfig`, write the wrapper + submit description via +`write_gpu_submit`, run `condor_submit`, and record what was submitted (and, +once known, the assigned cluster id) in a `CondorJobMeta` JSON sidecar next +to the job files — so the run directory alone is enough to understand what +ran, on what resources, and how to look it up later +(``condor_history ``). +""" + +from __future__ import annotations + +import json +import re +from dataclasses import asdict, dataclass +from pathlib import Path + + +@dataclass +class GpuSubmitConfig: + """Everything needed to write one single-job GPU submit description.""" + + run_dir: Path + accounting_group: str + repo_dir: Path + command: str + docker_image: str = "mschnepf/slc7-condocker" + request_memory_mb: int = 16384 + request_gpus: int = 1 + gpu_type: str | None = None # -> TARGET.GPUs_DeviceName =?= "..." + gpu_memory_mb: int | None = None # -> TARGET.GPUs_GlobalMemoryMb >= N + request_walltime_s: int = 172800 # 2 days + + +def _gpu_requirements(cfg: GpuSubmitConfig) -> str: + """`TARGET.ProvidesEtpCeph` (remote /ceph access) ANDed with any GPU pin.""" + clauses = ["TARGET.ProvidesEtpCeph =?= True"] + if cfg.gpu_type is not None: + clauses.append(f'TARGET.GPUs_DeviceName =?= "{cfg.gpu_type}"') + if cfg.gpu_memory_mb is not None: + clauses.append(f"TARGET.GPUs_GlobalMemoryMb >= {cfg.gpu_memory_mb}") + return " && ".join(clauses) + + +def _gpu_submit_description(cfg: GpuSubmitConfig, wrapper: Path) -> str: + return ( + "universe = docker\n" + f"docker_image = {cfg.docker_image}\n" + f"executable = {wrapper}\n" + "should_transfer_files = YES\n" + "when_to_transfer_output = ON_EXIT\n" + f"request_memory = {cfg.request_memory_mb}\n" + f"RequestGPUs = {cfg.request_gpus}\n" + f"+RequestWalltime = {cfg.request_walltime_s}\n" + f"accounting_group = {cfg.accounting_group}\n" + "+RemoteJob = True\n" + f"requirements = ({_gpu_requirements(cfg)})\n" + f"output = {cfg.run_dir}/logs/job.out\n" + f"error = {cfg.run_dir}/logs/job.err\n" + f"log = {cfg.run_dir}/logs/job.log\n" + "queue 1\n" + ) + + +def write_gpu_submit(cfg: GpuSubmitConfig, wrapper_body: str) -> Path: + """Write the wrapper script + submit description under ``cfg.run_dir``. + + Returns the submit description path (``/job.sub``). Does not + call ``condor_submit`` — that's the caller's job (``giant/cli.py``), same + contract as ``giant.analysis.condor.write_submit``. + """ + run_dir = cfg.run_dir + (run_dir / "logs").mkdir(parents=True, exist_ok=True) + + wrapper = run_dir / "run.sh" + wrapper.write_text(wrapper_body) + wrapper.chmod(0o755) + + sub = run_dir / "job.sub" + sub.write_text(_gpu_submit_description(cfg, wrapper)) + return sub + + +_CLUSTER_ID_RE = re.compile(r"submitted to cluster (\d+)") + + +def parse_cluster_id(condor_submit_stdout: str) -> int | None: + """Extract the assigned cluster id from `condor_submit`'s stdout, if present.""" + m = _CLUSTER_ID_RE.search(condor_submit_stdout) + return int(m.group(1)) if m else None + + +@dataclass +class CondorJobMeta: + """What was submitted — written to ``/submission.json``. + + ``cluster_id`` starts unset and is filled in by the caller once + ``condor_submit``'s stdout has been parsed, so the run directory alone is + enough to ``condor_history `` this job later. + """ + + accounting_group: str + request_gpus: int + gpu_type: str | None + gpu_memory_mb: int | None + request_memory_mb: int + request_walltime_s: int + docker_image: str + repo_dir: str + command: str + submitted_at: str + cluster_id: int | None = None + + def save(self, path: str | Path) -> None: + Path(path).write_text(json.dumps(asdict(self), indent=2)) + + @classmethod + def load(cls, path: str | Path) -> "CondorJobMeta": + return cls(**json.loads(Path(path).read_text())) diff --git a/giant/config.py b/giant/config.py index cfd7a1b..47f92c7 100644 --- a/giant/config.py +++ b/giant/config.py @@ -2,7 +2,8 @@ import random import subprocess import sys import tomllib -from datetime import datetime, timezone +import uuid +from datetime import date, datetime, timezone from pathlib import Path import numpy as np @@ -279,6 +280,26 @@ def save_config(cfg: dict, out_dir: Path, meta: dict) -> None: (out_dir / "config.toml").write_text("\n".join(lines)) +def default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path: + """Auto-derived checkpoint dir: hyperparams for readability + a short + uuid tag so same-day/same-hyperparam runs (including a repeat condor + submission) never collide on an existing directory. + """ + t, m = cfg["train"], cfg["model"] + tag = uuid.uuid4().hex[:8] + return base / ( + f"{date.today().strftime('%Y%m%d')}" + f"_{t['mode']}" + f"_h{m['hidden_dim']}" + f"_b{m['n_blocks']}" + f"_e{m['emb_dim']}" + f"_c{m['conditioning']}" + f"_lr{t['lr']}" + f"_bs{t['batch_size']}" + f"_{tag}" + ) + + def build_run_meta( data: Path, seed: int, diff --git a/tests/test_cli_new_run.py b/tests/test_cli_new_run.py new file mode 100644 index 0000000..31fe5b2 --- /dev/null +++ b/tests/test_cli_new_run.py @@ -0,0 +1,118 @@ +"""Tests for `giant new-run` (config.toml + run-dir scaffolding).""" + +from __future__ import annotations + +import tomllib +from pathlib import Path + +from typer.testing import CliRunner + +from giant.cli import app + +runner = CliRunner() + + +def test_writes_config_with_overrides_applied(tmp_path: Path): + out_dir = tmp_path / "run1" + result = runner.invoke( + app, + [ + "new-run", + "--out", + str(out_dir), + "--mode", + "ddpm", + "--hidden-dim", + "128", + "--n-blocks", + "4", + "--lr", + "0.0005", + ], + ) + assert result.exit_code == 0, result.output + + config_path = out_dir / "config.toml" + assert config_path.exists() + with open(config_path, "rb") as f: + cfg = tomllib.load(f) + + assert cfg["train"]["mode"] == "ddpm" + assert cfg["train"]["lr"] == 0.0005 + assert cfg["model"]["hidden_dim"] == 128 + assert cfg["model"]["n_blocks"] == 4 + # untouched defaults still present + assert cfg["train"]["epochs"] == 100 + assert "router" in cfg["model"] + + assert str(out_dir) in result.output + assert "" in result.output + assert "giant train-submit" in result.output + + +def test_comment_and_provenance_recorded_in_meta(tmp_path: Path): + out_dir = tmp_path / "run2" + result = runner.invoke( + app, + ["new-run", "--out", str(out_dir), "--comment", "quick test"], + ) + assert result.exit_code == 0, result.output + + with open(out_dir / "config.toml", "rb") as f: + cfg = tomllib.load(f) + + assert cfg["meta"]["comment"] == "quick test" + assert cfg["meta"]["created_by"] == "giant new-run" + assert "created_at" in cfg["meta"] + assert "git_hash" in cfg["meta"] + + +def test_data_flag_fills_printed_next_step_commands(tmp_path: Path): + out_dir = tmp_path / "run3" + result = runner.invoke( + app, + ["new-run", "--out", str(out_dir), "--data", "/ceph/lbogner/train.parquet"], + ) + assert result.exit_code == 0, result.output + assert "/ceph/lbogner/train.parquet" in result.output + assert "" not in result.output + + +def test_dry_run_writes_nothing(tmp_path: Path): + out_dir = tmp_path / "run4" + result = runner.invoke( + app, + ["new-run", "--out", str(out_dir), "--hidden-dim", "512", "--dry-run"], + ) + assert result.exit_code == 0, result.output + assert "dry-run" in result.output + assert "hidden_dim = 512" in result.output + assert not out_dir.exists() + + +def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path): + out_dir = tmp_path / "run5" + out_dir.mkdir() + (out_dir / "last.pt").touch() + + result = runner.invoke(app, ["new-run", "--out", str(out_dir), "--mode", "ddpm"]) + assert result.exit_code != 0 + assert "already has last.pt" in result.output + assert not (out_dir / "config.toml").exists() + + result = runner.invoke( + app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"] + ) + assert result.exit_code == 0, result.output + assert (out_dir / "config.toml").exists() + + +def test_default_out_dir_used_when_out_omitted(tmp_path: Path, monkeypatch): + monkeypatch.chdir(tmp_path) + result = runner.invoke(app, ["new-run", "--hidden-dim", "64"]) + assert result.exit_code == 0, result.output + + checkpoints_dir = tmp_path / "checkpoints" + run_dirs = list(checkpoints_dir.iterdir()) if checkpoints_dir.exists() else [] + assert len(run_dirs) == 1 + assert (run_dirs[0] / "config.toml").exists() diff --git a/tests/test_cli_predict.py b/tests/test_cli_predict.py index a6f74fb..6a00b2c 100644 --- a/tests/test_cli_predict.py +++ b/tests/test_cli_predict.py @@ -58,6 +58,17 @@ def test_each_call_produces_a_distinct_uuid(tmp_path): assert uuid1 != uuid2 +def test_pinned_pred_uuid_is_used_as_is(tmp_path): + # `rollout-submit` pins the uuid at submit time so its condor run-dir, + # the output filename, and the eventual YAML sidecar all agree. + data = tmp_path / "data.parquet" + pinned = "abcd1234-abcd-4abc-9abc-abcdabcdabcd" + out, _, pred_uuid = _resolve_prediction_output(data, None, pred_uuid=pinned) + + assert pred_uuid == pinned + assert out.name == f"{pinned}.parquet" + + def test_dataset_path_is_resolved(tmp_path): data = tmp_path / "data.parquet" _, dataset_path, _ = _resolve_prediction_output(data, None) diff --git a/tests/test_condor_gpu.py b/tests/test_condor_gpu.py new file mode 100644 index 0000000..bbe898e --- /dev/null +++ b/tests/test_condor_gpu.py @@ -0,0 +1,211 @@ +"""Tests for GPU-job HTCondor submission (`giant train-submit` / `giant rollout-submit`).""" + +from __future__ import annotations + +from pathlib import Path + +from typer.testing import CliRunner + +from giant import config as gconfig +from giant.cli import app +from giant.condor import CondorJobMeta, GpuSubmitConfig, parse_cluster_id, write_gpu_submit + +runner = CliRunner() + +_MINIMAL_TOML = """\ +[train] +mode = "flow" +epochs = 1 + +[model] +hidden_dim = 8 +n_blocks = 2 +""" + + +# --------------------------------------------------------------------------- +# default_out_dir +# --------------------------------------------------------------------------- + + +def test_default_out_dir_encodes_hyperparams(): + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}, {}) + out_dir = gconfig.default_out_dir(cfg) + name = out_dir.name + assert f"_h{cfg['model']['hidden_dim']}_" in name + assert f"_b{cfg['model']['n_blocks']}_" in name + assert f"_c{cfg['model']['conditioning']}_" in name + + +def test_default_out_dir_avoids_collisions(): + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}, {}) + a = gconfig.default_out_dir(cfg) + b = gconfig.default_out_dir(cfg) + assert a != b + # same hyperparam-derived prefix, differing only in the uuid tag + assert a.name.rsplit("_", 1)[0] == b.name.rsplit("_", 1)[0] + + +# --------------------------------------------------------------------------- +# giant/condor.py: GpuSubmitConfig / write_gpu_submit +# --------------------------------------------------------------------------- + + +def test_write_gpu_submit_description(tmp_path: Path): + cfg = GpuSubmitConfig( + run_dir=tmp_path / "condor", + accounting_group="cms", + repo_dir=tmp_path, + command="uv run giant train data.parquet --config c.toml --out out --device cuda", + ) + sub = write_gpu_submit(cfg, "#!/bin/bash\necho hi\n") + txt = sub.read_text() + + assert "universe = docker" in txt + assert "docker_image = mschnepf/slc7-condocker" in txt + assert "+RemoteJob = True" in txt + assert "RequestGPUs = 1" in txt + assert "accounting_group = cms" in txt + assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in txt + assert "ProvidesETPResources" not in txt + assert "queue 1" in txt + + wrapper = cfg.run_dir / "run.sh" + assert wrapper.exists() + assert wrapper.stat().st_mode & 0o111 + assert wrapper.read_text() == "#!/bin/bash\necho hi\n" + + +def test_gpu_requirements_include_type_and_memory_pins(tmp_path: Path): + cfg = GpuSubmitConfig( + run_dir=tmp_path / "condor", + accounting_group="cms", + repo_dir=tmp_path, + command="uv run giant train ...", + request_gpus=2, + gpu_type="Tesla V100-PCIE-32GB", + gpu_memory_mb=16000, + ) + txt = write_gpu_submit(cfg, "#!/bin/bash\n").read_text() + + assert "RequestGPUs = 2" in txt + assert 'TARGET.GPUs_DeviceName =?= "Tesla V100-PCIE-32GB"' in txt + assert "TARGET.GPUs_GlobalMemoryMb >= 16000" in txt + + +# --------------------------------------------------------------------------- +# CondorJobMeta / parse_cluster_id +# --------------------------------------------------------------------------- + + +def test_condor_job_meta_round_trip(tmp_path: Path): + meta = CondorJobMeta( + accounting_group="cms", + request_gpus=1, + gpu_type=None, + gpu_memory_mb=None, + request_memory_mb=16384, + request_walltime_s=172800, + docker_image="mschnepf/slc7-condocker", + repo_dir="/ceph/lbogner/giant", + command="uv run giant train ...", + submitted_at="2026-07-24T00:00:00+00:00", + ) + path = tmp_path / "submission.json" + meta.save(path) + loaded = CondorJobMeta.load(path) + assert loaded == meta + assert loaded.cluster_id is None + + loaded.cluster_id = 123456 + loaded.save(path) + assert CondorJobMeta.load(path).cluster_id == 123456 + + +def test_parse_cluster_id(): + assert parse_cluster_id("1 job(s) submitted to cluster 123456.\n") == 123456 + assert parse_cluster_id("ERROR: something went wrong\n") is None + + +# --------------------------------------------------------------------------- +# CLI: giant train-submit / giant rollout-submit (--dry-run only, no cluster contact) +# --------------------------------------------------------------------------- + + +def test_train_submit_dry_run_writes_condor_dir(tmp_path: Path): + config = tmp_path / "config.toml" + config.write_text(_MINIMAL_TOML) + data = tmp_path / "train.parquet" + out_dir = tmp_path / "ckpt" + + result = runner.invoke( + app, + [ + "train-submit", + str(data), + "--config", + str(config), + "--accounting-group", + "cms", + "--out", + str(out_dir), + "--dry-run", + ], + ) + assert result.exit_code == 0, result.output + + run_dir = out_dir / "condor" + sub_txt = (run_dir / "job.sub").read_text() + assert "+RemoteJob = True" in sub_txt + assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in sub_txt + + wrapper = (run_dir / "run.sh").read_text() + assert "--device cuda" in wrapper + assert "last.pt" in wrapper + assert "RESUME" in wrapper + + meta = CondorJobMeta.load(run_dir / "submission.json") + assert meta.accounting_group == "cms" + assert meta.cluster_id is None + + +def test_rollout_submit_dry_run_writes_condor_dir(tmp_path: Path): + data = tmp_path / "seed.parquet" + checkpoint = tmp_path / "best.pt" + geometry = tmp_path / "oracle.pkl" + out = tmp_path / "rollout.parquet" + + result = runner.invoke( + app, + [ + "rollout-submit", + str(data), + "--checkpoint", + str(checkpoint), + "--geometry", + str(geometry), + "--accounting-group", + "cms", + "--out", + str(out), + "--dry-run", + ], + ) + assert result.exit_code == 0, result.output + + condor_dirs = list(tmp_path.glob("condor_*")) + assert len(condor_dirs) == 1 + run_dir = condor_dirs[0] + + sub_txt = (run_dir / "job.sub").read_text() + assert "+RemoteJob = True" in sub_txt + assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in sub_txt + + wrapper = (run_dir / "run.sh").read_text() + assert "--device cuda" in wrapper + assert "--prediction-id" in wrapper + assert "last.pt" not in wrapper + assert "RESUME" not in wrapper + + meta = CondorJobMeta.load(run_dir / "submission.json") + assert meta.accounting_group == "cms"