Skip to content

Commit e45d73a

Browse files
crazyguitarcrazyguitar
authored
feat: add ucx example (#14)
Signed-off-by: chang-ning <spiderpower02@gmail.com>
1 parent fd82e6a commit e45d73a

5 files changed

Lines changed: 255 additions & 1 deletion

File tree

‎README.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,8 @@ rdmatop
5454
Use `rdmatop` to monitor RDMA traffic while running GPU
5555
communication benchmarks:
5656

57-
- [IB Perftest](examples/ib/) — two-node `ib_write_bw` benchmark via slurm
57+
- [IB Perftest](examples/ib/) — two-node `ib_write_bw` benchmark via Slurm `srun`
58+
- [UCX Perftest](examples/ucx/) — two-node `ucx_perftest` bandwidth / latency via Slurm `srun`
5859
- [NCCL](examples/nccl/) — collective communication
5960
- [NIXL](examples/nixl/) — point-to-point KV cache transfer
6061
- [NVSHMEM](examples/nvshmem/) — one-sided GPU communication

‎examples/ucx/Dockerfile‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
ARG BASE_IMAGE=ubuntu:22.04
2+
FROM ${BASE_IMAGE}
3+
4+
ENV DEBIAN_FRONTEND=noninteractive
5+
6+
RUN apt-get update \
7+
&& apt-get install -y --no-install-recommends \
8+
ca-certificates \
9+
perftest ucx-utils \
10+
libibverbs1 librdmacm1 ibverbs-utils ibverbs-providers \
11+
rdma-core \
12+
libnuma1 \
13+
iproute2 iputils-ping \
14+
&& rm -rf /var/lib/apt/lists/*
15+
16+
CMD ["/bin/bash"]

‎examples/ucx/Makefile‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
IMAGE_NAME ?= ucx
2+
IMAGE_TAG ?= latest
3+
BASE_IMAGE ?= ubuntu:22.04
4+
OUTPUT_DIR ?= $(PWD)
5+
6+
.PHONY: all docker sqsh clean
7+
8+
all: sqsh
9+
10+
docker:
11+
docker build --build-arg BASE_IMAGE=$(BASE_IMAGE) -t $(IMAGE_NAME):$(IMAGE_TAG) .
12+
13+
sqsh: docker
14+
rm -f $(OUTPUT_DIR)/$(IMAGE_NAME)+$(IMAGE_TAG).sqsh
15+
enroot import -o $(OUTPUT_DIR)/$(IMAGE_NAME)+$(IMAGE_TAG).sqsh dockerd://$(IMAGE_NAME):$(IMAGE_TAG)
16+
17+
clean:
18+
rm -f $(OUTPUT_DIR)/$(IMAGE_NAME)+$(IMAGE_TAG).sqsh
19+
docker rmi -f $(IMAGE_NAME):$(IMAGE_TAG) 2>/dev/null || true

‎examples/ucx/README.md‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
# UCX Perftest
2+
3+
[UCX](https://openucx.org/) (Unified Communication X) provides a
4+
high-performance messaging API that abstracts over verbs, TCP, shared
5+
memory, CUDA-IPC, ROCm-IPC, and other transports. `ucx_perftest` is its
6+
micro-benchmark — bandwidth and latency tests for tag matching, active
7+
messages, RMA put/get, streams, atomics, and more. Like `perftest`, each
8+
test is a two-process program: one rank is the server, the other dials in.
9+
10+
`examples/ucx/ucx.sh` is a thin wrapper that launches both ranks across
11+
two hosts with `srun`, so you can benchmark inter-node UCX bandwidth or
12+
latency in one command and watch the traffic in `rdmatop` from another
13+
shell.
14+
15+
## Prerequisites
16+
17+
- A working Slurm cluster with both hosts as compute nodes, and
18+
`ucx_perftest` installed on both:
19+
```bash
20+
sudo apt install -y ucx-utils
21+
```
22+
- `srun` reachable from where you launch the script (any login or
23+
compute node in the cluster).
24+
- An RDMA device that is ACTIVE on both ends (`rdma link`, `ucx_info -d`)
25+
and a netdev with an IPv4 address reachable between the two hosts —
26+
UCX uses it for the wire-up handshake.
27+
28+
## Usage
29+
30+
```bash
31+
./examples/ucx/ucx.sh <host1> <host2> [test_name]
32+
# host1 -> rank 0 -> server
33+
# host2 -> rank 1 -> client (dials host1)
34+
# test_name defaults to tag_bw; run `ucx_perftest -h` for the full list
35+
# host1/host2 must be Slurm NodeNames (see `sinfo -N -h -o %N`)
36+
```
37+
38+
`ucx_perftest` is **iteration-based** (`-n N`), not duration-based like
39+
`ib_write_bw`'s `-D 30`. The default `PERF_ITER=5000000` is sized for
40+
roughly 30 s at 64 KB messages on a 100 GbE link.
41+
42+
## Examples
43+
44+
```bash
45+
# UCP tag matching bandwidth (default test), ~30 s
46+
./examples/ucx/ucx.sh host1 host2
47+
48+
# UCP tag matching latency
49+
./examples/ucx/ucx.sh host1 host2 tag_lat
50+
51+
# UCP RMA put bandwidth, 1 MiB messages, 1 M iterations (~10 s)
52+
PERF_SIZE=1048576 PERF_ITER=1000000 \
53+
./examples/ucx/ucx.sh host1 host2 ucp_put_bw
54+
55+
# UCP RMA get bandwidth
56+
./examples/ucx/ucx.sh host1 host2 ucp_get
57+
58+
# UCP active message bandwidth at 1 MiB messages
59+
PERF_SIZE=1048576 ./examples/ucx/ucx.sh host1 host2 ucp_am_bw
60+
61+
# Sweep small-message tag latency (more iters because each one is cheap)
62+
PERF_SIZE=8 PERF_ITER=50000000 \
63+
./examples/ucx/ucx.sh host1 host2 tag_lat
64+
65+
# Pin to a different RDMA device + matching netdev
66+
UCX_DEV=mlx5_0:1 UCX_NETDEV=enp49s0f0np0 \
67+
./examples/ucx/ucx.sh host1 host2
68+
69+
# Pin to a Slurm partition / extra srun args
70+
SRUN_EXTRA="-p all" ./examples/ucx/ucx.sh host1 host2
71+
72+
# Run inside the enroot container image (see Container Image below)
73+
SRUN_EXTRA="--container-image=$HOME/changning/rdmatop/examples/ucx/ucx+latest.sqsh" \
74+
./examples/ucx/ucx.sh host1 host2
75+
```
76+
77+
## Container Image
78+
79+
`examples/ucx/Dockerfile` is a slim 87 MB Ubuntu 22.04 image with
80+
`perftest` + `ucx-utils` and the verbs runtime libs.
81+
82+
```bash
83+
# Build the image and import to ucx+latest.sqsh next to the Makefile.
84+
# Run on each node (no shared FS):
85+
cd examples/ucx && make
86+
87+
# Run inside the container:
88+
SRUN_EXTRA="--container-image=$PWD/ucx+latest.sqsh" \
89+
./ucx.sh host1 host2 tag_bw
90+
91+
# For HIP / rocm-smi inside, rebuild with the ROCm SDK base:
92+
BASE_IMAGE=rocm/dev-ubuntu-22.04:6.2.4 make
93+
```
94+
95+
## Environment Variables
96+
97+
| Variable | Default | Purpose |
98+
|-----------------------------|--------------------------------------|---------------------------------------------------------------------------------------------------------------------------|
99+
| `UCX_DEV` | `mlx5_1:1` | Exported as `UCX_NET_DEVICES` |
100+
| `UCX_NETDEV` | `enp49s0f1np1` | Netdev paired with `UCX_DEV`; used to look up host1's IP |
101+
| `PERF_SIZE` | `65536` | Message size, passed as `-s` |
102+
| `PERF_ITER` | `5000000` | Iterations, passed as `-n` (~30 s at 64 KB / 100 GbE) |
103+
| `PERF_EXTRA` | (empty) | Extra args appended verbatim to `ucx_perftest` |
104+
| `SRUN_EXTRA` | (empty) | Extra args passed to `srun` (e.g. `--gres=gpu:1`, `--container-image=...`) |
105+
| `UCX_TLS` | `rc_verbs,ud_verbs,self,sm,tcp` | Transports UCX is allowed to use. |
106+
| `UCX_SOCKADDR_TLS_PRIORITY` | `tcp` | Wire-up handshake transport. TCP stays robust when multiple HCAs are present. |
107+
| `UCX_WARN_UNUSED_ENV_VARS` | `n` | Silences UCX's warning about unrecognised `UCX_*` vars when you set additional ones in the launching shell. |
108+
109+
## Related Links
110+
111+
- [UCX](https://openucx.org/) — Unified Communication X framework
112+
- [`ucx_perftest` source](https://github.com/openucx/ucx/tree/master/src/tools/perf)
113+
- [Slurm](https://slurm.schedmd.com/) — `srun` launcher
114+
- [enroot](https://github.com/NVIDIA/enroot) + [pyxis](https://github.com/NVIDIA/pyxis) — used for the `--container-image` path

‎examples/ucx/ucx.sh‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
#!/usr/bin/env bash
2+
# ucx.sh - launch a ucx_perftest benchmark between two hosts via Slurm (srun).
3+
#
4+
# Usage:
5+
# ./ucx.sh <host1> <host2> [test_name]
6+
# host1 -> rank 0 -> ucx_perftest server
7+
# host2 -> rank 1 -> ucx_perftest client (connects to host1)
8+
# test_name defaults to tag_bw (see `ucx_perftest -h` for the full list)
9+
# host1/host2 must be Slurm NodeNames (e.g. `sinfo -N -h -o %N`)
10+
#
11+
# Env overrides:
12+
# UCX_DEV UCX_NET_DEVICES value (default: mlx5_1:1)
13+
# UCX_NETDEV netdev paired with UCX_DEV (default: enp49s0f1np1)
14+
# used to look up host1's IP so host2 can dial it
15+
# PERF_SIZE ucx_perftest -s message size (default: 65536)
16+
# PERF_ITER ucx_perftest -n iterations (default: 5000000)
17+
# ucx_perftest is iteration-based (no -D duration flag);
18+
# 5M iters @ 64KB ~= 30s on a 100GbE RoCE link.
19+
# PERF_EXTRA extra args appended verbatim (default: empty)
20+
# SRUN_EXTRA extra args passed to srun (default: empty)
21+
#
22+
# UCX_TLS transports (default: rc_verbs,ud_verbs,self,sm,tcp)
23+
# We skip mlx5 DV accelerated transports by default
24+
# because they were crashing with REM_ACCESS_ERR on
25+
# this RoCE setup. Override to "all" or
26+
# "rc_mlx5,..." to re-enable.
27+
# UCX_SOCKADDR_TLS_PRIORITY wireup transport (default: tcp). Forces the
28+
# control handshake to go over TCP instead of
29+
# rdma_cm, which is more robust when multiple
30+
# HCAs are present.
31+
32+
set -euo pipefail
33+
34+
UCX_DEV=${UCX_DEV:-mlx5_1:1}
35+
UCX_NETDEV=${UCX_NETDEV:-enp49s0f1np1}
36+
PERF_SIZE=${PERF_SIZE:-65536}
37+
PERF_ITER=${PERF_ITER:-5000000}
38+
PERF_EXTRA=${PERF_EXTRA:-}
39+
SRUN_EXTRA=${SRUN_EXTRA:-}
40+
41+
# UCX defaults — overridable by exporting before calling this script.
42+
: ${UCX_TLS:=rc_verbs,ud_verbs,self,sm,tcp}
43+
: ${UCX_SOCKADDR_TLS_PRIORITY:=tcp}
44+
: ${UCX_WARN_UNUSED_ENV_VARS:=n}
45+
export UCX_TLS UCX_SOCKADDR_TLS_PRIORITY UCX_WARN_UNUSED_ENV_VARS
46+
47+
if [[ $# -lt 2 ]]; then
48+
sed -n '2,21p' "$0"
49+
exit 2
50+
fi
51+
52+
host1=$1
53+
host2=$2
54+
test_name=${3:-tag_bw}
55+
56+
# Resolve host1's IP on the RDMA-adjacent netdev so host2 can connect to it.
57+
peer_ip=$(srun --nodelist="$host1" --nodes=1 --ntasks=1 \
58+
bash -c "ip -br -4 addr show $UCX_NETDEV 2>/dev/null | awk '{print \$3}' | cut -d/ -f1" \
59+
2>/dev/null | tr -d '[:space:]')
60+
if [[ -z "$peer_ip" ]]; then
61+
echo "ERROR: could not resolve IPv4 on $UCX_NETDEV at $host1" >&2
62+
echo " set UCX_NETDEV to a netdev that has an IP on both hosts." >&2
63+
exit 1
64+
fi
65+
66+
echo "=== ucx.sh (slurm) ==="
67+
echo " test : $test_name"
68+
echo " hosts : $host1 (server) -> $host2 (client)"
69+
echo " device : UCX_NET_DEVICES=$UCX_DEV via netdev $UCX_NETDEV"
70+
echo " peer ip : $peer_ip"
71+
echo " msg size : $PERF_SIZE iter=$PERF_ITER"
72+
[[ -n "$PERF_EXTRA" ]] && echo " extra : $PERF_EXTRA"
73+
[[ -n "$SRUN_EXTRA" ]] && echo " srun extra: $SRUN_EXTRA"
74+
echo
75+
76+
# Pin rank 0 -> host1 (server), rank 1 -> host2 (client) via SLURM_HOSTFILE + arbitrary.
77+
hostfile=$(mktemp)
78+
trap 'rm -f "$hostfile"' EXIT
79+
printf '%s\n%s\n' "$host1" "$host2" > "$hostfile"
80+
export SLURM_HOSTFILE="$hostfile"
81+
82+
export PERF_TEST="$test_name"
83+
export PEER_IP="$peer_ip"
84+
export UCX_NET_DEVICES="$UCX_DEV"
85+
export PERF_SIZE PERF_ITER PERF_EXTRA
86+
87+
exec srun \
88+
--ntasks=2 \
89+
--distribution=arbitrary \
90+
--label \
91+
$SRUN_EXTRA \
92+
bash -c '
93+
common=(-t "$PERF_TEST" -s "$PERF_SIZE" -n "$PERF_ITER")
94+
common+=($PERF_EXTRA)
95+
if [[ "$SLURM_PROCID" == "0" ]]; then
96+
echo "[$(hostname)] rank 0 server: ucx_perftest ${common[*]}"
97+
exec ucx_perftest "${common[@]}"
98+
else
99+
# Tiny delay so the server'\''s listener is up before we dial.
100+
sleep 1
101+
echo "[$(hostname)] rank 1 client: ucx_perftest ${common[*]} $PEER_IP"
102+
exec ucx_perftest "${common[@]}" "$PEER_IP"
103+
fi
104+
'

0 commit comments

Comments
 (0)