From ea4aa9aa84812505650e698471daceb864d0b450 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 28 Feb 2023 00:46:31 -0800 Subject: [PATCH] Update RMM allocators Following https://github.com/rapidsai/rmm/pull/1221, RMM allocators now live under the `rmm.allocators` submodule and using old paths causes warnings to be raised. This change updates to the new submodule and prevents such warnings. --- python/cudf/cudf/__init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/cudf/cudf/__init__.py b/python/cudf/cudf/__init__.py index 04b64e18594..7bab131a85a 100644 --- a/python/cudf/cudf/__init__.py +++ b/python/cudf/cudf/__init__.py @@ -8,6 +8,8 @@ from numba import config as numba_config, cuda import rmm +from rmm.allocators.cupy import rmm_cupy_allocator +from rmm.allocators.numba import RMMNumbaManager from cudf import api, core, datasets, testing from cudf.api.extensions import ( @@ -96,8 +98,8 @@ del patch_numba_linker_if_needed -cuda.set_memory_manager(rmm.RMMNumbaManager) -cupy.cuda.set_allocator(rmm.rmm_cupy_allocator) +cuda.set_memory_manager(RMMNumbaManager) +cupy.cuda.set_allocator(rmm_cupy_allocator) try: # Numba 0.54: Disable low occupancy warnings