HydraGNN perf characterization + iterative sysopt-loop recipe (MI355X)

.claude/commands/init-cluster.md

@@ -67,6 +67,12 @@ echo "HTTPS_PROXY=${HTTPS_PROXY:-<unset>}"
 
 # 13. Existing SIF files
 find "$HOME" /scratch /projects /opt -maxdepth 4 -name "*.sif" 2>/dev/null | head -20
+
+# 14. Performance tooling (perf-analysis / perf-optimizer-loop recipes — optional)
+# Look for an existing shared perf-tools dir (omnistat + TraceLens venv "perf-inspect/").
+for d in /shared/*/tools /shared/perf-tools "$HOME/perf-tools"; do
+    [[ -d "$d/perf-inspect" || -d "$d/omnistat-src" ]] && echo "perf_tools: $d"
+done
 ```
 
 ## Step 2 — Present findings and ask user to confirm or choose
@@ -119,7 +125,12 @@ Explain the expected layout under that root:
 Show result of connectivity test. Ask to confirm.
 If no internet, ask for proxy settings.
 
-**Q9. Config file location**
+**Q9. Performance tooling** (optional — only for `perf-analysis` / `perf-optimizer-loop` recipes)
+If a perf-tools dir was discovered (step 14), show it and ask to confirm; otherwise ask for the path
+or leave blank. This is exported to perf scripts as `PERF_TOOLS_DIR`; expected layout under it:
+`perf-inspect/` (omnistat + TraceLens venv), `omnistat-src/`, `victoriametrics/victoria-metrics-prod`.
+
+**Q10. Config file location**
 - `.cluster-config.yaml` in this repo (recommended — per-checkout)
 - `~/.config/ai4science-studio/cluster.yaml` (user-level, shared across clones)
 
@@ -156,10 +167,15 @@ gpu:
 network:
   internet_access: <true|false>
   proxy: "<proxy URL or empty>"
+  mgmt_iface: "<interface for RCCL OOB, e.g. enp193s0f1np1>"  # discover: ip -o link show up
+  ib_hca: "<IB HCA devices, e.g. ionic_0,ionic_1,...>"        # discover: ibstat | grep CA
+  rccl_anp_plugin: ""    # path to librccl-anp.so on host (default: /opt/rocm/lib/librccl-anp.so)
+  libionic_path: ""      # path to libionic.so.1 (default: /usr/lib/x86_64-linux-gnu/libionic.so.1)
 
-rccl:
-  socket_ifname: "<interface for RCCL OOB, e.g. enp193s0f1np1>"  # discover: ip link show up
-  ib_hca: "<IB HCA devices, e.g. ionic_0,ionic_1,...>"           # discover: ibstat | grep CA
+# Performance tooling — only if using perf-analysis / perf-optimizer-loop recipes (Q9).
+# Exported to scripts as PERF_TOOLS_DIR. Leave dir blank if not used.
+perf_tools:
+  dir: "<perf-tools dir or empty>"   # layout: perf-inspect/, omnistat-src/, victoriametrics/
 
 discovered_sifs:
   # SIF files found during init (informational — confirm they are under sif_cache)
@@ -177,6 +193,7 @@ After writing, confirm to the user:
 Suggest what the user can do next:
 - Export `AI4S_SHARED_DIR=<scratch>` (or add to `~/.bashrc`) so all model scripts resolve paths correctly
 - Run a model: `/run-stormcast`, `/run-orbit2`, etc.
+- If **Q9 (Performance tooling)** was configured: run ORBIT-2 or HydraGNN perf recipes with `/run-perf-orbit2` or `/run-perf-hydragnn`
 - The run commands will read cluster config automatically for SLURM partition/account defaults
 
 $ARGUMENTS

.claude/commands/run-perf-hydragnn.md

@@ -0,0 +1,146 @@
+# Run HydraGNN perf-analysis or perf-optimizer-loop on an AMD cluster
+
+Guide the user through **performance engineering** for HydraGNN Predictive GFM: one-shot **perf-analysis** (2-node training + TraceLens + Omnistat diagnosis) or iterative **perf-optimizer-loop** (accept/revert on **`epoch_time_s`**). This is **not** ensemble inference — use `/run-hydragnn` for that.
+
+**Recipe docs:** [material_science/models/HydraGNN/recipes/perf-analysis/README.md](../material_science/models/HydraGNN/recipes/perf-analysis/README.md), [material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md](../material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md).
+
+**Attribution / levers:** [dispatch-attribution.md](../material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md), [lever_catalog.yaml](../material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml).
+
+**Orchestration reference:** `.cursor/skills/ai4science-perf-analysis/SKILL.md` (launcher → parallel analysts → parallel verifiers → synthesizer).
+
+## Step 0 — Cluster config check (required)
+
+Read `.cluster-config.yaml` (repo root) or `~/.config/ai4science-studio/cluster.yaml`.
+
+**Required for perf recipes:**
+
+- Set **`AI4S_SHARED_DIR`** from cluster config (`paths.scratch` or documented scratch root).
+- Set **`PERF_TOOLS_DIR`** from `perf_tools.dir`. If missing, stop and send the user to **`/init-cluster`** Q9 (Performance tooling).
+
+Pre-fill SLURM partition/account from config.
+
+**Runtime:** Apptainer + SLURM only (`sbatch_train_perf_amd.sh`).
+
+---
+
+## Step 1 — Questionnaire (ask ALL before acting)
+
+**Q0. Mode**
+
+- **perf-analysis** — One 2-node perf job + post-hoc multi-subagent analysis (default topology for this recipe).
+- **perf-optimizer-loop** — Iterative tuning; primary FOM **`epoch_time_s`** (lower is better); see `recipes/perf-optimizer-loop/README.md` for stop rules and `do_not_retry.json`.
+
+**Q1. HydraGNN checkout (`HG_REPO_DIR`)**
+
+Default in sbatch: `$AI4S_SHARED_DIR/models/HydraGNN/code/HydraGNN` (cloned at pinned SHA if missing).
+
+- **Use default** / **Yes, custom path** / **No** (clone instructions) / **Auto-discover** — `find "$HOME" /scratch /projects /opt -maxdepth 5 -type d -name "HydraGNN" 2>/dev/null | head -20` and confirm `examples/multidataset_hpo_sc26` exists.
+
+**Q2. SIF (`HG_SIF`)**
+
+Default: `$AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif`.
+
+- **Yes** / **No** / **Auto-discover** — same SIF discovery pattern as `/run-orbit2`.
+
+**Q3. Overlay (`HG_OVERLAY`)**
+
+Default: `$AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img`.
+
+- **Yes** / **No, build one** — `sbatch material_science/models/HydraGNN/examples/build_overlay_amd.sh` (edit `#SBATCH` partition/account first) / **Auto-discover** — `find ... -name "*hydragnn*overlay*"`.
+
+**Q4. Datasets (`HG_DATASETS`, `HG_DATA_DIR`)**
+
+Default datasets and paths are set in `sbatch_train_perf_amd.sh` (`ANI1x,Alexandria` under `HG_DATA_DIR`). Confirm `.bp` files exist or user will stage per [recipes/data/README.md](../material_science/models/HydraGNN/recipes/data/README.md).
+
+**Q5. Precision**
+
+Perf recipe defaults to **`fp64`** / **`HG_PRECISION=fp64`** (MACE-style equivariance). **`bf16` is not used** for this GFM path — do not suggest bf16 as a lever; the lever catalog documents blocked options.
+
+**Q6. Blocked or high-risk levers (set expectations)**
+
+Point users to `lever_catalog.yaml` entries marked **`blocked`** or **`accepted-baked-in`**, including at minimum:
+
+- **`torch.compile`** on HydraGNN MLIP MACE — blocked (double-backward / AOT autograd).
+- **`torch.jit.script`** — blocked (keyword-arg expansion at MACE interaction block; see lever catalog).
+- **TunableOp live or warmup-then-use** — blocked on this stack (documented GPU faults during tuning); do not re-propose without a stack update and catalog change.
+
+**Q7. Partition and account**
+
+- **Provide manually** / **Auto-discover** — `sinfo` + `sacctmgr` as in `/run-orbit2`.
+
+**Q8. (perf-optimizer-loop only) Iteration budget**
+
+Default **5**; run inside **`tmux`**. Optional: `export ANTHROPIC_API_KEY=...` for Claude Code CLI orchestration inside `run_optimizer_loop.sh`.
+
+---
+
+## Step 2 — Auto-discovery
+
+Reuse SIF/overlay/SLURM snippets from `.claude/commands/run-orbit2.md`.
+
+**HydraGNN repo:**
+
+```bash
+find "$HOME" /scratch /projects /opt -maxdepth 5 -type d -name "HydraGNN" 2>/dev/null | head -20
+```
+
+Confirm branch **`Predictive_GFM_2024`** if user cares about upstream alignment (see model README).
+
+---
+
+## Step 3 — Run
+
+```bash
+export AI4S_SHARED_DIR=<from cluster config>
+export PERF_TOOLS_DIR=<from cluster config perf_tools.dir>
+# Optional overrides:
+# export HG_SIF=...  HG_OVERLAY=...  HG_REPO_DIR=...  HG_DATASETS=...  HG_DATA_DIR=...
+```
+
+Edit **`#SBATCH`** in **`material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh`** for partition/account.
+
+**Default perf-analysis job is 2-node** — ensure `#SBATCH --nodes=2` (or user override) matches intent.
+
+### perf-analysis
+
+```bash
+sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh
+```
+
+After completion, drive subagents per **`ai4science-perf-analysis` SKILL** using prompts under `material_science/models/HydraGNN/recipes/perf-analysis/agents/`.
+
+### perf-optimizer-loop
+
+Preflight only (smoke; no Claude invocation):
+
+```bash
+bash material_science/models/HydraGNN/examples/run_optimizer_loop.sh "$(uuidgen)" 5 --preflight-only
+```
+
+Full loop:
+
+```bash
+tmux new -s hg-perf-loop
+bash material_science/models/HydraGNN/examples/run_optimizer_loop.sh "$(uuidgen)" 5
+```
+
+Graceful stop: `touch "$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/loop-<uuid>/STOP"`.
+
+---
+
+## Step 4 — Monitor and artifacts
+
+```bash
+squeue -j <job_id>
+tail -f "$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/<job_id>/"*.out
+```
+
+**perf-analysis:** `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/<jobid>/` — `combined_report.md`, `manifest.json`, `tracelens/`, `omnistat/`, traces under `logs/`.
+
+**perf-optimizer-loop:** `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/loop-<uuid>/` — `STATUS.txt`, `foms.csv`, `story.md`, per-iteration job subdirs.
+
+---
+
+## Arguments
+
+$ARGUMENTS

.claude/commands/run-perf-orbit2.md

@@ -0,0 +1,164 @@
+# Run ORBIT-2 perf-analysis or perf-optimizer-loop on an AMD cluster
+
+Guide the user through **performance engineering** workflows for ORBIT-2 (Bayes-CAST EDM / same-dir timing configs): one-shot **perf-analysis** (TraceLens + Omnistat diagnosis) or iterative **perf-optimizer-loop** (accept/revert on throughput). This is **not** inference/visualization — use `/run-orbit2` for that.
+
+**Recipe docs:** [earth_science/models/ORBIT-2/recipes/perf-analysis/README.md](../earth_science/models/ORBIT-2/recipes/perf-analysis/README.md), [earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md](../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md).
+
+**Orchestration reference (Cursor / any agent):** read `.cursor/skills/ai4science-perf-analysis/SKILL.md` for the launcher → analyst → verifier → synthesizer loop.
+
+## Step 0 — Cluster config check (required)
+
+Read `.cluster-config.yaml` (repo root) or `~/.config/ai4science-studio/cluster.yaml`.
+
+**Required for perf recipes:**
+
+- `paths.scratch` (or equivalent) → use as **`AI4S_SHARED_DIR`** (same convention as other studio scripts).
+- `perf_tools.dir` → use as **`PERF_TOOLS_DIR`** (layout: `perf-inspect/`, `omnistat-src/`, `victoriametrics/`). If missing or empty, stop and tell the user to re-run **`/init-cluster`** and answer **Q9 (Performance tooling)**.
+
+Pre-fill SLURM partition/account from `slurm.partition` / `slurm.account` when present.
+
+**Runtime note:** Perf workflows here assume **Apptainer + SLURM** (same as `sbatch_train_perf_amd.sh`). There is no Docker perf path in the recipes.
+
+---
+
+## Step 1 — Questionnaire (ask ALL questions before doing anything)
+
+Ask the following. Do not assume defaults. Wait for answers before proceeding.
+
+**Q0. Mode**
+
+- **perf-analysis** — Submit one training perf job; after it finishes, drive subagents per `ai4science-perf-analysis` SKILL (`agents/launcher.md`, analysts, verifiers, `synthesizer.md`) to produce `combined_report.md`.
+- **perf-optimizer-loop** — Iterative loop: one lever per iteration, accept/revert on **`throughput_samples_per_s`** with **`loss_sanity_pass`** as a control (see perf-optimizer-loop README).
+
+**Q1. Code root (`ORBIT2_ROOT`)**
+
+Bayes-CAST checkout used for training (must contain `launch/train_edm.py` for the default perf path).
+
+- **Yes** — user provides full path.
+- **No** — generate clone instructions (user may use Bayes-CAST fork or path documented in perf-analysis README).
+- **Auto-discover** — search for a directory containing `launch/train_edm.py`.
+
+**Q2. SIF**
+
+Apptainer image for ROCm PyTorch `rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0` (or user override).
+
+- **Yes** / **No** (pull command) / **Auto-discover** — same pattern as `/run-orbit2`.
+
+**Q3. ORBIT-2 overlay**
+
+Pre-built `orbit2-overlay.img` (recommended).
+
+- **Yes** / **No, build one** (`sbatch earth_science/models/ORBIT-2/examples/build_overlay_amd.sh`) / **Auto-discover** — same pattern as `/run-orbit2`.
+
+**Q4. Data**
+
+- **`ORBIT2_DATA_ROOT`** — Staged ERA5 NPZ tree (perf defaults use Bayes-CAST template + ERA5; see perf-analysis README).
+- **`ORBIT2_CONFIG_TEMPLATE`** — Optional; default **`edm_8m_era5_1x8.yaml`** in perf docs. For HBM saturation / staging notes: [STAGING_ERA5_FOR_HBM.md](../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md).
+- **`interm_8m_prism.yaml` / `interm_8m_era5.yaml`** — Public ORBIT-2 same-dir configs (timing-only); document if used instead of Bayes-CAST EDM.
+
+**Q5. Topology**
+
+- **1 node × 8 GPUs** (default for phase-1 loop and perf-analysis quick start).
+- **Multi-node** — e.g. `sbatch --nodes=2 ...`; RCCL / GID / `LD_LIBRARY_PATH` caveats in perf-analysis README and `sbatch_train_perf_amd.sh`. After one-node loop converges, see perf-optimizer-loop README **Two-node gate**; optional driver: `earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh`.
+
+**Q6. Partition and account**
+
+- **Provide manually** / **Auto-discover** (same as `/run-orbit2`).
+
+**Q7. (perf-optimizer-loop only) Iteration budget**
+
+Max iterations (default **5**). Confirm use of **`tmux`** on the login node so SSH disconnect does not kill the driver.
+
+**Q8. (perf-optimizer-loop only) Claude Code CLI driver**
+
+`run_optimizer_loop.sh` can invoke the orchestrator via Claude Code CLI when **`ANTHROPIC_API_KEY`** is set. If the user has no API key, explain they can still run **preflight-only** steps and manual `sbatch` per iteration using `recipes/perf-optimizer-loop/agents/orchestrator.md` as a checklist.
+
+---
+
+## Step 2 — Auto-discovery (when chosen)
+
+**Code root containing `launch/train_edm.py`:**
+
+```bash
+find "$HOME" /scratch /projects /opt -maxdepth 6 -type f -path "*/launch/train_edm.py" 2>/dev/null | head -20
+```
+
+Derive `ORBIT2_ROOT` as the parent of `launch/`.
+
+**SIF / overlay / SLURM:** reuse the same `find` / `sinfo` / `sacctmgr` snippets as in `.claude/commands/run-orbit2.md` (SIF: `*.sif` filtered for rocm/pytorch; overlay: `*orbit2*overlay*`).
+
+Always confirm discovered paths with the user before `sbatch`.
+
+---
+
+## Step 3 — Run
+
+Export from cluster config + user answers:
+
+```bash
+export AI4S_SHARED_DIR=<from cluster config paths.scratch>
+export PERF_TOOLS_DIR=<from cluster config perf_tools.dir>
+export ORBIT2_ROOT=<path to Bayes-CAST or ORBIT-2 tree with launch/train_edm.py>
+export ORBIT2_SIF=<path>                    # if not using default under $AI4S_SHARED_DIR/images/...
+export ORBIT2_OVERLAY=<path>                # optional if script defaults suffice
+export ORBIT2_DATA_ROOT=<staged era5 root>
+# export ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml   # optional
+```
+
+Edit **`#SBATCH`** lines in **`earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh`** for partition/account (and `--nodes` if multi-node).
+
+### perf-analysis
+
+```bash
+# From ai4science-studio repo root
+sbatch earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh
+```
+
+After the job completes, follow **`.cursor/skills/ai4science-perf-analysis/SKILL.md`**: dispatch launcher (if not already folded into job), `tracelens_analyst` + `omnistat_analyst` in parallel, verifiers in parallel, then `synthesizer`. Prompts live under `earth_science/models/ORBIT-2/recipes/perf-analysis/agents/`.
+
+### perf-optimizer-loop
+
+Optional repo file check (does not validate cluster env):
+
+```bash
+bash earth_science/models/ORBIT-2/examples/validate_orbit2_optimizer_loop_recipe.sh
+```
+
+Preflight only (recommended first):
+
+```bash
+bash earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh "$(uuidgen)" 5 --preflight-only
+```
+
+Full loop (inside **`tmux`**):
+
+```bash
+tmux new -s orbit2-perf-loop
+bash earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh "$(uuidgen)" 5
+# detach: Ctrl-b d
+```
+
+Graceful stop: `touch "$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/loop-<uuid>/STOP"`.
+
+---
+
+## Step 4 — Monitor and artifacts
+
+**SLURM job:**
+
+```bash
+squeue -j <job_id>
+tail -f "$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/<job_id>/"*.out
+```
+
+**perf-analysis deliverables** (under `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/<jobid>/`): `manifest.json`, `omnistat-db/`, `tracelens/`, `combined_report.md`, `foms.json` (when extractor run), traces under `traces/`.
+
+**perf-optimizer-loop** (under `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/loop-<uuid>/`): `STATUS.txt`, `foms.csv`, `story.md`, `do_not_retry.json`, `known_bad_nodes.txt`; per-iter job dirs mirror perf-analysis layout.
+
+**Primary FOM (loop):** `throughput_samples_per_s` (higher is better); keep **`loss_sanity_pass`**. GEMM / lever evidence: [gemm-attribution.md](../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/gemm-attribution.md), [lever_catalog.yaml](../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/lever_catalog.yaml).
+
+---
+
+## Arguments
+
+$ARGUMENTS

.cluster-config.example.yaml

@@ -40,3 +40,10 @@ network:
   ib_hca: ""             # IB HCA device list for RCCL (discover: ibstat | grep 'CA ')
   rccl_anp_plugin: ""    # Path to librccl-anp.so on host (default: /opt/rocm/lib/librccl-anp.so)
   libionic_path: ""      # Path to libionic.so.1 on host (default: /usr/lib/x86_64-linux-gnu/libionic.so.1)
+
+# Performance tooling (perf-analysis / perf-optimizer-loop recipes)
+# Shared location of the omnistat + TraceLens perf tools, exported to scripts as
+# PERF_TOOLS_DIR. Expected layout under this dir: perf-inspect/ (Python venv with
+# omnistat + TraceLens), omnistat-src/, victoriametrics/victoria-metrics-prod.
+perf_tools:
+  dir: ""                # e.g. "/shared/perf-tools"; leave blank if not using perf recipes