From 8fb27bd2dadbf0dea24b1b78f4977b0d65211b40 Mon Sep 17 00:00:00 2001 From: Sophie Date: Fri, 18 Sep 2026 08:34:12 +0200 Subject: [PATCH] Report a dead peer in JACCL instead of spinning forever; name the link in init errors On Apple's Thunderbolt RDMA stack a UC queue pair never reports a dead peer (no error completion, no event) and RC queue pairs are not available (ibv_create_qp fails with EOPNOTSUPP), so when a rank dies every other rank polls its completion queue forever at 100% CPU (#4278). The TCP side channel already lives as long as the group. Expose its socket fds (TCPAllGather::fds, SideChannel::liveness_fds) and, in every completion wait loop of the mesh and ring implementations, check them while no completion arrives: at most every 250 ms, poll() with a zero timeout and a one-byte MSG_PEEK on POLLIN. HUP/ERR/NVAL or a zero-byte peek means the peer process is gone and the loop throws a runtime_error naming the peer and the operation. No reads, no writes, no locks, so it is safe from the ring wire threads. Rank 0 watches every peer and the other ranks watch rank 0, whose throw closes its sockets, so a non-zero rank's death reaches everyone in two hops. Measured on two M3 Ultras: 0.25 s from SIGKILL to the exception on the survivor, mesh and ring, either rank. A progress timeout (JACCL_PROGRESS_TIMEOUT_S / MLX_JACCL_PROGRESS_TIMEOUT_S, default 600 s, 0 disables) backstops the cases the side channel cannot see: a lost UC frame (no retransmission) or a peer that is alive but wedged. Diagnostics: on a device whose port is PORT_DOWN ibv_open_device succeeds but ibv_alloc_pd returns NULL, and a peer whose IPv4-mapped GID is not reachable on this link fails the RTR transition with errno 60; both errors now name the device, its port state and the peer address. TCPSocket::recv reports "peer closed the side channel" when recv() returns 0 instead of a stale errno (the "Recv failed with errno=2" that every other rank printed when one rank failed during init). The stock 3-argument SideChannel constructor is kept as an overload so the exported symbol set of libjaccl is a superset of the previous one. --- mlx/distributed/jaccl/lib/jaccl/jaccl.cpp | 4 +- mlx/distributed/jaccl/lib/jaccl/mesh.cpp | 8 +- mlx/distributed/jaccl/lib/jaccl/mesh_impl.h | 20 ++- mlx/distributed/jaccl/lib/jaccl/rdma.cpp | 155 +++++++++++++++++++- mlx/distributed/jaccl/lib/jaccl/rdma.h | 69 +++++++++ mlx/distributed/jaccl/lib/jaccl/ring.cpp | 9 +- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 15 +- mlx/distributed/jaccl/lib/jaccl/tcp.cpp | 9 +- 8 files changed, 276 insertions(+), 13 deletions(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp b/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp index 7715022fcc..489001a094 100644 --- a/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp @@ -215,12 +215,14 @@ SideChannel Config::get_side_channel() const { auto tcp = std::make_shared(rank_, size_, get_coordinator().c_str()); + auto fds = tcp->fds(); // watched for peer death return SideChannel( rank_, size_, [tcp = std::move(tcp)](const char* src, char* dst, size_t n_bytes) { (*tcp)(src, dst, n_bytes); - }); + }, + std::move(fds)); } Config Config::from_env() { diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp index 5b749f2bd9..3fecac470c 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp @@ -28,7 +28,13 @@ MeshGroup::MeshGroup( side_channel_.barrier(); // Create the mesh implementation object - mesh_ = MeshImpl(rank_, size_, connections_, buffers_, scatter_buffers_); + mesh_ = MeshImpl( + rank_, + size_, + connections_, + buffers_, + scatter_buffers_, + side_channel_.liveness_fds()); } void MeshGroup::initialize() { diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index 75c540e917..c430ebd8d1 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -20,14 +20,15 @@ class MeshImpl { int size, std::vector& conns, std::vector& buffers, - std::vector& scatter_buffers) + std::vector& scatter_buffers, + std::vector liveness_fds = {}) : rank_(rank), size_(size), connections_(conns), buffers_(buffers), scatter_buffers_(scatter_buffers), - staging_mem_( - std::make_unique(MESH_PIPELINE * MAX_BUFFER_SIZE)) {} + staging_mem_(std::make_unique(MESH_PIPELINE * MAX_BUFFER_SIZE)), + liveness_fds_(std::move(liveness_fds)) {} MeshImpl() : rank_(0), size_(1) {} @@ -86,6 +87,7 @@ class MeshImpl { } // Main loop + ProgressGuard _pg1(liveness_fds_, rank_, "all_reduce"); while (reduce_chunk < total_chunks) { // Poll the hardware for completions. // @@ -100,6 +102,7 @@ class MeshImpl { // receives. ibv_wc wc[WC_NUM]; int n = poll(connections_, WC_NUM, wc); + _pg1.tick(n > 0); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -194,9 +197,11 @@ class MeshImpl { } // Drain remaining in-flight completions (outstanding sends). + ProgressGuard _pg2(liveness_fds_, rank_, "all_reduce"); while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll(connections_, WC_NUM, wc); + _pg2.tick(n > 0); in_flight -= n; } } @@ -242,9 +247,11 @@ class MeshImpl { // Main loop // // Keep going until we have no longer data in flight. + ProgressGuard _pg3(liveness_fds_, rank_, "all_gather"); while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll(connections_, WC_NUM, wc); + _pg3.tick(n > 0); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -342,9 +349,11 @@ class MeshImpl { // Main loop // // Keep going until we have no longer data in flight. + ProgressGuard _pg4(liveness_fds_, rank_, "sum_scatter"); while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll(connections_, WC_NUM, wc); + _pg4.tick(n > 0); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -465,6 +474,7 @@ class MeshImpl { } // Main loop + ProgressGuard _pg5(liveness_fds_, rank_, "send"); while (in_flight > 0) { // Poll the hardware for completions. // @@ -472,6 +482,7 @@ class MeshImpl { // and send them. ibv_wc wc[WC_NUM]; int n = connections_[dst].poll(WC_NUM, wc); + _pg5.tick(n > 0); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; int rank = wc[i].wr_id & 0xff; @@ -510,6 +521,7 @@ class MeshImpl { } // Main loop + ProgressGuard _pg6(liveness_fds_, rank_, "recv"); while (in_flight > 0) { // Poll the hardware for completions. // @@ -517,6 +529,7 @@ class MeshImpl { // data to fetch post another recv. ibv_wc wc[WC_NUM]; int n = connections_[src].poll(WC_NUM, wc); + _pg6.tick(n > 0); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; int rank = wc[i].wr_id & 0xff; @@ -627,6 +640,7 @@ class MeshImpl { std::span buffers_; std::span scatter_buffers_; std::unique_ptr staging_mem_; + std::vector liveness_fds_; }; } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp index f449d8c814..802caa9372 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp @@ -1,6 +1,8 @@ // Copyright © 2025 Apple Inc. #include +#include +#include #include #include #include @@ -140,10 +142,44 @@ Connection::~Connection() { } } +namespace { +// name the device and its port state in the +// errors that a Thunderbolt link problem produces. Measured on macOS 26.6.1: +// on a device whose port is PORT_DOWN, ibv_open_device succeeds but +// ibv_alloc_pd returns NULL; a peer whose IPv4-mapped GID is not reachable on +// this link fails the RTR transition with errno 60 (ETIMEDOUT). +std::string describe_device(ibv_context* ctx) { + std::ostringstream s; + if (ctx == nullptr) { + return ""; + } + const char* name = ctx->device ? ibv().get_device_name(ctx->device) : nullptr; + s << (name ? name : ""); + ibv_port_attr pa; + if (ibv().query_port(ctx, 1, &pa) == 0) { + static const char* states[] = { + "NOP", + "PORT_DOWN", + "PORT_INIT", + "PORT_ARMED", + "PORT_ACTIVE", + "ACTIVE_DEFER"}; + int st = static_cast(pa.state); + s << " (port " << (st >= 0 && st < 6 ? states[st] : "?") << ")"; + } + return s.str(); +} +} // namespace + void Connection::allocate_protection_domain() { protection_domain = ibv().alloc_pd(ctx); if (protection_domain == nullptr) { - throw std::runtime_error("[jaccl] Couldn't allocate protection domain"); + int err = errno; + std::ostringstream msg; + msg << "[jaccl] Couldn't allocate protection domain on " + << describe_device(ctx) << " (errno " << err + << "): is the Thunderbolt link up?"; + throw std::runtime_error(msg.str()); } } @@ -263,7 +299,21 @@ void Connection::queue_pair_rtr(const Destination& dst) { if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) { std::ostringstream msg; - msg << "[jaccl] Changing queue pair to RTR failed with errno " << status; + msg << "[jaccl] Changing queue pair to RTR failed with errno " << status + << " on " << describe_device(ctx) << " towards peer "; + if (attr.ah_attr.is_global) { + const uint8_t* g = dst.global_identifier.raw; + msg << (int)g[12] << "." << (int)g[13] << "." << (int)g[14] << "." + << (int)g[15]; + } else { + msg << "lid " << dst.local_id; + } + if (status == 60) { + msg << " (peer unreachable on this link: wrong cable/device or the " + << "peer's link-local address is gone)"; + } else if (status == 22) { + msg << " (invalid GID/attributes for this link)"; + } throw std::invalid_argument(msg.str()); } } @@ -376,14 +426,111 @@ void TCPAllGather::operator()(const char* src, char* dst, size_t n_bytes) { } SideChannel::SideChannel(int rank, int size, AllGatherFn agf) - : rank_(rank), size_(size), all_gather_fn_(std::move(agf)) {} + : SideChannel(rank, size, std::move(agf), std::vector{}) {} + +SideChannel::SideChannel( + int rank, + int size, + AllGatherFn agf, + std::vector liveness_fds) + : rank_(rank), + size_(size), + all_gather_fn_(std::move(agf)), + liveness_fds_(std::move(liveness_fds)) {} SideChannel::SideChannel(SideChannel&& sc) : rank_(sc.rank_), size_(sc.size_), - all_gather_fn_(std::move(sc.all_gather_fn_)) { + all_gather_fn_(std::move(sc.all_gather_fn_)), + liveness_fds_(std::move(sc.liveness_fds_)) { sc.rank_ = -1; sc.size_ = -1; } +// ── liveness + progress guard +// ───────────────────────────────────────────────── + +double progress_timeout_s() { + static double timeout = [] { + const char* v = std::getenv("JACCL_PROGRESS_TIMEOUT_S"); + if (v == nullptr) { + v = std::getenv("MLX_JACCL_PROGRESS_TIMEOUT_S"); + } + if (v == nullptr || *v == 0) { + return 600.0; + } + return std::atof(v); + }(); + return timeout; +} + +void check_peers_alive(std::span fds, int rank, const char* what) { + if (fds.empty()) { + return; + } + std::vector pfds; + pfds.reserve(fds.size()); + for (int fd : fds) { + pfds.push_back({fd, POLLIN, 0}); + } + int r = ::poll(pfds.data(), pfds.size(), 0); + if (r <= 0) { + return; // nothing to report (EINTR/EAGAIN included) + } + for (size_t i = 0; i < pfds.size(); i++) { + bool dead = (pfds[i].revents & (POLLHUP | POLLERR | POLLNVAL)) != 0; + if (!dead && (pfds[i].revents & POLLIN)) { + char b; + ssize_t n = ::recv(pfds[i].fd, &b, 1, MSG_PEEK | MSG_DONTWAIT); + dead = (n == 0) || (n < 0 && errno != EAGAIN && errno != EINTR); + } + if (dead) { + std::ostringstream msg; + msg << IBV_TAG << " peer is gone: side channel to "; + if (rank == 0) { + msg << "rank " << (i + 1); + } else { + msg << "rank 0 (coordinator)"; + } + msg << " closed while waiting in " << what + << " (a rank died or exited; UC RDMA never reports this by itself)"; + throw std::runtime_error(msg.str()); + } + } +} + +ProgressGuard::ProgressGuard( + std::span fds, + int rank, + const char* what) + : fds_(fds), + rank_(rank), + what_(what), + last_progress_(std::chrono::steady_clock::now()), + last_check_(last_progress_) {} + +void ProgressGuard::slow_tick() { + auto now = std::chrono::steady_clock::now(); + if (progressed_) { + last_progress_ = now; + progressed_ = false; + } + if (now - last_check_ < std::chrono::milliseconds(250)) { + return; + } + last_check_ = now; + check_peers_alive(fds_, rank_, what_); + double timeout = progress_timeout_s(); + if (timeout > 0) { + double idle = std::chrono::duration(now - last_progress_).count(); + if (idle > timeout) { + std::ostringstream msg; + msg << IBV_TAG << " no progress in " << what_ << " for " << (int)idle + << " s (JACCL_PROGRESS_TIMEOUT_S=" << (int)timeout + << "): lost frame or wedged peer"; + throw std::runtime_error(msg.str()); + } + } +} + } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.h b/mlx/distributed/jaccl/lib/jaccl/rdma.h index 01d9ed37d6..7432e8c628 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.h +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.h @@ -4,6 +4,8 @@ #include +#include +#include #include #include #include @@ -271,6 +273,48 @@ inline int poll( using AllGatherFn = std::function; +// ── liveness + progress guard +// ───────────────────────────────────────────────── UC queue pairs give no +// error completion when a peer dies (measured on Apple's Thunderbolt RDMA +// stack, macOS 26.6.1, scripts/jaccl/rdma_pair.c): a survivor polls forever. +// The side channel TCP sockets live as long as the group, so a closed socket IS +// the death signal. ProgressGuard is ticked from every completion wait loop; +// while no completion arrives it checks the side channel at most every 250 ms +// (poll()/MSG_PEEK only: no reads, no locks, safe from several wire threads at +// once) and enforces a hard progress timeout (JACCL_PROGRESS_TIMEOUT_S / +// MLX_JACCL_PROGRESS_TIMEOUT_S, default 600 s, 0 = disabled) as a backstop for +// a lost UC frame or a wedged peer. +void check_peers_alive(std::span fds, int rank, const char* what); +double progress_timeout_s(); + +class ProgressGuard { + public: + ProgressGuard(std::span fds, int rank, const char* what); + + inline void tick(bool progressed) { + if (progressed) { + idle_polls_ = 0; + progressed_ = true; + return; + } + if ((++idle_polls_ & 255) != 0) { + return; + } + slow_tick(); + } + + private: + void slow_tick(); + + std::span fds_; + int rank_; + const char* what_; + std::chrono::steady_clock::time_point last_progress_; + std::chrono::steady_clock::time_point last_check_; + uint32_t idle_polls_ = 0; + bool progressed_ = false; +}; + class TCPAllGather { public: TCPAllGather(int rank, int size, const char* addr); @@ -282,6 +326,17 @@ class TCPAllGather { void operator()(const char* src, char* dst, size_t n_bytes); + // the socket fds (rank 0: one per peer in rank order; other + // ranks: the coordinator). Values only, the sockets stay owned here. + std::vector fds() const { + std::vector r; + r.reserve(sockets_.size()); + for (auto& s : sockets_) { + r.push_back(static_cast(s)); + } + return r; + } + private: int rank_; int size_; @@ -299,8 +354,21 @@ class TCPAllGather { class SideChannel { public: SideChannel(int rank, int size, AllGatherFn agf); + // same, with the side channel socket fds to watch (kept as a + // separate overload so the stock 3-argument symbol stays exported: ABI). + SideChannel( + int rank, + int size, + AllGatherFn agf, + std::vector liveness_fds); SideChannel(SideChannel&& sc); + // side channel socket fds to watch for peer death (empty when + // the side channel is not TCP based, e.g. a user supplied all-gather). + const std::vector& liveness_fds() const { + return liveness_fds_; + } + SideChannel(const SideChannel&) = delete; SideChannel& operator=(const SideChannel&) = delete; @@ -358,6 +426,7 @@ class SideChannel { int rank_; int size_; AllGatherFn all_gather_fn_; + std::vector liveness_fds_; }; } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.cpp b/mlx/distributed/jaccl/lib/jaccl/ring.cpp index a08035f3f3..0ef4773094 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/ring.cpp @@ -34,7 +34,14 @@ RingGroup::RingGroup( // Create the ring implementation object ring_ = RingImpl( - rank_, size_, left_, right_, send_buffers_, recv_buffers_, &pool_); + rank_, + size_, + left_, + right_, + send_buffers_, + recv_buffers_, + &pool_, + side_channel_.liveness_fds()); } void RingGroup::initialize() { diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index d6305d9d4b..269a9e0e59 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -21,7 +21,8 @@ class RingImpl { std::vector& right, std::vector& send_buffers, std::vector& recv_buffers, - ThreadPool* pool = nullptr) + ThreadPool* pool = nullptr, + std::vector liveness_fds = {}) : rank_(rank), size_(size), n_conns_(left.size()), @@ -29,7 +30,8 @@ class RingImpl { right_(right), send_buffers_(send_buffers), recv_buffers_(recv_buffers), - pool_(pool) {} + pool_(pool), + liveness_fds_(std::move(liveness_fds)) {} RingImpl( int rank, @@ -401,9 +403,11 @@ class RingImpl { } // Main loop + ProgressGuard _pg1(liveness_fds_, rank_, "reduce_scatter_wire"); while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll_wire(lw, WC_NUM, wc); + _pg1.tick(n > 0); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -529,9 +533,11 @@ class RingImpl { // Main loop // // Keep going until we have no longer data in flight. + ProgressGuard _pg2(liveness_fds_, rank_, "reduce_scatter_wire"); while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll_wire(lw, WC_NUM, wc); + _pg2.tick(n > 0); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -636,6 +642,7 @@ class RingImpl { } // Main loop + ProgressGuard _pg3(liveness_fds_, rank_, "send_wire"); while (in_flight > 0) { // Poll the hardware for completions. // @@ -643,6 +650,7 @@ class RingImpl { // and send them. ibv_wc wc[WC_NUM]; int n = conns[lw].poll(WC_NUM, wc); + _pg3.tick(n > 0); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; @@ -710,6 +718,7 @@ class RingImpl { } // Main loop + ProgressGuard _pg4(liveness_fds_, rank_, "recv_wire"); while (in_flight > 0) { // Poll the hardware for completions. // @@ -717,6 +726,7 @@ class RingImpl { // data to fetch post another recv. ibv_wc wc[WC_NUM]; int n = conns[lw].poll(WC_NUM, wc); + _pg4.tick(n > 0); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; @@ -820,6 +830,7 @@ class RingImpl { std::span send_buffers_; std::span recv_buffers_; ThreadPool* pool_; + std::vector liveness_fds_; }; } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/tcp.cpp b/mlx/distributed/jaccl/lib/jaccl/tcp.cpp index 08f2338c60..01ee7eddfe 100644 --- a/mlx/distributed/jaccl/lib/jaccl/tcp.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/tcp.cpp @@ -152,7 +152,14 @@ void TCPSocket::recv(const char* tag, void* data, size_t len) { auto n = ::recv(sock_, data, len, 0); if (n <= 0) { std::ostringstream msg; - msg << tag << " Recv failed with errno=" << errno; + // 0 means the peer closed the side channel + // (a rank failed during init or exited); errno is stale in that case. + if (n == 0) { + msg << tag << " Recv failed: peer closed the side channel" + << " (another rank failed during init or exited)"; + } else { + msg << tag << " Recv failed with errno=" << errno; + } throw std::runtime_error(msg.str()); } len -= n;