is_shutting_down

Author: brandon-b-miller

cuda_bindings/cuda/bindings/utils/__init__.py

@@ -1,6 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
-
+import atexit
 from typing import Any, Callable
 
 from ._ptx_utils import get_minimal_required_cuda_ver_from_ptx_ver, get_ptx_ver

cuda_core/cuda/core/experimental/_event.pyx

@@ -17,6 +17,7 @@ from cuda.core.experimental._utils.cuda_utils import (
     CUDAError,
     driver,
     handle_return,
+    is_shutting_down
 )
 
 if TYPE_CHECKING:
@@ -108,10 +109,11 @@ cdef class Event:
         self._ctx_handle = ctx_handle
         return self
 
-    cpdef close(self):
+    cpdef close(self, is_shutting_down=is_shutting_down):
         """Destroy the event."""
         if self._handle is not None:
-            err, = driver.cuEventDestroy(self._handle)
+            if not is_shutting_down()
+                err, = driver.cuEventDestroy(self._handle)
             self._handle = None
             raise_if_driver_error(err)
 

cuda_core/cuda/core/experimental/_memory.pyx

@@ -15,7 +15,7 @@ from typing import TypeVar, Union
 
 from cuda.core.experimental._dlpack import DLDeviceType, make_py_capsule
 from cuda.core.experimental._stream import Stream, default_stream
-from cuda.core.experimental._utils.cuda_utils import driver
+from cuda.core.experimental._utils.cuda_utils import driver, is_shutting_down
 
 # TODO: define a memory property mixin class and make Buffer and
 # MemoryResource both inherit from it
@@ -59,7 +59,7 @@ cdef class Buffer:
     def __del__(self):
         self.close()
 
-    cpdef close(self, stream: Stream = None):
+    cpdef close(self, stream: Stream = None, is_shutting_down=is_shutting_down):
         """Deallocate this buffer asynchronously on the given stream.
 
         This buffer is released back to their memory resource
@@ -72,7 +72,8 @@ cdef class Buffer:
             the behavior depends on the underlying memory resource.
         """
         if self._ptr and self._mr is not None:
-            self._mr.deallocate(self._ptr, self._size, stream)
+            if not is_shutting_down():
+                self._mr.deallocate(self._ptr, self._size, stream)
             self._ptr = 0
             self._mr = None
             self._ptr_obj = None

cuda_core/cuda/core/experimental/_stream.pyx

@@ -25,6 +25,7 @@ from cuda.core.experimental._utils.cuda_utils import (
     driver,
     get_device_from_ctx,
     handle_return,
+    is_shutting_down
 )
 
 
@@ -187,7 +188,7 @@ cdef class Stream:
     def __del__(self):
         self.close()
 
-    cpdef close(self):
+    cpdef close(self, is_shutting_down=is_shutting_down):
         """Destroy the stream.
 
         Destroy the stream if we own it. Borrowed foreign stream
@@ -196,7 +197,8 @@ cdef class Stream:
         """
         if self._owner is None:
             if self._handle and not self._builtin:
-                handle_return(driver.cuStreamDestroy(self._handle))
+                if not is_shutting_down():
+                    handle_return(driver.cuStreamDestroy(self._handle))
         else:
             self._owner = None
         self._handle = None

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

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import atexit
 import functools
 import importlib.metadata
 from collections import namedtuple
@@ -222,3 +223,25 @@ def get_binding_version():
     except importlib.metadata.PackageNotFoundError:
         major_minor = importlib.metadata.version("cuda-python").split(".")[:2]
     return tuple(int(v) for v in major_minor)
+
+# This code is to signal when the interpreter is in shutdown mode
+# to prevent using globals that could be already deleted in
+# objects `__del__` method
+#
+# This solution is taken from the Numba/llvmlite code
+_shutting_down = [False]
+
+
+@atexit.register
+def _at_shutdown():
+    _shutting_down[0] = True
+
+
+def is_shutting_down(_shutting_down=_shutting_down):
+    """
+    Whether the interpreter is currently shutting down.
+    For use in finalizers, __del__ methods, and similar; it is advised
+    to early bind this function rather than look it up when calling it,
+    since at shutdown module globals may be cleared.
+    """
+    return _shutting_down[0]