Adapt to rapidsai/rmm#1221 which moves allocator callbacks (#5249)

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.

Authors:
  - Lawrence Mitchell (https://github.com/wence-)

Approvers:
  - Carl Simon Adorf (https://github.com/csadorf)
  - William Hicks (https://github.com/wphicks)

URL: #5249

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

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):

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)