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](rapidsai/rmm#1676). Authors: - Matthew Murray (https://github.com/Matt711) Approvers: - Ben Frederickson (https://github.com/benfred) URL: #2451
rapidsai · Oct 8, 2024 · b7d0e98 · b7d0e98
1 parent 94cc2c2
commit b7d0e98
Show file tree

Hide file tree

Showing 12 changed files with 25 additions and 28 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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