Add Python bindings for limiting_resource_adaptor

python/rmm/_lib/memory_resource.pxd

@@ -69,6 +69,9 @@ cdef class CallbackMemoryResource(DeviceMemoryResource):
     cdef object _allocate_func
     cdef object _deallocate_func
 
+cdef class LimitingResourceAdaptor(UpstreamResourceAdaptor):
+    pass
+
 cdef class LoggingResourceAdaptor(UpstreamResourceAdaptor):
     cdef object _log_file_name
     cpdef get_file_name(self)

python/rmm/_lib/memory_resource.pyx

@@ -164,6 +164,16 @@ cdef extern from "rmm/mr/device/binning_memory_resource.hpp" \
             size_t allocation_size,
             device_memory_resource* bin_resource) except +
 
+cdef extern from "rmm/mr/device/limiting_resource_adaptor.hpp" \
+        namespace "rmm::mr" nogil:
+    cdef cppclass limiting_resource_adaptor[Upstream](device_memory_resource):
+        limiting_resource_adaptor(
+            Upstream* upstream_mr,
+            size_t allocation_limit) except +
+
+        size_t get_allocated_bytes() except +
+        size_t get_allocation_limit() except +
+
 cdef extern from "rmm/mr/device/logging_resource_adaptor.hpp" \
         namespace "rmm::mr" nogil:
     cdef cppclass logging_resource_adaptor[Upstream](device_memory_resource):
@@ -649,6 +659,61 @@ def _append_id(filename, id):
     return f"{name}.dev{id}{ext}"
 
 
+cdef class LimitingResourceAdaptor(UpstreamResourceAdaptor):
+
+    def __cinit__(
+        self,
+        DeviceMemoryResource upstream_mr,
+        size_t allocation_limit
+    ):
+        self.c_obj.reset(
+            new limiting_resource_adaptor[device_memory_resource](
+                upstream_mr.get_mr(),
+                allocation_limit
+            )
+        )
+
+    def __init__(
+        self,
+        DeviceMemoryResource upstream_mr,
+        size_t allocation_limit
+    ):
+        """
+        Memory resource that limits the total allocation amount possible
+        performed by an upstream memory resource.
+
+        Parameters
+        ----------
+        upstream_mr : DeviceMemoryResource
+            The upstream memory resource.
+        allocation_limit : size_t
+            Maximum memory allowed for this allocator.
+        """
+        pass
+
+    def get_allocated_bytes(self) -> size_t:
+        """
+        Query the number of bytes that have been allocated. Note that this can
+        not be used to know how large of an allocation is possible due to both
+        possible fragmentation and also internal page sizes and alignment that
+        is not tracked by this allocator.
+        """
+        return (<limiting_resource_adaptor[device_memory_resource]*>(
+            self.c_obj.get())
+        )[0].get_allocated_bytes()
+
+    def get_allocation_limit(self) -> size_t:
+        """
+        Query the maximum number of bytes that this allocator is allowed to
+        allocate. This is the limit on the allocator and not a representation
+        of the underlying device. The device may not be able to support this
+        limit.
+        """
+        return (<limiting_resource_adaptor[device_memory_resource]*>(
+            self.c_obj.get())
+        )[0].get_allocation_limit()
+
+
 cdef class LoggingResourceAdaptor(UpstreamResourceAdaptor):
     def __cinit__(
         self,

python/rmm/mr.py

@@ -19,6 +19,7 @@
     DeviceMemoryResource,
     FailureCallbackResourceAdaptor,
     FixedSizeMemoryResource,
+    LimitingResourceAdaptor,
     LoggingResourceAdaptor,
     ManagedMemoryResource,
     PoolMemoryResource,
@@ -46,6 +47,7 @@
     "CudaMemoryResource",
     "DeviceMemoryResource",
     "FixedSizeMemoryResource",
+    "LimitingResourceAdaptor",
     "LoggingResourceAdaptor",
     "ManagedMemoryResource",
     "PoolMemoryResource",

python/rmm/tests/test_rmm.py

@@ -586,6 +586,41 @@ def test_cuda_async_memory_resource_threshold(nelem, alloc):
     array_tester("u1", 2 * nelem, alloc)  # should trigger release
 
 
+@pytest.mark.parametrize(
+    "mr",
+    [
+        rmm.mr.CudaMemoryResource,
+        pytest.param(
+            rmm.mr.CudaAsyncMemoryResource,
+            marks=pytest.mark.skipif(
+                not _CUDAMALLOC_ASYNC_SUPPORTED,
+                reason="cudaMallocAsync not supported",
+            ),
+        ),
+    ],
+)
+def test_limiting_resource_adaptor(mr):
+    cuda_mr = mr()
+
+    allocation_limit = 1 << 20
+    num_buffers = 2
+    buffer_size = allocation_limit / num_buffers
+
+    mr = rmm.mr.LimitingResourceAdaptor(
+        cuda_mr, allocation_limit=allocation_limit
+    )
+    assert mr.get_allocation_limit() == allocation_limit
+
+    rmm.mr.set_current_device_resource(mr)
+
+    buffers = [rmm.DeviceBuffer(size=buffer_size) for _ in range(num_buffers)]
+
+    assert mr.get_allocated_bytes() == buffer_size * num_buffers
+
+    with pytest.raises(MemoryError):
+        rmm.DeviceBuffer(size=1)
+
+
 def test_statistics_resource_adaptor(stats_mr):
 
     buffers = [rmm.DeviceBuffer(size=1000) for _ in range(10)]