You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[FEA]: Raise a MemoryError subclass on CUDA out-of-memory errors #3260
Is your feature request related to a problem? Please describe.
Out-of-memory errors from cuda.core raise CUDAError, which derives only from Exception. CuPy (cupy.cuda.memory.OutOfMemoryError) and RMM raise a MemoryError subclass when device memory runs out, and NumPy raises one for host allocations, so code that retries with a smaller request catches MemoryError. CUDAError has no error code attribute, so the only way to catch an out-of-memory error from cuda.core today is to match CUDA_ERROR_OUT_OF_MEMORY in the message.
Use case: A CUDA-accelerated simulation library I am developing, cuNIBS halves its chunk size on except MemoryError (simulation.py). Its native code raises a MemoryError subclass for this reason (bindings.cpp). Out-of-memory errors from cuda.core bypass the retry.
Describe the solution you'd like
Add cuda.core.CUDAOutOfMemoryError(CUDAError, MemoryError) and raise it for CUDA_ERROR_OUT_OF_MEMORY and cudaErrorMemoryAllocation. Every driver and runtime status that can carry these codes goes through _check_driver_error or _check_runtime_error, so the change is limited to those two functions. The message stays the same, and except CUDAError still catches it. Nothing in cuda.core, cuda.bindings or cuda.pathfinder catches MemoryError.
NVRTC's NVRTC_ERROR_OUT_OF_MEMORY is a host allocation failure in the compiler and still raises NVRTCError. It can follow separately.
I opened #3235 with this change before opening an issue. I'll revise or close it depending on the outcome here.
Describe alternatives you've considered
An error code attribute on CUDAError would let callers check for out-of-memory, but code that catches MemoryError still wouldn't see it. PyTorch raises torch.OutOfMemoryError, a RuntimeError, which has the same problem.
Additional context
Related: #2761 discusses a dedicated exception type for sticky errors. Its alternative, a CUDAError subclass chosen from a set of error codes in _check_driver_error, is the mechanism used here.
Activity