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/docs/src/usage/distributed.rst b/docs/src/usage/distributed.rst index 861525dcbf..6566e42f24 100644 --- a/docs/src/usage/distributed.rst +++ b/docs/src/usage/distributed.rst @@ -401,6 +401,32 @@ 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` (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 +such wait, or set it to ``0`` to remove the limit. + .. _nccl_section: Getting Started with NCCL @@ -573,6 +599,13 @@ 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 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 b43c053d52..47e35dde0d 100644 --- a/docs/src/usage/environment_variables.rst +++ b/docs/src/usage/environment_variables.rst @@ -63,6 +63,28 @@ 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 ``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 + 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 1526db9d42..222eefd8dd 100644 --- a/fork.yaml +++ b/fork.yaml @@ -30,6 +30,13 @@ def: control, declared-mutable custom-kernel inputs for paged KV, and a use-after-free fix in 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. + + 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, 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. sub: - title: "Metal eval: use-after-free with mid-eval synchronization" description: | @@ -243,6 +250,74 @@ def: globs: - "python/tests/test_metal_threadgroup_bounds.py" - "docs/src/dev/custom_metal_kernels.rst" + - 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 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 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. 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, 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. 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 + 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 + 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; 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: + - "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" + - "mlx/distributed/jaccl/lib/jaccl/threadpool.h" + - "docs/src/usage/distributed.rst" + - "docs/src/usage/environment_variables.rst" + - "tests/jaccl/*" + - "tests/jaccl/fake_verbs/infiniband/verbs.h" - title: "Quantized tile bounds and distinct loader symbols" description: | Sorted gathered matrix products use the remaining reduction width for a @@ -290,6 +365,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 @@ -415,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/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 d682a06bbc..92d80c3199 100644 --- a/mlx/distributed/jaccl/lib/README.md +++ b/mlx/distributed/jaccl/lib/README.md @@ -102,6 +102,27 @@ 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 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 @@ -321,6 +342,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/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp index fe856faf97..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); } } @@ -47,7 +53,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 7beed48e64..7b62f3e746 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; @@ -21,14 +24,16 @@ 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)) {} + staging_mem_(std::make_unique(MESH_PIPELINE * MAX_BUFFER_SIZE)), + failed_(std::move(failed)) {} MeshImpl() : rank_(0), size_(1) {} @@ -42,6 +47,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; @@ -74,10 +80,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); @@ -101,6 +104,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; @@ -112,10 +116,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; @@ -198,6 +199,7 @@ 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; } } @@ -207,6 +209,8 @@ class MeshImpl { char* out_ptr, int64_t n_bytes, int64_t stride) { + 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_ * stride). @@ -225,18 +229,16 @@ 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; 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++; @@ -250,6 +252,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; @@ -261,10 +264,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; @@ -310,6 +311,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_) * stride; auto [sz, buffer_size] = buffer_size_from_message(count * sizeof(T)); @@ -340,8 +342,7 @@ class MeshImpl { continue; } const T* src = in + static_cast(p) * stride + 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); @@ -356,6 +357,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; @@ -375,10 +377,7 @@ class MeshImpl { } const T* src = in + static_cast(p) * stride + 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); @@ -456,6 +455,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); @@ -466,10 +466,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++; @@ -485,6 +483,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; @@ -492,10 +491,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; @@ -506,6 +503,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); @@ -530,6 +528,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; @@ -553,6 +552,48 @@ 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. 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) + : ProgressGuard(op, *mesh.failed_), + mesh_(mesh), + exceptions_(std::uncaught_exceptions()) { + require_open(*mesh.failed_); + } + + ~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()); + } + } + } + + 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); @@ -640,6 +681,10 @@ class MeshImpl { std::span buffers_; std::span scatter_buffers_; std::unique_ptr staging_mem_; + // 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/progress_guard.h b/mlx/distributed/jaccl/lib/jaccl/progress_guard.h new file mode 100644 index 0000000000..a2009dd972 --- /dev/null +++ b/mlx/distributed/jaccl/lib/jaccl/progress_guard.h @@ -0,0 +1,261 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace jaccl { + +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; + +/** + * 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; +} + +// 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() + : 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. The + * owner of the group closes it and then, unless failure_action() is Throw, + * calls teardown_exit(). + */ +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(); + } + } + + protected: + // The error that check() threw, or empty. + const std::string& failure() const { + return failure_; + } + + 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."; + failure_ = msg.str(); + throw std::runtime_error(failure_); + } + + const char* op_; + std::atomic& failed_; + 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.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 01d9ed37d6..2193d8de74 100644 --- a/mlx/distributed/jaccl/lib/jaccl/rdma.h +++ b/mlx/distributed/jaccl/lib/jaccl/rdma.h @@ -5,12 +5,14 @@ #include #include +#include #include #include #include #include #include +#include "jaccl/send_frame.h" #include "jaccl/tcp.h" constexpr const char* IBV_TAG = "[jaccl]"; @@ -18,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; @@ -109,6 +115,21 @@ class SharedBuffer { void register_to_protection_domain(ibv_pd* protection_domain); + // 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 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); + } + } + } + size_t size() const { return num_bytes_; } @@ -125,6 +146,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_); @@ -159,6 +186,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); @@ -217,6 +246,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( @@ -239,6 +289,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; } @@ -250,13 +304,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.cpp b/mlx/distributed/jaccl/lib/jaccl/ring.cpp index ae25d80904..c161ab8f41 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 5863de4006..b7fb0134cf 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -2,8 +2,13 @@ #pragma once +#include +#include +#include +#include #include +#include "jaccl/progress_guard.h" #include "jaccl/rdma.h" #include "jaccl/threadpool.h" @@ -89,7 +94,7 @@ class RingImpl { // Split the reduce scatter + all gather across the available wires. Each // wire handles a contiguous slice of each chunk in every direction. - dispatch_wires(pool_, n_wires, [&](int lw) { + dispatch_wires(n_wires, [&](int lw) { all_reduce_wire( in_ptr, out_ptr, @@ -195,7 +200,7 @@ class RingImpl { // Split the all gather across the available wires. Each wire handles a // contiguous slice of every rank's data in both directions. size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires); - dispatch_wires(pool_, n_wires, [&](int lw) { + dispatch_wires(n_wires, [&](int lw) { all_gather_wire( out_ptr, n_bytes, @@ -272,7 +277,7 @@ class RingImpl { // Two directional regions, each split across the wires. int64_t size_per_wire = (chunk + (2 * n_wires) - 1) / (2 * n_wires); - dispatch_wires(pool_, n_wires, [&](int lw) { + dispatch_wires(n_wires, [&](int lw) { reduce_scatter_wire( in_ptr, out_ptr, chunk, size_per_wire, n_wires, lw, reduce_op); }); @@ -368,6 +373,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. @@ -387,11 +393,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); } @@ -404,6 +412,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; @@ -414,11 +423,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++; @@ -500,6 +511,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. @@ -513,11 +525,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); } @@ -532,6 +545,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; @@ -541,11 +555,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]++; @@ -595,7 +610,7 @@ class RingImpl { // Split the send across the available wires. Each wire handles the // contiguous slice [lw * bytes_per_wire, (lw + 1) * bytes_per_wire). - dispatch_wires(pool_, n_wires, [&](int lw) { + dispatch_wires(n_wires, [&](int lw) { send_wire(in_ptr, n_bytes, dir, bytes_per_wire, lw); }); } @@ -624,10 +639,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++; @@ -636,6 +649,7 @@ class RingImpl { } // Main loop + ProgressGuard guard("ring send", *failed_); while (in_flight > 0) { // Poll the hardware for completions. // @@ -643,16 +657,16 @@ 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; 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; @@ -674,7 +688,7 @@ class RingImpl { // Split the recv across the available wires. Each wire handles the // contiguous slice [lw * bytes_per_wire, (lw + 1) * bytes_per_wire). - dispatch_wires(pool_, n_wires, [&](int lw) { + dispatch_wires(n_wires, [&](int lw) { recv_wire(out_ptr, n_bytes, dir, bytes_per_wire, lw); }); } @@ -710,6 +724,7 @@ class RingImpl { } // Main loop + ProgressGuard guard("ring recv", *failed_); while (in_flight > 0) { // Poll the hardware for completions. // @@ -717,6 +732,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; @@ -780,6 +796,86 @@ class RingImpl { work_completions); } + // 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, 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_); + + auto run = [&](int lw) { + try { + fn(lw); + } catch (...) { + failed_->store(true); + throw; + } + }; + + 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(); + } + } + }; + + 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) { + wait([&] { f.get(); }); + } + } + + if (error || stopped) { + close(); + if (failure_action() == FailureAction::TeardownExit) { + teardown_exit(describe(error ? error : stopped).c_str()); + } + 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(); + } + } + int rank_; int size_; int n_conns_; @@ -788,6 +884,8 @@ class RingImpl { std::span send_buffers_; std::span recv_buffers_; ThreadPool* pool_; + std::unique_ptr> failed_ = + std::make_unique>(false); }; } // namespace jaccl 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 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/CMakeLists.txt b/tests/jaccl/CMakeLists.txt new file mode 100644 index 0000000000..2f7eec9eee --- /dev/null +++ b/tests/jaccl/CMakeLists.txt @@ -0,0 +1,46 @@ +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 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 + disabled + alias + precedence + negative + not-a-number + default + teardown-exit + default-exit) + 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..963b946b39 --- /dev/null +++ b/tests/jaccl/fake_rdma.cpp @@ -0,0 +1,312 @@ +// 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++; + if (fixture::ledger().refuse_deregister) { + return 16; + } + 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..f7cb8af345 --- /dev/null +++ b/tests/jaccl/fake_rdma.h @@ -0,0 +1,94 @@ +// 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}; + std::atomic refuse_deregister{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..d7e01ef409 --- /dev/null +++ b/tests/jaccl/fixture.h @@ -0,0 +1,198 @@ +// 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); + } + } + } + } +} + +// 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({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); + } + 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], failed[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..4cf652c187 --- /dev/null +++ b/tests/jaccl/progress_guard_tests.cpp @@ -0,0 +1,1330 @@ +// Copyright © 2026 Apple Inc. + +// Tests of the progress guard with simulated verbs. Each call goes through +// 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 +#include +#include + +extern char** environ; + +#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, 256); }}, + {"mesh sum_scatter", + [&](auto& node) { + node.sum_scatter(b.in.data(), b.out.data(), 64, 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"); + }); +} + +// 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", [&] { + 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(); + }); +} + +// 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() { + 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"); + }); +} + +// 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: +// 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 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] : ""; + 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); + silent_peer(); + lost_completion(); + failed_completion(); + refused_destroy(); + refused_deregister(); + mesh_wires_share_failure(); + below_limit(); + pool_wire_error(); + } 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(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 | " + "teardown-exit | default-exit\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..849fbf8ffc --- /dev/null +++ b/tests/jaccl/send_frame_tests.cpp @@ -0,0 +1,363 @@ +// 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), + count * sizeof(float)); + } else { + pair.nodes[rank]->sum_scatter( + input[rank].data(), + output[rank].data(), + count, + 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 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 outstanding = MAX_SEND_WR + MAX_RECV_WR; + if (depth <= outstanding) { + std::cout << "PASS completion queue high-water mark " << depth << " of " + << COMPLETION_QUEUE_DEPTH << " entries (outstanding bound " + << outstanding << ")\n"; + } else { + std::cout << "FAIL completion queue high-water mark " << depth + << " exceeds the outstanding bound " << outstanding << '\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; +}