connex-rs 0.1.7

The Connex data plane: Rust over UCX, peer-death events in milliseconds, same-identity rejoin

About

connex-rs moves pipeline-step data between Connex endpoints and tells them, within milliseconds, when a peer dies. It is a Rust crate over UCX with a Python binding (connex_rs) and is the default data plane of connex.pp in Connex on hosts with a usable RDMA device (CONNEX_PP_BACKEND=auto, or force it with CONNEX_PP_BACKEND=connex-rs; Mooncake Transfer Engine stays available with CONNEX_PP_BACKEND=mooncake).

What it provides:

• One-sided RMA on registered memory, host or GPU (GPUDirect RDMA): transfer_sync_read(url, dst, src, len) and the rest of the TransferEngine surface that connex.pp uses, so the endpoint's design (sender stages bytes in a registered ring, signals, receiver does one read, ack frees the slot) maps one to one. Plus one-sided writes, put_signal (write, then a message that arrives only after the data is in place), and plain messages.
• Identity-keyed peers. A peer URL is ip:port#model|stage|shard|replica. Connections are bound to the logical id, so a stage that restarts under the same id and address is the same peer to everyone else: it rejoins without anyone else restarting.
• Peer-failure events. Every UCX endpoint runs with peer error handling; a dead process is reported in a few ms (its kernel tears the connection down). A heartbeat on the progress thread (off unless heartbeat_ms is set; the timeout defaults to 3x the interval) and per-op timeouts report hung or unreachable peers; those are suspicions that connex.pp confirms with a probe.
• A dedicated progress thread. It owns the UCX worker; callers submit commands and wait on completions, and the Python binding releases the GIL while waiting. Heartbeats and failure handling keep running while the inference engine holds the GIL.

Performance vs Mooncake Transfer Engine

Same cluster and deployment: Llama-3-8B split into 4 pipeline stages on A40 GPUs, one stage per GPU, InfiniBand between nodes, 2 pipeline replicas with 4 steps in flight per replica (CONNEX_PP_INFLIGHT=4); only CONNEX_PP_BACKEND differs. The steady-state and closed-loop Mooncake numbers were re-measured in the same session as connex-rs; the Mooncake transfer and handover numbers come from the earlier validation of the same setup. Full tables, setup and reproduction commands: docs/results.md in the Connex repository.

Raw transfers, GPU to GPU across nodes, one-sided read of registered GPU memory (GPUDirect RDMA):

Serving, steady state, trace replay, 100% goodput for both:

Failure handling, where the two differ:

Where connex-rs is not better: closed-loop throughput is 3% lower (not yet profiled), its TCP path (no RDMA) is several times slower than Mooncake's, which is why auto keeps Mooncake on hosts without RDMA, and individual requests can still stall for up to about 1 s (worst request 1.03 s at 1 step in flight, 0.92 s at 4) even though the per-run medians are far lower.

Data-plane microbenchmarks

cxrma (src/bin/cxrma.rs) measures the engine itself: blocking get/put latency, streaming bandwidth, put_chain latency, time from connect to the first read, and idle CPU. Two hosts with A40 GPUs and one ConnectX EDR port each (100 Gb/s, 12.5 GB/s), host memory, UCX 1.20, p50 over 3000 operations, median of 3 runs. "Before" is 0.1.5, "after" this release.

Large transfers already run at the link rate (95% of 12.5 GB/s, with or without several ops in flight), so the engine does not split them. GPU memory is limited by PCIe, not by the engine: on one host through the HCA, GPU to GPU reads reach 9.2 GB/s with connex-rs and 9.3 GB/s with ucx_perftest between the two GPUs that sit behind the NIC's PCIe host bridge, and ucx_perftest reaches only 7.7 GB/s between GPUs attached to the other CPU socket. Over TCP (UCX_TLS=tcp, loopback) a 1 MB read went from 0.9 to 3.3 GB/s and a 4 KB put from 46 to 41 µs.

In a full pipeline (vLLM PP=2 across two hosts, this build), the first put_chain on a freshly opened push channel completes 0.4-17 ms after the channel opens, and stage 0 switches to a spare replica 1-3 ms after connex-rs reports the active one failed.

Reproduce: cargo build --release, then target/release/cxrma local --ip <ib ip> on one host, or cxrma owner on one host and cxrma reader --url <url> --addr <addr> on the other (--gpu-owner D --gpu-reader D for GPU memory, --verify-chain checks that a consumer never sees a chain's flag before its data).

Layout

Install, build, test

pip install connex-rs   # prebuilt wheel: links the UCX of the build host by rpath; on
                        # RDMA/GPU hosts build from source against your UCX:

# Rust toolchain (rustup) and UCX with the transports you need (verbs/rdmacm
# for RDMA, CUDA for GPU memory). UCX_HOME points at its prefix.
cargo build --release
cargo test --release                       # UCX over TCP on loopback
CONNEX_RS_TEST_TLS=rc_mlx5 CONNEX_RS_TEST_IP=<IPoIB address> UCX_NET_DEVICES=<hca>:1 \
  cargo test --release                     # same tests over RDMA

# Python wheel (abi3, Python >= 3.8); links UCX by rpath, does not bundle it
pip install maturin
maturin build --release -o dist && pip install dist/connex_rs-*.whl
# backend tests over real Connex endpoints live in the Connex repository:
#   (in a Connex checkout) python -m pytest tests/test_connexrs_backend.py
python python/rma_bench.py local --ip <IPoIB address> --dev cuda:0 --peer-dev cuda:1

Python API

import connex_rs
e = connex_rs.Engine("llama-8b", stage=1, shard=0, replica=0,
                     listen_ip="10.0.0.2", port=51001,
                     heartbeat_ms=1000, heartbeat_timeout_ms=3000)
e.local_url                     # "10.0.0.2:51001#llama-8b|1|0|0"
e.register_memory(t.data_ptr(), t.numel())        # host or GPU memory
e.transfer_sync_read(peer_url, dst_ptr, src_ptr, n)   # 0, or < 0 on failure
h = e.put_signal(peer_url, src_ptr, dst_ptr, n, b"step 7")  # h.wait()
e.send_msg(peer_url_or_id_tuple, b"...");  e.recv_msg(timeout_ms=100)
e.wait_event(timeout_ms=100)    # {"type": "up"|"joined"|"failed", "peer": (...), "reason": ...}
e.stats()                       # endpoints, regions, rkeys, in-flight, failures ...
e.close()

Return codes of transfer_sync_read / transfer_sync_write: 0 ok, -2 peer failed, -3 timeout (the connection is then closed, so no late DMA lands in the destination), -4 invalid (e.g. no registered region at that address, or a GPU transfer below UCX_ZCOPY_THRESH), -5 canceled (engine closed), -1 other UCX error.

UCX settings that matter

• UCX_NET_DEVICES: pin the RDMA device and its IPoIB interface (connex.pp does this automatically, CONNEX_RS_NET_DEVICES=auto). On hosts with container or VPN interfaces UCX otherwise tries those and the connection fails. HCA names can differ between hosts. The list must also name the interface that carries the listen address (lo for 127.0.0.1) unless that is the IPoIB interface of a listed RDMA device: UCX's connection manager only uses listed devices. listen rejects such an address with an error that names the interface; earlier versions joined and then failed every connect with "Destination is unreachable".
• Protocol: when RDMA is usable the engine sets PROTO_ENABLE=n on its own UCX configuration. UCX 1.16's protocol v2 selects software emulation for get/put on CUDA memory over client/server endpoints, while v1 does GPUDirect zero-copy. Without RDMA it keeps v2, because v1 would emulate RMA on CUDA memory by touching the device pointer from the CPU. An explicit UCX_PROTO_ENABLE always wins.
• UCX_TCP_CM_REUSEADDR=y (set by connex.pp): a restarted stage can listen on its port again at once.
• CONNEX_RS_DEBUG=1 prints every event and failed transfer.
• Defaults the engine sets on its own UCX configuration unless the variable is set: IB_ALLOC=mmap,heap when RDMA is usable (the IB transports' buffer pools otherwise use transparent huge pages, and on a host with fragmented memory the compaction stall made a connect take up to 1.7 s), and TCP_TX_SEG_SIZE / TCP_RX_SEG_SIZE=256k when TCP is in use (reads over TCP are emulated in segments of that size; 8 KB gave 0.9 GB/s).
• On hosts with many network interfaces (containers, VPN), a TCP connect spends 100-400 ms in UCX computing configurations for every interface; UCX_NET_DEVICES cuts that as well.

Limits

• Node loss or a network partition (no RST, no ACK) is detected by the heartbeat or an op timeout and confirmed by the probe, not by UCX in ms. This path was not tested by taking a link down.
• Without RDMA (TCP), GPU memory is staged through host memory by UCX and is slow; host memory over TCP works as a fallback.
• One endpoint per peer pair: if two peers connect to each other at the same moment, both keep the connection opened by the lower logical id (no extra message exchange). connex.pp still links only towards peers ordered after itself, which avoids the case altogether.
• RMA from or into registered GPU memory must be at least UCX_ZCOPY_THRESH bytes (4096 unless set); smaller transfers fail with -4 instead of crashing in UCX's CPU-copy protocols.

Measurements: docs/results.md, section 6, in the Connex repository.

Used by

Connex, the communication library for endpoint mobility in distributed LLM serving, installs it with pip install "connex-ai[rs]" and selects it with CONNEX_PP_BACKEND=auto (or connex-rs). Changes are listed in CHANGELOG.md.

License

Apache-2.0, see LICENSE.