fix: store Metal completion handler errors instead of throwing

Metal completion handlers run on dispatch queues where C++ exceptions
cannot propagate — throwing causes std::terminate → SIGABRT, crashing
the process with no diagnostic information.

Instead, store the error message atomically in the CommandEncoder and
check it at the next synchronous point (commit, synchronize). This
converts fatal crashes into catchable runtime_error exceptions that
the application can handle gracefully.

Root cause analysis: the crash at 262K+ context reported as mlx#3216
was actually TWO separate issues:

1. Thread safety in stream management (fixed by PR #3281)
2. C++ exceptions thrown from Metal completion handler callbacks
   (fixed by this commit)

The GPU watchdog error (kIOGPUCommandBufferCallbackErrorImpactingInteractivity)
is a separate concern — macOS kills command buffers that block the GPU
beyond the watchdog threshold. This commit ensures that error is reported
as a Python RuntimeError instead of SIGABRT.

Author: Thump604

mlx/backend/metal/device.cpp

@@ -424,7 +424,28 @@ bool CommandEncoder::needs_commit() const {
   return (buffer_ops_ > max_ops) || ((buffer_sizes_ >> 20) > max_mb);
 }
 
+void CommandEncoder::store_error(const std::string& msg) {
+  std::lock_guard<std::mutex> lock(error_mtx_);
+  if (!has_error_.load(std::memory_order_relaxed)) {
+    error_msg_ = msg;
+    has_error_.store(true, std::memory_order_release);
+  }
+}
+
+void CommandEncoder::check_stored_error() {
+  if (has_error_.load(std::memory_order_acquire)) {
+    std::string msg;
+    {
+      std::lock_guard<std::mutex> lock(error_mtx_);
+      msg = std::move(error_msg_);
+      has_error_.store(false, std::memory_order_release);
+    }
+    throw std::runtime_error(msg);
+  }
+}
+
 void CommandEncoder::commit() {
+  check_stored_error();
   buffer_->commit();
   buffer_ = NS::RetainPtr(queue_->commandBufferWithUnretainedReferences());
   buffer_ops_ = 0;

mlx/backend/metal/device.h

@@ -3,6 +3,7 @@
 #pragma once
 
 #include <Metal/Metal.hpp>
+#include <atomic>
 #include <functional>
 #include <mutex>
 #include <shared_mutex>
@@ -90,6 +91,13 @@ class MLX_API CommandEncoder {
   bool needs_commit() const;
   void commit();
 
+  // Deferred error handling for completion handler errors.
+  // Metal completion handlers run on dispatch queues where C++ exceptions
+  // cannot propagate — throwing causes std::terminate → SIGABRT.
+  // Instead, errors are stored and checked at the next synchronous point.
+  void store_error(const std::string& msg);
+  void check_stored_error();
+
   MTL::CommandQueue* get_command_queue() const {
     return queue_.get();
   }
@@ -125,6 +133,12 @@ class MLX_API CommandEncoder {
   // A map of prior command encoder outputs to their corresponding fence.
   std::unordered_map<const void*, NS::SharedPtr<MTL::Fence>> prev_ce_outputs_;
   std::mutex outputs_mtx_;
+
+  // Deferred error from Metal completion handlers (set from dispatch queue,
+  // checked from eval/synchronize thread).
+  std::atomic<bool> has_error_{false};
+  std::mutex error_mtx_;
+  std::string error_msg_;
 };
 
 class MLX_API Device {

mlx/backend/metal/eval.cpp

@@ -17,7 +17,8 @@ void new_stream(Stream stream) {
   }
 }
 
-inline void check_error(MTL::CommandBuffer* cbuf) {
+// check_error_throw: for synchronous contexts where exceptions propagate.
+inline void check_error_throw(MTL::CommandBuffer* cbuf) {
   if (cbuf->status() == MTL::CommandBufferStatusError) {
     std::ostringstream msg;
     msg << "[METAL] Command buffer execution failed: "
@@ -56,19 +57,34 @@ void eval(array& arr) {
     buffers.erase(it);
   }
 
+  // Capture a pointer to the encoder for deferred error storage.
+  // Completion handlers run on Metal's dispatch queue where C++ exceptions
+  // cannot propagate (throwing → std::terminate → SIGABRT). Instead, store
+  // the error and check it at the next synchronous point (commit, synchronize).
+  auto* enc = &d.get_command_encoder(s.index);
   if (d.command_buffer_needs_commit(s.index)) {
     d.end_encoding(s.index);
     scheduler::notify_new_task(s);
     command_buffer->addCompletedHandler(
-        [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
+        [s, enc, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
           scheduler::notify_task_completion(s);
-          check_error(cbuf);
+          if (cbuf->status() == MTL::CommandBufferStatusError) {
+            std::ostringstream msg;
+            msg << "[METAL] Command buffer execution failed: "
+                << cbuf->error()->localizedDescription()->utf8String();
+            enc->store_error(msg.str());
+          }
         });
     d.commit_command_buffer(s.index);
   } else {
     command_buffer->addCompletedHandler(
-        [buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
-          check_error(cbuf);
+        [enc, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
+          if (cbuf->status() == MTL::CommandBufferStatusError) {
+            std::ostringstream msg;
+            msg << "[METAL] Command buffer execution failed: "
+                << cbuf->error()->localizedDescription()->utf8String();
+            enc->store_error(msg.str());
+          }
         });
   }
 }
@@ -77,20 +93,32 @@ void finalize(Stream s) {
   auto pool = metal::new_scoped_memory_pool();
   auto& d = metal::device(s.device);
   auto cb = d.get_command_buffer(s.index);
+  auto* enc = &d.get_command_encoder(s.index);
   d.end_encoding(s.index);
-  cb->addCompletedHandler([](MTL::CommandBuffer* cbuf) { check_error(cbuf); });
+  cb->addCompletedHandler([enc](MTL::CommandBuffer* cbuf) {
+    if (cbuf->status() == MTL::CommandBufferStatusError) {
+      std::ostringstream msg;
+      msg << "[METAL] Command buffer execution failed: "
+          << cbuf->error()->localizedDescription()->utf8String();
+      enc->store_error(msg.str());
+    }
+  });
   d.commit_command_buffer(s.index);
 }
 
 void synchronize(Stream s) {
   auto pool = metal::new_scoped_memory_pool();
   auto& d = metal::device(s.device);
+  auto& enc = d.get_command_encoder(s.index);
   auto cb = d.get_command_buffer(s.index);
   cb->retain();
   d.end_encoding(s.index);
   d.commit_command_buffer(s.index);
   cb->waitUntilCompleted();
-  check_error(cb);
+  // Check both this command buffer and any deferred errors from earlier
+  // completion handlers that ran on Metal's dispatch queue.
+  check_error_throw(cb);
+  enc.check_stored_error();
   cb->release();
 }