diff --git a/README.md b/README.md index 98d0ff04..f9320614 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,30 @@ fm = nvforest.load_model("/path/to/xgboost_model.ubj", device="gpu", pred = fm.predict(X) ``` +By default, nvForest creates a new +[CUDA stream](https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/asynchronous-execution.html#cuda-streams) +for the model and retains it. To use a custom stream, pass in a :external+cuda-python:py:class:`cuda.core.Stream` or +other stream-like objects (*): + +(*) A stream-like object is an object that exposes a method named `__cuda_stream__`. See [Stream Protocol]( +https://nvidia.github.io/cuda-python/cuda-core/latest/interoperability.html#cuda-stream-protocol) for more details. + +```python +from cuda.core import Device + +device = Device(0) +device.set_current() +stream = device.create_stream() + +fm = nvforest.load_model( + "/path/to/xgboost_model.ubj", + device="gpu", + device_id=0, + stream=stream, +) +pred = fm.predict(X) # Note: automatically syncs the stream +``` + Load a scikit-learn random forest model and get class probabilities: ```python diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 9707e20b..e4216f5c 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -11,6 +11,7 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=12.9.3,<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 12324e3d..6b5d28d2 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -11,6 +11,7 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=12.9.3,<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 82ddd30d..ac8b604a 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -11,6 +11,7 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=13.0.2,<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 68406e9d..0491de1d 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -11,6 +11,7 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=13.0.2,<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 fe178168..0450e9e1 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -208,7 +208,6 @@ 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) @@ -261,7 +260,8 @@ if(NVFOREST_ENABLE_GPU) src/infer8.cu src/infer9.cu src/infer10.cu - src/infer11.cu) + src/infer11.cu + src/detail/device_buffer.cu) endif() target_sources( ${NVFOREST_CPP_TARGET} @@ -321,7 +321,7 @@ elseif(NVFOREST_EXPORT_TREELITE_LINKAGE) endif() if(NVFOREST_ENABLE_GPU) - list(APPEND _nvforest_cpp_public_libs rmm::rmm raft::raft CUDA::cudart_static) + list(APPEND _nvforest_cpp_public_libs raft::raft CUDA::cudart_static CCCL::CCCL) endif() # These are always private: diff --git a/cpp/cmake/thirdparty/get_rmm.cmake b/cpp/cmake/thirdparty/get_rmm.cmake deleted file mode 100644 index 519b9905..00000000 --- a/cpp/cmake/thirdparty/get_rmm.cmake +++ /dev/null @@ -1,13 +0,0 @@ -# ============================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# 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/README.md b/cpp/include/nvforest/README.md index 74b3cbd6..8464ab9e 100644 --- a/cpp/include/nvforest/README.md +++ b/cpp/include/nvforest/README.md @@ -94,25 +94,24 @@ auto output = static_cast(nullptr); // Loaded as single cudaMalloc((void**)&output, num_rows * num_outputs * sizeof(float)); // Assuming that input is a float* pointing to data already located on-device - -auto handle = nvforest::handle_t{}; - nvforest_model.predict( - handle, + stream, output, input, num_rows, nvforest::device_type::gpu, // out_mem_type nvforest::device_type::gpu, // in_mem_type + nvforest::infer_kind::default_kind, 4 // chunk_size ); ``` -**handle**: To provide a unified interface on CPU and GPU, we introduce -`nvforest::handle_t` as a wrapper for `raft::handle_t`. This is currently just a -placeholder in CPU-only builds, and using it does not require any CUDA -functionality. For GPU-enabled builds, you can construct a -`nvforest::handle_t` directly from the `raft::handle_t` you wish to use. +**stream**: The CUDA stream on which inference and any required data copies +will run. It must be associated with the same GPU device as the model. Reusing +the stream passed during model import preserves the ordering between model +initialization and inference. Synchronize the stream before consuming the +output on the host or releasing resources used by the operation. Keep the stream +alive until the model is destroyed and all queued work has completed. **output**: Pointer to pre-allocated buffer where results should be written. If the model has been loaded at single precision, this should be a @@ -145,4 +144,4 @@ will be selected based on heuristics. ## Learning More While the above usage summary should be enough to get started using nvForest in another C++ application, you can learn more about the details of this -implementation by reading TODO(wphicks). +implementation in [Implementation.md](Implementation.md). diff --git a/cpp/include/nvforest/detail/device_id/gpu.hpp b/cpp/include/nvforest/detail/device_id/gpu.hpp index 49edfdd6..10d10119 100644 --- a/cpp/include/nvforest/detail/device_id/gpu.hpp +++ b/cpp/include/nvforest/detail/device_id/gpu.hpp @@ -7,23 +7,20 @@ #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 = value_type{}; + auto raw_id = int{}; cuda_check(cudaGetDevice(&raw_id)); return raw_id; }()} {}; - device_id(value_type dev_id) noexcept : id_{dev_id} {}; + device_id(int dev_id) noexcept : id_{dev_id} {}; - auto value() const noexcept { return id_.value(); } + auto value() const noexcept { return id_; } private: - rmm::cuda_device_id id_; + int 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 d69a0ff5..07ba9a34 100644 --- a/cpp/include/nvforest/detail/device_setter/gpu.hpp +++ b/cpp/include/nvforest/detail/device_setter/gpu.hpp @@ -8,10 +8,22 @@ #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) + namespace nvforest::detail { /** Struct for setting current device within a code block */ @@ -27,7 +39,7 @@ struct device_setter { cuda_check(cudaSetDevice(device.value())); } - ~device_setter() { RAFT_CUDA_TRY_NO_THROW(cudaSetDevice(prev_device_.value())); } + ~device_setter() { NVFOREST_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 744f48cf..4249cff1 100644 --- a/cpp/include/nvforest/detail/infer/gpu.cuh +++ b/cpp/include/nvforest/detail/infer/gpu.cuh @@ -21,6 +21,8 @@ #include #include +#include + #include #include #include @@ -41,6 +43,18 @@ 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) +{ + if (stream != cuda::stream_ref{cudaStream_t{cudaStreamDefault}} && + stream != cuda::stream_ref{cudaStreamLegacy} && + stream != cuda::stream_ref{cudaStreamPerThread} && 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 * @@ -99,6 +113,8 @@ 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 360dfdee..2e217d64 100644 --- a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp +++ b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp @@ -3,38 +3,51 @@ * 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() : data_{} {} - + owning_buffer() = default; owning_buffer(device_id device_id, std::size_t size, - 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), cuda::stream_ref{stream}}; - }()} + cuda_stream stream) noexcept(false) + : data_{device_id, size * sizeof(value_type), cuda::stream_ref{stream}} { } - auto* get() const { return reinterpret_cast(data_.data()); } + auto* get() const { return reinterpret_cast(data_.get()); } private: - mutable rmm::device_buffer data_; + mutable owning_device_buffer_type_erased data_; }; } // namespace nvforest::detail diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index 143881f4..346a48e4 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -6,13 +6,11 @@ #include #include #include -#include #include #include #include #include #include -#include #include #ifdef NVFOREST_ENABLE_GPU @@ -20,6 +18,8 @@ #endif #include +#include +#include #include #include #include @@ -46,11 +46,14 @@ struct forest_model { decision_forest_); } - /** The number of outputs per row generated by the model */ - auto num_outputs() + /** The number of outputs per row generated by the model for the given type of inference */ + auto num_outputs(infer_kind inference_kind = infer_kind::default_kind) { - return std::visit([](auto&& concrete_forest) { return concrete_forest.num_outputs(); }, - decision_forest_); + return std::visit( + [inference_kind](auto&& concrete_forest) { + return concrete_forest.num_outputs(inference_kind); + }, + decision_forest_); } /** The number of trees in the model */ @@ -118,12 +121,12 @@ struct forest_model { /** * Perform inference on given input * - * @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[out] output The buffer where model output should be stored. + * This must be of size at least ROWS x num_outputs(predict_type). + * @param[in] input The buffer containing input data. * @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. @@ -139,19 +142,46 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(buffer& output, + void predict(cuda_stream stream, + 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); + 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); + } else { + 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(), input.device_index()}; + } + auto buffer_out = buffer{}; + if (output.memory_type() != memory_type()) { + buffer_out = buffer{ + row_count * num_outputs(predict_type), memory_type(), device_index(), stream}; + } else { + buffer_out = buffer{output.data(), + row_count * num_outputs(predict_type), + memory_type(), + output.device_index()}; + } + 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); + } + } } else { throw type_error("Input type does not match model_type"); } @@ -160,14 +190,17 @@ struct forest_model { } /** - * Perform inference on given input + * Perform inference on given input, using multiple CUDA streams. + * A stream pool lets you overlap memory transfers with partitioned inference. * - * @param[in] handle The nvforest::handle_t (wrapper for raft::handle_t - * on GPU) which will be used to provide streams for evaluation. + * @param[in] stream_pool_next A callback function that returns a CUDA stream + * from the pool. + * @param[in] stream_pool_size The number of CUDA streams in the pool. This + * must be greater than zero. * @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(). + * to this output location. This must be of size at least ROWS x num_outputs(predict_type). * @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 @@ -187,36 +220,46 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(handle_t const& handle, + void predict(std::function const& stream_pool_next, + std::size_t stream_pool_size, buffer& output, buffer const& input, infer_kind predict_type = infer_kind::default_kind, std::optional specified_chunk_size = std::nullopt) { + if (stream_pool_size == 0) { + throw std::invalid_argument("stream_pool_size must be greater than zero"); + } 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; + [this, + predict_type, + &stream_pool_next, + stream_pool_size, + &output, + &input, + &specified_chunk_size](auto&& concrete_forest) { + using model_io_t = 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, handle.get_next_usable_stream(), predict_type, specified_chunk_size); + output, input, stream_pool_next(), 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 partition_size = - std::max(detail::ceildiv(row_count, handle.get_usable_stream_count()), + std::max(detail::ceildiv(row_count, stream_pool_size), 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 stream = stream_pool_next(); 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()}; + partition_in = buffer{ + rows_in_this_partition * num_features(), memory_type(), device_index(), stream}; copy(partition_in, input, 0, @@ -226,22 +269,28 @@ struct forest_model { } else { partition_in = buffer{input.data() + i * partition_size * num_features(), rows_in_this_partition * num_features(), - memory_type()}; + memory_type(), + input.device_index()}; } auto partition_out = buffer{}; if (output.memory_type() != memory_type()) { - partition_out = buffer{rows_in_this_partition * num_outputs(), memory_type()}; + partition_out = buffer{rows_in_this_partition * num_outputs(predict_type), + memory_type(), + device_index(), + stream}; } else { - partition_out = buffer{output.data() + i * partition_size * num_outputs(), - rows_in_this_partition * num_outputs(), - memory_type()}; + partition_out = + buffer{output.data() + i * partition_size * num_outputs(predict_type), + rows_in_this_partition * num_outputs(predict_type), + memory_type(), + output.device_index()}; } 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(), + i * partition_size * num_outputs(predict_type), 0, partition_out.size(), stream); @@ -258,8 +307,65 @@ struct forest_model { /** * 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[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[out] output Pointer to the memory location where output should end + * up + * @param[in] input Pointer to the input data + * @param[in] num_rows Number of rows in input + * @param[in] out_mem_type The memory type (device/host) of the output + * buffer + * @param[in] in_mem_type The memory type (device/host) of the input buffer + * @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(cuda_stream stream, + io_t* output, + io_t* input, + std::size_t num_rows, + device_type out_mem_type, + device_type in_mem_type, + infer_kind predict_type = infer_kind::default_kind, + std::optional specified_chunk_size = std::nullopt) + { + int current_device_id; + if (out_mem_type == device_type::gpu || in_mem_type == device_type::gpu) { +#ifdef NVFOREST_ENABLE_GPU + detail::cuda_check(cudaGetDevice(¤t_device_id)); +#else + throw detail::gpu_unsupported("Tried to use GPU memory in CPU-only build"); +#endif + } else { + current_device_id = -1; + } + auto out_buffer = + buffer{output, num_rows * num_outputs(predict_type), 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); + } + + /** + * Perform inference on given input, using multiple CUDA streams. + * A stream pool lets you overlap memory transfers with partitioned inference. + * + * @param[in] stream_pool_next A callback function that returns a CUDA stream + * from the pool. + * @param[in] stream_pool_size The number of CUDA streams in the pool. This + * must be greater than zero. * @param[out] output Pointer to the memory location where output should end * up * @param[in] input Pointer to the input data @@ -282,7 +388,8 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(handle_t const& handle, + void predict(std::function const& stream_pool_next, + std::size_t stream_pool_size, io_t* output, io_t* input, std::size_t num_rows, @@ -301,9 +408,15 @@ struct forest_model { } else { current_device_id = -1; } - 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(handle, out_buffer, in_buffer, predict_type, specified_chunk_size); + auto out_buffer = + buffer{output, num_rows * num_outputs(predict_type), out_mem_type, current_device_id}; + auto in_buffer = buffer{input, num_rows * num_features(), in_mem_type, current_device_id}; + predict(stream_pool_next, + stream_pool_size, + out_buffer, + in_buffer, + predict_type, + specified_chunk_size); } private: diff --git a/cpp/include/nvforest/handle.hpp b/cpp/include/nvforest/handle.hpp deleted file mode 100644 index 34e7d794..00000000 --- a/cpp/include/nvforest/handle.hpp +++ /dev/null @@ -1,43 +0,0 @@ -/* - * 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().get()}; - } - 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 6d45017e..c8b242f0 100644 --- a/cpp/include/nvforest/treelite_importer.hpp +++ b/cpp/include/nvforest/treelite_importer.hpp @@ -18,7 +18,6 @@ #include #include -#include #include #include diff --git a/cpp/src/detail/device_buffer.cu b/cpp/src/detail/device_buffer.cu new file mode 100644 index 00000000..87fb5ddd --- /dev/null +++ b/cpp/src/detail/device_buffer.cu @@ -0,0 +1,50 @@ +/* + * 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 96f5297e..3eea830e 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -51,10 +51,6 @@ 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}>") @@ -85,6 +81,8 @@ endfunction() ConfigureTest(NAME HOST_BUFFER_TEST buffer/buffer.cpp) if(NVFOREST_ENABLE_GPU) ConfigureTest(NAME DEVICE_BUFFER_TEST buffer/buffer.cu) + ConfigureTest(NAME INFERENCE_TEST inference.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/inference.cu b/cpp/tests/inference.cu new file mode 100644 index 00000000..a849c960 --- /dev/null +++ b/cpp/tests/inference.cu @@ -0,0 +1,239 @@ +/* + * 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 +#include + +#include +#include +#include +#include +#include + +namespace nvforest { + +std::unique_ptr get_tree_model() +{ + 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{2, {0, 0}, {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(); + + 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(); + + return model_builder->CommitModel(); +} + +std::vector repeat(std::vector const& x, std::size_t n_repeat) +{ + std::vector ret(x.size() * n_repeat); + for (std::size_t i = 0; i < n_repeat; ++i) { + for (std::size_t j = 0; j < x.size(); j++) { + ret[i * x.size() + j] = x[j]; + } + } + return ret; +} + +// Exercise every input/output memory combination on GPU 0 and, when available, GPU 1. +class GPUModelInferenceTest + : public testing::TestWithParam> {}; + +TEST_P(GPUModelInferenceTest, Predict) +{ + auto const [device_id, in_mem_type, out_mem_type] = GetParam(); + int device_count; + detail::cuda_check(cudaGetDeviceCount(&device_count)); + if (device_id >= device_count) { GTEST_SKIP() << "Requires GPU " << device_id; } + auto device_context = detail::device_setter{detail::device_id{device_id}}; + + constexpr std::size_t n_streams = 4; + std::vector stream_pool; + for (std::size_t i = 0; i < n_streams; ++i) { + stream_pool.emplace_back(cuda::device_ref{device_id}); + } + auto tree_model = get_tree_model(); + auto nvforest_model = import_from_treelite_model(*tree_model, + tree_layout::breadth_first, + index_type{}, + false, + device_type::gpu, + device_id, + stream_pool[0].get()); + // Model initialization must finish before inference can use the other streams. + stream_pool[0].sync(); + + // Larger than four minimum-size partitions, with an uneven final partition. + constexpr std::size_t n_repeat = 16385; + auto host_input = repeat(std::vector{-1.0f, 1.0f}, n_repeat); + auto input_view = buffer{host_input.data(), host_input.size()}; + auto input = buffer{input_view, in_mem_type, device_id, stream_pool[0].get()}; + stream_pool[0].sync(); + auto const cases = std::array{ + std::pair{infer_kind::default_kind, std::vector{-2.0f, 2.0f}}, + std::pair{infer_kind::per_tree, std::vector{-1.0f, -1.0f, 1.0f, 1.0f}}, + std::pair{infer_kind::leaf_id, std::vector{1.0f, 1.0f, 2.0f, 2.0f}}, + }; + std::size_t stream_calls = 0; + auto stream_pool_next = [&]() { + return cuda_stream{stream_pool[stream_calls++ % n_streams].get()}; + }; + for (auto const& [infer_type, row_pair] : cases) { + SCOPED_TRACE(static_cast(infer_type)); + auto expected = repeat(row_pair, n_repeat); + for (bool use_pool : {false, true}) { + SCOPED_TRACE(use_pool); + for (bool use_buffers : {false, true}) { + SCOPED_TRACE(use_buffers); + auto output = buffer{expected.size(), out_mem_type, device_id, stream_pool[0].get()}; + stream_pool[0].sync(); + auto const previous_stream_calls = stream_calls; + if (use_pool) { + if (use_buffers) { + nvforest_model.predict(stream_pool_next, n_streams, output, input, infer_type); + } else { + nvforest_model.predict(stream_pool_next, + n_streams, + output.data(), + input.data(), + host_input.size(), + out_mem_type, + in_mem_type, + infer_type); + } + auto const cross_memory = + in_mem_type != device_type::gpu || out_mem_type != device_type::gpu; + EXPECT_EQ(stream_calls - previous_stream_calls, cross_memory ? n_streams : 1); + } else if (use_buffers) { + nvforest_model.predict(stream_pool[0].get(), output, input, infer_type); + } else { + nvforest_model.predict(stream_pool[0].get(), + output.data(), + input.data(), + host_input.size(), + out_mem_type, + in_mem_type, + infer_type); + } + for (auto& stream : stream_pool) { + stream.sync(); + } + auto actual = buffer{output, device_type::cpu, 0, stream_pool[0].get()}; + stream_pool[0].sync(); + EXPECT_EQ(std::vector(actual.data(), actual.data() + actual.size()), expected); + } + } + } +} + +INSTANTIATE_TEST_SUITE_P(MemoryLocations, + GPUModelInferenceTest, + testing::Combine(testing::Values(0, 1), + testing::Values(device_type::cpu, device_type::gpu), + testing::Values(device_type::cpu, device_type::gpu))); + +TEST(Inference, EmptyStreamPool) +{ + auto tree_model = get_tree_model(); + auto model = import_from_treelite_model(*tree_model, tree_layout::breadth_first); + auto input = buffer{1}; + auto output = buffer{1}; + auto next_stream = []() -> cuda_stream { + ADD_FAILURE() << "An empty stream pool must be rejected before requesting a stream"; + return cuda_stream{}; + }; + EXPECT_THROW(model.predict(next_stream, 0, output, input), std::invalid_argument); + EXPECT_THROW( + model.predict( + next_stream, 0, output.data(), input.data(), 1, device_type::cpu, device_type::cpu), + std::invalid_argument); +} + +// TODO(chyunsu3): Remove once we remove the RAFT dependency in 27.02. +TEST(Inference, GPUModelWithRAFTHandle) +{ + auto tree_model = get_tree_model(); + auto nvforest_model = import_from_treelite_model( + *tree_model, tree_layout::breadth_first, index_type{}, false, device_type::gpu); + + std::size_t n_repeat = 256; + auto input = thrust::device_vector(repeat(std::vector{-1.0f, 1.0f}, n_repeat)); + auto const cases = std::array{ + std::pair{infer_kind::default_kind, + thrust::device_vector(repeat(std::vector{-2.0f, 2.0f}, n_repeat))}, + std::pair{ + infer_kind::per_tree, + thrust::device_vector(repeat(std::vector{-1.0f, -1.0f, 1.0f, 1.0f}, n_repeat))}, + std::pair{ + infer_kind::leaf_id, + thrust::device_vector(repeat(std::vector{1.0f, 1.0f, 2.0f, 2.0f}, n_repeat))}, + }; + + std::size_t n_streams = 4; + raft::handle_t handle{cuda::stream_ref{cudaStreamPerThread}, + std::make_shared(n_streams)}; + for (auto const& [infer_type, expected] : cases) { + SCOPED_TRACE(static_cast(infer_type)); + auto output = thrust::device_vector(expected.size()); + nvforest_model.predict([&] { return cuda_stream{handle.get_next_usable_stream().get()}; }, + handle.get_stream_pool_size(), + thrust::raw_pointer_cast(output.data()), + thrust::raw_pointer_cast(input.data()), + input.size(), + device_type::gpu, + device_type::gpu, + infer_type); + handle.sync_stream_pool(); + EXPECT_EQ(output, expected); + } +} + +} // namespace nvforest diff --git a/cpp/tests/invalid_stream.cu b/cpp/tests/invalid_stream.cu new file mode 100644 index 00000000..2777472e --- /dev/null +++ b/cpp/tests/invalid_stream.cu @@ -0,0 +1,87 @@ +/* + * 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 +#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(); + + auto device_context = detail::device_setter{detail::device_id{0}}; + cuda::stream import_stream{cuda::device_ref{0}}; + cuda::stream wrong_device_stream{cuda::device_ref{1}}; + + // Import on a valid stream before testing prediction with a wrong-device stream. + 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, + import_stream.get()); + import_stream.sync(); + auto input = thrust::device_vector{1.0f}; + auto output = thrust::device_vector(1); + EXPECT_THAT( + [&]() { + nvforest_model.predict(wrong_device_stream.get(), + 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 9c4c1452..664ce810 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 @@ -385,16 +385,10 @@ 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(handle, + nvforest_model.predict(nvforest::cuda_stream{}, preds.data(), X.data(), 1, @@ -411,16 +405,10 @@ 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(handle, + nvforest_model.predict(nvforest::cuda_stream{}, preds.data(), X.data(), 1, @@ -514,12 +502,6 @@ TEST(TreeliteImporter, SinglePrecisionDoubleThresholds) { auto constexpr eps = double{std::numeric_limits::epsilon()}; auto constexpr inf = std::numeric_limits::infinity(); -#ifdef NVFOREST_ENABLE_GPU - auto raft_handle = raft::handle_t{}; - auto handle = nvforest::handle_t{raft_handle}; -#else - auto handle = nvforest::handle_t{}; -#endif for (auto threshold : {1.0 + 1.5 * eps, 1.0 + 0.25 * eps, -1e300, 1e300}) { auto tl_model = make_double_threshold_forest(threshold); auto nearest = static_cast(threshold); @@ -544,7 +526,7 @@ TEST(TreeliteImporter, SinglePrecisionDoubleThresholds) auto nvforest_model = import_from_treelite_model(*tl_model, layout, index_type{}, false); ASSERT_FALSE(nvforest_model.is_double_precision()); auto preds = std::vector(X.size(), 0.0f); - nvforest_model.predict(handle, + nvforest_model.predict(nvforest::cuda_stream{}, preds.data(), X.data(), X.size(), diff --git a/dependencies.yaml b/dependencies.yaml index fdb2643a..5bdbfeda 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -13,6 +13,7 @@ files: - cuda - cuda_version - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - depends_on_libraft - depends_on_librmm @@ -36,6 +37,7 @@ files: - cuda - cuda_version - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - depends_on_libraft - depends_on_librmm @@ -131,6 +133,7 @@ files: table: project includes: - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - depends_on_libnvforest - depends_on_pylibraft @@ -433,6 +436,11 @@ dependencies: - pytest-cov - pytest-xdist - *scikit_learn + 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/docs/source/conf.py b/docs/source/conf.py index afaf139a..f0aedd0c 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -205,6 +205,10 @@ "python": ("https://docs.python.org/3", None), "sklearn": ("https://scikit-learn.org/stable/", None), "treelite": ("https://treelite.readthedocs.io/en/latest/", None), + "cuda-python": ( + "https://nvidia.github.io/cuda-python/cuda-core/", + None, + ), } # Config numpydoc diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index f29a191d..be738f85 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -129,6 +129,33 @@ You can inspect the type of the model by printing its type: Each model object is associated with a single device. Use the ``device_id`` property to look up which device the model object is located on. +CUDA stream selection +--------------------- + +By default, nvForest creates a new +`CUDA stream `_ +for a GPU model. To use a custom stream, pass in a :external+cuda-python:py:class:`cuda.core.Stream` (or +any stream-like object [#]_) via the ``stream`` parameter. A supplied stream must be associated with the same device as ``device_id``. + +.. code-block:: python + + from cuda.core import Device + + device = Device(0) + device.set_current() + stream = device.create_stream() + + fm = nvforest.load_from_sklearn( + skl_model, device="gpu", device_id=0, stream=stream + ) + predictions = fm.predict(X) + +nvForest retains a reference to a caller-supplied stream. Python prediction methods synchronize +the model's stream before returning, so their results are ready for immediate use. + +.. [#] A stream-like object is an object that exposes a method named `__cuda_stream__`. + See `Stream Protocol `_ for more details. + Running inference ----------------- @@ -223,29 +250,45 @@ Once the tree model is available as a Treelite object, pass it to the #include #include #include + #include #include + constexpr int device_id = 0; + cudaSetDevice(device_id); + cudaStream_t stream{}; + cudaStreamCreate(&stream); + auto fm = nvforest::import_from_treelite_model( *treelite_model, nvforest::preferred_tree_layout, nvforest::index_type{}, std::nullopt, - nvforest::device_type::gpu); + nvforest::device_type::gpu, + device_id, + stream); Now that the tree model is fully imported into nvForest, let's run inference: .. code-block:: cpp - #include - #include - - raft::handle_t raft_handle{}; - nvforest::handle_t handle{raft_handle}; - // Assumption: // * Both output and input are in the GPU memory. // * The input buffer should be of dimension (num_rows, num_features) // * The output buffer should be of dimension (num_rows, fm.num_outputs()) - fm.predict(handle, output, input, num_rows, + fm.predict(stream, output, input, num_rows, nvforest::device_type::gpu, nvforest::device_type::gpu, nvforest::infer_kind::default_kind); + + // Wait before reading or freeing output. + cudaStreamSynchronize(stream); + + // Use output + // ... + + // Clean up by destroying stream + // Also free input, output buffers + +Use a stream associated with the same device as the model. Re-use the stream for import and +inference to preserve the ordering between model initialization and prediction without a global +synchronization. Keep the stream alive until the model is destroyed and all queued +work has completed. diff --git a/python/nvforest/CMakeLists.txt b/python/nvforest/CMakeLists.txt index 258aa558..46bed124 100644 --- a/python/nvforest/CMakeLists.txt +++ b/python/nvforest/CMakeLists.txt @@ -80,3 +80,9 @@ 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 f92070eb..1bd8d2f9 100644 --- a/python/nvforest/nvforest/__init__.py +++ b/python/nvforest/nvforest/__init__.py @@ -25,6 +25,7 @@ GPUForestInferenceRegressor, ) from nvforest._handle import Handle +from nvforest._typing import StreamLike from nvforest._version import __git_commit__, __version__ __all__ = [ @@ -33,6 +34,7 @@ "GPUForestInferenceClassifier", "GPUForestInferenceRegressor", "Handle", + "StreamLike", "load_model", "load_from_sklearn", "load_from_treelite_model", diff --git a/python/nvforest/nvforest/_factory.py b/python/nvforest/nvforest/_factory.py index 870317be..46ae1eea 100644 --- a/python/nvforest/nvforest/_factory.py +++ b/python/nvforest/nvforest/_factory.py @@ -18,6 +18,8 @@ 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: @@ -39,7 +41,7 @@ def make_forest_inference_object( treelite_model: treelite.Model, device: str, device_id: Optional[int], - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -50,7 +52,7 @@ def make_forest_inference_object( kwargs = dict( treelite_model=treelite_model, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -72,6 +74,7 @@ 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. @@ -113,10 +116,15 @@ 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 27.02. Please use + the ``stream`` argument instead. For GPU execution, the nvForest handle containing the stream or stream - pool to use during loading and inference. If not given, a new - handle will be constructed. + pool to use during loading and inference. """ model_path = pathlib.Path(model_file) if not model_path.exists(): @@ -153,11 +161,12 @@ 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, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -174,6 +183,7 @@ 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 @@ -208,18 +218,24 @@ 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 27.02. Please use + the ``stream`` argument instead. For GPU execution, the nvForest handle containing the stream or stream - pool to use during loading and inference. If not given, a new - handle will be constructed. + pool to use during loading and inference. """ 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, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -236,6 +252,7 @@ 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 @@ -270,16 +287,22 @@ 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 27.02. Please use + the ``stream`` argument instead. For GPU execution, the nvForest handle containing the stream or stream - pool to use during loading and inference. If not given, a new - handle will be constructed. + pool to use during loading and inference. """ + stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) return make_forest_inference_object( treelite_model=tl_model, device=device, device_id=device_id, - handle=handle, + stream=stream, 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 c892d50b..88852756 100644 --- a/python/nvforest/nvforest/_forest_inference.py +++ b/python/nvforest/nvforest/_forest_inference.py @@ -8,6 +8,9 @@ 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: @@ -18,8 +21,7 @@ from cuda.bindings import runtime from nvforest._base import ForestInferenceClassifier, ForestInferenceRegressor -from nvforest._handle import Handle -from nvforest._typing import DataType +from nvforest._typing import DataType, StreamLike from nvforest.detail.forest_inference import ForestInferenceImpl @@ -123,7 +125,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -200,7 +202,7 @@ def _optimize( else: test_instances[layout] = type(self)._create_with_layout( treelite_model_bytes=self.forest.treelite_model_bytes, - handle=self.forest.handle, + stream=self.forest.stream, layout=layout, default_chunk_size=None, align_bytes=self.forest.align_bytes, @@ -243,7 +245,7 @@ def _optimize( # Return a new instance with optimal settings return type(self)._create_with_layout( treelite_model_bytes=self.forest.treelite_model_bytes, - handle=self.forest.handle, + stream=self.forest.stream, layout=optimal_layout, default_chunk_size=optimal_chunk_size, align_bytes=self.forest.align_bytes, @@ -260,19 +262,21 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = 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, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -284,7 +288,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -296,7 +300,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -397,19 +401,21 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = 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, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -421,7 +427,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -433,7 +439,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -526,20 +532,22 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = 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, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -551,7 +559,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -563,7 +571,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -665,20 +673,22 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = 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, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -690,7 +700,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -702,7 +712,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - handle=handle, + stream=stream, 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 41c89160..5f897fbd 100644 --- a/python/nvforest/nvforest/_typing.py +++ b/python/nvforest/nvforest/_typing.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 # -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Protocol, Union, runtime_checkable import numpy as np @@ -11,3 +11,10 @@ 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 690c0c4c..0fea6cab 100644 --- a/python/nvforest/nvforest/detail/cuda_stream.pxd +++ b/python/nvforest/nvforest/detail/cuda_stream.pxd @@ -2,6 +2,8 @@ # 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: - cdef cppclass cuda_stream: - pass + ctypedef cudaStream_t cuda_stream diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index 3d6eb3b4..295722eb 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -7,19 +7,17 @@ from typing import Optional, Union import numpy as np import treelite +from cuda.core import Device, Stream -from nvforest._handle import Handle -from nvforest._typing import DataType +from nvforest._typing import DataType, StreamLike 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 @@ -33,7 +31,7 @@ from nvforest.detail.treelite cimport ( cdef extern from "nvforest/forest_model.hpp" namespace "nvforest" nogil: cdef cppclass forest_model: void predict[io_t]( - const nvforest_handle_t&, + nvforest_stream_t, io_t*, io_t*, size_t, @@ -65,13 +63,13 @@ cdef extern from "nvforest/treelite_importer.hpp" namespace "nvforest" nogil: cdef class ForestInference_impl(): cdef forest_model model - cdef nvforest_handle_t nvforest_handle - cdef object raft_handle + cdef nvforest_stream_t stream_handle + cdef object stream cdef object device def __cinit__( self, - raft_handle: object, + stream: object, tl_model_bytes: Union[bytes, bytearray], *, layout: str = "depth_first", @@ -80,12 +78,21 @@ cdef class ForestInference_impl(): device: str = "cpu", device_id: Optional[int] = None, ): - # 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() - ) + # 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 cdef optional[bool] use_double_precision_c cdef bool use_double_precision_bool @@ -130,7 +137,7 @@ cdef class ForestInference_impl(): use_double_precision_c, dev_type, device_id, - self.nvforest_handle.get_next_usable_stream() + self.stream_handle ) safe_treelite_call( @@ -240,7 +247,7 @@ cdef class ForestInference_impl(): if model_dtype == np.float32: self.model.predict[float]( - self.nvforest_handle, + self.stream_handle, out_ptr, in_ptr, n_rows, @@ -251,7 +258,7 @@ cdef class ForestInference_impl(): ) else: self.model.predict[double]( - self.nvforest_handle, + self.stream_handle, out_ptr, in_ptr, n_rows, @@ -262,7 +269,7 @@ cdef class ForestInference_impl(): ) if self.device == "gpu": - self.nvforest_handle.synchronize() + self.stream.sync() return preds @@ -273,7 +280,7 @@ class ForestInferenceImpl: treelite_model: treelite.Model, device: str, device_id: int, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -281,7 +288,28 @@ 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. - self.handle = Handle() if handle is None else handle + 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" + stream = None + 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._layout = layout self.precision = precision self.default_chunk_size = default_chunk_size @@ -309,7 +337,7 @@ class ForestInferenceImpl: self._treelite_model_bytes = treelite_model.serialize_bytes() self.impl = ForestInference_impl( - self.handle, + self.stream, 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 deleted file mode 100644 index b80ad6b0..00000000 --- a/python/nvforest/nvforest/detail/handle.pxd +++ /dev/null @@ -1,17 +0,0 @@ -# -# 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 new file mode 100644 index 00000000..f3cc504b --- /dev/null +++ b/python/nvforest/nvforest/detail/raft_stream.pyx @@ -0,0 +1,40 @@ +# 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 + +from cuda.bindings.cyruntime cimport cudaStream_t +from libc.stdint cimport uintptr_t +from pylibraft.common.handle cimport handle_t + + +def _get_stream_from_raft_handle(handle: RaftHandle) -> Stream: + cdef handle_t* handle_ = handle.getHandle() + cdef cudaStream_t stream = handle_.get_stream().get() + return Stream.from_handle(stream) + + +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 27.02. " + "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 c5c7bb4f..24a25143 100644 --- a/python/nvforest/pyproject.toml +++ b/python/nvforest/pyproject.toml @@ -28,6 +28,7 @@ license = "Apache-2.0" requires-python = ">=3.12" dependencies = [ "cuda-bindings>=13.0.2,<14.0", + "cuda-core>=1.0.0,<2.0.0", "cupy-cuda13x[ctk]>=14.0.1,!=14.1.0", "libnvforest==26.12.*,>=0.0.0a0", "numpy>=2.0,<3.0a0", diff --git a/python/nvforest/tests/test_compat.py b/python/nvforest/tests/test_compat.py new file mode 100644 index 00000000..6063e216 --- /dev/null +++ b/python/nvforest/tests/test_compat.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Compatibility test to ensure that legacy APIs are working +""" + +import cupy as cp +import numpy as np +import pytest +import treelite +from cuda.core import Device + +import nvforest + + +def test_deprecated_handle(): + """Test use of RAFT handle with non-default stream""" + builder = treelite.model_builder.ModelBuilder( + threshold_type="float64", + leaf_output_type="float64", + metadata=treelite.model_builder.Metadata( + num_feature=1, + task_type="kRegressor", + average_tree_output=False, + num_target=1, + num_class=[1], + leaf_vector_shape=(1, 1), + ), + tree_annotation=treelite.model_builder.TreeAnnotation( + num_tree=1, target_id=[0], class_id=[0] + ), + postprocessor=treelite.model_builder.PostProcessorFunc( + name="identity" + ), + base_scores=[0.0], + ) + builder.start_tree() + builder.start_node(0) + builder.numerical_test( + feature_id=0, + threshold=0.0, + default_left=False, + opname="<", + left_child_key=1, + right_child_key=2, + ) + builder.end_node() + builder.start_node(1) + builder.leaf(-1.0) + builder.end_node() + builder.start_node(2) + builder.leaf(1.0) + builder.end_node() + builder.end_tree() + + treelite_model = builder.commit() + + device = Device() + device.set_current() + stream = device.create_stream() + handle = nvforest.Handle(stream=stream) + + rng = np.random.default_rng(42) + X = rng.standard_normal(size=(1000, 1), dtype="float64") + + with pytest.warns( + FutureWarning, + match=".*`handle` parameter is deprecated and will be removed in 27.02.*", + ): + fm = nvforest.load_from_treelite_model( + treelite_model, device="gpu", handle=handle + ) + pred_with_handle = cp.asnumpy(fm.predict(X)) + + stream2 = cp.cuda.Stream(non_blocking=True) + fm2 = nvforest.load_from_treelite_model( + treelite_model, device="gpu", stream=stream2 + ) + pred_with_stream = cp.asnumpy(fm2.predict(X)) + np.testing.assert_almost_equal(pred_with_handle, pred_with_stream) diff --git a/python/nvforest/tests/test_nvforest.py b/python/nvforest/tests/test_nvforest.py index 97088f62..1da75763 100644 --- a/python/nvforest/tests/test_nvforest.py +++ b/python/nvforest/tests/test_nvforest.py @@ -977,3 +977,37 @@ 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 + )