From a044866ef16456b593deb818a6a9cedb3db847ff Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 18 Jun 2026 08:45:37 +0000 Subject: [PATCH 1/4] Fix HydraGNN launcher.md placeholder name to match omnistat template The doc referenced a @OMNIHUB_TOOLS_DIR@ placeholder that does not exist; the template and sbatch sed substitution use @PERF_TOOLS_DIR@. Co-Authored-By: Claude --- .../models/HydraGNN/recipes/perf-analysis/agents/launcher.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md index 63ea2dc..20888c5 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -127,7 +127,7 @@ The sbatch wrapper expects the `.so` at `${PERF_TOOLS_DIR}/omnistat-src/build-tr ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) -If `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/omnistat.config` doesn't exist, write it from the template at `$REPO_ROOT/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template`. The template uses `@JOB_DIR@` / `@OMNIHUB_TOOLS_DIR@` placeholders, which the sbatch wrapper substitutes (via `sed`) into a per-job config under the perf-run directory. +If `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/omnistat.config` doesn't exist, write it from the template at `$REPO_ROOT/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template`. The template uses `@JOB_DIR@` / `@PERF_TOOLS_DIR@` placeholders, which the sbatch wrapper substitutes (via `sed`) into a per-job config under the perf-run directory. ### 3. Generate the per-job profile config From 83e9e9f9d8aa92dcf3a06a66d0f3c35ed639558c Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 18 Jun 2026 08:56:49 +0000 Subject: [PATCH 2/4] ORBIT-2: honour ORBIT2_MAX_BATCHES on the Bayes-CAST EDM training path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The EDM path exec'd launch/train_edm.py directly, whose batch loop is inline in main() with no cap — so ORBIT2_MAX_BATCHES was silently ignored (unlike the studio intermediate_downscaling path via run_orbit2_train.py, and HydraGNN). Add run_orbit2_train_edm.py, which wraps train_edm.main() and caps each epoch by patching IterDataModule.train_dataloader. Point both sbatch scripts' EDM exec at the wrapper; the profiler hook-runner call is unchanged. Co-Authored-By: Claude --- .../ORBIT-2/examples/run_orbit2_train_edm.py | 142 ++++++++++++++++++ .../ORBIT-2/examples/sbatch_train_amd.sh | 4 +- .../ORBIT-2/examples/sbatch_train_perf_amd.sh | 4 +- 3 files changed, 148 insertions(+), 2 deletions(-) create mode 100755 earth_science/models/ORBIT-2/examples/run_orbit2_train_edm.py diff --git a/earth_science/models/ORBIT-2/examples/run_orbit2_train_edm.py b/earth_science/models/ORBIT-2/examples/run_orbit2_train_edm.py new file mode 100755 index 0000000..e554911 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/run_orbit2_train_edm.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Launch Bayes-CAST EDM launch/train_edm.py with optional training caps. + +train_edm.py runs its batch loop inline in main() (no run_training_epochs to +patch like the studio intermediate_downscaling path). So the cap is applied at +the dataloader: IterDataModule.train_dataloader is wrapped to yield at most +ORBIT2_MAX_BATCHES batches per epoch. Mirrors run_orbit2_train.py. + +Config path is passed as argv[1] after this script's own args. + +Environment: + ORBIT2_ROOT Bayes-CAST checkout (must contain launch/train_edm.py) + ORBIT2_MAX_BATCHES Cap batches per epoch (0 = unlimited) +""" + +from __future__ import annotations + +import itertools +import os +import sys +import types +from datetime import timedelta +from pathlib import Path + + +def _stub_gptl4py() -> None: + """train_edm may import gptl4py (Frontier GPTL); stub for ROCm containers.""" + if "gptl4py" in sys.modules: + return + + class _Stub: + @staticmethod + def start(_name: str) -> None: + return None + + @staticmethod + def stop(_name: str) -> None: + return None + + mod = types.ModuleType("gptl4py") + mod.start = _Stub.start + mod.stop = _Stub.stop + sys.modules["gptl4py"] = mod + + +def _cap_dataloader(dl, limit: int): + class _Capped: + def __init__(self, inner, n: int): + self._inner = inner + self._n = n + + def __iter__(self): + return iter(itertools.islice(iter(self._inner), self._n)) + + def __len__(self): + try: + return min(len(self._inner), self._n) + except TypeError: + return self._n + + return _Capped(dl, limit) + + +def _apply_batch_cap() -> None: + max_batches = int(os.environ.get("ORBIT2_MAX_BATCHES", "0") or 0) + if max_batches <= 0: + return + + from climate_learn.data.itermodule import IterDataModule + + _orig = IterDataModule.train_dataloader + + def _capped(self, *args, **kwargs): + dl = _orig(self, *args, **kwargs) + if int(os.getenv("SLURM_PROCID", "0")) == 0: + print(f"[run_orbit2_train_edm] ORBIT2_MAX_BATCHES cap: {max_batches}", flush=True) + return _cap_dataloader(dl, max_batches) + + IterDataModule.train_dataloader = _capped + + +def main() -> None: + if len(sys.argv) < 2: + print("usage: run_orbit2_train_edm.py ", file=sys.stderr) + sys.exit(2) + + config = sys.argv[1] + orbit2_root = os.environ.get("ORBIT2_ROOT") + if not orbit2_root: + print("error: ORBIT2_ROOT must be set", file=sys.stderr) + sys.exit(2) + + launch_dir = Path(orbit2_root) / "launch" + if not (launch_dir / "train_edm.py").is_file(): + print(f"error: missing launch/train_edm.py under: {orbit2_root}", file=sys.stderr) + sys.exit(2) + + src_dir = Path(orbit2_root) / "src" + for p in (str(src_dir), str(launch_dir), orbit2_root): + if p not in sys.path: + sys.path.insert(0, p) + + os.chdir(launch_dir) + _stub_gptl4py() + + import torch + import torch.distributed as dist + + import train_edm # noqa: WPS433 + + _apply_batch_cap() + + # train_edm.main() reads sys.argv[1] as the config path. + sys.argv = ["train_edm.py", config] + + os.environ["MASTER_ADDR"] = str(os.environ["HOSTNAME"]) + os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "29500") + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + os.environ["RANK"] = os.environ["SLURM_PROCID"] + + local_rank = int(os.environ["SLURM_LOCALID"]) + os.environ["LOCAL_RANK"] = str(local_rank) + torch.cuda.set_device(local_rank) + device = torch.cuda.current_device() + + dist.init_process_group( + "nccl", + timeout=timedelta(seconds=7200000), + rank=int(os.environ["SLURM_PROCID"]), + world_size=int(os.environ["SLURM_NTASKS"]), + ) + + print("Using dist.init_process_group. world_size ", os.environ["SLURM_NTASKS"], flush=True) + + try: + train_edm.main(device) + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh index 723df79..aec09fd 100755 --- a/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh @@ -253,7 +253,9 @@ export ORBIT2_RANK_PRE_TRAIN_HOOK="\${ORBIT2_RANK_PRE_TRAIN_HOOK:-}" if [[ "${LAUNCH_EDM_DIRECT}" -eq 1 ]]; then cd /orbit2/launch python3 /examples/orbit2_rank_hook_runner.py - exec python3 train_edm.py /config/config.yaml + # run_orbit2_train_edm.py wraps train_edm.main() to honour ORBIT2_MAX_BATCHES + # (train_edm.py runs its batch loop inline, so there is no function to patch). + exec python3 /examples/run_orbit2_train_edm.py /config/config.yaml elif [[ -n "${LAUNCH_IC}" ]]; then cd /orbit2/examples python3 /examples/orbit2_rank_hook_runner.py diff --git a/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh index 7d2f093..751fef9 100755 --- a/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh @@ -328,7 +328,9 @@ export ORBIT2_RANK_PRE_TRAIN_HOOK="${_RANK_PRE_TRAIN_HOOK}" if [[ "${LAUNCH_EDM_DIRECT}" -eq 1 ]]; then cd /orbit2/launch python3 /examples/orbit2_rank_hook_runner.py - exec python3 train_edm.py /config/config.yaml + # run_orbit2_train_edm.py wraps train_edm.main() to honour ORBIT2_MAX_BATCHES + # (train_edm.py runs its batch loop inline, so there is no function to patch). + exec python3 /examples/run_orbit2_train_edm.py /config/config.yaml elif [[ -n "${LAUNCH_IC}" ]]; then cd /orbit2/examples python3 /examples/orbit2_rank_hook_runner.py From a27b8f06a7f5b4823c8fd542443198c226295cab Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 18 Jun 2026 19:19:06 +0000 Subject: [PATCH 3/4] ORBIT-2 perf-analysis: align agent prompts with sbatch manifest + venv The analyst/verifier prompts read manifest keys that the sbatch launcher does not write (perf_run_dir/jobid/trace_paths/omnistat_db_path) and hardcoded the perf-inspect venv name, so a fresh run needed manual fix-ups. Align to the actual keys (job_dir/job_id/trace_dir/omnistat_db), resolve the venv via OMNISTAT_VENV with the perf-inspect default, document that omnistat-inspect writes its payload to --scratch-dir (stdout is query-stats only), and fix the verifier's PromQL example to use the rmsjob_info join. Co-Authored-By: Claude --- .../recipes/perf-analysis/agents/omnistat_analyst.md | 12 +++++++----- .../perf-analysis/agents/omnistat_verifier.md | 5 +++-- .../perf-analysis/agents/tracelens_analyst.md | 8 +++++--- .../perf-analysis/agents/tracelens_verifier.md | 2 +- 4 files changed, 16 insertions(+), 11 deletions(-) diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md index 5ff92cf..d97cd5d 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md @@ -5,8 +5,8 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo ## Inputs - `/manifest.json` -- The Omnistat DB at `manifest.omnistat_db_path`. -- The `perf-inspect` venv at `${PERF_TOOLS_DIR}/perf-inspect/`. +- The Omnistat DB at `manifest.omnistat_db`. +- The omnistat/TraceLens venv at `${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}/` (the venv name is site-specific; honour `OMNISTAT_VENV` when set at submit time). - The VictoriaMetrics binary at `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -27,7 +27,7 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo Following the [load-database SKILL](https://github.com/ROCm/omnistat/blob/jorda/skills/skills/load-database/SKILL.md), with one addition for memory-constrained login nodes: pass `-fs.disableMmap` so VM can open a job-scoped DB inside the login-node cgroup (without it, even a 1.6 MB DB panics with "cannot mmap file with size 4096 bytes ... no such device"). ```bash -DB=$(jq -r .omnistat_db_path ) +DB=$(jq -r .omnistat_db ) # sbatch manifest key is omnistat_db (not omnistat_db_path) [[ -d "$DB/data" ]] || { echo "STATUS=fail; reason=no_db_data"; exit 1; } # Pick free port @@ -62,10 +62,10 @@ done ### 2. Run analyze-job phases (per the SKILL) ```bash -. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate +. "${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}/bin/activate" SCRATCH=$PERF_RUN/omnistat/scratch mkdir -p "$SCRATCH" -JOBID=$(jq -r .jobid ) +JOBID=$(jq -r .job_id ) # sbatch manifest key is job_id (not jobid) OUTD=$PERF_RUN/omnistat/inspect_outputs @@ -76,6 +76,8 @@ omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID health > omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID stats --level global > $OUTD/stats_global.json ``` +**Where the real payload lands:** for `job ... info|data-check|health|stats`, stdout (the `>` redirect above) captures only a small *query-stats* blob — `omnistat-inspect` writes the full analysis JSON under `--scratch-dir`. After the commands run, copy the scratch payloads into `$OUTD/` so the `evidence_paths` in your `claims.json` resolve to real data (e.g. `cp "$SCRATCH"//*.json "$OUTD/"`, matching each phase's output file). Verify each `$OUTD/*.json` actually contains the analysis (not just `query_stats`) before writing claims against it. + ### 3. Identify anomalous categories and drill down For each category in `stats_global.json` (`gpu`, `host`, `network`, `xgmi`, `vendor`): diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md index 1919bad..6a8fae0 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md @@ -40,9 +40,10 @@ Also: `start_time`/`end_time` from `job_info.json` are **UTC**. Convert with `ca | "Power throttling events" | sum increase of throttle counters over the job span, joined as above | ```bash -PORT=$(jq -r '.[0].port' ) # or pull from analyst stdout +PORT=$(cat "$PERF_RUN/omnistat/vm.pid" >/dev/null 2>&1 && grep -oE ':[0-9]+' "$PERF_RUN/omnistat/vm.log" | head -1 | tr -d :) # or take the port from the analyst's STATUS line +# Must join via rmsjob_info — a bare {jobid="..."} on rocm_* returns 0 series (see note above). curl -sG "http://127.0.0.1:$PORT/api/v1/query_range" \ - --data-urlencode "query=avg_over_time(rocm_utilization_percentage{jobid=\"$JOBID\"}[1m])" \ + --data-urlencode "query=avg(rocm_utilization_percentage * on (instance) group_left() (max by (instance) (rmsjob_info{jobid=\"$JOBID\"})))" \ --data-urlencode "start=$START" \ --data-urlencode "end=$END" \ --data-urlencode "step=$STEP" \ diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md index ea556d0..0522aaf 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md @@ -39,9 +39,11 @@ Every entry: ### 1. Sanity check ```bash -. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate -PERF_RUN=$(jq -r .perf_run_dir ) -TRACE=$(jq -r '.trace_paths[0] // empty' ) +# Venv name is site-specific; honour OMNISTAT_VENV (set at submit time), else the portable default. +. "${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}/bin/activate" +# sbatch-written manifest uses job_dir / trace_dir (not perf_run_dir / trace_paths[]). +PERF_RUN=$(jq -r .job_dir ) +TRACE=$(find "$(jq -r .trace_dir )" -name '*.pt.trace.json' 2>/dev/null | sort | head -1) [[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } mkdir -p "$PERF_RUN/tracelens" ``` diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md index 8861684..8205808 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md @@ -6,7 +6,7 @@ Independently re-derive the top 2-3 claims from `tracelens/claims.json` against - `/manifest.json` - `/tracelens/claims.json` -- The raw trace at `manifest.trace_paths[0]`. +- The raw trace: the single `*.pt.trace.json` under `manifest.trace_dir` (find with `find "$(jq -r .trace_dir )" -name '*.pt.trace.json' | sort | head -1`). ## Outputs From 255daf73bfe2928efda8e235dd8241f73b3cb0d4 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 18 Jun 2026 12:26:57 -0700 Subject: [PATCH 4/4] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../models/ORBIT-2/examples/run_orbit2_train_edm.py | 8 ++++---- .../recipes/perf-analysis/agents/omnistat_verifier.md | 3 ++- .../recipes/perf-analysis/agents/tracelens_verifier.md | 2 +- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/earth_science/models/ORBIT-2/examples/run_orbit2_train_edm.py b/earth_science/models/ORBIT-2/examples/run_orbit2_train_edm.py index e554911..f7976ba 100755 --- a/earth_science/models/ORBIT-2/examples/run_orbit2_train_edm.py +++ b/earth_science/models/ORBIT-2/examples/run_orbit2_train_edm.py @@ -49,6 +49,9 @@ def __init__(self, inner, n: int): self._inner = inner self._n = n + def __getattr__(self, name: str): + return getattr(self._inner, name) + def __iter__(self): return iter(itertools.islice(iter(self._inner), self._n)) @@ -57,7 +60,6 @@ def __len__(self): return min(len(self._inner), self._n) except TypeError: return self._n - return _Capped(dl, limit) @@ -102,14 +104,12 @@ def main() -> None: os.chdir(launch_dir) _stub_gptl4py() + _apply_batch_cap() import torch import torch.distributed as dist import train_edm # noqa: WPS433 - - _apply_batch_cap() - # train_edm.main() reads sys.argv[1] as the config path. sys.argv = ["train_edm.py", config] diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md index 6a8fae0..f324534 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md @@ -40,7 +40,8 @@ Also: `start_time`/`end_time` from `job_info.json` are **UTC**. Convert with `ca | "Power throttling events" | sum increase of throttle counters over the job span, joined as above | ```bash -PORT=$(cat "$PERF_RUN/omnistat/vm.pid" >/dev/null 2>&1 && grep -oE ':[0-9]+' "$PERF_RUN/omnistat/vm.log" | head -1 | tr -d :) # or take the port from the analyst's STATUS line +PORT=$(grep -oE '127\.0\.0\.1:[0-9]+' "$PERF_RUN/omnistat/vm.log" 2>/dev/null | head -1 | cut -d: -f2) # or take the port from the analyst's STATUS line +[[ -z "$PORT" ]] && { echo "STATUS=fail; reason=cannot_determine_vm_port"; exit 1; } # Must join via rmsjob_info — a bare {jobid="..."} on rocm_* returns 0 series (see note above). curl -sG "http://127.0.0.1:$PORT/api/v1/query_range" \ --data-urlencode "query=avg(rocm_utilization_percentage * on (instance) group_left() (max by (instance) (rmsjob_info{jobid=\"$JOBID\"})))" \ diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md index 8205808..db068d9 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md @@ -6,7 +6,7 @@ Independently re-derive the top 2-3 claims from `tracelens/claims.json` against - `/manifest.json` - `/tracelens/claims.json` -- The raw trace: the single `*.pt.trace.json` under `manifest.trace_dir` (find with `find "$(jq -r .trace_dir )" -name '*.pt.trace.json' | sort | head -1`). +- The raw trace: the single `*.pt.trace.json` under `manifest.trace_dir` (find with `find "$(jq -r .trace_dir )" -name '*.pt.trace.json' 2>/dev/null | sort | head -1`). ## Outputs