Make the stream module a part of the public API (#1775)

Closes #1770.

Authors:
  - Matthew Murray (https://github.com/Matt711)
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Bradley Dice (https://github.com/bdice)

URL: #1775

python/rmm/docs/guide.md

@@ -93,7 +93,7 @@ destination device ID and stream are optional parameters.
 ```python
 >>> import rmm
 >>> rmm.reinitialize(managed_memory=True)
->>> from rmm._cuda.stream import Stream
+>>> from rmm.pylibrmm.stream import Stream
 >>> stream = Stream()
 >>> buf = rmm.DeviceBuffer(size=100, stream=stream)
 >>> buf.prefetch(device=3, stream=stream) # prefetch to device on stream.

python/rmm/rmm/_cuda/CMakeLists.txt

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Copyright (c) 2022-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 # in compliance with the License. You may obtain a copy of the License at
@@ -12,7 +12,6 @@
 # the License.
 # =============================================================================
 
-set(cython_sources stream.pyx)
 set(linked_libraries rmm::rmm cpp_logger)
 
 rapids_cython_create_modules(SOURCE_FILES "${cython_sources}" LINKED_LIBRARIES "${linked_libraries}"

python/rmm/rmm/_cuda/stream.pxd

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2024, NVIDIA CORPORATION.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,22 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cuda.bindings.cyruntime cimport cudaStream_t
-from libc.stdint cimport uintptr_t
-from libcpp cimport bool
-
-from rmm.librmm.cuda_stream_view cimport cuda_stream_view
-
-
-cdef class Stream:
-    cdef cudaStream_t _cuda_stream
-    cdef object _owner
-
-    @staticmethod
-    cdef Stream _from_cudaStream_t(cudaStream_t s, object owner=*)
-
-    cdef cuda_stream_view view(self) except * nogil
-    cdef void c_synchronize(self) except * nogil
-    cdef bool c_is_default(self) except * nogil
-    cdef void _init_with_new_cuda_stream(self) except *
-    cdef void _init_from_stream(self, Stream stream) except *
+from rmm.pylibrmm.stream cimport Stream

python/rmm/rmm/_cuda/stream.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+
+from rmm.pylibrmm.stream import (
+    DEFAULT_STREAM,
+    LEGACY_DEFAULT_STREAM,
+    PER_THREAD_DEFAULT_STREAM,
+    Stream,
+)
+
+__all__ = [
+    "DEFAULT_STREAM",
+    "LEGACY_DEFAULT_STREAM",
+    "PER_THREAD_DEFAULT_STREAM",
+    "Stream",
+]
+
+warnings.warn(
+    "The `rmm._cuda.stream` module is deprecated in 25.02 and will be removed in a future release. Use `rmm.pylibrmm.stream` instead.",
+    FutureWarning,
+    stacklevel=2,
+)

python/rmm/rmm/allocators/cupy.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023-2024, NVIDIA CORPORATION.
+# Copyright (c) 2023-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from rmm import pylibrmm
-from rmm._cuda.stream import Stream
+from rmm.pylibrmm.stream import Stream
 
 try:
     import cupy

python/rmm/rmm/pylibrmm/CMakeLists.txt

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Copyright (c) 2022-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 # in compliance with the License. You may obtain a copy of the License at
@@ -12,7 +12,8 @@
 # the License.
 # =============================================================================
 
-set(cython_sources device_buffer.pyx logger.pyx memory_resource.pyx cuda_stream.pyx helper.pyx)
+set(cython_sources device_buffer.pyx logger.pyx memory_resource.pyx cuda_stream.pyx helper.pyx
+                   stream.pyx)
 set(linked_libraries rmm::rmm cpp_logger)
 
 # Build all of the Cython targets

python/rmm/rmm/pylibrmm/device_buffer.pxd

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2024, NVIDIA CORPORATION.
+# Copyright (c) 2019-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,9 +15,9 @@
 from libc.stdint cimport uintptr_t
 from libcpp.memory cimport unique_ptr
 
-from rmm._cuda.stream cimport Stream
 from rmm.librmm.device_buffer cimport device_buffer
 from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
+from rmm.pylibrmm.stream cimport Stream
 
 
 cdef class DeviceBuffer:

python/rmm/rmm/pylibrmm/device_buffer.pyx

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2024, NVIDIA CORPORATION.
+# Copyright (c) 2019-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,9 +19,9 @@ from libc.stdint cimport uintptr_t
 from libcpp.memory cimport unique_ptr
 from libcpp.utility cimport move
 
-from rmm._cuda.stream cimport Stream
+from rmm.pylibrmm.stream cimport Stream
 
-from rmm._cuda.stream import DEFAULT_STREAM
+from rmm.pylibrmm.stream import DEFAULT_STREAM
 
 from cuda.bindings.cyruntime cimport (
     cudaError,

python/rmm/rmm/pylibrmm/memory_resource.pyx

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2024, NVIDIA CORPORATION.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -32,9 +32,9 @@ from cuda.bindings.runtime import cudaError_t
 
 from rmm._cuda.gpu import CUDARuntimeError, getDevice, setDevice
 
-from rmm._cuda.stream cimport Stream
+from rmm.pylibrmm.stream cimport Stream
 
-from rmm._cuda.stream import DEFAULT_STREAM
+from rmm.pylibrmm.stream import DEFAULT_STREAM
 
 from rmm.librmm.cuda_stream_view cimport cuda_stream_view
 from rmm.librmm.per_device_resource cimport (

python/rmm/rmm/pylibrmm/stream.pxd

@@ -0,0 +1,33 @@
+# Copyright (c) 2020-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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from cuda.bindings.cyruntime cimport cudaStream_t
+from libc.stdint cimport uintptr_t
+from libcpp cimport bool
+
+from rmm.librmm.cuda_stream_view cimport cuda_stream_view
+
+
+cdef class Stream:
+    cdef cudaStream_t _cuda_stream
+    cdef object _owner
+
+    @staticmethod
+    cdef Stream _from_cudaStream_t(cudaStream_t s, object owner=*)
+
+    cdef cuda_stream_view view(self) except * nogil
+    cdef void c_synchronize(self) except * nogil
+    cdef bool c_is_default(self) except * nogil
+    cdef void _init_with_new_cuda_stream(self) except *
+    cdef void _init_from_stream(self, Stream stream) except *

python/rmm/rmm/tests/test_rmm.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2024, NVIDIA CORPORATION.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -26,10 +26,10 @@
 from numba import cuda
 
 import rmm
-import rmm._cuda.stream
 from rmm.allocators.cupy import rmm_cupy_allocator
 from rmm.allocators.numba import RMMNumbaManager
 from rmm.pylibrmm.logger import level_enum
+from rmm.pylibrmm.stream import Stream
 
 cuda.set_memory_manager(RMMNumbaManager)
 
@@ -348,8 +348,8 @@ def test_rmm_device_buffer_prefetch(pool, managed):
 def test_rmm_pool_numba_stream(stream):
     rmm.reinitialize(pool_allocator=True)
 
-    stream = rmm._cuda.stream.Stream(stream)
-    a = rmm.pylibrmm.device_buffer.DeviceBuffer(size=3, stream=stream)
+    stream = Stream(stream)
+    a = rmm.DeviceBuffer(size=3, stream=stream)
 
     assert a.size == 3
     assert a.ptr != 0
@@ -600,7 +600,7 @@ def test_mr_devicebuffer_lifetime():
     )
 
     # Creates a new non-default stream
-    stream = rmm._cuda.stream.Stream()
+    stream = rmm.pylibrmm.stream.Stream()
 
     # Allocate DeviceBuffer with Pool and Stream
     a = rmm.DeviceBuffer(size=10, stream=stream)
@@ -695,7 +695,7 @@ def test_cuda_async_memory_resource_stream(nelems):
     # with a non-default stream works
     mr = rmm.mr.CudaAsyncMemoryResource()
     rmm.mr.set_current_device_resource(mr)
-    stream = rmm._cuda.stream.Stream()
+    stream = Stream()
     expected = np.full(nelems, 5, dtype="u1")
     dbuf = rmm.DeviceBuffer.to_device(expected, stream=stream)
     result = np.asarray(dbuf.copy_to_host())