Skip to content

Commit 90abbed

Browse files
srimanachantaclaude
andcommitted
cuda.core: raise CUDAOutOfMemoryError, a MemoryError, on out-of-memory
CUDA_ERROR_OUT_OF_MEMORY and cudaErrorMemoryAllocation now raise CUDAOutOfMemoryError, a subclass of both CUDAError and MemoryError, so `except MemoryError` works for cuda.core allocations as it does for CuPy, RMM and NumPy. Code that catches CUDAError is unaffected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent d4c71ef commit 90abbed

8 files changed

Lines changed: 109 additions & 10 deletions

File tree

‎cuda_core/cuda/core/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -168,10 +168,11 @@ class _PatchedProperty(metaclass=_PatchedPropMeta):
168168
from cuda.core._stream import __all__ as _stream_all
169169
from cuda.core._tensor_map import *
170170
from cuda.core._tensor_map import __all__ as _tensor_map_all
171-
from cuda.core._utils.cuda_utils import CUDAError, CUDAWarning, NVRTCError
171+
from cuda.core._utils.cuda_utils import CUDAError, CUDAOutOfMemoryError, CUDAWarning, NVRTCError
172172

173173
__all__ = [
174174
"CUDAError",
175+
"CUDAOutOfMemoryError",
175176
"CUDAWarning",
176177
"NVRTCError",
177178
*_context_all,

‎cuda_core/cuda/core/_utils/cuda_utils.pyi‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,17 @@ class CUDAError(Exception):
2323
documentation for the guarantees an exception provides.
2424
"""
2525

26+
class CUDAOutOfMemoryError(CUDAError, MemoryError):
27+
"""Raised when a CUDA driver or runtime call fails for lack of memory.
28+
29+
``CUDA_ERROR_OUT_OF_MEMORY`` and ``cudaErrorMemoryAllocation`` raise this
30+
class. It is both a :class:`CUDAError` and a :class:`MemoryError`, so code
31+
that catches :class:`MemoryError` to retry with a smaller request handles
32+
``cuda.core`` the same way as other array libraries.
33+
34+
.. versionadded:: 1.3.0
35+
"""
36+
2637
class CUDAWarning(RuntimeWarning):
2738
"""Warning issued when ``cuda.core`` hits a CUDA error it cannot raise.
2839

‎cuda_core/cuda/core/_utils/cuda_utils.pyx‎

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,18 @@ class CUDAError(Exception):
4545
"""
4646

4747

48+
class CUDAOutOfMemoryError(CUDAError, MemoryError):
49+
"""Raised when a CUDA driver or runtime call fails for lack of memory.
50+
51+
``CUDA_ERROR_OUT_OF_MEMORY`` and ``cudaErrorMemoryAllocation`` raise this
52+
class. It is both a :class:`CUDAError` and a :class:`MemoryError`, so code
53+
that catches :class:`MemoryError` to retry with a smaller request handles
54+
``cuda.core`` the same way as other array libraries.
55+
56+
.. versionadded:: 1.3.0
57+
"""
58+
59+
4860
class CUDAWarning(RuntimeWarning):
4961
"""Warning issued when ``cuda.core`` hits a CUDA error it cannot raise.
5062
@@ -163,6 +175,7 @@ cdef int _raise_nvjitlink_error(
163175

164176

165177
_RUNTIME_SUCCESS = runtime.cudaError_t.cudaSuccess
178+
cdef object _RUNTIME_OUT_OF_MEMORY = runtime.cudaError_t.cudaErrorMemoryAllocation
166179
_NVRTC_SUCCESS = nvrtc.nvrtcResult.NVRTC_SUCCESS
167180

168181

@@ -182,14 +195,18 @@ cdef int _raise_driver_error(cydriver.CUresult error) except -1 nogil:
182195
raise CUDAError(f"UNEXPECTED ERROR CODE: {error}")
183196
desc_err = cydriver.cuGetErrorString(error, &desc)
184197
with gil:
198+
exc_type = (
199+
CUDAOutOfMemoryError if error == cydriver.CUresult.CUDA_ERROR_OUT_OF_MEMORY
200+
else CUDAError
201+
)
185202
# TODO: consider lower this to Cython
186203
expl = DRIVER_CU_RESULT_EXPLANATIONS.get(int(error))
187204
if expl is not None:
188-
exc = CUDAError(f"{name.decode()}: {expl}")
205+
exc = exc_type(f"{name.decode()}: {expl}")
189206
elif desc_err != cydriver.CUresult.CUDA_SUCCESS:
190-
exc = CUDAError(name.decode())
207+
exc = exc_type(name.decode())
191208
else:
192-
exc = CUDAError(f"{name.decode()}: {desc.decode()}")
209+
exc = exc_type(f"{name.decode()}: {desc.decode()}")
193210
if detail != NULL:
194211
_attach_detail(exc, detail.decode())
195212
raise exc
@@ -199,14 +216,15 @@ cdef int _raise_runtime_error(error) except -1:
199216
# `_check_error()` reaches this path only for `runtime.cudaError_t` values.
200217
# Use the enum name directly because Windows hybrid cudart can lag that table.
201218
name = error.name
219+
exc_type = CUDAOutOfMemoryError if error == _RUNTIME_OUT_OF_MEMORY else CUDAError
202220
expl = RUNTIME_CUDA_ERROR_EXPLANATIONS.get(int(error))
203221
if expl is not None:
204-
raise CUDAError(f"{name}: {expl}")
222+
raise exc_type(f"{name}: {expl}")
205223
desc_err, desc = runtime.cudaGetErrorString(error)
206224
if desc_err != _RUNTIME_SUCCESS:
207-
raise CUDAError(f"{name}")
225+
raise exc_type(f"{name}")
208226
desc = desc.decode()
209-
raise CUDAError(f"{name}: {desc}")
227+
raise exc_type(f"{name}: {desc}")
210228

211229

212230
cdef int _raise_nvrtc_status_error(error, handle) except -1:

‎cuda_core/docs/source/api.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -326,6 +326,7 @@ reported as a warning instead.
326326
:toctree: generated/
327327

328328
CUDAError
329+
CUDAOutOfMemoryError
329330
NVRTCError
330331
CUDAWarning
331332

‎cuda_core/docs/source/error_handling.rst‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,12 @@ A CUDA driver, runtime, NVRTC, NVVM or nvJitLink call that fails raises an
1818
exception whose message contains the error name and its description:
1919
:class:`CUDAError` for driver and runtime failures, its subclass
2020
:class:`NVRTCError` for NVRTC failures, and the ``cuda.bindings`` error types
21-
for NVVM and nvJitLink failures. Both ``cuda.core`` classes are importable from
22-
the top-level ``cuda.core`` namespace. Invalid arguments and misuse raise the
23-
usual Python exception types (``TypeError``, ``ValueError``, ``RuntimeError``).
21+
for NVVM and nvJitLink failures. A driver or runtime call that runs out of
22+
memory raises :class:`CUDAOutOfMemoryError`, which derives from both
23+
:class:`CUDAError` and :class:`MemoryError`. These exception classes are
24+
importable from the top-level ``cuda.core`` namespace. Invalid arguments and
25+
misuse raise the usual Python exception types (``TypeError``, ``ValueError``,
26+
``RuntimeError``).
2427

2528
When a ``cuda.core`` call raises, the following hold:
2629

‎cuda_core/docs/source/release/1.3.0-notes.rst‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,12 @@ New features
7676
top-level ``cuda.core`` namespace, alongside :class:`CUDAWarning`. Previously
7777
they were only available from a private module.
7878

79+
- Added :class:`CUDAOutOfMemoryError`, raised for ``CUDA_ERROR_OUT_OF_MEMORY``
80+
and ``cudaErrorMemoryAllocation``. It derives from both :class:`CUDAError`
81+
and :class:`MemoryError`, so ``except MemoryError`` catches an allocation
82+
that fails in ``cuda.core`` as it does in CuPy, RMM and NumPy. Code that
83+
catches :class:`CUDAError` is unaffected.
84+
7985
- Added :class:`VirtualMemoryBuffer`, the :class:`Buffer` subclass that
8086
:class:`VirtualMemoryResource` returns. Closing it on a capturing stream
8187
other than a default stream raises, because virtual memory deallocation

‎cuda_core/tests/test_cuda_utils.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,40 @@ def test_check_runtime_error():
5050
assert error.name in msg
5151

5252

53+
@pytest.mark.agent_authored(model="claude-opus-5-5")
54+
def test_check_driver_error_out_of_memory_is_memory_error():
55+
with pytest.raises(cuda_utils.CUDAOutOfMemoryError) as e:
56+
cuda_utils._check_driver_error(driver.CUresult.CUDA_ERROR_OUT_OF_MEMORY)
57+
assert isinstance(e.value, cuda_utils.CUDAError)
58+
assert isinstance(e.value, MemoryError)
59+
assert "CUDA_ERROR_OUT_OF_MEMORY" in str(e.value)
60+
61+
62+
@pytest.mark.agent_authored(model="claude-opus-5-5")
63+
def test_check_runtime_error_out_of_memory_is_memory_error():
64+
with pytest.raises(cuda_utils.CUDAOutOfMemoryError) as e:
65+
cuda_utils._check_runtime_error(runtime.cudaError_t.cudaErrorMemoryAllocation)
66+
assert isinstance(e.value, cuda_utils.CUDAError)
67+
assert isinstance(e.value, MemoryError)
68+
assert "cudaErrorMemoryAllocation" in str(e.value)
69+
70+
71+
@pytest.mark.agent_authored(model="claude-opus-5-5")
72+
def test_only_out_of_memory_errors_are_memory_errors():
73+
for error in driver.CUresult:
74+
if error in (driver.CUresult.CUDA_SUCCESS, driver.CUresult.CUDA_ERROR_OUT_OF_MEMORY):
75+
continue
76+
with pytest.raises(cuda_utils.CUDAError) as e:
77+
cuda_utils._check_driver_error(error)
78+
assert not isinstance(e.value, MemoryError), error
79+
for error in runtime.cudaError_t:
80+
if error in (runtime.cudaError_t.cudaSuccess, runtime.cudaError_t.cudaErrorMemoryAllocation):
81+
continue
82+
with pytest.raises(cuda_utils.CUDAError) as e:
83+
cuda_utils._check_runtime_error(error)
84+
assert not isinstance(e.value, MemoryError), error
85+
86+
5387
def test_driver_error_enum_has_non_empty_docstring():
5488

5589
doc = driver.CUresult.CUDA_ERROR_INVALID_VALUE.__doc__

‎cuda_core/tests/test_memory.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444

4545
from cuda.core import (
4646
Buffer,
47+
CUDAOutOfMemoryError,
4748
CUDAWarning,
4849
Device,
4950
DeviceMemoryResource,
@@ -2223,6 +2224,30 @@ def test_vmm_host_location_grow_and_close(init_cuda):
22232224
assert moved.size == 0
22242225

22252226

2227+
@pytest.mark.agent_authored(model="claude-opus-5-5")
2228+
def test_device_memory_resource_out_of_memory_is_memory_error(init_cuda):
2229+
device = Device()
2230+
if not device.properties.memory_pools_supported:
2231+
pytest.skip("Device does not support mempool operations")
2232+
device.set_current()
2233+
2234+
mr = DeviceMemoryResource(device, options=DeviceMemoryResourceOptions(max_size=POOL_SIZE))
2235+
try:
2236+
# Larger than any device, and than the pool's max_size, which the
2237+
# driver rounds up to its allocation granularity.
2238+
with pytest.raises(MemoryError) as e:
2239+
mr.allocate(1 << 46, stream=device.default_stream)
2240+
assert isinstance(e.value, CUDAOutOfMemoryError)
2241+
assert isinstance(e.value, CUDAError)
2242+
assert "CUDA_ERROR_OUT_OF_MEMORY" in str(e.value)
2243+
2244+
# The failed request leaves the pool usable.
2245+
buffer = mr.allocate(POOL_SIZE // 2, stream=device.default_stream)
2246+
buffer.close()
2247+
finally:
2248+
mr.close()
2249+
2250+
22262251
def test_device_memory_resource_with_options(init_cuda):
22272252
device = Device()
22282253
if not device.properties.memory_pools_supported:

0 commit comments

Comments
 (0)