From d92347bae9e42572d483d5e35cd57fb753dd173e Mon Sep 17 00:00:00 2001 From: Drifter4242 Date: Tue, 15 Sep 2026 18:42:14 -0500 Subject: [PATCH 01/19] Fix jaccl ring all_gather: direction 1 slice was not mirrored (backport of ml-explore/mlx #4443) Upstream fix. RingImpl::all_gather_wire gave both directions the same slice of the region of each rank, so the upper half of the region of each peer was not transferred. Each direction now has its own half, as in all_reduce_wire. The mesh does not change. Applied without changes, with the upstream test test_all_gather_rank_data. Co-authored-by: Dave Hind Co-authored-by: Anastasiia Filippova (cherry picked from commit 7241f12e631440387a2156ec1018d1c9fe8e56b9) Co-Authored-By: Claude Opus 5.5 --- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 29 ++++++++++++++------- python/tests/mlx_distributed_tests.py | 16 ++++++++++++ 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index 15b81a2177..d6305d9d4b 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -198,7 +198,11 @@ class RingImpl { size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires); dispatch_wires(n_wires, [&](int lw) { all_gather_wire( - out_ptr, n_bytes, static_cast(n_bytes_per_wire), lw); + out_ptr, + n_bytes, + static_cast(n_bytes_per_wire), + n_wires, + lw); }); } @@ -215,14 +219,21 @@ class RingImpl { char* out_ptr, int64_t n_bytes, int64_t n_bytes_per_wire, + int n_wires, int lw) { - // Both directions send the same contiguous slice of each rank's region. - int64_t slice = static_cast(lw) * n_bytes_per_wire; - int64_t wire_offset[2] = {slice, slice}; - // Clamp to this wire's slice end so the last frame can't spill into the - // next wire's slice. - int64_t wire_end_bytes = std::min(n_bytes, slice + n_bytes_per_wire); - int64_t wire_end[2] = {wire_end_bytes, wire_end_bytes}; + // Each rank's region is divided into 2 directional halves, each split + // into n_wires contiguous slices - mirroring all_reduce_wire. Without + // the lr term both directions cover the same lower-half slice and the + // upper half of every peer's region is never transferred. + int64_t wire_offset[2]; + int64_t wire_end[2]; + for (int lr = 0; lr < 2; lr++) { + wire_offset[lr] = lr * n_wires * n_bytes_per_wire + + static_cast(lw) * n_bytes_per_wire; + int64_t region_end = + std::min(n_bytes, (lr + 1) * n_wires * n_bytes_per_wire); + wire_end[lr] = std::min(region_end, wire_offset[lr] + n_bytes_per_wire); + } int64_t send_offset[2] = {rank_ * n_bytes, rank_ * n_bytes}; int64_t recv_offset[2] = { ((rank_ + size_ - 1) % size_) * n_bytes, @@ -230,7 +241,7 @@ class RingImpl { ring_pass<2, char>( lw, - n_bytes, + n_bytes * size_, n_bytes, n_bytes * size_, n_bytes_per_wire, diff --git a/python/tests/mlx_distributed_tests.py b/python/tests/mlx_distributed_tests.py index cbb9663046..7d6e56aaac 100644 --- a/python/tests/mlx_distributed_tests.py +++ b/python/tests/mlx_distributed_tests.py @@ -325,6 +325,22 @@ def test_all_gather(self): self.assertEqual(y.shape, (world.size() * 2, 2, 4)) self.assertTrue(mx.all(y == 1)) + def test_all_gather_rank_data(self): + # Every rank contributes distinct data so that a missing, stale, or + # misplaced peer region is detected (an all-ones gather cannot tell). + # Sizes include one large enough that backends which slice transfers + # across directions/wires (e.g. the jaccl ring) exercise every slice, + # and odd sizes that exercise the tail clamping. + world = mx.distributed.init() + for size in [7, 1024, 1000003, 4 * 1024 * 1024]: + base = mx.arange(size, dtype=mx.int32) % 100003 + x = base + world.rank() + y = mx.distributed.all_gather(x) + self.assertEqual(y.shape, (world.size() * size,)) + regions = y.reshape(world.size(), size) + for r in range(world.size()): + self.assertTrue(mx.all(regions[r] == base + r).item()) + def test_clip_grad_norm_sharded(self): world = mx.distributed.init() N = world.size() From d37393071ba45e1102ab7d75311c4d392b0ab127 Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Mon, 28 Sep 2026 06:10:36 -0700 Subject: [PATCH 02/19] chore: Fix assertion in ReduceScatter::eval_cpu (backport of ml-explore/mlx #4557) Upstream fix. ReduceScatter::eval_cpu asserted that it has no inputs and then read inputs[0]. A build with assertions stopped at each sum_scatter. The assertion now expects one input. A build without assertions does not change. Applied without changes. (cherry picked from commit 42987b66be853721c941897247c2e333efe9e2d9) Co-Authored-By: Claude Opus 5.5 --- mlx/backend/cpu/distributed.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx/backend/cpu/distributed.cpp b/mlx/backend/cpu/distributed.cpp index 6b788cde25..88ec269fe4 100644 --- a/mlx/backend/cpu/distributed.cpp +++ b/mlx/backend/cpu/distributed.cpp @@ -98,7 +98,7 @@ void Recv::eval_cpu( void ReduceScatter::eval_cpu( const std::vector& inputs, std::vector& outputs) { - assert(inputs.size() == 0); + assert(inputs.size() == 1); assert(outputs.size() == 1); auto [in, copied] = ensure_row_contiguous(inputs[0], stream()); From 4614968831eb71e02d514f9b8bc1cd489d44e02b Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 8 Sep 2026 20:14:23 +0900 Subject: [PATCH 03/19] Add a scoped_env utility for tests (partial backport of ml-explore/mlx #4474) Only the scoped_env helper in python/tests/mlx_tests.py. The test of the next commit (ml-explore/mlx #4558) uses it. The upstream commit also changes test_conv.py, test_conv_transpose.py and test_fast_sdpa.py to use the helper. Those changes are not included. (cherry picked from commit ec5f6ff7b1e0b314fc13f9ba8fb7e422ed65484d) Co-Authored-By: Claude Opus 5.5 --- python/tests/mlx_tests.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/python/tests/mlx_tests.py b/python/tests/mlx_tests.py index 2b60f46154..ba5595389d 100644 --- a/python/tests/mlx_tests.py +++ b/python/tests/mlx_tests.py @@ -1,5 +1,6 @@ # Copyright © 2023 Apple Inc. +import contextlib import os import platform import sys @@ -10,6 +11,26 @@ import numpy as np +@contextlib.contextmanager +def scoped_env(**environ): + """ + Temporarily set the process environment variables. + + Passing a value of None removes the variable for the duration of the context. + """ + old_environ = dict(os.environ) + for key, value in environ.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + try: + yield + finally: + os.environ.clear() + os.environ.update(old_environ) + + class MLXTestRunner(unittest.TestProgram): def __init__(self, *args, **kwargs): # Do not exit in runTests From 7112cbf8e08e496c996f0a7b67c1bc2746e57534 Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Thu, 1 Oct 2026 05:11:48 -0700 Subject: [PATCH 04/19] Do not cache failed distributed backend initialization (backport of ml-explore/mlx #4558) Upstream fix. When a backend could not be initialized and strict was false, init registered an empty group under the name of the backend. A later init of the same backend returned that empty group, also when strict was true. A failed initialization is no longer registered, so the later call tries again and a strict call throws. A group that was initialized stays registered until the process exits. This is also true for a JACCL group that a failure closed. Applied without changes, with the upstream test test_strict_init_after_fallback. (cherry picked from commit 8803e524567ea612dcfc4df703014631296603ec) Co-Authored-By: Claude Opus 5.5 --- mlx/distributed/distributed.cpp | 7 +++---- python/tests/ring_test_distributed.py | 10 ++++++++++ 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/mlx/distributed/distributed.cpp b/mlx/distributed/distributed.cpp index 34be2ec600..d8fca8b160 100644 --- a/mlx/distributed/distributed.cpp +++ b/mlx/distributed/distributed.cpp @@ -148,12 +148,11 @@ get_backends() { } Group register_group(std::shared_ptr group, std::string bk) { - auto& backends = get_backends(); if (group == nullptr) { - group = std::make_shared(); - } else { - backends.insert({"any", group}); + return Group(std::make_shared()); } + auto& backends = get_backends(); + backends.insert({"any", group}); backends.insert({std::move(bk), group}); return Group(group); } diff --git a/python/tests/ring_test_distributed.py b/python/tests/ring_test_distributed.py index 31db2ed5f7..8b4de9ccae 100644 --- a/python/tests/ring_test_distributed.py +++ b/python/tests/ring_test_distributed.py @@ -23,6 +23,16 @@ def test_groups(self): with self.assertRaises(RuntimeError): sub = world.split(world.rank() % 2) + def test_strict_init_after_fallback(self): + with mlx_tests.scoped_env( + JACCL_IBV_DEVICES=None, + MLX_IBV_DEVICES=None, + ): + self.assertEqual(mx.distributed.init(backend="jaccl").size(), 1) + + with self.assertRaises(RuntimeError): + mx.distributed.init(backend="jaccl", strict=True) + def test_all_reduce_extra(self): world = mx.distributed.init() dtypes = [ From 95e893852ecbb09f9247f37dfdaa014ba0a482f7 Mon Sep 17 00:00:00 2001 From: Gajesh Naik <26431906+Gajesh2007@users.noreply.github.com> Date: Mon, 28 Sep 2026 00:14:41 -0700 Subject: [PATCH 05/19] fix(jaccl): clear the unused tail of each send frame JACCL posts every frame at the full size of its scratch buffer. When a payload was shorter than the buffer, the bytes of earlier sends stayed after it and went to the peer. stage_send_frame() copies the payload and sets the remaining bytes of the frame to zero. All mesh and ring send paths now stage through SharedBuffer::stage_send(). The posted sizes and the receive paths do not change. This is 4e89c2e3 ("Preserve cluster JACCL send-frame clearing experiment"). It applied without conflicts: the three changed headers are identical in cb77239b and in the parent of that commit. clang-format with the repository style was then applied to the four touched files. Checked on CPU only, with the NativeSendChecks harness (the real JACCL headers, simulated verbs, ASan and UBSan). With the cb77239b headers the stale-tail executable finds old bytes after the payload. With this commit the five check groups pass and the stale-tail executable no longer finds them. No RDMA hardware was used. (cherry picked from commit 4e89c2e3edf031c3880b54f0f3d20ea1b3f39d79) Co-Authored-By: Claude Opus 5.5 --- mlx/distributed/jaccl/lib/jaccl/mesh_impl.h | 42 ++++---------- mlx/distributed/jaccl/lib/jaccl/rdma.h | 7 +++ mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 59 ++++++++++---------- mlx/distributed/jaccl/lib/jaccl/send_frame.h | 33 +++++++++++ 4 files changed, 83 insertions(+), 58 deletions(-) create mode 100644 mlx/distributed/jaccl/lib/jaccl/send_frame.h diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index 75c540e917..88467921a5 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -73,10 +73,7 @@ class MeshImpl { int64_t elems = std::min(N, total - read_offset); std::copy( in + read_offset, in + read_offset + elems, local_staging(buff)); - std::copy( - in + read_offset, - in + read_offset + elems, - send_buffer(sz, buff).begin()); + send_buffer(sz, buff).stage_send(in + read_offset, elems); recv_end[rank_]++; post_send_all(sz, buff); @@ -111,10 +108,7 @@ class MeshImpl { completed_send_count[buff]++; if (completed_send_count[buff] == num_peers) { int64_t elems = std::min(N, total - read_offset); - std::copy( - in + read_offset, - in + read_offset + elems, - send_buffer(sz, buff).begin()); + send_buffer(sz, buff).stage_send(in + read_offset, elems); post_send_all(sz, buff); completed_send_count[buff] = 0; @@ -228,10 +222,8 @@ class MeshImpl { int buff = 0; while (read_offset < total && buff < PIPELINE) { post_recv_all(sz, buff); - std::copy( - our_data + read_offset, - our_data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); + send_buffer(sz, buff).stage_send( + our_data + read_offset, std::min(N, total - read_offset)); post_send_all(sz, buff); buff++; @@ -256,10 +248,8 @@ class MeshImpl { if (work_type == SEND_WR && read_offset < total) { completed_send_count[buff]++; if (completed_send_count[buff] == num_peers) { - std::copy( - our_data + read_offset, - our_data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); + send_buffer(sz, buff).stage_send( + our_data + read_offset, std::min(N, total - read_offset)); post_send_all(sz, buff); completed_send_count[buff] = 0; @@ -329,8 +319,7 @@ class MeshImpl { continue; } const T* src = in + static_cast(p) * count + read_offset; - std::copy( - src, src + elems, scatter_send_buffer(sz, buff, p).begin()); + scatter_send_buffer(sz, buff, p).stage_send(src, elems); } scatter_post_send_all(sz, buff); @@ -363,10 +352,7 @@ class MeshImpl { continue; } const T* src = in + static_cast(p) * count + read_offset; - std::copy( - src, - src + elems, - scatter_send_buffer(sz, buff, p).begin()); + scatter_send_buffer(sz, buff, p).stage_send(src, elems); } scatter_post_send_all(sz, buff); @@ -453,10 +439,8 @@ class MeshImpl { // Prefill the pipeline int buff = 0; while (read_offset < n_bytes && buff < PIPELINE) { - std::copy( - in_ptr + read_offset, - in_ptr + std::min(read_offset + N, n_bytes), - send_buffer(sz, buff).begin()); + send_buffer(sz, buff).stage_send( + in_ptr + read_offset, std::min(N, n_bytes - read_offset)); send_to(sz, dst, buff); buff++; @@ -479,10 +463,8 @@ class MeshImpl { in_flight--; if (read_offset < n_bytes) { - std::copy( - in_ptr + read_offset, - in_ptr + std::min(read_offset + N, n_bytes), - send_buffer(sz, buff).begin()); + send_buffer(sz, buff).stage_send( + in_ptr + read_offset, std::min(N, n_bytes - read_offset)); send_to(sz, dst, buff); read_offset += N; diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.h b/mlx/distributed/jaccl/lib/jaccl/rdma.h index 01d9ed37d6..3f7d801026 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.h +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.h @@ -11,6 +11,7 @@ #include #include +#include "jaccl/send_frame.h" #include "jaccl/tcp.h" constexpr const char* IBV_TAG = "[jaccl]"; @@ -125,6 +126,12 @@ class SharedBuffer { return entry; } + template + void stage_send(const T* source, int64_t count) { + stage_send_frame( + std::span(static_cast(data_), size()), source, count); + } + template T* data() { return static_cast(data_); diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index d6305d9d4b..80f7e96922 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -387,11 +387,13 @@ class RingImpl { } for (int lr = 0; lr < MAX_DIR; lr++) { int64_t offset = wire_offset[lr] + send_count[lr] * N; - std::copy( - send_base + send_base_offset[lr] + offset, - send_base + send_base_offset[lr] + - std::max(offset, std::min(offset + N, wire_end[lr])), - send_buffer(sz, buff, lr, lw).template begin()); + const int64_t elems = + std::max(0, std::min(N, wire_end[lr] - offset)); + send_buffer(sz, buff, lr, lw) + .stage_send( + elems > 0 ? send_base + send_base_offset[lr] + offset + : nullptr, + elems); send_count[lr]++; send_to(sz, buff, lr, lw); } @@ -414,11 +416,13 @@ class RingImpl { if (work_type == SEND_WR) { if (send_count[lr] < n_steps) { int64_t offset = wire_offset[lr] + send_count[lr] * N; - std::copy( - send_base + send_base_offset[lr] + offset, - send_base + send_base_offset[lr] + - std::max(offset, std::min(offset + N, wire_end[lr])), - send_buffer(sz, buff, lr, lw).template begin()); + const int64_t elems = + std::max(0, std::min(N, wire_end[lr] - offset)); + send_buffer(sz, buff, lr, lw) + .stage_send( + elems > 0 ? send_base + send_base_offset[lr] + offset + : nullptr, + elems); send_count[lr]++; send_to(sz, buff, lr, lw); in_flight++; @@ -513,11 +517,12 @@ class RingImpl { } for (int lr = 0; lr < MAX_DIR; lr++) { int64_t offset = wire_offset[lr] + send_count[lr] * N; - std::copy( - send_base + send_offset[lr] + offset, - send_base + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).template begin()); + const int64_t elems = + std::max(0, std::min(N, send_limits[lr] - offset)); + send_buffer(sz, buff, lr, lw) + .stage_send( + elems > 0 ? send_base + send_offset[lr] + offset : nullptr, + elems); send_count[lr]++; send_to(sz, buff, lr, lw); } @@ -541,11 +546,12 @@ class RingImpl { if (work_type == SEND_WR && send_count[lr] < n_steps) { int64_t offset = wire_offset[lr] + send_count[lr] * N; - std::copy( - send_base + send_offset[lr] + offset, - send_base + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).template begin()); + const int64_t elems = + std::max(0, std::min(N, send_limits[lr] - offset)); + send_buffer(sz, buff, lr, lw) + .stage_send( + elems > 0 ? send_base + send_offset[lr] + offset : nullptr, + elems); send_to(sz, buff, lr, lw); in_flight++; send_count[lr]++; @@ -624,10 +630,8 @@ class RingImpl { // Prefill the pipeline int buff = 0; while (read_offset < limit && buff < PIPELINE) { - std::copy( - in_ptr + read_offset, - in_ptr + std::min(read_offset + N, limit), - send_buffer(sz, buff, dir, lw).begin()); + send_buffer(sz, buff, dir, lw) + .stage_send(in_ptr + read_offset, std::min(N, limit - read_offset)); send_to(sz, buff, dir, lw); buff++; @@ -649,10 +653,9 @@ class RingImpl { in_flight--; if (read_offset < limit) { - std::copy( - in_ptr + read_offset, - in_ptr + std::min(read_offset + N, limit), - send_buffer(sz, buff, dir, lw).begin()); + send_buffer(sz, buff, dir, lw) + .stage_send( + in_ptr + read_offset, std::min(N, limit - read_offset)); send_to(sz, buff, dir, lw); read_offset += N; diff --git a/mlx/distributed/jaccl/lib/jaccl/send_frame.h b/mlx/distributed/jaccl/lib/jaccl/send_frame.h new file mode 100644 index 0000000000..ded3e4bd35 --- /dev/null +++ b/mlx/distributed/jaccl/lib/jaccl/send_frame.h @@ -0,0 +1,33 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include +#include +#include +#include +#include + +namespace jaccl { + +// The posted receive/send size stays fixed. Clear every byte after this +// payload. +template +inline void +stage_send_frame(std::span destination, const T* source, int64_t count) { + if (destination.empty() || destination.data() == nullptr || count < 0 || + static_cast(count) > destination.size() / sizeof(T) || + (count > 0 && source == nullptr)) { + throw std::invalid_argument( + "[jaccl] Send payload exceeds its scratch frame"); + } + const size_t bytes = static_cast(count) * sizeof(T); + if (count > 0) { + std::copy_n(source, count, reinterpret_cast(destination.data())); + } + if (bytes < destination.size()) { + std::memset(destination.data() + bytes, 0, destination.size() - bytes); + } +} + +} // namespace jaccl From cbe63d1460e0d84ec828a8b62c7f3b9f03eb1207 Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Thu, 8 Oct 2026 17:35:05 -0700 Subject: [PATCH 06/19] fix(jaccl): bound the polling loops and release a group that fails The mesh and ring polling loops waited for completions without a limit. When a peer died or a completion was lost, the rank used a full CPU core until it was killed. In the incident, a rank that was killed while it waited in MeshImpl::recv left about 160 GB of wired memory with no owning process. The loops also did not read the status of a work completion, so a failed completion was used as data. Every polling loop now calls ProgressGuard::check() after each poll: - A completion with a status other than IBV_WC_SUCCESS is an error. A negative return of ibv_poll_cq is an error, and the shared poll() helpers no longer add it to the count. - When no completion arrives for longer than the limit, the call fails. The limit is JACCL_PROGRESS_TIMEOUT_MS / MLX_JACCL_PROGRESS_TIMEOUT_MS in milliseconds, read once. The default is 120000. Zero or a negative value removes the limit. A value that is not a number keeps the default. In recv the wait includes the time the peer computes, so the default is large. A caller that knows its longest wait can lower it. - The fast path has no clock read. A counter counts the polls without a completion and the clock is read once per 1024 of them. Measured without sanitizers this adds 0.7 ns per poll. A clock read on each poll adds 12.7 ns. The call fails with a std::runtime_error that starts with "[jaccl]" and names the call and the cause. It does not exit the process. MLX runs the collectives on a CPU stream, and Scheduler::enqueue catches the exception and reports it on the thread of the caller, so the caller can do its own cleanup. The errors of the ring pool wires now reach the caller too. Before this change dispatch_wires dropped them because it used future::wait(). State of the group after a failure. It is the same for each exception that leaves a call: - The queue pair and the completion queue of each connection are destroyed, which cancels the posted work, and all buffers are deregistered. The ring does this after all wires returned. The wire that fails sets a flag, and the other wires stop at their next clock sample. - The protection domains and the device contexts stay open until the group is destroyed. MLX keeps the group until the process exits. - Each later call on the group throws immediately. The group cannot be opened again. A new process is necessary. - The peer gets no message. Its own guard ends its wait. The size of the completion queue does not change. Each connection has at most two sends and two receives posted (PIPELINE and NUM_BUFFERS are 2) and its queue has 64 entries. The harness measured a maximum of 4 queued completions. Checked with simulated verbs (the real mesh, ring and rdma headers, ASan, UBSan and TSan; see tests/jaccl in a later commit): a peer that does not answer in each mesh and ring call, a lost completion, a failed completion, a failed poll, calls after a failure, and traffic with waits below the limit. The same cases fail on the parent commit. Checked on two Macs with a Thunderbolt RDMA link and a two-rank mesh, before this commit was moved onto the upstream backports: normal traffic passes with the guard on, so the driver reports a success status on good completions. A rank whose peer ended reported the timeout at the limit and exited, with a model loaded and without one, and neither Mac kept a process or wired memory. The same case without the guard polled for more than nine minutes. A failed completion and the ring path were not run on the real driver. Co-Authored-By: Claude Opus 5.5 --- mlx/distributed/jaccl/lib/README.md | 10 ++ mlx/distributed/jaccl/lib/jaccl/mesh_impl.h | 53 +++++++ .../jaccl/lib/jaccl/progress_guard.h | 135 ++++++++++++++++++ mlx/distributed/jaccl/lib/jaccl/rdma.h | 44 +++++- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 93 +++++++++--- 5 files changed, 313 insertions(+), 22 deletions(-) create mode 100644 mlx/distributed/jaccl/lib/jaccl/progress_guard.h diff --git a/mlx/distributed/jaccl/lib/README.md b/mlx/distributed/jaccl/lib/README.md index d682a06bbc..91f20d2c24 100644 --- a/mlx/distributed/jaccl/lib/README.md +++ b/mlx/distributed/jaccl/lib/README.md @@ -102,6 +102,16 @@ variables: (rank 0 listener) - **JACCL_RING** / **MLX_JACCL_RING**: (Optional) Prefer ring topology over mesh +- **JACCL_PROGRESS_TIMEOUT_MS** / **MLX_JACCL_PROGRESS_TIMEOUT_MS**: (Optional) + The longest time in milliseconds that a call waits for an RDMA completion. + The default is 120000. Zero or a negative value removes the limit. + +A call that gets no completion for this time, or that gets a failed work +completion, fails with an exception. Before the caller sees the exception, the +group destroys its queue pairs and completion queues and deregisters its +buffers. Every later call on that group throws immediately. Set the limit +above the longest time that a rank can wait for a peer, for example in `recv` +while the peer computes. ### Device File Format diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index 88467921a5..bbe99dcb50 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -2,9 +2,12 @@ #pragma once +#include +#include #include #include +#include "jaccl/progress_guard.h" #include "jaccl/rdma.h" constexpr int MESH_MAX_PEERS = 8; @@ -41,6 +44,7 @@ class MeshImpl { // Our own data is copied to a staging buffer to ensure we can reduce it in // the output when needed. + Call call(*this, "mesh all_reduce"); auto [sz, buffer_size] = buffer_size_from_message(size * sizeof(T)); int64_t N = buffer_size / sizeof(T); constexpr int PIPELINE = 2; @@ -97,6 +101,7 @@ class MeshImpl { // receives. ibv_wc wc[WC_NUM]; int n = poll(connections_, WC_NUM, wc); + call.check(wc, n); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -191,11 +196,14 @@ class MeshImpl { while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll(connections_, WC_NUM, wc); + call.check(wc, n); in_flight -= n; } } void all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes) { + Call call(*this, "mesh all_gather"); + // Copy our data to the appropriate place. Skip when in place (the scatter // gather all reduce passes our own reduced shard which already lives at // out_ptr + rank_ * n_bytes). @@ -237,6 +245,7 @@ class MeshImpl { while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll(connections_, WC_NUM, wc); + call.check(wc, n); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -289,6 +298,7 @@ class MeshImpl { // we use the dedicated scatter buffers: per (sz, buff) tile there are size_ // send slots (slot p -> peer p) and size_ recv slots (slot p <- peer p). + Call call(*this, "mesh sum_scatter"); const T* our_chunk = in + static_cast(rank_) * count; auto [sz, buffer_size] = buffer_size_from_message(count * sizeof(T)); @@ -334,6 +344,7 @@ class MeshImpl { while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll(connections_, WC_NUM, wc); + call.check(wc, n); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -429,6 +440,7 @@ class MeshImpl { } void send(const char* in_ptr, int64_t n_bytes, int dst) { + Call call(*this, "mesh send"); constexpr int PIPELINE = 2; constexpr int WC_NUM = PIPELINE; auto [sz, N] = buffer_size_from_message(n_bytes); @@ -456,6 +468,7 @@ class MeshImpl { // and send them. ibv_wc wc[WC_NUM]; int n = connections_[dst].poll(WC_NUM, wc); + call.check(wc, n); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; int rank = wc[i].wr_id & 0xff; @@ -475,6 +488,7 @@ class MeshImpl { } void recv(char* out_ptr, int64_t n_bytes, int src) { + Call call(*this, "mesh recv"); constexpr int PIPELINE = 2; constexpr int WC_NUM = PIPELINE; auto [sz, N] = buffer_size_from_message(n_bytes); @@ -499,6 +513,7 @@ class MeshImpl { // data to fetch post another recv. ibv_wc wc[WC_NUM]; int n = connections_[src].poll(WC_NUM, wc); + call.check(wc, n); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; int rank = wc[i].wr_id & 0xff; @@ -522,6 +537,42 @@ class MeshImpl { } private: + // One call into the group. It refuses a closed group, and it closes the + // group when the call ends with an exception. + class Call : public ProgressGuard { + public: + Call(MeshImpl& mesh, const char* op) + : ProgressGuard(op, *mesh.failed_), + mesh_(mesh), + exceptions_(std::uncaught_exceptions()) { + require_open(*mesh.failed_); + } + + ~Call() { + if (std::uncaught_exceptions() > exceptions_) { + mesh_.close(); + } + } + + private: + MeshImpl& mesh_; + int exceptions_; + }; + + // Cancel the posted work and unpin the buffers. + void close() { + failed_->store(true); + for (auto& c : connections_) { + c.release(); + } + for (auto& b : buffers_) { + b.deregister(); + } + for (auto& b : scatter_buffers_) { + b.deregister(); + } + } + void send_to(int sz, int rank, int buff) { connections_[rank].post_send( send_buffer(sz, buff), SEND_WR << 16 | buff << 8 | rank); @@ -609,6 +660,8 @@ class MeshImpl { std::span buffers_; std::span scatter_buffers_; std::unique_ptr staging_mem_; + std::unique_ptr> failed_ = + std::make_unique>(false); }; } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/progress_guard.h b/mlx/distributed/jaccl/lib/jaccl/progress_guard.h new file mode 100644 index 0000000000..6a18c549cc --- /dev/null +++ b/mlx/distributed/jaccl/lib/jaccl/progress_guard.h @@ -0,0 +1,135 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace jaccl { + +constexpr int64_t DEFAULT_PROGRESS_TIMEOUT_MS = 120000; +// The clock is read once per this number of polls without a completion. +constexpr int64_t PROGRESS_SAMPLE_POLLS = 1024; + +/** + * The longest time in milliseconds that a call waits for a completion. It is + * read once from JACCL_PROGRESS_TIMEOUT_MS / MLX_JACCL_PROGRESS_TIMEOUT_MS. + * Zero or a negative value removes the limit. + */ +inline int64_t progress_timeout_ms() { + static const int64_t timeout = []() -> int64_t { + const char* value = std::getenv("JACCL_PROGRESS_TIMEOUT_MS"); + if (value == nullptr) { + value = std::getenv("MLX_JACCL_PROGRESS_TIMEOUT_MS"); + } + if (value == nullptr) { + return DEFAULT_PROGRESS_TIMEOUT_MS; + } + // Keep the default when the value is not a number. + char* end = nullptr; + errno = 0; + long long parsed = std::strtoll(value, &end, 10); + if (end == value || *end != '\0' || errno == ERANGE) { + return DEFAULT_PROGRESS_TIMEOUT_MS; + } + return parsed; + }(); + return timeout; +} + +// Thrown in a wire when a different wire of the same call failed first. +struct WireStopped : std::runtime_error { + WireStopped() + : std::runtime_error( + "[jaccl] Stopped because a different wire of the call failed.") {} +}; + +// Refuse a group that an earlier failure closed. +inline void require_open(const std::atomic& failed) { + if (failed.load()) { + throw std::runtime_error( + "[jaccl] The group is closed after an earlier failure and cannot be " + "used again."); + } +} + +/** + * Stops a polling loop that makes no progress. Call check() after each poll. + * + * check() throws when a completion reports a failure, when the poll fails or + * when no completion arrived for longer than the timeout. It sets `failed` + * before it throws, which also stops the other wires of the same call. + */ +class ProgressGuard { + public: + ProgressGuard(const char* op, std::atomic& failed) + : op_(op), failed_(failed), timeout_ms_(progress_timeout_ms()) {} + + void check(const ibv_wc* wc, int n) { + if (n > 0) { + idle_polls_ = 0; + for (int i = 0; i < n; i++) { + if (wc[i].status != IBV_WC_SUCCESS) { + std::ostringstream msg; + msg << "a work completion failed with status " << wc[i].status + << " (vendor error " << wc[i].vendor_err << ", id 0x" << std::hex + << wc[i].wr_id << ")"; + fail(msg.str()); + } + } + } else if (n < 0) { + fail("the completion queue could not be polled"); + } else if (++idle_polls_ % PROGRESS_SAMPLE_POLLS == 0) { + check_idle(); + } + } + + private: + void check_idle() { + if (failed_.load()) { + throw WireStopped(); + } + if (timeout_ms_ <= 0) { + return; + } + auto now = std::chrono::steady_clock::now(); + // The first sample after a completion starts the measurement. + if (idle_polls_ == PROGRESS_SAMPLE_POLLS) { + idle_since_ = now; + return; + } + int64_t idle_ms = + std::chrono::duration_cast(now - idle_since_) + .count(); + if (idle_ms > timeout_ms_) { + std::ostringstream msg; + msg << "no completion for " << idle_ms + << " ms (JACCL_PROGRESS_TIMEOUT_MS is " << timeout_ms_ << ")"; + fail(msg.str()); + } + } + + [[noreturn]] void fail(const std::string& reason) { + failed_.store(true); + std::ostringstream msg; + msg << "[jaccl] " << op_ << ": " << reason + << ". The group is closed and cannot be used again."; + throw std::runtime_error(msg.str()); + } + + const char* op_; + std::atomic& failed_; + int64_t timeout_ms_; + int64_t idle_polls_{0}; + std::chrono::steady_clock::time_point idle_since_; +}; + +} // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.h b/mlx/distributed/jaccl/lib/jaccl/rdma.h index 3f7d801026..7e96c67afb 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.h +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.h @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -110,6 +111,14 @@ class SharedBuffer { void register_to_protection_domain(ibv_pd* protection_domain); + // Unpin the memory. The buffer cannot be posted after this call. + void deregister() { + for (auto& [pd, mr] : memory_regions_) { + ibv().dereg_mr(mr); + } + memory_regions_.clear(); + } + size_t size() const { return num_bytes_; } @@ -224,6 +233,27 @@ struct Connection { int poll(int num_completions, ibv_wc* work_completions) { return ibv_poll_cq(completion_queue, num_completions, work_completions); } + + // Destroy the queue pair and the completion queue, which cancels the posted + // work. The connection cannot be used after this call. + void release() { + if (queue_pair != nullptr) { + if (int status = ibv().destroy_qp(queue_pair); status != 0) { + std::cerr << IBV_TAG << " Could not destroy a queue pair (" << status + << ")" << std::endl; + } else { + queue_pair = nullptr; + } + } + if (completion_queue != nullptr) { + if (int status = ibv().destroy_cq(completion_queue); status != 0) { + std::cerr << IBV_TAG << " Could not destroy a completion queue (" + << status << ")" << std::endl; + } else { + completion_queue = nullptr; + } + } + } }; std::vector create_connections( @@ -246,6 +276,10 @@ inline int poll( c.completion_queue, num_completions - completions, work_completions + completions); + // Report a failed poll instead of adding it to the count. + if (n < 0) { + return n; + } completions += n; } @@ -257,13 +291,15 @@ inline int poll( std::span connections_2, int num_completions, ibv_wc* work_completions) { - int completions = 0; - completions += poll(connections_1, num_completions, work_completions); - completions += poll( + int completions = poll(connections_1, num_completions, work_completions); + if (completions < 0) { + return completions; + } + int n = poll( connections_2, num_completions - completions, work_completions + completions); - return completions; + return n < 0 ? n : completions + n; } /** diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index 80f7e96922..aa942d4283 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -2,9 +2,13 @@ #pragma once +#include +#include #include +#include #include +#include "jaccl/progress_guard.h" #include "jaccl/rdma.h" #include "jaccl/threadpool.h" @@ -368,6 +372,7 @@ class RingImpl { } }; + ProgressGuard guard("ring reduce_scatter", *failed_); for (int k = 0; k < size_ - 1; k++) { // Step 0 forwards this rank's own input; later steps forward the partial // accumulated in the output by the previous step. @@ -406,6 +411,7 @@ class RingImpl { while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll_wire(lw, WC_NUM, wc); + guard.check(wc, n); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -504,6 +510,7 @@ class RingImpl { }; set_limits(); + ProgressGuard guard("ring all_reduce/all_gather", *failed_); for (int k = 0; k < size_ - 1; k++) { // Step 0 sends this rank's own input; later steps send the accumulated // partial from out_ptr. @@ -537,6 +544,7 @@ class RingImpl { while (in_flight > 0) { ibv_wc wc[WC_NUM]; int n = poll_wire(lw, WC_NUM, wc); + guard.check(wc, n); for (int i = 0; i < n; i++) { int work_type = wc[i].wr_id >> 16; int buff = (wc[i].wr_id >> 8) & 0xff; @@ -640,6 +648,7 @@ class RingImpl { } // Main loop + ProgressGuard guard("ring send", *failed_); while (in_flight > 0) { // Poll the hardware for completions. // @@ -647,6 +656,7 @@ class RingImpl { // and send them. ibv_wc wc[WC_NUM]; int n = conns[lw].poll(WC_NUM, wc); + guard.check(wc, n); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; @@ -713,6 +723,7 @@ class RingImpl { } // Main loop + ProgressGuard guard("ring recv", *failed_); while (in_flight > 0) { // Poll the hardware for completions. // @@ -720,6 +731,7 @@ class RingImpl { // data to fetch post another recv. ibv_wc wc[WC_NUM]; int n = conns[lw].poll(WC_NUM, wc); + guard.check(wc, n); for (int i = 0; i < n; i++) { int buff = (wc[i].wr_id >> 8) & 0xff; @@ -785,33 +797,76 @@ class RingImpl { // Run fn(lw) for each wire, the first n_wires - 1 on the pool and the last // inline, then wait for the pool calls before returning. + // + // A wire that throws sets failed_, which stops the other wires. When all + // wires returned the group is closed and the first error is thrown. template void dispatch_wires(int n_wires, Fn&& fn) { - if (n_wires <= 1 || pool_ == nullptr) { - for (int lw = 0; lw < n_wires; lw++) { + require_open(*failed_); + + auto run = [&](int lw) { + try { fn(lw); + } catch (...) { + failed_->store(true); + throw; } - return; - } + }; - std::vector> futures; - futures.reserve(n_wires - 1); - for (int lw = 0; lw < n_wires - 1; lw++) { - futures.emplace_back(pool_->enqueue(fn, lw)); - } + std::exception_ptr error; + std::exception_ptr stopped; + auto wait = [&](auto&& call) { + try { + call(); + } catch (const WireStopped&) { + stopped = std::current_exception(); + } catch (...) { + if (!error) { + error = std::current_exception(); + } + } + }; - // Wait for the pool calls even if the inline one throws, so they never - // outlive this frame. - try { - fn(n_wires - 1); - } catch (...) { + if (n_wires <= 1 || pool_ == nullptr) { + wait([&] { + for (int lw = 0; lw < n_wires; lw++) { + run(lw); + } + }); + } else { + std::vector> futures; + futures.reserve(n_wires - 1); + wait([&] { + for (int lw = 0; lw < n_wires - 1; lw++) { + futures.emplace_back(pool_->enqueue(run, lw)); + } + run(n_wires - 1); + }); for (auto& f : futures) { - f.wait(); + wait([&] { f.get(); }); } - throw; } - for (auto& f : futures) { - f.wait(); + + if (error || stopped) { + close(); + std::rethrow_exception(error ? error : stopped); + } + } + + // Cancel the posted work and unpin the buffers. No wire may be running. + void close() { + failed_->store(true); + for (auto& c : left_) { + c.release(); + } + for (auto& c : right_) { + c.release(); + } + for (auto& b : send_buffers_) { + b.deregister(); + } + for (auto& b : recv_buffers_) { + b.deregister(); } } @@ -823,6 +878,8 @@ class RingImpl { std::span send_buffers_; std::span recv_buffers_; ThreadPool* pool_; + std::unique_ptr> failed_ = + std::make_unique>(false); }; } // namespace jaccl From 8785a85e221517386bf6ad8a2dc4bc55e4c28500 Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:03:08 -0700 Subject: [PATCH 07/19] fix(jaccl): send from the GID that was advertised to the peer Connection::info() searches the port's GID table for the IPv4-mapped GID and sends that to the peer as this rank's address. queue_pair_rtr then sets the source GID by a fixed index, 1. The two agree only when the mapped GID happens to sit at index 1. Its position depends on what other addresses the interface carries. On a Thunderbolt port that also has a unique-local IPv6 address the table reads unique-local, link-local, IPv4-mapped, and the mapped GID is at index 2, so the queue pair would send from a link-local IPv6 GID while the peer was told the IPv4-mapped one. Remember the index where the advertised GID was found and use it as the source index. Behaviour is unchanged where the mapped GID is at index 1. On two Macs where the mapped GID is at index 2, normal traffic passes with this commit and also without it, so the Thunderbolt driver tolerates the mismatch there and those runs do not show that the change is necessary. The table shape above was read on an M3 Ultra running macOS 27.2. Co-Authored-By: Claude Opus 5.5 --- mlx/distributed/jaccl/lib/jaccl/rdma.cpp | 6 +++++- mlx/distributed/jaccl/lib/jaccl/rdma.h | 2 ++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp index f449d8c814..478a4d9f57 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp @@ -123,6 +123,7 @@ Connection::Connection(Connection&& c) : Connection(nullptr) { std::swap(completion_queue, c.completion_queue); std::swap(queue_pair, c.queue_pair); std::swap(src, c.src); + std::swap(source_gid_index, c.source_gid_index); } Connection::~Connection() { @@ -196,6 +197,7 @@ const Destination& Connection::info() { if (*(uint64_t*)&tmp.raw[0] == 0 && *(uint16_t*)&tmp.raw[8] == 0 && *(uint16_t*)&tmp.raw[10] == 0xffff) { gid = tmp; + source_gid_index = i; found_gid = true; break; } @@ -255,7 +257,9 @@ void Connection::queue_pair_rtr(const Destination& dst) { attr.ah_attr.is_global = 1; attr.ah_attr.grh.hop_limit = 1; attr.ah_attr.grh.dgid = dst.global_identifier; - attr.ah_attr.grh.sgid_index = 1; + // Send from the GID the peer was told about. Its position depends on + // which other addresses the interface has; it is not always index 1. + attr.ah_attr.grh.sgid_index = source_gid_index; } int mask = IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.h b/mlx/distributed/jaccl/lib/jaccl/rdma.h index 7e96c67afb..317e3953b1 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.h +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.h @@ -175,6 +175,8 @@ struct Connection { ibv_cq* completion_queue; ibv_qp* queue_pair; Destination src; // holds the local information + // Index in the port's GID table of the GID advertised in `src`. + int source_gid_index = 1; Connection(ibv_context* ctx_); Connection(Connection&& c); From c9dcfb509535f0fa241162517649e92710ac74cc Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:53:38 -0700 Subject: [PATCH 08/19] docs(jaccl): describe the progress timeout Add MLX_JACCL_PROGRESS_TIMEOUT_MS and its JACCL_PROGRESS_TIMEOUT_MS alias to the list of environment variables, with the default and how to remove the limit. Add a section to the JACCL part of the distributed guide that tells what occurs when a peer does not answer: the call raises an error, the group is closed and each later call on it raises an error. The JACCL README now also says that the group cannot be opened again in the same process. Co-Authored-By: Claude Opus 5.5 --- docs/src/usage/distributed.rst | 22 ++++++++++++++++++++++ docs/src/usage/environment_variables.rst | 11 +++++++++++ mlx/distributed/jaccl/lib/README.md | 7 ++++--- 3 files changed, 37 insertions(+), 3 deletions(-) diff --git a/docs/src/usage/distributed.rst b/docs/src/usage/distributed.rst index 431aef7102..44acd3763f 100644 --- a/docs/src/usage/distributed.rst +++ b/docs/src/usage/distributed.rst @@ -378,6 +378,24 @@ concatenated in rank order. all_gather_factory=make_side_channel, ) +A peer that does not answer +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +JACCL waits for RDMA completions in a polling loop. When the peer stops or the +link fails, no completion arrives. A JACCL call that gets no completion for +longer than :envvar:`MLX_JACCL_PROGRESS_TIMEOUT_MS` (120 seconds by default) +raises an error. A call that gets a failed completion raises an error +immediately. + +After such an error the communication group is closed. JACCL destroys the queue +pairs and the completion queues of the group and deregisters its buffers, and +each later call on the group raises an error immediately. The group cannot be +opened again. Start a new process to communicate again. + +The time that a rank waits includes the time that its peer computes, for +example in :func:`mlx.core.distributed.recv`. Set the limit above the longest +such wait, or set it to ``0`` to remove the limit. + .. _nccl_section: Getting Started with NCCL @@ -550,6 +568,10 @@ the following: ["rdma_en3", "rdma_en4", "rdma_en5", null] ] +:envvar:`MLX_JACCL_PROGRESS_TIMEOUT_MS` is optional. It contains the longest time +in milliseconds that a call waits for an RDMA completion before it raises an +error and closes the group. The default is 120000 and 0 removes the limit. + NCCL ^^^^^ diff --git a/docs/src/usage/environment_variables.rst b/docs/src/usage/environment_variables.rst index 3584f4c22e..5e06f103d1 100644 --- a/docs/src/usage/environment_variables.rst +++ b/docs/src/usage/environment_variables.rst @@ -63,6 +63,17 @@ sets them. Prefer a ring topology for JACCL. This variable is enabled by its presence. ``JACCL_RING`` is accepted as a higher-priority alias. +.. envvar:: MLX_JACCL_PROGRESS_TIMEOUT_MS + + The longest time in milliseconds that a JACCL call waits for an RDMA + completion. The default is ``120000``. Set it to ``0`` or a negative value + to remove the limit. ``JACCL_PROGRESS_TIMEOUT_MS`` is accepted as a + higher-priority alias. When no completion arrives for this time, the call + raises an error and the group is closed: its queue pairs and completion + queues are destroyed, its buffers are deregistered, and each later call on + the group raises an error immediately. See the :ref:`JACCL section + `. + .. envvar:: MLX_WORLD_SIZE The total number of processes in an NCCL group. diff --git a/mlx/distributed/jaccl/lib/README.md b/mlx/distributed/jaccl/lib/README.md index 91f20d2c24..94120c6b05 100644 --- a/mlx/distributed/jaccl/lib/README.md +++ b/mlx/distributed/jaccl/lib/README.md @@ -109,9 +109,10 @@ variables: A call that gets no completion for this time, or that gets a failed work completion, fails with an exception. Before the caller sees the exception, the group destroys its queue pairs and completion queues and deregisters its -buffers. Every later call on that group throws immediately. Set the limit -above the longest time that a rank can wait for a peer, for example in `recv` -while the peer computes. +buffers. Every later call on that group throws immediately, and the group +cannot be opened again in the same process. Set the limit above the longest +time that a rank can wait for a peer, for example in `recv` while the peer +computes. ### Device File Format From fc71f7fbc4db05d754bddc7cd0dc5f21cef6bc68 Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:54:39 -0700 Subject: [PATCH 09/19] test(jaccl): add tests with simulated verbs tests/jaccl is a separate CMake project, in the same way as mlx/distributed/jaccl/lib/examples. The MLX build does not build it. It is not below mlx/, because the install rule of MLX copies each header below mlx/. The tests compile the mesh, ring and rdma headers with a stand-in for . A fake fabric completes a send when it meets a posted receive of the peer. It can lose or fail one completion and it can make the poll fail. - send_frame_tests: each posted frame holds its payload and then only zero bytes, for the mesh and the ring calls. One group does the ring all gather with different data on each rank (ml-explore/mlx #4443). - progress_guard_tests: a peer that does not answer in each mesh and ring call, a lost completion, a failed completion, a failed poll, calls after a failure, a refused destroy, waits below the limit, and the values of the environment variable. Each mode is one CTest test because the limit is read once in a process. Run them with ASan and UBSan, or add -DJACCL_TESTS_TSAN=ON for TSan: cmake -S tests/jaccl -B build/jaccl-tests cmake --build build/jaccl-tests ctest --test-dir build/jaccl-tests --output-on-failure The two configurations pass, 8 of 8 tests each. The tests do not use RDMA hardware. No workflow runs them. Co-Authored-By: Claude Opus 5.5 --- mlx/distributed/jaccl/lib/README.md | 20 + tests/jaccl/CMakeLists.txt | 35 + tests/jaccl/fake_rdma.cpp | 309 +++++++++ tests/jaccl/fake_rdma.h | 93 +++ tests/jaccl/fake_verbs/infiniband/verbs.h | 69 ++ tests/jaccl/fixture.h | 189 +++++ tests/jaccl/progress_guard_tests.cpp | 810 ++++++++++++++++++++++ tests/jaccl/send_frame_tests.cpp | 354 ++++++++++ 8 files changed, 1879 insertions(+) create mode 100644 tests/jaccl/CMakeLists.txt create mode 100644 tests/jaccl/fake_rdma.cpp create mode 100644 tests/jaccl/fake_rdma.h create mode 100644 tests/jaccl/fake_verbs/infiniband/verbs.h create mode 100644 tests/jaccl/fixture.h create mode 100644 tests/jaccl/progress_guard_tests.cpp create mode 100644 tests/jaccl/send_frame_tests.cpp diff --git a/mlx/distributed/jaccl/lib/README.md b/mlx/distributed/jaccl/lib/README.md index 94120c6b05..fc7bd84edb 100644 --- a/mlx/distributed/jaccl/lib/README.md +++ b/mlx/distributed/jaccl/lib/README.md @@ -332,6 +332,26 @@ class Config { } ``` +## Tests + +`tests/jaccl` in the MLX repository has tests that do not use RDMA hardware. +They compile the mesh, ring and RDMA headers with a small stand-in for +`` and simulate the verbs. They check that each send frame +holds its payload and then only zero bytes, and that a call fails and the group +is closed when a completion does not arrive or reports a failure. The directory +is a separate CMake project, so the MLX build does not build it: + +```bash +cmake -S tests/jaccl -B build/jaccl-tests +cmake --build build/jaccl-tests +ctest --test-dir build/jaccl-tests --output-on-failure +``` + +The tests use AddressSanitizer and UndefinedBehaviorSanitizer. Configure with +`-DJACCL_TESTS_TSAN=ON` to use ThreadSanitizer. The guard tests measure time: a +stalled call must fail between its limit of 250 ms and one second after it, so +a machine under heavy load can make them fail. + ## License JACCL is part of MLX and is released under the same license. diff --git a/tests/jaccl/CMakeLists.txt b/tests/jaccl/CMakeLists.txt new file mode 100644 index 0000000000..31ec33ca0e --- /dev/null +++ b/tests/jaccl/CMakeLists.txt @@ -0,0 +1,35 @@ +cmake_minimum_required(VERSION 3.24) +project(jaccl_tests LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +option(JACCL_TESTS_TSAN "Use ThreadSanitizer in the JACCL tests" OFF) + +if(JACCL_TESTS_TSAN) + set(SANITIZE -fsanitize=thread) +else() + set(SANITIZE -fsanitize=address,undefined) +endif() + +enable_testing() + +set(JACCL_LIB ${CMAKE_CURRENT_SOURCE_DIR}/../../mlx/distributed/jaccl/lib) + +# The tests use the JACCL headers and the fake verbs, not the jaccl library. +function(jaccl_test name) + add_executable(${name} ${name}.cpp fake_rdma.cpp) + target_include_directories(${name} PRIVATE fake_verbs ${JACCL_LIB}) + target_compile_options(${name} PRIVATE -O1 -g ${SANITIZE}) + target_link_options(${name} PRIVATE ${SANITIZE}) +endfunction() + +jaccl_test(send_frame_tests) +jaccl_test(progress_guard_tests) + +add_test(NAME send_frame COMMAND send_frame_tests) + +# The guard reads its limit once in a process, so each mode is one test. +foreach(mode checks disabled alias precedence negative not-a-number default) + add_test(NAME progress_guard_${mode} COMMAND progress_guard_tests ${mode}) +endforeach() diff --git a/tests/jaccl/fake_rdma.cpp b/tests/jaccl/fake_rdma.cpp new file mode 100644 index 0000000000..e38a00a823 --- /dev/null +++ b/tests/jaccl/fake_rdma.cpp @@ -0,0 +1,309 @@ +// Copyright © 2026 Apple Inc. + +#include +#include + +#include "fake_rdma.h" + +namespace { +std::vector bytes(const ibv_sge& entry) { + const auto* pointer = reinterpret_cast(entry.addr); + return {pointer, pointer + entry.length}; +} + +void check_unchanged(const fixture::Pending& pending) { + fixture::require( + bytes(pending.entry) == pending.original, + "Send buffer reused before its completion was polled"); +} + +void deliver( + fixture::Endpoint& endpoint, + int index, + int drop, + int fail, + fixture::Completion completion) { + if (index == drop) { + return; + } + if (index == fail) { + completion.value.status = IBV_WC_WR_FLUSH_ERR; + } + endpoint.completions.push_back(std::move(completion)); + auto& high_water = fixture::ledger().completion_high_water; + high_water = std::max(high_water.load(), endpoint.completions.size()); +} + +void drop(ibv_qp* pair) { + if (pair) { + delete pair; + fixture::ledger().queue_pairs--; + } +} + +void drop(ibv_cq* queue) { + if (queue) { + delete queue; + fixture::ledger().completion_queues--; + } +} + +void drop(ibv_mr* region) { + if (region) { + delete region; + fixture::ledger().regions--; + } +} + +// The three verbs that a failed group calls through jaccl::ibv(). +int destroy_queue_pair(ibv_qp* pair) { + fixture::ledger().destroy_queue_pair_calls++; + if (fixture::ledger().refuse_destroy) { + return 16; + } + auto& endpoint = *static_cast(pair->fixture); + { + // The posted work is cancelled. The peer cannot match this endpoint and + // its completions are gone. + std::lock_guard lock(endpoint.fabric->mutex); + endpoint.closed = true; + endpoint.sends.clear(); + endpoint.receives.clear(); + endpoint.completions.clear(); + } + drop(pair); + return 0; +} + +int destroy_completion_queue(ibv_cq* queue) { + fixture::ledger().destroy_completion_queue_calls++; + if (fixture::ledger().refuse_destroy) { + return 16; + } + drop(queue); + return 0; +} + +int deregister(ibv_mr* region) { + fixture::ledger().deregister_calls++; + drop(region); + return 0; +} +} // namespace + +namespace jaccl { +// Only allocation and registration are fake. The staging methods of +// SharedBuffer and the post and poll methods of Connection are the real ones. +SharedBuffer::SharedBuffer(size_t count) : data_(nullptr), num_bytes_(count) { + if (posix_memalign(&data_, 4096, count) != 0) { + throw std::bad_alloc(); + } + std::memset(data_, fixture::sentinel, count); +} + +SharedBuffer::SharedBuffer(SharedBuffer&& other) + : data_(nullptr), num_bytes_(0) { + std::swap(data_, other.data_); + std::swap(num_bytes_, other.num_bytes_); + memory_regions_.swap(other.memory_regions_); +} + +SharedBuffer::~SharedBuffer() { + for (const auto& entry : memory_regions_) { + drop(entry.second); + } + std::free(data_); +} + +void SharedBuffer::register_to_protection_domain(ibv_pd* domain) { + if (!memory_regions_.contains(domain)) { + memory_regions_[domain] = new ibv_mr{1}; + fixture::ledger().regions++; + } +} + +Connection::Connection(ibv_context* context) + : ctx(context), + protection_domain(nullptr), + completion_queue(nullptr), + queue_pair(nullptr), + src{} {} +Connection::Connection(Connection&& other) : Connection(nullptr) { + std::swap(ctx, other.ctx); + std::swap(protection_domain, other.protection_domain); + std::swap(completion_queue, other.completion_queue); + std::swap(queue_pair, other.queue_pair); + std::swap(src, other.src); + std::swap(source_gid_index, other.source_gid_index); +} + +Connection::~Connection() { + delete ctx; + delete protection_domain; + drop(completion_queue); + drop(queue_pair); +} + +// The real wrapper loads librdma.dylib. The fake one has only the three verbs +// that a failed group calls. +IBVWrapper::IBVWrapper() : librdma_handle_(nullptr) { + get_device_list = nullptr; + get_device_name = nullptr; + open_device = nullptr; + free_device_list = nullptr; + close_device = nullptr; + alloc_pd = nullptr; + create_qp = nullptr; + create_cq = nullptr; + dealloc_pd = nullptr; + query_port = nullptr; + query_gid = nullptr; + modify_qp = nullptr; + reg_mr = nullptr; + destroy_qp = destroy_queue_pair; + destroy_cq = destroy_completion_queue; + dereg_mr = deregister; +} + +IBVWrapper& ibv() { + static IBVWrapper wrapper; + return wrapper; +} +} // namespace jaccl + +namespace fixture { +Ledger& ledger() { + static Ledger value; + return value; +} + +Endpoint& endpoint(const jaccl::Connection& connection) { + return *static_cast(connection.queue_pair->fixture); +} + +void Fabric::pair( + jaccl::Connection& a, + jaccl::Connection& b, + int wire, + int side) { + auto first = std::make_unique(); + auto second = std::make_unique(); + first->fabric = second->fabric = this; + first->rank = 0; + second->rank = 1; + first->side = side; + second->side = 1 - side; + first->wire = second->wire = wire; + first->peer = second.get(); + second->peer = first.get(); + auto attach = [](jaccl::Connection& connection, Endpoint* endpoint) { + connection.ctx = new ibv_context{}; + connection.protection_domain = new ibv_pd{}; + connection.completion_queue = new ibv_cq{endpoint}; + connection.queue_pair = new ibv_qp{endpoint}; + ledger().completion_queues++; + ledger().queue_pairs++; + }; + attach(a, first.get()); + attach(b, second.get()); + endpoints.push_back(std::move(first)); + endpoints.push_back(std::move(second)); +} + +void Fabric::match(Endpoint& sender) { + auto& receiver = *sender.peer; + while (!sender.closed && !receiver.closed && !sender.sends.empty() && + !receiver.receives.empty()) { + auto send = std::move(sender.sends.front()); + sender.sends.pop_front(); + auto receive = std::move(receiver.receives.front()); + receiver.receives.pop_front(); + check_unchanged(send); + require( + send.entry.length == receive.entry.length, + "The posted send and receive lengths differ"); + std::memcpy( + reinterpret_cast(receive.entry.addr), + reinterpret_cast(send.entry.addr), + send.entry.length); + frames.push_back({sender.rank, sender.side, sender.wire, send.original}); + const ibv_wc sent{send.id, IBV_WC_SUCCESS, 0, send.entry.length}; + const ibv_wc received{receive.id, IBV_WC_SUCCESS, 0, receive.entry.length}; + deliver( + sender, + sender.sent++, + sender.fault.drop_send, + sender.fault.fail_send, + {sent, std::move(send)}); + deliver( + receiver, + receiver.received++, + receiver.fault.drop_receive, + receiver.fault.fail_receive, + {received, std::nullopt}); + } + changed.notify_all(); +} + +void Fabric::require_drained() { + std::lock_guard lock(mutex); + for (const auto& endpoint : endpoints) { + require( + endpoint->sends.empty() && endpoint->receives.empty() && + endpoint->completions.empty(), + "A posted request or a completion was not consumed"); + } +} + +void Fabric::expire_after(std::chrono::milliseconds duration) { + std::lock_guard lock(mutex); + deadline = std::chrono::steady_clock::now() + duration; +} +} // namespace fixture + +int ibv_post_send(ibv_qp* qp, ibv_send_wr* wr, ibv_send_wr**) { + fixture::require( + wr->num_sge == 1 && wr->next == nullptr && wr->opcode == IBV_WR_SEND && + wr->send_flags == IBV_SEND_SIGNALED, + "Unexpected send work request"); + auto& endpoint = *static_cast(qp->fixture); + std::lock_guard lock(endpoint.fabric->mutex); + endpoint.sends.push_back({wr->sg_list[0], wr->wr_id, bytes(wr->sg_list[0])}); + endpoint.fabric->match(endpoint); + return 0; +} + +int ibv_post_recv(ibv_qp* qp, ibv_recv_wr* wr, ibv_recv_wr**) { + fixture::require( + wr->num_sge == 1 && wr->next == nullptr, + "Unexpected receive work request"); + auto& endpoint = *static_cast(qp->fixture); + std::lock_guard lock(endpoint.fabric->mutex); + endpoint.receives.push_back({wr->sg_list[0], wr->wr_id, {}}); + endpoint.fabric->match(*endpoint.peer); + return 0; +} + +int ibv_poll_cq(ibv_cq* cq, int maximum, ibv_wc* out) { + auto& endpoint = *static_cast(cq->fixture); + std::unique_lock lock(endpoint.fabric->mutex); + fixture::require( + std::chrono::steady_clock::now() < endpoint.fabric->deadline, + "The fake poll reached its deadline"); + if (endpoint.fault.poll_error) { + return -1; + } + if (endpoint.completions.empty() && endpoint.fabric->idle_wait.count() > 0) { + endpoint.fabric->changed.wait_for(lock, endpoint.fabric->idle_wait); + } + int count = 0; + while (count < maximum && !endpoint.completions.empty()) { + auto completion = std::move(endpoint.completions.front()); + endpoint.completions.pop_front(); + if (completion.send) { + check_unchanged(*completion.send); + } + out[count++] = completion.value; + } + return count; +} diff --git a/tests/jaccl/fake_rdma.h b/tests/jaccl/fake_rdma.h new file mode 100644 index 0000000000..202d65229f --- /dev/null +++ b/tests/jaccl/fake_rdma.h @@ -0,0 +1,93 @@ +// Copyright © 2026 Apple Inc. + +// Simulated verbs for the tests. A send completes when it meets a posted +// receive of the peer endpoint. Each endpoint can lose or fail one completion. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "jaccl/rdma.h" + +namespace fixture { + +constexpr unsigned char sentinel = 0xa7; +struct Frame { + int rank, side, wire; + std::vector bytes; +}; + +struct Pending { + ibv_sge entry; + uint64_t id; + std::vector original; +}; + +struct Completion { + ibv_wc value; + std::optional send; +}; + +// The index, from 0, of the completion that an endpoint loses or fails. +struct Fault { + int drop_send = -1, fail_send = -1, drop_receive = -1, fail_receive = -1; + bool poll_error = false; +}; + +// The fake verbs objects that exist and the calls made through jaccl::ibv(). +struct Ledger { + std::atomic queue_pairs{0}, completion_queues{0}, regions{0}; + std::atomic destroy_queue_pair_calls{0}, + destroy_completion_queue_calls{0}; + std::atomic deregister_calls{0}; + std::atomic completion_high_water{0}; + std::atomic refuse_destroy{false}; +}; + +Ledger& ledger(); +struct Endpoint; +struct Fabric { + std::mutex mutex; + std::condition_variable changed; + std::vector> endpoints; + std::vector frames; + // The fake poll throws after this time, so that no test can run forever. + std::chrono::steady_clock::time_point deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(8); + // The sleep of an empty poll. With 0 the polling loop spins, as it does on + // hardware. + std::chrono::milliseconds idle_wait{1}; + void + pair(jaccl::Connection&, jaccl::Connection&, int wire, int first_side = 0); + void match(Endpoint&); + void require_drained(); + void expire_after(std::chrono::milliseconds); +}; + +struct Endpoint { + Fabric* fabric; + Endpoint* peer = nullptr; + int rank, side, wire; + std::deque sends, receives; + std::deque completions; + Fault fault; + int sent = 0, received = 0; + bool closed = false; +}; + +// Valid only while the connection has its queue pair. +Endpoint& endpoint(const jaccl::Connection&); +inline void require(bool value, const char* message) { + if (!value) { + throw std::runtime_error(message); + } +} +} // namespace fixture diff --git a/tests/jaccl/fake_verbs/infiniband/verbs.h b/tests/jaccl/fake_verbs/infiniband/verbs.h new file mode 100644 index 0000000000..a823b0b0a0 --- /dev/null +++ b/tests/jaccl/fake_verbs/infiniband/verbs.h @@ -0,0 +1,69 @@ +// Copyright © 2026 Apple Inc. + +// A stand-in for for the tests. It declares only the +// types that the JACCL headers use. fake_rdma.cpp defines the three verbs. + +#pragma once + +#include +#include + +struct ibv_device {}; +struct ibv_context {}; +struct ibv_pd {}; +struct ibv_comp_channel {}; +struct ibv_qp_init_attr {}; +struct ibv_port_attr {}; +struct ibv_qp_attr {}; +struct ibv_gid { + uint8_t raw[16]; +}; + +struct ibv_mr { + uint32_t lkey; +}; + +struct ibv_cq { + void* fixture; +}; + +struct ibv_qp { + void* fixture; +}; + +struct ibv_sge { + uint64_t addr; + uint32_t length; + uint32_t lkey; +}; + +struct ibv_send_wr { + uint64_t wr_id; + ibv_send_wr* next; + ibv_sge* sg_list; + int num_sge; + int opcode; + int send_flags; +}; + +struct ibv_recv_wr { + uint64_t wr_id; + ibv_recv_wr* next; + ibv_sge* sg_list; + int num_sge; +}; + +enum ibv_wc_status { IBV_WC_SUCCESS = 0, IBV_WC_WR_FLUSH_ERR = 5 }; +struct ibv_wc { + uint64_t wr_id; + ibv_wc_status status; + uint32_t vendor_err; + uint32_t byte_len; +}; + +constexpr int IBV_WR_SEND = 0; +constexpr int IBV_SEND_SIGNALED = 1; + +int ibv_post_send(ibv_qp*, ibv_send_wr*, ibv_send_wr**); +int ibv_post_recv(ibv_qp*, ibv_recv_wr*, ibv_recv_wr**); +int ibv_poll_cq(ibv_cq*, int, ibv_wc*); diff --git a/tests/jaccl/fixture.h b/tests/jaccl/fixture.h new file mode 100644 index 0000000000..9e1e5a48ae --- /dev/null +++ b/tests/jaccl/fixture.h @@ -0,0 +1,189 @@ +// Copyright © 2026 Apple Inc. + +// Two ranks that are connected through the simulated verbs, as a mesh or as a +// ring, and helpers to check the frames that went to the peer. + +#pragma once + +#include +#include + +#include "fake_rdma.h" +#include "jaccl/mesh_impl.h" +#include "jaccl/ring_impl.h" + +namespace fixture { + +struct StaleTail : std::runtime_error { + StaleTail() + : std::runtime_error("Nonzero bytes after the payload of a frame") {} +}; + +inline void zero_tail(const Frame& frame, size_t logical) { + require( + logical <= frame.bytes.size(), "The payload is larger than the frame"); + for (size_t index = logical; index < frame.bytes.size(); ++index) { + if (frame.bytes[index] != 0) { + throw StaleTail(); + } + } +} + +template +void both(First first, Second second) { + auto a = std::async(std::launch::async, first); + auto b = std::async(std::launch::async, second); + std::exception_ptr error; + try { + a.get(); + } catch (...) { + error = std::current_exception(); + } + try { + b.get(); + } catch (...) { + if (!error) { + error = std::current_exception(); + } + } + if (error) { + std::rethrow_exception(error); + } +} + +inline void connections(std::vector& values, int count) { + values.reserve(count); + for (int i = 0; i < count; ++i) { + values.emplace_back(nullptr); + } +} + +inline void buffers( + std::vector& values, + int slots, + const std::vector& first, + const std::vector* second = nullptr) { + values.reserve(BUFFER_SIZES * slots); + for (int size = 0; size < BUFFER_SIZES; ++size) { + for (int slot = 0; slot < slots; ++slot) { + values.emplace_back(FRAME_SIZE * (1 << size)); + for (const auto& connection : first) { + if (connection.ctx) { + values.back().register_to_protection_domain( + connection.protection_domain); + } + } + if (second) { + for (const auto& connection : *second) { + values.back().register_to_protection_domain( + connection.protection_domain); + } + } + } + } +} + +struct MeshPair { + Fabric fabric; + std::array, 2> links; + std::array, 2> scratch, scatter; + std::array, 2> nodes; + MeshPair() { + for (auto& rank : links) { + connections(rank, 2); + } + fabric.pair(links[0][1], links[1][0], 0); + for (int rank = 0; rank < 2; ++rank) { + buffers(scratch[rank], NUM_BUFFERS * 2, links[rank]); + buffers(scatter[rank], NUM_BUFFERS * 4, links[rank]); + nodes[rank] = std::make_unique( + rank, 2, links[rank], scratch[rank], scatter[rank]); + } + } +}; + +struct RingPair { + Fabric fabric; + int wires; + std::array, 2> left, right; + std::array, 2> sends, receives; + std::array, 2> pools; + std::array, 2> nodes; + explicit RingPair(int count) : wires(count) { + for (int rank = 0; rank < 2; ++rank) { + connections(left[rank], wires); + connections(right[rank], wires); + } + for (int wire = 0; wire < wires; ++wire) { + fabric.pair(right[0][wire], left[1][wire], wire, 0); + fabric.pair(left[0][wire], right[1][wire], wire, 1); + } + for (int rank = 0; rank < 2; ++rank) { + buffers(sends[rank], NUM_BUFFERS * wires * 2, left[rank], &right[rank]); + buffers( + receives[rank], NUM_BUFFERS * wires * 2, left[rank], &right[rank]); + if (wires > 1) { + pools[rank] = std::make_unique(wires - 1); + } + nodes[rank] = std::make_unique( + rank, + 2, + left[rank], + right[rank], + sends[rank], + receives[rank], + pools[rank].get()); + } + } +}; + +inline std::vector +selected(const Fabric& fabric, int rank, int side, int wire) { + std::vector result; + for (const auto& frame : fabric.frames) { + if (frame.rank == rank && frame.side == side && frame.wire == wire) { + result.push_back(frame); + } + } + return result; +} + +inline void payload_frames( + const std::vector& frames, + std::span input) { + const size_t capacity = buffer_size_from_message(input.size()).second; + require( + frames.size() == (input.size() + capacity - 1) / capacity, + "Unexpected number of frames"); + size_t offset = 0; + for (const auto& frame : frames) { + require(frame.bytes.size() == capacity, "Unexpected frame size"); + const size_t count = std::min(capacity, input.size() - offset); + require( + std::equal( + input.begin() + offset, + input.begin() + offset + count, + frame.bytes.begin()), + "Valid payload bytes changed"); + zero_tail(frame, count); + offset += count; + } +} + +struct Add { + void operator()(const float* source, float* target, int64_t count) const { + for (int64_t i = 0; i < count; ++i) { + target[i] += source[i]; + } + } + void operator()( + const float* source, + const float* base, + float* target, + int64_t count) const { + for (int64_t i = 0; i < count; ++i) { + target[i] = source[i] + base[i]; + } + } +}; +} // namespace fixture diff --git a/tests/jaccl/progress_guard_tests.cpp b/tests/jaccl/progress_guard_tests.cpp new file mode 100644 index 0000000000..bdd167fd8e --- /dev/null +++ b/tests/jaccl/progress_guard_tests.cpp @@ -0,0 +1,810 @@ +// Copyright © 2026 Apple Inc. + +// Tests of the progress guard with simulated verbs. Each call goes through +// the real mesh and ring headers. The limit is read once in a process, so each +// mode of this program sets the environment first and is one CTest test. + +#include +#include +#include +#include +#include +#include + +#include "fixture.h" +#include "jaccl/progress_guard.h" + +namespace { +using Clock = std::chrono::steady_clock; +using std::chrono::milliseconds; + +// The limit of this process in milliseconds. +int64_t limit = 0; +// A call may fail this much after the limit. The fake poll then stops it. +constexpr int64_t slack = 1000; +const char* const watchdog = "The fake poll reached its deadline"; +const char* const closed = + "[jaccl] The group is closed after an earlier failure and cannot be used again."; +const int64_t frame = FRAME_SIZE * (1 << (BUFFER_SIZES - 1)); + +struct Outcome { + bool threw = false; + std::string message; + int64_t ms = 0; +}; + +template +Outcome attempt(Call call) { + Outcome outcome; + const auto start = Clock::now(); + try { + call(); + } catch (const std::exception& error) { + outcome.threw = true; + outcome.message = error.what(); + } + outcome.ms = + std::chrono::duration_cast(Clock::now() - start).count(); + return outcome; +} + +template +std::array together(First first, Second second) { + auto a = std::async(std::launch::async, [&] { return attempt(first); }); + auto b = std::async(std::launch::async, [&] { return attempt(second); }); + Outcome x = a.get(), y = b.get(); + return {x, y}; +} + +// Print each kind of error once, as the caller sees it. +void show(bool& shown, const Outcome& outcome) { + if (!shown) { + std::cout << "INFO error: " << outcome.message << '\n'; + shown = true; + } +} + +[[noreturn]] void mismatch(const std::string& what, const Outcome& outcome) { + std::ostringstream text; + text << what << " (after " << outcome.ms << " ms: " + << (outcome.threw ? outcome.message : std::string("returned normally")) + << ")"; + throw std::runtime_error(text.str()); +} + +void expect_success(const Outcome& outcome) { + if (outcome.threw) { + mismatch("normal traffic failed", outcome); + } +} + +// The polling loop must stop by itself, with its own error, near the limit. +void expect_timeout(const Outcome& outcome, const std::string& op) { + if (!outcome.threw) { + mismatch("a stalled call returned without an error", outcome); + } + if (outcome.message == watchdog) { + mismatch( + "the polling loop did not stop by itself; the watchdog of the fake poll ended it", + outcome); + } + if (outcome.message.rfind("[jaccl] " + op + ": no completion for ", 0) != 0) { + mismatch("not the no-progress error of " + op, outcome); + } + if (outcome.ms < limit || outcome.ms > limit + slack) { + mismatch("the failure is outside [limit, limit + slack]", outcome); + } + static bool shown = false; + show(shown, outcome); +} + +// A failed completion must be an error at once; it must not wait for the limit. +void expect_failed_completion(const Outcome& outcome, const std::string& op) { + if (!outcome.threw) { + mismatch("a failed completion was accepted as data", outcome); + } + if (outcome.message.rfind( + "[jaccl] " + op + ": a work completion failed with status 5", 0) != + 0) { + mismatch("not the failed-completion error of " + op, outcome); + } + if (limit > 0 && outcome.ms >= limit / 2) { + mismatch("the failure was not immediate", outcome); + } + static bool shown = false; + show(shown, outcome); +} + +void expect_closed(const Outcome& outcome) { + if (!outcome.threw || outcome.message != closed) { + mismatch("a closed group accepted a call", outcome); + } + if (outcome.ms > 100) { + mismatch("a closed group did not refuse at once", outcome); + } + static bool shown = false; + show(shown, outcome); +} + +void arm(fixture::Fabric& fabric, int64_t ms = -1) { + fabric.idle_wait = milliseconds(0); + fabric.expire_after(milliseconds(ms < 0 ? limit + slack : ms)); +} + +// The failed rank destroyed its queue pairs and completion queues through the +// verbs and deregistered every region. `before` is the ledger before the call. +struct Counts { + int queue_pairs, completion_queues, regions, destroy_pairs, destroy_queues, + deregister; + static Counts now() { + const auto& l = fixture::ledger(); + return { + l.queue_pairs, + l.completion_queues, + l.regions, + l.destroy_queue_pair_calls, + l.destroy_completion_queue_calls, + l.deregister_calls}; + } +}; + +void expect_released( + const Counts& before, + std::initializer_list*> links, + int ranks_failed, + int ranks) { + int connections = 0; + for (const auto* group : links) { + for (const auto& connection : *group) { + if (connection.ctx == nullptr) { + continue; + } + connections++; + fixture::require( + connection.queue_pair == nullptr && + connection.completion_queue == nullptr, + "A failed group kept a queue pair or a completion queue"); + } + } + const Counts after = Counts::now(); + fixture::require( + after.destroy_pairs - before.destroy_pairs == connections && + after.destroy_queues - before.destroy_queues == connections && + before.queue_pairs - after.queue_pairs == connections && + before.completion_queues - after.completion_queues == connections, + "Queue pairs/completion queues were not destroyed exactly once"); + const int regions = before.regions * ranks_failed / ranks; + fixture::require( + regions > 0 && before.regions - after.regions == regions && + after.deregister - before.deregister == regions, + "The buffers of a failed group were not deregistered exactly once"); +} + +std::vector floats(size_t count, float value) { + return std::vector(count, value); +} + +const char* raw(const std::vector& values) { + return reinterpret_cast(values.data()); +} + +char* raw(std::vector& values) { + return reinterpret_cast(values.data()); +} + +using MeshCall = std::function; +using RingCall = std::function; +struct Buffers { + std::vector in = floats(4096, 1), out = floats(8192, -19); +}; + +std::vector> mesh_calls(Buffers& b) { + return { + {"mesh recv", [&](auto& node) { node.recv(raw(b.out), 100, 1); }}, + {"mesh send", [&](auto& node) { node.send(raw(b.in), 100, 1); }}, + {"mesh all_reduce", + [&](auto& node) { + node.all_reduce(b.in.data(), b.out.data(), 64, fixture::Add{}); + }}, + {"mesh all_gather", + [&](auto& node) { node.all_gather(raw(b.in), raw(b.out), 256); }}, + {"mesh sum_scatter", + [&](auto& node) { + node.sum_scatter(b.in.data(), b.out.data(), 64, fixture::Add{}); + }}, + // The scatter-gather all reduce starts with a reduce scatter. + {"mesh sum_scatter", + [&](auto& node) { + node.all_reduce_scatter_gather( + b.in.data(), b.out.data(), 64, fixture::Add{}); + }}, + }; +} + +std::vector> ring_calls(Buffers& b) { + return { + {"ring recv", + [&](auto& node, int wires) { node.recv(raw(b.out), 100, 1, wires); }}, + {"ring send", + [&](auto& node, int wires) { node.send(raw(b.in), 100, 1, wires); }}, + {"ring all_reduce/all_gather", + [&](auto& node, int wires) { + if (wires == 1) { + node.template all_reduce<1>( + b.in.data(), b.out.data(), 64, 1, fixture::Add{}); + } else { + node.template all_reduce<2>( + b.in.data(), b.out.data(), 64, wires, fixture::Add{}); + } + }}, + {"ring all_reduce/all_gather", + [&](auto& node, int wires) { + node.all_gather(raw(b.in), raw(b.out), 256, wires); + }}, + {"ring reduce_scatter", + [&](auto& node, int wires) { + node.reduce_scatter( + b.in.data(), b.out.data(), 128, wires, fixture::Add{}); + }}, + }; +} + +int failures = 0, passes = 0; +template +void check(const std::string& name, Check body) { + try { + body(); + std::cout << "PASS " << name << '\n'; + passes++; + } catch (const std::exception& error) { + std::cout << "FAIL " << name << ": " << error.what() << '\n'; + failures++; + } +} + +// A peer that never answers: every polling loop must stop by itself, release +// the group and refuse every later call. +void silent_peer() { + Buffers b; + const auto mesh = mesh_calls(b); + const char* mesh_names[] = { + "recv", + "send", + "all_reduce", + "all_gather", + "sum_scatter", + "all_reduce_scatter_gather"}; + for (size_t index = 0; index < mesh.size(); ++index) { + check( + std::string("silent peer: mesh ") + mesh_names[index] + + " stops, releases, then refuses", + [&] { + fixture::MeshPair pair; + arm(pair.fabric); + const Counts before = Counts::now(); + expect_timeout( + attempt([&] { mesh[index].second(*pair.nodes[0]); }), + mesh[index].first); + expect_released(before, {&pair.links[0]}, 1, 2); + for (const auto& later : mesh) { + expect_closed(attempt([&] { later.second(*pair.nodes[0]); })); + } + }); + } + const auto ring = ring_calls(b); + const char* ring_names[] = { + "recv", "send", "all_reduce", "all_gather", "reduce_scatter"}; + for (int wires : {1, 2}) { + for (size_t index = 0; index < ring.size(); ++index) { + check( + "silent peer: ring " + std::string(ring_names[index]) + " on " + + std::to_string(wires) + " wire(s) stops, releases, then refuses", + [&] { + fixture::RingPair pair(wires); + arm(pair.fabric); + const Counts before = Counts::now(); + expect_timeout( + attempt([&] { ring[index].second(*pair.nodes[0], wires); }), + ring[index].first); + expect_released(before, {&pair.left[0], &pair.right[0]}, 1, 2); + for (const auto& later : ring) { + expect_closed( + attempt([&] { later.second(*pair.nodes[0], wires); })); + } + }); + } + } +} + +// One completion is lost while the peer keeps working. +void lost_completion() { + check( + "lost completion: mesh all_reduce drain stops after the reduction finished", + [&] { + fixture::MeshPair pair; + arm(pair.fabric); + fixture::endpoint(pair.links[0][1]).fault.drop_send = 0; + std::array, 2> in{floats(16, 1), floats(16, 2)}, + out{floats(17, -19), floats(17, -19)}; + const Counts before = Counts::now(); + auto result = together( + [&] { + pair.nodes[0]->all_reduce( + in[0].data(), out[0].data(), 16, fixture::Add{}); + }, + [&] { + pair.nodes[1]->all_reduce( + in[1].data(), out[1].data(), 16, fixture::Add{}); + }); + expect_timeout(result[0], "mesh all_reduce"); + expect_success(result[1]); + for (const auto& values : out) { + fixture::require( + std::all_of( + values.begin(), + values.begin() + 16, + [](float v) { return v == 3; }) && + values[16] == -19, + "The reduction before the drain changed"); + } + expect_released(before, {&pair.links[0]}, 1, 2); + }); + check( + "lost completion: mesh recv stops in the middle of a five-frame transfer", + [&] { + fixture::MeshPair pair; + arm(pair.fabric); + fixture::endpoint(pair.links[1][0]).fault.drop_receive = 2; + std::vector in(frame * 5, 0x53), out(frame * 5, 0x29); + const Counts before = Counts::now(); + auto result = together( + [&] { pair.nodes[0]->send(in.data(), in.size(), 1); }, + [&] { pair.nodes[1]->recv(out.data(), out.size(), 0); }); + expect_success(result[0]); + expect_timeout(result[1], "mesh recv"); + expect_released(before, {&pair.links[1]}, 1, 2); + }); + // Wire 1 runs inline and wire 0 on the pool, so both orders are covered. + for (int stalled : {0, 1}) { + check( + "lost completion: ring send on 2 wires stops when wire " + + std::to_string(stalled) + " loses its completion", + [&] { + fixture::RingPair pair(2); + arm(pair.fabric); + // Rank 0 sends to rank 1 on its left connections (direction 1). + fixture::endpoint(pair.left[0][stalled]).fault.drop_send = 0; + std::vector in(200, 0x64), out(201, 0x29); + const Counts before = Counts::now(); + auto result = together( + [&] { pair.nodes[0]->send(in.data(), 200, 1, 2); }, + [&] { pair.nodes[1]->recv(out.data(), 200, 0, 2); }); + expect_timeout(result[0], "ring send"); + expect_success(result[1]); + fixture::require( + std::equal(in.begin(), in.end(), out.begin()) && + out.back() == 0x29, + "The receiver of the completed transfer changed"); + expect_released(before, {&pair.left[0], &pair.right[0]}, 1, 2); + }); + } +} + +// A completion with a failure status is an error, never data. +void failed_completion() { + check( + "failed completion: mesh recv throws at once and leaves the output alone", + [&] { + fixture::MeshPair pair; + arm(pair.fabric); + fixture::endpoint(pair.links[1][0]).fault.fail_receive = 0; + std::vector in(100, 0x53), out(100, 0x29); + const Counts before = Counts::now(); + auto result = together( + [&] { pair.nodes[0]->send(in.data(), 100, 1); }, + [&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + expect_success(result[0]); + expect_failed_completion(result[1], "mesh recv"); + fixture::require( + std::all_of( + out.begin(), out.end(), [](char v) { return v == 0x29; }), + "A failed completion was copied to the output"); + expect_released(before, {&pair.links[1]}, 1, 2); + Buffers b; + for (const auto& later : mesh_calls(b)) { + expect_closed(attempt([&] { later.second(*pair.nodes[1]); })); + } + }); + check("failed completion: mesh send throws at once", [&] { + fixture::MeshPair pair; + arm(pair.fabric); + fixture::endpoint(pair.links[0][1]).fault.fail_send = 0; + std::vector in(100, 0x53), out(100, 0x29); + auto result = together( + [&] { pair.nodes[0]->send(in.data(), 100, 1); }, + [&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + expect_failed_completion(result[0], "mesh send"); + expect_success(result[1]); + }); + Buffers b; + const auto mesh = mesh_calls(b); + for (size_t index : {size_t(2), size_t(3), size_t(4)}) { + check("failed completion: " + mesh[index].first + " throws at once", [&] { + fixture::MeshPair pair; + arm(pair.fabric); + fixture::endpoint(pair.links[0][1]).fault.fail_receive = 0; + Buffers peer; + const auto peer_calls = mesh_calls(peer); + auto result = together( + [&] { mesh[index].second(*pair.nodes[0]); }, + [&] { peer_calls[index].second(*pair.nodes[1]); }); + expect_failed_completion(result[0], mesh[index].first); + expect_success(result[1]); + }); + } + check("failed poll: mesh recv throws at once", [&] { + fixture::MeshPair pair; + arm(pair.fabric); + fixture::endpoint(pair.links[1][0]).fault.poll_error = true; + std::vector out(100, 0x29); + const Outcome outcome = + attempt([&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + if (!outcome.threw || + outcome.message.rfind( + "[jaccl] mesh recv: the completion queue could not be polled", 0) != + 0 || + outcome.ms > 100) { + mismatch("a failed poll was not an immediate error", outcome); + } + static bool shown = false; + show(shown, outcome); + }); + // The collectives poll several connections through the shared helpers. + check("failed poll: mesh all_reduce throws at once", [&] { + fixture::MeshPair pair; + arm(pair.fabric); + fixture::endpoint(pair.links[0][1]).fault.poll_error = true; + Buffers b; + const Counts before = Counts::now(); + const Outcome outcome = attempt([&] { + pair.nodes[0]->all_reduce(b.in.data(), b.out.data(), 64, fixture::Add{}); + }); + if (!outcome.threw || + outcome.message.rfind( + "[jaccl] mesh all_reduce: the completion queue could not be polled", + 0) != 0 || + outcome.ms > 100) { + mismatch("a failed poll was not an immediate error", outcome); + } + expect_released(before, {&pair.links[0]}, 1, 2); + }); + for (int side : {0, 1}) { + check( + std::string("failed poll: ring all_reduce throws at once (") + + (side ? "right" : "left") + " connection)", + [&] { + fixture::RingPair pair(1); + arm(pair.fabric); + fixture::endpoint(side ? pair.right[0][0] : pair.left[0][0]) + .fault.poll_error = true; + Buffers b; + const Counts before = Counts::now(); + const Outcome outcome = attempt([&] { + pair.nodes[0]->all_reduce<1>( + b.in.data(), b.out.data(), 64, 1, fixture::Add{}); + }); + if (!outcome.threw || + outcome.message.rfind( + "[jaccl] ring all_reduce/all_gather: the completion queue could not be polled", + 0) != 0 || + outcome.ms > 100) { + mismatch("a failed poll was not an immediate error", outcome); + } + expect_released(before, {&pair.left[0], &pair.right[0]}, 1, 2); + }); + } + // Two wires: one gets the failed completion, the other has no traffic. The + // idle wire must stop, the cause must be reported, and it must not take the + // limit. `failing` selects the wire; wire 1 runs inline, wire 0 on the pool. + for (int failing : {0, 1}) { + check( + "failed completion: ring recv on 2 wires, wire " + + std::to_string(failing) + + " fails, the idle wire stops and the cause is reported", + [&] { + fixture::RingPair pair(2); + arm(pair.fabric); + // Rank 1 receives from rank 0 on its right connections (direction 1). + fixture::endpoint(pair.right[1][failing]).fault.fail_receive = 0; + std::vector in(100, 0x64), out(100, 0x29); + const Counts before = Counts::now(); + auto result = together( + [&] { pair.nodes[0]->send_wire(in.data(), 100, 1, 50, failing); }, + [&] { pair.nodes[1]->recv(out.data(), 100, 0, 2); }); + expect_success(result[0]); + expect_failed_completion(result[1], "ring recv"); + fixture::require( + std::all_of( + out.begin(), out.end(), [](char v) { return v == 0x29; }), + "A failed completion was copied to the output"); + expect_released(before, {&pair.left[1], &pair.right[1]}, 1, 2); + Buffers b; + for (const auto& later : ring_calls(b)) { + expect_closed(attempt([&] { later.second(*pair.nodes[1], 2); })); + } + }); + } + check( + "failed completion: ring all_reduce on 2 wires fails on one rank, the peer then stops by itself", + [&] { + fixture::RingPair pair(2); + arm(pair.fabric); + fixture::endpoint(pair.left[0][0]).fault.fail_receive = 0; + std::array, 2> in{floats(64, 1), floats(64, 2)}, + out{floats(96, -19), floats(96, -19)}; + const Counts before = Counts::now(); + auto result = together( + [&] { + pair.nodes[0]->all_reduce<2>( + in[0].data(), out[0].data(), 64, 2, fixture::Add{}); + }, + [&] { + pair.nodes[1]->all_reduce<2>( + in[1].data(), out[1].data(), 64, 2, fixture::Add{}); + }); + expect_failed_completion(result[0], "ring all_reduce/all_gather"); + expect_timeout(result[1], "ring all_reduce/all_gather"); + expect_released( + before, + {&pair.left[0], &pair.right[0], &pair.left[1], &pair.right[1]}, + 2, + 2); + }); +} + +// When the verbs refuse to destroy, the failure is still reported, the call +// still fails, and the group still refuses later calls. +void refused_destroy() { + check( + "refused destroy: the error is printed, the handles stay for the destructor", + [&] { + fixture::MeshPair pair; + arm(pair.fabric); + fixture::endpoint(pair.links[1][0]).fault.poll_error = true; + std::vector out(100, 0x29); + std::ostringstream captured; + auto* previous = std::cerr.rdbuf(captured.rdbuf()); + fixture::ledger().refuse_destroy = true; + const Outcome outcome = + attempt([&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + const Outcome later = + attempt([&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + fixture::ledger().refuse_destroy = false; + std::cerr.rdbuf(previous); + if (!outcome.threw || + outcome.message.rfind("[jaccl] mesh recv: ", 0) != 0) { + mismatch("the call did not fail", outcome); + } + expect_closed(later); + fixture::require( + captured.str() == + "[jaccl] Could not destroy a queue pair (16)\n" + "[jaccl] Could not destroy a completion queue (16)\n", + "The refused destroy was not reported"); + fixture::require( + pair.links[1][0].queue_pair != nullptr && + pair.links[1][0].completion_queue != nullptr, + "A handle that was not destroyed was dropped"); + }); +} + +// Traffic that waits, but less than the limit, is not touched. +void below_limit() { + check("below the limit: a peer that starts late is waited for", [&] { + fixture::MeshPair pair; + arm(pair.fabric, 8000); + std::vector in(100, 0x53), out(101, 0x29); + auto result = together( + [&] { + std::this_thread::sleep_for(milliseconds(limit / 2)); + pair.nodes[0]->send(in.data(), 100, 1); + }, + [&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + expect_success(result[0]); + expect_success(result[1]); + fixture::require( + std::equal(in.begin(), in.end(), out.begin()) && out.back() == 0x29, + "Late transfer changed"); + pair.fabric.require_drained(); + }); + check( + "below the limit: three gaps longer than the limit in total, each one shorter", + [&] { + fixture::MeshPair pair; + arm(pair.fabric, 8000); + std::vector in(frame * 3), out(frame * 3 + 1, 0x29); + for (size_t index = 0; index < in.size(); ++index) { + in[index] = char(index / frame + 1); + } + auto result = together( + [&] { + for (int part = 0; part < 3; ++part) { + std::this_thread::sleep_for(milliseconds(limit / 2)); + pair.nodes[0]->send(in.data() + part * frame, frame, 1); + } + }, + [&] { pair.nodes[1]->recv(out.data(), frame * 3, 0); }); + expect_success(result[0]); + expect_success(result[1]); + fixture::require( + result[1].ms >= limit * 3 / 2, + "The receiver did not wait across the gaps"); + fixture::require( + std::equal(in.begin(), in.end(), out.begin()) && out.back() == 0x29, + "Paused transfer changed"); + pair.fabric.require_drained(); + }); + check( + "below the limit: ring all_reduce on 2 wires with a peer that starts late", + [&] { + fixture::RingPair pair(2); + arm(pair.fabric, 8000); + std::array, 2> in{floats(4096, 1), floats(4096, 2)}, + out{floats(4097, -19), floats(4097, -19)}; + auto result = together( + [&] { + pair.nodes[0]->all_reduce<2>( + in[0].data(), out[0].data(), 4096, 2, fixture::Add{}); + }, + [&] { + std::this_thread::sleep_for(milliseconds(limit / 2)); + pair.nodes[1]->all_reduce<2>( + in[1].data(), out[1].data(), 4096, 2, fixture::Add{}); + }); + expect_success(result[0]); + expect_success(result[1]); + for (const auto& values : out) { + fixture::require( + std::all_of( + values.begin(), + values.begin() + 4096, + [](float v) { return v == 3; }) && + values[4096] == -19, + "Late ring reduction changed"); + } + pair.fabric.require_drained(); + }); +} + +// With the limit removed the loops wait as before; a failed completion is +// still an error. +void disabled() { + check( + "disabled: a stalled mesh recv keeps waiting until the watchdog of the fake poll", + [&] { + fixture::MeshPair pair; + arm(pair.fabric, 1200); + std::vector out(100, 0x29); + const Outcome outcome = + attempt([&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + if (!outcome.threw || outcome.message != watchdog || + outcome.ms < 1200) { + mismatch("the loop stopped although the limit is removed", outcome); + } + }); + check( + "disabled: a stalled ring recv on 2 wires keeps waiting until the watchdog of the fake poll", + [&] { + fixture::RingPair pair(2); + arm(pair.fabric, 1200); + std::vector out(100, 0x29); + const Outcome outcome = + attempt([&] { pair.nodes[1]->recv(out.data(), 100, 0, 2); }); + if (!outcome.threw || outcome.message != watchdog || + outcome.ms < 1200) { + mismatch("the loop stopped although the limit is removed", outcome); + } + }); + check("disabled: a failed completion is still an immediate error", [&] { + fixture::MeshPair pair; + arm(pair.fabric, 8000); + fixture::endpoint(pair.links[1][0]).fault.fail_receive = 0; + std::vector in(100, 0x53), out(100, 0x29); + auto result = together( + [&] { pair.nodes[0]->send(in.data(), 100, 1); }, + [&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + expect_success(result[0]); + expect_failed_completion(result[1], "mesh recv"); + if (result[1].ms > 500) { + mismatch("the failure was not immediate", result[1]); + } + }); +} + +// Check the limit that the guard read from the environment of this process. +void value(int64_t expected) { + limit = expected; + check( + "the limit of this process is " + std::to_string(expected) + " ms", [&] { + if (jaccl::progress_timeout_ms() != expected) { + throw std::runtime_error( + "the guard read " + std::to_string(jaccl::progress_timeout_ms())); + } + }); +} + +// One stalled call that must stop at the limit of this process. +void stall() { + check("silent peer: mesh recv stops at the limit of this process", [&] { + fixture::MeshPair pair; + arm(pair.fabric); + std::vector out(100, 0x29); + expect_timeout( + attempt([&] { pair.nodes[1]->recv(out.data(), 100, 0); }), "mesh recv"); + }); +} + +} // namespace + +// One mode in each process: +// checks all cases with a limit of 250 ms +// disabled a limit of 0: a stalled call continues to wait +// alias the limit comes from MLX_JACCL_PROGRESS_TIMEOUT_MS +// precedence JACCL_PROGRESS_TIMEOUT_MS is used before the alias +// negative a negative value removes the limit +// not-a-number a value that is not a number keeps the default +// default no variable: the default +int main(int count, char** arguments) { + const char* name = "JACCL_PROGRESS_TIMEOUT_MS"; + const char* alias = "MLX_JACCL_PROGRESS_TIMEOUT_MS"; + const std::string mode = count == 2 ? arguments[1] : ""; + // The guard reads the variables when the first call starts. + unsetenv(name); + unsetenv(alias); + if (mode == "checks") { + setenv(name, "250", 1); + value(250); + silent_peer(); + lost_completion(); + failed_completion(); + refused_destroy(); + below_limit(); + } else if (mode == "disabled") { + setenv(name, "0", 1); + value(0); + disabled(); + } else if (mode == "alias") { + setenv(alias, "400", 1); + value(400); + stall(); + } else if (mode == "precedence") { + setenv(name, "300", 1); + setenv(alias, "60000", 1); + value(300); + stall(); + } else if (mode == "negative") { + setenv(name, "-1", 1); + value(-1); + } else if (mode == "not-a-number") { + setenv(name, "30s", 1); + value(jaccl::DEFAULT_PROGRESS_TIMEOUT_MS); + } else if (mode == "default") { + value(jaccl::DEFAULT_PROGRESS_TIMEOUT_MS); + } else { + std::cout << "FAIL usage: progress_guard_tests checks | disabled | alias | " + "precedence | negative | not-a-number | default\n"; + return 2; + } + const auto& ledger = fixture::ledger(); + check( + "every fake queue pair, completion queue and region was freed once", [&] { + fixture::require( + ledger.queue_pairs == 0 && ledger.completion_queues == 0 && + ledger.regions == 0, + "Fake verbs objects remain after the checks"); + }); + std::cout << (failures ? "FAIL " : "PASS ") << passes << " passed, " + << failures << " failed; limit " << limit + << " ms; simulated verbs\n"; + return failures ? 1 : 0; +} diff --git a/tests/jaccl/send_frame_tests.cpp b/tests/jaccl/send_frame_tests.cpp new file mode 100644 index 0000000000..46e420dd65 --- /dev/null +++ b/tests/jaccl/send_frame_tests.cpp @@ -0,0 +1,354 @@ +// Copyright © 2026 Apple Inc. + +// Tests of the send staging with simulated verbs. The scratch buffers start +// with a sentinel in each byte, and the fake endpoint keeps each posted frame. +// Each frame must hold the payload and then only zero bytes. + +#include +#include + +#include "fixture.h" +#include "jaccl/send_frame.h" + +namespace { + +void staging_bounds() { + alignas(double) std::array frame; + const std::array input{0x01020304, 0x05060708}; + frame.fill(char(fixture::sentinel)); + jaccl::stage_send_frame(std::span(frame), input.data(), 2); + fixture::require( + std::memcmp(frame.data(), input.data(), 8) == 0, + "Typed payload bytes changed"); + fixture::zero_tail( + {0, 0, 0, std::vector(frame.begin(), frame.end())}, 8); + jaccl::stage_send_frame( + std::span(frame), static_cast(nullptr), 0); + fixture::zero_tail( + {0, 0, 0, std::vector(frame.begin(), frame.end())}, 0); + for (int64_t count : + {int64_t(-1), int64_t(9), std::numeric_limits::max()}) { + frame.fill(char(fixture::sentinel)); + bool refused = false; + try { + jaccl::stage_send_frame(std::span(frame), input.data(), count); + } catch (const std::invalid_argument&) { + refused = true; + } + fixture::require( + refused && + std::all_of( + frame.begin(), + frame.end(), + [](char value) { + return static_cast(value) == fixture::sentinel; + }), + "Invalid count must refuse before any scratch write"); + } + bool refused = false; + try { + jaccl::stage_send_frame( + std::span(frame), static_cast(nullptr), 1); + } catch (const std::invalid_argument&) { + refused = true; + } + fixture::require(refused, "Null nonempty source refused"); +} + +void mesh_point_to_point() { + fixture::MeshPair pair; + const int64_t maximum = FRAME_SIZE * (1 << (BUFFER_SIZES - 1)); + for (int64_t count : + {int64_t(1), + int64_t(24), + int64_t(4096), + int64_t(17), + int64_t(4097), + maximum * 3 + 17}) { + pair.fabric.frames.clear(); + std::vector input(count, 0x53), output(count + 1, 0x29); + fixture::both( + [&] { pair.nodes[0]->send(input.data(), count, 1); }, + [&] { pair.nodes[1]->recv(output.data(), count, 0); }); + fixture::require( + std::equal(input.begin(), input.end(), output.begin()) && + output.back() == 0x29, + "Mesh receive changed payload or wrote past end"); + pair.fabric.require_drained(); + fixture::payload_frames(fixture::selected(pair.fabric, 0, 0, 0), input); + } +} + +void ring_point_to_point() { + const int64_t maximum = FRAME_SIZE * (1 << (BUFFER_SIZES - 1)); + for (int wires : {1, 2}) { + fixture::RingPair pair(wires); + for (int64_t count : + {int64_t(1), int64_t(4096), int64_t(17), maximum * 3 * wires + 13}) { + pair.fabric.frames.clear(); + std::vector input(count, 0x64), output(count + 1, 0x29); + fixture::both( + [&] { pair.nodes[0]->send(input.data(), count, 1, wires); }, + [&] { pair.nodes[1]->recv(output.data(), count, 0, wires); }); + fixture::require( + std::equal(input.begin(), input.end(), output.begin()) && + output.back() == 0x29, + "Ring receive changed payload or wrote past end"); + pair.fabric.require_drained(); + const size_t per_wire = (count + wires - 1) / wires; + for (int wire = 0; wire < wires; ++wire) { + const size_t start = std::min(size_t(count), wire * per_wire); + const size_t end = std::min(size_t(count), (wire + 1) * per_wire); + auto frames = fixture::selected(pair.fabric, 0, 1, wire); + fixture::payload_frames( + frames, std::span(input).subspan(start, end - start)); + } + } + } +} + +std::span byte_view( + const std::vector& values, + size_t count) { + return {reinterpret_cast(values.data()), count * sizeof(float)}; +} + +void all_equal(const std::vector& values, size_t count, float expected) { + fixture::require( + std::all_of( + values.begin(), + values.begin() + count, + [&](float value) { return value == expected; }), + "Collective logical output changed"); + fixture::require( + values[count] == -19, "Collective wrote outside logical output"); +} + +void mesh_collectives() { + const size_t count = + (FRAME_SIZE * (1 << (BUFFER_SIZES - 1)) * 3) / sizeof(float) + 7; + for (int operation : {0, 1, 2}) { + fixture::MeshPair pair; + std::array, 2> input, output; + for (int rank = 0; rank < 2; ++rank) { + input[rank].assign(count * 2, float(rank + 1)); + output[rank].assign(count * 2 + 1, -19); + } + auto run = [&](int rank) { + if (operation == 0) { + pair.nodes[rank]->all_reduce( + input[rank].data(), output[rank].data(), count, fixture::Add{}); + } else if (operation == 1) { + pair.nodes[rank]->all_gather( + reinterpret_cast(input[rank].data()), + reinterpret_cast(output[rank].data()), + count * sizeof(float)); + } else { + pair.nodes[rank]->sum_scatter( + input[rank].data(), output[rank].data(), count, fixture::Add{}); + } + }; + fixture::both([&] { run(0); }, [&] { run(1); }); + pair.fabric.require_drained(); + for (int rank = 0; rank < 2; ++rank) { + if (operation == 1) { + fixture::require( + std::all_of( + output[rank].begin(), + output[rank].begin() + count, + [](float value) { return value == 1; }), + "Gather rank0 payload changed"); + fixture::require( + std::all_of( + output[rank].begin() + count, + output[rank].begin() + count * 2, + [](float value) { return value == 2; }) && + output[rank].back() == -19, + "Gather rank1 payload or output bound changed"); + } else { + all_equal(output[rank], count, 3); + } + fixture::payload_frames( + fixture::selected(pair.fabric, rank, rank, 0), + byte_view(input[rank], count)); + } + } +} + +void ring_collectives() { + // A chunk of 1 gives slices with no payload. The long chunk fills each + // pipeline slot again. buffer_size_from_message selects the frame sizes. + const int64_t maximum = FRAME_SIZE * (1 << (BUFFER_SIZES - 1)); + for (bool scatter : {false, true}) { + for (int64_t chunk : + {int64_t(1), maximum * 3 / int64_t(sizeof(float)) * 4 + 7}) { + const int wires = 2; + fixture::RingPair pair(wires); + const int64_t count = scatter ? chunk * 2 : chunk; + const int64_t stage_chunk = scatter ? chunk : (count + 1) / 2; + const int64_t per_wire = (stage_chunk + 2 * wires - 1) / (2 * wires); + const int64_t capacity = + buffer_size_from_message(per_wire * sizeof(float)).second; + const int64_t elements = capacity / sizeof(float); + const int64_t steps = (per_wire + elements - 1) / elements; + std::array, 2> input, output; + for (int rank = 0; rank < 2; ++rank) { + input[rank].assign(count, float(rank + 1)); + output[rank].assign((scatter ? chunk : count) + 32, -19); + } + auto run = [&](int rank) { + if (scatter) { + pair.nodes[rank]->reduce_scatter( + input[rank].data(), + output[rank].data(), + count, + wires, + fixture::Add{}); + } else { + pair.nodes[rank]->all_reduce<2>( + input[rank].data(), + output[rank].data(), + count, + wires, + fixture::Add{}); + } + }; + fixture::both([&] { run(0); }, [&] { run(1); }); + pair.fabric.require_drained(); + for (int rank = 0; rank < 2; ++rank) { + all_equal(output[rank], scatter ? chunk : count, 3); + for (int direction = 0; direction < 2; ++direction) { + for (int wire = 0; wire < wires; ++wire) { + auto frames = fixture::selected(pair.fabric, rank, direction, wire); + const int passes = scatter ? 1 : 2; + fixture::require( + frames.size() == size_t(steps * passes), + "Ring completion/frame count changed"); + const int64_t offset = + direction * wires * per_wire + wire * per_wire; + const int64_t end = std::min(stage_chunk, offset + per_wire); + for (int pass = 0; pass < passes; ++pass) { + const int64_t source_offset = + (pass == 0 ? rank : 1 - rank) * stage_chunk; + const int64_t limit = scatter + ? end + : std::min(end, std::max(0, count - source_offset)); + for (int64_t step = 0; step < steps; ++step) { + const int64_t valid = std::max( + 0, std::min(elements, limit - offset - step * elements)); + const auto& frame = frames[pass * steps + step]; + fixture::require( + frame.bytes.size() == size_t(capacity), + "Ring posted SGE length changed"); + fixture::zero_tail(frame, valid * sizeof(float)); + } + } + } + } + } + } + } +} + +// Ring all gather with different data on each rank. A gather of equal data +// cannot find a region that was not transferred. +void ring_gather() { + const int64_t maximum = FRAME_SIZE * (1 << (BUFFER_SIZES - 1)); + for (int wires : {1, 2}) { + for (int64_t bytes : + {int64_t(1), + int64_t(7), + int64_t(4096), + int64_t(4097), + maximum * 3 * wires + 13}) { + fixture::RingPair pair(wires); + std::array, 2> input, output; + for (int rank = 0; rank < 2; ++rank) { + input[rank].resize(bytes); + for (int64_t index = 0; index < bytes; ++index) { + input[rank][index] = char((index * 31 + rank * 101 + 7) % 251); + } + output[rank].assign(bytes * 2 + 1, 0x29); + } + fixture::both( + [&] { + pair.nodes[0]->all_gather( + input[0].data(), output[0].data(), bytes, wires); + }, + [&] { + pair.nodes[1]->all_gather( + input[1].data(), output[1].data(), bytes, wires); + }); + pair.fabric.require_drained(); + for (int rank = 0; rank < 2; ++rank) { + for (int source = 0; source < 2; ++source) { + fixture::require( + std::equal( + input[source].begin(), + input[source].end(), + output[rank].begin() + source * bytes), + "Ring all gather lost or misplaced a rank region"); + } + fixture::require( + output[rank].back() == 0x29, + "Ring all gather wrote past the output"); + } + } + } +} + +// Run all groups, so that a failure does not hide the next ones. +bool group(void (*check)(), const char* name) { + try { + check(); + std::cout << "PASS " << name << '\n'; + return true; + } catch (const std::exception& error) { + std::cout << "FAIL " << name << ": " << error.what() << '\n'; + return false; + } +} + +} // namespace + +int main() { + bool passed = + group(staging_bounds, "bounded typed staging and zero-logical frame"); + passed &= group( + mesh_point_to_point, + "mesh short/full/reused/multiframe and actual receive"); + passed &= group( + ring_point_to_point, + "ring one/two-wire short/full/reused/multiframe receive"); + passed &= group( + mesh_collectives, "mesh gather/reduce/scatter valid payload and tail"); + passed &= group( + ring_collectives, "ring reduction/scatter full/partial/zero slices"); + passed &= + group(ring_gather, "ring all_gather distinct rank data, one/two wires"); + // The completion queues of the groups have MAX_SEND_WR + MAX_RECV_WR entries. + // Report how many of them this traffic used at the same time. + const size_t depth = fixture::ledger().completion_high_water; + const size_t capacity = MAX_SEND_WR + MAX_RECV_WR; + if (depth <= capacity) { + std::cout << "PASS completion queue high-water mark " << depth << " of " + << capacity << " entries\n"; + } else { + std::cout << "FAIL completion queue high-water mark " << depth + << " exceeds " << capacity << '\n'; + passed = false; + } + const auto& ledger = fixture::ledger(); + if (ledger.queue_pairs == 0 && ledger.completion_queues == 0 && + ledger.regions == 0) { + std::cout + << "PASS every fake queue pair, completion queue and region was freed once\n"; + } else { + std::cout << "FAIL fake verbs objects remain after the checks\n"; + passed = false; + } + if (!passed) { + return 1; + } + return 0; +} From d0f8b0a89ae70ceb0b8c4aaa2744d78483412f42 Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:57:09 -0700 Subject: [PATCH 10/19] docs(forkdiff): describe the JACCL changes and the upstream JACCL fixes Add two sections to fork.yaml so that check_forkdiff.py passes. The first one describes the fork work in JACCL: the zeroed send-frame tails, the progress guard, the advertised GID, the documentation of the new environment variable and the tests in tests/jaccl. The second one lists the cherry-picks of ml-explore/mlx #4443, #4557 and #4558 and the scoped_env helper of #4474 as upstream work. It also says why #4607 is not ported. base.hash does not change, because a cherry-pick does not move the merge-base. The overview now also names the JACCL changes, which are not gated, and counts three sections that are not fork work. Co-Authored-By: Claude Opus 5.5 --- fork.yaml | 90 +++++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 87 insertions(+), 3 deletions(-) diff --git a/fork.yaml b/fork.yaml index 3a72e886dd..961bc83c17 100644 --- a/fork.yaml +++ b/fork.yaml @@ -30,9 +30,16 @@ def: Metal eval. Every change is meant to be additive and rebasable; each is gated by shape, device or environment checks so the upstream path stays the default elsewhere. - One section is not fork work at all: the fork carries upstream's own v0.32.2 changes - (applied in PR #10 rather than merged), and those files are listed as such so the page - stays honest about what Layr-Labs actually changed. + The fork also changes JACCL, the RDMA-over-Thunderbolt backend that runs one model across + two Macs: it zeroes the unused tail of every send frame, and a progress guard fails a call + that gets no completion and closes its group. These two are not gated. The zeroing is + always on, and the guard is on with a 120 s limit that `JACCL_PROGRESS_TIMEOUT_MS=0` + removes. + + Three sections are not fork work at all: the fork carries upstream's own v0.32.2 changes + (applied in PR #10 rather than merged) and two sets of cherry-picked upstream fixes, and + those files are listed as such so the page stays honest about what Layr-Labs actually + changed. sub: - title: "Metal eval: use-after-free with mid-eval synchronization" description: | @@ -169,6 +176,54 @@ def: - "mlx/backend/metal/kernels/binary_ops.h" - "tests/ops_tests.cpp" - "python/tests/test_ops.py" + - title: "JACCL: zeroed send-frame tails, a progress guard and the advertised GID" + description: | + JACCL is the RDMA-over-Thunderbolt backend. Upstream has changed none of the three + things below. + + Send frames: JACCL posts every frame at the full size of its scratch buffer, so a + payload shorter than the buffer went to the peer followed by the bytes of earlier + sends. `stage_send_frame` (`send_frame.h`) copies the payload and zeroes the rest of + the frame, and every mesh and ring send path stages through + `SharedBuffer::stage_send`. The posted sizes and the receive paths are unchanged. + + Progress guard: the polling loops waited for completions without a limit and never + read `wc.status`, so a rank whose peer had died spun on a full core until it was + killed, and a failed completion was used as data. `ProgressGuard` + (`progress_guard.h`) is called after every poll of every mesh and ring loop. A + completion with a failure status, or a failed `ibv_poll_cq`, is an error at once; no + completion for longer than `JACCL_PROGRESS_TIMEOUT_MS` / + `MLX_JACCL_PROGRESS_TIMEOUT_MS` (default 120000 ms; `0` or a negative value removes + the limit) is an error. The clock is read once per 1024 polls that return nothing, so + the polling fast path gains a counter and no clock read. The call throws a `[jaccl]` + `std::runtime_error` rather than exiting the process. Before the caller sees it, the + group destroys its queue pairs and completion queues and deregisters its buffers; the + ring does so only after all wires have returned. Every later call on that group throws + immediately and the group cannot be reopened in the process. `dispatch_wires` also + reports the errors of the ring pool wires, which it used to drop. The completion-queue + size is unchanged. + + Advertised GID: `Connection::info()` advertises the IPv4-mapped GID of the port, but + `queue_pair_rtr` sent from GID index 1 whatever the index of that GID. The connection + now remembers the index where the advertised GID was found and sends from it; + behaviour is unchanged where that index is 1. + + `tests/jaccl` is a separate CMake project that the MLX build does not build. It + compiles the real mesh, ring and RDMA headers against simulated verbs and checks the + frame tails, the guard and the release path under ASan, UBSan and TSan. `ring_impl.h` + also carries the upstream ring `all_gather` fix described further down. + globs: + - "mlx/distributed/jaccl/lib/README.md" + - "mlx/distributed/jaccl/lib/jaccl/mesh_impl.h" + - "mlx/distributed/jaccl/lib/jaccl/progress_guard.h" + - "mlx/distributed/jaccl/lib/jaccl/rdma.cpp" + - "mlx/distributed/jaccl/lib/jaccl/rdma.h" + - "mlx/distributed/jaccl/lib/jaccl/ring_impl.h" + - "mlx/distributed/jaccl/lib/jaccl/send_frame.h" + - "docs/src/usage/distributed.rst" + - "docs/src/usage/environment_variables.rst" + - "tests/jaccl/*" + - "tests/jaccl/fake_verbs/infiniband/verbs.h" - title: "Upstream v0.32.2 sync carried in PR #10" description: | **Not fork work.** PR #10 moved the fork from an 0.32.0-era upstream main to official @@ -306,6 +361,35 @@ def: - "mlx/backend/metal/kernels/gather_mm_offsets.metal" - "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h" - "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h" + - title: "Upstream JACCL and distributed fixes (ml-explore/mlx #4443, #4557, #4558)" + description: | + **Not fork work.** Cherry-picks of upstream commits `7241f12e` + ([ml-explore/mlx#4443](https://github.com/ml-explore/mlx/pull/4443)), `42987b66` + ([ml-explore/mlx#4557](https://github.com/ml-explore/mlx/pull/4557)) and `8803e524` + ([ml-explore/mlx#4558](https://github.com/ml-explore/mlx/pull/4558)), applied without + changes. #4443: the JACCL ring `all_gather` gave both directions the same slice of each + rank's region, so the upper half of every peer's region was never transferred. The fix + is in `ring_impl.h`, which is listed in the JACCL section above, and comes with the + test `test_all_gather_rank_data`. #4557: `ReduceScatter::eval_cpu` asserted that it + had no inputs and then read `inputs[0]`, so a build with assertions stopped at every + `sum_scatter` (`mlx/backend/cpu/distributed.cpp`, listed in the v0.32.2 section). + #4558: a backend initialization that failed with `strict` off registered an empty + group, and a later `init` of that backend returned it, even with `strict` on. Its test + `test_strict_init_after_fallback` needs the `scoped_env` test helper, so that helper + alone is taken from `ec5f6ff7` + ([ml-explore/mlx#4474](https://github.com/ml-explore/mlx/pull/4474)) into + `python/tests/mlx_tests.py`, also listed in the v0.32.2 section. + + Upstream `f4232c7f` + ([ml-explore/mlx#4607](https://github.com/ml-explore/mlx/pull/4607), a second cable + for the JACCL mesh) is not ported: it is a feature that restructures the mesh into one + `MeshImpl` per wire and moves `dispatch_wires`, so the progress guard above has to be + reworked for it first. The section disappears once the fork rebases onto an upstream + commit at or past `8803e524`. + globs: + - "mlx/distributed/distributed.cpp" + - "python/tests/mlx_distributed_tests.py" + - "python/tests/ring_test_distributed.py" - title: "Fork tooling" description: | The fork deletes the upstream workflows that cannot run here: the PyPI release build and From 4241bc454332339a7fc1423c856523d4e0328b99 Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:14:44 -0700 Subject: [PATCH 11/19] style(jaccl): format tests/jaccl/CMakeLists.txt with cmake-format The Check Lint job runs cmake-format on all files and wants one argument per line in this foreach. No test changes. Co-Authored-By: Claude Opus 5.5 --- tests/jaccl/CMakeLists.txt | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/jaccl/CMakeLists.txt b/tests/jaccl/CMakeLists.txt index 31ec33ca0e..f3da64e2ae 100644 --- a/tests/jaccl/CMakeLists.txt +++ b/tests/jaccl/CMakeLists.txt @@ -30,6 +30,14 @@ jaccl_test(progress_guard_tests) add_test(NAME send_frame COMMAND send_frame_tests) # The guard reads its limit once in a process, so each mode is one test. -foreach(mode checks disabled alias precedence negative not-a-number default) +foreach( + mode + checks + disabled + alias + precedence + negative + not-a-number + default) add_test(NAME progress_guard_${mode} COMMAND progress_guard_tests ${mode}) endforeach() From b76589642c19cab803e27ed800f5cd7989f2a7db Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Fri, 9 Oct 2026 05:04:11 -0700 Subject: [PATCH 12/19] fix(jaccl): do the ring all_gather slice arithmetic in 64 bits The slice bounds that came with the upstream all_gather fix multiply lr and n_wires as int before the byte count widens the product. Both are small, so nothing overflows today, but the code scanner reports the pattern as a high-severity alert on the pull request. Cast the first factor so the whole product is int64_t. No behaviour change; the simulated-verbs tests pass. Co-Authored-By: Claude Opus 5.5 --- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index aa942d4283..776b51e2f8 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -232,10 +232,12 @@ class RingImpl { int64_t wire_offset[2]; int64_t wire_end[2]; for (int lr = 0; lr < 2; lr++) { - wire_offset[lr] = lr * n_wires * n_bytes_per_wire + + // 64-bit from the first factor: lr * n_wires would otherwise be an + // int product before it meets the byte count. + wire_offset[lr] = static_cast(lr) * n_wires * n_bytes_per_wire + static_cast(lw) * n_bytes_per_wire; - int64_t region_end = - std::min(n_bytes, (lr + 1) * n_wires * n_bytes_per_wire); + int64_t region_end = std::min( + n_bytes, static_cast(lr + 1) * n_wires * n_bytes_per_wire); wire_end[lr] = std::min(region_end, wire_offset[lr] + n_bytes_per_wire); } int64_t send_offset[2] = {rank_ * n_bytes, rank_ * n_bytes}; From ccf517a04bb168258009be33042b4acd677dc913 Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:28:16 -0700 Subject: [PATCH 13/19] Revert the NAX gather_qmm simdgroup skip (#4352) on this tree Upstream ml-explore/mlx c7ff35d97 ("Skip unnecessary simdgroup computations for quantised MOE matmuls on NAX", #4352) reached this fork through the v0.32.2 squash (734241bb), not by ancestry, in the form that #4572 later gave it: `const bool sg_active = sgp_sm > 0;` and two `if (sg_active)` blocks around the MMA loop in each of fp_gather_qmm_rhs_nax (fp_quantized_nax.h) and affine_gather_qmm_rhs_nax (quantized_nax.h). This removes the flag and runs the MMA in every simdgroup again, as before #4352. It is the same change as the earlier revert of #4352 on 0.32.2 (e9debbeaa) and its re-application on the 0.32.3 kernels, hand-ported because the surrounding lines differ here. Why: these kernels run only on GPUs with NAX. On a two-Mac JACCL pair with one NAX GPU, mixture-of-experts serving with the skip present lost all-reduce completions, and 92 to 96 GiB stayed wired on the host until it was rebooted. The builds that served that pair afterwards carried the revert. It is therefore kept out of every build that runs distributed on mixed hardware. Output is unchanged: stores stay bounded by sgp_sm, and the rows past the tile that the MMA now reads are either inside M (rows_in_bounds) or loaded as zeros by load_safe with sgp_sm rows. Co-Authored-By: Claude Opus 5.5 --- fork.yaml | 9 ++ mlx/backend/metal/kernels/fp_quantized_nax.h | 100 +++++++++---------- mlx/backend/metal/kernels/quantized_nax.h | 88 ++++++++-------- 3 files changed, 95 insertions(+), 102 deletions(-) diff --git a/fork.yaml b/fork.yaml index 961bc83c17..ccff2ca3c1 100644 --- a/fork.yaml +++ b/fork.yaml @@ -351,6 +351,15 @@ def: floating-point kernels have no global-scale input, and the NAX kernels keep the fork's bound for the last K tile. The section disappears once the fork rebases onto an upstream commit at or past `90089440`. + + One upstream change in these NAX kernels is reverted on purpose: the simdgroup skip + of [ml-explore/mlx#4352](https://github.com/ml-explore/mlx/pull/4352) (`c7ff35d9`, + which arrived with the v0.32.2 changes; `#4572` kept it as + `sg_active = sgp_sm > 0`). `fp_gather_qmm_rhs_nax` and `affine_gather_qmm_rhs_nax` run + the MMA in every simdgroup again. With the skip, mixture-of-experts serving on a + two-Mac JACCL pair with one NAX GPU lost all-reduce completions and left 92 to 96 GiB + wired until a reboot. Stores stay bounded by `sgp_sm`, so the output is unchanged. + Keep the revert at every rebase onto upstream. globs: - ".gitignore" - "benchmarks/python/swiglu_bench.py" diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 39fbfd0d54..c7f9f8fcc7 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -882,8 +882,6 @@ template < using AccumType = float; - const bool sg_active = sgp_sm > 0; - NAXTile Dtile; Dtile.clear(); @@ -913,35 +911,31 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); } + + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } xn += BK; @@ -955,32 +949,28 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short psk = min(int(SK), max(0, (BK - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 0921d007a3..bf90bb84fa 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -1561,8 +1561,6 @@ template < using AccumType = float; - const bool sg_active = sgp_sm > 0; - NAXTile Dtile; Dtile.clear(); @@ -1593,33 +1591,31 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); } + + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } xn += BK; @@ -1633,30 +1629,28 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; - volatile int compiler_barrier; + volatile int compiler_barrier; - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + const short psk = min(int(SK), max(0, (BK - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - (void)compiler_barrier; - } + (void)compiler_barrier; } } From fdae317492c37972f99b1d5e8cc1237c5338b4ec Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:42:17 -0700 Subject: [PATCH 14/19] fix(jaccl): 30 s progress limit, teardown exit by default, 4x completion queues A JACCL call that failed (no completion for the progress limit, a failed work completion, a failed poll) closed its group and threw. The thrown error reaches the program only through the stream that ran the collective, and a rank whose peer died sat in the guard for two minutes before it could even start to release its memory. RDMA work left outstanding on a dying process has been seen to keep tens of GiB wired until a reboot. - The default progress limit is 30000 ms (was 120000). - After the group is released (queue pairs and completion queues destroyed, buffers deregistered; the ring waits for all its wires), the default action is now a teardown exit: print the reason, run the memory release hook, exit with status 75 (EX_TEMPFAIL, restartable) through std::_Exit. JACCL_TIMEOUT_ACTION / MLX_JACCL_TIMEOUT_ACTION = throw keeps the previous behaviour; any other value keeps the default. Only the first failing thread exits; a second one waits for it. - MLX's jaccl.cpp registers the hook when it creates a group: set_wired_limit(0), then clear_cache(), each best effort. The hook gets 10 s, then the process exits anyway. - Completion queues have 4 * (MAX_SEND_WR + MAX_RECV_WR) entries, four times what a queue pair can have outstanding, so that a provider never sees a full queue. The queue pair attributes are not touched (the driver refuses changes to the RTR mask) and no global collective lock is added. Tests (tests/jaccl, simulated verbs, ASan/UBSan and TSan, all pass): the existing modes select the throw action and are otherwise unchanged. New teardown-exit mode runs each case as a child process and checks status 75, the reason, one release after the group was released and the time from the start of the call to the exit: mesh recv and all_reduce, ring recv on two wires, a failed completion (at once), no hook, a hook that never returns (cut off at 10 s), two groups failing together (one release, one exit), an unknown action, and the throw action by name and by alias. default-exit runs a stalled call with no variable set and checks the exit at 30 s. send_frame reports the completion high-water mark against the new depth. Co-Authored-By: Claude Opus 5.5 --- docs/src/usage/distributed.rst | 31 +- docs/src/usage/environment_variables.rst | 19 +- fork.yaml | 27 +- mlx/distributed/jaccl/jaccl.cpp | 39 +- mlx/distributed/jaccl/lib/README.md | 28 +- mlx/distributed/jaccl/lib/jaccl/mesh.cpp | 2 +- mlx/distributed/jaccl/lib/jaccl/mesh_impl.h | 8 +- .../jaccl/lib/jaccl/progress_guard.h | 132 +++++- mlx/distributed/jaccl/lib/jaccl/rdma.h | 4 + mlx/distributed/jaccl/lib/jaccl/ring.cpp | 4 +- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 7 +- tests/jaccl/CMakeLists.txt | 7 +- tests/jaccl/progress_guard_tests.cpp | 432 +++++++++++++++++- tests/jaccl/send_frame_tests.cpp | 16 +- 14 files changed, 690 insertions(+), 66 deletions(-) diff --git a/docs/src/usage/distributed.rst b/docs/src/usage/distributed.rst index 44acd3763f..43ffc865dd 100644 --- a/docs/src/usage/distributed.rst +++ b/docs/src/usage/distributed.rst @@ -383,14 +383,22 @@ A peer that does not answer JACCL waits for RDMA completions in a polling loop. When the peer stops or the link fails, no completion arrives. A JACCL call that gets no completion for -longer than :envvar:`MLX_JACCL_PROGRESS_TIMEOUT_MS` (120 seconds by default) -raises an error. A call that gets a failed completion raises an error -immediately. - -After such an error the communication group is closed. JACCL destroys the queue -pairs and the completion queues of the group and deregisters its buffers, and -each later call on the group raises an error immediately. The group cannot be -opened again. Start a new process to communicate again. +longer than :envvar:`MLX_JACCL_PROGRESS_TIMEOUT_MS` (30 seconds by default) +fails. A call that gets a failed completion fails immediately. + +A failed call closes the communication group: JACCL destroys the queue pairs +and the completion queues of the group and deregisters its buffers. Then, by +default, it releases the wired GPU memory of the process (it sets the wired +limit to zero and clears the buffer cache) and exits the process with status +75, so that a supervisor can restart it. This teardown exit is the default +because a rank that dies with RDMA work outstanding can leave its memory wired +until the machine restarts, and because the call often runs on a thread where +an exception cannot reach the program. + +With :envvar:`MLX_JACCL_TIMEOUT_ACTION` set to ``throw`` the call raises an +error instead and the process continues. Each later call on the closed group +raises an error immediately, and the group cannot be opened again. Start a new +process to communicate again. The time that a rank waits includes the time that its peer computes, for example in :func:`mlx.core.distributed.recv`. Set the limit above the longest @@ -569,8 +577,11 @@ the following: ] :envvar:`MLX_JACCL_PROGRESS_TIMEOUT_MS` is optional. It contains the longest time -in milliseconds that a call waits for an RDMA completion before it raises an -error and closes the group. The default is 120000 and 0 removes the limit. +in milliseconds that a call waits for an RDMA completion before it fails and +closes the group. The default is 30000 and 0 removes the limit. +:envvar:`MLX_JACCL_TIMEOUT_ACTION` is optional. It selects what a failed call +does after it closed the group: ``teardown-exit`` (the default) releases the +GPU memory and exits the process with status 75, ``throw`` raises an error. NCCL diff --git a/docs/src/usage/environment_variables.rst b/docs/src/usage/environment_variables.rst index 5e06f103d1..ec537b8b2f 100644 --- a/docs/src/usage/environment_variables.rst +++ b/docs/src/usage/environment_variables.rst @@ -66,14 +66,25 @@ sets them. .. envvar:: MLX_JACCL_PROGRESS_TIMEOUT_MS The longest time in milliseconds that a JACCL call waits for an RDMA - completion. The default is ``120000``. Set it to ``0`` or a negative value + completion. The default is ``30000``. Set it to ``0`` or a negative value to remove the limit. ``JACCL_PROGRESS_TIMEOUT_MS`` is accepted as a higher-priority alias. When no completion arrives for this time, the call - raises an error and the group is closed: its queue pairs and completion - queues are destroyed, its buffers are deregistered, and each later call on - the group raises an error immediately. See the :ref:`JACCL section + fails and the group is closed: its queue pairs and completion queues are + destroyed and its buffers are deregistered. What follows is set by + :envvar:`MLX_JACCL_TIMEOUT_ACTION`. See the :ref:`JACCL section `. +.. envvar:: MLX_JACCL_TIMEOUT_ACTION + + What a JACCL call does after it failed (no completion within + :envvar:`MLX_JACCL_PROGRESS_TIMEOUT_MS`, a failed completion, or a failed + poll) and closed its group. ``teardown-exit``, the default, sets the wired + limit to zero, clears the buffer cache and exits the process with status 75 + without running static destructors. ``throw`` raises an error; each later + call on the group then raises an error immediately. Any other value keeps + the default. ``JACCL_TIMEOUT_ACTION`` is accepted as a higher-priority + alias. + .. envvar:: MLX_WORLD_SIZE The total number of processes in an NCCL group. diff --git a/fork.yaml b/fork.yaml index ccff2ca3c1..024b4ab920 100644 --- a/fork.yaml +++ b/fork.yaml @@ -32,9 +32,10 @@ def: The fork also changes JACCL, the RDMA-over-Thunderbolt backend that runs one model across two Macs: it zeroes the unused tail of every send frame, and a progress guard fails a call - that gets no completion and closes its group. These two are not gated. The zeroing is - always on, and the guard is on with a 120 s limit that `JACCL_PROGRESS_TIMEOUT_MS=0` - removes. + that gets no completion, closes its group and by default releases the GPU memory and exits + the process with status 75. These two are not gated. The zeroing is always on, and the + guard is on with a 30 s limit that `JACCL_PROGRESS_TIMEOUT_MS=0` removes; + `JACCL_TIMEOUT_ACTION=throw` keeps the process alive and throws instead. Three sections are not fork work at all: the fork carries upstream's own v0.32.2 changes (applied in PR #10 rather than merged) and two sets of cherry-picked upstream fixes, and @@ -193,15 +194,19 @@ def: (`progress_guard.h`) is called after every poll of every mesh and ring loop. A completion with a failure status, or a failed `ibv_poll_cq`, is an error at once; no completion for longer than `JACCL_PROGRESS_TIMEOUT_MS` / - `MLX_JACCL_PROGRESS_TIMEOUT_MS` (default 120000 ms; `0` or a negative value removes + `MLX_JACCL_PROGRESS_TIMEOUT_MS` (default 30000 ms; `0` or a negative value removes the limit) is an error. The clock is read once per 1024 polls that return nothing, so - the polling fast path gains a counter and no clock read. The call throws a `[jaccl]` - `std::runtime_error` rather than exiting the process. Before the caller sees it, the - group destroys its queue pairs and completion queues and deregisters its buffers; the - ring does so only after all wires have returned. Every later call on that group throws - immediately and the group cannot be reopened in the process. `dispatch_wires` also - reports the errors of the ring pool wires, which it used to drop. The completion-queue - size is unchanged. + the polling fast path gains a counter and no clock read. The failed call first + destroys the group's queue pairs and completion queues and deregisters its buffers; + the ring does so only after all wires have returned. Then, by default + (`JACCL_TIMEOUT_ACTION` / `MLX_JACCL_TIMEOUT_ACTION` = `teardown-exit`), it prints the + reason, runs the memory release hook that MLX's `jaccl.cpp` registers + (`set_wired_limit(0)`, then `clear_cache()`, at most 10 s) and leaves through + `std::_Exit(75)`, a restartable status. With `throw` it throws a `[jaccl]` + `std::runtime_error` instead; every later call on that group throws immediately and + the group cannot be reopened in the process. `dispatch_wires` also reports the errors + of the ring pool wires, which it used to drop. Each completion queue has four times + the entries its queue pair can have outstanding (`COMPLETION_QUEUE_DEPTH`). Advertised GID: `Connection::info()` advertises the IPv4-mapped GID of the port, but `queue_pair_rtr` sent from GID index 1 whatever the index of that GID. The connection diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp index 99a4ac3be3..e13fc549f5 100644 --- a/mlx/distributed/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/jaccl.cpp @@ -4,9 +4,11 @@ #include "mlx/backend/cpu/encoder.h" #include "mlx/distributed/distributed_impl.h" #include "mlx/dtype_utils.h" +#include "mlx/memory.h" #include #include +#include #include #include @@ -168,6 +170,31 @@ class JACCLGroup : public GroupImpl { std::shared_ptr<::jaccl::Group> group_; }; +/** + * Run by a JACCL teardown exit (see progress_guard.h) after the group released + * its queue pairs. The exit skips static destructors, so release the wired GPU + * memory here: unwire first (the residency set shrinks to nothing), then drop + * the buffer cache. Best effort; it never throws into the exit path. + */ +void release_memory_before_exit() { + try { + set_wired_limit(0); + } catch (...) { + } + try { + clear_cache(); + } catch (...) { + } +} + +std::shared_ptr make_group(std::shared_ptr<::jaccl::Group> group) { + if (group == nullptr) { + return nullptr; + } + ::jaccl::set_memory_release(&release_memory_before_exit); + return std::make_shared(std::move(group)); +} + } // namespace bool is_available() { @@ -175,19 +202,11 @@ bool is_available() { } std::shared_ptr init(bool strict /* = false */) { - auto group = ::jaccl::init(strict); - if (group == nullptr) { - return nullptr; - } - return std::make_shared(std::move(group)); + return make_group(::jaccl::init(strict)); } std::shared_ptr init(bool strict, AllGatherFactory factory) { - auto group = ::jaccl::init(strict, factory); - if (group == nullptr) { - return nullptr; - } - return std::make_shared(std::move(group)); + return make_group(::jaccl::init(strict, factory)); } } // namespace mlx::core::distributed::jaccl diff --git a/mlx/distributed/jaccl/lib/README.md b/mlx/distributed/jaccl/lib/README.md index fc7bd84edb..92d80c3199 100644 --- a/mlx/distributed/jaccl/lib/README.md +++ b/mlx/distributed/jaccl/lib/README.md @@ -104,15 +104,25 @@ variables: mesh - **JACCL_PROGRESS_TIMEOUT_MS** / **MLX_JACCL_PROGRESS_TIMEOUT_MS**: (Optional) The longest time in milliseconds that a call waits for an RDMA completion. - The default is 120000. Zero or a negative value removes the limit. - -A call that gets no completion for this time, or that gets a failed work -completion, fails with an exception. Before the caller sees the exception, the -group destroys its queue pairs and completion queues and deregisters its -buffers. Every later call on that group throws immediately, and the group -cannot be opened again in the same process. Set the limit above the longest -time that a rank can wait for a peer, for example in `recv` while the peer -computes. + The default is 30000. Zero or a negative value removes the limit. +- **JACCL_TIMEOUT_ACTION** / **MLX_JACCL_TIMEOUT_ACTION**: (Optional) What a + failed call does: `teardown-exit` (the default) or `throw`. Any other value + keeps the default. + +A call fails when it gets no completion for this time, a failed work +completion or a failed poll. It first destroys the queue pairs and completion +queues of the group and deregisters its buffers (the ring waits until all of +its wires have returned). With `teardown-exit` it then prints the reason, runs +the memory release hook registered with `jaccl::set_memory_release` (MLX +registers one that sets the wired limit to zero and clears the buffer cache; +it gets at most 10 seconds) and exits the process with status 75 (EX_TEMPFAIL, +restartable) through `std::_Exit`. With `throw` the caller gets the exception; +every later call on that group throws immediately, and the group cannot be +opened again in the same process. Set the limit above the longest time that a +rank can wait for a peer, for example in `recv` while the peer computes. + +Each completion queue has `4 * (MAX_SEND_WR + MAX_RECV_WR)` entries, four times +what its queue pair can have outstanding. ### Device File Format diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp index 5b749f2bd9..c5ee1aa5a5 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp @@ -38,7 +38,7 @@ void MeshGroup::initialize() { continue; } conn.allocate_protection_domain(); - conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); + conn.create_completion_queue(COMPLETION_QUEUE_DEPTH); conn.create_queue_pair(); } diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index bbe99dcb50..fd8e21ca15 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -538,7 +538,8 @@ class MeshImpl { private: // One call into the group. It refuses a closed group, and it closes the - // group when the call ends with an exception. + // group when the call ends with an exception. Unless the failure action is + // Throw, it then leaves the process through teardown_exit(). class Call : public ProgressGuard { public: Call(MeshImpl& mesh, const char* op) @@ -551,6 +552,11 @@ class MeshImpl { ~Call() { if (std::uncaught_exceptions() > exceptions_) { mesh_.close(); + if (failure_action() == FailureAction::TeardownExit) { + teardown_exit( + failure().empty() ? "[jaccl] A mesh call failed with an error." + : failure().c_str()); + } } } diff --git a/mlx/distributed/jaccl/lib/jaccl/progress_guard.h b/mlx/distributed/jaccl/lib/jaccl/progress_guard.h index 6a18c549cc..a2009dd972 100644 --- a/mlx/distributed/jaccl/lib/jaccl/progress_guard.h +++ b/mlx/distributed/jaccl/lib/jaccl/progress_guard.h @@ -8,14 +8,18 @@ #include #include #include +#include #include +#include +#include #include #include #include +#include namespace jaccl { -constexpr int64_t DEFAULT_PROGRESS_TIMEOUT_MS = 120000; +constexpr int64_t DEFAULT_PROGRESS_TIMEOUT_MS = 30000; // The clock is read once per this number of polls without a completion. constexpr int64_t PROGRESS_SAMPLE_POLLS = 1024; @@ -45,6 +49,118 @@ inline int64_t progress_timeout_ms() { return timeout; } +// What a call does after it failed and closed its group. +enum class FailureAction { + // Release the group, run the memory release hook, exit with status 75. + TeardownExit, + // Throw to the caller; the process continues with the group closed. + Throw, +}; + +// EX_TEMPFAIL: the failure is temporary and the process can be restarted. +constexpr int TEARDOWN_EXIT_STATUS = 75; +// The memory release hook gets this long, then the process exits anyway. +constexpr int64_t TEARDOWN_RELEASE_LIMIT_MS = 10000; + +/** + * The action after a failed call. It is read once from JACCL_TIMEOUT_ACTION / + * MLX_JACCL_TIMEOUT_ACTION: "teardown-exit" (the default) or "throw". Any + * other value keeps the default. + */ +inline FailureAction failure_action() { + static const FailureAction action = []() { + const char* value = std::getenv("JACCL_TIMEOUT_ACTION"); + if (value == nullptr) { + value = std::getenv("MLX_JACCL_TIMEOUT_ACTION"); + } + if (value != nullptr && std::strcmp(value, "throw") == 0) { + return FailureAction::Throw; + } + return FailureAction::TeardownExit; + }(); + return action; +} + +// Releases the memory of the embedding library before a teardown exit, for +// example the wired GPU buffers that would otherwise stay wired after the +// process is gone. It must not throw. +using MemoryRelease = void (*)(); + +inline std::atomic& memory_release_hook() { + static std::atomic hook{nullptr}; + return hook; +} + +// Register the hook that teardown_exit() runs. The last registration wins. +inline void set_memory_release(MemoryRelease release) { + memory_release_hook().store(release); +} + +/** + * Leave the process after a failed call when the action is teardown-exit. + * The caller has already destroyed the queue pairs and completion queues of + * the group and deregistered its buffers. This prints the reason, runs the + * memory release hook (at most TEARDOWN_RELEASE_LIMIT_MS), then exits with + * TEARDOWN_EXIT_STATUS without running atexit handlers or static destructors. + * Only the first caller does so; a thread that fails at the same time waits + * for that exit. + */ +[[noreturn]] inline void teardown_exit(const char* reason) noexcept { + static std::atomic exiting{false}; + if (exiting.exchange(true)) { + for (;;) { + std::this_thread::sleep_for(std::chrono::seconds(1)); + } + } + std::fprintf( + stderr, + "%s\n[jaccl] The group is released; releasing memory, then " + "exiting with status %d (JACCL_TIMEOUT_ACTION=teardown-exit).\n", + reason, + TEARDOWN_EXIT_STATUS); + std::fflush(stderr); + MemoryRelease release = memory_release_hook().load(); + if (release != nullptr) { + try { + std::thread([] { + std::this_thread::sleep_for( + std::chrono::milliseconds(TEARDOWN_RELEASE_LIMIT_MS)); + std::fprintf( + stderr, + "[jaccl] The memory release did not finish in %lld ms; exiting.\n", + static_cast(TEARDOWN_RELEASE_LIMIT_MS)); + std::fflush(stderr); + std::_Exit(TEARDOWN_EXIT_STATUS); + }).detach(); + } catch (...) { + // Without the time limit the release still runs. + } + try { + release(); + } catch (...) { + // The hook must not throw; the exit goes ahead. + } + } + std::fprintf( + stderr, + "[jaccl] %s; exiting with status %d.\n", + release != nullptr ? "Memory released" : "No memory release registered", + TEARDOWN_EXIT_STATUS); + std::fflush(stderr); + std::_Exit(TEARDOWN_EXIT_STATUS); +} + +// The text of an exception, for teardown_exit(). +inline std::string describe(const std::exception_ptr& error) { + try { + std::rethrow_exception(error); + } catch (const std::exception& e) { + return e.what(); + } catch (...) { + return "[jaccl] A call failed with an exception that is not std::exception."; + } +} + // Thrown in a wire when a different wire of the same call failed first. struct WireStopped : std::runtime_error { WireStopped() @@ -66,7 +182,9 @@ inline void require_open(const std::atomic& failed) { * * check() throws when a completion reports a failure, when the poll fails or * when no completion arrived for longer than the timeout. It sets `failed` - * before it throws, which also stops the other wires of the same call. + * before it throws, which also stops the other wires of the same call. The + * owner of the group closes it and then, unless failure_action() is Throw, + * calls teardown_exit(). */ class ProgressGuard { public: @@ -92,6 +210,12 @@ class ProgressGuard { } } + protected: + // The error that check() threw, or empty. + const std::string& failure() const { + return failure_; + } + private: void check_idle() { if (failed_.load()) { @@ -122,7 +246,8 @@ class ProgressGuard { std::ostringstream msg; msg << "[jaccl] " << op_ << ": " << reason << ". The group is closed and cannot be used again."; - throw std::runtime_error(msg.str()); + failure_ = msg.str(); + throw std::runtime_error(failure_); } const char* op_; @@ -130,6 +255,7 @@ class ProgressGuard { int64_t timeout_ms_; int64_t idle_polls_{0}; std::chrono::steady_clock::time_point idle_since_; + std::string failure_; }; } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.h b/mlx/distributed/jaccl/lib/jaccl/rdma.h index 317e3953b1..f2023ca515 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.h +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.h @@ -20,6 +20,10 @@ constexpr int SEND_WR = 1; constexpr int RECV_WR = 2; constexpr int MAX_SEND_WR = 32; constexpr int MAX_RECV_WR = 32; +// Each completion queue has four times the entries its queue pair can have +// outstanding. A queue sized exactly to the outstanding work has no headroom, +// and a provider may drop a completion on a full queue without an error. +constexpr int COMPLETION_QUEUE_DEPTH = 4 * (MAX_SEND_WR + MAX_RECV_WR); constexpr int BUFFER_SIZES = 8; constexpr int NUM_BUFFERS = 2; constexpr int FRAME_SIZE = 4096; diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.cpp b/mlx/distributed/jaccl/lib/jaccl/ring.cpp index a08035f3f3..8b45a9dea9 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/ring.cpp @@ -41,12 +41,12 @@ void RingGroup::initialize() { // Create the queue pairs for (auto& conn : left_) { conn.allocate_protection_domain(); - conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); + conn.create_completion_queue(COMPLETION_QUEUE_DEPTH); conn.create_queue_pair(); } for (auto& conn : right_) { conn.allocate_protection_domain(); - conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); + conn.create_completion_queue(COMPLETION_QUEUE_DEPTH); conn.create_queue_pair(); } diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index 776b51e2f8..d2c6f9b957 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -801,7 +801,9 @@ class RingImpl { // inline, then wait for the pool calls before returning. // // A wire that throws sets failed_, which stops the other wires. When all - // wires returned the group is closed and the first error is thrown. + // wires returned the group is closed and the first error is thrown, or, + // unless the failure action is Throw, the process leaves through + // teardown_exit(). template void dispatch_wires(int n_wires, Fn&& fn) { require_open(*failed_); @@ -851,6 +853,9 @@ class RingImpl { if (error || stopped) { close(); + if (failure_action() == FailureAction::TeardownExit) { + teardown_exit(describe(error ? error : stopped).c_str()); + } std::rethrow_exception(error ? error : stopped); } } diff --git a/tests/jaccl/CMakeLists.txt b/tests/jaccl/CMakeLists.txt index f3da64e2ae..2f7eec9eee 100644 --- a/tests/jaccl/CMakeLists.txt +++ b/tests/jaccl/CMakeLists.txt @@ -29,7 +29,8 @@ jaccl_test(progress_guard_tests) add_test(NAME send_frame COMMAND send_frame_tests) -# The guard reads its limit once in a process, so each mode is one test. +# The guard reads its limit and action once in a process, so each mode is one +# test. teardown-exit and default-exit start this program again as children. foreach( mode checks @@ -38,6 +39,8 @@ foreach( precedence negative not-a-number - default) + default + teardown-exit + default-exit) add_test(NAME progress_guard_${mode} COMMAND progress_guard_tests ${mode}) endforeach() diff --git a/tests/jaccl/progress_guard_tests.cpp b/tests/jaccl/progress_guard_tests.cpp index bdd167fd8e..4e3536d5ee 100644 --- a/tests/jaccl/progress_guard_tests.cpp +++ b/tests/jaccl/progress_guard_tests.cpp @@ -1,15 +1,29 @@ // Copyright © 2026 Apple Inc. // Tests of the progress guard with simulated verbs. Each call goes through -// the real mesh and ring headers. The limit is read once in a process, so each -// mode of this program sets the environment first and is one CTest test. +// the real mesh and ring headers. The limit and the failure action are read +// once in a process, so each mode of this program sets the environment first +// and is one CTest test. The modes that check the exceptions select the +// "throw" action; the teardown-exit modes run each case in a child process of +// this program and check how it ended. +#include +#include +#include +#include + +#include #include +#include #include #include +#include #include #include #include +#include + +extern char** environ; #include "fixture.h" #include "jaccl/progress_guard.h" @@ -744,6 +758,389 @@ void stall() { }); } +// Teardown exit. A child runs one failing call; its memory release hook +// reports the ledger. The parent checks the exit status, the time from the +// start of the call to the exit, and what the child printed. + +// What the release hook of a child expects to find already destroyed. +int expected_pairs = 0, expected_regions = 0; +bool at_least = false; +std::atomic releases{0}; + +void report_release() { + const auto& l = fixture::ledger(); + const int pairs = l.destroy_queue_pair_calls, + queues = l.destroy_completion_queue_calls, + regions = l.deregister_calls; + const bool ok = at_least ? pairs >= expected_pairs && + queues >= expected_pairs && regions >= expected_regions + : pairs == expected_pairs && + queues == expected_pairs && regions == expected_regions; + std::printf( + "RELEASE %d %s pairs=%d queues=%d regions=%d expected=%d/%d\n", + ++releases, + ok ? "ok" : "mismatch", + pairs, + queues, + regions, + expected_pairs, + expected_regions); + std::fflush(stdout); +} + +void hang() { + for (;;) { + std::this_thread::sleep_for(std::chrono::seconds(1)); + } +} + +void call_start() { + std::printf( + "CALL %lld\n", + static_cast( + std::chrono::duration_cast( + Clock::now().time_since_epoch()) + .count())); + std::fflush(stdout); +} + +// The connections of one rank that hold a queue pair, and the regions of one +// rank (both ranks register the same number). +int live(const std::vector& links) { + int count = 0; + for (const auto& connection : links) { + count += connection.ctx != nullptr; + } + return count; +} + +// Run the call; reaching the end means the process did not exit. +template +int child_call(Call call) { + call_start(); + const Outcome outcome = attempt(call); + if (outcome.threw) { + std::printf("THREW %s\n", outcome.message.c_str()); + } else { + std::printf("RETURNED\n"); + } + return 0; +} + +int child(const std::string& name) { + limit = jaccl::progress_timeout_ms(); + if (name != "no-hook") { + jaccl::set_memory_release(name == "hung-hook" ? &hang : &report_release); + } + if (name == "mesh-recv" || name == "no-hook" || name == "hung-hook") { + fixture::MeshPair pair; + arm(pair.fabric); + expected_pairs = live(pair.links[1]); + expected_regions = fixture::ledger().regions / 2; + std::vector out(100, 0x29); + return child_call([&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + } + if (name == "mesh-all-reduce") { + fixture::MeshPair pair; + arm(pair.fabric); + expected_pairs = live(pair.links[0]); + expected_regions = fixture::ledger().regions / 2; + Buffers b; + return child_call([&] { + pair.nodes[0]->all_reduce(b.in.data(), b.out.data(), 64, fixture::Add{}); + }); + } + if (name == "mesh-failed-completion") { + fixture::MeshPair pair; + arm(pair.fabric); + fixture::endpoint(pair.links[1][0]).fault.fail_receive = 0; + expected_pairs = live(pair.links[1]); + expected_regions = fixture::ledger().regions / 2; + std::vector in(100, 0x53), out(100, 0x29); + std::thread sender( + [&] { attempt([&] { pair.nodes[0]->send(in.data(), 100, 1); }); }); + sender.detach(); + return child_call([&] { pair.nodes[1]->recv(out.data(), 100, 0); }); + } + if (name == "ring-recv-2") { + fixture::RingPair pair(2); + arm(pair.fabric); + expected_pairs = live(pair.left[1]) + live(pair.right[1]); + expected_regions = fixture::ledger().regions / 2; + std::vector out(100, 0x29); + return child_call([&] { pair.nodes[1]->recv(out.data(), 100, 0, 2); }); + } + if (name == "two-groups") { + // Two groups fail at the same time on two threads: one exit, one release. + fixture::MeshPair first, second; + arm(first.fabric); + arm(second.fabric); + expected_pairs = live(first.links[1]); + expected_regions = fixture::ledger().regions / 4; + at_least = true; + std::vector a(100, 0x29), b(100, 0x29); + std::thread other( + [&] { attempt([&] { second.nodes[1]->recv(b.data(), 100, 0); }); }); + other.detach(); + return child_call([&] { first.nodes[1]->recv(a.data(), 100, 0); }); + } + std::printf("FAIL unknown child case %s\n", name.c_str()); + return 2; +} + +struct Ended { + int status = -1; // the exit status, or -1 when the child did not exit + int64_t ms = -1; // from the start of the call to the exit + std::string out, err; +}; + +std::string slurp(const std::string& path) { + std::ifstream file(path); + std::stringstream text; + text << file.rdbuf(); + return text.str(); +} + +// Run this program as a child with `name` and only `variables` of the four +// guard variables set. +Ended run_child( + const char* program, + const std::string& name, + const std::map& variables) { + std::vector keep; + for (char** e = environ; *e != nullptr; ++e) { + std::string entry(*e); + bool guard = false; + for (const char* v : + {"JACCL_PROGRESS_TIMEOUT_MS=", + "MLX_JACCL_PROGRESS_TIMEOUT_MS=", + "JACCL_TIMEOUT_ACTION=", + "MLX_JACCL_TIMEOUT_ACTION="}) { + guard |= entry.rfind(v, 0) == 0; + } + if (!guard) { + keep.push_back(entry); + } + } + for (const auto& [key, value] : variables) { + keep.push_back(key + "=" + value); + } + // ThreadSanitizer sleeps 1 s in _exit by default; the timing checks measure + // the exit, so the children skip that sleep. Other sanitizers ignore it. + bool tsan_options = false; + for (auto& entry : keep) { + if (entry.rfind("TSAN_OPTIONS=", 0) == 0) { + entry += ":atexit_sleep_ms=0"; + tsan_options = true; + } + } + if (!tsan_options) { + keep.push_back("TSAN_OPTIONS=atexit_sleep_ms=0"); + } + std::vector envp; + for (auto& entry : keep) { + envp.push_back(entry.data()); + } + envp.push_back(nullptr); + + const char* dir = std::getenv("TMPDIR"); + std::string base = std::string(dir ? dir : "/tmp") + "/jaccl-guard-child-" + + std::to_string(getpid()) + "-" + name; + const std::string out_path = base + ".out", err_path = base + ".err"; + posix_spawn_file_actions_t actions; + posix_spawn_file_actions_init(&actions); + posix_spawn_file_actions_addopen( + &actions, 1, out_path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0600); + posix_spawn_file_actions_addopen( + &actions, 2, err_path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0600); + std::string self(program), mode("child"); + std::string which(name); + char* argv[] = {self.data(), mode.data(), which.data(), nullptr}; + pid_t pid = 0; + const int spawned = + posix_spawn(&pid, program, &actions, nullptr, argv, envp.data()); + posix_spawn_file_actions_destroy(&actions); + fixture::require(spawned == 0, "Could not start the child process"); + int status = 0; + while (waitpid(pid, &status, 0) < 0 && errno == EINTR) { + } + const auto end = Clock::now(); + Ended ended; + ended.status = WIFEXITED(status) ? WEXITSTATUS(status) : -1; + ended.out = slurp(out_path); + ended.err = slurp(err_path); + std::remove(out_path.c_str()); + std::remove(err_path.c_str()); + const auto call = ended.out.find("CALL "); + if (call != std::string::npos) { + const auto started = Clock::time_point( + std::chrono::nanoseconds( + std::stoll( + ended.out.substr(call + 5, ended.out.find('\n', call))))); + ended.ms = std::chrono::duration_cast(end - started).count(); + } + return ended; +} + +int count(const std::string& text, const std::string& part) { + int n = 0; + for (auto at = text.find(part); at != std::string::npos; + at = text.find(part, at + part.size())) { + n++; + } + return n; +} + +[[noreturn]] void child_mismatch(const std::string& what, const Ended& e) { + std::ostringstream text; + text << what << " (status " << e.status << ", " << e.ms + << " ms; stdout: " << e.out << "; stderr: " << e.err << ")"; + throw std::runtime_error(text.str()); +} + +// The child left through teardown_exit: status 75, the reason first, the +// release hook once after the group was released (unless `hook` is false), +// within [at, at + slack] ms of the start of the call. +void expect_teardown_exit( + const Ended& e, + const std::string& reason, + int64_t at, + bool hook = true) { + if (e.status != jaccl::TEARDOWN_EXIT_STATUS) { + child_mismatch("the process did not exit with status 75", e); + } + if (count(e.err, reason) != 1 || + count( + e.err, + "[jaccl] The group is released; releasing memory, then exiting " + "with status 75 (JACCL_TIMEOUT_ACTION=teardown-exit).") != 1) { + child_mismatch("the reason or the exit notice is missing", e); + } + if (hook) { + if (count(e.out, "RELEASE ") != 1 || count(e.out, "RELEASE 1 ok ") != 1 || + count(e.err, "[jaccl] Memory released; exiting with status 75.") != 1) { + child_mismatch( + "the release hook did not run once after the group was released", e); + } + } else if ( + count(e.out, "RELEASE ") != 0 || + count( + e.err, + "[jaccl] No memory release registered; exiting with status 75.") != + 1) { + child_mismatch("a process without a release hook did not say so", e); + } + if (count(e.out, "THREW") != 0 || count(e.out, "RETURNED") != 0) { + child_mismatch("the failed call came back to the caller", e); + } + if (e.ms < at || e.ms > at + slack) { + child_mismatch("the exit is outside [at, at + slack] ms", e); + } + static bool shown = false; + if (!shown) { + std::cout << "INFO teardown exit:\n" << e.err; + shown = true; + } +} + +// The child kept running: the call threw `reason` and nothing was released. +void expect_throw(const Ended& e, const std::string& reason) { + if (e.status != 0 || count(e.out, "THREW " + reason) != 1 || + count(e.out, "RELEASE ") != 0 || + count(e.err, "exiting with status") != 0) { + child_mismatch("the call did not throw to the caller", e); + } +} + +void teardown_exit_checks(const char* program) { + limit = 250; + const std::string limit_text = std::to_string(limit); + const std::string no_progress = "[jaccl] mesh recv: no completion for "; + using Env = std::map; + const Env ms = {{"JACCL_PROGRESS_TIMEOUT_MS", limit_text}}; + check("teardown exit is the default: a stalled mesh recv exits with 75", [&] { + expect_teardown_exit( + run_child(program, "mesh-recv", ms), no_progress, limit); + }); + check("teardown exit by name: a stalled mesh all_reduce exits with 75", [&] { + Env env = ms; + env["JACCL_TIMEOUT_ACTION"] = "teardown-exit"; + expect_teardown_exit( + run_child(program, "mesh-all-reduce", env), + "[jaccl] mesh all_reduce: no completion for ", + limit); + }); + check( + "teardown exit: a stalled ring recv on 2 wires exits after both wires stopped", + [&] { + expect_teardown_exit( + run_child(program, "ring-recv-2", ms), + "[jaccl] ring recv: no completion for ", + limit); + }); + check("teardown exit: a failed completion exits at once", [&] { + const Env env = {{"JACCL_PROGRESS_TIMEOUT_MS", "4000"}}; + const Ended e = run_child(program, "mesh-failed-completion", env); + expect_teardown_exit( + e, "[jaccl] mesh recv: a work completion failed with status 5", 0); + }); + check("teardown exit without a release hook still exits with 75", [&] { + expect_teardown_exit( + run_child(program, "no-hook", ms), no_progress, limit, false); + }); + check( + "teardown exit: a release hook that never returns is cut off after 10 s", + [&] { + const Ended e = run_child(program, "hung-hook", ms); + if (e.status != jaccl::TEARDOWN_EXIT_STATUS || + count( + e.err, + "[jaccl] The memory release did not finish in 10000 ms; " + "exiting.") != 1 || + e.ms < limit + jaccl::TEARDOWN_RELEASE_LIMIT_MS || + e.ms > limit + jaccl::TEARDOWN_RELEASE_LIMIT_MS + slack) { + child_mismatch("the release was not cut off at its limit", e); + } + }); + check( + "teardown exit: two groups that fail together give one release and one exit", + [&] { + expect_teardown_exit( + run_child(program, "two-groups", ms), no_progress, limit); + }); + check("an unknown action keeps the default (teardown exit)", [&] { + Env env = ms; + env["JACCL_TIMEOUT_ACTION"] = "abort"; + expect_teardown_exit( + run_child(program, "mesh-recv", env), no_progress, limit); + }); + check("JACCL_TIMEOUT_ACTION=throw is used before the alias", [&] { + Env env = ms; + env["JACCL_TIMEOUT_ACTION"] = "throw"; + env["MLX_JACCL_TIMEOUT_ACTION"] = "teardown-exit"; + expect_throw(run_child(program, "mesh-recv", env), no_progress); + }); + check("MLX_JACCL_TIMEOUT_ACTION=throw is read when the name is unset", [&] { + Env env = ms; + env["MLX_JACCL_TIMEOUT_ACTION"] = "throw"; + expect_throw(run_child(program, "mesh-recv", env), no_progress); + }); +} + +// No variable at all: the 30 s default limit and the teardown exit. +void default_exit_check(const char* program) { + limit = 30000; + check( + "default: with no variable a stalled mesh recv exits with 75 after 30 s", + [&] { + const Ended e = run_child(program, "mesh-recv", {}); + expect_teardown_exit(e, "[jaccl] mesh recv: no completion for ", limit); + if (count(e.err, "(JACCL_PROGRESS_TIMEOUT_MS is 30000)") != 1) { + child_mismatch("the reason does not name the 30 s limit", e); + } + }); +} + } // namespace // One mode in each process: @@ -753,14 +1150,25 @@ void stall() { // precedence JACCL_PROGRESS_TIMEOUT_MS is used before the alias // negative a negative value removes the limit // not-a-number a value that is not a number keeps the default -// default no variable: the default +// default no variable: the default limit (30 s) and action +// teardown-exit the failure action, each case in a child process +// default-exit no variable: a stalled call exits with 75 after 30 s +// All modes but the last three select JACCL_TIMEOUT_ACTION=throw. int main(int count, char** arguments) { const char* name = "JACCL_PROGRESS_TIMEOUT_MS"; const char* alias = "MLX_JACCL_PROGRESS_TIMEOUT_MS"; - const std::string mode = count == 2 ? arguments[1] : ""; + const std::string mode = count >= 2 ? arguments[1] : ""; + if (mode == "child" && count == 3) { + return child(arguments[2]); + } // The guard reads the variables when the first call starts. unsetenv(name); unsetenv(alias); + unsetenv("JACCL_TIMEOUT_ACTION"); + unsetenv("MLX_JACCL_TIMEOUT_ACTION"); + if (mode != "default" && mode != "teardown-exit" && mode != "default-exit") { + setenv("JACCL_TIMEOUT_ACTION", "throw", 1); + } if (mode == "checks") { setenv(name, "250", 1); value(250); @@ -789,10 +1197,22 @@ int main(int count, char** arguments) { setenv(name, "30s", 1); value(jaccl::DEFAULT_PROGRESS_TIMEOUT_MS); } else if (mode == "default") { - value(jaccl::DEFAULT_PROGRESS_TIMEOUT_MS); + value(30000); + check( + "the default limit is 30 s and the default action teardown-exit", [&] { + fixture::require( + jaccl::DEFAULT_PROGRESS_TIMEOUT_MS == 30000 && + jaccl::failure_action() == jaccl::FailureAction::TeardownExit, + "unexpected default limit or action"); + }); + } else if (mode == "teardown-exit") { + teardown_exit_checks(arguments[0]); + } else if (mode == "default-exit") { + default_exit_check(arguments[0]); } else { std::cout << "FAIL usage: progress_guard_tests checks | disabled | alias | " - "precedence | negative | not-a-number | default\n"; + "precedence | negative | not-a-number | default | " + "teardown-exit | default-exit\n"; return 2; } const auto& ledger = fixture::ledger(); diff --git a/tests/jaccl/send_frame_tests.cpp b/tests/jaccl/send_frame_tests.cpp index 46e420dd65..6b8b7cbc53 100644 --- a/tests/jaccl/send_frame_tests.cpp +++ b/tests/jaccl/send_frame_tests.cpp @@ -326,16 +326,20 @@ int main() { ring_collectives, "ring reduction/scatter full/partial/zero slices"); passed &= group(ring_gather, "ring all_gather distinct rank data, one/two wires"); - // The completion queues of the groups have MAX_SEND_WR + MAX_RECV_WR entries. - // Report how many of them this traffic used at the same time. + // The completion queues of the groups have COMPLETION_QUEUE_DEPTH entries, + // four times the work a queue pair can have outstanding. Report how many + // completions this traffic queued at the same time; it must stay within the + // outstanding bound, so that the headroom is real. + static_assert(COMPLETION_QUEUE_DEPTH == 4 * (MAX_SEND_WR + MAX_RECV_WR)); const size_t depth = fixture::ledger().completion_high_water; - const size_t capacity = MAX_SEND_WR + MAX_RECV_WR; - if (depth <= capacity) { + const size_t outstanding = MAX_SEND_WR + MAX_RECV_WR; + if (depth <= outstanding) { std::cout << "PASS completion queue high-water mark " << depth << " of " - << capacity << " entries\n"; + << COMPLETION_QUEUE_DEPTH << " entries (outstanding bound " + << outstanding << ")\n"; } else { std::cout << "FAIL completion queue high-water mark " << depth - << " exceeds " << capacity << '\n'; + << " exceeds the outstanding bound " << outstanding << '\n'; passed = false; } const auto& ledger = fixture::ledger(); From 56a20430714c75136a9da9d0332feb954d74cb6b Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Tue, 13 Oct 2026 18:04:28 -0700 Subject: [PATCH 15/19] docs(forkdiff): list jaccl.cpp, mesh.cpp and ring.cpp under the JACCL section fdae3174 changed three files that no fork.yaml glob covered, so the "fork.yaml describes the fork" check failed on this branch: mlx/distributed/jaccl/jaccl.cpp (registers the memory release hook that the teardown exit runs) and lib/jaccl/mesh.cpp / ring.cpp (create each completion queue at COMPLETION_QUEUE_DEPTH). Add the three paths to the JACCL section's globs and name them in its description. No code change. scripts/check_forkdiff.py --upstream-ref upstream/main: 205 changed files, 205 described, exit 0 (was 3 uncovered, exit 1). Co-Authored-By: Claude Opus 5.5 --- fork.yaml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/fork.yaml b/fork.yaml index 024b4ab920..60c109f672 100644 --- a/fork.yaml +++ b/fork.yaml @@ -206,7 +206,9 @@ def: `std::runtime_error` instead; every later call on that group throws immediately and the group cannot be reopened in the process. `dispatch_wires` also reports the errors of the ring pool wires, which it used to drop. Each completion queue has four times - the entries its queue pair can have outstanding (`COMPLETION_QUEUE_DEPTH`). + the entries its queue pair can have outstanding (`COMPLETION_QUEUE_DEPTH`, used by + `mesh.cpp` and `ring.cpp` when they create the queues). MLX's `jaccl.cpp` registers + the release hook each time it wraps a new JACCL group. Advertised GID: `Connection::info()` advertises the IPv4-mapped GID of the port, but `queue_pair_rtr` sent from GID index 1 whatever the index of that GID. The connection @@ -218,11 +220,14 @@ def: frame tails, the guard and the release path under ASan, UBSan and TSan. `ring_impl.h` also carries the upstream ring `all_gather` fix described further down. globs: + - "mlx/distributed/jaccl/jaccl.cpp" - "mlx/distributed/jaccl/lib/README.md" + - "mlx/distributed/jaccl/lib/jaccl/mesh.cpp" - "mlx/distributed/jaccl/lib/jaccl/mesh_impl.h" - "mlx/distributed/jaccl/lib/jaccl/progress_guard.h" - "mlx/distributed/jaccl/lib/jaccl/rdma.cpp" - "mlx/distributed/jaccl/lib/jaccl/rdma.h" + - "mlx/distributed/jaccl/lib/jaccl/ring.cpp" - "mlx/distributed/jaccl/lib/jaccl/ring_impl.h" - "mlx/distributed/jaccl/lib/jaccl/send_frame.h" - "docs/src/usage/distributed.rst" From 187b6c67bc13851911fed826fa26346d0ae990ee Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Sun, 25 Oct 2026 22:52:22 -0700 Subject: [PATCH 16/19] jaccl: report pool-wire errors in the shared dispatch_wires The merge of main brought upstream #4607 (a second cable for the JACCL mesh). Its MeshGroup::split_wires runs the wires through the free dispatch_wires in threadpool.h, which waits on the pool futures but never calls get(), so an error thrown in a pool wire was dropped and the call returned with that wire's part of the output unwritten. This branch had fixed the same drop in the ring's dispatch_wires; this applies the fix to the copy that main moved into threadpool.h. After every wire returned, the first pool-wire error is thrown. Nothing changes when no wire fails or when there is one wire (the two-Mac cluster uses one). fork.yaml: the JACCL section describes it and lists threadpool.h. Tests: tests/jaccl progress_guard_tests gains "dispatch_wires throws the error of a pool wire" (fails without the fix, passes with it). tests/jaccl with ASan and UBSan: 10 of 10 passed; with TSan: 10 of 10 passed. check_forkdiff.py --upstream-ref upstream/main: 112 changed files, all described. clang-format 21: no changes. Co-Authored-By: Claude Opus 5.5 --- fork.yaml | 12 ++++++--- mlx/distributed/jaccl/lib/jaccl/threadpool.h | 6 ++++- tests/jaccl/progress_guard_tests.cpp | 28 ++++++++++++++++++++ 3 files changed, 41 insertions(+), 5 deletions(-) diff --git a/fork.yaml b/fork.yaml index 1dd94021d5..5a7f26db3a 100644 --- a/fork.yaml +++ b/fork.yaml @@ -281,10 +281,13 @@ def: `dispatch_wires`, which also reports the errors of the ring pool wires that it used to drop. A mesh with more than one wire (upstream's second cable) has one `MeshImpl`, and so one guard, for each wire: a failed wire closes its own queue pairs and buffers, - and later calls that use that wire are refused. Each completion queue has four times - the entries its queue pair can have outstanding (`COMPLETION_QUEUE_DEPTH`, used by - `mesh.cpp` and `ring.cpp` when they create the queues). MLX's `jaccl.cpp` registers - the release hook each time it wraps a new JACCL group. + and later calls that use that wire are refused. The mesh runs its wires through the + shared `dispatch_wires` in `threadpool.h`, which upstream lets drop the error of a + pool wire; the fork throws that error after all wires returned. Each completion + queue has four times the entries its queue pair can have outstanding + (`COMPLETION_QUEUE_DEPTH`, used by `mesh.cpp` and `ring.cpp` when they create the + queues). MLX's `jaccl.cpp` registers the release hook each time it wraps a new JACCL + group. Advertised GID: `Connection::info()` advertises the IPv4-mapped GID of the port, but `queue_pair_rtr` sent from GID index 1 whatever the index of that GID. The connection @@ -305,6 +308,7 @@ def: - "mlx/distributed/jaccl/lib/jaccl/ring.cpp" - "mlx/distributed/jaccl/lib/jaccl/ring_impl.h" - "mlx/distributed/jaccl/lib/jaccl/send_frame.h" + - "mlx/distributed/jaccl/lib/jaccl/threadpool.h" - "docs/src/usage/distributed.rst" - "docs/src/usage/environment_variables.rst" - "tests/jaccl/*" diff --git a/mlx/distributed/jaccl/lib/jaccl/threadpool.h b/mlx/distributed/jaccl/lib/jaccl/threadpool.h index 6e96152cb5..123233f5e8 100644 --- a/mlx/distributed/jaccl/lib/jaccl/threadpool.h +++ b/mlx/distributed/jaccl/lib/jaccl/threadpool.h @@ -135,7 +135,8 @@ inline void ThreadPool::start_threads(size_t threads) { } // Run fn(lw) for each wire, the first n_wires - 1 on the pool and the last -// inline, then wait for the pool calls before returning. +// inline, then wait for the pool calls before returning. An error of a pool +// call is thrown after all calls returned. template void dispatch_wires(ThreadPool* pool, int n_wires, Fn&& fn) { if (n_wires <= 1 || pool == nullptr) { @@ -164,6 +165,9 @@ void dispatch_wires(ThreadPool* pool, int n_wires, Fn&& fn) { for (auto& f : futures) { f.wait(); } + for (auto& f : futures) { + f.get(); + } } } // namespace jaccl diff --git a/tests/jaccl/progress_guard_tests.cpp b/tests/jaccl/progress_guard_tests.cpp index 3aa02d51f2..8d527fcc88 100644 --- a/tests/jaccl/progress_guard_tests.cpp +++ b/tests/jaccl/progress_guard_tests.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -19,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -690,6 +692,31 @@ void below_limit() { }); } +// The mesh runs its wires through the shared dispatch_wires. The error of a +// pool wire must reach the caller, after every wire returned. +void pool_wire_error() { + check("dispatch_wires throws the error of a pool wire", [&] { + jaccl::ThreadPool pool(1); + std::atomic returned{0}; + std::string error; + try { + jaccl::dispatch_wires(&pool, 2, [&](int lw) { + if (lw == 0) { + std::this_thread::sleep_for(milliseconds(50)); + returned++; + throw std::runtime_error("pool wire failed"); + } + returned++; + }); + } catch (const std::runtime_error& e) { + error = e.what(); + } + fixture::require( + error == "pool wire failed", "The pool wire error was dropped"); + fixture::require(returned == 2, "A wire was still running"); + }); +} + // With the limit removed the loops wait as before; a failed completion is // still an error. void disabled() { @@ -1177,6 +1204,7 @@ int main(int count, char** arguments) { failed_completion(); refused_destroy(); below_limit(); + pool_wire_error(); } else if (mode == "disabled") { setenv(name, "0", 1); value(0); From d376bfe689db1ccaea24456ebf4f9ab281ae26ba Mon Sep 17 00:00:00 2001 From: jonathan308 <5055076+jonathan308@users.noreply.github.com> Date: Tue, 27 Oct 2026 13:24:27 -0700 Subject: [PATCH 17/19] Describe the NAX split attention JIT fix in fork.yaml main's #48 changed mlx/backend/metal/jit_kernels.cpp without a fork.yaml section, so the fork.yaml gate failed on main and on every PR based on it. The new section explains the change: the JIT source passes attention_nax_dsplit the same template arguments as its metallib instantiation. Co-Authored-By: Claude Opus 5.5 --- fork.yaml | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/fork.yaml b/fork.yaml index 5a7f26db3a..a8b8c87a1a 100644 --- a/fork.yaml +++ b/fork.yaml @@ -360,6 +360,20 @@ def: globs: - "mlx/backend/metal/kernels/quantized_nax.h" - "mlx/backend/metal/kernels/fp_quantized_nax.h" + - title: "JIT instantiation of the NAX head-dim split attention kernel" + description: | + `get_steel_attention_nax_kernel` builds the JIT source of both `attention_nax` and + `attention_nax_dsplit`, and upstream passes `bv` to both. `attention_nax_dsplit` has + no `BV` template parameter (its value head dim is `BD`), so every later argument + shifts by one, an integer lands in the `MaskType` slot, and the library does not + compile. JIT builds such as mlx-swift's compile this kernel at runtime for full + attention with head dim 256 or 512 and at least 1,024 query tokens on NAX GPUs, and + the caller's process exited with "Unable to build metal library from source". The + fork passes the split kernel the same arguments as its instantiation in + `steel_attention_nax.metal`; the non-split kernel is unchanged. Builds that use the + precompiled metallib were not affected. This is a latent upstream bug (PR #48). + globs: + - "mlx/backend/metal/jit_kernels.cpp" - title: "64-bit products in the JACCL ring and the vector SDPA layout check" description: | The JACCL ring computes each wire's byte offset as `lr * n_wires * size`. Upstream From 10d25ed620d3d6143b8c6c7b60988616e2b23856 Mon Sep 17 00:00:00 2001 From: David Tai Date: Tue, 27 Oct 2026 17:50:49 -0500 Subject: [PATCH 18/19] jaccl: share the mesh failure flag across wires; 64-bit all_gather offsets; report refused deregistration - The wires of a MeshGroup now share one failure flag. A failed wire stops a call that waits on another wire (WireStopped), and every later call on the group is refused, as the README and the guard comment say. Before, each MeshImpl had its own flag, so a failure on wire 1 left wire 0 open. - MeshImpl::all_gather keeps read_offset and write_offset in int64_t. With int, more than 2 GiB per rank overflowed and indexed out of bounds. - SharedBuffer::deregister reports a region that the verbs refuse to deregister and keeps it for the destructor, as release() does for queue pairs and completion queues. Before, it dropped the handle silently. - tests/jaccl: a refused-deregister case and a two-wire mesh case with a shared flag. Fork Tests build-cpp now builds and runs tests/jaccl, which no CI job built before. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/fork-tests.yml | 16 ++++- fork.yaml | 16 +++-- mlx/distributed/jaccl/lib/jaccl/mesh.cpp | 10 ++- mlx/distributed/jaccl/lib/jaccl/mesh_impl.h | 17 +++-- mlx/distributed/jaccl/lib/jaccl/rdma.h | 15 +++-- tests/jaccl/fake_rdma.cpp | 3 + tests/jaccl/fake_rdma.h | 1 + tests/jaccl/fixture.h | 13 +++- tests/jaccl/progress_guard_tests.cpp | 73 +++++++++++++++++++++ 9 files changed, 144 insertions(+), 20 deletions(-) diff --git a/.github/workflows/fork-tests.yml b/.github/workflows/fork-tests.yml index 8af6aac956..8cad0b2499 100644 --- a/.github/workflows/fork-tests.yml +++ b/.github/workflows/fork-tests.yml @@ -6,7 +6,8 @@ name: Fork Tests # The jobs: # build-wheel builds the Python wheel and uploads it. # build-cpp builds the C++ library and tests, and runs the C++ tests -# with the GPU and then with the CPU as the default device. +# with the GPU and then with the CPU as the default device, +# and the JACCL tests with simulated verbs (tests/jaccl). # Then it runs the standalone allocation cap tests, one # process for each case. # python-gpu runs the Python tests with the GPU as the default device. @@ -140,6 +141,19 @@ jobs: LLVM_PROFILE_FILE: ${{ runner.temp }}/profiles/cpu-%p-%m.profraw run: ./build/tests/tests + - name: Run the JACCL tests with simulated verbs + # tests/jaccl is its own CMake project, which the build above does + # not build: the real mesh, ring and RDMA headers against fake verbs, + # with ASan and UBSan. The step clears the coverage flags. + if: ${{ !cancelled() }} + env: + CFLAGS: '' + CXXFLAGS: '' + run: | + cmake -S tests/jaccl -B build-jaccl -DCMAKE_BUILD_TYPE=Debug + cmake --build build-jaccl -j $(sysctl -n hw.physicalcpu) + ctest --test-dir build-jaccl --output-on-failure --no-tests=error + - name: Run the standalone allocation cap tests # The process allocation cap is install-once in each process. ctest # runs each entry in a fresh process, and the Metal test runs one case diff --git a/fork.yaml b/fork.yaml index a8b8c87a1a..222eefd8dd 100644 --- a/fork.yaml +++ b/fork.yaml @@ -280,10 +280,14 @@ def: the group cannot be reopened in the process. The ring keeps its own guarded `dispatch_wires`, which also reports the errors of the ring pool wires that it used to drop. A mesh with more than one wire (upstream's second cable) has one `MeshImpl`, - and so one guard, for each wire: a failed wire closes its own queue pairs and buffers, - and later calls that use that wire are refused. The mesh runs its wires through the + and so one guard, for each wire, and the wires of a `MeshGroup` share one failure + flag: a failed wire closes its own queue pairs and buffers, a wire that waits in the + same call stops with `WireStopped` and closes its own, and every later call on the + group is refused. The mesh runs its wires through the shared `dispatch_wires` in `threadpool.h`, which upstream lets drop the error of a - pool wire; the fork throws that error after all wires returned. Each completion + pool wire; the fork throws that error after all wires returned. A region that the + verbs refuse to deregister is reported and kept for the destructor, as a refused + queue pair or completion queue is. Each completion queue has four times the entries its queue pair can have outstanding (`COMPLETION_QUEUE_DEPTH`, used by `mesh.cpp` and `ring.cpp` when they create the queues). MLX's `jaccl.cpp` registers the release hook each time it wraps a new JACCL @@ -294,7 +298,8 @@ def: now remembers the index where the advertised GID was found and sends from it; behaviour is unchanged where that index is 1. - `tests/jaccl` is a separate CMake project that the MLX build does not build. It + `tests/jaccl` is a separate CMake project that the MLX build does not build; the + `build-cpp` job of Fork Tests builds it and runs it. It compiles the real mesh, ring and RDMA headers against simulated verbs and checks the frame tails, the guard and the release path under ASan, UBSan and TSan. globs: @@ -499,7 +504,8 @@ def: `fork-tests.yml` builds the fork with Metal and runs its C++ and Python tests on GitHub-hosted macOS runners. Two jobs build in parallel. One job builds the Python wheel. The other job builds the C++ library and tests, and runs the C++ tests with the - GPU and then with the CPU as the default device. Then it runs the standalone allocation + GPU and then with the CPU as the default device, and builds and runs the JACCL tests + with simulated verbs (`tests/jaccl`). Then it runs the standalone allocation cap tests with ctest, one process for each case. The Python tests run in shard jobs: one set with the GPU as the default device, and one set with the CPU. `.github/scripts/fork-tests/python-shards.json` gives the shards. diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp index 6a34b2f377..5470ff1342 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp @@ -32,10 +32,16 @@ MeshGroup::MeshGroup( // Make sure every node has reached here before continuing side_channel_.barrier(); - // Create the mesh implementation objects + // Create the mesh implementation objects. The wires share one failure flag. + auto failed = std::make_shared>(false); for (int w = 0; w < n_wires_; w++) { meshes_.emplace_back( - rank_, size_, connections_[w], buffers_[w], scatter_buffers_[w]); + rank_, + size_, + connections_[w], + buffers_[w], + scatter_buffers_[w], + failed); } } diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index ed23d952e8..2cf71f57d2 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -24,14 +24,17 @@ class MeshImpl { int size, std::vector& conns, std::vector& buffers, - std::vector& scatter_buffers) + std::vector& scatter_buffers, + std::shared_ptr> failed = + std::make_shared>(false)) : rank_(rank), size_(size), connections_(conns), buffers_(buffers), scatter_buffers_(scatter_buffers), staging_mem_( - std::make_unique(MESH_PIPELINE * MAX_BUFFER_SIZE)) {} + std::make_unique(MESH_PIPELINE * MAX_BUFFER_SIZE)), + failed_(std::move(failed)) {} MeshImpl() : rank_(0), size_(1) {} @@ -227,9 +230,9 @@ class MeshImpl { // Counters to maintain the state of transfers int in_flight = 0; - int read_offset = 0; + int64_t read_offset = 0; int completed_send_count[PIPELINE] = {0}; - int write_offset[MESH_MAX_PEERS] = {0}; + int64_t write_offset[MESH_MAX_PEERS] = {0}; // Prefill the pipeline int buff = 0; @@ -679,8 +682,10 @@ class MeshImpl { std::span buffers_; std::span scatter_buffers_; std::unique_ptr staging_mem_; - std::unique_ptr> failed_ = - std::make_unique>(false); + // Shared by the wires of one group, so that a failure on one wire stops the + // others and closes the group for every later call. + std::shared_ptr> failed_ = + std::make_shared>(false); }; } // namespace jaccl diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.h b/mlx/distributed/jaccl/lib/jaccl/rdma.h index f2023ca515..2193d8de74 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.h +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.h @@ -115,12 +115,19 @@ class SharedBuffer { void register_to_protection_domain(ibv_pd* protection_domain); - // Unpin the memory. The buffer cannot be posted after this call. + // Unpin the memory. The buffer cannot be posted after this call. A region + // that the verbs refuse to deregister is reported and kept, so that the + // destructor tries again. void deregister() { - for (auto& [pd, mr] : memory_regions_) { - ibv().dereg_mr(mr); + for (auto it = memory_regions_.begin(); it != memory_regions_.end();) { + if (int status = ibv().dereg_mr(it->second); status != 0) { + std::cerr << IBV_TAG << " Could not deregister a memory region (" + << status << ")" << std::endl; + ++it; + } else { + it = memory_regions_.erase(it); + } } - memory_regions_.clear(); } size_t size() const { diff --git a/tests/jaccl/fake_rdma.cpp b/tests/jaccl/fake_rdma.cpp index e38a00a823..963b946b39 100644 --- a/tests/jaccl/fake_rdma.cpp +++ b/tests/jaccl/fake_rdma.cpp @@ -86,6 +86,9 @@ int destroy_completion_queue(ibv_cq* queue) { int deregister(ibv_mr* region) { fixture::ledger().deregister_calls++; + if (fixture::ledger().refuse_deregister) { + return 16; + } drop(region); return 0; } diff --git a/tests/jaccl/fake_rdma.h b/tests/jaccl/fake_rdma.h index 202d65229f..f7cb8af345 100644 --- a/tests/jaccl/fake_rdma.h +++ b/tests/jaccl/fake_rdma.h @@ -50,6 +50,7 @@ struct Ledger { std::atomic deregister_calls{0}; std::atomic completion_high_water{0}; std::atomic refuse_destroy{false}; + std::atomic refuse_deregister{false}; }; Ledger& ledger(); diff --git a/tests/jaccl/fixture.h b/tests/jaccl/fixture.h index 9e1e5a48ae..d7e01ef409 100644 --- a/tests/jaccl/fixture.h +++ b/tests/jaccl/fixture.h @@ -83,12 +83,21 @@ inline void buffers( } } +// The failure flag of a mesh group, which its wires share. +using Flag = std::shared_ptr>; +inline Flag flag() { + return std::make_shared>(false); +} + struct MeshPair { Fabric fabric; std::array, 2> links; std::array, 2> scratch, scatter; std::array, 2> nodes; - MeshPair() { + MeshPair() : MeshPair({flag(), flag()}) {} + // `failed` holds the failure flag of each rank. Two pairs that get the same + // flags are two wires of one group. + explicit MeshPair(std::array failed) { for (auto& rank : links) { connections(rank, 2); } @@ -97,7 +106,7 @@ struct MeshPair { buffers(scratch[rank], NUM_BUFFERS * 2, links[rank]); buffers(scatter[rank], NUM_BUFFERS * 4, links[rank]); nodes[rank] = std::make_unique( - rank, 2, links[rank], scratch[rank], scatter[rank]); + rank, 2, links[rank], scratch[rank], scatter[rank], failed[rank]); } } }; diff --git a/tests/jaccl/progress_guard_tests.cpp b/tests/jaccl/progress_guard_tests.cpp index 8d527fcc88..896f505e9a 100644 --- a/tests/jaccl/progress_guard_tests.cpp +++ b/tests/jaccl/progress_guard_tests.cpp @@ -614,6 +614,77 @@ void refused_destroy() { }); } +// When the verbs refuse to deregister a region, the failure is reported and +// the region stays registered, so that the destructor tries again. +void refused_deregister() { + check( + "refused deregister: the error is printed, the region stays for the destructor", + [&] { + fixture::MeshPair pair; + auto& buffer = pair.scratch[0].front(); + auto* domain = pair.links[0][1].protection_domain; + std::ostringstream captured; + auto* previous = std::cerr.rdbuf(captured.rdbuf()); + fixture::ledger().refuse_deregister = true; + buffer.deregister(); + fixture::ledger().refuse_deregister = false; + std::cerr.rdbuf(previous); + fixture::require( + captured.str() == + "[jaccl] Could not deregister a memory region (16)\n", + "The refused deregistration was not reported"); + fixture::require( + buffer.local_key(domain) == 1, + "A region that was not deregistered was dropped"); + buffer.deregister(); + bool dropped = false; + try { + buffer.local_key(domain); + } catch (const std::out_of_range&) { + dropped = true; + } + fixture::require(dropped, "A deregistered region was kept"); + }); +} + +// The wires of one mesh group share the failure flag: a failure on one wire +// stops a call that waits on the other wire, and the group then refuses +// every later call on every wire. +void mesh_wires_share_failure() { + check( + "mesh wires: a failed wire stops the other wire, then both refuse", + [&] { + const auto failed = std::array{fixture::flag(), fixture::flag()}; + fixture::MeshPair first(failed), second(failed); + arm(first.fabric); + arm(second.fabric); + fixture::endpoint(second.links[1][0]).fault.poll_error = true; + std::vector out(100, 0x29), other(100, 0x29); + // Rank 0 sends nothing, so the receive on the first wire waits. + auto waiting = std::async(std::launch::async, [&] { + return attempt([&] { first.nodes[1]->recv(out.data(), 100, 0); }); + }); + std::this_thread::sleep_for(milliseconds(limit / 5)); + const Outcome failure = + attempt([&] { second.nodes[1]->recv(other.data(), 100, 0); }); + const Outcome stopped = waiting.get(); + if (!failure.threw || + failure.message.rfind("[jaccl] mesh recv: ", 0) != 0) { + mismatch("the call on the failed wire did not fail", failure); + } + if (!stopped.threw || stopped.message != jaccl::WireStopped().what()) { + mismatch("the waiting wire was not stopped", stopped); + } + if (stopped.ms >= limit) { + mismatch("the waiting wire ran into the limit", stopped); + } + expect_closed( + attempt([&] { first.nodes[1]->recv(out.data(), 100, 0); })); + expect_closed( + attempt([&] { second.nodes[1]->recv(other.data(), 100, 0); })); + }); +} + // Traffic that waits, but less than the limit, is not touched. void below_limit() { check("below the limit: a peer that starts late is waited for", [&] { @@ -1203,6 +1274,8 @@ int main(int count, char** arguments) { lost_completion(); failed_completion(); refused_destroy(); + refused_deregister(); + mesh_wires_share_failure(); below_limit(); pool_wire_error(); } else if (mode == "disabled") { From 28e32c79f96abfbf82f3a3926fe8fa69cba22e96 Mon Sep 17 00:00:00 2001 From: David Tai Date: Tue, 27 Oct 2026 18:00:59 -0500 Subject: [PATCH 19/19] Apply clang-format to the shared failure flag change Co-Authored-By: Claude Opus 5.5 --- mlx/distributed/jaccl/lib/jaccl/mesh_impl.h | 3 +-- tests/jaccl/progress_guard_tests.cpp | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index 2cf71f57d2..7b62f3e746 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -32,8 +32,7 @@ class MeshImpl { 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)), failed_(std::move(failed)) {} MeshImpl() : rank_(0), size_(1) {} diff --git a/tests/jaccl/progress_guard_tests.cpp b/tests/jaccl/progress_guard_tests.cpp index 896f505e9a..4cf652c187 100644 --- a/tests/jaccl/progress_guard_tests.cpp +++ b/tests/jaccl/progress_guard_tests.cpp @@ -652,8 +652,7 @@ void refused_deregister() { // every later call on every wire. void mesh_wires_share_failure() { check( - "mesh wires: a failed wire stops the other wire, then both refuse", - [&] { + "mesh wires: a failed wire stops the other wire, then both refuse", [&] { const auto failed = std::array{fixture::flag(), fixture::flag()}; fixture::MeshPair first(failed), second(failed); arm(first.fabric);