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 unique_ptr, shared_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,10 @@ 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 device_memory_resource used for allocation.
+    # Ensures the memory resource does not get destroyed before
+    # this DeviceBuffer as it is needed for deallocation.
+    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().get()
         self.stream = stream
 
     def __len__(self):

python/rmm/_lib/memory_resource.pxd

@@ -26,6 +26,7 @@ cdef extern from "rmm/mr/device/device_memory_resource.hpp" \
 cdef class DeviceMemoryResource:
     cdef shared_ptr[device_memory_resource] c_obj
 
+    cdef shared_ptr[device_memory_resource] get(self)
     cdef device_memory_resource* get_mr(self)
 
 cdef class UpstreamResourceAdaptor(DeviceMemoryResource):

python/rmm/_lib/memory_resource.pyx

@@ -162,6 +162,9 @@ cdef extern from "rmm/mr/device/per_device_resource.hpp" namespace "rmm" nogil:
 
 cdef class DeviceMemoryResource:
 
+    cdef shared_ptr[device_memory_resource] get(self):
+        return self.c_obj
+
     cdef device_memory_resource* get_mr(self):
         return self.c_obj.get()