Propagate CPU errors to events

[zcbenz]

This PR implements exception handling for errors happened in eval_cpu. Similar to #3523, the cpu scheduler would poison all pending events in the stream whenever an error happened, and an exception would throw when the poisoned event is synchronized.

Most of this PR is doing refactoring:

1. Move the error handling from metal::EventImpl to the public Event class.
2. Add methods to Scheduler to make it capable of setting errors in events.
3. Refactor platform event implementations to use the new Scheduler methods to signal/wait events.

Note that most of the errors happened in eval_cpu would be fatal and not recoverable, so this PR does not catch all errors, instead we have to catch the expected errors and pass to the scheduler explicitly, this PR handles the IO error in Load::eval_cpu as example.

[aleroot]

This PR adds errors to Event, but array::is_available() can now silently discard them.

If a CPU load fails, the event may have both error != nullptr and is_signaled() == true by the time the caller reaches array::wait(). In that case is_available() takes this branch, detaches the event, marks the array available, and never calls Event::wait() / check_error().

A concrete fast-failure interleaving is: the event is inserted, set_error() poisons it, the stream later signals it, and all of that finishes before the main thread calls eval_impl(...).wait(). The final wait then sees a signaled event and swallows the error.

I think either array::is_available() must check/take the event error before detaching a signaled event, or Event::is_signaled() needs to surface poisoned events somehow.

For context, these issues would prevent me from reliably landing ml-explore/mlx-swift#427, which is why I opened my original MLX PR.

That Swift PR depends on CPU lazy-load read failures propagating deterministically to eval. If those errors can be dropped or swallowed, the progress API can work for the happy path but still cannot safely handle truncated or failed safetensors reads.

[zcbenz]

Thanks a lot for reviewing this!

I updated the PR with a different strategy: the error happened in eval_cpu is now persistent in scheduler per stream, until the eval ends. All signaled events in the stream would be poisoned by the error in stream, and all waited events would poison the stream if an error happened.

On the race condition of error() I made method private and added a thread-safe load_error() to replace it.

On array::is_available() swallowing the error, I made array::detach_event check error before detaching.

[aleroot]

Thank you for this work, once released I will definitely make use of it in my apps.

[aleroot]

@zcbenz Sorry to disturb, but can this be merged so that we can unlock the mlx-swift development as well ? Thanks.

[sashko-zakharchuk]

Ran into a case this doesn't cover yet, in the ring backend after #4060.

When a peer disconnects, SocketThread fails the pending promises and the exception is
rethrown by future::get() inside the closures the collectives dispatch onto the stream
thread. Nothing catches on that thread, so the process dies with terminate called after throwing an instance of 'std::runtime_error' rather than surfacing the error. Deterministic
on a 2-rank localhost ring: hard-kill one rank while the other runs ring send/recv on one
thread and a CPU-stream all_sum on another; the survivor aborts mid-run, every run, CPU and
CUDA backends both (sm_120). Same with this branch merged into current main; it merges
cleanly.

Wiring the ring dispatch sites into this PR's mechanism closes it: the dispatched closures
get wrapped so a comm failure goes to scheduler::set_error instead of escaping the thread.
Two smaller pieces fell out of that: the segment futures need draining so no socket task
still points into array buffers after a failure, and a broken SocketThread fails fast so a
collective on a dead group raises rather than hanging. The survivor then gets
RuntimeError: [ring] connection to a peer was lost out of mx.eval and can keep running.

What I checked: a CPU-stream all_sum feeding a GPU matmul in one synchronous eval raises
out of mx.eval, so the error-carrying events in this PR do cross streams through the fence
path; a second collective on the dead group raises immediately; ring_test_distributed.py
passes 13/13 on both ranks under both DEVICE=cpu and DEVICE=gpu. The functional delta is
+71/-22 in ring.cpp (git diff -w; clang-format reindentation on top of that). Happy to open
it as a follow-up PR on top of this branch, or fold it in here, whichever you prefer.

One note: even with the fix, a process that exits while the ring is still up can hit
terminate called without an active exception at teardown. That is pre-existing on main for
clean runs as well and looks like #4110's territory, so I left it alone.

repro (2-rank localhost ring, kill one rank mid-run)

# repro.py, one process per rank:
#   MLX_RANK=0 MLX_HOSTFILE=hosts.json python repro.py
#   MLX_RANK=1 MLX_HOSTFILE=hosts.json python repro.py
# hosts.json: [["127.0.0.1:15500"], ["127.0.0.1:15501"]]
# kill -9 the rank-1 process mid-run; rank 0 dies with SIGABRT on main.
import time
import mlx.core as mx

group = mx.distributed.init(backend="ring")
cpu = mx.default_stream(mx.Device(mx.cpu))
a = mx.ones((256, 256))
expected = float(group.size()) * a.size
for it in range(2000):
    z = mx.distributed.all_sum(a, group=group, stream=cpu)
    y = (z @ mx.eye(256)).sum()  # gpu consumer, fence-linked to z
    try:
        mx.eval(y)
    except RuntimeError as e:
        print(f"it {it}: caught {e}")
        break
    assert float(y) == expected
    time.sleep(0.01)

[erwinzhang7]

Confirmed this myself, since #4060 is what introduced the gap.

Reproduced on an M5 Max, macOS 26.6, 2-rank localhost ring on current main, using your repro. Rank 0 dies
with exit 134:

[ring] Socket 3 was closed by the peer
[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 except RuntimeError never runs. So the detection from #4060 is doing its job and the
infinite hang is gone, but the exception escapes the stream thread and the process aborts
instead of the caller catching it. #4060 claimed pending operations fail "with an exception the
caller can observe", and that isn't happening on this path.

Your read of the mechanism matches what I'm seeing: the f.get() calls live inside the
encoder.dispatch closures in ring.cpp, so there is no handler on that thread.

[zcbenz]

Thanks for testing this PR. I have rebased it on #4174 and change the code to simply catch and transfer all exceptions in CPU streams.

[erwinzhang7]

Tested the rebase against the peer-loss case. It closes it, with no change to ring.cpp.

Built 1add2cfe4, M5 Max, macOS 26.6, two rank localhost ring, one rank hard killed mid-run
while the survivor runs a CPU-stream all_sum in a loop inside try/except RuntimeError:

build	outcome
v0.32.0 (released)	hangs, no exception, killed at 40 s, 3/3
current main 9ab977b56	aborts, exit 134, except never runs, 4/4
this PR 1add2cfe4	except RuntimeError runs, exit 0, 8/8

[ring] Too many send/recv errors. Failing pending tasks...
rank 0 CAUGHT RuntimeError: [ring] connection to a peer was lost
RESULT: exception surfaced to the caller

The currently released version hangs rather than aborts, since #4060 has not shipped in a
tag yet. So on 0.32.0 a peer loss is an infinite wait, on main it is an abort, and with this it
is a catchable error. The last is the one #4060 claimed and did not deliver.

Why no ring change is needed

Scheduler::enqueue wraps every task now:

st.enqueue([&st, task = std::move(task)]() mutable {
  try {
    task();
  } catch (const std::exception& error) {
    if (!st.error.valid()) { ... }
  }
});

and cpu::CommandEncoder::dispatch goes through scheduler::enqueue on both of its branches,
which is where ring's four encoder.dispatch sites end up. The f.get() rethrow that had no
handler on the stream thread now has one, so it becomes a stream error and surfaces at eval.

One thing I looked for and did not find

In dispatch, every tenth call wraps the task so that notify_task_completion runs after it:

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

The new catch sits outside that lambda, in enqueue, so a throw from task() skips
notify_task_completion and I expected an unbalanced counter to turn the abort into a hang
whenever the failing dispatch landed on that boundary. Eight consecutive runs all exited
cleanly, so if it is reachable it is rarer than this repro reaches, and it may not be reachable
at all. It's clearly not urgent, but may potentially show up later as a flake rather than a failure.

Not touching @sashko-zakharchuk's ring work: the two pieces beyond the abort, draining segment
futures so no socket task still points into array buffers and failing fast on a dead group, are
separate problems that a catch does not solve.

[zcbenz]

In dispatch, every tenth call wraps the task so that notify_task_completion runs after it:

Thanks for noticing that, I think it is a valid concern that totally could happen. I added a fix for that.