From b7d0e986d843b28e7e8be4b4b99fca02dfa62f91 Mon Sep 17 00:00:00 2001 From: Matthew Murray <41342305+Matt711@users.noreply.github.com> Date: Mon, 7 Oct 2024 22:04:26 -0400 Subject: [PATCH] Update all rmm imports to use pylibrmm/librmm (#2451) This PR updates all the RMM imports to use pylibrmm/librmm now that `rmm._lib` is deprecated . It should be merged after [rmm/1676](https://github.com/rapidsai/rmm/pull/1676). Authors: - Matthew Murray (https://github.com/Matt711) Approvers: - Ben Frederickson (https://github.com/benfred) URL: https://github.com/rapidsai/raft/pull/2451 --- python/pylibraft/pylibraft/common/handle.pxd | 6 +++--- python/pylibraft/pylibraft/common/handle.pyx | 7 +++++-- python/pylibraft/pylibraft/common/interruptible.pxd | 4 ++-- python/pylibraft/pylibraft/common/interruptible.pyx | 4 ++-- python/pylibraft/pylibraft/neighbors/cagra/cagra.pyx | 6 ++---- .../pylibraft/pylibraft/neighbors/cagra/cpp/c_cagra.pxd | 2 +- python/pylibraft/pylibraft/neighbors/cpp/brute_force.pxd | 2 +- python/pylibraft/pylibraft/neighbors/cpp/rbc.pxd | 2 +- .../pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd | 4 ++-- .../pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx | 8 +++----- .../pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd | 2 +- python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx | 6 ++---- 12 files changed, 25 insertions(+), 28 deletions(-) diff --git a/python/pylibraft/pylibraft/common/handle.pxd b/python/pylibraft/pylibraft/common/handle.pxd index c090663547..78a07ac8bf 100644 --- a/python/pylibraft/pylibraft/common/handle.pxd +++ b/python/pylibraft/pylibraft/common/handle.pxd @@ -1,5 +1,5 @@ # -# Copyright (c) 2022-2023, NVIDIA CORPORATION. +# Copyright (c) 2022-2024, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,8 +22,8 @@ from libcpp.memory cimport shared_ptr, unique_ptr -from rmm._lib.cuda_stream_pool cimport cuda_stream_pool -from rmm._lib.cuda_stream_view cimport cuda_stream_view +from rmm.librmm.cuda_stream_pool cimport cuda_stream_pool +from rmm.librmm.cuda_stream_view cimport cuda_stream_view # Keeping `handle_t` around for backwards compatibility at the diff --git a/python/pylibraft/pylibraft/common/handle.pyx b/python/pylibraft/pylibraft/common/handle.pyx index 7e3dc289e0..d256e671bf 100644 --- a/python/pylibraft/pylibraft/common/handle.pyx +++ b/python/pylibraft/pylibraft/common/handle.pyx @@ -1,5 +1,5 @@ # -# Copyright (c) 2022-2023, NVIDIA CORPORATION. +# Copyright (c) 2022-2024, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,7 +24,10 @@ import functools from cuda.ccudart cimport cudaStream_t from libc.stdint cimport uintptr_t -from rmm._lib.cuda_stream_view cimport cuda_stream_per_thread, cuda_stream_view +from rmm.librmm.cuda_stream_view cimport ( + cuda_stream_per_thread, + cuda_stream_view, +) from .cuda cimport Stream diff --git a/python/pylibraft/pylibraft/common/interruptible.pxd b/python/pylibraft/pylibraft/common/interruptible.pxd index aaccf8aeab..27259571ea 100644 --- a/python/pylibraft/pylibraft/common/interruptible.pxd +++ b/python/pylibraft/pylibraft/common/interruptible.pxd @@ -1,5 +1,5 @@ # -# Copyright (c) 2021-2022, NVIDIA CORPORATION. +# Copyright (c) 2021-2024, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -21,7 +21,7 @@ from libcpp.memory cimport shared_ptr -from rmm._lib.cuda_stream_view cimport cuda_stream_view +from rmm.librmm.cuda_stream_view cimport cuda_stream_view cdef extern from "raft/core/interruptible.hpp" namespace "raft" nogil: diff --git a/python/pylibraft/pylibraft/common/interruptible.pyx b/python/pylibraft/pylibraft/common/interruptible.pyx index bb5415428f..c489f2ee20 100644 --- a/python/pylibraft/pylibraft/common/interruptible.pyx +++ b/python/pylibraft/pylibraft/common/interruptible.pyx @@ -1,5 +1,5 @@ # -# Copyright (c) 2021-2022, NVIDIA CORPORATION. +# Copyright (c) 2021-2024, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -25,7 +25,7 @@ import signal from cuda.ccudart cimport cudaStream_t from cython.operator cimport dereference -from rmm._lib.cuda_stream_view cimport cuda_stream_view +from rmm.librmm.cuda_stream_view cimport cuda_stream_view from .cuda cimport Stream diff --git a/python/pylibraft/pylibraft/neighbors/cagra/cagra.pyx b/python/pylibraft/pylibraft/neighbors/cagra/cagra.pyx index 0e488a51ca..9b376f5f0a 100644 --- a/python/pylibraft/pylibraft/neighbors/cagra/cagra.pyx +++ b/python/pylibraft/pylibraft/neighbors/cagra/cagra.pyx @@ -52,10 +52,8 @@ from pylibraft.common.handle cimport device_resources from pylibraft.common.handle import auto_sync_handle from pylibraft.common.input_validation import is_c_contiguous -from rmm._lib.memory_resource cimport ( - DeviceMemoryResource, - device_memory_resource, -) +from rmm.librmm.memory_resource cimport device_memory_resource +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource cimport pylibraft.neighbors.cagra.cpp.c_cagra as c_cagra from pylibraft.common.optional cimport make_optional, optional diff --git a/python/pylibraft/pylibraft/neighbors/cagra/cpp/c_cagra.pxd b/python/pylibraft/pylibraft/neighbors/cagra/cpp/c_cagra.pxd index 1dffd40186..75ace7f1a8 100644 --- a/python/pylibraft/pylibraft/neighbors/cagra/cpp/c_cagra.pxd +++ b/python/pylibraft/pylibraft/neighbors/cagra/cpp/c_cagra.pxd @@ -27,7 +27,7 @@ from libc.stdint cimport int8_t, int64_t, uint8_t, uint32_t, uint64_t from libcpp cimport bool, nullptr from libcpp.string cimport string -from rmm._lib.memory_resource cimport device_memory_resource +from rmm.librmm.memory_resource cimport device_memory_resource from pylibraft.common.cpp.mdspan cimport ( device_matrix_view, diff --git a/python/pylibraft/pylibraft/neighbors/cpp/brute_force.pxd b/python/pylibraft/pylibraft/neighbors/cpp/brute_force.pxd index 5f6a83a9dc..f513517868 100644 --- a/python/pylibraft/pylibraft/neighbors/cpp/brute_force.pxd +++ b/python/pylibraft/pylibraft/neighbors/cpp/brute_force.pxd @@ -28,7 +28,7 @@ from libcpp cimport bool, nullptr from libcpp.string cimport string from libcpp.vector cimport vector -from rmm._lib.memory_resource cimport device_memory_resource +from rmm.librmm.memory_resource cimport device_memory_resource from pylibraft.common.cpp.mdspan cimport ( device_matrix_view, diff --git a/python/pylibraft/pylibraft/neighbors/cpp/rbc.pxd b/python/pylibraft/pylibraft/neighbors/cpp/rbc.pxd index 531c0dc2c1..d544797119 100644 --- a/python/pylibraft/pylibraft/neighbors/cpp/rbc.pxd +++ b/python/pylibraft/pylibraft/neighbors/cpp/rbc.pxd @@ -28,7 +28,7 @@ from libcpp cimport bool, nullptr from libcpp.string cimport string from libcpp.vector cimport vector -from rmm._lib.memory_resource cimport device_memory_resource +from rmm.librmm.memory_resource cimport device_memory_resource from pylibraft.common.cpp.mdspan cimport ( device_matrix_view, diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd index a281d33310..22b08b3f19 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd @@ -1,5 +1,5 @@ # -# Copyright (c) 2023, NVIDIA CORPORATION. +# Copyright (c) 2023-2024, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,7 +27,7 @@ from libc.stdint cimport int8_t, int64_t, uint8_t, uint32_t, uintptr_t from libcpp cimport bool, nullptr from libcpp.string cimport string -from rmm._lib.memory_resource cimport device_memory_resource +from rmm.librmm.memory_resource cimport device_memory_resource from pylibraft.common.cpp.mdspan cimport ( device_matrix_view, diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx index d8fbdc74da..6826b2bc59 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx @@ -1,5 +1,5 @@ # -# Copyright (c) 2023, NVIDIA CORPORATION. +# Copyright (c) 2023-2024, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -51,10 +51,8 @@ from pylibraft.common.handle cimport device_resources from pylibraft.common.handle import auto_sync_handle from pylibraft.common.input_validation import is_c_contiguous -from rmm._lib.memory_resource cimport ( - DeviceMemoryResource, - device_memory_resource, -) +from rmm.librmm.memory_resource cimport device_memory_resource +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource cimport pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat as c_ivf_flat from pylibraft.common.cpp.optional cimport optional diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd b/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd index 895abbadca..18319bf452 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd @@ -27,7 +27,7 @@ from libc.stdint cimport int8_t, int64_t, uint8_t, uint32_t, uintptr_t from libcpp cimport bool, nullptr from libcpp.string cimport string -from rmm._lib.memory_resource cimport device_memory_resource +from rmm.librmm.memory_resource cimport device_memory_resource from pylibraft.common.cpp.mdspan cimport ( device_matrix_view, diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx index 5b89f0d9a5..f467957fd6 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx @@ -44,10 +44,8 @@ from pylibraft.common.handle cimport device_resources from pylibraft.common.handle import auto_sync_handle from pylibraft.common.input_validation import is_c_contiguous -from rmm._lib.memory_resource cimport ( - DeviceMemoryResource, - device_memory_resource, -) +from rmm.librmm.memory_resource cimport device_memory_resource +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource cimport pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat as c_ivf_flat cimport pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq as c_ivf_pq