Adapt to rapidsai/rmm#1221 which moves allocator callbacks

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 · eb0c7f5 · eb0c7f5
1 parent d399edd
commit eb0c7f5
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/python/cuml/internals/api_context_managers.py b/python/cuml/internals/api_context_managers.py
@@ -40,7 +40,9 @@
 cupy_using_allocator = gpu_only_import_from(
     "cupy.cuda", "using_allocator", alt=UnavailableNullContext
 )
-rmm_cupy_allocator = gpu_only_import_from("rmm", "rmm_cupy_allocator")
+rmm_cupy_allocator = gpu_only_import_from(
+    "rmm.allocators.cupy", "rmm_cupy_allocator"
+)
 
 
 @contextlib.contextmanager

diff --git a/python/cuml/internals/memory_utils.py b/python/cuml/internals/memory_utils.py
@@ -39,7 +39,9 @@
 )
 PandasSeries = cpu_only_import_from("pandas", "Series")
 PandasDataFrame = cpu_only_import_from("pandas", "DataFrame")
-rmm_cupy_allocator = gpu_only_import_from("rmm", "rmm_cupy_allocator")
+rmm_cupy_allocator = gpu_only_import_from(
+    "rmm.allocators.cupy", "rmm_cupy_allocator"
+)
 
 
 def set_global_memory_type(memory_type):

diff --git a/wiki/python/ESTIMATOR_GUIDE.md b/wiki/python/ESTIMATOR_GUIDE.md
@@ -706,7 +706,9 @@ Functionally, you can think of these decorators equivalent to the following pseu
 ```python
 def my_func(self, X):
    with cuml.using_ouput_type("mirror"):
-      with cupy.cuda.cupy_using_allocator(rmm.rmm_cupy_allocator):
+      with cupy.cuda.cupy_using_allocator(
+          rmm.allocators.cupy.rmm_cupy_allocator
+      ):
          # Set the input properties
          self._set_base_attributes(output_type=X, n_features=X)