Add Buffer.release()

cuda_core/cuda/core/experimental/_memory.py

@@ -34,13 +34,13 @@ class Buffer:
     """
 
     class _MembersNeededForFinalize:
-        __slots__ = ("ptr", "size", "mr")
+        __slots__ = ("ptr", "size", "mr", "finalizer")
 
         def __init__(self, buffer_obj, ptr, size, mr):
             self.ptr = ptr
             self.size = size
             self.mr = mr
-            weakref.finalize(buffer_obj, self.close)
+            self.finalizer = weakref.finalize(buffer_obj, self.close)
 
         def close(self, stream=None):
             if self.ptr and self.mr is not None:
@@ -74,6 +74,15 @@ def close(self, stream: Stream = None):
         """
         self._mnff.close(stream)
 
+    def release(self):
+        """Release this buffer from being subject to the garbage collection.
+
+        After this method is called, the caller is responsible for calling :meth:`close`
+        when done using this buffer, otherwise there would be a memory leak!
+        """
+        self._mnff.finalizer.detach()
+        self._mnff.finalizer = None
+
     @property
     def handle(self) -> DevicePointerT:
         """Return the buffer handle object.

cuda_core/docs/source/release/0.4.0-notes.rst

@@ -0,0 +1,31 @@
+.. SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-License-Identifier: Apache-2.0
+
+.. currentmodule:: cuda.core.experimental
+
+``cuda.core`` 0.4.0 Release Notes
+=================================
+
+Released on MM DD, 2025
+
+
+Highlights
+----------
+
+
+Breaking Changes
+----------------
+
+
+New features
+------------
+
+- :class:`Buffer` adds a method :meth:`Buffer.release` allowing users to have full control over its lifetime without interfered by the Python garbage collector
+
+
+New examples
+------------
+
+
+Fixes and enhancements
+----------------------

cuda_core/tests/test_memory.py

@@ -1,18 +1,14 @@
 # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-try:
-    from cuda.bindings import driver
-except ImportError:
-    from cuda import cuda as driver
-
 import ctypes
+import gc
 
 import pytest
 
 from cuda.core.experimental import Buffer, Device, MemoryResource
 from cuda.core.experimental._memory import DLDeviceType
-from cuda.core.experimental._utils.cuda_utils import handle_return
+from cuda.core.experimental._utils.cuda_utils import driver, handle_return
 
 
 class DummyDeviceMemoryResource(MemoryResource):
@@ -257,3 +253,42 @@ def test_buffer_dunder_dlpack_device_failure():
     buffer = dummy_mr.allocate(size=1024)
     with pytest.raises(BufferError, match=r"^buffer is neither device-accessible nor host-accessible$"):
         buffer.__dlpack_device__()
+
+
+class DummyTrackingMemoryResource(MemoryResource):
+    def __init__(self):
+        self.live_counts = 0
+
+    def allocate(self, size, stream=None) -> Buffer:
+        self.live_counts += 1
+        return Buffer(ptr=1, size=size, mr=self)
+
+    def deallocate(self, ptr, size, stream=None):
+        self.live_counts -= 1
+
+    @property
+    def is_device_accessible(self) -> bool:
+        return False
+
+    @property
+    def is_host_accessible(self) -> bool:
+        return False
+
+    @property
+    def device_id(self) -> int:
+        return 0
+
+
+@pytest.mark.parametrize("close_buffer", (True, False))
+def test_buffer_release_closed(close_buffer):
+    mr = DummyTrackingMemoryResource()
+    buf = mr.allocate(123)
+    buf.release()
+    if close_buffer:
+        buf.close()
+    del buf
+    gc.collect()
+    # If Buffer.close() is explicitly called, it would end up calling MR.deallocate()
+    # which then decrease the count. If the buffer does not have the finalizer attached,
+    # this is the only way to trigger deallocation.
+    assert mr.live_counts == 0 if close_buffer else 1