ring: fail on peer disconnect instead of hanging forever

[erwinzhang7]

Proposed changes

Fixes #3862.

Killing one rank of a ring group leaves every surviving rank hung with no way to recover.
Two separate defects combine to produce it.

An orderly peer close is invisible. recv() reports it by returning 0 and leaves
errno untouched, so the errno != EAGAIN test reads a stale value. On a non-blocking
socket that value is almost always EAGAIN, because every earlier call with no data
available set it, so the failure is skipped entirely. error_count never increments, the
abort is never reached, and the worker spins on a dead socket forever at 100% CPU without
logging anything. sendAll() and recvAll() in the nccl backend already treat <= 0 as
failure.

This is the more common failure in practice. In a 4 rank reproduction, two of the three
survivors ended up here rather than on the reported abort path, which is why the issue
describes ranks hanging with no diagnostic at all.

Reaching the abort does not release waiters. The worker returns, leaving every queued
task's promise unsatisfied. Because the SocketThread outlives its worker those promises
are not destroyed either, so no broken_promise is delivered, the futures never become
ready, and every .wait() blocks forever.

The change treats r == 0 as an error on both the send and recv paths, and exits rather
than returning once the error threshold is hit.

On stopping the process

The log line already said Aborting...; this makes the behaviour match it.

Stopping is deliberate. I first implemented the alternative the issue suggests, resolving
the pending promises so the surviving ranks continue. It does unblock them, but the
collectives then complete without the dead rank's contribution and return wrong results
silently. With a workload whose input changes each iteration, the survivors produced
got=2552 against want=4600 and kept running. For numerical work that seemed worse than
failing, so I did not propose it.

Rejecting the promises with set_exception and switching the internal .wait() calls to
.get() does not work either: the throw would surface on the stream thread, and
StreamThread::thread_fn calls task() with no try/catch, so it unwinds out of the thread
entry point into std::terminate. I believe this is what the second reporter on the issue
saw as a crash on Metal.

That leaves failing fast, which is what the nccl backend already does for unrecoverable
communication errors. If you would prefer a different failure mode the mechanism is
confined to one line and I am happy to change it.

Verification

4 rank loopback ring, all_sum in a loop, one rank killed with SIGKILL mid run.
Ranks launched directly rather than through mlx.launch, since the launcher tears down the
survivors on a rank exit and hides the behaviour.

	before	after
surviving ranks	3/3 alive and frozen indefinitely	3/3 log the error and exit
the spinning rank	0 log lines, 100% CPU	11 log lines, then exits

Fork CI is green across all 22 jobs (Linux x86_64 and aarch64 on cpu and CUDA 12.6/12.9/13.0,
Windows x86_64 and aarch64, macOS 14/15/26.2 on cpu, metal and jit, plus lint). That includes
the new test, which unittest discover picks up and runs on the runners, and
mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py, which passes with no new
log output, as expected when no fault is injected.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

On the test. python/tests/test_ring_peer_loss.py spawns a 3 rank loopback ring,
kills one rank, and asserts the survivors exit instead of blocking forever. Against 0.32.0
it fails after 61s with the survivors still running; with this change it passes in about a
second. I ran it 8 consecutive times on the fix with no flakes.

It is named test_*.py rather than *_test_distributed.py deliberately. The
*_test_distributed.py files are excluded from unittest discover so they only run under
mpirun or mlx.launch; this one needs the opposite. It has to manage its own processes,
because mlx.launch terminates the surviving ranks when one exits, which is exactly the
behaviour under test, and running it inside ring_test_distributed.py would disturb that
file's eight rank run.

Two things I did for CI robustness: ports come from binding to port 0 and releasing rather
than a fixed range that could collide on a shared runner, and a failure to form the ring at
all skips rather than fails, so a broken environment does not look like a regression. Only
the survivors-still-running case fails.

[zcbenz]

I think the error can be passed to the main thread by setting exception in the promise when error_count >= 10. Aborting the process is not the best choice in this case.

[erwinzhang7]

Thanks, I thought about this before opening the PR but only reasoned about why it wouldn't work. This time I built it, and I think the outcome still argues against it, though the reasoning changed.

Implemented exactly as suggested: on error_count >= 10, reject every queued task's promise
with set_exception instead of stopping, and switch the ten internal .wait() calls to
.get() so the exception is observed rather than silently ignored.

Same reproduction as the PR: 4 ranks, all_sum in a loop, SIGKILL rank 2. All three
survivors end like this:

[ring] Socket 3 was closed by the peer
... (x10)
[ring] Too many send/recv errors. Failing pending tasks...
libc++abi: terminating due to uncaught exception of type std::runtime_error: [ring] connection to a peer was lost

The exception never reaches the main thread. The ring's .wait() calls run inside a lambda
handed to encoder.dispatch(...) (ring.cpp:525), which executes on the stream thread, and
StreamThread::thread_fn (scheduler.h:48) calls task() with no try/catch. Throwing there
unwinds straight out of the thread entry point.

The try/catch in eval() (transforms.cpp:228, from #3675) does not cover it, that
runs on the calling thread and only guards the inline portion of eval_cpu, not work that
has been dispatched.

So the net change is that a controlled exit(1) with a logged reason becomes an uncatchable
SIGABRT with no Python traceback. The process stops either way; likely what the
second reporter on #3862 saw as a crash on Metal.

There is a second problem even setting terminate aside. CommandEncoder::dispatch wraps
work as:

auto task_wrap = [s = stream_, task = std::move(task)]() mutable {
  task();
  scheduler::notify_task_completion(s);
};

If task() throws, notify_task_completion never runs, so the scheduler's active-task count
leaks and wait_for_one() can block. Failing a promise from inside dispatched work is not
currently safe regardless of what the ring does with it.

What it would take to do it properly

I do like the shape of what you are asking for, and it is achievable, just not from inside
ring.cpp. It needs the exception caught where the work is dispatched, stored as an
exception_ptr, and rethrown on the calling thread at a synchronization point. There is no
such mechanism today: grep exception_ptr across mlx/scheduler.h, mlx/backend/cpu/ and
mlx/backend/metal/ returns nothing, and CommandEncoder::dispatch has no handler.

That would be a useful addition, every backend that can fail inside dispatched
work would benefit, not just the ring, but it is a scheduler change with its own design
questions (which thread observes it, what happens to the other queued tasks on that stream,
how it interacts with notify_task_completion), and it felt wrong to bundle into a
peer-disconnect fix.

Proposed Fixes:

1. Keep exit(1) here, matching the nccl backend's handling of unrecoverable communication
   errors and the Aborting... the log line already claims.
2. Leave this PR at the r == 0 detection half only, and fix the wedge separately once
   there is a way to surface errors from dispatched work.
3. If you would like, I will prototype the exception_ptr plumbing in the scheduler as its
   own PR, and rebase this on top.

I'd do 1 for now and 3 as a follow-up, but I can also try to tackle the propagation path.

[zcbenz]

We have a PR propagating the error from stream thread to main thread #3742, I can rebase it after this PR properly throw exception.

[erwinzhang7]

Done, and thanks for the pointer to #3742.

The pending promises are now rejected with set_exception and the internal waits are get()
so the exception is observed rather than discarded. exit(1) is gone.

Also rebased onto current main while making this change.

[zcbenz]

Thanks!