From 55ffa8295eb72b4ac40a150ab4db52daab64d372 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Tue, 17 Feb 2026 09:04:59 -0800 Subject: [PATCH 01/10] move/deprecate --- python/rmm/rmm/pylibrmm/cuda_stream.pxd | 18 ++----- python/rmm/rmm/pylibrmm/cuda_stream.pyx | 63 ++++++++----------------- python/rmm/rmm/pylibrmm/stream.pxd | 10 ++++ python/rmm/rmm/pylibrmm/stream.pyi | 11 +++++ python/rmm/rmm/pylibrmm/stream.pyx | 42 ++++++++++++++++- 5 files changed, 86 insertions(+), 58 deletions(-) diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pxd b/python/rmm/rmm/pylibrmm/cuda_stream.pxd index 0cabbf933..08854bb49 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pxd +++ b/python/rmm/rmm/pylibrmm/cuda_stream.pxd @@ -1,16 +1,6 @@ -# SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2025, 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 +# Re-export from stream for backward compatibility. +# Prefer: from rmm.pylibrmm.stream cimport CudaStream +from rmm.pylibrmm.stream cimport CudaStream diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pyx b/python/rmm/rmm/pylibrmm/cuda_stream.pyx index 4287de7bd..1d504af95 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pyx +++ b/python/rmm/rmm/pylibrmm/cuda_stream.pyx @@ -1,46 +1,23 @@ # SPDX-FileCopyrightText: Copyright (c) 2020-2025, 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 = (cuda_stream_flags.sync_default) - NON_BLOCKING = (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() +import warnings + +# Re-export from stream for backward compat +# prefer from rmm.pylibrmm.stream import CudaStream, CudaStreamFlags +from rmm.pylibrmm.stream import ( + CudaStream, + CudaStreamFlags, +) + +warnings.warn( + "rmm.pylibrmm.cuda_stream is deprecated; use rmm.pylibrmm.stream for " + "CudaStream and CudaStreamFlags.", + DeprecationWarning, + stacklevel=2, +) + +__all__ = [ + "CudaStream", + "CudaStreamFlags", +] diff --git a/python/rmm/rmm/pylibrmm/stream.pxd b/python/rmm/rmm/pylibrmm/stream.pxd index ec0facf08..8aef2a975 100644 --- a/python/rmm/rmm/pylibrmm/stream.pxd +++ b/python/rmm/rmm/pylibrmm/stream.pxd @@ -1,13 +1,23 @@ # SPDX-FileCopyrightText: Copyright (c) 2020-2025, 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 CudaStream: + 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 diff --git a/python/rmm/rmm/pylibrmm/stream.pyi b/python/rmm/rmm/pylibrmm/stream.pyi index eccf632c2..edd37e603 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyi +++ b/python/rmm/rmm/pylibrmm/stream.pyi @@ -1,8 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 +from enum import IntEnum from typing import Any, Optional + +class CudaStreamFlags(IntEnum): + SYNC_DEFAULT: int + NON_BLOCKING: int + + +class CudaStream: + def __init__(self) -> None: ... + + class Stream: def __init__(self, obj: Optional[Any] = None) -> None: ... def __cuda_stream__(self) -> tuple[int, int]: ... diff --git a/python/rmm/rmm/pylibrmm/stream.pyx b/python/rmm/rmm/pylibrmm/stream.pyx index 7b4f3abf3..416ef428a 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyx +++ b/python/rmm/rmm/pylibrmm/stream.pyx @@ -1,17 +1,57 @@ # SPDX-FileCopyrightText: Copyright (c) 2020-2025, 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 libcpp.memory cimport unique_ptr +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 = (cuda_stream_flags.sync_default) + NON_BLOCKING = (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() cdef class Stream: From 88af3f47772f89a9c73827c5963a720ff576b58b Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Tue, 17 Feb 2026 13:59:26 -0800 Subject: [PATCH 02/10] move code around, pass tests --- python/rmm/rmm/pylibrmm/__init__.py | 4 ++-- python/rmm/rmm/pylibrmm/cuda_stream.pxd | 2 +- python/rmm/rmm/pylibrmm/cuda_stream.pyx | 7 ++----- python/rmm/rmm/pylibrmm/stream.pxd | 2 +- python/rmm/rmm/pylibrmm/stream.pyi | 7 ++----- python/rmm/rmm/pylibrmm/stream.pyx | 3 +-- python/rmm/rmm/tests/test_stream.py | 7 +++---- 7 files changed, 12 insertions(+), 20 deletions(-) diff --git a/python/rmm/rmm/pylibrmm/__init__.py b/python/rmm/rmm/pylibrmm/__init__.py index f82c5fb5f..a626ffa4e 100644 --- a/python/rmm/rmm/pylibrmm/__init__.py +++ b/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__ = [ diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pxd b/python/rmm/rmm/pylibrmm/cuda_stream.pxd index 08854bb49..4fd6fb599 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pxd +++ b/python/rmm/rmm/pylibrmm/cuda_stream.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # Re-export from stream for backward compatibility. diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pyx b/python/rmm/rmm/pylibrmm/cuda_stream.pyx index 1d504af95..542a16570 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pyx +++ b/python/rmm/rmm/pylibrmm/cuda_stream.pyx @@ -1,14 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 import warnings # Re-export from stream for backward compat # prefer from rmm.pylibrmm.stream import CudaStream, CudaStreamFlags -from rmm.pylibrmm.stream import ( - CudaStream, - CudaStreamFlags, -) +from rmm.pylibrmm.stream import CudaStream, CudaStreamFlags warnings.warn( "rmm.pylibrmm.cuda_stream is deprecated; use rmm.pylibrmm.stream for " diff --git a/python/rmm/rmm/pylibrmm/stream.pxd b/python/rmm/rmm/pylibrmm/stream.pxd index 8aef2a975..8dc4bed67 100644 --- a/python/rmm/rmm/pylibrmm/stream.pxd +++ b/python/rmm/rmm/pylibrmm/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 cimport cython diff --git a/python/rmm/rmm/pylibrmm/stream.pyi b/python/rmm/rmm/pylibrmm/stream.pyi index edd37e603..27f92fca5 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyi +++ b/python/rmm/rmm/pylibrmm/stream.pyi @@ -4,16 +4,13 @@ from enum import IntEnum from typing import Any, Optional - class CudaStreamFlags(IntEnum): - SYNC_DEFAULT: int - NON_BLOCKING: int - + SYNC_DEFAULT = ... + NON_BLOCKING = ... class CudaStream: def __init__(self) -> None: ... - class Stream: def __init__(self, obj: Optional[Any] = None) -> None: ... def __cuda_stream__(self) -> tuple[int, int]: ... diff --git a/python/rmm/rmm/pylibrmm/stream.pyx b/python/rmm/rmm/pylibrmm/stream.pyx index 416ef428a..231dd66cd 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyx +++ b/python/rmm/rmm/pylibrmm/stream.pyx @@ -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 cimport cython @@ -6,7 +6,6 @@ from enum import IntEnum 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, cuda_stream_flags from rmm.librmm.cuda_stream_view cimport ( diff --git a/python/rmm/rmm/tests/test_stream.py b/python/rmm/rmm/tests/test_stream.py index 5d7de48ed..8a6f123d2 100644 --- a/python/rmm/rmm/tests/test_stream.py +++ b/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): From 12a45d30f7b2a161c3d7be1703a68c561acd9831 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 18 Feb 2026 07:29:36 -0800 Subject: [PATCH 03/10] address reviews --- python/rmm/rmm/pylibrmm/cuda_stream.pxd | 4 +--- python/rmm/rmm/pylibrmm/cuda_stream.pyx | 8 +++----- python/rmm/rmm/pylibrmm/stream.pxd | 2 +- python/rmm/rmm/pylibrmm/stream.pyx | 9 ++++++--- 4 files changed, 11 insertions(+), 12 deletions(-) diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pxd b/python/rmm/rmm/pylibrmm/cuda_stream.pxd index 4fd6fb599..38205b1b1 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pxd +++ b/python/rmm/rmm/pylibrmm/cuda_stream.pxd @@ -1,6 +1,4 @@ # SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 -# Re-export from stream for backward compatibility. -# Prefer: from rmm.pylibrmm.stream cimport CudaStream -from rmm.pylibrmm.stream cimport CudaStream +from rmm.pylibrmm.stream cimport _OwningStream as CudaStream diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pyx b/python/rmm/rmm/pylibrmm/cuda_stream.pyx index 542a16570..6429a6696 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pyx +++ b/python/rmm/rmm/pylibrmm/cuda_stream.pyx @@ -3,18 +3,16 @@ import warnings -# Re-export from stream for backward compat -# prefer from rmm.pylibrmm.stream import CudaStream, CudaStreamFlags -from rmm.pylibrmm.stream import CudaStream, CudaStreamFlags +# Re-export from stream for backward compat. +from rmm.pylibrmm.stream import CudaStreamFlags warnings.warn( "rmm.pylibrmm.cuda_stream is deprecated; use rmm.pylibrmm.stream for " - "CudaStream and CudaStreamFlags.", + "CudaStreamFlags.", DeprecationWarning, stacklevel=2, ) __all__ = [ - "CudaStream", "CudaStreamFlags", ] diff --git a/python/rmm/rmm/pylibrmm/stream.pxd b/python/rmm/rmm/pylibrmm/stream.pxd index 8dc4bed67..d55698e73 100644 --- a/python/rmm/rmm/pylibrmm/stream.pxd +++ b/python/rmm/rmm/pylibrmm/stream.pxd @@ -12,7 +12,7 @@ from rmm.librmm.cuda_stream_view cimport cuda_stream_view @cython.final -cdef class CudaStream: +cdef class _OwningStream: cdef unique_ptr[cuda_stream] c_obj cdef cudaStream_t value(self) except * nogil cdef bool is_valid(self) except * nogil diff --git a/python/rmm/rmm/pylibrmm/stream.pyx b/python/rmm/rmm/pylibrmm/stream.pyx index 231dd66cd..73ba7fc8c 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyx +++ b/python/rmm/rmm/pylibrmm/stream.pyx @@ -32,10 +32,10 @@ class CudaStreamFlags(IntEnum): @cython.final -cdef class CudaStream: +cdef class _OwningStream: """ Wrapper around a CUDA stream with RAII semantics. - When a CudaStream instance is GC'd, the underlying + When an _OwningStream instance is GC'd, the underlying CUDA stream is destroyed. """ def __cinit__(self): @@ -53,6 +53,9 @@ cdef class CudaStream: return self.c_obj.get()[0].is_valid() +CudaStream = _OwningStream + + cdef class Stream: def __init__(self, obj=None): """ @@ -172,7 +175,7 @@ cdef class Stream: return hash(int((self._cuda_stream))) cdef void _init_with_new_cuda_stream(self) except *: - cdef CudaStream stream = CudaStream() + cdef _OwningStream stream = _OwningStream() self._cuda_stream = stream.value() self._owner = stream From ce2a938e697054eeaf211ed87a87e38131e169f6 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 18 Feb 2026 08:47:01 -0800 Subject: [PATCH 04/10] put module in namespace directly to workaround cython --- python/rmm/rmm/pylibrmm/cuda_stream.pyx | 12 ++++++++++-- python/rmm/rmm/pylibrmm/stream.pyi | 2 +- python/rmm/rmm/pylibrmm/stream.pyx | 3 --- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pyx b/python/rmm/rmm/pylibrmm/cuda_stream.pyx index 6429a6696..61a43ef21 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pyx +++ b/python/rmm/rmm/pylibrmm/cuda_stream.pyx @@ -1,10 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 +# backwards compat file reexporting CudaStream and CudaStreamFlags + import warnings -# Re-export from stream for backward compat. -from rmm.pylibrmm.stream import CudaStreamFlags +from rmm.pylibrmm.stream cimport _OwningStream as CudaStream + +import sys + +from rmm.pylibrmm.stream import CudaStreamFlags, _OwningStream + +sys.modules[__name__].CudaStream = _OwningStream warnings.warn( "rmm.pylibrmm.cuda_stream is deprecated; use rmm.pylibrmm.stream for " @@ -14,5 +21,6 @@ warnings.warn( ) __all__ = [ + "CudaStream", "CudaStreamFlags", ] diff --git a/python/rmm/rmm/pylibrmm/stream.pyi b/python/rmm/rmm/pylibrmm/stream.pyi index 27f92fca5..8dd329c55 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyi +++ b/python/rmm/rmm/pylibrmm/stream.pyi @@ -8,7 +8,7 @@ class CudaStreamFlags(IntEnum): SYNC_DEFAULT = ... NON_BLOCKING = ... -class CudaStream: +class _OwningStream: def __init__(self) -> None: ... class Stream: diff --git a/python/rmm/rmm/pylibrmm/stream.pyx b/python/rmm/rmm/pylibrmm/stream.pyx index 73ba7fc8c..8a4d2fc11 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyx +++ b/python/rmm/rmm/pylibrmm/stream.pyx @@ -53,9 +53,6 @@ cdef class _OwningStream: return self.c_obj.get()[0].is_valid() -CudaStream = _OwningStream - - cdef class Stream: def __init__(self, obj=None): """ From 5c7b71ada00e70f1ab656b9e0f9e17360d3effe0 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Tue, 24 Feb 2026 06:46:11 -0800 Subject: [PATCH 05/10] warning update --- python/rmm/rmm/pylibrmm/cuda_stream.pyx | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pyx b/python/rmm/rmm/pylibrmm/cuda_stream.pyx index 61a43ef21..debc0c1fc 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pyx +++ b/python/rmm/rmm/pylibrmm/cuda_stream.pyx @@ -14,8 +14,9 @@ from rmm.pylibrmm.stream import CudaStreamFlags, _OwningStream sys.modules[__name__].CudaStream = _OwningStream warnings.warn( - "rmm.pylibrmm.cuda_stream is deprecated; use rmm.pylibrmm.stream for " - "CudaStreamFlags.", + "rmm.pylibrmm.cuda_stream is deprecated; " + "use the new rmm.pylibrmm.stream module for " + "CudaStreamFlags and CudaStream.", DeprecationWarning, stacklevel=2, ) From 6346a4c5a088b0da6edaffc864b62fa1b8cdaed7 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Tue, 24 Feb 2026 17:46:19 -0600 Subject: [PATCH 06/10] Apply copyright checks to pyi files --- .pre-commit-config.yaml | 2 +- python/rmm/rmm/pylibrmm/stream.pyi | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cadb26a61..6a529b8a7 100644 --- a/.pre-commit-config.yaml +++ b/.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$| diff --git a/python/rmm/rmm/pylibrmm/stream.pyi b/python/rmm/rmm/pylibrmm/stream.pyi index 8dd329c55..a994ea991 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyi +++ b/python/rmm/rmm/pylibrmm/stream.pyi @@ -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 enum import IntEnum From 7368d65d5705e1bfc774fdc46f7c3aec0c402cc3 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Tue, 24 Feb 2026 18:47:52 -0600 Subject: [PATCH 07/10] Fix deprecation message in cuda_stream shim and clean up _OwningStream leak --- python/rmm/rmm/pylibrmm/cuda_stream.pyx | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pyx b/python/rmm/rmm/pylibrmm/cuda_stream.pyx index debc0c1fc..6d5acb6b1 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pyx +++ b/python/rmm/rmm/pylibrmm/cuda_stream.pyx @@ -12,11 +12,12 @@ 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 the new rmm.pylibrmm.stream module for " - "CudaStreamFlags and CudaStream.", + "use rmm.pylibrmm.stream.CudaStreamFlags for stream flags, " + "and rmm.pylibrmm.stream.Stream for stream objects.", DeprecationWarning, stacklevel=2, ) From ecebb9af070d37c7809757886a2391f417726dd7 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Tue, 24 Feb 2026 20:13:42 -0600 Subject: [PATCH 08/10] Expose cuda_stream flags constructor through Cython binding and Stream API --- python/rmm/rmm/librmm/cuda_stream.pxd | 3 ++- python/rmm/rmm/pylibrmm/stream.pxd | 2 +- python/rmm/rmm/pylibrmm/stream.pyi | 6 ++++-- python/rmm/rmm/pylibrmm/stream.pyx | 18 ++++++++++++------ 4 files changed, 19 insertions(+), 10 deletions(-) diff --git a/python/rmm/rmm/librmm/cuda_stream.pxd b/python/rmm/rmm/librmm/cuda_stream.pxd index 0be62f804..0a1171eac 100644 --- a/python/rmm/rmm/librmm/cuda_stream.pxd +++ b/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 + diff --git a/python/rmm/rmm/pylibrmm/stream.pxd b/python/rmm/rmm/pylibrmm/stream.pxd index d55698e73..8f8ed62fe 100644 --- a/python/rmm/rmm/pylibrmm/stream.pxd +++ b/python/rmm/rmm/pylibrmm/stream.pxd @@ -28,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 * diff --git a/python/rmm/rmm/pylibrmm/stream.pyi b/python/rmm/rmm/pylibrmm/stream.pyi index a994ea991..100964822 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyi +++ b/python/rmm/rmm/pylibrmm/stream.pyi @@ -9,10 +9,12 @@ class CudaStreamFlags(IntEnum): NON_BLOCKING = ... class _OwningStream: - def __init__(self) -> None: ... + 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: ... diff --git a/python/rmm/rmm/pylibrmm/stream.pyx b/python/rmm/rmm/pylibrmm/stream.pyx index 8a4d2fc11..4c746a890 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyx +++ b/python/rmm/rmm/pylibrmm/stream.pyx @@ -38,9 +38,10 @@ cdef class _OwningStream: When an _OwningStream instance is GC'd, the underlying CUDA stream is destroyed. """ - def __cinit__(self): + def __cinit__(self, flags=CudaStreamFlags.SYNC_DEFAULT): + cdef cuda_stream_flags c_flags = (flags) with nogil: - self.c_obj.reset(new cuda_stream()) + self.c_obj.reset(new cuda_stream(c_flags)) def __dealloc__(self): with nogil: @@ -54,7 +55,7 @@ cdef class _OwningStream: cdef class Stream: - def __init__(self, obj=None): + def __init__(self, obj=None, flags=CudaStreamFlags.SYNC_DEFAULT): """ A Stream represents a CUDA stream. @@ -64,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: @@ -171,8 +175,10 @@ cdef class Stream: def __hash__(self): return hash(int((self._cuda_stream))) - cdef void _init_with_new_cuda_stream(self) except *: - cdef _OwningStream stream = _OwningStream() + 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 From e43195c00ed691dedaa0afd24971e22edd60b0fc Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Tue, 24 Feb 2026 20:13:51 -0600 Subject: [PATCH 09/10] Add test for DeprecationWarning on import of rmm.pylibrmm.cuda_stream --- python/rmm/rmm/tests/test_stream.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/rmm/rmm/tests/test_stream.py b/python/rmm/rmm/tests/test_stream.py index 8a6f123d2..ef13565b1 100644 --- a/python/rmm/rmm/tests/test_stream.py +++ b/python/rmm/rmm/tests/test_stream.py @@ -117,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() From f4f5ea56e49f2a5aeaa633625be7219074ce235d Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Tue, 24 Feb 2026 20:13:51 -0600 Subject: [PATCH 10/10] Add test for DeprecationWarning on import of rmm.pylibrmm.cuda_stream --- python/rmm/rmm/pylibrmm/stream.pyx | 33 ++++++++++++++++++----------- python/rmm/rmm/tests/test_stream.py | 24 +++++++++++++++++++++ 2 files changed, 45 insertions(+), 12 deletions(-) diff --git a/python/rmm/rmm/pylibrmm/stream.pyx b/python/rmm/rmm/pylibrmm/stream.pyx index 4c746a890..694995744 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyx +++ b/python/rmm/rmm/pylibrmm/stream.pyx @@ -55,26 +55,35 @@ cdef class _OwningStream: cdef class Stream: - def __init__(self, obj=None, flags=CudaStreamFlags.SYNC_DEFAULT): + def __init__(self, obj=None, flags=None): """ A Stream represents a CUDA stream. Parameters ---------- - obj: optional - * 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. + obj : object, optional + If None (the default), a new CUDA stream is created. + Otherwise a thin wrapper is created around an existing + stream from Numba, CuPy, or any object implementing the + CUDA stream protocol (``__cuda_stream__``). + flags : CudaStreamFlags, optional + Stream creation flags. Only valid when ``obj`` is None; + raises ``ValueError`` if provided alongside ``obj``. + Defaults to ``CudaStreamFlags.SYNC_DEFAULT``. """ if obj is None: - self._init_with_new_cuda_stream(flags) - elif isinstance(obj, Stream): - self._init_from_stream(obj) + self._init_with_new_cuda_stream( + CudaStreamFlags.SYNC_DEFAULT if flags is None else flags + ) else: - if hasattr(obj, "__cuda_stream__"): + if flags is not None: + raise ValueError( + "flags may only be specified when obj is None; " + "flags cannot be applied to an existing stream." + ) + if isinstance(obj, Stream): + self._init_from_stream(obj) + elif hasattr(obj, "__cuda_stream__"): self._init_from_cuda_stream_protocol(obj) else: try: diff --git a/python/rmm/rmm/tests/test_stream.py b/python/rmm/rmm/tests/test_stream.py index 8a6f123d2..12ac94cc8 100644 --- a/python/rmm/rmm/tests/test_stream.py +++ b/python/rmm/rmm/tests/test_stream.py @@ -117,6 +117,30 @@ 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_flags_with_obj_raises(): + existing = rmm.pylibrmm.stream.Stream() + with pytest.raises( + ValueError, match="flags may only be specified when obj is None" + ): + rmm.pylibrmm.stream.Stream( + existing, flags=rmm.pylibrmm.stream.CudaStreamFlags.NON_BLOCKING + ) + + def test_hashable(): a = rmm.pylibrmm.stream.Stream() b = rmm.pylibrmm.stream.Stream()