Fix silent Metal OOM error by sticky stream

mlx/array.cpp

@@ -143,13 +143,18 @@ bool array::is_available() const {
 }
 
 void array::wait() {
-  if (!is_available()) {
-    if (event().valid()) {
-      event().wait();
-      detach_event();
-    }
-    set_status(Status::available);
+  if (status() == Status::available) {
+    return;
+  }
+  // Wait on the event even when it is already signaled. Going through
+  // Event::wait is what reports an asynchronous backend failure, and taking
+  // the is_available() shortcut here would skip it in exactly the case the
+  // GPU raced ahead of the error being recorded.
+  if (event().valid()) {
+    event().wait();
+    detach_event();
   }
+  set_status(Status::available);
 }
 
 void array::eval() {

mlx/backend/metal/device.cpp

@@ -1,5 +1,6 @@
 // Copyright © 2023-2024 Apple Inc.
 
+#include <atomic>
 #include <cstdlib>
 #include <sstream>
 
@@ -306,11 +307,62 @@ MTL::Library* load_library(
 
 } // namespace
 
+// Whether any stream has ever parked an error. Errors are rare and terminal,
+// so waiters gate on this single flag instead of paying the registry mutex
+// and shared_ptr atomics on every wait. Monotonic: never cleared.
+static std::atomic<bool> any_stream_error{false};
+
+void StreamError::set_error(std::shared_ptr<std::string> error) {
+  if (!error) {
+    return;
+  }
+  // Preserve the earliest error.
+  std::shared_ptr<std::string> expected{nullptr};
+  std::atomic_compare_exchange_strong(&error_, &expected, std::move(error));
+  any_stream_error.store(true, std::memory_order_release);
+}
+
+void StreamError::check_error() {
+  auto error = std::atomic_exchange(&error_, {});
+  if (error) {
+    throw std::runtime_error(*error);
+  }
+}
+
+void StreamError::clear_error() {
+  std::atomic_store(&error_, std::shared_ptr<std::string>{});
+}
+
+std::shared_ptr<std::string> StreamError::error() const {
+  return std::atomic_load(&error_);
+}
+
+std::shared_ptr<StreamError> get_stream_error(int stream_index) {
+  static std::mutex mtx;
+  static std::unordered_map<int, std::shared_ptr<StreamError>> errors;
+  std::lock_guard<std::mutex> lk(mtx);
+  auto& slot = errors[stream_index];
+  if (!slot) {
+    slot = std::make_shared<StreamError>();
+  }
+  return slot;
+}
+
+void check_stream_error(Stream s) {
+  if (s.device != mlx::core::Device::gpu) {
+    return;
+  }
+  if (!any_stream_error.load(std::memory_order_acquire)) {
+    return;
+  }
+  get_stream_error(s.index)->check_error();
+}
+
 CommandEncoder::CommandEncoder(
     Device& d,
     int index,
     ResidencySet& residency_set)
-    : device_(d) {
+    : device_(d), stream_error_(get_stream_error(index)) {
   auto pool = new_scoped_memory_pool();
   queue_ = NS::TransferPtr(device_.mtl_device()->newCommandQueue());
   if (!queue_) {
@@ -501,6 +553,13 @@ void CommandEncoder::signal_event(
     std::shared_ptr<EventImpl> event,
     uint64_t value) {
   end_encoding();
+  // Carry a known failure onto the event as early as possible. Streams that
+  // wait on this event pick their error up from it (see wait_events_ in
+  // commit), and doing it here rather than in the completion handler narrows
+  // the window where a downstream buffer commits without seeing it.
+  if (auto error = stream_error_->error()) {
+    event->set_error(error);
+  }
   buffer_->encodeSignalEvent(event->mtl_event(), value);
   signal_events_.push_back({std::move(event), value});
 }
@@ -520,31 +579,29 @@ bool CommandEncoder::needs_commit() const {
 
 void CommandEncoder::commit(std::function<void()> completion) {
   buffer_->addCompletedHandler(
-      [&error_ = error_,
+      [stream_error = stream_error_,
        wait_events = std::move(wait_events_),
        signal_events = std::move(signal_events_),
        completion = std::move(completion)](MTL::CommandBuffer* cbuf) {
         if (completion) {
           completion();
         }
-        // If any of the waited event has error in it, poison the encoder.
+        // If any of the waited event has error in it, poison the stream.
+        // set_error ignores null and keeps the earliest error.
         for (auto& event : wait_events) {
-          if (event->error()) {
-            error_ = event->error();
-            break;
-          }
+          stream_error->set_error(event->error());
         }
-        // Set error only when no error happended before, to preserve the
-        // earliest error.
-        if (!error_ && cbuf->status() == MTL::CommandBufferStatusError) {
-          error_ = std::make_shared<std::string>(fmt::format(
-              "[METAL] Command buffer execution failed: {}.",
-              cbuf->error()->localizedDescription()->utf8String()));
+        // set_error preserves the earliest error.
+        if (cbuf->status() == MTL::CommandBufferStatusError) {
+          stream_error->set_error(
+              std::make_shared<std::string>(fmt::format(
+                  "[METAL] Command buffer execution failed: {}.",
+                  cbuf->error()->localizedDescription()->utf8String())));
         }
         // Poison all the signaled events when error happened.
-        if (error_) {
+        if (auto error = stream_error->error()) {
           for (auto& [event, value] : signal_events) {
-            event->set_error(error_);
+            event->set_error(error);
           }
         }
         // Metal won't signal the events for us on error, manually signal them
@@ -568,9 +625,10 @@ void CommandEncoder::synchronize() {
   commit();
   cbuf->waitUntilCompleted();
 
-  if (error_ && !exiting_) {
-    auto error = std::move(error_);
-    throw std::runtime_error(*error);
+  // Don't throw during teardown, and don't clear either: the slot is shared
+  // by every user of this stream index, not owned by this encoder.
+  if (!exiting_) {
+    stream_error_->check_error();
   }
 }
 
@@ -579,9 +637,9 @@ MTL::ComputeCommandEncoder* CommandEncoder::get_command_encoder() {
     encoder_ = NS::RetainPtr(
         buffer_->computeCommandEncoder(MTL::DispatchTypeConcurrent));
     fence_ = NS::TransferPtr(device_.mtl_device()->newFence());
-    // Reset error when user starts to encode new commands, they are supposed to
-    // have handled the error in synchronize() or Event::wait().
-    error_.reset();
+    // Do not clear stream_error_ here. It is sticky until synchronize() or
+    // Event::wait() reports it — mid-eval commits may fail with no signal
+    // event attached, and resetting would silently drop the error.
   }
   return encoder_.get();
 }

mlx/backend/metal/device.h

@@ -23,6 +23,23 @@ using MTLFCList =
 class Device;
 class EventImpl;
 
+// Sticky per-stream GPU error. Lives by stream index so it outlives
+// CommandEncoder encode cycles and is visible from any thread that waits on
+// that stream. Set from a command buffer completion handler, consumed by the
+// first synchronization point that reports it.
+struct MLX_API StreamError {
+  void set_error(std::shared_ptr<std::string> error);
+  void check_error();
+  void clear_error();
+  std::shared_ptr<std::string> error() const;
+
+ private:
+  std::shared_ptr<std::string> error_;
+};
+
+MLX_API std::shared_ptr<StreamError> get_stream_error(int stream_index);
+void check_stream_error(Stream s);
+
 class MLX_API CommandEncoder {
  public:
   CommandEncoder(Device& d, int index, ResidencySet& residency_set);
@@ -118,8 +135,8 @@ class MLX_API CommandEncoder {
   std::vector<std::shared_ptr<EventImpl>> wait_events_;
   std::vector<std::tuple<std::shared_ptr<EventImpl>, uint64_t>> signal_events_;
 
-  // Error from previous commited command buffer.
-  std::shared_ptr<std::string> error_;
+  // Sticky error for this stream (shared by stream index).
+  std::shared_ptr<StreamError> stream_error_;
 
   // Encoder for issuing GPU commands.
   // The members are used within a single ComputeCommandEncoder and will be

mlx/backend/metal/event.cpp

@@ -26,9 +26,7 @@ EventImpl::~EventImpl() {
 }
 
 void EventImpl::wait(uint64_t value) {
-  check_error();
   mtl_event_->waitUntilSignaledValue(value, -1); // never times out
-  check_error();
 }
 
 void EventImpl::signal(uint64_t value) {
@@ -39,13 +37,6 @@ void EventImpl::set_error(std::shared_ptr<std::string> error) {
   std::atomic_store(&error_, std::move(error));
 }
 
-void EventImpl::check_error() {
-  auto error = std::atomic_exchange(&error_, {});
-  if (error) {
-    throw std::runtime_error(*error);
-  }
-}
-
 } // namespace metal
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -58,11 +49,20 @@ Event::Event(Stream stream) : stream_(stream) {
 
 void Event::wait() {
   static_cast<metal::EventImpl*>(event_.get())->wait(value());
+  // The event being signaled says nothing about whether the work succeeded:
+  // Metal signals on the GPU timeline, ahead of the completion handler that
+  // records the failure, and a command buffer can die with no event attached
+  // at all. The stream error covers both.
+  metal::check_stream_error(stream_);
 }
 
 void Event::wait(Stream stream) {
   auto impl = std::static_pointer_cast<metal::EventImpl>(event_);
   if (stream.device == Device::cpu) {
+    // Deliberately does not report errors. This runs on a scheduler worker
+    // with no handler above it, so throwing here would terminate the process
+    // instead of reaching the user. The failure stays parked on the signaling
+    // stream and surfaces at the next synchronization point.
     scheduler::enqueue(stream, [impl = std::move(impl), value = value()]() {
       impl->wait(value);
     });

mlx/backend/metal/event.h

@@ -10,10 +10,11 @@ class EventImpl {
   EventImpl(Device& d);
   ~EventImpl();
 
+  // Blocks until signaled. Does not throw: a failure that reaches this event
+  // was parked on its stream first, so the stream is where it gets reported.
   void wait(uint64_t value);
   void signal(uint64_t value);
   void set_error(std::shared_ptr<std::string> error);
-  void check_error();
 
   const auto& error() const {
     return error_;

tests/gpu_tests.cpp

@@ -4,10 +4,17 @@
 #include <chrono>
 #include <cmath>
 #include <future>
+#include <memory>
+#include <stdexcept>
+#include <string>
 
 #include "doctest/doctest.h"
 #include "mlx/mlx.h"
 
+#if defined(__APPLE__)
+#include "mlx/backend/metal/device.h"
+#endif
+
 using namespace mlx::core;
 
 static const std::array<Dtype, 5> types =
@@ -686,3 +693,93 @@ TEST_CASE("test layer norm vjp bias grad race") {
   }
   CHECK(worst <= 1e-5);
 }
+
+#if defined(__APPLE__)
+// A command buffer that runs out of the working set is killed by the driver
+// without writing its outputs. Reproducing that needs more memory than a test
+// machine can be assumed to have, so these inject the error the completion
+// handler would have recorded and check that it reaches the user from there.
+namespace {
+
+constexpr const char* kInjected =
+    "[METAL] Command buffer execution failed: test sticky error.";
+
+void inject_stream_error(Stream s) {
+  metal::get_stream_error(s.index)->set_error(
+      std::make_shared<std::string>(kInjected));
+}
+
+// Drop anything left parked so one failing case cannot leak into the next.
+struct ClearStreamErrorOnExit {
+  Stream s;
+  ~ClearStreamErrorOnExit() {
+    metal::get_stream_error(s.index)->clear_error();
+  }
+};
+
+} // namespace
+
+TEST_CASE("test stream error reported by synchronize") {
+  auto s = default_stream(Device::gpu);
+  ClearStreamErrorOnExit cleanup{s};
+
+  inject_stream_error(s);
+  CHECK_THROWS_WITH_AS(synchronize(s), kInjected, std::runtime_error);
+  // Consumed by the report, so the next one is clean.
+  synchronize(s);
+}
+
+TEST_CASE("test stream error survives encoding more work") {
+  // The failing command buffer is committed mid-eval and signals no event, so
+  // the error is only on the stream. Encoding the next kernel must not drop it.
+  auto s = default_stream(Device::gpu);
+  ClearStreamErrorOnExit cleanup{s};
+
+  auto a = full({4, 4}, 1.0f, float32, s);
+  eval({a});
+
+  inject_stream_error(s);
+  auto b = add(a, a, s);
+  CHECK_THROWS_WITH_AS(eval({b}), kInjected, std::runtime_error);
+  synchronize(s);
+}
+
+TEST_CASE("test stream error reported through an already signaled event") {
+  // Metal signals shared events on the GPU timeline, ahead of the completion
+  // handler that records the failure, so a waiter can find the event signaled
+  // and the error not yet posted. Reading the array must still report it.
+  auto s = default_stream(Device::gpu);
+  ClearStreamErrorOnExit cleanup{s};
+
+  auto a = full({4, 4}, 1.0f, float32, s);
+  auto x = add(a, a, s);
+  async_eval({x});
+  synchronize(s); // event is signaled and the work is done
+
+  inject_stream_error(s);
+  CHECK_THROWS_WITH_AS(eval({x}), kInjected, std::runtime_error);
+}
+
+TEST_CASE("test stream error does not escape a scheduler worker") {
+  // A cpu-stream consumer waits on the gpu event from a scheduler thread that
+  // has no handler above it. Reporting there would terminate the process, so
+  // the error has to stay parked for the next synchronization point.
+  auto gs = default_stream(Device::gpu);
+  auto cs = default_stream(Device::cpu);
+  ClearStreamErrorOnExit cleanup{gs};
+
+  auto a = full({512, 512}, 1.0f, float32, gs);
+  auto x = add(a, a, gs);
+  for (int i = 0; i < 32; ++i) {
+    x = add(x, a, gs); // keep the gpu busy so the event is still unsignaled
+  }
+  async_eval({x});
+
+  inject_stream_error(gs);
+  auto y = add(x, x, cs);
+  CHECK_NOTHROW(eval({y}));
+
+  // Still reportable on the stream that failed.
+  CHECK_THROWS_WITH_AS(synchronize(gs), kInjected, std::runtime_error);
+}
+#endif