diff --git a/docs/02-user-guide/cli-reference.md b/docs/02-user-guide/cli-reference.md index e7928ba8f..530308182 100644 --- a/docs/02-user-guide/cli-reference.md +++ b/docs/02-user-guide/cli-reference.md @@ -165,6 +165,7 @@ These run under `primus/cli/main.py` unless you change `--script` in direct mode | `train posttrain --config ` | Post-training (SFT or LoRA-style workflows; same top-level flags as pretrain in the parser). | | `benchmark [args]` | Performance microbenchmarks (see table below). | | `preflight [--host] [--gpu] [--network] [--perf-test]` | Cluster and node diagnostics. | +| `preflight --mori [MORI options]` | Run NIC configuration checking, build MORI and do MORI single/cross node test. Recommended if you plan to enable MORI for training. | | `projection memory --config ` | Memory estimation from a merged config. | | `projection performance --config ` | Performance projection from a merged config. | | `projection both --config ` | Single benchmark → both performance and memory projections (cluster sizing). | @@ -208,6 +209,7 @@ Within a chosen file, nested keys follow normal YAML structure. Slurm and contai | Container pretrain | `./runner/primus-cli container --volume /data:/data -- train pretrain --config /data/exp.yaml` | | Slurm training | `./runner/primus-cli slurm srun -N 4 -- train pretrain --config exp.yaml` | | Preflight (fast) | `./runner/primus-cli slurm srun -N 4 -- preflight --host --gpu --network` | +| MORI preflight | `./runner/primus-cli direct -- preflight --mori` | | Inspect launch command | `./runner/primus-cli --dry-run direct -- train pretrain --config exp.yaml` | | Dry-run Slurm | `./runner/primus-cli --dry-run slurm srun -N 2 -- train pretrain --config exp.yaml` | diff --git a/docs/02-user-guide/node-smoke-test-instruction.md b/docs/02-user-guide/node-smoke-test-instruction.md index 41f163be1..661b280e2 100644 --- a/docs/02-user-guide/node-smoke-test-instruction.md +++ b/docs/02-user-guide/node-smoke-test-instruction.md @@ -5,7 +5,7 @@ A lightweight, distributed-rendezvous-free preflight check that runs on every no Use it to **screen a cluster fast and exclude bad nodes before launching a real training job**. A bad GPU, NIC, wedged driver, or leaked process on any node surfaces as a node FAIL — without a single global rendezvous, so a stuck node can't wedge its peers. - **Recommended launcher**: `runner/primus-cli slurm srun -- direct -- node_smoke ...` (auto-resolves the distributed env, applies `slurm.*` config defaults, same pattern as `train` / `benchmark`). The shorter `runner/primus-cli direct -- node_smoke ...` (bare `srun` + `direct`) is equivalent and handy for ad-hoc runs. -- **Companion tool**: [`preflight`](./preflight.md) — the heavier diagnostic with a global rendezvous and inter-node bandwidth tests. The recommended workflow is **node-smoke first, preflight second** (see [§10](#10-comparison-with-the-full-preflight)). +- **Companion tool**: [`preflight`](./preflight.md) — the heavier diagnostic with a global rendezvous and inter-node bandwidth tests. `preflight --mori` mode additionally performs NIC configuration checking and MORI single/cross node test. The recommended workflow is **node-smoke first, preflight second** (see [§10](#10-comparison-with-the-full-preflight)). --- diff --git a/docs/02-user-guide/preflight.md b/docs/02-user-guide/preflight.md index eb1b325b4..da7469b39 100644 --- a/docs/02-user-guide/preflight.md +++ b/docs/02-user-guide/preflight.md @@ -21,17 +21,19 @@ Preflight has two report types, controlled by a single precedence rule: | Mode | Triggered by | What it does | |---|---|---| +| **MORI runtime preflight** | `--mori` | Runs on every selected node, prints NIC/RDMA details, pulls the base image, builds pinned MORI with live NIC detection, runs an 8-GPU correctness smoke per node, verifies matching fingerprints, and optionally runs one `8 × N`-rank all-gather. Exclusive with the standard selectors below. | | **Info-only** | `--host`, `--gpu`, `--network` (in any combination) | Lightweight host / GPU / network introspection. Emits a per-node report **without requiring a rendezvous**; multi-node aggregation then uses a **timeout-bounded** rendezvous (`--dist-timeout-sec`), so it never hangs indefinitely on network misconfig. | | **Perf-only** | `--perf-test`, `--tests ...`, or `--quick` | Runs the configured perf tests under a global rendezvous. **Implied** by `--tests` and `--quick`. | | **Default (info + perf)** | No flags at all | Runs the info report first, then every perf test. | ### Mode precedence -1. **Any of `--perf-test` / `--tests` / `--quick` is set → perf-only mode.** +1. **`--mori` is set → MORI runtime mode.** +2. **Any of `--perf-test` / `--tests` / `--quick` is set → perf-only mode.** If info selectors (`--host`/`--gpu`/`--network`) are also present, they are dropped and a `WARN` is emitted (also written as a `> Note:` at the top of the perf report). To get both reports, run two invocations. -2. **Otherwise, any of `--host`/`--gpu`/`--network` is set → info-only mode.** +3. **Otherwise, any of `--host`/`--gpu`/`--network` is set → info-only mode.** Perf-only tuning knobs (e.g. `--comm-sizes-mb`) are inert in this mode and trigger a single `WARN` listing them. -3. **Otherwise (no flags) → default**: info report **first** (no rendezvous), then perf tests. +4. **Otherwise (no flags) → default**: info report **first** (no rendezvous), then perf tests. The default order ensures you always get a report even if `torch.distributed` initialization later hangs. @@ -63,6 +65,51 @@ primus-cli direct -- preflight --perf-test primus-cli direct -- preflight --quick ``` +### MORI runtime build and local correctness smoke + +```bash +primus-cli direct -- preflight --mori +``` + +MORI mode must use the host/direct launcher because it starts its own +privileged temporary container. Under Slurm, include the explicit `direct` +entry: + +```bash +primus-cli slurm srun -N 1 --ntasks-per-node=1 \ + -- direct -- preflight --mori +``` + +General multi-node preflight builds/tests every listed node, verifies that their +NIC-stack fingerprints match, then runs one all-gather across all GPUs: + +```bash +primus-cli direct -- preflight --mori \ + --mori-nodes node1,node2,node3,node4 \ + --mori-socket-ifname fenic \ + --mori-gid-index 1 +``` + +Detailed phase behavior, timing, and validated MI355X commands are documented +in [`docs/04-technical-guides/sdma-allgather.md`](../04-technical-guides/sdma-allgather.md#mori-for-primus-fsdp). + +MORI mode options: + +| Flag | Default | Purpose | +|---|---|---| +| `--mori-base-image` | ROCm 7.15 Primus nightly | Base image pulled on every run. | +| `--mori-repo` | `https://github.com/ROCm/mori.git` | MORI source repository. | +| `--mori-ref` | pinned validated commit | Revision built by preflight. | +| `--mori-max-jobs` | `32` | Parallel source-build jobs. | +| `--mori-smoke-numel` | `67108864` | BF16 elements/rank in local and N-node smokes (128 MiB/rank). | +| `--mori-keep-container` | off | Keep the temporary build container for debugging. | +| `--mori-log-dir DIR` | under `--dump-path` | Override timed phase-log directory. | +| `--mori-nodes NODES` | current node | Comma-separated hosts, Slurm hostlist, or `@file`. Each node runs full local preflight before the N-node smoke. | +| `--mori-master-addr IP` | auto | Override master bootstrap address. | +| `--mori-master-port PORT` | `29610` | N-node torchrun port. | +| `--mori-socket-ifname IFACE` | auto | Override bootstrap interface. | +| `--mori-gid-index N` | auto | Override RoCEv2 GID index. | + Equivalent on SLURM via `primus-cli slurm`: ```bash @@ -314,6 +361,11 @@ sudo sysctl --system | `--report-file-name NAME` | auto-generated `preflight-${NNODES}N-YYYYMMDD-HHMMSS` | Base name for report files. Omit to let preflight auto-generate a unique timestamped name (prevents stale leftovers from prior runs being mistaken for fresh output). Pass an explicit value when you want a stable / well-known filename. | | `--disable-pdf` | enabled | Skip PDF generation (Markdown only). Useful when `weasyprint`/`markdown2` aren't installed. | +MORI mode writes timed phase logs and container diagnostics under +`/mori-preflight--/`, or the directory supplied by +`--mori-log-dir`. It does not generate the standard Markdown/PDF performance +report. + Output files: | File | Produced when | Notes | diff --git a/docs/04-technical-guides/README.md b/docs/04-technical-guides/README.md index 2f45c99c3..51e898be8 100644 --- a/docs/04-technical-guides/README.md +++ b/docs/04-technical-guides/README.md @@ -5,6 +5,7 @@ Deep technical topics for advanced users. - [Parallelism strategies](parallelism-strategies.md): DP, TP, PP, SP, CP, EP, FSDP explained - [Parallelism configuration](parallelism-configuration.md): per-backend parallelism setup and batch size relationships - [Collective operations](collective-operations.md): NCCL/RCCL operations and their role in each parallelism strategy +- [SDMA and MORI AllGather for FSDP](sdma-allgather.md): RCCL symmetric-memory SDMA and MORI hierarchical FSDP2 communication paths - [Performance tuning](performance-tuning.md): HipBLASLt, Primus-Turbo, FP8, MoE optimization - [MoE training deep-dive](moe-training.md): bottlenecks and Primus-Turbo optimizations for Mixture-of-Experts models - [MegaMoE fused MoE layer](mega-moe.md): FlyDSL-based fused MoE layer for EP-only bf16 training, setup and reproduction diff --git a/docs/04-technical-guides/sdma-allgather.md b/docs/04-technical-guides/sdma-allgather.md new file mode 100644 index 000000000..e38ea9e13 --- /dev/null +++ b/docs/04-technical-guides/sdma-allgather.md @@ -0,0 +1,252 @@ +# SDMA and MORI AllGather for FSDP + +Primus provides two custom PyTorch FSDP2 communication paths: + +| `FSDP_ALL_GATHER_BACKEND` | FSDP all-gather implementation | Intra-node path | Cross-node path | +|---|---|---|---| +| `rccl_sdma` | PyTorch symmetric memory over RCCL | RCCL with SDMA | RDMA without SDMA | +| `mori` | MORI `HierAllGather` | MORI SDMA | RDMA & SDMA | + +Both paths move all-gather traffic away from CU-resident RCCL kernels so that +FSDP communication can overlap with GEMM-heavy forward compute. Leave +`FSDP_ALL_GATHER_BACKEND` unset to use the framework default. + +## RCCL symmetric-memory SDMA + +### Enablement + +Set one user-facing switch before launching Primus: + +```bash +export FSDP_ALL_GATHER_BACKEND=rccl_sdma + +runner/primus-cli direct -- train pretrain --config +``` + +When `FSDP_ALL_GATHER_BACKEND=rccl_sdma` is set, the Primus hook +`runner/helpers/hooks/06_enable_sdma_all_gather.sh` emits the runtime +environment needed by the training container and torchrun children. + +The hook sets: + +```bash +NCCL_CTA_POLICY=2 +NCCL_CUMEM_ENABLE=1 +NCCL_LOCAL_REGISTER=0 +TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK=true +FSDP_ALL_GATHER_BACKEND=rccl_sdma +LD_PRELOAD=/tmp/libhip_attr_drain.so +``` + +It also rebuilds `runner/helpers/hooks/sdma/hip_attr_drain_preload.c` into +`/tmp/libhip_attr_drain.so`. The interposer drains a stale HIP TLS error from +RCCL's cuMem capability probe on ROCm builds that do not have the upstream +fix. It does not change RCCL return values. + +### What the Primus patch does + +The Python backend patch is gated by +`FSDP_ALL_GATHER_BACKEND=rccl_sdma`. It wires PyTorch FSDP2 modules to use +symmetric-memory collectives: + +```python +from torch.distributed.fsdp._fully_shard._fsdp_collectives import ( + SymmMemAllGather, + SymmMemReduceScatter, +) + +module.set_custom_all_gather(SymmMemAllGather(group)) +module.set_custom_reduce_scatter(SymmMemReduceScatter(group)) +``` + +`SymmMemAllGather` allocates all-gather buffers from PyTorch symmetric memory. +Those buffers are cuMem-backed and rendezvoused across ranks. With zero-CTA +policy enabled, RCCL can dispatch the all-gather through the ROCm copy-engine +path (`__amd_rocclr_batchMemOp.kd` / `hsa_amd_memory_async_batch_copy`) instead +of running the data movement inside `ncclDevKernel_Generic_2` on CUs. + +The important discriminator is the buffer provenance: + +| Buffer source | Expected data path | +|---|---| +| `symm_mem.empty` / FSDP `SymmMemAllGather` | SDMA / copy engine | +| regular `torch.empty` / default FSDP all-gather | CU-resident RCCL kernel | + +The environment variables make the copy-engine path legal and observable, but +they do not by themselves turn regular `torch.empty` FSDP buffers into SDMA +buffers. The FSDP custom all-gather hook is the key. + +### Validation + +For low-level validation, use a symmetric-memory probe and verify that +`symm_mem.rendezvous()` completes and that a symmetric-memory all-gather runs. +For profiling validation, count HSA API calls or inspect traces for: + +```text +hsa_amd_memory_async_batch_copy +__amd_rocclr_batchMemOp.kd +``` + +Non-zero counts during all-gather indicate the SDMA copy-engine path. A trace +showing only `ncclDevKernel_Generic_2` for the data movement is not the SDMA +path, even if the communicator reports cuMem transport setup. + +### Driver and runtime compatibility + +The SDMA/FSDP path depends on PyTorch symmetric-memory rendezvous, which uses +ROCr virtual-memory APIs under the hood. The ROCm runtime and loaded amdgpu +driver must be compatible. + +One known failure mode was observed with a ROCm 7.15 nightly image where +`symm_mem.rendezvous()` hung inside: + +```text +hsa_amd_vmem_set_access + -> hsaKmtMemoryVaMap + -> driver ioctl / timeline wait +``` + +The userspace change involved: + +```text +b58362f60ff4f0b2b31a32a2a368db6bffdd5883 +ROCM-21775 Use DRM_IOCTL_SYNCOBJ_TIMELINE_WAIT ioctl in hsaKmt map/unmap ops +``` + +With an older loaded amdgpu driver, the relevant ioctl did not return, causing +`torch.distributed._symmetric_memory.rendezvous()` to hang. Updating and +reloading the amdgpu driver fixed the hang on the affected MI300X system: + +```text +$ sudo dkms status +amdgpu/7.1.3-2377367.22.04, 6.5.0-45-generic, x86_64: installed +$ uname -r +uname -r: 6.5.0-45-generic +``` + +If an SDMA run hangs before any FSDP forward progress, please try to dump stack and see where it hangs. + +## MORI hierarchical all-gather + +MORI replaces FSDP2 all-gather with `mori.ccl.HierAllGather`. It uses SDMA for +the intra-node comm and vendor direct verbs for the cross-node RDMA comm. + +### Enablement + +Set the single user-facing switch before launching Primus: + +```bash +export FSDP_ALL_GATHER_BACKEND=mori + +runner/primus-cli direct -- train pretrain --config +``` + +TorchTitan applies MORI to each compatible FSDP2 module. Megatron applies it to +FSDP2 transformer layers and additionally requires: + +```bash +--use_torch_fsdp2 true +``` + + +### What the Primus patch does + +The TorchTitan and Megatron patches wrap `fully_shard()` and attach one shared +adapter to compatible modules: + +```python +from primus.backends.common.mori_allgather import MoriAllGather + +mori_all_gather = MoriAllGather() +module.set_custom_all_gather(mori_all_gather) +``` + +The adapter: + +1. Initializes MORI SHMEM once from torchrun's default c10d `TCPStore`. This + avoids creating an eager cross-node RCCL transport solely for MORI + bootstrap. +2. Derives ranks per node from `LOCAL_WORLD_SIZE`. +3. Builds and caches `HierAllGather` for the FSDP process group and largest + observed per-rank input. +4. Launches MORI on the current CUDA stream and returns a Work-like object when + FSDP requests asynchronous completion. + + +### Runtime preflight + +MORI is sensitive to the live NIC driver, firmware, direct-verbs library, GID, +and capabilities such as Ionic CCQE. Any slight misalignment / misconfig will likly cause MORI to fail. To mitigate this issue, we provide an unified Primus preflight command, which can detect all the known critical configs and do a test all-gather on every target node: + +```bash +runner/primus-cli direct -- preflight --mori +``` + +For multi-node validation: + +```bash +runner/primus-cli direct -- preflight --mori \ + --mori-nodes node1,node2 \ + --mori-socket-ifname \ + --mori-gid-index +``` + +The CLI invokes `primus/tools/preflight/mori_preflight.py`, which runs +`mori_preflight.sh` on every selected node. The shell worker: + +1. Prints host identity, GPU, IP, RDMA links, valid GIDs, NIC + driver/firmware, vendor-library hash, and required DV symbols. +2. Starts a privileged temporary container from the pinned Primus CI image. +3. Mounts the detected host vendor library into that container. +4. Calls `runner/helpers/mori/install_mori.sh` to install dependencies, clone + the pinned source/submodules, and build MORI with live RDMA visibility. +5. Runs an 8-GPU bit-exact all-gather smoke. +6. When `--mori-nodes` is set, keeps the temporary containers for this same + information/build/local smoke on + every node, verifies matching node fingerprints, then launches one + all-gather over all `8 × N` ranks before removing them. + +The mori source version pinning is required for now, till this PR is stablized in our base rocm docker: https://github.com/ROCm/mori/pull/441 + +Logs and phase timing are written under +`/tmp/primus-mori-preflight--/`. + +### Vendor library names + +The library names used by preflight come directly from MORI: + +| NIC | MORI runtime loader names | +|---|---| +| Ionic / AINIC | `libionic.so` | +| Broadcom BNXT | `libbnxt_re.so`, then `libbnxt_re-rdmav59.so`, then `libbnxt_re-rdmav34.so` | +| Mellanox mlx5 | `libmlx5.so` | + +MORI's +[`dv_loader.hpp`](https://github.com/ROCm/mori/blob/dc4bc75a8ae63cb79a3ce17e55f2be3d8aa692c2/include/mori/application/transport/rdma/providers/dv_loader.hpp#L133) +uses these exact `dlopen()` names. Its +[`MoriDetectDevice.cmake`](https://github.com/ROCm/mori/blob/dc4bc75a8ae63cb79a3ce17e55f2be3d8aa692c2/cmake/MoriDetectDevice.cmake#L140) +uses the same names for build-time `find_library()` detection. Preflight mounts +the host's detected vendor library under these aliases so build-time detection +and runtime loading use the same library. + +### Automatic MORI installation + +When `FSDP_ALL_GATHER_BACKEND=mori` is set, the launcher hook checks whether +`mori.ccl.HierAllGather` is available. If it is missing, the hook calls +`runner/helpers/mori/install_mori.sh` before torchrun starts, so no separate +user installation command is needed. Automatic installation requires root +inside the training container. + +Useful overrides are `MORI_REPO`, `MORI_REF`, `MORI_SOURCE_DIR`, `MAX_JOBS`, +and `ROCM_PATH`. The installer clears `MORI_DEVICE_NIC` so MORI detects the +live NIC and mounted vendor library. + + +### Troubleshooting + +- `ccqe=True` on one node and `ccqe=False` on the other: choose nodes with + matching ionic stacks. +- `local GID N/A`: inspect `/sys/class/infiniband/ionic_*/ports/1/gids/`; + this pair uses `NCCL_IB_GID_INDEX=1`, not `3`. +- BNXT `231.x`: unsupported for MORI IBGDA; use supported firmware/userspace or + a validated mlx5/ionic pair. diff --git a/primus/backends/common/__init__.py b/primus/backends/common/__init__.py new file mode 100644 index 000000000..14fbae675 --- /dev/null +++ b/primus/backends/common/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared backend integration helpers.""" diff --git a/primus/backends/common/mori_allgather.py b/primus/backends/common/mori_allgather.py new file mode 100644 index 000000000..730f54e56 --- /dev/null +++ b/primus/backends/common/mori_allgather.py @@ -0,0 +1,378 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""MORI-backed FSDP2 all-gather adapter. + +The implementation mirrors the public FSDP2 adapter shape from +ROCm/mori's ``examples/fsdp_sdma/mori_allgather.py`` while keeping the +Primus backend patches small. It is intentionally all-gather only: +FSDP reduce-scatter stays on the framework default path. +""" + +from __future__ import annotations + +import importlib +import os +from collections.abc import Sequence +from typing import Any + +import torch +import torch.distributed as dist + +from primus.core.utils.module_utils import log_rank_0 + +try: + from torch.distributed.fsdp._fully_shard._fsdp_api import ( + AllGather as _FSDPAllGather, + ) +except Exception as e: # pragma: no cover - depends on torch internal version + _FSDPAllGather = object + _FSDP_ALL_GATHER_IMPORT_ERROR = e +else: + _FSDP_ALL_GATHER_IMPORT_ERROR = None + +_MORI_SHMEM_INITIALIZED = False + + +def _safe_log_rank_0(message: str) -> None: + """Log through Primus when initialized; otherwise fall back to print.""" + try: + log_rank_0(message) + except Exception: + if not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0: + print(message, flush=True) + + +def mori_all_gather_enabled() -> bool: + """Return whether Primus should install the MORI FSDP all-gather backend.""" + return os.environ.get("FSDP_ALL_GATHER_BACKEND", "") == "mori" + + +def ensure_mori_shmem_initialized(pg_name: str = "default") -> None: + """Initialize MORI SHMEM from torchrun's c10d store once. + + MORI's convenience process-group initializer broadcasts its unique ID via + the process-group backend. For NCCL groups that eagerly creates a second + cross-node transport before MORI is ready. Use the rendezvous TCPStore + instead so MORI remains the only RDMA data path. + """ + global _MORI_SHMEM_INITIALIZED + + if _MORI_SHMEM_INITIALIZED: + return + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("MORI FSDP all-gather requires torch.distributed to be initialized") + + shmem = importlib.import_module("mori.shmem") + rank = dist.get_rank() + world_size = dist.get_world_size() + store = dist.distributed_c10d._get_default_store() + uid_key = f"primus_mori_shmem_uid_{pg_name}_{world_size}" + if rank == 0: + store.set(uid_key, shmem.shmem_get_unique_id()) + uid = store.get(uid_key) + shmem.shmem_init_attr( + shmem.MORI_SHMEM_INIT_WITH_UNIQUEID, + rank, + world_size, + uid, + ) + + if shmem.shmem_mype() != rank or shmem.shmem_npes() != world_size: + raise RuntimeError( + "MORI SHMEM PE mapping must match the FSDP process group: " + f"rank/world_size={rank}/{world_size}, " + f"mype/npes={shmem.shmem_mype()}/{shmem.shmem_npes()}" + ) + + _MORI_SHMEM_INITIALIZED = True + _safe_log_rank_0("[MORI:FSDP] initialized MORI SHMEM from torch process group") + + +class _CudaEventWork: + """Small Work-like object for async FSDP all-gather calls.""" + + def __init__(self, event: torch.cuda.Event, device: torch.device) -> None: + self._event = event + self._device = device + self._waited = False + + def wait(self) -> bool: + if not self._waited: + torch.cuda.current_stream(self._device).wait_event(self._event) + self._waited = True + return True + + +class _DeviceDeferredHostSyncWork(dist.distributed_c10d.Work): + """Defer MORI's reliable host landing fence until FSDP consumes the AG.""" + + def __init__(self, stream: torch.cuda.Stream, event: torch.cuda.Event | None = None) -> None: + super().__init__() + self._stream = stream + self._event = event + self._done = False + + def wait(self, timeout=None) -> bool: # noqa: ARG002 + if not self._done: + if self._event is not None: + self._event.synchronize() + else: + self._stream.synchronize() + self._done = True + return True + + def is_completed(self) -> bool: + return self._done + + +class _HostProxyDeferredWork(dist.distributed_c10d.Work): + """Defer host-proxy completion to FSDP's wait/copy-out point.""" + + def __init__(self, collective: Any, handle: Any, drain: bool = False) -> None: + super().__init__() + self._collective = collective + self._handle = handle + self._drain = drain + self._done = False + + def wait(self, timeout=None) -> bool: # noqa: ARG002 + if not self._done: + self._collective._complete(self._handle) + if self._drain: + self._handle["stream"].synchronize() + self._collective._pending = None + self._done = True + return True + + def is_completed(self) -> bool: + return self._done + + +class MoriAllGather(_FSDPAllGather): + """FSDP2 custom all-gather backed by ``mori.ccl.HierAllGather``.""" + + def __init__(self, ranks_per_node: int | None = None) -> None: + if _FSDP_ALL_GATHER_IMPORT_ERROR is not None: + raise ImportError( + "MORI FSDP all-gather requires PyTorch FSDP2's internal " "AllGather API" + ) from _FSDP_ALL_GATHER_IMPORT_ERROR + + os.environ.setdefault("MORI_ENABLE_SDMA", "1") + os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G") + os.environ.setdefault("MORI_HIER_CUDA_GRAPH", "0") + if "MORI_SOCKET_IFNAME" not in os.environ and "NCCL_SOCKET_IFNAME" in os.environ: + os.environ["MORI_SOCKET_IFNAME"] = os.environ["NCCL_SOCKET_IFNAME"].lstrip("=") + + self._ranks_per_node = ranks_per_node + self._collective: Any | None = None + self._rank: int | None = None + self._world_size: int | None = None + self._cap_bytes = 0 + self._output_buffer: torch.Tensor | None = None + + world = int(os.environ.get("WORLD_SIZE", "0") or "0") + if world > 0: + rpn = self._ranks_per_node_value(world) + num_nodes = world // rpn if rpn else 1 + if num_nodes >= 2: + # These defaults are copied from MORI's FSDP example and only + # apply if the user did not explicitly tune the same variables. + setdefault = os.environ.setdefault + setdefault("MORI_HIER_FUSE_LOCAL", "1") + setdefault("MORI_HIER_FUSE_REMOTE", "1") + setdefault("MORI_HIER_LOCAL_PUSHONLY", "1") + if rpn < 8: + setdefault("MORI_HIER_DEEP_PIPE", "auto") + setdefault("MORI_SDMA_NUM_CHANNELS", "8") + else: + setdefault("MORI_HIER_DEBUG_SYNC", "1") + setdefault("MORI_HIER_CUDA_GRAPH", "0") + setdefault("MORI_FSDP_DEFER_HOSTSYNC", "1") + setdefault("MORI_FSDP_EVENT_FENCE", "1") + setdefault("MORI_FSDP_FWD_PREFETCH", "1") + + self._host_proxy = os.environ.get("MORI_FSDP_HOST_PROXY", "") not in ( + "", + "0", + "false", + "False", + ) + self._hostproxy_async = os.environ.get("MORI_HOSTPROXY_ASYNC", "") not in ( + "", + "0", + "false", + "False", + ) + if self._hostproxy_async: + os.environ.setdefault("MORI_HOSTPROXY_ASYNC_DRAIN", "1") + os.environ.setdefault("MORI_HOSTPROXY_ASYNC_RING", "2") + self._hostproxy_async_drain = os.environ.get("MORI_HOSTPROXY_ASYNC_DRAIN", "") not in ( + "", + "0", + "false", + "False", + ) + self._defer_hostsync = os.environ.get("MORI_FSDP_DEFER_HOSTSYNC", "") not in ( + "", + "0", + "false", + "False", + ) + self._event_fence = os.environ.get("MORI_FSDP_EVENT_FENCE", "") not in ( + "", + "0", + "false", + "False", + ) + + def allocate( + self, + size: Sequence[int | torch.SymInt], + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + numel = 1 + for dim in size: + numel *= int(dim) + if ( + self._output_buffer is not None + and self._output_buffer.dtype == dtype + and self._output_buffer.device == device + and self._output_buffer.numel() >= numel + ): + return self._output_buffer.narrow(0, 0, numel) + self._output_buffer = torch.empty(numel, dtype=dtype, device=device) + return self._output_buffer + + def _ranks_per_node_value(self, world_size: int) -> int: + if self._ranks_per_node is not None: + return self._ranks_per_node + env_value = os.environ.get("LOCAL_WORLD_SIZE") + if env_value: + return int(env_value) + return min(torch.cuda.device_count(), world_size) + + def _get_collective(self, group: dist.ProcessGroup, per_rank_bytes: int) -> Any: + rank, world_size = group.rank(), group.size() + if ( + self._collective is not None + and self._rank == rank + and self._world_size == world_size + and self._cap_bytes >= per_rank_bytes + ): + return self._collective + + ensure_mori_shmem_initialized("default") + + shmem = importlib.import_module("mori.shmem") + ccl = importlib.import_module("mori.ccl") + my_pe = shmem.shmem_mype() + npes = shmem.shmem_npes() + if my_pe != rank or npes != world_size: + raise RuntimeError( + "MORI FSDP HierAllGather requires the FSDP process group to match " + f"SHMEM PEs, got rank/world_size={rank}/{world_size} and " + f"my_pe/npes={my_pe}/{npes}" + ) + + cap = max(per_rank_bytes, self._cap_bytes) + ranks_per_node = self._ranks_per_node_value(world_size) + if self._host_proxy: + cap_floor = int(os.environ.get("MORI_FSDP_HOSTPROXY_CAP_MB", "160")) * (1 << 20) + cap = max(cap, cap_floor) + if self._collective is not None: + raise RuntimeError( + "HostProxyHierAllGather built with cap " + f"{self._cap_bytes} B but a {per_rank_bytes} B AG arrived; " + "raise MORI_FSDP_HOSTPROXY_CAP_MB" + ) + collective = ccl.HostProxyHierAllGather( + rank, + world_size, + ranks_per_node, + output_buffer_size=cap * world_size, + ) + else: + collective = ccl.HierAllGather( + my_pe, + npes, + input_buffer_size=cap, + output_buffer_size=cap * world_size, + copy_output_to_user=True, + ranks_per_node=ranks_per_node, + ) + + self._collective = collective + self._rank = rank + self._world_size = world_size + self._cap_bytes = cap + return collective + + def _validate( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: dist.ProcessGroup, + ) -> None: + if not input_tensor.is_cuda or not output_tensor.is_cuda: + raise RuntimeError("MORI FSDP HierAllGather requires CUDA tensors") + if input_tensor.device != output_tensor.device: + raise RuntimeError("MORI FSDP HierAllGather requires tensors on the same device") + if input_tensor.dtype != output_tensor.dtype: + raise RuntimeError("MORI FSDP HierAllGather requires matching dtypes") + expected = input_tensor.numel() * group.size() + if output_tensor.numel() != expected: + raise RuntimeError( + f"MORI FSDP HierAllGather expected output numel {expected}, " f"got {output_tensor.numel()}" + ) + if (input_tensor.numel() * input_tensor.element_size()) % 4 != 0: + raise RuntimeError("MORI FSDP HierAllGather requires 4-byte-aligned input bytes") + + def __call__( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: dist.ProcessGroup, + async_op: bool = False, + ) -> Any | None: + self._validate(output_tensor, input_tensor, group) + per_rank_bytes = input_tensor.numel() * input_tensor.element_size() + collective = self._get_collective(group, per_rank_bytes) + device = input_tensor.device + stream = torch.cuda.current_stream(device) + + input_tensor.record_stream(stream) + output_tensor.record_stream(stream) + + if self._host_proxy and self._hostproxy_async: + pending = getattr(collective, "_pending", None) + if pending is not None: + pending.wait() + handle = collective.call_async(input_tensor, output_tensor, input_tensor.numel(), stream=stream) + if handle is None: + return None + work = _HostProxyDeferredWork(collective, handle, drain=self._hostproxy_async_drain) + collective._pending = work + return work + + ok = collective(input_tensor, output_tensor, input_tensor.numel(), stream=stream) + if not ok: + raise RuntimeError("MORI HierAllGather call failed") + + if self._defer_hostsync: + event = None + if self._event_fence: + event = torch.cuda.Event() + event.record(stream) + return _DeviceDeferredHostSyncWork(stream, event) + + if async_op: + event = torch.cuda.Event() + event.record(stream) + return _CudaEventWork(event, device) + return None diff --git a/primus/backends/megatron/patches/mori_allgather_patches.py b/primus/backends/megatron/patches/mori_allgather_patches.py new file mode 100644 index 000000000..4fcf8e971 --- /dev/null +++ b/primus/backends/megatron/patches/mori_allgather_patches.py @@ -0,0 +1,151 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Megatron FSDP2 MORI all-gather patch. + +Activation: + Export ``FSDP_ALL_GATHER_BACKEND=mori`` and run Megatron with + ``use_torch_fsdp2: true``. This patch wraps the ``fully_shard`` symbol + used by Megatron's Torch FSDP2 wrapper and attaches MORI's FSDP2 + ``AllGather`` backend to transformer-layer units. + +Scope: + Transformer layers only for the first integration pass. Embedding, + lm_head, rotary, and reduce-scatter stay on framework defaults. +""" + +import functools +import os +import sys + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 + + +def _mori_all_gather_enabled(ctx: PatchContext) -> bool: + """Gate on the MORI backend and the Megatron FSDP2 path.""" + if os.environ.get("FSDP_ALL_GATHER_BACKEND", "") != "mori": + return False + return getattr(get_args(ctx), "use_torch_fsdp2", False) + + +@register_patch( + "megatron.fsdp.mori_allgather", + backend="megatron", + phase="before_train", + description=( + "Attach MORI HierAllGather to Megatron FSDP2 transformer-layer " + "modules so all-gather uses intra-node SDMA and cross-node RDMA. " + "Gated on FSDP_ALL_GATHER_BACKEND=mori and use_torch_fsdp2." + ), + condition=_mori_all_gather_enabled, + priority=40, +) +def patch_megatron_fsdp_mori_allgather(ctx: PatchContext) -> None: + """Install MORI's FSDP2 all-gather backend for Megatron transformer layers.""" + import megatron.core.distributed.torch_fully_sharded_data_parallel as _mfsdp_mod + import torch.distributed.fsdp as _fsdp_pkg + from torch.distributed.fsdp._fully_shard import _fully_shard as _ffs_mod + + from primus.backends.common.mori_allgather import MoriAllGather + + if not getattr(_mfsdp_mod, "HAVE_FSDP", False) or not hasattr(_mfsdp_mod, "fully_shard"): + warning_rank_0( + "[Patch:megatron.fsdp.mori_allgather] " + "torch_fully_sharded_data_parallel.fully_shard not found; skipping." + ) + return + + try: + from megatron.core.transformer.transformer_layer import TransformerLayer + except Exception as e: + warning_rank_0( + "[Patch:megatron.fsdp.mori_allgather] could not import " f"TransformerLayer; skipping: {e}" + ) + return + + orig_fully_shard = _fsdp_pkg.fully_shard + mori_all_gather = MoriAllGather() + log_all = os.environ.get("MORI_LOG_ATTACH", "0") == "1" + + def _attach_mori_all_gather(fsdp_module) -> None: + """Attach MoriAllGather to transformer-layer units only.""" + if not isinstance(fsdp_module, TransformerLayer): + if log_all: + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] skip non-transformer " + f"module: {type(fsdp_module).__name__}" + ) + return + try: + state = fsdp_module._get_fsdp_state() + except Exception as e: + if log_all: + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] skip " + f"{type(fsdp_module).__name__} without FSDP state: {e}" + ) + return + + groups = getattr(state, "_fsdp_param_groups", None) or [] + if len(groups) != 1: + if log_all: + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] skip " + f"{type(fsdp_module).__name__} with {len(groups)} param groups" + ) + return + + try: + fsdp_module.set_custom_all_gather(mori_all_gather) + except (AttributeError, ValueError, AssertionError) as e: + warning_rank_0( + f"[Patch:megatron.fsdp.mori_allgather] WARN: failed to " + f"attach MoriAllGather to {type(fsdp_module).__name__}: {e}" + ) + return + if log_all: + pg = groups[0]._all_gather_process_group + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] attached MoriAllGather " + f"to {type(fsdp_module).__name__} (group={pg.group_name})" + ) + + @functools.wraps(orig_fully_shard) + def wrapped_fully_shard(module, *args, **kwargs): + result = orig_fully_shard(module, *args, **kwargs) + _attach_mori_all_gather(result if result is not None else module) + return result + + for attr in dir(orig_fully_shard): + if attr.startswith("__"): + continue + try: + setattr(wrapped_fully_shard, attr, getattr(orig_fully_shard, attr)) + except (AttributeError, TypeError): + pass + + _fsdp_pkg.fully_shard = wrapped_fully_shard + if hasattr(_ffs_mod, "fully_shard"): + _ffs_mod.fully_shard = wrapped_fully_shard + _mfsdp_mod.fully_shard = wrapped_fully_shard + + patched_aliases = 0 + for module in tuple(sys.modules.values()): + if module is None: + continue + try: + if getattr(module, "fully_shard", None) is orig_fully_shard: + setattr(module, "fully_shard", wrapped_fully_shard) + patched_aliases += 1 + except (AttributeError, TypeError): + continue + + log_rank_0( + "[Patch:megatron.fsdp.mori_allgather] installed: FSDP2 " + "transformer-layer modules in Megatron's TorchFullyShardedDataParallel " + f"will use MoriAllGather; patched {patched_aliases} loaded aliases." + ) diff --git a/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py b/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py index 8f6bf5708..bc07aca4e 100644 --- a/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py +++ b/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py @@ -25,8 +25,8 @@ FSDP default. Activation: - Export ``SDMA_ALL_GATHER=1`` AND run with ``use_torch_fsdp2: true``. - No-op otherwise. The companion hook + Export ``FSDP_ALL_GATHER_BACKEND=rccl_sdma`` AND run with + ``use_torch_fsdp2: true``. No-op otherwise. The companion hook ``runner/helpers/hooks/06_enable_sdma_all_gather.sh`` exports the zero-CTA env (``NCCL_CTA_POLICY=2``, ...) and the LD_PRELOAD interposer so no YAML changes are required to opt in. @@ -45,8 +45,8 @@ def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: - """Gate on SDMA_ALL_GATHER=1 AND the use_torch_fsdp2 path.""" - if os.environ.get("SDMA_ALL_GATHER", "0") != "1": + """Gate on the RCCL SDMA backend and the use_torch_fsdp2 path.""" + if os.environ.get("FSDP_ALL_GATHER_BACKEND", "") != "rccl_sdma": return False return getattr(get_args(ctx), "use_torch_fsdp2", False) @@ -59,7 +59,8 @@ def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: "Attach SymmMemAllGather to FSDP2 transformer-layer modules in " "Megatron's TorchFullyShardedDataParallel so their all-gather uses " "the SDMA (copy-engine) dispatch path; other units stay on FSDP's " - "default all-gather. Gated on SDMA_ALL_GATHER=1 and use_torch_fsdp2." + "default all-gather. Gated on FSDP_ALL_GATHER_BACKEND=rccl_sdma " + "and use_torch_fsdp2." ), condition=_sdma_all_gather_enabled, ) diff --git a/primus/backends/torchtitan/patches/__init__.py b/primus/backends/torchtitan/patches/__init__.py index b4c6e2d9a..a0556f14a 100644 --- a/primus/backends/torchtitan/patches/__init__.py +++ b/primus/backends/torchtitan/patches/__init__.py @@ -34,6 +34,7 @@ metrics_output_format, mock_dataset_patches, model_override_patches, + mori_allgather, peak_flops_patches, pipelining_schedule_patches, sdma_symm_mem_collectives, diff --git a/primus/backends/torchtitan/patches/mori_allgather.py b/primus/backends/torchtitan/patches/mori_allgather.py new file mode 100644 index 000000000..00aaf64ed --- /dev/null +++ b/primus/backends/torchtitan/patches/mori_allgather.py @@ -0,0 +1,134 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""TorchTitan / FSDP2 MORI all-gather patch. + +Activation: + Export ``FSDP_ALL_GATHER_BACKEND=mori`` before launching a TorchTitan + pretrain. + The patch wraps ``torch.distributed.fsdp.fully_shard`` and attaches + MORI's FSDP2 ``AllGather`` backend to each fully-sharded module via + ``set_custom_all_gather``. + +MORI routes intra-node traffic over SDMA and cross-node traffic over RDMA. +Reduce-scatter stays on the framework default path. +""" + +import functools +import os +import sys + +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +def _mori_all_gather_enabled(ctx: PatchContext) -> bool: + """Gate on the MORI all-gather backend.""" + return os.environ.get("FSDP_ALL_GATHER_BACKEND", "") == "mori" + + +@register_patch( + "torchtitan.fsdp.mori_allgather", + backend="torchtitan", + phase="setup", + description=( + "Auto-attach MORI HierAllGather to every fully_shard'd module so " + "FSDP all-gather uses intra-node SDMA and cross-node RDMA. " + "Gated on FSDP_ALL_GATHER_BACKEND=mori." + ), + condition=_mori_all_gather_enabled, + priority=40, +) +def patch_torchtitan_fsdp_mori_allgather(ctx: PatchContext) -> None: + """Install MORI's FSDP2 all-gather backend for TorchTitan modules.""" + import torch.distributed.fsdp as _fsdp_pkg + from torch.distributed.fsdp._fully_shard import _fsdp_collectives as _ffsc + from torch.distributed.fsdp._fully_shard import _fully_shard as _ffs_mod + + from primus.backends.common.mori_allgather import MoriAllGather + + orig_fully_shard = _fsdp_pkg.fully_shard + mori_all_gather = MoriAllGather() + log_all = os.environ.get("MORI_LOG_ATTACH", "0") == "1" + + def _attach_mori_all_gather(fsdp_module) -> None: + try: + state = fsdp_module._get_fsdp_state() + except Exception as e: + if log_all: + log_rank_0( + f"[Patch:torchtitan.fsdp.mori_allgather] " + f"skip (no _fsdp_state): {type(fsdp_module).__name__}: {e}" + ) + return + + groups = getattr(state, "_fsdp_param_groups", None) or [] + if len(groups) != 1: + if log_all: + log_rank_0( + f"[Patch:torchtitan.fsdp.mori_allgather] " + f"skip multi-group module ({len(groups)} groups): " + f"{type(fsdp_module).__name__}" + ) + return + + try: + fsdp_module.set_custom_all_gather(mori_all_gather) + except (AttributeError, ValueError, AssertionError) as e: + log_rank_0( + f"[Patch:torchtitan.fsdp.mori_allgather] WARN: failed to " + f"attach MoriAllGather to {type(fsdp_module).__name__}: {e}" + ) + return + if log_all: + pg = groups[0]._all_gather_process_group + log_rank_0( + f"[Patch:torchtitan.fsdp.mori_allgather] attached " + f"MoriAllGather to {type(fsdp_module).__name__} " + f"(group={pg.group_name})" + ) + + @functools.wraps(orig_fully_shard) + def wrapped_fully_shard(module, *args, **kwargs): + result = orig_fully_shard(module, *args, **kwargs) + _attach_mori_all_gather(result if result is not None else module) + return result + + for _attr in dir(orig_fully_shard): + if _attr.startswith("__"): + continue + try: + setattr(wrapped_fully_shard, _attr, getattr(orig_fully_shard, _attr)) + except (AttributeError, TypeError): + pass + + _fsdp_pkg.fully_shard = wrapped_fully_shard + if hasattr(_ffs_mod, "fully_shard"): + _ffs_mod.fully_shard = wrapped_fully_shard + + # TorchTitan model modules may already hold a local + # ``from torch.distributed.fsdp import fully_shard`` alias by the time + # Primus runs setup patches. Replace every loaded alias that still points + # at the exact original function so those call sites also attach MORI. + patched_aliases = 0 + for module in tuple(sys.modules.values()): + if module is None: + continue + try: + if getattr(module, "fully_shard", None) is orig_fully_shard: + setattr(module, "fully_shard", wrapped_fully_shard) + patched_aliases += 1 + except (AttributeError, TypeError): + continue + + log_rank_0( + "[Patch:torchtitan.fsdp.mori_allgather] installed: every " + "fully_shard()-d module will use MoriAllGather for FSDP all-gather; " + f"patched {patched_aliases} loaded aliases." + ) + + _ffsc._primus_mori_orig_fully_shard = orig_fully_shard + _ffsc._primus_mori_attach = _attach_mori_all_gather diff --git a/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py b/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py index 60aa0a639..8af9907f7 100644 --- a/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py +++ b/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py @@ -14,18 +14,19 @@ immediately after construction. Activation: - Export ``SDMA_ALL_GATHER=1`` in the shell before launching any + Export ``FSDP_ALL_GATHER_BACKEND=rccl_sdma`` before launching any torchtitan pretrain (e.g. ``primus-cli direct -- train pretrain --config ``). The patch is a no-op otherwise. The companion hook ``runner/helpers/hooks/06_enable_sdma_all_gather.sh`` runs at - ``primus-cli`` startup and, when ``SDMA_ALL_GATHER=1``, exports + ``primus-cli`` startup and, when the backend is ``rccl_sdma``, exports the standard zero-CTA env (``NCCL_CTA_POLICY=2``, ``NCCL_CUMEM_ENABLE=1``, ...) and the LD_PRELOAD interposer so no YAML or script changes are required to opt in. - ``SDMA_ALL_GATHER`` is the only knob; there are no sub-options. + ``FSDP_ALL_GATHER_BACKEND`` is the only selector; there are no + backend-specific sub-options. Requirements: - PyTorch >= 2.12 (introduces ``SymmMemAllGather``). @@ -43,8 +44,8 @@ def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: - """Single env-driven gate. Triggered only by ``SDMA_ALL_GATHER=1``.""" - return os.environ.get("SDMA_ALL_GATHER", "0") == "1" + """Gate on the RCCL symmetric-memory SDMA backend.""" + return os.environ.get("FSDP_ALL_GATHER_BACKEND", "") == "rccl_sdma" @register_patch( @@ -54,7 +55,7 @@ def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: description=( "Auto-attach SymmMemAllGather to every fully_shard'd module " "so FSDP's all-gather uses the SDMA (copy-engine) dispatch " - "path. Gated on SDMA_ALL_GATHER=1." + "path. Gated on FSDP_ALL_GATHER_BACKEND=rccl_sdma." ), condition=_sdma_all_gather_enabled, ) @@ -73,7 +74,7 @@ def patch_torchtitan_fsdp_sdma_symm_mem(ctx: PatchContext) -> None: from torch.distributed.fsdp._fully_shard import _fully_shard as _ffs_mod from torch.distributed.fsdp._fully_shard._fsdp_collectives import SymmMemAllGather - # Hardcoded sensible defaults. SDMA_ALL_GATHER is the only knob. + # Hardcoded sensible defaults. FSDP_ALL_GATHER_BACKEND is the only selector. backend = "NCCL" log_all = False diff --git a/primus/cli/subcommands/preflight.py b/primus/cli/subcommands/preflight.py index fc3c581d8..a43c1b093 100644 --- a/primus/cli/subcommands/preflight.py +++ b/primus/cli/subcommands/preflight.py @@ -13,6 +13,7 @@ primus-cli preflight --gpu # GPU info only primus-cli preflight --network # Network info only primus-cli preflight --perf-test # Perf tests only (skip info) + primus-cli preflight --mori # MORI runtime build + smoke """ from __future__ import annotations @@ -24,6 +25,11 @@ def run(args: Any, extra_args: List[str]) -> None: """ Entry point for the 'preflight' subcommand. """ + if getattr(args, "mori", False): + from primus.tools.preflight.mori_preflight import run_mori_preflight + + raise SystemExit(run_mori_preflight(args, extra_args)) + from primus.tools.preflight.preflight_perf_test import run_preflight if extra_args: @@ -43,6 +49,7 @@ def register_subcommand(subparsers): primus-cli preflight --gpu # GPU info only primus-cli preflight --network # Network info only primus-cli preflight --perf-test # Perf only + primus-cli preflight --mori # MORI runtime build + smoke """ from primus.tools.preflight.preflight_args import add_preflight_parser diff --git a/primus/tools/preflight/mori_preflight.py b/primus/tools/preflight/mori_preflight.py new file mode 100644 index 000000000..9f7255257 --- /dev/null +++ b/primus/tools/preflight/mori_preflight.py @@ -0,0 +1,417 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""MORI per-node preflight and general N-node correctness orchestrator.""" + +from __future__ import annotations + +import concurrent.futures +import os +import shlex +import shutil +import socket +import subprocess +import sys +import time +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Sequence + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_PRIMUS_ROOT = SCRIPT_DIR.parents[2] +FINGERPRINT_PREFIX = "[preflight] NODE_FINGERPRINT " + + +@dataclass +class NodeResult: + node: str + returncode: int + fingerprint: str | None + log_file: Path + + +def _dedupe(items: list[str]) -> list[str]: + seen = set() + result = [] + for item in items: + item = item.strip() + if item and item not in seen: + seen.add(item) + result.append(item) + return result + + +def resolve_nodes(spec: str | None) -> list[str]: + if not spec: + return [socket.gethostname().split(".")[0]] + if spec.startswith("@"): + return _dedupe(Path(spec[1:]).read_text().splitlines()) + if "," in spec: + return _dedupe(spec.split(",")) + if shutil.which("scontrol"): + result = subprocess.run( + ["scontrol", "show", "hostnames", spec], + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + nodes = _dedupe(result.stdout.splitlines()) + if result.returncode == 0 and nodes: + return nodes + return [spec] + + +def is_local(node: str) -> bool: + names = { + "localhost", + "127.0.0.1", + socket.gethostname(), + socket.gethostname().split(".")[0], + socket.getfqdn(), + } + return node in names + + +def preflight_container_name(node: str) -> str: + short_node = node.split(".")[0] + return f"primus_mori_preflight_{os.environ.get('USER', 'user')}_{short_node}" + + +def shell_join_env(env: dict[str, str], command: list[str]) -> str: + assignments = [f"{name}={value}" for name, value in env.items()] + return shlex.join(["env", *assignments, *command]) + + +def local_preflight_command( + args: Any, + node: str, + node_log_dir: Path, + repo_root: Path, +) -> list[str]: + env = { + "BASE_IMAGE": args.mori_base_image, + "MORI_REPO": args.mori_repo, + "MORI_REF": args.mori_ref, + "MAX_JOBS": str(args.mori_max_jobs), + "SMOKE_NUMEL": str(args.mori_smoke_numel), + # Reuse these containers for the N-node smoke and remove them after. + "KEEP_CONTAINER": "1", + "CONTAINER_NAME": preflight_container_name(node), + "LOG_DIR": str(node_log_dir), + "PRIMUS_ROOT": str(repo_root), + } + helper = str(repo_root / "primus" / "tools" / "preflight" / "mori_preflight.sh") + command = ["bash", helper] + if is_local(node): + return ["env", *[f"{k}={v}" for k, v in env.items()], *command] + + remote_command = f"cd {shlex.quote(str(repo_root))} && {shell_join_env(env, command)}" + return ["ssh", "-o", "BatchMode=yes", node, remote_command] + + +def run_node_preflight( + args: Any, + node: str, + root_log_dir: Path, + repo_root: Path, +) -> NodeResult: + short_node = node.split(".")[0] + node_log_dir = root_log_dir / "nodes" / short_node + node_log_dir.mkdir(parents=True, exist_ok=True) + launcher_log = node_log_dir / "launcher.log" + command = local_preflight_command(args, node, node_log_dir, repo_root) + fingerprint = None + + with launcher_log.open("w", encoding="utf-8") as log: + process = subprocess.Popen( + command, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + bufsize=1, + ) + assert process.stdout is not None + for line in process.stdout: + log.write(line) + log.flush() + print(f"[{short_node}] {line}", end="", flush=True) + if line.startswith(FINGERPRINT_PREFIX): + fingerprint = line[len(FINGERPRINT_PREFIX) :].strip() + returncode = process.wait() + + return NodeResult(node, returncode, fingerprint, launcher_log) + + +def remote_output(node: str, command: str) -> str: + if is_local(node): + argv = ["bash", "-lc", command] + else: + argv = ["ssh", "-o", "BatchMode=yes", node, command] + result = subprocess.run( + argv, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + if result.returncode != 0: + raise RuntimeError( + f"{node}: command failed ({result.returncode}): " f"{result.stderr.strip() or command}" + ) + return result.stdout.strip() + + +def probe_network( + node: str, + interface_override: str | None, + gid_override: int | None, + address_override: str | None, +) -> tuple[str, int, str]: + if interface_override and gid_override is not None and address_override: + return interface_override, gid_override, address_override + + command = f""" +interface={shlex.quote(interface_override or "")} +gid_index={shlex.quote("" if gid_override is None else str(gid_override))} +master_addr={shlex.quote(address_override or "")} + +if [ -z "$interface" ]; then + if ip -o -4 addr show dev fenic scope global >/dev/null 2>&1; then + interface=fenic + else + for path in /sys/class/infiniband/*/device/net/*; do + [ -e "$path" ] || continue + candidate=${{path##*/}} + if ip -o -4 addr show dev "$candidate" scope global >/dev/null 2>&1; then + interface=$candidate + break + fi + done + fi +fi +[ -n "$interface" ] || {{ echo "no RDMA interface with IPv4" >&2; exit 1; }} + +if [ -z "$gid_index" ]; then + gid_index=0 + for dev in /sys/class/infiniband/*; do + [ -d "$dev" ] || continue + for path in "$dev"/ports/1/gids/*; do + [ -f "$path" ] || continue + idx=${{path##*/}} + gid=$(cat "$path") + type=$(cat "$dev/ports/1/gid_attrs/types/$idx" 2>/dev/null || true) + case "$type:$gid" in + "RoCE v2:"*":ffff:"*) gid_index=$idx; break 2 ;; + esac + done + done +fi + +if [ -z "$master_addr" ]; then + master_addr=$(ip -o -4 addr show dev "$interface" scope global | + awk 'NR==1{{split($4,a,"/"); print a[1]}}') +fi +[ -n "$master_addr" ] || {{ echo "no IPv4 address on $interface" >&2; exit 1; }} +printf '%s\\t%s\\t%s\\n' "$interface" "$gid_index" "$master_addr" +""" + fields = remote_output(node, command).split("\t") + if len(fields) != 3: + raise RuntimeError(f"{node}: malformed network probe output: {fields}") + return fields[0], int(fields[1]), fields[2] + + +def validate_fingerprints(results: list[NodeResult]) -> None: + missing = [result.node for result in results if not result.fingerprint] + if missing: + raise RuntimeError(f"Missing MORI fingerprint from nodes: {', '.join(missing)}") + grouped: dict[str, list[str]] = {} + for result in results: + assert result.fingerprint is not None + grouped.setdefault(result.fingerprint, []).append(result.node) + if len(grouped) != 1: + detail = "; ".join(f"{nodes}: {fingerprint}" for fingerprint, nodes in grouped.items()) + raise RuntimeError(f"MORI node-stack mismatch: {detail}") + + +def run_multinode(args: Any, nodes: list[str], log_dir: Path) -> int: + master_node = nodes[0] + interface, gid_index, master_addr = probe_network( + master_node, + args.mori_socket_ifname, + args.mori_gid_index, + args.mori_master_addr, + ) + multinode_log_dir = log_dir / "multinode" + multinode_log_dir.mkdir(parents=True, exist_ok=True) + + def launch(node: str, rank: int) -> tuple[int, Path]: + torchrun = shlex.join( + [ + "torchrun", + f"--nnodes={len(nodes)}", + "--nproc_per_node=8", + f"--node_rank={rank}", + f"--master_addr={master_addr}", + f"--master_port={args.mori_master_port}", + "/src/primus/runner/helpers/mori/multinode_allgather_smoke.py", + "--numel", + str(args.mori_smoke_numel), + ] + ) + docker_command = [ + "docker", + "exec", + "-e", + "PYTHONPATH=/src/primus", + "-e", + f"NCCL_SOCKET_IFNAME={interface}", + "-e", + f"NCCL_IB_GID_INDEX={gid_index}", + preflight_container_name(node), + "bash", + "-lc", + "export LD_LIBRARY_PATH=/opt/mori-host-libs:${LD_LIBRARY_PATH}; " f"exec {torchrun}", + ] + command = ( + docker_command + if is_local(node) + else ["ssh", "-o", "BatchMode=yes", node, shlex.join(docker_command)] + ) + rank_log = multinode_log_dir / f"node-{rank}-{node}.log" + with rank_log.open("w", encoding="utf-8") as output: + result = subprocess.run( + command, + stdout=output, + stderr=subprocess.STDOUT, + check=False, + ) + return result.returncode, rank_log + + print( + f"[MORI:Preflight] launching {len(nodes)} nodes via {master_node} " + f"({master_addr}, {interface}, gid={gid_index})", + flush=True, + ) + failures = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=len(nodes)) as executor: + futures = {} + for rank, node in enumerate(nodes): + futures[executor.submit(launch, node, rank)] = node + if rank == 0: + time.sleep(1) + for future, node in futures.items(): + returncode, rank_log = future.result() + if returncode: + failures.append((node, rank_log)) + for node, rank_log in failures: + print(f"[MORI:Preflight] FAIL {node}: {rank_log}", file=sys.stderr) + return int(bool(failures)) + + +def remove_preflight_containers(nodes: list[str]) -> None: + def remove(node: str) -> None: + name = preflight_container_name(node) + remote_output( + node, + f"docker rm -f {shlex.quote(name)} >/dev/null 2>&1 || true", + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=len(nodes)) as executor: + list(executor.map(remove, nodes)) + + +def run_orchestrator(args: Any, repo_root: Path, log_dir: Path) -> int: + nodes = resolve_nodes(args.mori_nodes) + log_dir.mkdir(parents=True, exist_ok=True) + print(f"[MORI:Preflight] nodes ({len(nodes)}): {', '.join(nodes)}") + print(f"[MORI:Preflight] logs: {log_dir}") + + try: + results = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=len(nodes)) as executor: + future_map = { + executor.submit(run_node_preflight, args, node, log_dir, repo_root): node for node in nodes + } + for future in concurrent.futures.as_completed(future_map): + results.append(future.result()) + + results.sort(key=lambda result: nodes.index(result.node)) + failed = [result for result in results if result.returncode != 0] + if failed: + for result in failed: + print( + f"[MORI:Preflight] FAIL {result.node}: {result.log_file}", + file=sys.stderr, + ) + return 1 + + validate_fingerprints(results) + print("[MORI:Preflight] all node fingerprints match", flush=True) + if len(nodes) > 1: + return run_multinode(args, nodes, log_dir) + return 0 + finally: + if not args.mori_keep_container: + remove_preflight_containers(nodes) + + +def _validate_mori_mode(args: Any, extra_args: Sequence[str]) -> int | None: + if extra_args: + print( + f"[Primus:Preflight] ERROR: unknown MORI arguments: {list(extra_args)}", + file=sys.stderr, + ) + return 2 + + incompatible = [] + for attr, flag in ( + ("check_host", "--host"), + ("check_gpu", "--gpu"), + ("check_network", "--network"), + ("perf_test", "--perf-test"), + ("plot", "--plot"), + ("tests", "--tests"), + ("comm_sizes_mb", "--comm-sizes-mb"), + ("intra_comm_sizes_mb", "--intra-comm-sizes-mb"), + ("inter_comm_sizes_mb", "--inter-comm-sizes-mb"), + ("intra_group_sizes", "--intra-group-sizes"), + ("inter_group_sizes", "--inter-group-sizes"), + ("ring_p2p_sizes_mb", "--ring-p2p-sizes-mb"), + ("quick", "--quick"), + ): + if getattr(args, attr, None): + incompatible.append(flag) + if getattr(args, "split_nodes_subgroup", True) is False: + incompatible.append("--no-split-nodes-subgroup") + if incompatible: + print( + "[Primus:Preflight] ERROR: --mori cannot be combined with " + ", ".join(incompatible), + file=sys.stderr, + ) + return 2 + return None + + +def run_mori_preflight(args: Any, extra_args: Sequence[str] = ()) -> int: + validation_rc = _validate_mori_mode(args, extra_args) + if validation_rc is not None: + return validation_rc + + repo_root = Path( + os.environ.get( + "PRIMUS_PATH", + str(DEFAULT_PRIMUS_ROOT), + ) + ).resolve() + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + short_host = socket.gethostname().split(".")[0] + report_name = args.report_file_name or f"mori-preflight-{short_host}-{timestamp}" + log_dir = Path(args.mori_log_dir or (Path(args.dump_path) / report_name)).resolve() + + print(f"[Primus:Preflight] MORI logs: {log_dir}", flush=True) + return run_orchestrator(args, repo_root, log_dir) diff --git a/primus/tools/preflight/mori_preflight.sh b/primus/tools/preflight/mori_preflight.sh new file mode 100755 index 000000000..ec736411d --- /dev/null +++ b/primus/tools/preflight/mori_preflight.sh @@ -0,0 +1,353 @@ +#!/usr/bin/env bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PRIMUS_ROOT="${PRIMUS_ROOT:-$(cd "${SCRIPT_DIR}/../../.." && pwd)}" + +BASE_IMAGE="${BASE_IMAGE:-unifiedtrainingdockers.azurecr.io/utd/nightly:primus_the_rock_rocm7.15_20260728}" +MORI_REPO="${MORI_REPO:-https://github.com/ROCm/mori.git}" +MORI_REF="${MORI_REF:-12d1bc32d0c93dcd5062e74f4e0f772e36e1aac4}" +MAX_JOBS="${MAX_JOBS:-32}" +SMOKE_NUMEL="${SMOKE_NUMEL:-67108864}" +KEEP_CONTAINER="${KEEP_CONTAINER:-0}" +LOG_DIR="${LOG_DIR:-/tmp/primus-mori-preflight-$(hostname -s)-$(date +%Y%m%d-%H%M%S)}" +CONTAINER_NAME="${CONTAINER_NAME:-primus_mori_preflight_${USER}_$(hostname -s)}" + +mkdir -p "${LOG_DIR}" + +CURRENT_PHASE="initialization" +PHASE_SUMMARY=() + +section() { + echo + echo "================================================================================" + echo "$*" + echo "================================================================================" +} + +format_seconds() { + local seconds="$1" + printf "%dm%02ds" "$((seconds / 60))" "$((seconds % 60))" +} + +run_phase() { + local name="$1" + shift + local start end elapsed rc + CURRENT_PHASE="${name}" + start="$(date +%s)" + section "Phase: ${name}" + + set +e + "$@" 2>&1 | tee "${LOG_DIR}/${name// /_}.log" + rc="${PIPESTATUS[0]}" + set -e + + end="$(date +%s)" + elapsed="$((end - start))" + PHASE_SUMMARY+=("${name}|${elapsed}|${rc}") + echo "[preflight] ${name}: $(format_seconds "${elapsed}") (rc=${rc})" + if [[ "${rc}" -ne 0 ]]; then + return "${rc}" + fi +} + +cleanup() { + local rc=$? + if [[ "${rc}" -ne 0 ]]; then + echo + echo "[preflight] FAILED during phase: ${CURRENT_PHASE}" >&2 + echo "[preflight] Logs: ${LOG_DIR}" >&2 + docker inspect "${CONTAINER_NAME}" >"${LOG_DIR}/container-inspect.json" 2>/dev/null || true + fi + + if [[ "${KEEP_CONTAINER}" != "1" ]]; then + docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true + else + echo "[preflight] Keeping container: ${CONTAINER_NAME}" + fi +} +trap cleanup EXIT INT TERM + +command -v docker >/dev/null || { + echo "[preflight] docker is required" >&2 + exit 1 +} + +section "Host identity" +echo "hostname : $(hostname -f)" +echo "kernel : $(uname -r)" +echo "user : $(id)" +echo "base image : ${BASE_IMAGE}" +echo "MORI revision : ${MORI_REF}" +echo "Primus root : ${PRIMUS_ROOT}" +echo "log directory : ${LOG_DIR}" + +section "GPU information" +if command -v rocm-smi >/dev/null; then + rocm-smi --showproductname --showuse --showmemuse --csv 2>&1 || true +else + echo "rocm-smi not found" +fi +if command -v rocminfo >/dev/null; then + rocminfo 2>/dev/null | + awk '/Name: +gfx/{print "GPU architecture : " $2}' | + sort -u || true +fi + +section "IP interfaces" +ip -o -4 addr show scope global 2>&1 || true + +section "RDMA devices and links" +if command -v ibv_devices >/dev/null; then + ibv_devices 2>&1 || true +fi +if command -v ibdev2netdev >/dev/null; then + ibdev2netdev 2>&1 || true +fi +if command -v rdma >/dev/null; then + rdma link show 2>&1 || true +fi + +section "Valid GIDs" +for dev_path in /sys/class/infiniband/*; do + [[ -d "${dev_path}" ]] || continue + dev="${dev_path##*/}" + for gid_path in "${dev_path}"/ports/1/gids/*; do + [[ -f "${gid_path}" ]] || continue + idx="${gid_path##*/}" + gid="$(<"${gid_path}")" + [[ "${gid}" != "0000:0000:0000:0000:0000:0000:0000:0000" ]] || continue + type="unknown" + ndev="unknown" + if [[ -r "${dev_path}/ports/1/gid_attrs/types/${idx}" ]]; then + type="$(<"${dev_path}/ports/1/gid_attrs/types/${idx}")" + fi + if [[ -r "${dev_path}/ports/1/gid_attrs/ndevs/${idx}" ]]; then + ndev="$(<"${dev_path}/ports/1/gid_attrs/ndevs/${idx}")" + fi + printf "%-12s index=%-3s type=%-8s netdev=%-12s gid=%s\n" \ + "${dev}" "${idx}" "${type}" "${ndev}" "${gid}" + done +done + +section "NIC driver and firmware" +mapfile -t NETDEVS < <( + for dev_path in /sys/class/infiniband/*; do + [[ -d "${dev_path}" ]] || continue + for ndev_path in "${dev_path}"/device/net/*; do + [[ -e "${ndev_path}" ]] && basename "${ndev_path}" + done + done | sort -u +) +STACK_DATA="" +for netdev in "${NETDEVS[@]:-}"; do + driver_info="$(ethtool -i "${netdev}" 2>&1 || true)" + echo "--- ${netdev} ---" + echo "${driver_info}" + driver_version="$(awk -F': ' '$1=="version"{print $2}' <<<"${driver_info}")" + firmware_version="$(awk -F': ' '$1=="firmware-version"{print $2}' <<<"${driver_info}")" + STACK_DATA+="${netdev}|${driver_version}|${firmware_version};" +done +STACK_SHA="$(printf "%s" "${STACK_DATA}" | sha256sum | awk '{print $1}')" + +detect_nic() { + if compgen -G "/sys/class/infiniband/ionic*" >/dev/null; then + echo ionic + elif compgen -G "/sys/class/infiniband/bnxt_re*" >/dev/null; then + echo bnxt + elif compgen -G "/sys/class/infiniband/mlx5*" >/dev/null; then + echo mlx5 + else + echo unknown + fi +} + +DETECTED_NIC="$(detect_nic)" + +section "Vendor library checks (detected NIC: ${DETECTED_NIC})" + +find_library() { + local name="$1" + local path + path="$(ldconfig -p 2>/dev/null | awk -v n="${name}" '$1==n{print $NF; exit}')" + if [[ -z "${path}" ]]; then + for candidate in \ + "/usr/local/lib/${name}" \ + "/usr/lib/x86_64-linux-gnu/${name}" \ + "/lib/x86_64-linux-gnu/${name}"; do + if [[ -e "${candidate}" ]]; then + path="${candidate}" + break + fi + done + fi + echo "${path}" +} + +check_symbols() { + local path="$1" + shift + local symbol + for symbol in "$@"; do + if nm -D "${path}" 2>/dev/null | awk -v s="${symbol}" '$3==s{found=1} END{exit !found}'; then + echo " ${symbol}: present" + else + echo " ${symbol}: MISSING" + fi + done +} + +VENDOR_LIB="" +VENDOR_NAMES=() +VENDOR_SHA="missing" +case "${DETECTED_NIC}" in + ionic) + VENDOR_LIB="$(find_library libionic.so)" + VENDOR_NAMES=(libionic.so) + ;; + bnxt) + VENDOR_LIB="$(find_library libbnxt_re.so)" + VENDOR_NAMES=(libbnxt_re.so libbnxt_re-rdmav34.so) + ;; + mlx5) + VENDOR_LIB="$(find_library libmlx5.so)" + VENDOR_NAMES=(libmlx5.so) + ;; +esac + +if [[ -n "${VENDOR_LIB}" ]]; then + VENDOR_LIB="$(readlink -f "${VENDOR_LIB}")" + VENDOR_SHA="$(sha256sum "${VENDOR_LIB}" | awk '{print $1}')" + echo "library : ${VENDOR_LIB}" + echo "sha256 : ${VENDOR_SHA}" + case "${DETECTED_NIC}" in + ionic) + check_symbols "${VENDOR_LIB}" \ + ionic_dv_get_ctx \ + ionic_dv_get_cq \ + ionic_dv_get_qp \ + ionic_dv_create_cq_ex + ;; + bnxt) + check_symbols "${VENDOR_LIB}" \ + bnxt_re_dv_umem_reg \ + bnxt_re_dv_umem_dereg \ + bnxt_re_dv_create_cq \ + bnxt_re_dv_destroy_cq \ + bnxt_re_dv_init_obj \ + bnxt_re_dv_create_qp \ + bnxt_re_dv_destroy_qp \ + bnxt_re_dv_modify_qp + ;; + esac +else + echo "vendor library : not found" +fi + +run_phase "pull base image" docker pull "${BASE_IMAGE}" + +docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true + +DOCKER_ARGS=( + docker run -d + --name "${CONTAINER_NAME}" + --device=/dev/kfd + --device=/dev/dri + --group-add video + --cap-add SYS_PTRACE + --security-opt seccomp=unconfined + --privileged + --ipc=host + --network=host + -e PYTHONDONTWRITEBYTECODE=1 + -v "${PRIMUS_ROOT}:/src/primus:ro" +) + +if [[ -n "${VENDOR_LIB}" ]]; then + for name in "${VENDOR_NAMES[@]}"; do + DOCKER_ARGS+=(-v "${VENDOR_LIB}:/opt/mori-host-libs/${name}:ro") + done +fi + +DOCKER_ARGS+=("${BASE_IMAGE}" sleep infinity) +run_phase "start container" "${DOCKER_ARGS[@]}" + +# shellcheck disable=SC2016 +run_phase "container network check" \ + docker exec "${CONTAINER_NAME}" bash -lc ' + set -e + echo "torch=$(python3 -c "import torch; print(torch.__version__)")" + echo "ROCM_PATH=${ROCM_PATH}" + ip -o -4 addr show scope global || true + ibv_devices || true + export LD_LIBRARY_PATH=/opt/mori-host-libs:${LD_LIBRARY_PATH} + python3 - <<'"'"'PY'"'"' +import ctypes +for name in ("libionic.so", "libbnxt_re.so", "libmlx5.so"): + try: + ctypes.CDLL(name) + print(f"{name}: loadable") + except OSError as exc: + print(f"{name}: unavailable ({exc})") +PY + ' + +run_phase "install MORI" \ + docker exec \ + -e MORI_REPO="${MORI_REPO}" \ + -e MORI_REF="${MORI_REF}" \ + -e MAX_JOBS="${MAX_JOBS}" \ + "${CONTAINER_NAME}" \ + bash /src/primus/runner/helpers/mori/install_mori.sh + +# shellcheck disable=SC2016 +run_phase "8-GPU MORI all-gather smoke" \ + docker exec \ + -e PYTHONPATH=/src/primus \ + -e NCCL_SOCKET_IFNAME=lo \ + -e SMOKE_NUMEL="${SMOKE_NUMEL}" \ + "${CONTAINER_NAME}" bash -lc ' + set -e + export LD_LIBRARY_PATH="/opt/mori-host-libs:${LD_LIBRARY_PATH}" + torchrun --standalone --nproc_per_node=8 \ + /src/primus/runner/helpers/mori/multinode_allgather_smoke.py \ + --numel "${SMOKE_NUMEL}" + ' + +CCQE="n/a" +if [[ "${DETECTED_NIC}" == "ionic" ]]; then + CCQE="$( + awk ' + /_ccqe_enabled:/ { + for (i = 1; i <= NF; i++) { + if ($i == "_ccqe_enabled:") { + print $(i + 1) + exit + } + } + } + ' "${LOG_DIR}/8-GPU_MORI_all-gather_smoke.log" + )" + CCQE="${CCQE:-unknown}" +fi +echo "[preflight] NODE_FINGERPRINT nic=${DETECTED_NIC} stack_sha=${STACK_SHA} vendor_sha=${VENDOR_SHA} ccqe=${CCQE}" + +section "Timing summary" +total=0 +for entry in "${PHASE_SUMMARY[@]}"; do + IFS="|" read -r name elapsed rc <<<"${entry}" + printf "%-32s %8s rc=%s\n" "${name}" "$(format_seconds "${elapsed}")" "${rc}" + total="$((total + elapsed))" +done +printf "%-32s %8s\n" "TOTAL" "$(format_seconds "${total}")" + +echo +echo "[preflight] PASS" +echo "[preflight] Logs: ${LOG_DIR}" diff --git a/primus/tools/preflight/preflight_args.py b/primus/tools/preflight/preflight_args.py index 1710e14b4..9db5c47f7 100644 --- a/primus/tools/preflight/preflight_args.py +++ b/primus/tools/preflight/preflight_args.py @@ -49,6 +49,7 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa primus-cli preflight --gpu # GPU info only primus-cli preflight --network # Network info only primus-cli preflight --gpu --network # GPU + Network info + primus-cli preflight --mori # MORI runtime build + smoke primus-cli preflight --perf-test # Perf only, all tests primus-cli preflight --quick # Perf only, fast preset primus-cli preflight --tests gemm # Perf only, GEMM only @@ -229,4 +230,81 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa action="store_false", help="Disable PDF report generation.", ) + + # MORI preflight is an exclusive orchestration mode. It runs once per node + # (not once per GPU), builds MORI inside a privileged runtime container, + # and executes local plus optional cross-node all-gather correctness tests. + mori = parser.add_argument_group("MORI runtime preflight") + mori.add_argument( + "--mori", + action="store_true", + help="Run MORI NIC diagnostics, runtime source build, and all-gather smoke. " + "Cannot be combined with standard info/perf selectors.", + ) + mori.add_argument( + "--mori-base-image", + default=("unifiedtrainingdockers.azurecr.io/utd/nightly:" "primus_the_rock_rocm7.15_20260728"), + help="Base image pulled by MORI preflight.", + ) + mori.add_argument( + "--mori-repo", + default="https://github.com/ROCm/mori.git", + help="MORI git repository.", + ) + mori.add_argument( + "--mori-ref", + default="12d1bc32d0c93dcd5062e74f4e0f772e36e1aac4", + help="Pinned MORI git revision to build.", + ) + mori.add_argument( + "--mori-max-jobs", + type=int, + default=32, + help="Maximum parallel MORI build jobs.", + ) + mori.add_argument( + "--mori-smoke-numel", + type=int, + default=67108864, + help="BF16 elements per rank in MORI all-gather smokes " "(default: 67108864 = 128 MiB).", + ) + mori.add_argument( + "--mori-keep-container", + action="store_true", + help="Keep the temporary build container after preflight.", + ) + mori.add_argument( + "--mori-log-dir", + default=None, + help="MORI phase-log directory. Defaults under --dump-path.", + ) + mori.add_argument( + "--mori-nodes", + default=None, + help="General multi-node target: comma-separated hosts, Slurm hostlist, " + "or @file. Each node runs full diagnostics/build/local smoke before " + "the N-node all-gather.", + ) + mori.add_argument( + "--mori-master-addr", + default=None, + help="Master IP for the N-node smoke (auto-detected when unset).", + ) + mori.add_argument( + "--mori-master-port", + type=int, + default=29610, + help="torchrun rendezvous port for the N-node MORI smoke.", + ) + mori.add_argument( + "--mori-socket-ifname", + default=None, + help="Bootstrap interface for the N-node smoke.", + ) + mori.add_argument( + "--mori-gid-index", + type=int, + default=None, + help="RCCL RoCEv2 GID index for the N-node smoke.", + ) return parser diff --git a/primus/tools/preflight/preflight_perf_test.py b/primus/tools/preflight/preflight_perf_test.py index 636dee728..fd97ad639 100644 --- a/primus/tools/preflight/preflight_perf_test.py +++ b/primus/tools/preflight/preflight_perf_test.py @@ -415,14 +415,20 @@ def run_preflight(args): Mode precedence (single rule): - 1. Any of --perf-test / --tests / --quick is set -> perf-only mode. + 1. --mori is set -> exclusive MORI runtime preflight mode. + 2. Any of --perf-test / --tests / --quick is set -> perf-only mode. If info selectors (--host/--gpu/--network) are also present, they are dropped with a WARN. - 2. Otherwise, any of --host/--gpu/--network is set -> info-only mode. + 3. Otherwise, any of --host/--gpu/--network is set -> info-only mode. Perf tuning knobs (e.g. --comm-sizes-mb), if set, are inert and a WARN is emitted. - 3. Otherwise (no flags) -> default: info AND all perf tests. + 4. Otherwise (no flags) -> default: info AND all perf tests. """ + if getattr(args, "mori", False): + from primus.tools.preflight.mori_preflight import run_mori_preflight + + return run_mori_preflight(args) + # R4: canonical normalization point for the report file name. Done here # before any downstream code reads args.report_file_name, so every code # path (info-only, perf-only, info+perf, dist-init failure) sees the same diff --git a/pyproject.toml b/pyproject.toml index 4851a3253..5d772f12e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,7 +94,10 @@ path = "primus/__init__.py" packages = ["primus"] # _thirdparty.lock is generated at build time (and git-ignored); force-include it # so the wheel ships it for `primus-cli deps sync`. -artifacts = ["primus/_thirdparty.lock"] +artifacts = [ + "primus/_thirdparty.lock", + "primus/tools/preflight/mori_preflight.sh", +] [tool.hatch.build.targets.wheel.force-include] # Ship the whole primus-cli bash toolkit as package data under primus/runner/. diff --git a/runner/helpers/hooks/06_enable_sdma_all_gather.sh b/runner/helpers/hooks/06_enable_sdma_all_gather.sh index 7364fdbf4..8f0ab15e6 100755 --- a/runner/helpers/hooks/06_enable_sdma_all_gather.sh +++ b/runner/helpers/hooks/06_enable_sdma_all_gather.sh @@ -8,15 +8,15 @@ # Global hook: opt into the SDMA/RCCL dispatch path for FSDP # all-gather. # -# Single trigger -- the only knob: +# Backend selector: # -# export SDMA_ALL_GATHER=1 +# export FSDP_ALL_GATHER_BACKEND=rccl_sdma # primus-cli direct -- train pretrain --config # -# When SDMA_ALL_GATHER=1, this hook: +# When the backend is rccl_sdma, this hook: # 1. Exports the zero-CTA env that RCCL needs to actually take the # copy-engine path (NCCL_CTA_POLICY=2, NCCL_CUMEM_ENABLE=1, ...). -# 2. Propagates SDMA_ALL_GATHER=1 into the launched torchrun children +# 2. Propagates FSDP_ALL_GATHER_BACKEND=rccl_sdma into torchrun children # so the companion Python patch's gate fires there too. See # primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py. # 3. Rebuilds the bundled LD_PRELOAD interposer @@ -33,7 +33,7 @@ set -euo pipefail -if [[ "${SDMA_ALL_GATHER:-0}" != "1" ]]; then +if [[ "${FSDP_ALL_GATHER_BACKEND:-}" != "rccl_sdma" ]]; then exit 0 fi @@ -48,11 +48,11 @@ echo "env.NCCL_CUMEM_ENABLE=1" echo "env.NCCL_LOCAL_REGISTER=0" echo "env.TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK=true" -# 2) Make the trigger visible to torchrun children so the Python patch +# 2) Make the selector visible to torchrun children so the Python patch # fires there. primus-cli direct doesn't inherit the host env into # the child unless it's either CLI-passed via --env or hook-emitted # via env.*. -echo "env.SDMA_ALL_GATHER=1" +echo "env.FSDP_ALL_GATHER_BACKEND=rccl_sdma" # 3) Always (re)build the interposer. The source is tiny and gcc is # typically <1s; we don't bother with a staleness check so the .so diff --git a/runner/helpers/hooks/07_enable_mori_all_gather.sh b/runner/helpers/hooks/07_enable_mori_all_gather.sh new file mode 100755 index 000000000..d16891eb5 --- /dev/null +++ b/runner/helpers/hooks/07_enable_mori_all_gather.sh @@ -0,0 +1,60 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Global hook: opt into MORI FSDP all-gather. +# +# Trigger: +# +# export FSDP_ALL_GATHER_BACKEND=mori +# primus-cli direct -- train pretrain --config +# +# When enabled, this hook installs MORI when needed and propagates MORI_* env +# vars into torchrun children. The Python patches then attach MoriAllGather to +# FSDP2 modules. + +set -euo pipefail + +if [[ "${FSDP_ALL_GATHER_BACKEND:-}" != "mori" ]]; then + exit 0 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +mori_installer="${SCRIPT_DIR}/../mori/install_mori.sh" + +if ! python3 -c "from mori.ccl import HierAllGather" >/dev/null 2>&1; then + echo "[MORI] HierAllGather is unavailable; installing MORI before launch." >&2 + if [[ ! -x "${mori_installer}" ]]; then + echo "[ERROR] MORI installer is not executable: ${mori_installer}" >&2 + exit 1 + fi + "${mori_installer}" >&2 + if ! python3 -c "from mori.ccl import HierAllGather" >/dev/null 2>&1; then + echo "[ERROR] MORI installation completed but HierAllGather is unavailable." >&2 + exit 1 + fi +fi + +# MORI all-gather uses SDMA for the intra-node leg. Allow an explicit caller +# value to win, but default it on for this feature. +export MORI_ENABLE_SDMA="${MORI_ENABLE_SDMA:-1}" +export MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-8G}" + +# MORI's single-node eager path is the correctness-safe default on the ROCm +# versions used by Primus v26.4. Explicit user settings still win. +export MORI_HIER_CUDA_GRAPH="${MORI_HIER_CUDA_GRAPH:-0}" + +if [[ -z "${MORI_SOCKET_IFNAME:-}" && -n "${NCCL_SOCKET_IFNAME:-}" ]]; then + export MORI_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME#=}" +fi + +# primus-cli direct does not implicitly propagate host env into torchrun +# children. Emit every MORI_* variable as env.* so user-selected MORI tuning +# knobs (host-proxy, RDMA devices, async, graph/debug flags, etc.) survive. +echo "env.FSDP_ALL_GATHER_BACKEND=mori" +for name in "${!MORI_@}"; do + echo "env.${name}=${!name}" +done diff --git a/runner/helpers/mori/install_mori.sh b/runner/helpers/mori/install_mori.sh new file mode 100755 index 000000000..4f32e5c8b --- /dev/null +++ b/runner/helpers/mori/install_mori.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +set -euo pipefail + +MORI_REPO="${MORI_REPO:-https://github.com/ROCm/mori.git}" +MORI_REF="${MORI_REF:-12d1bc32d0c93dcd5062e74f4e0f772e36e1aac4}" +MORI_SOURCE_DIR="${MORI_SOURCE_DIR:-/opt/mori}" +MAX_JOBS="${MAX_JOBS:-32}" +ROCM_PATH="${ROCM_PATH:-/opt/rocm}" + +if [[ "$(id -u)" -ne 0 ]]; then + echo "[MORI:Install] run as root inside the training container" >&2 + exit 1 +fi +if [[ -z "${MORI_SOURCE_DIR}" || "${MORI_SOURCE_DIR}" == "/" ]]; then + echo "[MORI:Install] unsafe MORI_SOURCE_DIR=${MORI_SOURCE_DIR@Q}" >&2 + exit 1 +fi +if [[ ! -d "${ROCM_PATH}" ]]; then + echo "[MORI:Install] ROCM_PATH does not exist: ${ROCM_PATH}" >&2 + exit 1 +fi + +echo "[MORI:Install] install dependencies" +apt-get update +apt-get install -y --no-install-recommends \ + git ibverbs-utils libibverbs-dev libnuma-dev libpci-dev +python3 -m pip install --no-cache-dir \ + "packaging<26" "setuptools==81.0.0" setuptools_scm Cython pybind11 ninja + +echo "[MORI:Install] clone ${MORI_REPO}@${MORI_REF}" +rm -rf "${MORI_SOURCE_DIR}" +git clone --filter=blob:none "${MORI_REPO}" "${MORI_SOURCE_DIR}" +cd "${MORI_SOURCE_DIR}" +git checkout "${MORI_REF}" +git submodule update --init --depth 1 3rdparty/msgpack-c 3rdparty/spdlog + +if [[ ! -e /usr/lib64/libc.so ]]; then + mkdir -p /usr/lib64 + ln -s /usr/lib/x86_64-linux-gnu/libc.so /usr/lib64/libc.so +fi + +echo "[MORI:Install] build and verify" +export ROCM_PATH +export CMAKE_PREFIX_PATH="${ROCM_PATH}:${ROCM_PATH}/lib/rocm_sysdeps${CMAKE_PREFIX_PATH:+:${CMAKE_PREFIX_PATH}}" +export CMAKE_LIBRARY_PATH="/opt/mori-host-libs${CMAKE_LIBRARY_PATH:+:${CMAKE_LIBRARY_PATH}}" +export LD_LIBRARY_PATH="/opt/mori-host-libs:${ROCM_PATH}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" +unset MORI_DEVICE_NIC +MAX_JOBS="${MAX_JOBS}" python3 -m pip install --no-build-isolation . +python3 -c \ + "import mori; from mori.ccl import HierAllGather; print(f'MORI {mori.__version__}: {HierAllGather}')" +echo "[MORI:Install] PASS" diff --git a/runner/helpers/mori/multinode_allgather_smoke.py b/runner/helpers/mori/multinode_allgather_smoke.py new file mode 100644 index 000000000..a8804e2b8 --- /dev/null +++ b/runner/helpers/mori/multinode_allgather_smoke.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""Two-node correctness smoke for Primus's MORI FSDP all-gather adapter.""" + +import argparse +import os + +import torch +import torch.distributed as dist + +from primus.backends.common.mori_allgather import MoriAllGather + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--numel", type=int, default=64 * 1024 * 1024) + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + # Keep RCCL lazy: this smoke validates MORI's RDMA path, and no torch + # collective is needed before MORI initializes from the default c10d store. + dist.init_process_group("nccl") + + rank = dist.get_rank() + world_size = dist.get_world_size() + ranks_per_node = int(os.environ["LOCAL_WORLD_SIZE"]) + + input_tensor = torch.full( + (args.numel,), + rank + 1, + dtype=torch.bfloat16, + device=device, + ) + output_tensor = torch.empty( + args.numel * world_size, + dtype=input_tensor.dtype, + device=device, + ) + + work = MoriAllGather(ranks_per_node=ranks_per_node)( + output_tensor, + input_tensor, + dist.group.WORLD, + async_op=True, + ) + if work is not None: + work.wait() + torch.cuda.synchronize(device) + + expected = torch.repeat_interleave( + torch.arange(1, world_size + 1, dtype=input_tensor.dtype, device=device), + args.numel, + ) + if not torch.equal(output_tensor, expected): + raise RuntimeError(f"MORI all-gather mismatch on rank {rank}") + + if rank == 0: + size_mb = args.numel * input_tensor.element_size() / (1 << 20) + print( + f"MORI_MULTINODE_ALLGATHER_PASS world={world_size} " + f"ranks_per_node={ranks_per_node} per_rank_mb={size_mb:.1f}", + flush=True, + ) + + import mori.shmem as shmem + + shmem.shmem_finalize() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/runner/primus-cli-direct.sh b/runner/primus-cli-direct.sh index c90a0d86d..4db1ec93a 100755 --- a/runner/primus-cli-direct.sh +++ b/runner/primus-cli-direct.sh @@ -100,6 +100,9 @@ Examples: # Run silently (back-pocket option; launcher errors and log file preserved) primus-cli direct --silent -- preflight --quick + # Runtime NIC inspection, MORI build, and local all-gather correctness smoke + primus-cli direct -- preflight --mori + Notes: - If --single is specified, Primus skips torchrun and uses python3 directly. - run_mode auto-detection: when the primus subcommand is 'node_smoke', run_mode @@ -415,6 +418,15 @@ if [[ -z "${direct_config[run_mode]:-}" ]]; then break fi done + if [[ "$_detected_subcmd" == "preflight" ]]; then + for _arg in "${primus_args[@]}"; do + if [[ "$_arg" == "--mori" ]]; then + _default_run_mode="single" + LOG_INFO_RANK0 "[direct] Auto-selected run_mode=single for 'preflight --mori'" + break + fi + done + fi direct_config[run_mode]="$_default_run_mode" fi diff --git a/tests/runner/test_primus_cli_direct.sh b/tests/runner/test_primus_cli_direct.sh index 86ea1a542..bc3381751 100755 --- a/tests/runner/test_primus_cli_direct.sh +++ b/tests/runner/test_primus_cli_direct.sh @@ -441,7 +441,18 @@ test_auto_single_for_node_smoke() { assert_contains "$out_preflight" "Run Mode : torchrun" "preflight defaults to torchrun" assert_contains "$out_preflight" "torchrun --nproc_per_node" "preflight uses torchrun command" - # Sub-test 13b: node_smoke auto-selects single mode. + # Sub-test 13b: MORI preflight auto-selects single mode because it + # launches its own temporary-container torchrun workload. + local out_mori + out_mori=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- preflight --mori 2>&1 || true) + assert_contains "$out_mori" "Auto-selected run_mode=single for 'preflight --mori'" \ + "MORI preflight auto-detect fires" + assert_contains "$out_mori" "Run Mode : single" "MORI preflight uses single mode" + assert_contains "$out_mori" "python3" "MORI preflight uses python3 launcher" + assert_not_contains "$out_mori" "torchrun --nproc_per_node" \ + "MORI preflight does NOT use launcher torchrun" + + # Sub-test 13c: node_smoke auto-selects single mode. local out_smoke out_smoke=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- node_smoke --tier2-perf 2>&1 || true) assert_contains "$out_smoke" "Auto-selected run_mode=single for subcommand 'node_smoke'" \ @@ -450,7 +461,7 @@ test_auto_single_for_node_smoke() { assert_contains "$out_smoke" "python3" "node_smoke uses python3 launcher" assert_not_contains "$out_smoke" "torchrun --nproc_per_node" "node_smoke does NOT use torchrun" - # Sub-test 13c: explicit --single on a non-node_smoke subcommand still + # Sub-test 13d: explicit --single on a non-node_smoke subcommand still # works (regression guard). local out_explicit out_explicit=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --single --dry-run -- benchmark gemm 2>&1 || true) diff --git a/tests/unit_tests/cli/test_mori_preflight_helper.py b/tests/unit_tests/cli/test_mori_preflight_helper.py new file mode 100644 index 000000000..9740b731e --- /dev/null +++ b/tests/unit_tests/cli/test_mori_preflight_helper.py @@ -0,0 +1,76 @@ +"""Unit tests for the native MORI multi-node preflight orchestrator.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from primus.tools.preflight import mori_preflight as helper + + +def test_resolve_nodes_csv_deduplicates(): + assert helper.resolve_nodes("node1,node2,node1,node3") == [ + "node1", + "node2", + "node3", + ] + + +def test_resolve_nodes_file(tmp_path): + node_file = tmp_path / "nodes.txt" + node_file.write_text("node1\nnode2\nnode1\n") + assert helper.resolve_nodes(f"@{node_file}") == ["node1", "node2"] + + +def test_network_probe_uses_one_remote_command(monkeypatch): + calls = [] + monkeypatch.setattr( + helper, + "remote_output", + lambda *args: calls.append(args) or "fenic\t1\t10.0.0.1", + ) + assert helper.probe_network("node1", None, None, None) == ("fenic", 1, "10.0.0.1") + assert len(calls) == 1 + + +def test_multinode_smoke_launches_every_rank(monkeypatch, tmp_path): + args = SimpleNamespace( + mori_socket_ifname="fenic", + mori_gid_index=1, + mori_master_addr="10.0.0.1", + mori_master_port=29610, + mori_smoke_numel=1024, + ) + commands = [] + + def fake_run(command, **_kwargs): + commands.append(command) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(helper.subprocess, "run", fake_run) + monkeypatch.setattr(helper, "is_local", lambda _node: True) + monkeypatch.setattr(helper.time, "sleep", lambda _seconds: None) + + assert helper.run_multinode(args, ["node1", "node2"], tmp_path) == 0 + launches = [command[-1] for command in commands] + assert len(launches) == 2 + assert any("--node_rank=0" in command for command in launches) + assert any("--node_rank=1" in command for command in launches) + + +def test_matching_fingerprints_pass(tmp_path): + results = [ + helper.NodeResult("node1", 0, "nic=ionic hash=abc ccqe=True", tmp_path / "1.log"), + helper.NodeResult("node2", 0, "nic=ionic hash=abc ccqe=True", tmp_path / "2.log"), + ] + helper.validate_fingerprints(results) + + +def test_mismatched_fingerprints_fail(tmp_path): + results = [ + helper.NodeResult("node1", 0, "nic=ionic hash=abc ccqe=True", tmp_path / "1.log"), + helper.NodeResult("node2", 0, "nic=ionic hash=def ccqe=False", tmp_path / "2.log"), + ] + with pytest.raises(RuntimeError, match="node-stack mismatch"): + helper.validate_fingerprints(results) diff --git a/tests/unit_tests/cli/test_preflight_subcommand.py b/tests/unit_tests/cli/test_preflight_subcommand.py index 47a997b5d..b9887b6e1 100644 --- a/tests/unit_tests/cli/test_preflight_subcommand.py +++ b/tests/unit_tests/cli/test_preflight_subcommand.py @@ -14,8 +14,10 @@ from __future__ import annotations import argparse +from pathlib import Path from primus.cli.subcommands import preflight +from primus.tools.preflight import mori_preflight def _build_parser(): @@ -43,6 +45,9 @@ def test_defaults(): # (preflight-{NNODES}N-{YYYYMMDD-HHMMSS}) at run time. assert args.report_file_name is None assert args.save_pdf is True + assert args.mori is False + assert args.mori_nodes is None + assert args.mori_smoke_numel == 67108864 def test_selection_flags(): @@ -71,3 +76,54 @@ def test_perf_test_flag(): args = parser.parse_args(["preflight", "--perf-test", "--plot"]) assert args.perf_test is True assert args.plot is True + + +def test_mori_mode_flags(): + parser, _ = _build_parser() + args = parser.parse_args( + [ + "preflight", + "--mori", + "--mori-nodes", + "node-1,node-2,node-3", + "--mori-socket-ifname", + "fenic", + "--mori-gid-index", + "1", + ] + ) + assert args.mori is True + assert args.mori_nodes == "node-1,node-2,node-3" + assert args.mori_socket_ifname == "fenic" + assert args.mori_gid_index == 1 + + +def test_mori_mode_maps_cli_to_helper_args(monkeypatch, tmp_path): + parser, _ = _build_parser() + args = parser.parse_args( + [ + "preflight", + "--mori", + "--dump-path", + str(tmp_path), + "--mori-nodes", + "node-1,node-2", + "--mori-max-jobs", + "12", + ] + ) + repo_root = Path(mori_preflight.__file__).resolve().parents[3] + monkeypatch.setenv("PRIMUS_PATH", str(repo_root)) + captured = [] + monkeypatch.setattr( + mori_preflight, + "run_orchestrator", + lambda *values: captured.extend(values) or 0, + ) + assert mori_preflight.run_mori_preflight(args, []) == 0 + native_args, native_repo_root, native_log_dir = captured + assert native_args is args + assert native_args.mori_nodes == "node-1,node-2" + assert native_args.mori_max_jobs == 12 + assert native_log_dir.parent == tmp_path + assert native_repo_root == repo_root