diff --git a/mlx/backend/cuda/event.cu b/mlx/backend/cuda/event.cu index d3b6f97f5d..b2bd8f57de 100644 --- a/mlx/backend/cuda/event.cu +++ b/mlx/backend/cuda/event.cu @@ -341,7 +341,8 @@ void Event::wait() { } else { event.atomic->wait(value()); } - CHECK_CUDA_ERROR(cudaPeekAtLastError()); + // Check for errors during kernel execution and reset the error state. + CHECK_CUDA_ERROR(cudaGetLastError()); check_error(); } diff --git a/mlx/backend/cuda/utils.cpp b/mlx/backend/cuda/utils.cpp index 82272b74b2..f81f589fff 100644 --- a/mlx/backend/cuda/utils.cpp +++ b/mlx/backend/cuda/utils.cpp @@ -12,6 +12,8 @@ namespace mlx::core { void check_cuda_error(const char* name, cudaError_t err) { if (err != cudaSuccess) { + // Clear the error so it does not resurface in later checks. + cudaGetLastError(); throw std::runtime_error( fmt::format("{} failed: {}", name, cudaGetErrorString(err))); } diff --git a/python/src/buffer.h b/python/src/buffer.h index 4b194b3d25..553bba80cd 100644 --- a/python/src/buffer.h +++ b/python/src/buffer.h @@ -88,9 +88,13 @@ extern "C" inline int getbuffer(PyObject* obj, Py_buffer* view, int flags) { std::memset(view, 0, sizeof(Py_buffer)); auto a = nb::cast(nb::handle(obj)); - { + // Exceptions can not propagate through the buffer protocol. + try { nb::gil_scoped_release nogil; a.eval(); + } catch (const std::exception& e) { + PyErr_SetString(PyExc_RuntimeError, e.what()); + return -1; } std::vector shape(a.shape().begin(), a.shape().end()); diff --git a/python/tests/test_array.py b/python/tests/test_array.py index e03803631a..3ec8616458 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -2065,6 +2065,12 @@ def test_buffer_protocol_ref_counting(self): mv = None self.assertIsNone(wr()) + def test_buffer_protocol_eval_error(self): + # Errors from evaluating the array are raised instead of aborting + a = mx.linalg.inv(mx.zeros((2, 2)), stream=mx.cpu) + with self.assertRaises(RuntimeError): + memoryview(a) + def test_array_view_ref_counting(self): a = mx.arange(3) wr = weakref.ref(a) diff --git a/tests/cuda_tests.cpp b/tests/cuda_tests.cpp index aca39ca881..9afa980cd7 100644 --- a/tests/cuda_tests.cpp +++ b/tests/cuda_tests.cpp @@ -109,3 +109,13 @@ TEST_CASE("test clear cache trims CUDA pool") { cudaSuccess); CHECK_LT(final_reserved, allocated_reserved); } + +TEST_CASE("test eval after cuda error") { + auto s = default_stream(Device::gpu); + // A failed allocation sets a CUDA error on the calling thread. + CHECK_THROWS(eval(zeros({1 << 20, 1 << 20}, float32, s))); + // The error must not resurface in later evaluations. + auto a = ones({4}, float32, s); + CHECK_NOTHROW(eval(a)); + CHECK(array_equal(a, ones({4})).item()); +}