address review

rapidsai · Jan 8, 2025 · f80d832 · f80d832
1 parent 205ff92
commit f80d832
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 46 deletions.
diff --git a/python/rmm/rmm/_cuda/stream.py b/python/rmm/rmm/_cuda/stream.py
@@ -14,15 +14,22 @@
 
 import warnings
 
-from rmm.pylibrmm.stream import (  # noqa: F401
+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.pylibrmm.stream` module is deprecated in will be removed in a future release. Use `rmm.pylibrmm.stream` instead.",
+    "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,
 )
diff --git a/python/rmm/rmm/pylibrmm/stream.pyx b/python/rmm/rmm/pylibrmm/stream.pyx
@@ -66,50 +66,6 @@ cdef class Stream:
         """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 *:
         """