diff --git a/benchmark/kernel/ep_bench/README.md b/benchmark/kernel/ep_bench/README.md new file mode 100644 index 000000000..accb51d1c --- /dev/null +++ b/benchmark/kernel/ep_bench/README.md @@ -0,0 +1,145 @@ +# MoRI EP-Bench (Slurm, no torchrun) + +This directory provides a Slurm-based launcher for the **MoRI** expert-parallel +dispatch/combine microbenchmarks on AMD Instinct GPUs. + + +## Benchmarks covered + +| Scope | Mode | Script | +|-----------|-------------|------------------------------------------------------------------| +| Intranode | Normal (bf16) | `tests/python/ops/bench_dispatch_combine.py` | +| Intranode | Low latency (fp8) | `tests/python/ops/bench_dispatch_combine.py --dtype fp8_e4m3_fnuz` | +| Internode | Normal (v1, bf16) | `examples/ops/dispatch_combine/test_dispatch_combine_internode.py --cmd bench` | +| Internode | Low latency (v1_ll) | `examples/ops/dispatch_combine/test_dispatch_combine_internode.py --cmd bench --kernel-type v1_ll` | + +Intranode runs trigger automatically when `NNODES == 1`; internode runs trigger +when `NNODES > 1`. (DeepEP tests are out of scope for this launcher.) + +## Prerequisites + +- Slurm cluster with AMD Instinct nodes. +- InfiniBand HCA(s) available on each node (default `mlx5_0`). +- Docker available on compute nodes. +- A Docker image that ships MoRI under `/app/mori`. Two options: + + ### Option A (recommended): build the slim MoRI image shipped here + + Built on the latest stable upstream + [`vllm/vllm-openai-rocm`](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) + image (which already provides ROCm, PyTorch, and vLLM) with **only MoRI** + added on top. + + ```bash + cd Primus/benchmark/kernel/ep_bench/docker + + # Defaults: BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0, GPU_ARCH=gfx942, + # MORI_REPO=https://github.com/ROCm/mori.git, MORI_REF=main + docker build -f Dockerfile.mori -t primus-mori-ep:latest . + + # MI355 / gfx950 + pinned MoRI ref: + docker build -f Dockerfile.mori \ + --build-arg GPU_ARCH=gfx950 \ + --build-arg MORI_REF=v0.x.y \ + -t primus-mori-ep:mi355 . + + # Precompile MoRI JIT kernels at build time (needs a GPU exposed to the build): + DOCKER_BUILDKIT=0 docker build -f Dockerfile.mori \ + --build-arg MORI_PRECOMPILE=1 \ + -t primus-mori-ep:precompiled . + ``` + + Build args supported by [`docker/Dockerfile.mori`](./docker/Dockerfile.mori): + + | Build arg | Default | Purpose | + |--------------------|------------------------------------------|----------------------------------------------------| + | `BASE_IMAGE` | `vllm/vllm-openai-rocm:v0.21.0` | Override to pin a different upstream vLLM ROCm tag | + | `GPU_ARCH` | `gfx942` | Sets `PYTORCH_ROCM_ARCH` for the MoRI install | + | `MORI_REPO` | `https://github.com/ROCm/mori.git` | MoRI git remote | + | `MORI_REF` | `main` | Branch / tag / SHA to check out | + | `MORI_PRECOMPILE` | `0` | If `1`, JIT-precompile MoRI kernels into the image | + + ### Option B: reuse the MAD `large_ep_benchmark` image + + Heavier image that also includes rocSHMEM + DeepEP + perftest. Build from + the MAD repository root: + + ```bash + # From the MAD repository root + docker build \ + -f docker/large_ep_benchmark.ubuntu.amd.Dockerfile \ + -t ep-benchmarking:latest . + ``` + + See + [`MAD/scripts/large-ep-benchmark/README.md`](https://github.com/ROCm/MAD/blob/main/scripts/large-ep-benchmark/README.md) + for cluster/Dockerfile customization (GPU arch, NIC arch). + +## Environment variables + +| Variable | Required | Default | Description | +|---------------------|----------|----------------------------------------|--------------------------------------------------------------| +| `DOCKER_IMAGE` | yes | — | Docker image to launch (e.g. `primus-mori-ep:latest` or `ep-benchmarking:latest`). | +| `NNODES` | no | `$SLURM_NNODES`, then `1` | Number of nodes; drives intranode vs internode branch. | +| `PARTITION` | no | unset (prefer `-p` on `sbatch`) | Slurm partition for child `srun` calls. | +| `GPUS_PER_NODE` | no | `8` | GPUs per node; forwarded as `GPU_PER_NODE` to MoRI. | +| `MASTER_PORT` | no | `2373` | Rendezvous port for `dist.init_process_group`. | +| `IBDEVICES` | no | `mlx5_0` | IB HCA(s) for `ROCSHMEM_USE_IB_HCA`. | +| `LOG_DIR` | no | `${SLURM_SUBMIT_DIR}/logs` | Host directory mounted as `/run_logs` inside the container. | +| `EXTRA_DOCKER_ARGS` | no | unset | Extra args appended to `docker run`. | + +## Usage + +### Single node (intranode only) + +```bash +cd Primus/benchmark/kernel/ep_bench + +DOCKER_IMAGE=primus-mori-ep:latest \ + sbatch -N 1 -p run_slurm.sh +``` + +### Multi-node (internode) + +```bash +cd Primus/benchmark/kernel/ep_bench + +DOCKER_IMAGE=primus-mori-ep:latest \ +IBDEVICES=mlx5_0 \ +NNODES=4 \ + sbatch -N 4 -p -w run_slurm.sh +``` + +### Overriding the rendezvous port or log directory + +```bash +DOCKER_IMAGE=primus-mori-ep:latest \ +MASTER_PORT=39566 \ +LOG_DIR=/shared/logs/mori-$(date +%Y%m%d) \ + sbatch -N 4 -p run_slurm.sh +``` + +## Output + +- Slurm job stdout/stderr go to the default `slurm-.out` (or whatever + you configure with `-o` / `-e` on `sbatch`). +- Per-test logs are written under `${LOG_DIR}` (default + `${SLURM_SUBMIT_DIR}/logs`): + + | File | Produced when | + |---------------------------------------------------|---------------| + | `mori_intranode_bf16.log` | `NNODES == 1` | + | `mori_intranode_fp8_ll.log` | `NNODES == 1` | + | `mori_internode_v1_rank.log` | `NNODES > 1` | + | `mori_internode_v1_ll_rank.log` | `NNODES > 1` | + +## Files + +- [`run_slurm.sh`](./run_slurm.sh) — sbatch entrypoint; pre-pulls the image + and `docker run`s the inner script on each node with one task per node. +- [`run_mori_bench.sh`](./run_mori_bench.sh) — inner script executed inside + the container; branches on `NNODES` and invokes MoRI bench scripts directly + with plain `python` (no torchrun). +- [`docker/Dockerfile.mori`](./docker/Dockerfile.mori) — slim image recipe + layered on the latest stable `vllm/vllm-openai-rocm`, installing only MoRI + (mirrors the MoRI portion of the MAD `large_ep_benchmark` Dockerfile). diff --git a/benchmark/kernel/ep_bench/docker/Dockerfile.mori b/benchmark/kernel/ep_bench/docker/Dockerfile.mori new file mode 100644 index 000000000..2024eb1b0 --- /dev/null +++ b/benchmark/kernel/ep_bench/docker/Dockerfile.mori @@ -0,0 +1,58 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +# +############################################################################### +# +# Slim image for the Primus MoRI EP microbenchmarks. +# +# Based on the upstream vLLM ROCm image (which already provides ROCm, PyTorch, +# and vLLM and MoRI) and inspired by ROCm/MAD's docker/large_ep_benchmark.ubuntu.amd.Dockerfile, +# +# Build (from this directory): +# docker build \ +# -f Dockerfile.mori \ +# --build-arg GPU_ARCH=gfx942 \ +# -t primus-mori-ep:latest . +# +# Use with the Slurm launcher: +# DOCKER_IMAGE=primus-mori-ep:latest sbatch -N 1 -p \ +# ../run_slurm.sh +# +############################################################################### + +# Latest stable upstream vLLM ROCm image. +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0 +FROM ${BASE_IMAGE} + +# +## The upstream vllm-openai-rocm image sets ENTRYPOINT to launch the vLLM API +## server; clear it so the container can be driven by our bench scripts. +ENTRYPOINT [] +CMD ["/bin/bash"] +# +WORKDIR /app +RUN sed -i 's/http/https/g' /etc/apt/sources.list + +RUN apt update && \ + apt install -y \ + autoconf pkg-config \ + libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ + libibverbs-dev ibverbs-utils libtool libboost-all-dev \ + libgrpc++-dev protobuf-compiler-grpc protobuf-compiler libprotobuf-dev \ + libaio-dev liburing-dev pybind11-dev ninja-build libgflags-dev \ + rdma-core infiniband-diags perftest openssh-server \ + psmisc vim cmake-curses-gui + +WORKDIR /app + +# not installing mori since its already installed in vllm container. +RUN pip install tqdm prettytable +RUN git clone --recursive $(grep '^MORI_REPO:' versions.txt | cut -d' ' -f2) && \ + cd mori && \ + git checkout $(grep '^MORI_BRANCH:' /app/versions.txt | cut -d' ' -f2) + +RUN cd /app/mori && git log | head diff --git a/benchmark/kernel/ep_bench/run_mori_bench.sh b/benchmark/kernel/ep_bench/run_mori_bench.sh new file mode 100755 index 000000000..0f35457dd --- /dev/null +++ b/benchmark/kernel/ep_bench/run_mori_bench.sh @@ -0,0 +1,121 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Inner script (runs inside the MAD large_ep_benchmark Docker container) that +# executes MoRI dispatch/combine microbenchmarks WITHOUT torchrun. +# +# Both MoRI bench scripts use torch.multiprocessing.spawn internally to spawn +# one process per GPU, so this script is invoked exactly once per node. +# +# For the internode bench, MoRI expects: +# RANK = node rank (0..NNODES-1) +# WORLD_SIZE = number of nodes +# MASTER_ADDR = rendezvous host (first node) +# MASTER_PORT = rendezvous port +# These are derived from SLURM_NODEID / SLURM_NNODES / etc. by the outer +# run_slurm.sh and forwarded as NODE_RANK / NNODES / MASTER_ADDR / MASTER_PORT. +# +# Expected environment (set by run_slurm.sh): +# NNODES, NODE_RANK, MASTER_ADDR, MASTER_PORT, GPUS_PER_NODE, IBDEVICES +# +# Output logs are written under /run_logs (host: LOG_DIR). +############################################################################### + +set -eo pipefail + +MASTER_ADDR="${MASTER_ADDR:-localhost}" +MASTER_PORT="${MASTER_PORT:-2373}" +NODE_RANK="${NODE_RANK:-0}" +NNODES="${NNODES:-1}" +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +IBDEVICES="${IBDEVICES:-mlx5_0}" + +LOG_DIR="${LOG_DIR_IN_CONTAINER:-/run_logs}" +mkdir -p "${LOG_DIR}" + +HOST_NAME="$(hostname)" +HOST_IP="$(hostname -I | awk '{print $1}')" + +# Resolve default-route interface for Gloo (TCP) rendezvous. +DEFAULT_IFACE="$(ip route | awk '/^default/ {print $5; exit}')" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${DEFAULT_IFACE}}" + +# rocSHMEM / MoRI RDMA env (shared with MAD run_ep_bench.sh defaults). +export ROCSHMEM_USE_IB_HCA="${ROCSHMEM_USE_IB_HCA:-${IBDEVICES}}" +export ROCSHMEM_MAX_NUM_CONTEXTS="${ROCSHMEM_MAX_NUM_CONTEXTS:-144}" + +MORI_PATH="${MORI_PATH:-/app/mori}" +if [[ ! -d "${MORI_PATH}" ]]; then + echo "[ERROR] MoRI not found at ${MORI_PATH}. Set MORI_PATH or use an image that ships MoRI." >&2 + exit 2 +fi +export PYTHONPATH="${MORI_PATH}:${PYTHONPATH:-}" + +cat <&1 | tee "${LOG_DIR}/mori_intranode_bf16.log" + + echo "[Node ${NODE_RANK}] Running MoRI INTRANODE low-latency benchmark (fp8_e4m3_fnuz)..." + python tests/python/ops/bench_dispatch_combine.py --dtype fp8_e4m3_fnuz \ + 2>&1 | tee "${LOG_DIR}/mori_intranode_fp8_ll.log" +else + # Internode: MoRI internode test reads RANK as node-rank and WORLD_SIZE as + # number of nodes (it spawns one process per GPU itself), so we just + # promote NODE_RANK / NNODES and call python directly. + export RANK="${NODE_RANK}" + export WORLD_SIZE="${NNODES}" + export MASTER_ADDR + export MASTER_PORT + export GPU_PER_NODE="${GPUS_PER_NODE}" + + INTERNODE_SCRIPT="examples/ops/dispatch_combine/test_dispatch_combine_internode.py" + + export GLOO_SOCKET_IFNAME=$(ip route | grep '^default' | awk '{print $NF}' | head -n 1) + + echo "[Node ${NODE_RANK}] Running MoRI INTERNODE dispatch/combine benchmark (v1, bf16)..." + torchrun --nnodes=$NNODES \ + --node_rank=$NODE_RANK \ + --nproc_per_node=1 \ + --master_addr=$MASTER_ADDR \ + --master_port=$MASTER_PORT \ + "${INTERNODE_SCRIPT}" --cmd bench \ + 2>&1 | tee "${LOG_DIR}/mori_internode_v1_rank${NODE_RANK}.log" + + sleep 10 + + echo "[Node ${NODE_RANK}] Running MoRI INTERNODE low-latency benchmark (v1_ll)..." + torchrun --nnodes=$NNODES \ + --node_rank=$NODE_RANK \ + --nproc_per_node=1 \ + --master_addr=$MASTER_ADDR \ + --master_port=$MASTER_PORT \ + "${INTERNODE_SCRIPT}" --cmd bench --kernel-type v1_ll \ + 2>&1 | tee "${LOG_DIR}/mori_internode_v1_ll_rank${NODE_RANK}.log" +fi + +echo "[Node ${NODE_RANK}] MoRI EP benchmarks complete. Logs under ${LOG_DIR}/" +exit 0 diff --git a/benchmark/kernel/ep_bench/run_slurm.sh b/benchmark/kernel/ep_bench/run_slurm.sh new file mode 100755 index 000000000..7eb9377df --- /dev/null +++ b/benchmark/kernel/ep_bench/run_slurm.sh @@ -0,0 +1,161 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Slurm launcher for MoRI EP (dispatch/combine) microbenchmarks inside a +# Docker container, WITHOUT torchrun. +# +# Both MoRI bench scripts spawn one process per GPU internally via +# torch.multiprocessing.spawn, so we only need ONE task per node. SLURM +# directly provides node rank / nnodes / master address; we forward those as +# RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT to the internode script. +# +# Usage: +# DOCKER_IMAGE=primus-mori-ep:latest sbatch -N 1 -p run_slurm.sh +# DOCKER_IMAGE=primus-mori-ep:latest NNODES=4 IBDEVICES=mlx5_0 \ +# sbatch -N 4 -p -w run_slurm.sh +# +# Environment variables (all optional except DOCKER_IMAGE): +# DOCKER_IMAGE Docker image to use (required). Two options: +# - primus-mori-ep:latest (built from +# docker/Dockerfile.mori in this directory) +# - ep-benchmarking:latest (MAD large_ep_benchmark) +# NNODES Number of nodes [default: SLURM_NNODES, then 1] +# PARTITION Slurm partition [default: unset; prefer -p on sbatch] +# GPUS_PER_NODE GPUs per node [default: 8] +# MASTER_PORT Port for distributed rendezvous [default: 2373] +# IBDEVICES InfiniBand HCA(s) for rocSHMEM / MoRI +# [default: mlx5_0] +# LOG_DIR Host directory for benchmark logs +# [default: ${SLURM_SUBMIT_DIR}/logs] +# EXTRA_DOCKER_ARGS Extra arguments passed to docker run +# +############################################################################### + +#SBATCH --exclusive +#SBATCH --ntasks-per-node=1 +#SBATCH --gpus-per-node=8 +#SBATCH --job-name=mori-ep-bench +#SBATCH --time=04:00:00 + +set -euo pipefail + +NNODES="${NNODES:-${SLURM_NNODES:-1}}" +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +MASTER_PORT="${MASTER_PORT:-2373}" +IBDEVICES="${IBDEVICES:-mlx5_0}" + +SCRIPT_DIR="${SLURM_SUBMIT_DIR:-$(cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")" && pwd)}" +LOG_DIR="${LOG_DIR:-${SCRIPT_DIR}/logs}" +mkdir -p "${LOG_DIR}" + +if [[ -z "${DOCKER_IMAGE:-}" ]]; then + echo "[ERROR] DOCKER_IMAGE is not set. Export it before submitting the job." + echo " Example: DOCKER_IMAGE=primus-mori-ep:latest sbatch -N 1 -p run_slurm.sh" + echo " (build the slim MoRI image via docker/Dockerfile.mori in this directory)" + exit 1 +fi + +# Build sbatch overrides from env vars +SBATCH_OVERRIDES=() +if [[ -n "${PARTITION:-}" ]]; then + SBATCH_OVERRIDES+=(-p "$PARTITION") +fi + +echo "================================================" +echo " MoRI EP-Bench - Slurm + Docker launcher" +echo "================================================" +echo " DOCKER_IMAGE : ${DOCKER_IMAGE}" +echo " NNODES : ${NNODES}" +echo " GPUS_PER_NODE : ${GPUS_PER_NODE}" +echo " MASTER_PORT : ${MASTER_PORT}" +echo " IBDEVICES : ${IBDEVICES}" +echo " SCRIPT_DIR : ${SCRIPT_DIR}" +echo " LOG_DIR : ${LOG_DIR}" +echo "================================================" + +# Pre-pull image + stop any stale containers on allocated nodes +# shellcheck disable=SC2016 +srun -N "${NNODES}" \ + --exclusive \ + --export=ALL \ + --ntasks-per-node=1 \ + "${SBATCH_OVERRIDES[@]}" \ + bash -c 'docker pull "${DOCKER_IMAGE}"; docker ps -q --filter "name=mori-ep-bench" | xargs -r docker stop >/dev/null 2>&1 || true; docker ps -aq --filter "name=mori-ep-bench" | xargs -r docker rm >/dev/null 2>&1 || true' + +# Run workload on allocated nodes +# shellcheck disable=SC2016 +srun -N "${NNODES}" \ + --exclusive \ + --export=ALL \ + --ntasks-per-node=1 \ + "${SBATCH_OVERRIDES[@]}" \ + bash -c ' + +set -euo pipefail + +# ---- Resolve master address from Slurm node list ---- +readarray -t NODE_ARRAY < <(scontrol show hostnames "$SLURM_JOB_NODELIST") +MASTER_ADDR="${NODE_ARRAY[0]}" +NODE_RANK="${SLURM_NODEID}" + +if [[ "$NODE_RANK" == "0" ]]; then + echo "========== Slurm cluster info ==========" + echo "SLURM_NODELIST : ${NODE_ARRAY[*]}" + echo "SLURM_NNODES : ${SLURM_NNODES}" + echo "MASTER_ADDR : ${MASTER_ADDR}" + echo "NODE_RANK : ${NODE_RANK}" + echo "" +fi + +SCRIPT_DIR='"${SCRIPT_DIR}"' +LOG_DIR='"${LOG_DIR}"' +DOCKER_IMAGE='"${DOCKER_IMAGE}"' +NNODES='"${NNODES}"' +GPUS_PER_NODE='"${GPUS_PER_NODE}"' +MASTER_PORT='"${MASTER_PORT}"' +IBDEVICES='"${IBDEVICES}"' +EXTRA_DOCKER_ARGS='"${EXTRA_DOCKER_ARGS:-}"' + +CONTAINER_NAME="mori-ep-bench-${SLURM_JOB_ID:-manual}-${NODE_RANK}" + +docker run --rm \ + --name "${CONTAINER_NAME}" \ + --network=host \ + --ipc=host \ + --device=/dev/kfd \ + --device=/dev/dri \ + --device=/dev/infiniband \ + --privileged \ + --cap-add=SYS_PTRACE \ + --cap-add=CAP_SYS_ADMIN \ + --security-opt seccomp=unconfined \ + --group-add video \ + --ulimit memlock=-1:-1 \ + --shm-size 64G \ + -v "${SCRIPT_DIR}:${SCRIPT_DIR}" \ + -v "${HOME}:${HOME}" \ + -v "${LOG_DIR}:/run_logs" \ + -v /dev/infiniband:/dev/infiniband \ + -v /sys/class/infiniband:/sys/class/infiniband:ro \ + -v /sys/class/net:/sys/class/net:ro \ + -v /sys/bus/pci:/sys/bus/pci:ro \ + -w "${SCRIPT_DIR}" \ + -e MASTER_ADDR="${MASTER_ADDR}" \ + -e MASTER_PORT="${MASTER_PORT}" \ + -e NNODES="${NNODES}" \ + -e NODE_RANK="${NODE_RANK}" \ + -e GPUS_PER_NODE="${GPUS_PER_NODE}" \ + -e IBDEVICES="${IBDEVICES}" \ + -e SLURM_JOB_ID="${SLURM_JOB_ID:-}" \ + -e SLURM_JOB_NODELIST="${SLURM_JOB_NODELIST:-}" \ + ${EXTRA_DOCKER_ARGS} \ + --entrypoint /bin/bash \ + "${DOCKER_IMAGE}" \ + "${SCRIPT_DIR}/run_mori_bench.sh" +' + +echo " Results written to: ${LOG_DIR}/" diff --git a/benchmark/kernel/rdma_perf/README.md b/benchmark/kernel/rdma_perf/README.md new file mode 100644 index 000000000..8b6bcf056 --- /dev/null +++ b/benchmark/kernel/rdma_perf/README.md @@ -0,0 +1,148 @@ +# Cluster RDMA Perf Tests + +Two-node `ib_write_bw` benchmark, launched via Slurm + Docker + +Rank 0 runs `ib_write_bw` as the server; rank 1 connects as the client. +A small `socket_barrier.py` coordination step ensures both containers are +up before the perftest handshake fires. + +## Contents + +| File | Purpose | +|---|---| +| `run_slurm.sh` | sbatch entrypoint (resolves per-node IPs, pre-pulls image, runs container). | +| `run_rdma_tests.sh` | Inner container script: cross-node barrier, then `ib_write_bw` server or client. | +| `socket_barrier.py` | TCP barrier — opens a local port, polls all peers until they accept. | +| `socket_wait.py` | Optional helper to wait while a remote port stays open. Not on the critical path. | + +## Prerequisites + +- Slurm with `sbatch` / `srun`. +- Docker on every allocated node, with access to `/dev/infiniband`, `/dev/kfd`, + `/dev/dri`. +- A Docker image that ships [`linux-rdma/perftest`](https://github.com/linux-rdma/perftest) + with the `ib_write_bw` binary in `PATH`. The default + (`lmsysorg/sglang:v0.5.7-rocm700-mi35x`) ships it. +- Healthy RoCE/IB fabric between the two nodes. See **Troubleshooting** below + if `ib_write_bw` hangs or fails to transition QPs. + +## Environment variables + +All are optional; only `sbatch -N 2 -p ` is strictly required. + +| Variable | Default | Description | +|---|---|---| +| `DOCKER_IMAGE` | `lmsysorg/sglang:v0.5.7-rocm700-mi35x` | Image with `ib_write_bw` available. | +| `NNODES` | `${SLURM_NNODES:-2}` | Number of allocated nodes (must be 2 for client/server). | +| `PARTITION` | _(unset)_ | Slurm partition; prefer `sbatch -p` instead. | +| `IBDEVICES` | `rdma0` | HCA passed to `ib_write_bw -d`. | +| `LOG_PATH` | `${SLURM_SUBMIT_DIR}/logs` | Host directory for `ib_write_bw_node{0,1}.log` and sbatch out/err. | +| `CONTAINER_NAME` | `primus-rdma-tests` | Docker container name. | +| `MASTER_PORT` | `39566` | Reserved rendezvous port (parity with other Primus launchers; not used by perftest). | +| `BARRIER_PORT` | `5000` | TCP port used by `socket_barrier.py`. | +| `IB_WRITE_BW_PORT` | `2000` | Data port passed as `ib_write_bw -p`. | +| `EXTRA_DOCKER_ARGS` | _(empty)_ | Appended to `docker run`. Use for site-local mounts (see below). | + +## Usage + +### Default 2-node single-NIC run + +```bash +LOG_PATH=/shared/logs sbatch -N 2 -p \ + Primus/benchmark/kernel/rdma_perf/run_slurm.sh +``` + +### Pick a different NIC + image + +```bash +IBDEVICES=rdma3 \ +DOCKER_IMAGE=my.registry/perftest:latest \ +LOG_PATH=/shared/logs \ + sbatch -N 2 -p Primus/benchmark/kernel/rdma_perf/run_slurm.sh +``` + +### Pin specific nodes + +```bash +sbatch -N 2 -p Compute-DCPT -w smci355-ccs-aus-n04-[25,29] \ + Primus/benchmark/kernel/rdma_perf/run_slurm.sh +``` + +## Output + +Inside the container the test writes to `/run_logs/ib_write_bw_node${NODE_RANK}.log` +which is the same as `${LOG_PATH}/ib_write_bw_node{0,1}.log` on the host. +Slurm stdout/stderr land in the submit directory as `primus_rdma_perf_.{out,err}`. + +A successful run prints a `ib_write_bw` results table on the client +(node 1) similar to: + +``` + #bytes #iterations BW peak[Gb/sec] BW average[Gb/sec] MsgRate[Mpps] + ... +``` + +## Site-specific overrides + +The default `docker run` mounts are intentionally minimal. If you may want to forward the vendor RDMA libraries, +Pass them via `EXTRA_DOCKER_ARGS`: + +```bash +EXTRA_DOCKER_ARGS="\ + -v /it-share:/it-share \ + -v /it-share/models:/data \ + -v \$HOME/.ssh:/root/.ssh \ + -v /usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:ro \ + -v /etc/libibverbs.d:/etc/libibverbs.d:ro \ + -v /usr/lib/x86_64-linux-gnu/libionic.so.1:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro \ + -v /usr/lib/x86_64-linux-gnu/libionic.so:/usr/lib/x86_64-linux-gnu/libionic.so:ro \ + -v /etc/libibverbs.d/ionic.driver:/etc/libibverbs.d/ionic.driver:ro \ + " \ +sbatch -N 2 -p Primus/benchmark/kernel/rdma_perf/run_slurm.sh +``` + +## Troubleshooting + +### `Failed to modify QP to RTR` / "Unable to Connect the HCA's through the link" + +Indicates a fabric-level problem before perftest even sends data. Common causes: + +- **Wrong GID index.** Check `ibv_devinfo -v -d $IBDEVICES` and inspect + `GID[*]` entries; only RoCEv2 GIDs that share the same VLAN/L3 routing on + both nodes will work. Pass `-x ` to `ib_write_bw` if the default + GID is unreachable. +- **Rail mismatch.** The HCA you picked on node 0 must be in the same rail + as the one on node 1. On multi-rail systems pick HCAs that share a + switch / L3 segment. +- **Ethernet OOB unreachable.** Add `-R` to `ib_write_bw` to use RDMA-CM for + connection establishment and bypass the OOB TCP handshake. +- **PFC / lossless config.** Verify the switch is configured for the + expected priority class and ECN is on. + +### Hang at "Waiting for nodes. . ." + +The `socket_barrier.py` is unable to reach `${BARRIER_PORT}` on one of the +peer hosts. Verify: + +- The container actually started on the peer (`docker ps` on that host). +- `BARRIER_PORT` is not blocked by a firewall on the management network. +- All peers can resolve each other's `hostname -I` first interface (which is + what `run_slurm.sh` records into `IPADDRS`). + +### `ib_write_bw: command not found` + +Your Docker image doesn't ship perftest. Either switch to a perftest-bearing +image (e.g. `lmsysorg/sglang:v0.5.7-rocm700-mi35x`) or add an install step +via `EXTRA_DOCKER_ARGS` + an entrypoint override. + +## Manual reference (without Slurm) + +For ad-hoc debugging without `sbatch`, on two interactively-allocated nodes: + +```bash +# Server (node A) +ib_write_bw -d rdma0 -q 4 -a --report_gbits -F -p 2000 + +# Client (node B), once server is listening +ib_write_bw -d rdma0 -q 4 -a --report_gbits -F -p 2000 +``` diff --git a/benchmark/kernel/rdma_perf/run_rdma_tests.sh b/benchmark/kernel/rdma_perf/run_rdma_tests.sh new file mode 100755 index 000000000..8078a9308 --- /dev/null +++ b/benchmark/kernel/rdma_perf/run_rdma_tests.sh @@ -0,0 +1,101 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Inner container script for Primus cluster RDMA perf tests. +# +# Expects to be launched by run_slurm.sh, which exports: +# NODE_RANK 0-based per-node index (rank 0 = server, rank 1 = client) +# NNODES number of allocated nodes +# MASTER_ADDR IB-routable IP of the rank-0 node (reserved; unused here) +# MASTER_PORT Reserved rendezvous port (kept for parity) +# IPADDRS comma-separated list of per-node IB-routable IPs +# IBDEVICES HCA passed to `ib_write_bw -d` +# BARRIER_PORT TCP port used by socket_barrier.py [default: 5000] +# IB_WRITE_BW_PORT Data port for ib_write_bw [default: 2000] +# RDMA_TESTS_PATH Path inside the container where this directory is mounted +# +# Ported from: +# ROCm/dist-inf-cookbook : cluster-sphere/cluster-rdma-tests/slurm_scripts/ +# run_rdma_tests.sh +# +# Logs are tee'd into /run_logs (= host LOG_PATH) so they survive the +# container teardown. +############################################################################### + +set -euo pipefail + +NODE_RANK="${NODE_RANK:-0}" +NNODES="${NNODES:-2}" +IPADDRS="${IPADDRS:-localhost}" +IBDEVICES="${IBDEVICES:-rdma0}" +BARRIER_PORT="${BARRIER_PORT:-5000}" +IB_WRITE_BW_PORT="${IB_WRITE_BW_PORT:-2000}" +RDMA_TESTS_PATH="${RDMA_TESTS_PATH:-/root/rdma-tests-repo}" + +LOG_DIR="/run_logs" +mkdir -p "${LOG_DIR}" +LOG_FILE="${LOG_DIR}/ib_write_bw_node${NODE_RANK}.log" + +HOST_NAME="$(hostname)" +HOST_IP="$(hostname -I | awk '{print $1}')" +SERVER_IP="$(echo "${IPADDRS}" | awk -F',' '{print $1}')" + +cat </dev/null 2>&1; then + echo "[ERROR] ib_write_bw not found in PATH inside container." | tee -a "${LOG_FILE}" >&2 + echo " Use a Docker image with linux-rdma/perftest preinstalled," | tee -a "${LOG_FILE}" >&2 + echo " or apt-install/build it via EXTRA_DOCKER_ARGS overrides." | tee -a "${LOG_FILE}" >&2 + exit 2 +fi + +# Cross-node container startup barrier (blocks until all peers are listening). +echo "[${HOST_NAME}] Waiting at the container creation barrier ..." | tee -a "${LOG_FILE}" +python "${RDMA_TESTS_PATH}/socket_barrier.py" \ + --local-ip "${HOST_IP}" \ + --local-port "${BARRIER_PORT}" \ + --enable-port \ + --node-ips "${IPADDRS}" \ + --node-ports "${BARRIER_PORT}" 2>&1 | tee -a "${LOG_FILE}" + +if [[ "${NODE_RANK}" -eq 0 ]]; then + echo "-------------------------------------------------" | tee -a "${LOG_FILE}" + echo "[${HOST_NAME}:${HOST_IP}] Running ib_write_bw as SERVER" | tee -a "${LOG_FILE}" + echo "-------------------------------------------------" | tee -a "${LOG_FILE}" + + ib_write_bw -d "${IBDEVICES}" -q 4 -a --report_gbits -F -p "${IB_WRITE_BW_PORT}" \ + 2>&1 | tee -a "${LOG_FILE}" +else + echo "-------------------------------------------------" | tee -a "${LOG_FILE}" + echo "[${HOST_NAME}:${HOST_IP}] Running ib_write_bw as CLIENT against ${SERVER_IP}" | tee -a "${LOG_FILE}" + echo "-------------------------------------------------" | tee -a "${LOG_FILE}" + + echo "[${HOST_NAME}] Waiting for server port to open..." | tee -a "${LOG_FILE}" + sleep 30 + + ib_write_bw -d "${IBDEVICES}" -q 4 -a --report_gbits -F "${SERVER_IP}" -p "${IB_WRITE_BW_PORT}" \ + 2>&1 | tee -a "${LOG_FILE}" +fi + +echo "[${HOST_NAME}] Finished running RDMA tests. Log: ${LOG_FILE}" | tee -a "${LOG_FILE}" +exit 0 diff --git a/benchmark/kernel/rdma_perf/run_slurm.sh b/benchmark/kernel/rdma_perf/run_slurm.sh new file mode 100755 index 000000000..bc5d2758f --- /dev/null +++ b/benchmark/kernel/rdma_perf/run_slurm.sh @@ -0,0 +1,194 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Slurm launcher for cluster RDMA perf tests (ib_write_bw between two nodes) +# inside a Docker container. +# - env-driven configuration with sensible defaults +# - separate pre-pull and run srun stages +# - minimal default bind-mounts; site-specific mounts go through +# EXTRA_DOCKER_ARGS +# - no hardcoded site-local paths (/it-share, libionic.*, etc.) +# +# Each node runs exactly one container; rank 0 hosts the ib_write_bw server, +# rank 1 connects as the client. +# +# Usage: +# sbatch -N 2 -p run_slurm.sh +# IBDEVICES=rdma3 LOG_PATH=/shared/logs sbatch -N 2 -p run_slurm.sh +# +# Environment variables (all optional): +# DOCKER_IMAGE Docker image with linux-rdma/perftest preinstalled +# [default: lmsysorg/sglang:v0.5.7-rocm700-mi35x] +# NNODES Number of nodes [default: SLURM_NNODES, then 2] +# PARTITION Slurm partition [default: unset; prefer -p on sbatch] +# IBDEVICES InfiniBand HCA passed to ib_write_bw -d +# [default: mlx5_0] +# LOG_PATH Host directory for benchmark logs and sbatch out/err +# [default: ${SLURM_SUBMIT_DIR}/logs] +# CONTAINER_NAME Docker container name [default: primus-rdma-tests] +# MASTER_PORT Reserved rendezvous port (currently unused by perftest; +# kept for parity with other Primus launchers) +# [default: 39566] +# BARRIER_PORT TCP port used by socket_barrier.py [default: 5000] +# IB_WRITE_BW_PORT Data port passed to ib_write_bw -p [default: 2000] +# EXTRA_DOCKER_ARGS Extra args appended to `docker run` (e.g. site-local +# bind mounts like /it-share, libionic.*) +# +############################################################################### + +#SBATCH --job-name=primus-rdma-perf +#SBATCH --ntasks-per-node=1 +#SBATCH --time=01:00:00 +#SBATCH --output=primus_rdma_perf_%j.out +#SBATCH --error=primus_rdma_perf_%j.err + +set -euo pipefail + +DOCKER_IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:v0.5.7-rocm700-mi35x}" +NNODES="${NNODES:-${SLURM_NNODES:-2}}" +IBDEVICES="${IBDEVICES:-mlx5_0}" +CONTAINER_NAME="${CONTAINER_NAME:-primus-rdma-tests}" +MASTER_PORT="${MASTER_PORT:-39566}" +BARRIER_PORT="${BARRIER_PORT:-5000}" +IB_WRITE_BW_PORT="${IB_WRITE_BW_PORT:-2000}" + +SCRIPT_DIR="$(cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")" && pwd)" +SUBMIT_DIR="${SLURM_SUBMIT_DIR:-${SCRIPT_DIR}}" +LOG_PATH="${LOG_PATH:-${SUBMIT_DIR}/logs}" +mkdir -p "${LOG_PATH}" + +RDMA_TESTS_REPO="${SCRIPT_DIR}" +RDMA_TESTS_PATH="/root/rdma-tests-repo" + +# --------------------------------------------------------------------------- +# Resolve master node + per-node IB-routable IPs from the Slurm allocation. +# We use `hostname -I` per node (same approach as upstream) because it returns +# the IP the IB stack will actually use; `getent` would give us the management +# IP which is often on a separate VLAN. +# --------------------------------------------------------------------------- +readarray -t NODE_ARRAY < <(scontrol show hostnames "${SLURM_JOB_NODELIST}") +MASTER_NODE="${NODE_ARRAY[0]}" + +IPS=() +for NODE in "${NODE_ARRAY[@]}"; do + IP="$(srun --nodes=1 --ntasks=1 -w "${NODE}" --export=ALL bash -c 'hostname -I' | awk 'NR==1{print $1}')" + IPS+=("${IP}") +done +MASTER_ADDR="${IPS[0]}" +IPADDRS="$(IFS=,; echo "${IPS[*]}")" + +# Build sbatch overrides from env vars +SBATCH_OVERRIDES=() +if [[ -n "${PARTITION:-}" ]]; then + SBATCH_OVERRIDES+=(-p "${PARTITION}") +fi + +echo "================================================" +echo " Primus Cluster RDMA Perf Tests - Slurm launcher" +echo "================================================" +echo " SLURM_JOB_ID : ${SLURM_JOB_ID:-N/A}" +echo " SLURM_JOB_NODES : ${NODE_ARRAY[*]}" +echo " NNODES : ${NNODES}" +echo " DOCKER_IMAGE : ${DOCKER_IMAGE}" +echo " CONTAINER_NAME : ${CONTAINER_NAME}" +echo " IBDEVICES : ${IBDEVICES}" +echo " MASTER_NODE : ${MASTER_NODE}" +echo " MASTER_ADDR : ${MASTER_ADDR}" +echo " MASTER_PORT : ${MASTER_PORT}" +echo " BARRIER_PORT : ${BARRIER_PORT}" +echo " IB_WRITE_BW_PORT : ${IB_WRITE_BW_PORT}" +echo " IPADDRS : ${IPADDRS}" +echo " SCRIPT_DIR : ${SCRIPT_DIR}" +echo " RDMA_TESTS_REPO : ${RDMA_TESTS_REPO}" +echo " RDMA_TESTS_PATH : ${RDMA_TESTS_PATH}" +echo " LOG_PATH : ${LOG_PATH}" +echo "================================================" + +export DOCKER_IMAGE NNODES IBDEVICES CONTAINER_NAME +export MASTER_ADDR MASTER_PORT BARRIER_PORT IB_WRITE_BW_PORT +export IPADDRS LOG_PATH RDMA_TESTS_REPO RDMA_TESTS_PATH +export EXTRA_DOCKER_ARGS="${EXTRA_DOCKER_ARGS:-}" + +# --------------------------------------------------------------------------- +# Stage A: pre-pull image + clean up any stale container of the same name. +# --------------------------------------------------------------------------- +# shellcheck disable=SC2016 +srun -N "${NNODES}" \ + --ntasks-per-node=1 \ + --export=ALL \ + "${SBATCH_OVERRIDES[@]}" \ + bash -c ' + echo "[$(hostname)] Pulling ${DOCKER_IMAGE} ..."; + docker pull "${DOCKER_IMAGE}"; + docker ps -q --filter "name=${CONTAINER_NAME}" | xargs -r docker stop >/dev/null 2>&1 || true; + docker ps -aq --filter "name=${CONTAINER_NAME}" | xargs -r docker rm >/dev/null 2>&1 || true; + ' + +# --------------------------------------------------------------------------- +# Stage B: run the ib_write_bw server (rank 0) / client (rank 1) on each node. +# --------------------------------------------------------------------------- +# shellcheck disable=SC2016 +srun -N "${NNODES}" \ + --ntasks-per-node=1 \ + --export=ALL \ + "${SBATCH_OVERRIDES[@]}" \ + bash -c ' + set -euo pipefail + NODE_RANK="${SLURM_NODEID:-${SLURM_PROCID:-0}}" + echo "[$(hostname)] NODE_RANK=${NODE_RANK} starting container ${CONTAINER_NAME}"; + + # shellcheck disable=SC2086 + docker run --rm \ + --name "${CONTAINER_NAME}" \ + --device /dev/dri \ + --device /dev/kfd \ + --device /dev/infiniband \ + --network host \ + --ipc host \ + --group-add video \ + --cap-add SYS_PTRACE \ + --privileged \ + --security-opt seccomp=unconfined \ + --ulimit memlock=-1:-1 \ + --shm-size 64G \ + -v "${RDMA_TESTS_REPO}:${RDMA_TESTS_PATH}" \ + -v "${LOG_PATH}:/run_logs" \ + -v /sys:/sys \ + -v /dev/infiniband:/dev/infiniband \ + -v /sys/class/infiniband:/sys/class/infiniband:ro \ + -v /sys/class/net:/sys/class/net:ro \ + -v /sys/bus/pci:/sys/bus/pci:ro \ + -e SLURM_JOB_ID="${SLURM_JOB_ID:-}" \ + -e SLURM_JOB_NODELIST="${SLURM_JOB_NODELIST:-}" \ + -e NODE_RANK="${NODE_RANK}" \ + -e NNODES="${NNODES}" \ + -e MASTER_ADDR="${MASTER_ADDR}" \ + -e MASTER_PORT="${MASTER_PORT}" \ + -e IPADDRS="${IPADDRS}" \ + -e IBDEVICES="${IBDEVICES}" \ + -e BARRIER_PORT="${BARRIER_PORT}" \ + -e IB_WRITE_BW_PORT="${IB_WRITE_BW_PORT}" \ + -e RDMA_TESTS_PATH="${RDMA_TESTS_PATH}" \ + ${EXTRA_DOCKER_ARGS} \ + "${DOCKER_IMAGE}" \ + bash -c "${RDMA_TESTS_PATH}/run_rdma_tests.sh" + ' + +# --------------------------------------------------------------------------- +# Stage C: defensive cleanup (matches upstream tail). +# --------------------------------------------------------------------------- +# shellcheck disable=SC2016 +srun -N "${NNODES}" \ + --ntasks-per-node=1 \ + --export=ALL \ + "${SBATCH_OVERRIDES[@]}" \ + bash -c ' + docker ps -q --filter "name=${CONTAINER_NAME}" | xargs -r docker stop >/dev/null 2>&1 || true; + docker ps -aq --filter "name=${CONTAINER_NAME}" | xargs -r docker rm >/dev/null 2>&1 || true; + ' + +echo " Results written to: ${LOG_PATH}/" diff --git a/benchmark/kernel/rdma_perf/socket_barrier.py b/benchmark/kernel/rdma_perf/socket_barrier.py new file mode 100644 index 000000000..fb9fd3981 --- /dev/null +++ b/benchmark/kernel/rdma_perf/socket_barrier.py @@ -0,0 +1,93 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Cross-node container startup barrier. +# +# Each node opens a TCP listen socket on (local-ip, local-port) and then polls +# every (node-ip, node-port) pair until all of them are accepting connections. +# This guarantees that the container processes on all peer nodes are up before +# the RDMA perftest handshake fires. +# +# Ported verbatim from: +# ROCm/dist-inf-cookbook : cluster-sphere/cluster-rdma-tests/slurm_scripts/ +# socket_barrier.py +############################################################################### +import socket +import time +import threading +import argparse + +parser = argparse.ArgumentParser(description="Optionally open and close a port on the local node.") +parser.add_argument("--local-ip", required=False, help="Local IP address to bind the server.") +parser.add_argument("--local-port", type=int, required=False, help="Port number to bind the server.") +parser.add_argument("--enable-port", action="store_true", help="Enable opening and closing of local port.") +parser.add_argument("--node-ips", required=True, help="Comma-separated list of node IPs.") +parser.add_argument("--node-ports", required=True, help="Comma-separated list of ports to check.") +args = parser.parse_args() + +NODE_IPS = [ip.strip() for ip in args.node_ips.split(",") if ip.strip()] +NODE_PORTS = [int(port.strip()) for port in args.node_ports.split(",") if port.strip()] + +if len(NODE_PORTS) == 1: + NODE_PORTS *= len(NODE_IPS) +elif len(NODE_PORTS) != len(NODE_IPS): + print("Error: Number of ports must match number of node IPs or only one port should be given for all.") + exit(1) + +server_socket = None + + +def is_port_open(ip, port): + """Check if a given IP and port are accessible.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.settimeout(2) + return s.connect_ex((ip, port)) == 0 + + +def wait_for_all_ports(): + """Wait until all nodes have opened the specified ports.""" + while True: + all_open = all(is_port_open(ip, port) for ip, port in zip(NODE_IPS, NODE_PORTS)) + if all_open: + break + print("Waiting for nodes. . .", flush=True) + time.sleep(5) + + +def open_port(): + """Open a listening socket on the current node.""" + global server_socket + server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server_socket.bind((args.local_ip, args.local_port)) + server_socket.listen(5) + print(f"Port {args.local_port} is now open on {args.local_ip}.") + while True: + conn, addr = server_socket.accept() + conn.close() + + +def close_port(): + """Close the opened port.""" + global server_socket + if server_socket: + server_socket.close() + print(f"Port {args.local_port} has been closed on {args.local_ip}.") + + +if __name__ == "__main__": + if not NODE_IPS: + print("Error: NODE_IPS argument is empty or not set.") + exit(1) + + if args.enable_port: + threading.Thread(target=open_port, daemon=True).start() + + wait_for_all_ports() + + if args.enable_port: + time.sleep(120) + close_port() diff --git a/benchmark/kernel/rdma_perf/socket_wait.py b/benchmark/kernel/rdma_perf/socket_wait.py new file mode 100644 index 000000000..45c6f3167 --- /dev/null +++ b/benchmark/kernel/rdma_perf/socket_wait.py @@ -0,0 +1,40 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Wait while a remote (host, port) endpoint stays open. Used as a teardown +# barrier for staged perftest runs. +# +# Ported verbatim from: +# ROCm/dist-inf-cookbook : cluster-sphere/cluster-rdma-tests/slurm_scripts/ +# socket_wait.py +############################################################################### +import socket +import time +import argparse + + +def is_port_open(host, port, timeout=2): + """Check if a given host and port are accessible.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.settimeout(timeout) + return s.connect_ex((host, port)) == 0 + + +def wait_while_port_open(host, port, check_interval=5): + """Wait while the remote port remains open.""" + print(f"Waiting while port {port} on {host} is open...") + while is_port_open(host, port): + time.sleep(check_interval) + print(f"Port {port} on {host} is now closed.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Wait while a remote port remains open.") + parser.add_argument("--remote-ip", required=True, help="Remote server IP address.") + parser.add_argument("--remote-port", type=int, required=True, help="Remote port number.") + args = parser.parse_args() + + wait_while_port_open(args.remote_ip, args.remote_port) diff --git a/tools/cluster-rdma-env-recommender/README.md b/tools/cluster-rdma-env-recommender/README.md new file mode 100644 index 000000000..508cb519f --- /dev/null +++ b/tools/cluster-rdma-env-recommender/README.md @@ -0,0 +1,32 @@ +# Cluster RDMA Environment Mapping & Recommender Tool + +Cluster RDMA Env Mapping & Recommender tool is an end to end RDMA environment discovery and recommendation tool that automatically inspects the host, maps the RDMA devices to physical NICs and generates production ready recommendations. This tool bridges the critical gap between bare metal RDMA hardware and Distributed Inference workloads by reducing manual trial and error. + +The tool contains the following features: + +* RDMA → PCI → NetDev mapping + +* Vendor Detection (AINIC/BNXT/MLNX) + +* Firmware & RoCEv2 GID Discovery + +* Recommended docker launch command. + +* Recommended Framework level environment variables required for workloads. + +* Human readable CLI report + +## Steps to run the tool + +On baremetal, run the following (Currently only works for Ubuntu) + +``` +cd tools/cluster-rdma-env-recommender +python3 cluster_rdma_env_recommender.py |& tee report.txt +``` + +The report.txt contains the details of the above mentioned features of a cluster. + + + + diff --git a/tools/cluster-rdma-env-recommender/cluster_rdma_env_recommender.py b/tools/cluster-rdma-env-recommender/cluster_rdma_env_recommender.py new file mode 100644 index 000000000..533f2861d --- /dev/null +++ b/tools/cluster-rdma-env-recommender/cluster_rdma_env_recommender.py @@ -0,0 +1,449 @@ +import os +import glob +import subprocess +import re +from typing import Tuple +import argparse + +LIB_SEARCH_PATHS = [ + "/usr/lib", + "/usr/lib64", + "/usr/lib/x86_64-linux-gnu", + "/usr/local/lib", + "/etc/libibverbs.d", +] + +class RDMAClusterMapper: + """RDMA cluster mapping utility for environment recommendation.""" + + def __init__(self): + self.rdma_devices = [] + + def get_pci_device(self, device_path: str) -> str: + try: + link = os.path.join(device_path, "device") + if os.path.islink(link): + return os.path.basename(os.readlink(link)) + except OSError: + pass + return "UNKNOWN_PCI" + + def find_netdev_for_pci(self, target_pci: str) -> str: + for netdev in glob.glob("/sys/class/net/*"): + try: + link = os.path.join(netdev, "device") + if os.path.islink(link): + if os.path.basename(os.readlink(link)) == target_pci: + return os.path.basename(netdev) + except OSError: + pass + return "NO_NETDEV" + + def get_socket_ifname_value(self) -> str: + try: + cmd = "ip route show default | awk '{print $5}'" + out = subprocess.check_output(cmd, shell=True, text=True).strip() + if not out: + print ("\n WARNING: no default route interfaces found.") + + ifnames = list(dict.fromkeys(out.splitlines())) + return ifnames[0] + + except Exception as e: + return "NA" + + # ------------------------ + # vendor mapping from pci + # ------------------------ + def rdma_vendor_from_pci(self, pci: str) -> str: + try: + pci_updated = pci.replace("0000:", "") + out = subprocess.check_output( + ["lspci", "-s", pci, "-nn"], + text=True + ).lower() + + if "pensando" in out: + return "AINIC" + elif "broadcom" in out: + return "BNXT" + elif "mellanox" in out: + return "MLNX" + else: + return "UNKNOWN" + except Exception: + pass + + return "UNKNOWN" + + # ------------------------- + # ibv devinfo parsing. + # ------------------------- + def _ibv_devinfo(self, rdma: str) -> str: + """Run ibv_devinfo once per device.""" + try: + result = subprocess.run( + ["ibv_devinfo", "-d", rdma, "-v"], + capture_output=True, + text=True, + timeout=10, + ) + if result.returncode == 0: + return result.stdout + except Exception: + pass + return "" + + def get_firmware_version(self, output: str) -> str: + """Extract fw_ver from ibv_devinfo output.""" + for line in output.splitlines(): + line = line.strip() + if line.startswith("fw_ver:"): + return line.split("fw_ver:", 1)[1].strip() + return "UNKNOWN" + + def get_gid_info(self, output: str) -> Tuple[str, str]: + """Extract IPv4-mapped RoCE GID.""" + for line in output.splitlines(): + if "::ffff:" in line and "GID[" in line: + idx = re.search(r"GID\[\s*(\d+)\]", line) + ip = re.search(r"(::ffff:[0-9.]+)", line) + if idx and ip: + return idx.group(1), ip.group(1) + return "-", "N/A" + + # ------------------------- + # Check RDMA devices + # ------------------------- + def scan_rdma_devices(self): + rdma_paths = sorted(glob.glob("/sys/class/infiniband/*")) + if not rdma_paths: + print ("No RDMA devices found.") + return + + for path in rdma_paths: + rdma = os.path.basename(path) + pci = self.get_pci_device(path) + netdev = self.find_netdev_for_pci(pci) + + ibv_out = self._ibv_devinfo(rdma) + firmware = self.get_firmware_version(ibv_out) + gid_index, gid_value = self.get_gid_info(ibv_out) + vendor = self.rdma_vendor_from_pci(pci) + + self.rdma_devices.append({ + "rdma": rdma, + "pci": pci, + "netdev": netdev, + "firmware": firmware, + "gid_index": gid_index, + "gid_value": gid_value, + "vendor": vendor, + }) + + # ------------------------- + # REPORTING + # ------------------------- + def print_table(self): + print ("\n RDMA CLUSTER MAPPING") + print ("=" * 100) + print (f"{'RDMA':<8} | {'PCI':<12} | {'NETDEV':<10} | {'FIRMWARE':<15} | {'GID_IDX':<7} | {'GID':<20} | VENDOR ") + print ("-" * 100) + + for d in self.rdma_devices: + print ( + f"{d['rdma']:<8} | {d['pci']:<12} | {d['netdev']:<10} | " + f"{d['firmware']:<15} | {d['gid_index']:<7} | {d['gid_value']:<20} | {d['vendor']}" + ) + + def print_detailed_info(self): + print ("\n DETAILED RDMA DEVICE ANALYSIS") + print ("=" * 80) + + for d in self.rdma_devices: + print (f"\nDevice: {d['rdma']}") + print (f" PCI: {d['pci']}") + print (f" NETDEV: {d['netdev']}") + print (f" Firmware: {d['firmware']}") + print (f" GID Index: {d['gid_index']}") + print (f" GID: {d['gid_value']}") + #print (f" Status: {self._analyze_firmware_version(d['firmware'])}") + print (f" VENDOR: {d['vendor']}") + + def _analyze_firmware_version(self, fw: str) -> str: + """ + TODO: analyze the firmware version and recommend the best based on NIC. + """ + + def generate_firmware_report(self): + print ("\n FIRMWARE VERSION REPORT") + print ("=" * 60) + + fw_map = {} + for d in self.rdma_devices: + fw_map.setdefault(d["firmware"], []).append(d["rdma"]) + + for fw, devs in fw_map.items(): + print (f"\nFirmware: {fw}") + #print (f"Status: {self._analyze_firmware_version(fw)}") + print (f"Devices: {', '.join(devs)}") + + if len(fw_map) > 1: + print ("\n Multiple firmware versions detected — standardization recommended") + + def _find_lib(self, patterns): + """ Find first matching library for given glob patterns.""" + for base in LIB_SEARCH_PATHS: + for pat in patterns: + matches = glob.glob(os.path.join(base, "**", pat), recursive=True) + if matches: + return matches[0] + + return None + + def _find_all_libs(self, patterns): + """Find all matching libraries.""" + found = [] + for base in LIB_SEARCH_PATHS: + for pat in patterns: + found.extend(glob.glob(os.path.join(base, "**", pat), recursive=True)) + return sorted(set(found)) + + def _docker_cmd_bnxt(self): + bnxt_rdma = self._find_lib(["libbnxt_re-rdmav*.so"]) + rdmacm = self._find_lib(["librdmacm.so.1"]) + ibverbs = self._find_lib(["libibverbs.so.1"]) + libnl3 = self._find_lib(["libnl-3.so.200"]) + libnl3_router = self._find_lib(["libnl-route-3.so.200"]) + + if not bnxt_rdma: + print ("Missing libbnxt_re-rdma*.so files \n") + if not rdmacm: + print ("Missing librdmacm.so* files \n") + if not libnl3: + print ("Missing libnl* files \n") + + print ("Libraries detected on host device:") + print (f"{bnxt_rdma:>5}") + print (f"{rdmacm:>5}") + print (f"{ibverbs:>5}") + print (f"{libnl3:>5}") + print (f"{libnl3_router:>5}") + + cmd_string = f""" + docker run --rm -it \\ + --device /dev/dri \\ + --device /dev/infiniband \\ + --device /dev/kfd \\ + --network host \\ + --ipc host \\ + --privileged \\ + --ulimit memlock=-1:-1 \\ + --group-add video \\ + --cap-add SYS_PTRACE \\ + --security-opt seccomp=unconfined \\ + --shm-size 64G \\ + -v /sys:/sys \\ + -v $HOME/.ssh:/root/.ssh \\ + -v $HOME:$HOME \\ + -v /dev/infiniband:/dev/infiniband \\ + -v /sys/class/infiniband:/sys/class/infiniband:ro \\ + -v /sys/class/net:/sys/class/net:ro \\ + -v /sys/bus/pci:/sys/bus/pci:ro \\ + -v /etc/libibverbs.d:/etc/libibverbs.d:ro \\ + -v /etc/rdma:/etc/rdma:ro \\ + """ + if bnxt_rdma: + cmd_string = cmd_string + (" " * 8) + f"-v {bnxt_rdma}:{bnxt_rdma}:ro \\\n" + if rdmacm: + cmd_string = cmd_string + (" " * 16) + f"-v {rdmacm}:{rdmacm}:ro \\\n" + if ibverbs: + cmd_string = cmd_string + (" " * 16) + f"-v {ibverbs}:{ibverbs}:ro \\\n" + if libnl3: + cmd_string = cmd_string + (" " * 16) + f"-v {libnl3}:{libnl3}:ro \\\n" + if libnl3_router: + cmd_string = cmd_string + (" " * 16) + f"-v {libnl3_router}:{libnl3_router}:ro \\\n" + + cmd_string = cmd_string + (" " * 16) + " \n" + + return cmd_string.strip() + + def _docker_cmd_mlnx(self): + cmd_string = f""" + docker run --rm -it \\ + --device /dev/dri \\ + --device /dev/infiniband \\ + --device /dev/kfd \\ + --network host \\ + --ipc host \\ + --privileged \\ + --ulimit memlock=-1:-1 \\ + --group-add video \\ + --cap-add SYS_PTRACE \\ + --security-opt seccomp=unconfined \\ + --shm-size 64G \\ + -v /sys:/sys \\ + -v $HOME/.ssh:/root/.ssh \\ + -v $HOME:$HOME \\ + -v /dev/infiniband:/dev/infiniband \\ + -v /sys/class/infiniband:/sys/class/infiniband:ro \\ + -v /sys/class/net:/sys/class/net:ro \\ + -v /sys/bus/pci:/sys/bus/pci:ro \\ + """ + cmd_string = cmd_string + (" " * 8) + " \n" + return cmd_string.strip() + + def _docker_cmd_ionic(self): + ionic_rdma = self._find_lib(["libionic-rdmav*.so"]) + ionic_so = self._find_all_libs(["libionic.so*"]) + ionic_driver = self._find_lib(["ionic.driver"]) + + if not ionic_rdma: + print ("Missing libionic-rdma*.so files \n") + if not ionic_so: + print ("Missing libionic.so files \n") + if not ionic_driver: + print ("Missing ionic.driver file \n") + + print ("\n") + print ("Libraries detected on host device:") + print (f"{ionic_rdma:>5}") + print (f"{ionic_driver:>5}") + for so_file in ionic_so: + print (f"{so_file:>5}") + + cmd_string = f""" + docker run --rm -it \\ + --device /dev/dri \\ + --device /dev/infiniband \\ + --device /dev/kfd \\ + --network host \\ + --ipc host \\ + --privileged \\ + --ulimit memlock=-1:-1 \\ + --group-add video \\ + --cap-add SYS_PTRACE \\ + --security-opt seccomp=unconfined \\ + --shm-size 64G \\ + -v /sys:/sys \\ + -v $HOME/.ssh:/root/.ssh \\ + -v $HOME:$HOME \\ + -v /dev/infiniband:/dev/infiniband \\ + -v /sys/class/infiniband:/sys/class/infiniband:ro \\ + -v /sys/class/net:/sys/class/net:ro \\ + -v /sys/bus/pci:/sys/bus/pci:ro \\ + """ + + if ionic_rdma: + cmd_string = cmd_string + (" " * 8) + f"-v {ionic_rdma}:{ionic_rdma}:ro \\\n" + if ionic_so: + for j in range(len(ionic_so)): + so_file = ionic_so[j] + cmd_string = cmd_string + (" " * 16) + f"-v {so_file}:{so_file}:ro \\\n" + if ionic_driver: + cmd_string = cmd_string + (" " * 16) + f"-v {ionic_driver}:{ionic_driver}:ro \\\n" + + cmd_string = cmd_string + (" " * 16) + " \n" + + return cmd_string.strip() + + def generate_docker_launch_command(self): + vendors = {d['vendor'] for d in self.rdma_devices} + + print ("\n RECOMMENDED DOCKER LAUNCH COMMAND") + print ("=" * 80) + + if len(vendors) > 1: + print ("\n WARNING: Multiple RDMA vendors detected.") + + print ("\n") + print ("Vendors detected: {}".format(vendors)) + docker_cmd = "" + if "AINIC" in vendors: + docker_cmd = self._docker_cmd_ionic() + print ("\n") + print ("Docker launch command:") + print (docker_cmd) + elif "BNXT" in vendors: + docker_cmd = self._docker_cmd_bnxt() + print ("\n") + print ("Docker launch command:") + print (docker_cmd) + elif "MLNX" in vendors: + docker_cmd = self._docker_cmd_mlnx() + print ("\n") + print ("Docker launch command:") + print (docker_cmd) + else: + print ("\n WARNING: this vendor docker cannot be generated, please verify manually.") + + return docker_cmd + + def _get_nccl_env_variables(self): + nccl_env_variables = [] + + # IB NCCL CPU affinity. + nccl_env_variables.append("export NCCL_IGNORE_CPU_AFFINITY=1") + # IB GID + gid_indexes = {d['gid_index'] for d in self.rdma_devices} + if (len(gid_indexes) > 1): + print (" \n WARNING: multiple GID indeces detected, please check detailed report for mapping the env variables.") + nccl_env_variables.append(f"export NCCL_IB_GID_INDEX={max(list(gid_indexes))}") + + # IB HCA + #rdma_devices = ",".join([d['rdma'] for d in self.rdma_devices]) + rdma_devices="" + firmware_version = max([d['firmware'] for d in self.rdma_devices]) + for d in self.rdma_devices: + device = d['rdma'] if d['gid_index'].isnumeric() and d['firmware'] == firmware_version else "" + if device: + rdma_devices = rdma_devices + device + "," + nccl_env_variables.append(f"export NCCL_IB_HCA={rdma_devices.rstrip(',')}") + + # NCCL/GLOO socket if name if framework requires. + socket_ifname = self.get_socket_ifname_value() + nccl_env_variables.append(f"export NCCL_SOCKET_IFNAME={socket_ifname}") + nccl_env_variables.append(f"export GLOO_SOCKET_IFNAME={socket_ifname}") + + return nccl_env_variables + + def generate_framework_env_variables(self): + print ("\n RECOMMENDED ENV VARIABLES") + print ("=" * 80) + + env_variables = [] + print ("Note: Please cross check before exporting in your scripts") + nccl_env_variables = self._get_nccl_env_variables() + print ("\n NCCL Env Variables:") + for var in nccl_env_variables: + print (var) + + print ("\n rocSHMEM Env variables:") + rocshmem_env =[] + rocshmem_env.append("export ROCSHMEM_HEAP_SIZE=7524589824") + rocshmem_env.append("export ROCSHMEM_MAX_NUM_CONTEXTS=256") + for var in rocshmem_env: + print (var) + + env_variables = nccl_env_variables + rocshmem_env + return env_variables + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--html", help="Generate HTML report", action="store_true") + args = parser.parse_args() + + mapper = RDMAClusterMapper() + mapper.scan_rdma_devices() + + if not mapper.rdma_devices: + return + + mapper.print_table() + mapper.print_detailed_info() + mapper.generate_firmware_report() + mapper.generate_docker_launch_command() + mapper.generate_framework_env_variables() + +if __name__ == "__main__": + main()