Adapt to github.com/rapidsai/rmm/pull/1221 which moves allocator call…

…backs The allocator callbacks now live in their own submodules (so that RMM does not, for example, import pytorch unless required) and so must be explicitly imported.
rapidsai · Feb 24, 2023 · 90b2bf6 · 90b2bf6
1 parent 75b6d80
commit 90b2bf6
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/python/examples/basic.py b/python/examples/basic.py
@@ -117,13 +117,14 @@ def main():
         import cupy as xp
 
         import rmm
+        from rmm.allocators.cupy import rmm_cupy_allocator
 
         rmm.reinitialize(
             pool_allocator=True,
             managed_memory=False,
         )
         xp.cuda.runtime.setDevice(0)
-        xp.cuda.set_allocator(rmm.rmm_cupy_allocator)
+        xp.cuda.set_allocator(rmm_cupy_allocator)
     else:
         import numpy as xp
 

diff --git a/python/ucxx/benchmarks/cudf_merge.py b/python/ucxx/benchmarks/cudf_merge.py
@@ -242,10 +242,12 @@ def _get_worker_command(
 
 async def worker(rank, eps, args):
     # Setting current device and make RMM use it
+    from rmm.allocators.cupy import rmm_cupy_allocator
+
     rmm.reinitialize(pool_allocator=True, initial_pool_size=args.rmm_init_pool_size)
 
     # Make cupy use RMM
-    cupy.cuda.set_allocator(rmm.rmm_cupy_allocator)
+    cupy.cuda.set_allocator(rmm_cupy_allocator)
 
     df1 = generate_chunk(rank, args.chunk_size, args.n_chunks, "build", args.frac_match)
     df2 = generate_chunk(rank, args.chunk_size, args.n_chunks, "other", args.frac_match)

diff --git a/python/ucxx/benchmarks/utils.py b/python/ucxx/benchmarks/utils.py
@@ -58,13 +58,14 @@ def get_allocator(
         import cupy as xp
 
         import rmm
+        from rmm.allocators.cupy import rmm_cupy_allocator
 
         rmm.reinitialize(
             pool_allocator=True,
             managed_memory=rmm_managed_memory,
             initial_pool_size=rmm_init_pool_size,
         )
-        xp.cuda.set_allocator(rmm.rmm_cupy_allocator)
+        xp.cuda.set_allocator(rmm_cupy_allocator)
 
     return xp