Avoid race due to deleted imports when cleaning up objects at interpreter shutdown

cuda_bindings/cuda/bindings/utils/__init__.py

@@ -1,6 +1,5 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
-
 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,6 +109,14 @@ cdef class Event:
         self._ctx_handle = ctx_handle
         return self
 
+    cpdef safe_close(self, is_shutting_down=is_shutting_down):
+        """Destroy the event."""
+        if self._handle is not None:
+            if not is_shutting_down():
+                err, = driver.cuEventDestroy(self._handle)
+            self._handle = None
+            raise_if_driver_error(err)
+
     cpdef close(self):
         """Destroy the event."""
         if self._handle is not None:
@@ -116,7 +125,7 @@ cdef class Event:
             raise_if_driver_error(err)
 
     def __del__(self):
-        self.close()
+        self.safe_close()
 
     def __isub__(self, other):
         return NotImplemented

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
@@ -57,7 +57,15 @@ cdef class Buffer:
         return self
 
     def __del__(self):
-        self.close()
+        self.safe_close()
+
+    cpdef safe_close(self, stream: Stream = None, is_shutting_down=is_shutting_down):
+        if self._ptr and self._mr is not None:
+            if not is_shutting_down():
+                self._mr.deallocate(self._ptr, self._size, stream)
+            self._ptr = 0
+            self._mr = None
+            self._ptr_obj = None
 
     cpdef close(self, stream: Stream = None):
         """Deallocate this buffer asynchronously on the given stream.

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
 )
 
 
@@ -185,9 +186,18 @@ cdef class Stream:
         return self
 
     def __del__(self):
-        self.close()
+        self.safe_close()
 
-    cpdef close(self):
+    cpdef safe_close(self, is_shutting_down=is_shutting_down):
+        if self._owner is None:
+            if self._handle and not self._builtin:
+                if not is_shutting_down():
+                    handle_return(driver.cuStreamDestroy(self._handle))
+        else:
+            self._owner = None
+        self._handle = None
+
+    cpdef close(self, is_shutting_down=is_shutting_down):
         """Destroy the stream.
 
         Destroy the stream if we own it. Borrowed foreign stream

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]