Move stream objects from cuda_stream cython to stream cython

.pre-commit-config.yaml

@@ -103,7 +103,7 @@ repos:
         args: [--fix, --spdx]
         files: |
           (?x)
-            [.](cmake|cpp|cu|cuh|h|hpp|sh|pxd|py|pyx|rs|java)$|
+            [.](cmake|cpp|cu|cuh|h|hpp|sh|pxd|py|pyi|pyx|rs|java)$|
             CMakeLists[.]txt$|
             CMakeLists_standalone[.]txt$|
             meta[.]yaml$|

python/rmm/rmm/librmm/cuda_stream.pxd

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 from cuda.bindings.cyruntime cimport cudaStream_t
@@ -15,6 +15,7 @@ cdef extern from "rmm/cuda_stream.hpp" namespace "rmm" nogil:
         non_blocking "rmm::cuda_stream::flags::non_blocking"
     cdef cppclass cuda_stream:
         cuda_stream() except +
+        cuda_stream(cuda_stream_flags flags) except +
         bool is_valid() except +
         cudaStream_t value() except +
         cuda_stream_view view() except +

python/rmm/rmm/pylibrmm/__init__.py

@@ -1,10 +1,10 @@
-# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 from rmm.pylibrmm import memory_resource
 
 from .cuda_stream_pool import CudaStreamPool
-from .cuda_stream import CudaStreamFlags
+from .stream import CudaStreamFlags
 from .device_buffer import DeviceBuffer
 
 __all__ = [

python/rmm/rmm/pylibrmm/cuda_stream.pxd

@@ -1,16 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
-cimport cython
-from cuda.bindings.cyruntime cimport cudaStream_t
-from libcpp cimport bool
-from libcpp.memory cimport unique_ptr
-
-from rmm.librmm.cuda_stream cimport cuda_stream
-
-
-@cython.final
-cdef class CudaStream:
-    cdef unique_ptr[cuda_stream] c_obj
-    cdef cudaStream_t value(self) except * nogil
-    cdef bool is_valid(self) except * nogil
+from rmm.pylibrmm.stream cimport _OwningStream as CudaStream

python/rmm/rmm/pylibrmm/cuda_stream.pyx

@@ -1,46 +1,28 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
-cimport cython
-from enum import IntEnum
-from cuda.bindings.cyruntime cimport cudaStream_t
-from libcpp cimport bool
-
-from rmm.librmm.cuda_stream cimport cuda_stream, cuda_stream_flags
-
-
-class CudaStreamFlags(IntEnum):
-    """
-    Enumeration of CUDA stream creation flags.
-
-    Attributes
-    ----------
-    SYNC_DEFAULT : int
-        Created stream synchronizes with the default stream.
-    NON_BLOCKING : int
-        Created stream does not synchronize with the default stream.
-    """
-    SYNC_DEFAULT = <int>(cuda_stream_flags.sync_default)
-    NON_BLOCKING = <int>(cuda_stream_flags.non_blocking)
-
-
-@cython.final
-cdef class CudaStream:
-    """
-    Wrapper around a CUDA stream with RAII semantics.
-    When a CudaStream instance is GC'd, the underlying
-    CUDA stream is destroyed.
-    """
-    def __cinit__(self):
-        with nogil:
-            self.c_obj.reset(new cuda_stream())
-
-    def __dealloc__(self):
-        with nogil:
-            self.c_obj.reset()
-
-    cdef cudaStream_t value(self) except * nogil:
-        return self.c_obj.get()[0].value()
-
-    cdef bool is_valid(self) except * nogil:
-        return self.c_obj.get()[0].is_valid()
+# backwards compat file reexporting CudaStream and CudaStreamFlags
+
+import warnings
+
+from rmm.pylibrmm.stream cimport _OwningStream as CudaStream
+
+import sys
+
+from rmm.pylibrmm.stream import CudaStreamFlags, _OwningStream
+
+sys.modules[__name__].CudaStream = _OwningStream
+del _OwningStream
+
+warnings.warn(
+    "rmm.pylibrmm.cuda_stream is deprecated; "
+    "use rmm.pylibrmm.stream.CudaStreamFlags for stream flags, "
+    "and rmm.pylibrmm.stream.Stream for stream objects.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+__all__ = [
+    "CudaStream",
+    "CudaStreamFlags",
+]

python/rmm/rmm/pylibrmm/stream.pxd

@@ -1,13 +1,23 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
+cimport cython
 from cuda.bindings.cyruntime cimport cudaStream_t
 from libc.stdint cimport uintptr_t
 from libcpp cimport bool
+from libcpp.memory cimport unique_ptr
 
+from rmm.librmm.cuda_stream cimport cuda_stream
 from rmm.librmm.cuda_stream_view cimport cuda_stream_view
 
 
+@cython.final
+cdef class _OwningStream:
+    cdef unique_ptr[cuda_stream] c_obj
+    cdef cudaStream_t value(self) except * nogil
+    cdef bool is_valid(self) except * nogil
+
+
 cdef class Stream:
     cdef cudaStream_t _cuda_stream
     cdef object _owner
@@ -18,5 +28,5 @@ cdef class Stream:
     cdef cuda_stream_view view(self) noexcept nogil
     cdef void c_synchronize(self) except * nogil
     cdef bool c_is_default(self) noexcept nogil
-    cdef void _init_with_new_cuda_stream(self) except *
+    cdef void _init_with_new_cuda_stream(self, flags=*) except *
     cdef void _init_from_stream(self, Stream stream) except *

python/rmm/rmm/pylibrmm/stream.pyi

@@ -1,10 +1,20 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
+from enum import IntEnum
 from typing import Any, Optional
 
+class CudaStreamFlags(IntEnum):
+    SYNC_DEFAULT = ...
+    NON_BLOCKING = ...
+
+class _OwningStream:
+    def __init__(self, flags: CudaStreamFlags = ...) -> None: ...
+
 class Stream:
-    def __init__(self, obj: Optional[Any] = None) -> None: ...
+    def __init__(
+        self, obj: Optional[Any] = None, flags: CudaStreamFlags = ...
+    ) -> None: ...
     def __cuda_stream__(self) -> tuple[int, int]: ...
     def synchronize(self) -> None: ...
     def is_default(self) -> bool: ...

python/rmm/rmm/pylibrmm/stream.pyx

@@ -1,21 +1,61 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
+cimport cython
+from enum import IntEnum
 from cuda.bindings.cyruntime cimport cudaStream_t
 from libc.stdint cimport uintptr_t
 from libcpp cimport bool
 
+from rmm.librmm.cuda_stream cimport cuda_stream, cuda_stream_flags
 from rmm.librmm.cuda_stream_view cimport (
     cuda_stream_default,
     cuda_stream_legacy,
     cuda_stream_per_thread,
     cuda_stream_view,
 )
-from rmm.pylibrmm.cuda_stream cimport CudaStream
+
+
+class CudaStreamFlags(IntEnum):
+    """
+    Enumeration of CUDA stream creation flags.
+
+    Attributes
+    ----------
+    SYNC_DEFAULT : int
+        Created stream synchronizes with the default stream.
+    NON_BLOCKING : int
+        Created stream does not synchronize with the default stream.
+    """
+    SYNC_DEFAULT = <int>(cuda_stream_flags.sync_default)
+    NON_BLOCKING = <int>(cuda_stream_flags.non_blocking)
+
+
+@cython.final
+cdef class _OwningStream:
+    """
+    Wrapper around a CUDA stream with RAII semantics.
+    When an _OwningStream instance is GC'd, the underlying
+    CUDA stream is destroyed.
+    """
+    def __cinit__(self, flags=CudaStreamFlags.SYNC_DEFAULT):
+        cdef cuda_stream_flags c_flags = <cuda_stream_flags>(<int>flags)
+        with nogil:
+            self.c_obj.reset(new cuda_stream(c_flags))
+
+    def __dealloc__(self):
+        with nogil:
+            self.c_obj.reset()
+
+    cdef cudaStream_t value(self) except * nogil:
+        return self.c_obj.get()[0].value()
+
+    cdef bool is_valid(self) except * nogil:
+        return self.c_obj.get()[0].is_valid()
 
 
 cdef class Stream:
-    def __init__(self, obj=None):
+    def __init__(self, obj=None, flags=CudaStreamFlags.SYNC_DEFAULT):
         """
         A Stream represents a CUDA stream.
 
@@ -25,9 +65,12 @@ cdef class Stream:
             * If None (the default), a new CUDA stream is created.
             * If a Numba or CuPy stream is provided, we make a thin
               wrapper around it.
+        flags: CudaStreamFlags, optional
+            Stream creation flags. Only used when obj is None.
+            Default is CudaStreamFlags.SYNC_DEFAULT.
         """
         if obj is None:
-            self._init_with_new_cuda_stream()
+            self._init_with_new_cuda_stream(flags)
         elif isinstance(obj, Stream):
             self._init_from_stream(obj)
         else:
@@ -132,8 +175,10 @@ cdef class Stream:
     def __hash__(self):
         return hash(int(<uintptr_t>(self._cuda_stream)))
 
-    cdef void _init_with_new_cuda_stream(self) except *:
-        cdef CudaStream stream = CudaStream()
+    cdef void _init_with_new_cuda_stream(
+        self, flags=CudaStreamFlags.SYNC_DEFAULT
+    ) except *:
+        cdef _OwningStream stream = _OwningStream(flags)
         self._cuda_stream = stream.value()
         self._owner = stream
 

python/rmm/rmm/tests/test_stream.py

@@ -1,10 +1,9 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 
 import pytest
 
-import rmm.pylibrmm.cuda_stream
 import rmm.pylibrmm.cuda_stream_pool
 import rmm.pylibrmm.stream
 
@@ -94,8 +93,8 @@ def test_cuda_core_buffer(current_device):
 @pytest.mark.parametrize(
     "flags",
     [
-        rmm.pylibrmm.cuda_stream.CudaStreamFlags.SYNC_DEFAULT,
-        rmm.pylibrmm.cuda_stream.CudaStreamFlags.NON_BLOCKING,
+        rmm.pylibrmm.stream.CudaStreamFlags.SYNC_DEFAULT,
+        rmm.pylibrmm.stream.CudaStreamFlags.NON_BLOCKING,
     ],
 )
 def test_cuda_stream_pool(current_device, flags):
@@ -118,6 +117,20 @@ def test_cuda_stream_pool(current_device, flags):
         assert streams[i] == stream_pool.get_stream(i)
 
 
+def test_cuda_stream_module_deprecation():
+    import importlib
+    import sys
+    import warnings
+
+    sys.modules.pop("rmm.pylibrmm.cuda_stream", None)
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        importlib.import_module("rmm.pylibrmm.cuda_stream")
+    assert any(
+        issubclass(warning.category, DeprecationWarning) for warning in w
+    )
+
+
 def test_hashable():
     a = rmm.pylibrmm.stream.Stream()
     b = rmm.pylibrmm.stream.Stream()