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/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/__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 0cabbf933..38205b1b1 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pxd +++ b/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 diff --git a/python/rmm/rmm/pylibrmm/cuda_stream.pyx b/python/rmm/rmm/pylibrmm/cuda_stream.pyx index 4287de7bd..6d5acb6b1 100644 --- a/python/rmm/rmm/pylibrmm/cuda_stream.pyx +++ b/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 = (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() +# 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", +] diff --git a/python/rmm/rmm/pylibrmm/stream.pxd b/python/rmm/rmm/pylibrmm/stream.pxd index ec0facf08..8f8ed62fe 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-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 * diff --git a/python/rmm/rmm/pylibrmm/stream.pyi b/python/rmm/rmm/pylibrmm/stream.pyi index eccf632c2..100964822 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyi +++ b/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: ... diff --git a/python/rmm/rmm/pylibrmm/stream.pyx b/python/rmm/rmm/pylibrmm/stream.pyx index 7b4f3abf3..694995744 100644 --- a/python/rmm/rmm/pylibrmm/stream.pyx +++ b/python/rmm/rmm/pylibrmm/stream.pyx @@ -1,37 +1,89 @@ -# 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 = (cuda_stream_flags.sync_default) + NON_BLOCKING = (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 = (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=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. + 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() - 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: @@ -132,8 +184,10 @@ cdef class Stream: def __hash__(self): return hash(int((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 diff --git a/python/rmm/rmm/tests/test_stream.py b/python/rmm/rmm/tests/test_stream.py index 5d7de48ed..f77fd2655 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): @@ -132,3 +131,27 @@ def test_hashable(): a2 = rmm.pylibrmm.stream.Stream(a) assert a2 == a assert hash(a2) == hash(a) + + +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 + )