From d071e21c8157bc3cd9599b53a1dd3d47d42d9f0f Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 6 Mar 2023 16:07:41 +0000 Subject: [PATCH] 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: https://github.com/rapidsai/cuml/pull/5249 --- python/cuml/internals/api_context_managers.py | 4 +++- python/cuml/internals/memory_utils.py | 4 +++- wiki/python/ESTIMATOR_GUIDE.md | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/python/cuml/internals/api_context_managers.py b/python/cuml/internals/api_context_managers.py index 3d043f66ab..6be137d5b2 100644 --- 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 index b3a2870965..921c17dac6 100644 --- 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 index d006286a6c..a3832f8def 100644 --- 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)