docs(ib): switch perftest example from mpirun to srun

Signed-off-by: chang-ning <spiderpower02@gmail.com>

Author: crazyguitar

examples/ib/README.md

@@ -6,22 +6,18 @@ suite of micro-benchmarks for RDMA verbs — `ib_write_bw`, `ib_read_bw`,
 program: one rank acts as a server, the other connects as a client.
 
 `examples/ib/ib.sh` is a thin wrapper that launches both ranks across
-two hosts with `mpirun`, so you can benchmark inter-node RDMA with a
+two hosts with `srun`, so you can benchmark inter-node RDMA with a
 single command and observe the traffic in `rdmatop` from another shell.
 
 ## Prerequisites
 
-- `perftest` and an MPI launcher installed on both hosts:
+- A working Slurm cluster with both hosts as compute nodes, and
+  `perftest` installed on both:
   ```bash
-  sudo apt install -y perftest openmpi-bin
-  ```
-- Passwordless SSH from the launcher host to both hosts (mpirun shells
-  in to spawn each rank):
-  ```bash
-  ssh-keygen -t ed25519 -N "" -f ~/.ssh/id_ed25519
-  ssh-copy-id user@host1
-  ssh-copy-id user@host2
+  sudo apt install -y perftest
   ```
+- `srun` reachable from where you launch the script (any login or
+  compute node in the cluster).
 - An RDMA device that is ACTIVE on both ends (`rdma link`) and a netdev
   with an IPv4 address reachable between the two hosts.
 
@@ -31,6 +27,7 @@ single command and observe the traffic in `rdmatop` from another shell.
 ./examples/ib/ib.sh <host1> <host2> [perftest_binary]
 #   host1 -> rank 0 -> server
 #   host2 -> rank 1 -> client (dials host1)
+#   host1/host2 must be Slurm NodeNames (see `sinfo -N -h -o %N`)
 ```
 
 `perftest_binary` defaults to `ib_write_bw`. Any perftest tool that
@@ -40,22 +37,25 @@ follows the `<tool> [opts] [server_ip]` convention works.
 
 ```bash
 # RDMA write bandwidth, 30s, 1 QP, mlx5_1
-./examples/ib/ib.sh amd0 amd1
+./examples/ib/ib.sh host1 host2
 
 # RDMA read bandwidth
-./examples/ib/ib.sh amd0 amd1 ib_read_bw
+./examples/ib/ib.sh host1 host2 ib_read_bw
 
 # Send latency
-./examples/ib/ib.sh amd0 amd1 ib_send_lat
+./examples/ib/ib.sh host1 host2 ib_send_lat
 
 # Longer run with 4 queue pairs (closer to line rate)
-IB_DURATION=30 IB_QPS=4 ./examples/ib/ib.sh amd0 amd1
+IB_DURATION=30 IB_QPS=4 ./examples/ib/ib.sh host1 host2
 
 # Sweep all message sizes
-IB_EXTRA="-a" ./examples/ib/ib.sh amd0 amd1
+IB_EXTRA="-a" ./examples/ib/ib.sh host1 host2
 
 # Use a different RDMA device + matching netdev
-IB_DEV=mlx5_0 IB_NETDEV=enp49s0f0np0 ./examples/ib/ib.sh amd0 amd1
+IB_DEV=mlx5_0 IB_NETDEV=enp49s0f0np0 ./examples/ib/ib.sh host1 host2
+
+# Pin to a specific Slurm partition
+SRUN_EXTRA="-p all" ./examples/ib/ib.sh host1 host2
 ```
 
 ## Environment Variables
@@ -67,13 +67,13 @@ IB_DEV=mlx5_0 IB_NETDEV=enp49s0f0np0 ./examples/ib/ib.sh amd0 amd1
 | `IB_DURATION` | `30`             | Seconds, passed as `-D`                                |
 | `IB_QPS`      | `1`              | Queue pairs, passed as `-q`                            |
 | `IB_EXTRA`    | (empty)          | Extra args appended verbatim to the perftest binary    |
-| `MPI_EXTRA`   | (empty)          | Extra args passed to `mpirun`                          |
+| `SRUN_EXTRA`  | (empty)          | Extra args passed to `srun` (e.g. `-p <partition>`)    |
 
 ## Related Links
 
 - [perftest](https://github.com/linux-rdma/perftest)
   — upstream RDMA verbs benchmarks
-- [Open MPI](https://www.open-mpi.org/)
-  — `mpirun` launcher
+- [Slurm](https://slurm.schedmd.com/)
+  — `srun` launcher
 - [rdma-core](https://github.com/linux-rdma/rdma-core)
   — userspace RDMA libraries and drivers

examples/ib/ib.sh

@@ -1,11 +1,12 @@
 #!/usr/bin/env bash
-# ib.sh - launch a perftest benchmark between two hosts via mpirun.
+# ib.sh - launch a perftest benchmark between two hosts via Slurm (srun).
 #
 # Usage:
 #   ./ib.sh <host1> <host2> [perftest_binary]
 #     host1 -> rank 0 -> perftest server
 #     host2 -> rank 1 -> perftest client (connects to host1)
 #     perftest_binary defaults to ib_write_bw
+#     host1/host2 must be Slurm NodeNames (e.g. `sinfo -N -h -o %N`)
 #
 # Env overrides:
 #   IB_DEV       RDMA device                   (default: mlx5_1)
@@ -14,13 +15,7 @@
 #   IB_DURATION  perftest -D seconds           (default: 30)
 #   IB_QPS       perftest -q queue pairs       (default: 1)
 #   IB_EXTRA     extra args appended verbatim  (default: empty)
-#   MPI_EXTRA    extra args passed to mpirun   (default: empty)
-#
-# Examples:
-#   ./ib.sh amd0 amd1
-#   ./ib.sh amd0 amd1 ib_read_bw
-#   IB_DURATION=30 IB_QPS=4 ./ib.sh amd0 amd1 ib_write_bw
-#   IB_DEV=mlx5_0 IB_NETDEV=enp49s0f0np0 ./ib.sh amd0 amd1
+#   SRUN_EXTRA   extra args passed to srun     (default: empty)
 
 set -euo pipefail
 
@@ -29,28 +24,10 @@ IB_NETDEV=${IB_NETDEV:-enp49s0f1np1}
 IB_DURATION=${IB_DURATION:-30}
 IB_QPS=${IB_QPS:-1}
 IB_EXTRA=${IB_EXTRA:-}
-MPI_EXTRA=${MPI_EXTRA:-}
-
-# Worker mode: mpirun re-invokes this script on each host with a rank set.
-rank=${OMPI_COMM_WORLD_RANK:-${PMI_RANK:-${PMIX_RANK:-}}}
-if [[ -n "${rank}" ]]; then
-  test_bin=$1
-  peer_ip=$2
-  common=(-d "$IB_DEV" -F -D "$IB_DURATION" -q "$IB_QPS" $IB_EXTRA)
-  if [[ "$rank" == "0" ]]; then
-    echo "[$(hostname)] rank 0: $test_bin ${common[*]}  (server)"
-    exec "$test_bin" "${common[@]}"
-  else
-    # Tiny delay so the server's TCP listener is up before we dial.
-    sleep 1
-    echo "[$(hostname)] rank 1: $test_bin ${common[*]} $peer_ip  (client)"
-    exec "$test_bin" "${common[@]}" "$peer_ip"
-  fi
-fi
+SRUN_EXTRA=${SRUN_EXTRA:-}
 
-# Launcher mode
 if [[ $# -lt 2 ]]; then
-  sed -n '2,25p' "$0"
+  sed -n '2,19p' "$0"
   exit 2
 fi
 
@@ -59,26 +36,49 @@ host2=$2
 test_bin=${3:-ib_write_bw}
 
 # Resolve host1's IP on the RDMA-adjacent netdev so host2 can connect to it.
-peer_ip=$(ssh -o BatchMode=yes "$host1" \
-  "ip -br -4 addr show $IB_NETDEV 2>/dev/null | awk '{print \$3}' | cut -d/ -f1")
+peer_ip=$(srun --nodelist="$host1" --nodes=1 --ntasks=1 \
+  bash -c "ip -br -4 addr show $IB_NETDEV 2>/dev/null | awk '{print \$3}' | cut -d/ -f1" \
+  2>/dev/null | tr -d '[:space:]')
 if [[ -z "$peer_ip" ]]; then
   echo "ERROR: could not resolve IPv4 on $IB_NETDEV at $host1" >&2
   echo "       set IB_NETDEV to a netdev that has an IP on both hosts." >&2
   exit 1
 fi
 
-script=$(readlink -f "$0")
-echo "=== ib.sh ==="
+echo "=== ib.sh (slurm) ==="
 echo "  test     : $test_bin"
 echo "  hosts    : $host1 (server) -> $host2 (client)"
 echo "  device   : $IB_DEV  via netdev $IB_NETDEV"
 echo "  peer ip  : $peer_ip"
 echo "  duration : ${IB_DURATION}s   qps=${IB_QPS}"
-[[ -n "$IB_EXTRA"  ]] && echo "  extra    : $IB_EXTRA"
-[[ -n "$MPI_EXTRA" ]] && echo "  mpi extra: $MPI_EXTRA"
+[[ -n "$IB_EXTRA"   ]] && echo "  extra     : $IB_EXTRA"
+[[ -n "$SRUN_EXTRA" ]] && echo "  srun extra: $SRUN_EXTRA"
 echo
 
-exec mpirun --host "$host1,$host2" -np 2 \
-  --tag-output \
-  $MPI_EXTRA \
-  "$script" "$test_bin" "$peer_ip"
+# Pin rank 0 -> host1, rank 1 -> host2 via SLURM_HOSTFILE + arbitrary distribution.
+hostfile=$(mktemp)
+trap 'rm -f "$hostfile"' EXIT
+printf '%s\n%s\n' "$host1" "$host2" > "$hostfile"
+export SLURM_HOSTFILE="$hostfile"
+
+export PERF_BIN="$test_bin"
+export PEER_IP="$peer_ip"
+export IB_DEV IB_DURATION IB_QPS IB_EXTRA
+
+exec srun \
+  --ntasks=2 \
+  --distribution=arbitrary \
+  --label \
+  $SRUN_EXTRA \
+  bash -c '
+    common=(-d "$IB_DEV" -F -D "$IB_DURATION" -q "$IB_QPS" $IB_EXTRA)
+    if [[ "$SLURM_PROCID" == "0" ]]; then
+      echo "[$(hostname)] rank 0 server: $PERF_BIN ${common[*]}"
+      exec "$PERF_BIN" "${common[@]}"
+    else
+      # Tiny delay so the server'\''s listener is up before we dial.
+      sleep 1
+      echo "[$(hostname)] rank 1 client: $PERF_BIN ${common[*]} $PEER_IP"
+      exec "$PERF_BIN" "${common[@]}" "$PEER_IP"
+    fi
+  '