From 82edb84e380d267833d5cb14920b926a5c7e43d4 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 8 Sep 2026 07:50:19 -0700 Subject: [PATCH 1/2] Revert "Remove RMM/RAFT dependencies, Part 1 (#195)" This reverts commit a7763a2d54ea54153d198f6c0e3a2f8c2cd3905e. --- .../all_cuda-129_arch-aarch64.yaml | 1 - .../all_cuda-129_arch-x86_64.yaml | 1 - .../all_cuda-133_arch-aarch64.yaml | 1 - .../all_cuda-133_arch-x86_64.yaml | 1 - cpp/CMakeLists.txt | 6 +- cpp/cmake/thirdparty/get_rmm.cmake | 13 ++ cpp/include/nvforest/detail/device_id/gpu.hpp | 11 +- .../nvforest/detail/device_setter/gpu.hpp | 18 +-- cpp/include/nvforest/detail/infer/gpu.cuh | 17 --- .../nvforest/detail/owning_buffer/gpu.hpp | 38 ++--- cpp/include/nvforest/forest_model.hpp | 135 ++++++++++++++---- cpp/include/nvforest/handle.hpp | 43 ++++++ cpp/include/nvforest/treelite_importer.hpp | 1 + cpp/src/detail/device_buffer.cu | 50 ------- cpp/tests/CMakeLists.txt | 5 +- cpp/tests/invalid_stream.cu | 82 ----------- cpp/tests/treelite_importer.cpp | 18 ++- dependencies.yaml | 8 -- python/nvforest/CMakeLists.txt | 6 - python/nvforest/nvforest/__init__.py | 6 +- python/nvforest/nvforest/_factory.py | 47 ++---- python/nvforest/nvforest/_forest_inference.py | 54 +++---- python/nvforest/nvforest/_typing.py | 11 +- .../nvforest/nvforest/detail/cuda_stream.pxd | 6 +- .../nvforest/detail/forest_inference.pyx | 69 +++------ python/nvforest/nvforest/detail/handle.pxd | 17 +++ .../nvforest/nvforest/detail/raft_stream.pyx | 35 ----- python/nvforest/pyproject.toml | 1 - python/nvforest/tests/test_nvforest.py | 34 ----- 29 files changed, 284 insertions(+), 451 deletions(-) create mode 100644 cpp/cmake/thirdparty/get_rmm.cmake create mode 100644 cpp/include/nvforest/handle.hpp delete mode 100644 cpp/src/detail/device_buffer.cu delete mode 100644 cpp/tests/invalid_stream.cu create mode 100644 python/nvforest/nvforest/detail/handle.pxd delete mode 100644 python/nvforest/nvforest/detail/raft_stream.pyx diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 056840f2..6dd75f5c 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -11,7 +11,6 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=12.9.2,<13.0 -- cuda-core>=1.0.0,<2.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-profiler-api diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 915d79ee..d9e43d6c 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -11,7 +11,6 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=12.9.2,<13.0 -- cuda-core>=1.0.0,<2.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-profiler-api diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index 16caada7..fdbe67eb 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -11,7 +11,6 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=13.0.1,<14.0 -- cuda-core>=1.0.0,<2.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-profiler-api diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index b43abc1e..eb3d2c2c 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -11,7 +11,6 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=13.0.1,<14.0 -- cuda-core>=1.0.0,<2.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-profiler-api diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index d0f8d2b2..885fa544 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -199,6 +199,7 @@ endif() # CCCL before RMM, and RMM before RAFT include(cmake/thirdparty/get_cccl.cmake) if(NVFOREST_ENABLE_GPU) + include(cmake/thirdparty/get_rmm.cmake) include(cmake/thirdparty/get_raft.cmake) endif() include(cmake/thirdparty/get_treelite.cmake) @@ -251,8 +252,7 @@ if(NVFOREST_ENABLE_GPU) src/infer8.cu src/infer9.cu src/infer10.cu - src/infer11.cu - src/detail/device_buffer.cu) + src/infer11.cu) endif() target_sources( ${NVFOREST_CPP_TARGET} @@ -312,7 +312,7 @@ elseif(NVFOREST_EXPORT_TREELITE_LINKAGE) endif() if(NVFOREST_ENABLE_GPU) - list(APPEND _nvforest_cpp_public_libs raft::raft CUDA::cudart_static CCCL::CCCL) + list(APPEND _nvforest_cpp_public_libs rmm::rmm raft::raft CUDA::cudart_static) endif() # These are always private: diff --git a/cpp/cmake/thirdparty/get_rmm.cmake b/cpp/cmake/thirdparty/get_rmm.cmake new file mode 100644 index 00000000..b476b09a --- /dev/null +++ b/cpp/cmake/thirdparty/get_rmm.cmake @@ -0,0 +1,13 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +function(find_and_configure_rmm) + include(${rapids-cmake-dir}/cpm/rmm.cmake) + rapids_cpm_rmm(BUILD_EXPORT_SET nvforest-exports INSTALL_EXPORT_SET nvforest-exports) +endfunction() + +find_and_configure_rmm() diff --git a/cpp/include/nvforest/detail/device_id/gpu.hpp b/cpp/include/nvforest/detail/device_id/gpu.hpp index 10d10119..49edfdd6 100644 --- a/cpp/include/nvforest/detail/device_id/gpu.hpp +++ b/cpp/include/nvforest/detail/device_id/gpu.hpp @@ -7,20 +7,23 @@ #include #include +#include + namespace nvforest::detail { template <> struct device_id { + using value_type = typename rmm::cuda_device_id::value_type; device_id() noexcept(false) : id_{[]() { - auto raw_id = int{}; + auto raw_id = value_type{}; cuda_check(cudaGetDevice(&raw_id)); return raw_id; }()} {}; - device_id(int dev_id) noexcept : id_{dev_id} {}; + device_id(value_type dev_id) noexcept : id_{dev_id} {}; - auto value() const noexcept { return id_; } + auto value() const noexcept { return id_.value(); } private: - int id_; + rmm::cuda_device_id id_; }; } // namespace nvforest::detail diff --git a/cpp/include/nvforest/detail/device_setter/gpu.hpp b/cpp/include/nvforest/detail/device_setter/gpu.hpp index 07ba9a34..d69a0ff5 100644 --- a/cpp/include/nvforest/detail/device_setter/gpu.hpp +++ b/cpp/include/nvforest/detail/device_setter/gpu.hpp @@ -8,21 +8,9 @@ #include #include -#include - -#include +#include -#define NVFOREST_CUDA_TRY_NO_THROW(call) \ - do { \ - cudaError_t const status = call; \ - if (cudaSuccess != status) { \ - printf("CUDA call='%s' at file=%s line=%d failed with %s\n", \ - #call, \ - __FILE__, \ - __LINE__, \ - cudaGetErrorString(status)); \ - } \ - } while (0) +#include namespace nvforest::detail { @@ -39,7 +27,7 @@ struct device_setter { cuda_check(cudaSetDevice(device.value())); } - ~device_setter() { NVFOREST_CUDA_TRY_NO_THROW(cudaSetDevice(prev_device_.value())); } + ~device_setter() { RAFT_CUDA_TRY_NO_THROW(cudaSetDevice(prev_device_.value())); } private: device_id prev_device_; diff --git a/cpp/include/nvforest/detail/infer/gpu.cuh b/cpp/include/nvforest/detail/infer/gpu.cuh index 31cc8eaa..744f48cf 100644 --- a/cpp/include/nvforest/detail/infer/gpu.cuh +++ b/cpp/include/nvforest/detail/infer/gpu.cuh @@ -21,8 +21,6 @@ #include #include -#include - #include #include #include @@ -43,19 +41,6 @@ inline auto compute_output_size(index_type row_output_size, return result; } -// If a non-default stream is provided, it must reside on the correct device. -inline void validate_stream(device_id device, cuda::stream_ref stream) -{ - cuda::stream_ref legacy_default_stream{cudaStreamLegacy}; - cuda::stream_ref per_thread_default_stream{cudaStreamPerThread}; - if (stream != legacy_default_stream && stream != per_thread_default_stream && - stream.device().get() != device.value()) { - throw std::runtime_error{std::string("Stream on the wrong device. ") + - "Expected: " + std::to_string(device.value()) + - ", Actual: " + std::to_string(stream.device().get())}; - } -} - /* A wrapper around the underlying inference kernels to support dispatching to * the right kernel * @@ -114,8 +99,6 @@ std::enable_if_t infer( { using output_t = typename forest_t::template raw_output_type; - validate_stream(device, cuda::stream_ref{stream}); - auto sm_count = get_sm_count(device); auto const max_shared_mem_per_block = get_max_shared_mem_per_block(device); auto const max_shared_mem_per_sm = get_max_shared_mem_per_sm(device); diff --git a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp index 2e217d64..ba6030ff 100644 --- a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp +++ b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp @@ -3,51 +3,37 @@ * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include #include +#include #include #include -#include +#include + #include -#include -#include #include namespace nvforest::detail { - -struct owning_device_buffer_type_erased_impl; - -struct owning_device_buffer_type_erased { - owning_device_buffer_type_erased(); - owning_device_buffer_type_erased(device_id device_id, - std::size_t size, - cuda::stream_ref stream); - owning_device_buffer_type_erased(owning_device_buffer_type_erased&& other) noexcept; - owning_device_buffer_type_erased& operator=(owning_device_buffer_type_erased&& other) noexcept; - ~owning_device_buffer_type_erased(); - std::byte* get(); - - private: - std::unique_ptr impl_; -}; - template struct owning_buffer { // TODO(wphicks): Assess need for buffers of const T using value_type = std::remove_const_t; - owning_buffer() = default; + owning_buffer() : data_{} {} + owning_buffer(device_id device_id, std::size_t size, - cuda_stream stream) noexcept(false) - : data_{device_id, size * sizeof(value_type), cuda::stream_ref{stream}} + cudaStream_t stream) noexcept(false) + : data_{[&device_id, &size, &stream]() { + auto device_context = device_setter{device_id}; + return rmm::device_buffer{size * sizeof(value_type), rmm::cuda_stream_view{stream}}; + }()} { } - auto* get() const { return reinterpret_cast(data_.get()); } + auto* get() const { return reinterpret_cast(data_.data()); } private: - mutable owning_device_buffer_type_erased data_; + mutable rmm::device_buffer data_; }; } // namespace nvforest::detail diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index dea577e5..143881f4 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -6,11 +6,13 @@ #include #include #include +#include #include #include #include #include #include +#include #include #ifdef NVFOREST_ENABLE_GPU @@ -116,12 +118,12 @@ struct forest_model { /** * Perform inference on given input * - * @param[in] stream A nvforest::cuda_stream, which (on GPU-enabled builds) is - * a transparent wrapper for the cudaStream_t or (on CPU-only builds) a - * CUDA-free placeholder object. * @param[out] output The buffer where model output should be stored. * This must be of size at least ROWS x num_outputs(). * @param[in] input The buffer containing input data. + * @param[in] stream A nvforest::cuda_stream, which (on GPU-enabled builds) is + * a transparent wrapper for the cudaStream_t or (on CPU-only builds) a + * CUDA-free placeholder object. * @param[in] predict_type Type of inference to perform. Defaults to summing * the outputs of all trees and produce an output per row. If set to * "per_tree", we will instead output all outputs of individual trees. @@ -137,40 +139,113 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(cuda_stream stream, - buffer& output, + void predict(buffer& output, buffer const& input, + cuda_stream stream = cuda_stream{}, infer_kind predict_type = infer_kind::default_kind, std::optional specified_chunk_size = std::nullopt) { std::visit( [this, predict_type, &output, &input, &stream, &specified_chunk_size]( + auto&& concrete_forest) { + if constexpr (std::is_same_v< + typename std::remove_reference_t::io_type, + io_t>) { + concrete_forest.predict(output, input, stream, predict_type, specified_chunk_size); + } else { + throw type_error("Input type does not match model_type"); + } + }, + decision_forest_); + } + + /** + * Perform inference on given input + * + * @param[in] handle The nvforest::handle_t (wrapper for raft::handle_t + * on GPU) which will be used to provide streams for evaluation. + * @param[out] output The buffer where model output should be stored. If + * this buffer is on host while the model is on device or vice versa, + * work will be distributed across available streams to copy the data back + * to this output location. This must be of size at least ROWS x num_outputs(). + * @param[in] input The buffer containing input data. If + * this buffer is on host while the model is on device or vice versa, + * work will be distributed across available streams to copy the input data + * to the appropriate location and perform inference. + * @param[in] predict_type Type of inference to perform. Defaults to summing + * the outputs of all trees and produce an output per row. If set to + * "per_tree", we will instead output all outputs of individual trees. + * If set to "leaf_id", we will output the integer ID of the leaf node + * for each tree. + * @param[in] specified_chunk_size: Specifies the mini-batch size for + * processing. This has different meanings on CPU and GPU, but on GPU it + * corresponds to the number of rows evaluated per inference iteration + * on a single block. It can take on any power of 2 from 1 to 32, and + * runtime performance is quite sensitive to the value chosen. In general, + * larger batches benefit from higher values, but it is hard to predict the + * optimal value a priori. If omitted, a heuristic will be used to select a + * reasonable value. On CPU, this argument can generally just be omitted. + */ + template + void predict(handle_t const& handle, + buffer& output, + buffer const& input, + infer_kind predict_type = infer_kind::default_kind, + std::optional specified_chunk_size = std::nullopt) + { + std::visit( + [this, predict_type, &handle, &output, &input, &specified_chunk_size]( auto&& concrete_forest) { using model_io_t = typename std::remove_reference_t::io_type; if constexpr (std::is_same_v) { if (output.memory_type() == memory_type() && input.memory_type() == memory_type()) { - concrete_forest.predict(output, input, stream, predict_type, specified_chunk_size); + concrete_forest.predict( + output, input, handle.get_next_usable_stream(), predict_type, specified_chunk_size); } else { + auto constexpr static const MIN_CHUNKS_PER_PARTITION = std::size_t{64}; + auto constexpr static const MAX_CHUNK_SIZE = std::size_t{64}; + auto row_count = input.size() / num_features(); - auto buffer_in = buffer{}; - if (input.memory_type() != memory_type()) { - buffer_in = - buffer{row_count * num_features(), memory_type(), device_index(), stream}; - copy(buffer_in, input, 0, 0, buffer_in.size(), stream); - } else { - buffer_in = buffer{input.data(), row_count * num_features(), memory_type()}; - } - auto buffer_out = buffer{}; - if (output.memory_type() != memory_type()) { - buffer_out = - buffer{row_count * num_outputs(), memory_type(), device_index(), stream}; - } else { - buffer_out = buffer{output.data(), row_count * num_outputs(), memory_type()}; - } - concrete_forest.predict( - buffer_out, buffer_in, stream, predict_type, specified_chunk_size); - if (output.memory_type() != memory_type()) { - copy(output, buffer_out, 0, 0, buffer_out.size(), stream); + auto partition_size = + std::max(detail::ceildiv(row_count, handle.get_usable_stream_count()), + specified_chunk_size.value_or(MAX_CHUNK_SIZE) * MIN_CHUNKS_PER_PARTITION); + auto partition_count = detail::ceildiv(row_count, partition_size); + for (auto i = std::size_t{}; i < partition_count; ++i) { + auto stream = handle.get_next_usable_stream(); + auto rows_in_this_partition = + std::min(partition_size, row_count - i * partition_size); + auto partition_in = buffer{}; + if (input.memory_type() != memory_type()) { + partition_in = buffer{rows_in_this_partition * num_features(), memory_type()}; + copy(partition_in, + input, + 0, + i * partition_size * num_features(), + partition_in.size(), + stream); + } else { + partition_in = buffer{input.data() + i * partition_size * num_features(), + rows_in_this_partition * num_features(), + memory_type()}; + } + auto partition_out = buffer{}; + if (output.memory_type() != memory_type()) { + partition_out = buffer{rows_in_this_partition * num_outputs(), memory_type()}; + } else { + partition_out = buffer{output.data() + i * partition_size * num_outputs(), + rows_in_this_partition * num_outputs(), + memory_type()}; + } + concrete_forest.predict( + partition_out, partition_in, stream, predict_type, specified_chunk_size); + if (output.memory_type() != memory_type()) { + copy(output, + partition_out, + i * partition_size * num_outputs(), + 0, + partition_out.size(), + stream); + } } } } else { @@ -183,10 +258,8 @@ struct forest_model { /** * Perform inference on given input * - * @param[in] stream A nvforest::cuda_stream, which (on GPU-enabled builds) is - * a transparent wrapper for the cudaStream_t or (on CPU-only builds) a - * CUDA-free placeholder object. If this argument is cudaStream_t, ensure - * that this stream is associated with the same GPU device as the model object. + * @param[in] handle The nvforest::handle_t (wrapper for raft::handle_t + * on GPU) which will be used to provide streams for evaluation. * @param[out] output Pointer to the memory location where output should end * up * @param[in] input Pointer to the input data @@ -209,7 +282,7 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(cuda_stream stream, + void predict(handle_t const& handle, io_t* output, io_t* input, std::size_t num_rows, @@ -230,7 +303,7 @@ struct forest_model { } auto out_buffer = buffer{output, num_rows * num_outputs(), out_mem_type, current_device_id}; auto in_buffer = buffer{input, num_rows * num_features(), in_mem_type, current_device_id}; - predict(stream, out_buffer, in_buffer, predict_type, specified_chunk_size); + predict(handle, out_buffer, in_buffer, predict_type, specified_chunk_size); } private: diff --git a/cpp/include/nvforest/handle.hpp b/cpp/include/nvforest/handle.hpp new file mode 100644 index 00000000..fd1f3bb7 --- /dev/null +++ b/cpp/include/nvforest/handle.hpp @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once +#include + +#include +#include +#ifdef NVFOREST_ENABLE_GPU +#include +#endif + +namespace nvforest { +#ifdef NVFOREST_ENABLE_GPU +struct handle_t { + handle_t(raft::handle_t const* handle_ptr = nullptr) : raft_handle_{handle_ptr} {} + handle_t(raft::handle_t const& raft_handle) : raft_handle_{&raft_handle} {} + auto get_next_usable_stream() const + { + return cuda_stream{raft_handle_->get_next_usable_stream().value()}; + } + auto get_stream_pool_size() const { return raft_handle_->get_stream_pool_size(); } + auto get_usable_stream_count() const { return std::max(get_stream_pool_size(), std::size_t{1}); } + void synchronize() const + { + raft_handle_->sync_stream_pool(); + raft_handle_->sync_stream(); + } + + private: + // Have to store a pointer because handle is not movable + raft::handle_t const* raft_handle_; +}; +#else +struct handle_t { + auto get_next_usable_stream() const { return cuda_stream{}; } + auto get_stream_pool_size() const { return std::size_t{}; } + auto get_usable_stream_count() const { return std::max(get_stream_pool_size(), std::size_t{1}); } + void synchronize() const {} +}; +#endif +} // namespace nvforest diff --git a/cpp/include/nvforest/treelite_importer.hpp b/cpp/include/nvforest/treelite_importer.hpp index b504403f..5402b880 100644 --- a/cpp/include/nvforest/treelite_importer.hpp +++ b/cpp/include/nvforest/treelite_importer.hpp @@ -18,6 +18,7 @@ #include #include +#include #include #include diff --git a/cpp/src/detail/device_buffer.cu b/cpp/src/detail/device_buffer.cu deleted file mode 100644 index 87fb5ddd..00000000 --- a/cpp/src/detail/device_buffer.cu +++ /dev/null @@ -1,50 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#include -#include -#include - -#include -#include -#include - -#include -#include - -namespace nvforest::detail { - -struct owning_device_buffer_type_erased_impl { - owning_device_buffer_type_erased_impl(int device_id, std::size_t size, cuda::stream_ref stream) - : buffer_{[&stream, device_id, size]() { - auto device = cuda::device_ref{device_id}; - auto mr = cuda::device_default_memory_pool(device); - return cuda::make_buffer(stream, mr, size, cuda::no_init); - }()} - { - } - - std::byte* get() { return buffer_.data(); } - - cuda::device_buffer buffer_; -}; - -owning_device_buffer_type_erased::owning_device_buffer_type_erased() : impl_{nullptr} {} - -owning_device_buffer_type_erased::owning_device_buffer_type_erased( - device_id device_id, std::size_t size, cuda::stream_ref stream) -{ - auto device_context = device_setter{device_id}; - impl_ = std::make_unique(device_id.value(), size, stream); -} - -owning_device_buffer_type_erased::owning_device_buffer_type_erased( - owning_device_buffer_type_erased&& other) noexcept = default; -owning_device_buffer_type_erased& owning_device_buffer_type_erased::operator=( - owning_device_buffer_type_erased&& other) noexcept = default; -owning_device_buffer_type_erased::~owning_device_buffer_type_erased() = default; - -std::byte* owning_device_buffer_type_erased::get() { return impl_ ? impl_->get() : nullptr; } - -} // namespace nvforest::detail diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 7046f9ba..96f5297e 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -51,6 +51,10 @@ function(ConfigureTest) Threads::Threads ${TREELITE_LIBS} $) + if(NVFOREST_ENABLE_GPU) + target_link_libraries(${_NVFOREST_TEST_NAME} PRIVATE rmm::rmm raft::raft) + endif() + target_compile_options(${_NVFOREST_TEST_NAME} PRIVATE "$<$:${NVFOREST_CXX_FLAGS}>" "$<$:${NVFOREST_CUDA_FLAGS}>") @@ -81,7 +85,6 @@ endfunction() ConfigureTest(NAME HOST_BUFFER_TEST buffer/buffer.cpp) if(NVFOREST_ENABLE_GPU) ConfigureTest(NAME DEVICE_BUFFER_TEST buffer/buffer.cu) - ConfigureTest(NAME INVALID_STREAM_TEST invalid_stream.cu) endif() ConfigureTest(NAME FOREST_TRAVERSAL_TEST forest/traversal_forest.cpp) ConfigureTest(NAME TREELITE_TRAVERSAL_TEST forest/treelite_traversal.cpp) diff --git a/cpp/tests/invalid_stream.cu b/cpp/tests/invalid_stream.cu deleted file mode 100644 index 0d9b028b..00000000 --- a/cpp/tests/invalid_stream.cu +++ /dev/null @@ -1,82 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include - -#include -#include - -namespace nvforest { - -TEST(InvalidStream, wrong_device) -{ - auto const num_devices = cuda::devices.size(); - if (num_devices < 2) { GTEST_SKIP() << "Test requires at least 2 GPU devices"; } - - // Make a stump tree - auto model_builder = - treelite::model_builder::GetModelBuilder(treelite::TypeInfo::kFloat32, - treelite::TypeInfo::kFloat32, - treelite::model_builder::Metadata{ - 1, - treelite::TaskType::kRegressor, - false, - 1, - {1}, - {1, 1}, - }, - treelite::model_builder::TreeAnnotation{1, {0}, {0}}, - treelite::model_builder::PostProcessorFunc{"identity"}, - std::vector{0.0}); - model_builder->StartTree(); - model_builder->StartNode(0); - model_builder->NumericalTest(0, 0.0, true, treelite::Operator::kLT, 1, 2); - model_builder->EndNode(); - model_builder->StartNode(1); - model_builder->LeafScalar(-1.0); - model_builder->EndNode(); - model_builder->StartNode(2); - model_builder->LeafScalar(1.0); - model_builder->EndNode(); - model_builder->EndTree(); - - cuda_stream stream; - { - auto device_context = detail::device_setter{detail::device_id{1}}; - detail::cuda_check(cudaStreamCreate(&stream)); - } - - // Loading a model on device 0 and inferencing on device 1 is an error. - auto tl_model = model_builder->CommitModel(); - auto nvforest_model = import_from_treelite_model( - *tl_model, tree_layout::breadth_first, index_type{}, false, device_type::gpu, 0, stream); - auto input = thrust::device_vector{1.0f}; - auto output = thrust::device_vector(1); - EXPECT_THAT( - [&]() { - nvforest_model.predict(stream, - thrust::raw_pointer_cast(output.data()), - thrust::raw_pointer_cast(input.data()), - 1, - device_type::gpu, - device_type::gpu, - infer_kind::default_kind); - }, - testing::ThrowsMessage( - testing::HasSubstr("Stream on the wrong device. Expected: 0, Actual: 1"))); -} - -} // namespace nvforest diff --git a/cpp/tests/treelite_importer.cpp b/cpp/tests/treelite_importer.cpp index 33403d05..37d3d9c8 100644 --- a/cpp/tests/treelite_importer.cpp +++ b/cpp/tests/treelite_importer.cpp @@ -3,8 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include #include +#include #include #include #include @@ -382,10 +382,16 @@ TEST(TreeliteImporter, DegenerateTree) auto nvforest_model = import_from_treelite_model(*tl_model, tree_layout::breadth_first); ASSERT_FALSE(nvforest_model.has_vector_leaves()); +#ifdef NVFOREST_ENABLE_GPU + auto raft_handle = raft::handle_t{}; + auto handle = nvforest::handle_t{raft_handle}; +#else + auto handle = nvforest::handle_t{}; +#endif auto X = std::vector{0.0}; auto preds = std::vector(1, 0.0); auto expected_preds = std::vector{1.0}; - nvforest_model.predict(nvforest::cuda_stream{}, + nvforest_model.predict(handle, preds.data(), X.data(), 1, @@ -402,10 +408,16 @@ TEST(TreeliteImporter, DegenerateTreeWithVectorLeaf) auto nvforest_model = import_from_treelite_model(*tl_model, tree_layout::breadth_first); ASSERT_TRUE(nvforest_model.has_vector_leaves()); +#ifdef NVFOREST_ENABLE_GPU + auto raft_handle = raft::handle_t{}; + auto handle = nvforest::handle_t{raft_handle}; +#else + auto handle = nvforest::handle_t{}; +#endif auto X = std::vector{0.0}; auto preds = std::vector(2, 0.0); auto expected_preds = std::vector{0.5, 0.5}; - nvforest_model.predict(nvforest::cuda_stream{}, + nvforest_model.predict(handle, preds.data(), X.data(), 1, diff --git a/dependencies.yaml b/dependencies.yaml index 13dd6145..ddcbbe24 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -13,7 +13,6 @@ files: - cuda - cuda_version - depends_on_cuda_python - - depends_on_cuda_core - depends_on_cupy - depends_on_libraft - depends_on_librmm @@ -38,7 +37,6 @@ files: - cuda - cuda_version - depends_on_cuda_python - - depends_on_cuda_core - depends_on_cupy - depends_on_libraft - depends_on_librmm @@ -129,7 +127,6 @@ files: table: project includes: - depends_on_cuda_python - - depends_on_cuda_core - depends_on_cupy - depends_on_libnvforest - depends_on_pylibraft @@ -458,11 +455,6 @@ dependencies: - matrix: packages: - xgboost>=3.3.0; python_version >= '3.12' - depends_on_cuda_core: - common: - - output_types: [conda, requirements, pyproject] - packages: - - cuda-core>=1.0.0,<2.0.0 depends_on_cuda_python: specific: - output_types: [conda, requirements, pyproject] diff --git a/python/nvforest/CMakeLists.txt b/python/nvforest/CMakeLists.txt index ce399e66..5ebd872c 100644 --- a/python/nvforest/CMakeLists.txt +++ b/python/nvforest/CMakeLists.txt @@ -77,9 +77,3 @@ rapids_cython_create_modules( SOURCE_FILES "nvforest/detail/treelite.pyx" LINKED_LIBRARIES "${linked_libraries}" INSTALL_DIR nvforest/detail MODULE_PREFIX nvforest_) - -rapids_cython_create_modules( - CXX - SOURCE_FILES "nvforest/detail/raft_stream.pyx" - LINKED_LIBRARIES "${linked_libraries}" - INSTALL_DIR nvforest/detail MODULE_PREFIX nvforest_) diff --git a/python/nvforest/nvforest/__init__.py b/python/nvforest/nvforest/__init__.py index e5a01302..4fc9b152 100644 --- a/python/nvforest/nvforest/__init__.py +++ b/python/nvforest/nvforest/__init__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # @@ -24,7 +24,7 @@ GPUForestInferenceClassifier, GPUForestInferenceRegressor, ) -from nvforest._typing import StreamLike +from nvforest._handle import Handle from nvforest._version import __git_commit__, __version__ __all__ = [ @@ -32,7 +32,7 @@ "CPUForestInferenceRegressor", "GPUForestInferenceClassifier", "GPUForestInferenceRegressor", - "StreamLike", + "Handle", "load_model", "load_from_sklearn", "load_from_treelite_model", diff --git a/python/nvforest/nvforest/_factory.py b/python/nvforest/nvforest/_factory.py index f5418227..cc0700a8 100644 --- a/python/nvforest/nvforest/_factory.py +++ b/python/nvforest/nvforest/_factory.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # @@ -18,8 +18,6 @@ infer_is_classifier, ) from nvforest._handle import Handle -from nvforest._typing import StreamLike -from nvforest.detail.raft_stream import _handle_deprecated_handle_arg def get_forest_inference_class(device, is_classifier) -> type: @@ -41,7 +39,7 @@ def make_forest_inference_object( treelite_model: treelite.Model, device: str, device_id: Optional[int], - stream: Optional[StreamLike], + handle: Optional[Handle], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -52,7 +50,7 @@ def make_forest_inference_object( kwargs = dict( treelite_model=treelite_model, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -74,7 +72,6 @@ def load_model( align_bytes: Optional[int] = None, precision: Optional[str] = None, device_id: Optional[int] = None, - stream: Optional[StreamLike] = None, handle: Optional[Handle] = None, ) -> ForestInference: """Load a model into nvForest from a serialized model file. @@ -116,15 +113,10 @@ def load_model( device_id : int or None, default=None For GPU execution, the device on which to load and execute this model. For CPU execution, this value is currently ignored. - stream : StreamLike or None, default=None - For GPU execution, the CUDA stream to use during model loading and - inference. If not given, a new stream will be created. For CPU - execution, this value is ignored. handle : nvforest.Handle or None - **Deprecated**. This argument will be removed in 26.12. Please use - the ``stream`` argument instead. For GPU execution, the nvForest handle containing the stream or stream - pool to use during loading and inference. + pool to use during loading and inference. If not given, a new + handle will be constructed. """ model_path = pathlib.Path(model_file) if not model_path.exists(): @@ -161,12 +153,11 @@ def load_model( case _: raise ValueError(f"Unknown model type: {model_type}") - stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) return make_forest_inference_object( treelite_model=tl_model, device=device, device_id=device_id, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -183,7 +174,6 @@ def load_from_sklearn( align_bytes: Optional[int] = None, precision: Optional[str] = None, device_id: Optional[int] = None, - stream: Optional[StreamLike] = None, handle: Optional[Handle] = None, ) -> ForestInference: """Load a Scikit-Learn forest model to nvForest @@ -218,24 +208,18 @@ def load_from_sklearn( device_id : int or None, default=None For GPU execution, the device on which to load and execute this model. For CPU execution, this value is currently ignored. - stream : StreamLike or None, default=None - For GPU execution, the CUDA stream to use during model loading and - inference. If not given, a new stream will be created. For CPU - execution, this value is ignored. handle : nvforest.Handle or None - **Deprecated**. This argument will be removed in 26.12. Please use - the ``stream`` argument instead. For GPU execution, the nvForest handle containing the stream or stream - pool to use during loading and inference. + pool to use during loading and inference. If not given, a new + handle will be constructed. """ tl_model = treelite.sklearn.import_model(skl_model) - stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) return make_forest_inference_object( treelite_model=tl_model, device=device, device_id=device_id, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -252,7 +236,6 @@ def load_from_treelite_model( align_bytes: Optional[int] = None, precision: Optional[str] = None, device_id: Optional[int] = None, - stream: Optional[StreamLike] = None, handle: Optional[Handle] = None, ) -> ForestInference: """Load a Treelite forest model to nvForest @@ -287,22 +270,16 @@ def load_from_treelite_model( device_id : int or None, default=None For GPU execution, the device on which to load and execute this model. For CPU execution, this value is currently ignored. - stream : StreamLike or None, default=None - For GPU execution, the CUDA stream to use during model loading and - inference. If not given, a new stream will be created. For CPU - execution, this value is ignored. handle : nvforest.Handle or None - **Deprecated**. This argument will be removed in 26.12. Please use - the ``stream`` argument instead. For GPU execution, the nvForest handle containing the stream or stream - pool to use during loading and inference. + pool to use during loading and inference. If not given, a new + handle will be constructed. """ - stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) return make_forest_inference_object( treelite_model=tl_model, device=device, device_id=device_id, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, diff --git a/python/nvforest/nvforest/_forest_inference.py b/python/nvforest/nvforest/_forest_inference.py index 88852756..f6bbbd34 100644 --- a/python/nvforest/nvforest/_forest_inference.py +++ b/python/nvforest/nvforest/_forest_inference.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # @@ -8,9 +8,6 @@ from time import perf_counter from typing import Optional -from nvforest._handle import Handle -from nvforest.detail.raft_stream import _handle_deprecated_handle_arg - try: from typing import Self except ImportError: @@ -21,7 +18,8 @@ from cuda.bindings import runtime from nvforest._base import ForestInferenceClassifier, ForestInferenceRegressor -from nvforest._typing import DataType, StreamLike +from nvforest._handle import Handle +from nvforest._typing import DataType from nvforest.detail.forest_inference import ForestInferenceImpl @@ -125,7 +123,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[StreamLike], + handle: Optional[Handle], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -202,7 +200,7 @@ def _optimize( else: test_instances[layout] = type(self)._create_with_layout( treelite_model_bytes=self.forest.treelite_model_bytes, - stream=self.forest.stream, + handle=self.forest.handle, layout=layout, default_chunk_size=None, align_bytes=self.forest.align_bytes, @@ -245,7 +243,7 @@ def _optimize( # Return a new instance with optimal settings return type(self)._create_with_layout( treelite_model_bytes=self.forest.treelite_model_bytes, - stream=self.forest.stream, + handle=self.forest.handle, layout=optimal_layout, default_chunk_size=optimal_chunk_size, align_bytes=self.forest.align_bytes, @@ -262,21 +260,19 @@ def __init__( self, *, treelite_model: treelite.Model, - stream: Optional[StreamLike] = None, + handle: Optional[Handle] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, precision: Optional[str] = None, - handle: Optional[Handle] = None, ): if not infer_is_classifier(treelite_model): raise ValueError("treelite_model must be a classifier.") - stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) self.forest = ForestInferenceImpl( treelite_model=treelite_model, device="cpu", device_id=-1, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -288,7 +284,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[StreamLike], + handle: Optional[Handle], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -300,7 +296,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -401,21 +397,19 @@ def __init__( self, *, treelite_model: treelite.Model, - stream: Optional[StreamLike] = None, + handle: Optional[Handle] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, precision: Optional[str] = None, - handle: Optional[Handle] = None, ): if infer_is_classifier(treelite_model): raise ValueError("treelite_model must be a regressor.") - stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) self.forest = ForestInferenceImpl( treelite_model=treelite_model, device="cpu", device_id=-1, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -427,7 +421,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[StreamLike], + handle: Optional[Handle], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -439,7 +433,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -532,22 +526,20 @@ def __init__( self, *, treelite_model: treelite.Model, - stream: Optional[StreamLike] = None, + handle: Optional[Handle] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, precision: Optional[str] = None, - handle: Optional[Handle] = None, device_id: int, ): if not infer_is_classifier(treelite_model): raise ValueError("treelite_model must be a classifier.") - stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) self.forest = ForestInferenceImpl( treelite_model=treelite_model, device="gpu", device_id=device_id, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -559,7 +551,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[StreamLike], + handle: Optional[Handle], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -571,7 +563,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -673,22 +665,20 @@ def __init__( self, *, treelite_model: treelite.Model, - stream: Optional[StreamLike] = None, + handle: Optional[Handle] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, precision: Optional[str] = None, - handle: Optional[Handle] = None, device_id: int, ): if infer_is_classifier(treelite_model): raise ValueError("treelite_model must be a regressor.") - stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) self.forest = ForestInferenceImpl( treelite_model=treelite_model, device="gpu", device_id=device_id, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -700,7 +690,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[StreamLike], + handle: Optional[Handle], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -712,7 +702,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - stream=stream, + handle=handle, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, diff --git a/python/nvforest/nvforest/_typing.py b/python/nvforest/nvforest/_typing.py index 5f897fbd..548a05b7 100644 --- a/python/nvforest/nvforest/_typing.py +++ b/python/nvforest/nvforest/_typing.py @@ -1,9 +1,9 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # -from typing import TYPE_CHECKING, Protocol, Union, runtime_checkable +from typing import TYPE_CHECKING, Union import numpy as np @@ -11,10 +11,3 @@ import cupy DataType = Union[np.ndarray, "cupy.ndarray"] - - -@runtime_checkable -class StreamLike(Protocol): - """Duck typing for all stream-like objects""" - - def __cuda_stream__(self) -> tuple[int, int]: ... diff --git a/python/nvforest/nvforest/detail/cuda_stream.pxd b/python/nvforest/nvforest/detail/cuda_stream.pxd index 0fea6cab..690c0c4c 100644 --- a/python/nvforest/nvforest/detail/cuda_stream.pxd +++ b/python/nvforest/nvforest/detail/cuda_stream.pxd @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -from cuda.bindings.cyruntime cimport cudaStream_t - - cdef extern from "nvforest/cuda_stream.hpp" namespace "nvforest" nogil: - ctypedef cudaStream_t cuda_stream + cdef cppclass cuda_stream: + pass diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index 8d55a7a7..3d6eb3b4 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -7,17 +7,19 @@ from typing import Optional, Union import numpy as np import treelite -from cuda.core import Device, Stream -from nvforest._typing import DataType, StreamLike +from nvforest._handle import Handle +from nvforest._typing import DataType from nvforest.detail.treelite import safe_treelite_call from libc.stdint cimport uint32_t, uintptr_t from libcpp cimport bool from libcpp.optional cimport nullopt, optional +from pylibraft.common.handle cimport handle_t as raft_handle_t from nvforest.detail.cuda_stream cimport cuda_stream as nvforest_stream_t from nvforest.detail.device_type cimport device_type as nvforest_device_t +from nvforest.detail.handle cimport handle_t as nvforest_handle_t from nvforest.detail.infer_kind cimport infer_kind from nvforest.detail.postprocessing cimport element_op, row_op from nvforest.detail.tree_layout cimport tree_layout as nvforest_tree_layout @@ -31,7 +33,7 @@ from nvforest.detail.treelite cimport ( cdef extern from "nvforest/forest_model.hpp" namespace "nvforest" nogil: cdef cppclass forest_model: void predict[io_t]( - nvforest_stream_t, + const nvforest_handle_t&, io_t*, io_t*, size_t, @@ -63,13 +65,13 @@ cdef extern from "nvforest/treelite_importer.hpp" namespace "nvforest" nogil: cdef class ForestInference_impl(): cdef forest_model model - cdef nvforest_stream_t stream_handle - cdef object stream + cdef nvforest_handle_t nvforest_handle + cdef object raft_handle cdef object device def __cinit__( self, - stream: object, + raft_handle: object, tl_model_bytes: Union[bytes, bytearray], *, layout: str = "depth_first", @@ -78,21 +80,12 @@ cdef class ForestInference_impl(): device: str = "cpu", device_id: Optional[int] = None, ): - # Assumption: The caller needs to pass in correct (device, device_id) pair - # This function will not contain any logic for auto-detecting device. - # Assumption: The caller should convert a user-provided stream-like object - # into cuda.core.Stream before constructing ForestInference_impl() - cdef uintptr_t stream_ptr = 0 - if stream is not None: - # Use assertion here, since the failure here indicates a bug, not - # a user error. - assert isinstance(stream, Stream), ( - "stream must be converted to cuda.core.Stream before " - "building ForestInference_impl()" - ) - stream_ptr = int(stream.handle) - self.stream = stream - self.stream_handle = stream_ptr + # Store reference to RAFT handle to control lifetime, since + # nvforest_handle keeps a pointer to it + self.raft_handle = raft_handle + self.nvforest_handle = nvforest_handle_t( + self.raft_handle.getHandle() + ) cdef optional[bool] use_double_precision_c cdef bool use_double_precision_bool @@ -137,7 +130,7 @@ cdef class ForestInference_impl(): use_double_precision_c, dev_type, device_id, - self.stream_handle + self.nvforest_handle.get_next_usable_stream() ) safe_treelite_call( @@ -247,7 +240,7 @@ cdef class ForestInference_impl(): if model_dtype == np.float32: self.model.predict[float]( - self.stream_handle, + self.nvforest_handle, out_ptr, in_ptr, n_rows, @@ -258,7 +251,7 @@ cdef class ForestInference_impl(): ) else: self.model.predict[double]( - self.stream_handle, + self.nvforest_handle, out_ptr, in_ptr, n_rows, @@ -269,7 +262,7 @@ cdef class ForestInference_impl(): ) if self.device == "gpu": - self.stream.sync() + self.nvforest_handle.synchronize() return preds @@ -280,7 +273,7 @@ class ForestInferenceImpl: treelite_model: treelite.Model, device: str, device_id: int, - stream: Optional[StreamLike] = None, + handle: Optional[Handle] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -288,27 +281,7 @@ class ForestInferenceImpl: ): # Assumption: The caller needs to pass in correct (device, device_id) pair # This function will not contain any logic for auto-detecting device. - if stream is not None and not isinstance(stream, StreamLike): - raise TypeError("stream must be a stream-like object or None") - if device == "gpu": - previous_device = Device() - try: - cuda_device = Device(device_id) - cuda_device.set_current() - if stream is None: - stream = cuda_device.create_stream() - else: - stream = cuda_device.create_stream(stream) - finally: - previous_device.set_current() - else: - assert device == "cpu" - if device == "gpu" and stream.device.device_id != device_id: - raise ValueError( - f"stream is associated with device {stream.device.device_id}, " - f"but device_id is {device_id}" - ) - self.stream = stream + self.handle = Handle() if handle is None else handle self._layout = layout self.precision = precision self.default_chunk_size = default_chunk_size @@ -336,7 +309,7 @@ class ForestInferenceImpl: self._treelite_model_bytes = treelite_model.serialize_bytes() self.impl = ForestInference_impl( - self.stream, + self.handle, self._treelite_model_bytes, layout=self._layout, align_bytes=self.align_bytes, diff --git a/python/nvforest/nvforest/detail/handle.pxd b/python/nvforest/nvforest/detail/handle.pxd new file mode 100644 index 00000000..b80ad6b0 --- /dev/null +++ b/python/nvforest/nvforest/detail/handle.pxd @@ -0,0 +1,17 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +from pylibraft.common.handle cimport handle_t as raft_handle_t + +from nvforest.detail.cuda_stream cimport cuda_stream as nvforest_stream_t + + +cdef extern from "nvforest/handle.hpp" namespace "nvforest" nogil: + cdef cppclass handle_t: + handle_t() except + + handle_t(const raft_handle_t* handle_ptr) except + + handle_t(const raft_handle_t& handle) except + + nvforest_stream_t get_next_usable_stream() except + + void synchronize() except+ diff --git a/python/nvforest/nvforest/detail/raft_stream.pyx b/python/nvforest/nvforest/detail/raft_stream.pyx deleted file mode 100644 index e47be556..00000000 --- a/python/nvforest/nvforest/detail/raft_stream.pyx +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import warnings -from typing import Optional - -from cuda.core import Stream -from pylibraft.common.handle import Handle as RaftHandle - -from nvforest._typing import StreamLike - - -def _get_stream_from_raft_handle(handle: RaftHandle) -> Stream: - cdef int stream_ptr = handle.c_obj.get_stream() - return Stream.from_handle(stream_ptr) - - -def _handle_deprecated_handle_arg( - *, - handle: Optional[RaftHandle] = None, - stream: Optional[StreamLike] = None, -) -> Optional[StreamLike]: - if handle is None: - return stream - - if stream is not None: - raise ValueError("Cannot set `stream` and `handle` at the same time.") - warnings.warn( - "`handle` parameter is deprecated and will be removed in 26.12. " - "Please use `stream` instead", - FutureWarning, - stacklevel=2, - ) - stream = _get_stream_from_raft_handle(handle) - return stream diff --git a/python/nvforest/pyproject.toml b/python/nvforest/pyproject.toml index 6a6bc434..437fb9d4 100644 --- a/python/nvforest/pyproject.toml +++ b/python/nvforest/pyproject.toml @@ -28,7 +28,6 @@ license = "Apache-2.0" requires-python = ">=3.11" dependencies = [ "cuda-bindings>=13.0.1,<14.0", - "cuda-core>=1.0.0,<2.0.0", "cupy-cuda13x[ctk]>=14.0.1,!=14.1.0", "libnvforest==26.10.*,>=0.0.0a0", "numpy>=2.0,<3.0a0", diff --git a/python/nvforest/tests/test_nvforest.py b/python/nvforest/tests/test_nvforest.py index 1d5bb4e2..a7282ae4 100644 --- a/python/nvforest/tests/test_nvforest.py +++ b/python/nvforest/tests/test_nvforest.py @@ -927,37 +927,3 @@ def test_incorrect_data_shape(input_size, predict_func): with pytest.raises(ValueError, match=f"Expected {n_features} features"): X_test = np.zeros((1, input_size)) _ = predict_func(fm, X_test) - - -@pytest.mark.skipif( - cp.cuda.runtime.getDeviceCount() < 2, reason="Requires at least 2 GPUs" -) -def test_incompatible_stream(tmp_path): - random_state = np.random.RandomState(43210) - - X, y = _simulate_data( - 1000, - 2, - random_state=random_state, - classification=False, - bias=10.0, - ) - - model_path = tmp_path / "xgb_class.ubj" - _ = _build_and_save_xgboost( - model_path, - X, - y, - num_rounds=10, - classification=False, - ) - with cp.cuda.Device(1): - stream = cp.cuda.Stream(non_blocking=True) - assert stream.device_id == 1 - with pytest.raises( - ValueError, - match=r".*stream is associated with device 1, but device_id is 0.*", - ): - _ = nvforest.load_model( - model_path, device="gpu", device_id=0, stream=stream - ) From fee870ff425ef04f15d00ba74b461a9fb6e09974 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 8 Sep 2026 08:03:56 -0700 Subject: [PATCH 2/2] Formatting fixes --- cpp/cmake/thirdparty/get_rmm.cmake | 2 +- python/nvforest/nvforest/__init__.py | 2 +- python/nvforest/nvforest/_factory.py | 2 +- python/nvforest/nvforest/_forest_inference.py | 2 +- python/nvforest/nvforest/_typing.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/cpp/cmake/thirdparty/get_rmm.cmake b/cpp/cmake/thirdparty/get_rmm.cmake index b476b09a..519b9905 100644 --- a/cpp/cmake/thirdparty/get_rmm.cmake +++ b/cpp/cmake/thirdparty/get_rmm.cmake @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= diff --git a/python/nvforest/nvforest/__init__.py b/python/nvforest/nvforest/__init__.py index 4fc9b152..f92070eb 100644 --- a/python/nvforest/nvforest/__init__.py +++ b/python/nvforest/nvforest/__init__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/python/nvforest/nvforest/_factory.py b/python/nvforest/nvforest/_factory.py index cc0700a8..96b44546 100644 --- a/python/nvforest/nvforest/_factory.py +++ b/python/nvforest/nvforest/_factory.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/python/nvforest/nvforest/_forest_inference.py b/python/nvforest/nvforest/_forest_inference.py index f6bbbd34..c892d50b 100644 --- a/python/nvforest/nvforest/_forest_inference.py +++ b/python/nvforest/nvforest/_forest_inference.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/python/nvforest/nvforest/_typing.py b/python/nvforest/nvforest/_typing.py index 548a05b7..41c89160 100644 --- a/python/nvforest/nvforest/_typing.py +++ b/python/nvforest/nvforest/_typing.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 #