Prevent DeviceBuffer DeviceMemoryResource premature release

python/rmm/_lib/device_buffer.pxd

@@ -13,13 +13,19 @@
 # limitations under the License.
 
 from libc.stdint cimport uintptr_t
-from libcpp.memory cimport unique_ptr
+from libcpp.memory cimport shared_ptr, unique_ptr
 
 from rmm._cuda.stream cimport Stream
 from rmm._lib.cuda_stream_view cimport cuda_stream_view
 from rmm._lib.memory_resource cimport DeviceMemoryResource
 
 
+cdef extern from "rmm/mr/device/device_memory_resource.hpp" \
+        namespace "rmm::mr" nogil:
+    cdef cppclass device_memory_resource:
+        pass
+
+
 cdef extern from "rmm/device_buffer.hpp" namespace "rmm" nogil:
     cdef cppclass device_buffer:
         device_buffer()
@@ -36,10 +42,12 @@ cdef extern from "rmm/device_buffer.hpp" namespace "rmm" nogil:
 cdef class DeviceBuffer:
     cdef unique_ptr[device_buffer] c_obj
 
-    # Holds a reference to the DeviceMemoryResource used for allocation.
-    # Ensures the MR does not get destroyed before this DeviceBuffer. `mr` is
-    # needed for deallocation
-    cdef DeviceMemoryResource mr
+    # Holds a reference to the memory resource used for allocation
+    # and ensures that it is not destroyed before the device buffer
+    # is deallocated. A reference to a DeviceMemoryResource cannot
+    # be used as it could be released prematurely because of the circle
+    # reference breaking mechanism.
+    cdef shared_ptr[device_memory_resource] mr
 
     # Holds a reference to the stream used by the underlying `device_buffer`.
     # Ensures the stream does not get destroyed before this DeviceBuffer

python/rmm/_lib/device_buffer.pyx

@@ -89,7 +89,7 @@ cdef class DeviceBuffer:
                     stream.c_synchronize()
 
         # Save a reference to the MR and stream used for allocation
-        self.mr = get_current_device_resource()
+        self.mr = get_current_device_resource().c_obj
         self.stream = stream
 
     def __len__(self):

python/rmm/tests/test_rmm.py

@@ -691,3 +691,27 @@ def callback(nbytes: int) -> bool:
     with pytest.raises(MemoryError):
         rmm.DeviceBuffer(size=int(1e11))
     assert retried[0]
+
+
+def test_dev_buf_circle_ref_dealloc():
+    rmm.mr.set_current_device_resource(rmm.mr.CudaMemoryResource())
+
+    dbuf1 = rmm.DeviceBuffer(size=1_000_000)
+
+    # Make dbuf1 part of a reference cycle:
+    l1 = [dbuf1]
+    l1.append(l1)
+
+    # due to the reference cycle, the device buffer doesn't actually get
+    # cleaned up until later, when we invoke `gc.collect()`:
+    del dbuf1, l1
+
+    rmm.mr.set_current_device_resource(rmm.mr.CudaMemoryResource())
+
+    # by now, the only remaining reference to the *original* memory
+    # resource should be in `dbuf1`. However, the cyclic garbage collector
+    # will eliminate that reference when it clears the object via its
+    # `tp_clear` method.  Later, when `tp_dealloc` attemps to actually
+    # deallocate `dbuf1` (which needs the MR alive), a segfault occurs.
+
+    gc.collect()