Migrate stream APIs from rmm::cuda_stream_view to cuda::stream_ref

cpp/examples/hybrid_scan_io/hybrid_scan_multifile_single_step.cpp

@@ -16,6 +16,7 @@
 #include <rmm/cuda_stream_pool.hpp>
 #include <rmm/mr/statistics_resource_adaptor.hpp>
 
+#include <cassert>
 #include <ranges>
 #include <stdexcept>
 #include <string>
@@ -60,7 +61,8 @@ struct hybrid_scan_single_step_fn {
           input_sources[source_idx], {}, filters, false, stream, mr);
       }
     }
-    stream.synchronize_no_throw();
+    [[maybe_unused]] auto const status = cudaStreamSynchronize(stream.get());
+    assert(status == cudaSuccess);
     if (verbose) {
       std::cout << "Thread " << tid << " ";
       timer.print_elapsed_millis();

cpp/examples/hybrid_scan_io/hybrid_scan_multifile_two_step.cpp

@@ -16,6 +16,7 @@
 #include <rmm/cuda_stream_pool.hpp>
 #include <rmm/mr/statistics_resource_adaptor.hpp>
 
+#include <cassert>
 #include <ranges>
 #include <stdexcept>
 #include <string>
@@ -62,7 +63,8 @@ struct hybrid_scan_two_step_fn {
         input_sources[source_idx], filter_expression_opt, filters, false, stream, mr);
     }
 
-    stream.synchronize_no_throw();
+    [[maybe_unused]] auto const status = cudaStreamSynchronize(stream.get());
+    assert(status == cudaSuccess);
 
     if (verbose) {
       std::cout << "Thread " << tid << " ";
@@ -157,7 +159,7 @@ int main(int argc, char const** argv)
   // Create filter expressions (one per thread; reused circularly if needed)
   auto const column_reference = cudf::ast::column_name_reference(column_name);
   auto scalar                 = cudf::string_scalar(literal_value, true, default_stream);
-  default_stream.synchronize();
+  default_stream.sync();
   auto literal = cudf::ast::literal(scalar);
 
   std::vector<cudf::ast::operation> filter_expressions;

cpp/examples/parquet_io/parquet_io_multithreaded.cpp

@@ -147,7 +147,7 @@ std::vector<table_t> read_parquet_multithreaded(std::vector<io_source> const& in
   if (read_mode == read_mode::CONCATENATE_ALL) {
     auto stream    = stream_pool.get_stream();
     auto final_tbl = concatenate_tables(std::move(tables), stream);
-    stream.synchronize();
+    stream.sync();
     tables.clear();
     tables.emplace_back(std::move(final_tbl));
   }

cpp/libcudf_streaming/src/bloom_filter.cpp

@@ -118,7 +118,7 @@ rapidsmpf::streaming::Actor bloom_filter::apply(
   auto storage = (co_await bloom_filter->receive()).release<rmm::device_buffer>();
   RAPIDSMPF_EXPECTS((co_await bloom_filter->receive()).empty(),
                     "Bloom filter channel contained more than one message");
-  auto stream = cuda::stream_ref{storage.stream().get()};
+  auto stream = storage.stream();
   rapidsmpf::CudaEvent event;
   auto filter = cudf_streaming::detail::device_bloom_filter(filter_size_, seed_, storage.data());
   auto meta   = co_await ch_in->receive_metadata();

cpp/src/io/parquet/reader_impl_dict_transcode.cu

@@ -194,7 +194,7 @@ void remap_dict_indices_by_chunk(cudf::device_span<int32_t> indices,
                                  cudf::device_span<size_type const> row_offsets,
                                  cudf::device_span<size_type const> key_counts_prefix,
                                  cudf::device_span<int32_t const> stacked_to_unique,
-                                 rmm::cuda_stream_view stream)
+                                 cuda::stream_ref stream)
 {
   thrust::for_each(
     rmm::exec_policy_nosync(stream, get_current_device_resource_ref()),

cpp/tests/groupby/streaming_groupby_test.cpp

@@ -378,7 +378,7 @@ TEST_F(StreamingGroupbyTest, ConcurrentAggregate)
         std::this_thread::yield();
       }
       try {
-        streaming_agg.aggregate(batches[i], streams[i]->view());
+        streaming_agg.aggregate(batches[i], *streams[i]);
       } catch (...) {
         errors[i] = std::current_exception();
       }

cpp/tests/join/streaming_hash_join_tests.cpp

@@ -137,7 +137,7 @@ TEST_F(StreamingHashJoinTest, ConcurrentInsert)
                                    cudf::nullable_join::NO,
                                    cudf::null_equality::EQUAL,
                                    /*load_factor=*/0.5,
-                                   build_stream.view()};
+                                   build_stream};
   build_stream.synchronize();
 
   auto const device = rmm::get_current_cuda_device();
@@ -157,7 +157,7 @@ TEST_F(StreamingHashJoinTest, ConcurrentInsert)
         std::this_thread::yield();
       }
       try {
-        joiner.insert(right_partitions[i], streams[i]->view());
+        joiner.insert(right_partitions[i], *streams[i]);
       } catch (...) {
         errors[i] = std::current_exception();
       }

python/cudf_streaming/cudf_streaming/partition_utils.pyx

@@ -119,7 +119,7 @@ cpdef size_t partition_and_pack_cost(
     --------
     cudf_streaming.partition_utils.partition_and_pack
     """
-    cdef stream_ref _stream = stream_ref(stream.view().get())
+    cdef stream_ref _stream = stream.view()
     cdef cpp_BufferResource* _br = br.ptr()
     cdef table_view tbl = table.view()
     cdef size_t ret
@@ -177,7 +177,7 @@ cpdef object partition_and_pack(
     pylibcudf.contiguous_split.pack
     cudf_streaming.partition_utils.split_and_pack
     """
-    cdef stream_ref _stream = stream_ref(stream.view().get())
+    cdef stream_ref _stream = stream.view()
     cdef cpp_BufferResource* _br = br.ptr()
     cdef vector[size_type] _columns_to_hash = tuple(columns_to_hash)
     cdef unordered_map[uint32_t, cpp_PackedData] _ret
@@ -244,7 +244,7 @@ cpdef size_t split_and_pack_cost(
     --------
     cudf_streaming.partition_utils.split_and_pack
     """
-    cdef stream_ref _stream = stream_ref(stream.view().get())
+    cdef stream_ref _stream = stream.view()
     cdef cpp_BufferResource* _br = br.ptr()
     cdef table_view tbl = table.view()
     cdef size_t ret
@@ -299,7 +299,7 @@ cpdef object split_and_pack(
     pylibcudf.copying.split
     cudf_streaming.partition_utils.partition_and_pack
     """
-    cdef stream_ref _stream = stream_ref(stream.view().get())
+    cdef stream_ref _stream = stream.view()
     cdef cpp_BufferResource* _br = br.ptr()
     cdef vector[size_type] _splits = tuple(splits)
     cdef unordered_map[uint32_t, cpp_PackedData] _ret
@@ -463,7 +463,7 @@ cpdef object unpack_and_concat(
     cudf_streaming.partition_utils.unpack_and_concat_cost
     cudf_streaming.partition_utils.partition_and_pack
     """
-    cdef stream_ref _stream = stream_ref(stream.view().get())
+    cdef stream_ref _stream = stream.view()
     cdef cpp_BufferResource* _br = br.ptr()
     cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions)
     cdef unique_ptr[cpp_table] _ret
@@ -550,7 +550,7 @@ cpdef object packed_data_from_cudf_packed_columns(
     """
     if packed_columns is None or stream is None or br is None:
         raise TypeError("Arguments must not be None")
-    cdef stream_ref _stream = stream_ref(stream.view().get())
+    cdef stream_ref _stream = stream.view()
     cdef cpp_BufferResource* _br = br.ptr()
     cdef PackedData ret = PackedData.__new__(PackedData)
     with nogil:

python/cudf_streaming/cudf_streaming/table_chunk.pyx

@@ -183,7 +183,7 @@ cdef class TableChunk:
         persists even when the chunk is transferred through Channels.
 
         """
-        cdef stream_ref _stream = stream_ref(stream.view().get())
+        cdef stream_ref _stream = stream.view()
         cdef cpp_table_view view = table.view()
         return TableChunk.from_handle(
             cpp_from_table_view_with_owner(

python/pylibcudf/pylibcudf/typing.py

@@ -3,14 +3,16 @@
 
 from typing import Protocol, TypeAlias, TypedDict
 
+from cuda.bindings.runtime import cudaStream_t
+
 from rmm.pylibrmm.stream import Stream
 
 
 class HasCudaStream(Protocol):
     def __cuda_stream__(self) -> object: ...
 
 
-CudaStreamLike: TypeAlias = Stream | HasCudaStream
+CudaStreamLike: TypeAlias = Stream | HasCudaStream | cudaStream_t
 
 
 class ArrayInterfaceBase(TypedDict):

python/pylibcudf/pylibcudf/utils.pyx

@@ -3,7 +3,16 @@
 
 from cython.operator import dereference
 
-from libc.stdint cimport uint32_t
+from cuda.bindings import runtime
+
+from cuda.bindings.cyruntime cimport (
+    cudaError_t,
+    cudaFree,
+    cudaStream_t,
+    cudaSuccess,
+)
+
+from libc.stdint cimport uint32_t, uintptr_t
 from libcpp.functional cimport reference_wrapper
 from libcpp.optional cimport make_optional, nullopt, optional
 from libcpp.vector cimport vector
@@ -58,11 +67,25 @@ cdef vector[reference_wrapper[const scalar]] _as_vector(list source):
     return c_scalars
 
 
+cdef inline int _ensure_cuda_context() except -1:
+    cdef cudaError_t status
+    with nogil:
+        status = cudaFree(NULL)
+    if status != cudaSuccess:
+        raise RuntimeError(f"Failed to initialize CUDA context: {status}")
+    return 0
+
+
 cpdef Stream _get_stream(object stream: CudaStreamLike | None = None):
+    _ensure_cuda_context()
     if stream is None:
         return CUDF_DEFAULT_STREAM
     if isinstance(stream, Stream):
         return <Stream>stream
+    if isinstance(stream, runtime.cudaStream_t):
+        return Stream._from_cudaStream_t(
+            <cudaStream_t><uintptr_t>int(stream), owner=stream
+        )
     return Stream(stream)  # Handles __cuda_stream__ protocol
 
 

python/pylibcudf/tests/test_cuda_context.py

@@ -0,0 +1,89 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import subprocess
+import sys
+from concurrent.futures import ThreadPoolExecutor
+
+import pyarrow as pa
+import pytest
+from cuda.bindings import runtime
+
+import pylibcudf as plc
+
+
+@pytest.mark.parametrize(
+    "stream",
+    [
+        None,
+        runtime.cudaStream_t(runtime.cudaStreamDefault),
+        runtime.cudaStream_t(runtime.cudaStreamLegacy),
+        runtime.cudaStream_t(runtime.cudaStreamPerThread),
+    ],
+)
+def test_empty_like_on_fresh_thread(stream):
+    column = plc.Column.from_arrow(pa.array([1, 2, 3]))
+
+    def empty_like():
+        result = plc.copying.empty_like(column, stream=stream)
+        assert result.size() == 0
+        if stream is not None:
+            assert plc.utils._get_stream(stream).__cuda_stream__() == (
+                0,
+                int(stream),
+            )
+
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        pool.submit(empty_like).result()
+
+
+def test_get_stream_preserves_current_device():
+    status, count = runtime.cudaGetDeviceCount()
+    assert status == runtime.cudaError_t.cudaSuccess
+
+    def get_stream(device):
+        assert runtime.cudaSetDevice(device) == (
+            runtime.cudaError_t.cudaSuccess,
+        )
+        status, stream = runtime.cudaStreamCreate()
+        assert status == runtime.cudaError_t.cudaSuccess
+        try:
+            plc.utils._get_stream()
+            assert plc.utils._get_stream(stream).__cuda_stream__() == (
+                0,
+                int(stream),
+            )
+            assert runtime.cudaGetDevice() == (
+                runtime.cudaError_t.cudaSuccess,
+                device,
+            )
+        finally:
+            assert runtime.cudaStreamDestroy(stream) == (
+                runtime.cudaError_t.cudaSuccess,
+            )
+
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        for device in range(count):
+            pool.submit(get_stream, device).result()
+
+
+def test_get_stream_rejects_integer_handle():
+    with pytest.raises(TypeError):
+        plc.utils._get_stream(runtime.cudaStreamDefault)
+
+
+def test_empty_like_initializes_cuda(tmp_path):
+    subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            """
+import pylibcudf as plc
+
+column = plc.Column(plc.DataType(plc.TypeId.INT32), 0, None, None, 0, 0, [])
+assert plc.copying.empty_like(column).size() == 0
+""",
+        ],
+        check=True,
+        cwd=tmp_path,
+    )