Make the stream module a part of the public API

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,28 @@
+# 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 (  # noqa: F401
+    DEFAULT_STREAM,
+    LEGACY_DEFAULT_STREAM,
+    PER_THREAD_DEFAULT_STREAM,
+    Stream,
+)
+
+warnings.warn(
+    "The `rmm.pylibrmm.stream` module is deprecated in 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,7 +15,7 @@
 from libc.stdint cimport uintptr_t
 from libcpp.memory cimport unique_ptr
 
-from rmm._cuda.stream cimport Stream
+from rmm.pylibrmm.stream cimport Stream
 from rmm.librmm.device_buffer cimport device_buffer
 from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 

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,32 @@
+# 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.
+# 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 *

python/rmm/rmm/_cuda/stream.pyx → python/rmm/rmm/pylibrmm/stream.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.
@@ -34,19 +34,82 @@ cdef class Stream:
         ----------
         obj: optional
             * If None (the default), a new CUDA stream is created.
+            * If a stream that implements the __cuda_stream__ protocol
+              is provided, we use it.
             * If a Numba or CuPy stream is provided, we make a thin
               wrapper around it.
         """
         if obj is None:
             self._init_with_new_cuda_stream()
-        elif isinstance(obj, Stream):
-            self._init_from_stream(obj)
+            return
+        elif hasattr(obj, "__cuda_stream__"):
+            protocol = getattr(obj, "__cuda_stream__")
+            if protocol[0] != 0:
+                raise ValueError("Only protocol version 0 is supported")
+            self._cuda_stream = <cudaStream_t>obj
+            self.owner = obj
         else:
+            # TODO: Remove this branch when numba and cupy
+            # streams implement __cuda_stream__
             try:
                 self._init_from_numba_stream(obj)
             except TypeError:
                 self._init_from_cupy_stream(obj)
 
+    @property
+    def __cuda_stream__(self):
+        """Return an instance of a __cuda_stream__ protocol."""
+        return (0, self.handle)
+
+    @property
+    def handle(self) -> int:
+        """Return the underlying cudaStream_t pointer address as Python int."""
+        return int(<uintptr_t>self._cuda_stream)
+
+    # @singledispatchmethod
+    # def _init_from_stream(self, obj):
+    #     if obj is None:
+    #         self._init_with_new_cuda_stream()
+    #         return
+    #     try:
+    #         protocol = getattr(obj, "__cuda_stream__")
+    #     except AttributeError:
+    #         raise ValueError(
+    #             "Argument must be None, a Stream, or implement __cuda_stream__"
+    #         )
+    #     if protocol[0] != 0:
+    #         raise ValueError("Only protocol version 0 is supported")
+
+    #     self._cuda_stream = <cudaStream_t>obj
+    #     self.owner = obj
+
+    # @_init_from_stream.register
+    # def _(self, stream: Stream):
+    #     self._cuda_stream, self._owner = stream._cuda_stream, stream._owner
+
+    # try:
+    #     from numba import cuda
+    #     @_init_from_stream.register
+    #     def _(self, obj: cuda.cudadrv.driver.Stream):
+    #         self._cuda_stream = <cudaStream_t><uintptr_t>(int(obj))
+    #         self._owner = obj
+    # except ImportError:
+    #     pass
+
+    # try:
+    #     import cupy
+    #     @_init_from_stream.register(cupy.cuda.stream.Stream)
+    #     def _(self, obj):
+    #         self._cuda_stream = <cudaStream_t><uintptr_t>(obj.ptr)
+    #         self._owner = obj
+
+    #     @_init_from_stream.register(cupy.cuda.stream.ExternalStream)
+    #     def _(self, obj):
+    #         self._cuda_stream = <cudaStream_t><uintptr_t>(obj.ptr)
+    #         self._owner = obj
+    # except ImportError:
+    #     pass
+
     @staticmethod
     cdef Stream _from_cudaStream_t(cudaStream_t s, object owner=None) except *:
         """
@@ -117,9 +180,6 @@ cdef class Stream:
         self._cuda_stream = stream.value()
         self._owner = stream
 
-    cdef void _init_from_stream(self, Stream stream) except *:
-        self._cuda_stream, self._owner = stream._cuda_stream, stream._owner
-
 
 DEFAULT_STREAM = Stream._from_cudaStream_t(cuda_stream_default.value())
 LEGACY_DEFAULT_STREAM = Stream._from_cudaStream_t(cuda_stream_legacy.value())

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,7 +26,7 @@
 from numba import cuda
 
 import rmm
-import rmm._cuda.stream
+import rmm.pylibrmm.stream
 from rmm.allocators.cupy import rmm_cupy_allocator
 from rmm.allocators.numba import RMMNumbaManager
 from rmm.pylibrmm.logger import level_enum
@@ -348,7 +348,7 @@ 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)
+    stream = rmm.pylibrmm.stream.Stream(stream)
     a = rmm.pylibrmm.device_buffer.DeviceBuffer(size=3, stream=stream)
 
     assert a.size == 3
@@ -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 = rmm.pylibrmm.stream.Stream()
     expected = np.full(nelems, 5, dtype="u1")
     dbuf = rmm.DeviceBuffer.to_device(expected, stream=stream)
     result = np.asarray(dbuf.copy_to_host())