ORBIT-2: honor ORBIT2_MAX_BATCHES on the EDM training path

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 __getattr__(self, name: str):
+            return getattr(self._inner, name)
+
+        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 <config.yaml>", 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()
+    _apply_batch_cap()
+
+    import torch
+    import torch.distributed as dist
+
+    import train_edm  # noqa: WPS433
+    # 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()

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

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

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
 
 - `<perf_run_dir>/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 <manifest>)
+DB=$(jq -r .omnistat_db <manifest>)   # 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 <manifest>)
+JOBID=$(jq -r .job_id <manifest>)   # 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"/<job>/*.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`):

earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md

@@ -40,9 +40,11 @@ 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' <vm_info>)  # or pull from analyst stdout
+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_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" \

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 <manifest>)
-TRACE=$(jq -r '.trace_paths[0] // empty' <manifest>)
+# 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 <manifest>)
+TRACE=$(find "$(jq -r .trace_dir <manifest>)" -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"
 ```

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
 
 - `<perf_run_dir>/manifest.json`
 - `<perf_run_dir>/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 <manifest>)" -name '*.pt.trace.json' 2>/dev/null | sort | head -1`).
 
 ## Outputs
 

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