Ensure UpstreamResourceAdaptor is not cleared by the Python GC

python/rmm/_lib/memory_resource.pyx

@@ -16,6 +16,7 @@ import os
 import warnings
 from collections import defaultdict
 
+cimport cython
 from cython.operator cimport dereference as deref
 from libc.stdint cimport int8_t, int64_t, uintptr_t
 from libcpp cimport bool
@@ -247,6 +248,8 @@ cdef class DeviceMemoryResource:
         self.c_obj.get().deallocate(<void*>(ptr), nbytes)
 
 
+# See the note about `no_gc_clear` in `device_buffer.pyx`.
+@cython.no_gc_clear
 cdef class UpstreamResourceAdaptor(DeviceMemoryResource):
 
     def __cinit__(self, DeviceMemoryResource upstream_mr, *args, **kwargs):

python/rmm/tests/test_rmm.py

@@ -742,6 +742,13 @@ def callback(nbytes: int) -> bool:
 
 
 def test_dev_buf_circle_ref_dealloc():
+    # This test creates a reference cycle containing a `DeviceBuffer`
+    # and ensure that the garbage collector does not clear it, i.e.,
+    # that the GC does not remove all references to other Python
+    # objects from it. The `DeviceBuffer` needs to keep its reference
+    # to the `DeviceMemoryResource` that was used to create it in
+    # order to be cleaned up properly. See GH #931.
+
     rmm.mr.set_current_device_resource(rmm.mr.CudaMemoryResource())
 
     dbuf1 = rmm.DeviceBuffer(size=1_000_000)
@@ -765,6 +772,20 @@ def test_dev_buf_circle_ref_dealloc():
     gc.collect()
 
 
+def test_upstream_mr_circle_ref_dealloc():
+    # This test is just like the one above, except it tests that
+    # instances of `UpstreamResourceAdaptor` (such as
+    # `PoolMemoryResource`) are not cleared by the GC.
+
+    rmm.mr.set_current_device_resource(rmm.mr.CudaMemoryResource())
+    mr = rmm.mr.PoolMemoryResource(rmm.mr.get_current_device_resource())
+    l1 = [mr]
+    l1.append(l1)
+    del mr, l1
+    rmm.mr.set_current_device_resource(rmm.mr.CudaMemoryResource())
+    gc.collect()
+
+
 def test_mr_allocate_deallocate():
     mr = rmm.mr.TrackingResourceAdaptor(rmm.mr.get_current_device_resource())
     size = 1 << 23  # 8 MiB