From 3759515edc5dd3b325c621f0d36cfffd052cd967 Mon Sep 17 00:00:00 2001
From: Leo Fang <leof@nvidia.com>
Date: Tue, 20 May 2025 17:14:02 +0000
Subject: [PATCH] add Buffer.release

---
 cuda_core/cuda/core/experimental/_memory.py   | 13 ++++-
 cuda_core/docs/source/release/0.3.0-notes.rst |  2 +
 cuda_core/tests/test_memory.py                | 47 ++++++++++++++++---
 3 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/cuda_core/cuda/core/experimental/_memory.py b/cuda_core/cuda/core/experimental/_memory.py
index e214da0f8..c29d725e8 100644
--- a/cuda_core/cuda/core/experimental/_memory.py
+++ b/cuda_core/cuda/core/experimental/_memory.py
@@ -45,13 +45,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:
@@ -83,6 +83,15 @@ def close(self, 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.
diff --git a/cuda_core/docs/source/release/0.3.0-notes.rst b/cuda_core/docs/source/release/0.3.0-notes.rst
index dea108977..21ab2f57d 100644
--- a/cuda_core/docs/source/release/0.3.0-notes.rst
+++ b/cuda_core/docs/source/release/0.3.0-notes.rst
@@ -21,6 +21,8 @@ New features
 ------------
 
 - :class:`Kernel` adds :property:`Kernel.num_arguments` and :property:`Kernel.arguments_info` for introspection of kernel arguments. (#612)
+- :class:`Buffer` adds a method :meth:`Buffer.release` allowing users to have full control over its lifetime without the garbage collector's
+  interference
 
 New examples
 ------------
diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py
index 5bcc607da..43f2a1ae4 100644
--- a/cuda_core/tests/test_memory.py
+++ b/cuda_core/tests/test_memory.py
@@ -1,18 +1,14 @@
 # Copyright 2024 NVIDIA Corporation.  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 Device
 from cuda.core.experimental._memory import Buffer, DLDeviceType, MemoryResource
-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