Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 91aed19d4b | |||
| dde8b367a4 | |||
| 3d1fa7979e | |||
| 430917d8f2 |
@@ -10,12 +10,15 @@ 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 dev # add dev extras (pytest, etc.)
|
||||||
uv sync --extra cpu --extra geometry # add scikit-learn for the geometry oracle (giant rollout)
|
uv sync --extra cpu --extra geometry # add scikit-learn for the geometry oracle (giant rollout)
|
||||||
pytest # run tests
|
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 flow # train (flow matching)
|
||||||
giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline)
|
giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline)
|
||||||
giant train path/to/steps.parquet --mode wgan # train (WGAN-GP, single-pass eval; implemented, not yet tested)
|
giant train path/to/steps.parquet --mode wgan # train (WGAN-GP, single-pass eval; implemented, not yet tested)
|
||||||
giant train path/to/steps.parquet --router --router-type energy # MoE routing trunk (implemented; first rollout benchmark failed with lambda_balance=0, retrain needed — see Roadmap)
|
giant train path/to/steps.parquet --router --router-type energy # MoE routing trunk (implemented; first rollout benchmark failed with lambda_balance=0, retrain needed — see Roadmap)
|
||||||
|
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 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 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 submit rollout.yaml --accounting-group cms # parallel rollout-vs-reference analysis on HTCondor
|
||||||
giant analyze render <run_dir> --gallery # render PDFs + HTML gallery (run_dir from prep/submit)
|
giant analyze render <run_dir> --gallery # render PDFs + HTML gallery (run_dir from prep/submit)
|
||||||
dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
|
dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
|
||||||
@@ -76,6 +79,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.
|
**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_<tag>/` (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 <cluster_id>` 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
|
## 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).
|
**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).
|
||||||
|
|||||||
@@ -2,15 +2,19 @@
|
|||||||
|
|
||||||
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate — a play on *Geant4* and the step function being the computationally heaviest part of the simulation.
|
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate — a play on *Geant4* and the step function being the computationally heaviest part of the simulation.
|
||||||
|
|
||||||
Proof-of-concept surrogate model for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained conditional generative model.
|
Conditional generative surrogate for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the variable-length list of secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained generative model. A trained checkpoint autoregressively rolls out full showers, stepping each primary and pushing secondaries as new tracks.
|
||||||
|
|
||||||
Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
|
Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
|
||||||
|
|
||||||
## Architecture
|
## Architecture
|
||||||
|
|
||||||
A **two-stage conditional flow matching** model (Lipman et al. 2022): a small MLP learns a vector field mapping noise → step outcomes in ~10 ODE steps per sample. Falls back to DDPM for comparison.
|
A **two-stage model**, both stages checkpointed together, with a choice of generative mode per stage (`--mode`):
|
||||||
|
|
||||||
**Stage 1 — primary (9D, diffused):**
|
- **`flow`** (default) — conditional flow matching (Lipman et al. 2022): an MLP learns a vector field mapping noise → step outcomes, sampled via ODE integration in ~10 steps.
|
||||||
|
- **`ddpm`** — a standard denoising diffusion baseline for comparison (`giant/model/schedule.py:CosineSchedule`).
|
||||||
|
- **`wgan`** — a single-pass Wasserstein-GAN-GP generator/critic (`giant/model/wgan.py`), trading iterative sampling for one forward pass, evaluated against a ~10× native-Geant4 latency budget.
|
||||||
|
|
||||||
|
**Stage 1 — primary (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):**
|
||||||
|
|
||||||
| Index | Variable | Encoding |
|
| Index | Variable | Encoding |
|
||||||
|-------|----------|----------|
|
|-------|----------|----------|
|
||||||
@@ -19,13 +23,20 @@ A **two-stage conditional flow matching** model (Lipman et al. 2022): a small ML
|
|||||||
| 3–5 | `post_dir` in local frame | unit vector |
|
| 3–5 | `post_dir` in local frame | unit vector |
|
||||||
| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
|
| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
|
||||||
|
|
||||||
The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss. Stage 1 also has a classifier head predicting the number of secondaries `n_sec ∈ {0..15}` from the conditioning alone.
|
The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss (`energy_simplex_decode`). Stage 1 also has a classifier head (`predict_n_sec`) predicting the number of secondaries `n_sec ∈ {0..K_MAX}` (`K_MAX = 15`) from the conditioning alone, no diffusion noise involved.
|
||||||
|
|
||||||
Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric. `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, so the two stay consistent by construction instead of being learned (and potentially diverging) independently.
|
Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric. `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, so the two stay consistent by construction instead of being learned (and potentially diverging) independently.
|
||||||
|
|
||||||
**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second flow net generates all `K_MAX = 15` secondary slots at once. Each slot carries a stick-breaking energy fraction, a local-frame direction, and a continuous particle-type embedding (snapped to the nearest PDG at inference), ordered by descending energy; slots beyond the predicted `n_sec` are masked. The secondary energies are a stick-breaking partition of the `e_sec` budget from Stage 1, so the full chain conserves energy. Each secondary's momentum is reconstructed afterward from `(energy, direction, species)` rather than predicted.
|
**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second net generates all `K_MAX` secondary slots at once — `(stick-breaking energy logit, local-frame direction, log-mass, charge)` per slot, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a stick-breaking partition of the `e_sec` budget from Stage 1, so the whole chain conserves energy. A secondary's mass/charge are regressed directly against its ground-truth PDG code's physical values (`giant.particles.particle_mass_charge`) and used as-is at inference — including for its own conditioning if it takes further steps in a rollout. No snapping to a known PDG code happens in the model path; `giant.particles.nearest_known_pdg` is a reporting-only lookup used to populate a nominal `pdg` label on output rows.
|
||||||
|
|
||||||
**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID. (`n_sec` / `e_sec` are outputs now, not inputs.)
|
**Conditioning (`--conditioning`, per-checkpoint):** pre-step position, log(pre-energy), pre-step direction, layer ID, plus particle/material physical properties — mass/charge (`giant/particles.py`) and Z_eff/A_eff/density/X0/λ_int (`giant/materials.py`). Two mutually exclusive modes:
|
||||||
|
|
||||||
|
- **`physical`** (default) — the physical-property columns are routed through small MLPs, computable for any PDG code / material, letting the surrogate generalize to species/materials outside the training menu.
|
||||||
|
- **`embedding`** — the original design: a learned `nn.Embedding` per PDG code / material, kept as a generalization-comparison baseline (memorizes the training menu).
|
||||||
|
|
||||||
|
`n_sec` and `e_sec` are model outputs, not conditioning inputs — a rollout is self-contained and never injects ground truth.
|
||||||
|
|
||||||
|
**Mixture-of-experts routing (`--router`, opt-in):** `giant/model/network.py` also implements a pluggable `Router` contract (`ROUTER_REGISTRY`: `energy`, `pdg`, `process`, plus a `composed` router combining several axes) that splits `DenoisingMLP`/`SecondaryDecoder` into per-expert trunks (`RoutedDenoisingMLP`/`RoutedSecondaryDecoder`), soft-gated in training and hard-dispatched at eval. See `--router-type`/`--n-experts`/`--router-axis` on `giant train`/`giant new-run`.
|
||||||
|
|
||||||
## Roadmap
|
## Roadmap
|
||||||
|
|
||||||
@@ -33,11 +44,13 @@ Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pr
|
|||||||
|
|
||||||
**Phase 2 (implemented — baseline):** the two-stage model above predicts `n_sec` and each secondary's energy, direction, and species jointly with the primary, so a shower rollout is fully self-contained.
|
**Phase 2 (implemented — baseline):** the two-stage model above predicts `n_sec` and each secondary's energy, direction, and species jointly with the primary, so a shower rollout is fully self-contained.
|
||||||
|
|
||||||
**Next directions:** faster-eval architectures against a ~10× native-Geant4 budget (Wasserstein-GAN, mixture-of-experts routing tree), a multi-material sampling-calorimeter dataset, and physical-property conditioning over learned embeddings.
|
**Physical-property conditioning (implemented):** replaces learned PDG/material embeddings with physical-property MLPs (see above); Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. Not yet done: the held-out-material/species generalization comparison against the `embedding` baseline.
|
||||||
|
|
||||||
|
**In progress:** a WGAN-GP single-pass mode (`--mode wgan`) and mixture-of-experts routing (`--router`) are implemented and under evaluation as faster-eval alternatives to iterative flow/DDPM sampling, against a ~10× native-Geant4 latency budget. A multi-material sampling-calorimeter dataset is the target for both the routing/WGAN evaluation and the physical-property generalization comparison.
|
||||||
|
|
||||||
## Data
|
## Data
|
||||||
|
|
||||||
Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `uv run dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle) to avoid leaking correlated steps from the same shower.
|
Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle) to avoid leaking correlated steps from the same shower.
|
||||||
|
|
||||||
## Project structure
|
## Project structure
|
||||||
|
|
||||||
@@ -49,21 +62,33 @@ giant/
|
|||||||
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode
|
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode
|
||||||
│ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
|
│ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
|
||||||
│ ├── model/
|
│ ├── model/
|
||||||
│ │ ├── network.py # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP, SecondaryDecoder
|
│ │ ├── network.py # ConditionEncoder, DenoisingMLP, SecondaryDecoder, Router/MoE, WGAN generator/critic
|
||||||
│ │ └── schedule.py # CosineSchedule (DDPM) and flow matching utilities
|
│ │ ├── schedule.py # CosineSchedule (DDPM) and flow matching utilities
|
||||||
|
│ │ └── wgan.py # WGAN-GP gradient penalty / critic / generator losses
|
||||||
│ ├── constants.py # output/conditioning dims, K_MAX, secondary slot layout, schema keys
|
│ ├── constants.py # output/conditioning dims, K_MAX, secondary slot layout, schema keys
|
||||||
|
│ ├── particles.py # PDG → (mass, charge) decode, incl. nuclear/ion codes; nearest-known-PDG lookup
|
||||||
|
│ ├── materials.py # material name → (Z_eff, A_eff, density, X0, λ_int)
|
||||||
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect
|
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect
|
||||||
│ ├── pipeline.py # builds datasets/normalizers and kicks off a training run
|
│ ├── pipeline.py # builds datasets/normalizers and kicks off a training run
|
||||||
│ ├── train.py # two-stage training loop, checkpointing, graceful shutdown
|
│ ├── train.py # two-stage training loop, checkpointing, graceful shutdown
|
||||||
│ ├── sample.py # DDPM / DDIM / flow matching samplers + secondary sampling
|
│ ├── sample.py # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
|
||||||
│ ├── geometry.py # GeometryOracle: position → (material, layer_id, escaped) for rollout
|
│ ├── geometry.py # GeometryOracle: position → (material, layer_id, escaped) for rollout
|
||||||
│ ├── rollout.py # autoregressive shower rollout driver
|
│ ├── rollout.py # autoregressive shower rollout driver
|
||||||
│ ├── validate.py # step-level marginal + KL-divergence validation
|
│ ├── validate.py # step-level marginal + KL-divergence validation
|
||||||
│ ├── analysis.py # step- and shower-level diagnostics: marginals, correlations, rollout observables
|
│ ├── condor.py # HTCondor submission for `train-submit`/`rollout-submit` (remote GPU workers)
|
||||||
│ └── cli.py # `giant train` / `predict` / `rollout` Typer app
|
│ ├── analysis/ # rollout-vs-reference analysis pipeline (see `giant analyze` below)
|
||||||
├── scripts/ # dataset/tooling logic, unified under the `dwarf` CLI (`uv run dwarf --help`)
|
│ │ ├── sources.py # canonical LazyFrames + secondary view
|
||||||
|
│ │ ├── reduce.py # streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
|
||||||
|
│ │ ├── grouping.py # fixed bin edges + energy/pdg/material group sets
|
||||||
|
│ │ ├── context.py # resolves grouping into `shared.json` once per run
|
||||||
|
│ │ ├── catalog.py # declarative PlotSpec registry
|
||||||
|
│ │ ├── condor.py # prep / compute-one / submit-description plumbing
|
||||||
|
│ │ └── render.py # PDFs + HTML gallery (only module importing plotstyle/LaTeX)
|
||||||
|
│ └── cli.py # `giant train` / `predict` / `rollout` / `new-run` / `*-submit` / `analyze` Typer app
|
||||||
|
├── scripts/ # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
|
||||||
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
|
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
|
||||||
│ │ # update-manifest, create-manifest, make-root, hparam-scan
|
│ │ # update-manifest, create-manifest, make-root, build-geometry-oracle,
|
||||||
|
│ │ # hparam-scan
|
||||||
│ ├── steps_to_parquet.py # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
|
│ ├── steps_to_parquet.py # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
|
||||||
│ ├── steps_to_parquet_parallel.py # fan out conversion over several ROOT files — `dwarf convert --jobs N`
|
│ ├── steps_to_parquet_parallel.py # fan out conversion over several ROOT files — `dwarf convert --jobs N`
|
||||||
│ ├── migrate_geant_steps.py # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
|
│ ├── migrate_geant_steps.py # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
|
||||||
@@ -80,27 +105,48 @@ giant/
|
|||||||
```bash
|
```bash
|
||||||
uv sync --extra cpu # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
|
uv sync --extra cpu # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
|
||||||
uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
|
uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
|
||||||
|
uv sync --extra cpu --extra geometry # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
|
||||||
```
|
```
|
||||||
|
|
||||||
`cpu` and `cuda` are mutually exclusive extras selecting the torch build; plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
`cpu` and `cuda` are mutually exclusive extras selecting the torch build (pinned to 2.3.x); plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
||||||
|
|
||||||
## Training, prediction, and rollout
|
## Training, prediction, and rollout
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
giant train path/to/steps.parquet --mode flow
|
giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir for a new run
|
||||||
|
giant train path/to/steps.parquet --mode flow # train (flow matching; also --mode ddpm / wgan)
|
||||||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||||||
|
|
||||||
# Full-shower rollout needs a geometry oracle (position → material/layer_id):
|
# Full-shower rollout needs a geometry oracle (position → material/layer_id):
|
||||||
uv sync --extra cpu --extra geometry
|
|
||||||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
|
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
|
||||||
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
||||||
```
|
```
|
||||||
|
|
||||||
`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant rollout` seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up `material`/`layer_id` from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
|
`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant rollout` seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up `material`/`layer_id` from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
|
||||||
|
|
||||||
## Validation
|
### Remote-GPU training and rollout (HTCondor)
|
||||||
|
|
||||||
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`). For deeper diagnostics on a trained checkpoint — stratified marginals, correlation structure, physical-constraint violations, and shower-level rollout observables (longitudinal/transverse profiles, PDG energy shares) — see `giant.analysis`.
|
`giant train-submit`/`giant rollout-submit` run a `giant train`/`giant rollout` invocation as a remote-GPU HTCondor job (ETP's TOpAS V100/A100 or NEMO2 L40S workers). GPU workers are remote-only and reach `/ceph` (not `/work`/`/home`) — the checkout submitting these jobs, and its venv, must live under `/ceph`.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
giant train-submit path/to/steps.parquet --config run/config.toml --accounting-group cms
|
||||||
|
giant rollout-submit path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl --accounting-group cms
|
||||||
|
```
|
||||||
|
|
||||||
|
Each submission writes a self-contained `condor/`/`condor_<tag>/` directory (`run.sh`, `job.sub`, `submission.json`). Training jobs re-check `out_dir/last.pt` on every invocation and resume automatically if preempted. `train-submit` requires `--config` (rather than mirroring `train`'s hyperparameter flags) since the TOML is the reproducible source of truth, and passing the same one on every retry is what keeps a resumed run's architecture from drifting.
|
||||||
|
|
||||||
|
## Validation and analysis
|
||||||
|
|
||||||
|
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`).
|
||||||
|
|
||||||
|
For deeper rollout-vs-reference diagnostics — marginals stratified by energy/pdg/material, per-event totals, shower profiles, species share, leakage, and secondaries — `giant analyze` runs a streaming compute/render pipeline against a `giant rollout` YAML sidecar:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
giant analyze submit rollout.yaml --accounting-group cms # prep + one HTCondor job per plot (compute only)
|
||||||
|
giant analyze render <run_dir> --gallery # local: styled PDFs + HTML gallery (needs LaTeX)
|
||||||
|
```
|
||||||
|
|
||||||
|
`<run_dir>` is derived next to the rollout parquet (`analyze prep`/`submit` print it). Compute jobs are polars/numpy only; only `render` imports plotstyle/LaTeX, so it always runs locally.
|
||||||
|
|
||||||
## Development
|
## Development
|
||||||
|
|
||||||
|
|||||||
+538
-25
@@ -1,9 +1,10 @@
|
|||||||
from collections import Counter
|
from collections import Counter
|
||||||
from datetime import date, datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
import math
|
import math
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import re
|
import re
|
||||||
|
import shlex
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
import uuid as uuid_mod
|
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"]
|
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:
|
def _coerce_scalar(value: str) -> object:
|
||||||
"""Best-effort str -> bool/int/float, else leave as str.
|
"""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")
|
_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).
|
"""Return (out_path, resolved_dataset_path, pred_uuid).
|
||||||
|
|
||||||
When *out* is None the output path is derived from *data*:
|
When *out* is None the output path is derived from *data*:
|
||||||
- under /ceph/ → fixed central store with a UUID filename
|
- under /ceph/ → fixed central store with a UUID filename
|
||||||
- elsewhere → sibling of *data* 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()
|
dataset_path = data.resolve()
|
||||||
pred_uuid = str(uuid_mod.uuid4())
|
pred_uuid = pred_uuid or str(uuid_mod.uuid4())
|
||||||
if out is None:
|
if out is None:
|
||||||
if str(dataset_path).startswith("/ceph/"):
|
if str(dataset_path).startswith("/ceph/"):
|
||||||
out = _CEPH_PREDICTIONS / f"{pred_uuid}.parquet"
|
out = _CEPH_PREDICTIONS / f"{pred_uuid}.parquet"
|
||||||
@@ -451,17 +482,7 @@ def train(
|
|||||||
}.items()
|
}.items()
|
||||||
if v is not None
|
if v is not None
|
||||||
}
|
}
|
||||||
cli_router: dict[str, object] = {
|
cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
|
||||||
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))
|
|
||||||
if cli_router:
|
if cli_router:
|
||||||
cli_model["router"] = cli_router
|
cli_model["router"] = cli_router
|
||||||
cfg = gconfig.merge_cli_overrides(
|
cfg = gconfig.merge_cli_overrides(
|
||||||
@@ -484,16 +505,7 @@ def train(
|
|||||||
f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
|
f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
|
||||||
)
|
)
|
||||||
|
|
||||||
out_dir = out or Path(
|
out_dir = out or gconfig.default_out_dir(cfg)
|
||||||
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']}"
|
|
||||||
)
|
|
||||||
|
|
||||||
typer.echo(f"device: {_device}")
|
typer.echo(f"device: {_device}")
|
||||||
typer.echo(f"out_dir: {out_dir}")
|
typer.echo(f"out_dir: {out_dir}")
|
||||||
@@ -510,6 +522,149 @@ 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 "<data.parquet>"
|
||||||
|
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 <group>"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@app.command()
|
@app.command()
|
||||||
def predict(
|
def predict(
|
||||||
data: Annotated[
|
data: Annotated[
|
||||||
@@ -998,6 +1153,15 @@ def rollout(
|
|||||||
Optional[int],
|
Optional[int],
|
||||||
typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
|
typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
|
||||||
] = None,
|
] = 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:
|
) -> None:
|
||||||
"""Roll the surrogate forward into full showers (autoregressive)."""
|
"""Roll the surrogate forward into full showers (autoregressive)."""
|
||||||
if seed is not None:
|
if seed is not None:
|
||||||
@@ -1051,7 +1215,9 @@ def rollout(
|
|||||||
seeds = _seed_from_data(files, n_events)
|
seeds = _seed_from_data(files, n_events)
|
||||||
typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)")
|
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)
|
out.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
# Written incrementally as each batch of steps is produced, rather than
|
# Written incrementally as each batch of steps is produced, rather than
|
||||||
@@ -1317,5 +1483,352 @@ def analyze_submit(
|
|||||||
subprocess.run(["condor_submit", str(sub)], check=True)
|
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 = (
|
||||||
|
f"#!/bin/bash\nset -euo pipefail\ncd {q(str(repo_path))}\nexec {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__":
|
if __name__ == "__main__":
|
||||||
app()
|
app()
|
||||||
|
|||||||
+129
@@ -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 <cluster_id>``).
|
||||||
|
"""
|
||||||
|
|
||||||
|
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 (``<run_dir>/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 ``<run_dir>/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 <cluster_id>`` 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()))
|
||||||
+22
-1
@@ -2,7 +2,8 @@ import random
|
|||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
import tomllib
|
import tomllib
|
||||||
from datetime import datetime, timezone
|
import uuid
|
||||||
|
from datetime import date, datetime, timezone
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -294,6 +295,26 @@ def save_config(cfg: dict, out_dir: Path, meta: dict) -> None:
|
|||||||
(out_dir / "config.toml").write_text("\n".join(lines))
|
(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(
|
def build_run_meta(
|
||||||
data: Path,
|
data: Path,
|
||||||
seed: int,
|
seed: int,
|
||||||
|
|||||||
@@ -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 "<data.parquet>" 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 "<data.parquet>" 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()
|
||||||
@@ -58,6 +58,17 @@ def test_each_call_produces_a_distinct_uuid(tmp_path):
|
|||||||
assert uuid1 != uuid2
|
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):
|
def test_dataset_path_is_resolved(tmp_path):
|
||||||
data = tmp_path / "data.parquet"
|
data = tmp_path / "data.parquet"
|
||||||
_, dataset_path, _ = _resolve_prediction_output(data, None)
|
_, dataset_path, _ = _resolve_prediction_output(data, None)
|
||||||
|
|||||||
@@ -0,0 +1,216 @@
|
|||||||
|
"""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"
|
||||||
Reference in New Issue
Block a user