From 0eaaefc278284435de578e5258ba2c4ef9d14e91 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Tue, 1 Sep 2026 17:16:44 -0500 Subject: [PATCH 01/12] Adapt stream accessors to cuda::stream_ref --- .../cudf_streaming/partition_utils.pyx | 12 ++++++------ python/cudf_streaming/cudf_streaming/table_chunk.pyx | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/cudf_streaming/cudf_streaming/partition_utils.pyx b/python/cudf_streaming/cudf_streaming/partition_utils.pyx index 5646bec5b83b..a7a143324a37 100644 --- a/python/cudf_streaming/cudf_streaming/partition_utils.pyx +++ b/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: diff --git a/python/cudf_streaming/cudf_streaming/table_chunk.pyx b/python/cudf_streaming/cudf_streaming/table_chunk.pyx index ed6a41674e49..735fb78d8622 100644 --- a/python/cudf_streaming/cudf_streaming/table_chunk.pyx +++ b/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( From 42007899bfac153d6a5d1e31c2265d7b785f11f6 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 2 Sep 2026 17:32:59 -0500 Subject: [PATCH 02/12] Use cuda::stream_ref in reader_impl_dict_transcode.cu --- cpp/src/io/parquet/reader_impl_dict_transcode.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/parquet/reader_impl_dict_transcode.cu b/cpp/src/io/parquet/reader_impl_dict_transcode.cu index 6c45678d3803..4ee6af6460d8 100644 --- a/cpp/src/io/parquet/reader_impl_dict_transcode.cu +++ b/cpp/src/io/parquet/reader_impl_dict_transcode.cu @@ -194,7 +194,7 @@ void remap_dict_indices_by_chunk(cudf::device_span indices, cudf::device_span row_offsets, cudf::device_span key_counts_prefix, cudf::device_span 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()), From 26732d44ac1fdaab9187cdf94afd37c9b5c85695 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 2 Sep 2026 21:29:59 -0500 Subject: [PATCH 03/12] Fix remaining cuda::stream_ref build errors --- cpp/libcudf_streaming/src/bloom_filter.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/libcudf_streaming/src/bloom_filter.cpp b/cpp/libcudf_streaming/src/bloom_filter.cpp index da857b99773b..7f77697cb399 100644 --- a/cpp/libcudf_streaming/src/bloom_filter.cpp +++ b/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(); 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(); From 9eb23ba123b67c3b749ed3b8e79dedde0cfa1592 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 00:13:46 -0500 Subject: [PATCH 04/12] Fix copyright and deprecations --- cpp/tests/join/streaming_hash_join_tests.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/tests/join/streaming_hash_join_tests.cpp b/cpp/tests/join/streaming_hash_join_tests.cpp index 55ef7b400e74..e70a994e763a 100644 --- a/cpp/tests/join/streaming_hash_join_tests.cpp +++ b/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(); } From e58849b513b1cd9db83572a8b657b119155a6e63 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 08:52:28 -0500 Subject: [PATCH 05/12] Avoid deprecated stream view in streaming groupby test Signed-off-by: Bradley Dice --- cpp/tests/groupby/streaming_groupby_test.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/groupby/streaming_groupby_test.cpp b/cpp/tests/groupby/streaming_groupby_test.cpp index 343ca6663ef2..ebef2d70c3c7 100644 --- a/cpp/tests/groupby/streaming_groupby_test.cpp +++ b/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(); } From ba5f035c661ea9da0ed5d1af372103e082a87a14 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 10:28:16 -0500 Subject: [PATCH 06/12] Fix stream synchronization in C++ examples Signed-off-by: Bradley Dice --- .../hybrid_scan_io/hybrid_scan_multifile_single_step.cpp | 4 +++- .../hybrid_scan_io/hybrid_scan_multifile_two_step.cpp | 6 ++++-- cpp/examples/parquet_io/parquet_io_multithreaded.cpp | 2 +- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/cpp/examples/hybrid_scan_io/hybrid_scan_multifile_single_step.cpp b/cpp/examples/hybrid_scan_io/hybrid_scan_multifile_single_step.cpp index 5b0c038d9af1..dff6392a25a4 100644 --- a/cpp/examples/hybrid_scan_io/hybrid_scan_multifile_single_step.cpp +++ b/cpp/examples/hybrid_scan_io/hybrid_scan_multifile_single_step.cpp @@ -16,6 +16,7 @@ #include #include +#include #include #include #include @@ -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(); diff --git a/cpp/examples/hybrid_scan_io/hybrid_scan_multifile_two_step.cpp b/cpp/examples/hybrid_scan_io/hybrid_scan_multifile_two_step.cpp index caadece7abc9..087ab0ba0d16 100644 --- a/cpp/examples/hybrid_scan_io/hybrid_scan_multifile_two_step.cpp +++ b/cpp/examples/hybrid_scan_io/hybrid_scan_multifile_two_step.cpp @@ -16,6 +16,7 @@ #include #include +#include #include #include #include @@ -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 filter_expressions; diff --git a/cpp/examples/parquet_io/parquet_io_multithreaded.cpp b/cpp/examples/parquet_io/parquet_io_multithreaded.cpp index f6832fd2f4e9..0697c1b53245 100644 --- a/cpp/examples/parquet_io/parquet_io_multithreaded.cpp +++ b/cpp/examples/parquet_io/parquet_io_multithreaded.cpp @@ -147,7 +147,7 @@ std::vector read_parquet_multithreaded(std::vector 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)); } From b429dc7d5d15cd4bda0b5441154243364921b6f1 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 13:21:09 -0500 Subject: [PATCH 07/12] Ensure a CUDA context is current when resolving pylibcudf streams --- python/pylibcudf/pylibcudf/utils.pyx | 20 ++++ python/pylibcudf/tests/test_cuda_context.py | 101 ++++++++++++++++++++ 2 files changed, 121 insertions(+) create mode 100644 python/pylibcudf/tests/test_cuda_context.py diff --git a/python/pylibcudf/pylibcudf/utils.pyx b/python/pylibcudf/pylibcudf/utils.pyx index f38cba03c710..657cc114ef4f 100644 --- a/python/pylibcudf/pylibcudf/utils.pyx +++ b/python/pylibcudf/pylibcudf/utils.pyx @@ -3,6 +3,15 @@ from cython.operator import dereference +from cuda.bindings.cydriver cimport ( + CUDA_ERROR_NOT_INITIALIZED, + CUDA_SUCCESS, + CUcontext, + CUresult, + cuCtxGetCurrent, +) +from cuda.bindings.cyruntime cimport cudaError_t, cudaFree, cudaSuccess + from libc.stdint cimport uint32_t from libcpp.functional cimport reference_wrapper from libcpp.optional cimport make_optional, nullopt, optional @@ -59,6 +68,17 @@ cdef vector[reference_wrapper[const scalar]] _as_vector(list source): cpdef Stream _get_stream(object stream: CudaStreamLike | None = None): + cdef CUcontext context = NULL + cdef CUresult status = cuCtxGetCurrent(&context) + cdef cudaError_t runtime_status + if status != CUDA_SUCCESS and status != CUDA_ERROR_NOT_INITIALIZED: + raise RuntimeError(f"Failed to get current CUDA context: {status}") + if status == CUDA_ERROR_NOT_INITIALIZED or context == NULL: + with nogil: + runtime_status = cudaFree(NULL) + if runtime_status != cudaSuccess: + raise RuntimeError(f"Failed to initialize CUDA context: {runtime_status}") + if stream is None: return CUDF_DEFAULT_STREAM if isinstance(stream, Stream): diff --git a/python/pylibcudf/tests/test_cuda_context.py b/python/pylibcudf/tests/test_cuda_context.py new file mode 100644 index 000000000000..62f894f7bfa4 --- /dev/null +++ b/python/pylibcudf/tests/test_cuda_context.py @@ -0,0 +1,101 @@ +# 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 driver + +from rmm.pylibrmm.stream import DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM + +import pylibcudf as plc + + +@pytest.mark.parametrize( + "stream", [None, DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM] +) +def test_get_stream_initializes_thread_context(stream): + assert driver.cuInit(0) == (driver.CUresult.CUDA_SUCCESS,) + + def get_stream(): + status, context = driver.cuCtxGetCurrent() + assert status == driver.CUresult.CUDA_SUCCESS + assert int(context) == 0 + plc.utils._get_stream(stream) + status, context = driver.cuCtxGetCurrent() + assert status == driver.CUresult.CUDA_SUCCESS + assert int(context) != 0 + plc.utils._get_stream(stream) + assert driver.cuCtxGetCurrent() == (status, context) + assert driver.cuCtxPopCurrent() == (status, context) + plc.utils._get_stream(stream) + assert driver.cuCtxGetCurrent() == (status, context) + + with ThreadPoolExecutor(max_workers=1) as pool: + pool.submit(get_stream).result() + + +def test_get_stream_preserves_current_context(): + assert driver.cuInit(0) == (driver.CUresult.CUDA_SUCCESS,) + status, device = driver.cuDeviceGet(0) + assert status == driver.CUresult.CUDA_SUCCESS + status, context = driver.cuDevicePrimaryCtxRetain(device) + assert status == driver.CUresult.CUDA_SUCCESS + + def get_stream(): + assert driver.cuCtxPushCurrent(context) == ( + driver.CUresult.CUDA_SUCCESS, + ) + try: + plc.utils._get_stream() + assert driver.cuCtxGetCurrent() == ( + driver.CUresult.CUDA_SUCCESS, + context, + ) + finally: + assert driver.cuCtxPopCurrent() == ( + driver.CUresult.CUDA_SUCCESS, + context, + ) + status, current = driver.cuCtxGetCurrent() + assert status == driver.CUresult.CUDA_SUCCESS + assert int(current) == 0 + + try: + with ThreadPoolExecutor(max_workers=1) as pool: + pool.submit(get_stream).result() + finally: + assert driver.cuDevicePrimaryCtxRelease(device) == ( + driver.CUresult.CUDA_SUCCESS, + ) + + +def test_empty_like_on_fresh_thread(): + column = plc.Column.from_arrow(pa.array([1, 2, 3])) + with ThreadPoolExecutor(max_workers=1) as pool: + result = pool.submit(plc.copying.empty_like, column).result() + assert result.size() == 0 + + +def test_get_stream_initializes_cuda(): + subprocess.run( + [ + sys.executable, + "-c", + """ +from cuda.bindings import driver +import pylibcudf as plc + +status, _ = driver.cuCtxGetCurrent() +assert status == driver.CUresult.CUDA_ERROR_NOT_INITIALIZED +plc.utils._get_stream() +status, context = driver.cuCtxGetCurrent() +assert status == driver.CUresult.CUDA_SUCCESS +assert int(context) != 0 +""", + ], + check=True, + ) From 6ed804118ba649625c77f17dfe446ce967055c04 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 13:25:25 -0500 Subject: [PATCH 08/12] Extract CUDA context initialization into a helper --- python/pylibcudf/pylibcudf/utils.pyx | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/pylibcudf/pylibcudf/utils.pyx b/python/pylibcudf/pylibcudf/utils.pyx index 657cc114ef4f..7e824efcbcec 100644 --- a/python/pylibcudf/pylibcudf/utils.pyx +++ b/python/pylibcudf/pylibcudf/utils.pyx @@ -67,7 +67,7 @@ cdef vector[reference_wrapper[const scalar]] _as_vector(list source): return c_scalars -cpdef Stream _get_stream(object stream: CudaStreamLike | None = None): +cdef inline int _ensure_cuda_context() except -1: cdef CUcontext context = NULL cdef CUresult status = cuCtxGetCurrent(&context) cdef cudaError_t runtime_status @@ -78,7 +78,11 @@ cpdef Stream _get_stream(object stream: CudaStreamLike | None = None): runtime_status = cudaFree(NULL) if runtime_status != cudaSuccess: raise RuntimeError(f"Failed to initialize CUDA context: {runtime_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): From 18c6174ea32739b3059e43cde6e96295b75387cf Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 13:27:07 -0500 Subject: [PATCH 09/12] Use CUDA bindings stream handles in context tests --- python/pylibcudf/tests/test_cuda_context.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/python/pylibcudf/tests/test_cuda_context.py b/python/pylibcudf/tests/test_cuda_context.py index 62f894f7bfa4..e8990062631c 100644 --- a/python/pylibcudf/tests/test_cuda_context.py +++ b/python/pylibcudf/tests/test_cuda_context.py @@ -7,15 +7,18 @@ import pyarrow as pa import pytest -from cuda.bindings import driver - -from rmm.pylibrmm.stream import DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM +from cuda.bindings import driver, runtime import pylibcudf as plc @pytest.mark.parametrize( - "stream", [None, DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM] + "stream", + [ + None, + runtime.cudaStream_t(runtime.cudaStreamDefault), + runtime.cudaStream_t(runtime.cudaStreamPerThread), + ], ) def test_get_stream_initializes_thread_context(stream): assert driver.cuInit(0) == (driver.CUresult.CUDA_SUCCESS,) From 0fae2a3f76f23e2cc83b9f990e77683dadd4adfa Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 14:56:50 -0500 Subject: [PATCH 10/12] Add CUDA profiler headers to pylibcudf host requirements --- conda/recipes/pylibcudf/recipe.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/conda/recipes/pylibcudf/recipe.yaml b/conda/recipes/pylibcudf/recipe.yaml index 22f6fca29ae6..c5e82d41744a 100644 --- a/conda/recipes/pylibcudf/recipe.yaml +++ b/conda/recipes/pylibcudf/recipe.yaml @@ -79,6 +79,7 @@ requirements: - libcudf =${{ version }} - rmm =${{ minor_version }} - cuda-cudart-dev + - cuda-profiler-api - cuda-nvrtc - if: linux and x86_64 then: From ad4dd186a47ca2718d078d616d838bab4bba2c4f Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 15:37:20 -0500 Subject: [PATCH 11/12] Use CUDA runtime APIs for pylibcudf stream handling --- conda/recipes/pylibcudf/recipe.yaml | 1 - python/pylibcudf/pylibcudf/typing.py | 4 +- python/pylibcudf/pylibcudf/utils.pyx | 35 ++++---- python/pylibcudf/tests/test_cuda_context.py | 98 +++++++++------------ 4 files changed, 61 insertions(+), 77 deletions(-) diff --git a/conda/recipes/pylibcudf/recipe.yaml b/conda/recipes/pylibcudf/recipe.yaml index c5e82d41744a..22f6fca29ae6 100644 --- a/conda/recipes/pylibcudf/recipe.yaml +++ b/conda/recipes/pylibcudf/recipe.yaml @@ -79,7 +79,6 @@ requirements: - libcudf =${{ version }} - rmm =${{ minor_version }} - cuda-cudart-dev - - cuda-profiler-api - cuda-nvrtc - if: linux and x86_64 then: diff --git a/python/pylibcudf/pylibcudf/typing.py b/python/pylibcudf/pylibcudf/typing.py index 965bcd054054..1c60a65e3bda 100644 --- a/python/pylibcudf/pylibcudf/typing.py +++ b/python/pylibcudf/pylibcudf/typing.py @@ -3,6 +3,8 @@ from typing import Protocol, TypeAlias, TypedDict +from cuda.bindings.runtime import cudaStream_t + from rmm.pylibrmm.stream import Stream @@ -10,7 +12,7 @@ class HasCudaStream(Protocol): def __cuda_stream__(self) -> object: ... -CudaStreamLike: TypeAlias = Stream | HasCudaStream +CudaStreamLike: TypeAlias = Stream | HasCudaStream | cudaStream_t class ArrayInterfaceBase(TypedDict): diff --git a/python/pylibcudf/pylibcudf/utils.pyx b/python/pylibcudf/pylibcudf/utils.pyx index 7e824efcbcec..68b32659c552 100644 --- a/python/pylibcudf/pylibcudf/utils.pyx +++ b/python/pylibcudf/pylibcudf/utils.pyx @@ -3,16 +3,16 @@ from cython.operator import dereference -from cuda.bindings.cydriver cimport ( - CUDA_ERROR_NOT_INITIALIZED, - CUDA_SUCCESS, - CUcontext, - CUresult, - cuCtxGetCurrent, +from cuda.bindings import runtime + +from cuda.bindings.cyruntime cimport ( + cudaError_t, + cudaFree, + cudaStream_t, + cudaSuccess, ) -from cuda.bindings.cyruntime cimport cudaError_t, cudaFree, cudaSuccess -from libc.stdint cimport uint32_t +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 @@ -68,16 +68,11 @@ cdef vector[reference_wrapper[const scalar]] _as_vector(list source): cdef inline int _ensure_cuda_context() except -1: - cdef CUcontext context = NULL - cdef CUresult status = cuCtxGetCurrent(&context) - cdef cudaError_t runtime_status - if status != CUDA_SUCCESS and status != CUDA_ERROR_NOT_INITIALIZED: - raise RuntimeError(f"Failed to get current CUDA context: {status}") - if status == CUDA_ERROR_NOT_INITIALIZED or context == NULL: - with nogil: - runtime_status = cudaFree(NULL) - if runtime_status != cudaSuccess: - raise RuntimeError(f"Failed to initialize CUDA context: {runtime_status}") + cdef cudaError_t status + with nogil: + status = cudaFree(NULL) + if status != cudaSuccess: + raise RuntimeError(f"Failed to initialize CUDA context: {status}") return 0 @@ -87,6 +82,10 @@ cpdef Stream _get_stream(object stream: CudaStreamLike | None = None): return CUDF_DEFAULT_STREAM if isinstance(stream, Stream): return stream + if isinstance(stream, runtime.cudaStream_t): + return Stream._from_cudaStream_t( + int(stream), owner=stream + ) return Stream(stream) # Handles __cuda_stream__ protocol diff --git a/python/pylibcudf/tests/test_cuda_context.py b/python/pylibcudf/tests/test_cuda_context.py index e8990062631c..ab9b058d4a08 100644 --- a/python/pylibcudf/tests/test_cuda_context.py +++ b/python/pylibcudf/tests/test_cuda_context.py @@ -7,7 +7,7 @@ import pyarrow as pa import pytest -from cuda.bindings import driver, runtime +from cuda.bindings import runtime import pylibcudf as plc @@ -17,87 +17,71 @@ [ None, runtime.cudaStream_t(runtime.cudaStreamDefault), + runtime.cudaStream_t(runtime.cudaStreamLegacy), runtime.cudaStream_t(runtime.cudaStreamPerThread), ], ) -def test_get_stream_initializes_thread_context(stream): - assert driver.cuInit(0) == (driver.CUresult.CUDA_SUCCESS,) - - def get_stream(): - status, context = driver.cuCtxGetCurrent() - assert status == driver.CUresult.CUDA_SUCCESS - assert int(context) == 0 - plc.utils._get_stream(stream) - status, context = driver.cuCtxGetCurrent() - assert status == driver.CUresult.CUDA_SUCCESS - assert int(context) != 0 - plc.utils._get_stream(stream) - assert driver.cuCtxGetCurrent() == (status, context) - assert driver.cuCtxPopCurrent() == (status, context) - plc.utils._get_stream(stream) - assert driver.cuCtxGetCurrent() == (status, context) +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(get_stream).result() + pool.submit(empty_like).result() -def test_get_stream_preserves_current_context(): - assert driver.cuInit(0) == (driver.CUresult.CUDA_SUCCESS,) - status, device = driver.cuDeviceGet(0) - assert status == driver.CUresult.CUDA_SUCCESS - status, context = driver.cuDevicePrimaryCtxRetain(device) - assert status == driver.CUresult.CUDA_SUCCESS +def test_get_stream_preserves_current_device(): + status, count = runtime.cudaGetDeviceCount() + assert status == runtime.cudaError_t.cudaSuccess - def get_stream(): - assert driver.cuCtxPushCurrent(context) == ( - driver.CUresult.CUDA_SUCCESS, + 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 driver.cuCtxGetCurrent() == ( - driver.CUresult.CUDA_SUCCESS, - context, + assert plc.utils._get_stream(stream).__cuda_stream__() == ( + 0, + int(stream), + ) + assert runtime.cudaGetDevice() == ( + runtime.cudaError_t.cudaSuccess, + device, ) finally: - assert driver.cuCtxPopCurrent() == ( - driver.CUresult.CUDA_SUCCESS, - context, + assert runtime.cudaStreamDestroy(stream) == ( + runtime.cudaError_t.cudaSuccess, ) - status, current = driver.cuCtxGetCurrent() - assert status == driver.CUresult.CUDA_SUCCESS - assert int(current) == 0 - - try: - with ThreadPoolExecutor(max_workers=1) as pool: - pool.submit(get_stream).result() - finally: - assert driver.cuDevicePrimaryCtxRelease(device) == ( - driver.CUresult.CUDA_SUCCESS, - ) - -def test_empty_like_on_fresh_thread(): - column = plc.Column.from_arrow(pa.array([1, 2, 3])) with ThreadPoolExecutor(max_workers=1) as pool: - result = pool.submit(plc.copying.empty_like, column).result() - assert result.size() == 0 + 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_get_stream_initializes_cuda(): +def test_empty_like_initializes_cuda(): subprocess.run( [ sys.executable, "-c", """ -from cuda.bindings import driver import pylibcudf as plc -status, _ = driver.cuCtxGetCurrent() -assert status == driver.CUresult.CUDA_ERROR_NOT_INITIALIZED -plc.utils._get_stream() -status, context = driver.cuCtxGetCurrent() -assert status == driver.CUresult.CUDA_SUCCESS -assert int(context) != 0 +column = plc.Column(plc.DataType(plc.TypeId.INT32), 0, None, None, 0, 0, []) +assert plc.copying.empty_like(column).size() == 0 """, ], check=True, From 815439d1ca2d4b0c9600b46559335170d860ff27 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 16:34:46 -0500 Subject: [PATCH 12/12] Run CUDA initialization subprocess outside the source checkout --- python/pylibcudf/tests/test_cuda_context.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/pylibcudf/tests/test_cuda_context.py b/python/pylibcudf/tests/test_cuda_context.py index ab9b058d4a08..1d0c55414955 100644 --- a/python/pylibcudf/tests/test_cuda_context.py +++ b/python/pylibcudf/tests/test_cuda_context.py @@ -72,7 +72,7 @@ def test_get_stream_rejects_integer_handle(): plc.utils._get_stream(runtime.cudaStreamDefault) -def test_empty_like_initializes_cuda(): +def test_empty_like_initializes_cuda(tmp_path): subprocess.run( [ sys.executable, @@ -85,4 +85,5 @@ def test_empty_like_initializes_cuda(): """, ], check=True, + cwd=tmp_path, )