From a3f9d1805410092a5f9b0e26a5b9947133e99477 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 4 Aug 2026 01:22:07 -0700 Subject: [PATCH 1/8] Add missing headers from #457 --- src/nvforest_predict.h | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/nvforest_predict.h b/src/nvforest_predict.h index 76ddf1b..2602ca3 100644 --- a/src/nvforest_predict.h +++ b/src/nvforest_predict.h @@ -19,6 +19,11 @@ #include #include +#ifdef TRITON_ENABLE_GPU +#include +#endif +#include + #include #include #include From 5c35012349c8a0778beecd8279a3ba911516f43b Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 6 Aug 2026 15:40:49 -0700 Subject: [PATCH 2/8] Patch nvForest --- cmake/thirdparty/get_nvforest.cmake | 3 +- cmake/thirdparty/patches/nvforest.patch | 40 +++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) create mode 100644 cmake/thirdparty/patches/nvforest.patch diff --git a/cmake/thirdparty/get_nvforest.cmake b/cmake/thirdparty/get_nvforest.cmake index ae70842..92f75fb 100644 --- a/cmake/thirdparty/get_nvforest.cmake +++ b/cmake/thirdparty/get_nvforest.cmake @@ -28,6 +28,7 @@ function(find_and_configure_nvforest) GIT_REPOSITORY https://github.com/${PKG_FORK}/nvforest.git GIT_TAG ${PKG_PINNED_TAG} SOURCE_SUBDIR cpp + PATCHES "./patches/nvforest.patch" OPTIONS "BUILD_NVFOREST_TESTS OFF" "BUILD_SHARED_LIBS ON" @@ -42,7 +43,7 @@ endfunction() find_and_configure_nvforest(VERSION 26.08 FORK rapidsai - PINNED_TAG dd49f1a906654f59f20023849b2f1cae820f601f + PINNED_TAG 948c3526b37cea774697a02407705d7345105099 USE_TREELITE_STATIC ${TRITON_FIL_USE_TREELITE_STATIC} ENABLE_GPU ${TRITON_ENABLE_GPU} ) diff --git a/cmake/thirdparty/patches/nvforest.patch b/cmake/thirdparty/patches/nvforest.patch new file mode 100644 index 0000000..2d9885d --- /dev/null +++ b/cmake/thirdparty/patches/nvforest.patch @@ -0,0 +1,40 @@ +From 60c8df5adef69291d4978b4631d58b2ab9ac6429 Mon Sep 17 00:00:00 2001 +From: Hyunsu Cho +Date: Tue, 4 Aug 2026 03:06:43 -0700 +Subject: [PATCH] Check raft_handle_ for nullptr + +--- + cpp/include/nvforest/handle.hpp | 12 +++++++++--- + 1 file changed, 9 insertions(+), 3 deletions(-) + +diff --git a/cpp/include/nvforest/handle.hpp b/cpp/include/nvforest/handle.hpp +index fd1f3bb..42066c8 100644 +--- a/cpp/include/nvforest/handle.hpp ++++ b/cpp/include/nvforest/handle.hpp +@@ -18,14 +18,20 @@ struct handle_t { + handle_t(raft::handle_t const& raft_handle) : raft_handle_{&raft_handle} {} + auto get_next_usable_stream() const + { ++ if (!raft_handle_) { return cuda_stream{}; } + 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_stream_pool_size() const { ++ if (!raft_handle_) { return std::size_t{}; } ++ 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(); ++ if (raft_handle_) { ++ raft_handle_->sync_stream_pool(); ++ raft_handle_->sync_stream(); ++ } + } + + private: +-- +2.43.0 + From b299d4d8002db65c16707f3a2c99c96232402a5c Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 6 Aug 2026 15:56:27 -0700 Subject: [PATCH 3/8] stream_synchronize() should be a no-op for host mem --- src/rapids_triton/memory/buffer.hpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/rapids_triton/memory/buffer.hpp b/src/rapids_triton/memory/buffer.hpp index d5b5a98..87af7c9 100644 --- a/src/rapids_triton/memory/buffer.hpp +++ b/src/rapids_triton/memory/buffer.hpp @@ -178,7 +178,9 @@ struct Buffer { void stream_synchronize() const { if constexpr (IS_GPU_BUILD) { - cuda_check(cudaStreamSynchronize(stream_)); + if (mem_type() == DeviceMemory) { + cuda_check(cudaStreamSynchronize(stream_)); + } } } From 2a66c0becc874557626b944a448c78dfd5247866 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 6 Aug 2026 16:13:29 -0700 Subject: [PATCH 4/8] Provide version bumps for 26.08 --- CMakeLists.txt | 2 +- cmake/thirdparty/get_cuml.cmake | 2 +- cmake/thirdparty/get_nvforest.cmake | 2 +- conda/environments/triton_benchmark.yml | 2 +- conda/environments/triton_test.yml | 4 ++-- conda/environments/triton_test_no_client.yml | 4 ++-- ops/Dockerfile | 4 +--- 7 files changed, 9 insertions(+), 11 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 692cbb2..3048362 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -80,7 +80,7 @@ else() # - Prepare rapids-cmake ----------------------------------------------------- set(rapids-cmake-version 26.08) file(DOWNLOAD - https://raw.githubusercontent.com/rapidsai/rapids-cmake/691ccac03770627311b271e9dc5820f6ea2b6346/RAPIDS.cmake + https://raw.githubusercontent.com/rapidsai/rapids-cmake/refs/tags/v26.08.00/RAPIDS.cmake ${CMAKE_BINARY_DIR}/RAPIDS.cmake) include(${CMAKE_BINARY_DIR}/RAPIDS.cmake) include(rapids-cmake) diff --git a/cmake/thirdparty/get_cuml.cmake b/cmake/thirdparty/get_cuml.cmake index dd70e93..c6097af 100644 --- a/cmake/thirdparty/get_cuml.cmake +++ b/cmake/thirdparty/get_cuml.cmake @@ -55,6 +55,6 @@ endfunction() # CPM_raft_SOURCE=/path/to/local/raft find_and_configure_cuml(VERSION 26.08 FORK rapidsai - PINNED_TAG 91d972563e5f7994c3fa57d1cd7d5d6e4caefd4f + PINNED_TAG 265b9da6a0e75dbef071a3168398b993a5ff6f0e USE_TREELITE_STATIC ${TRITON_FIL_USE_TREELITE_STATIC} ) diff --git a/cmake/thirdparty/get_nvforest.cmake b/cmake/thirdparty/get_nvforest.cmake index 92f75fb..d58c944 100644 --- a/cmake/thirdparty/get_nvforest.cmake +++ b/cmake/thirdparty/get_nvforest.cmake @@ -43,7 +43,7 @@ endfunction() find_and_configure_nvforest(VERSION 26.08 FORK rapidsai - PINNED_TAG 948c3526b37cea774697a02407705d7345105099 + PINNED_TAG cef3a50da0f74b0015876b9d6d424c86141898dc USE_TREELITE_STATIC ${TRITON_FIL_USE_TREELITE_STATIC} ENABLE_GPU ${TRITON_ENABLE_GPU} ) diff --git a/conda/environments/triton_benchmark.yml b/conda/environments/triton_benchmark.yml index e110c71..59e9bcd 100644 --- a/conda/environments/triton_benchmark.yml +++ b/conda/environments/triton_benchmark.yml @@ -5,7 +5,7 @@ channels: - rapidsai dependencies: - cuda-version=12.8 - - cudf=26.06 + - cudf=26.08 - libcusolver - libcusparse - matplotlib diff --git a/conda/environments/triton_test.yml b/conda/environments/triton_test.yml index af71c85..b1f93d2 100644 --- a/conda/environments/triton_test.yml +++ b/conda/environments/triton_test.yml @@ -7,8 +7,8 @@ dependencies: - aws-sdk-cpp - clang-tools=19.1.7 - cuda-version=12.8 - - cudf=26.06 - - cuml=26.06 + - cudf=26.08 + - cuml=26.08 - flake8 - hypothesis - lightgbm diff --git a/conda/environments/triton_test_no_client.yml b/conda/environments/triton_test_no_client.yml index 0883531..79cf498 100644 --- a/conda/environments/triton_test_no_client.yml +++ b/conda/environments/triton_test_no_client.yml @@ -7,8 +7,8 @@ dependencies: - aws-sdk-cpp - clang-tools=19.1.7 - cuda-version=12.8 - - cudf=26.06 - - cuml=26.06 + - cudf=26.08 + - cuml=26.08 - flake8 - hypothesis - lightgbm diff --git a/ops/Dockerfile b/ops/Dockerfile index 6af05a4..1012b40 100644 --- a/ops/Dockerfile +++ b/ops/Dockerfile @@ -3,7 +3,7 @@ # Arguments for controlling build details ########################################################################################### # Version of Triton to use -ARG TRITON_VERSION=26.06 +ARG TRITON_VERSION=26.07 # Base container image ARG BASE_IMAGE=nvcr.io/nvidia/tritonserver:${TRITON_VERSION}-py3 # Whether or not to enable GPU build @@ -64,8 +64,6 @@ RUN conda run --no-capture-output -n triton_test \ && rm -r /sdk_install FROM wheel-install-${USE_CLIENT_WHEEL} as conda-test -RUN conda run --no-capture-output -n triton_test \ - pip install git+https://github.com/rapidsai/rapids-triton.git@v26.06.00#subdirectory=python RUN conda-pack --ignore-missing-files -n triton_test -o /tmp/env.tar \ && mkdir /conda/test/ \ && cd /conda/test/ \ From 7f7578de55a747571d25d64a4bbbb0ebb9ef59db Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 6 Aug 2026 17:06:42 -0700 Subject: [PATCH 5/8] Bring in Python wrapper for rapids_triton --- ops/Dockerfile | 15 +- rapids_triton_python/pyproject.toml | 3 + .../rapids_triton/__init__.py | 16 ++ rapids_triton_python/rapids_triton/client.py | 263 ++++++++++++++++++ .../rapids_triton/exceptions.py | 18 ++ rapids_triton_python/rapids_triton/logging.py | 18 ++ rapids_triton_python/rapids_triton/testing.py | 118 ++++++++ .../rapids_triton/triton/__init__.py | 13 + .../rapids_triton/triton/client.py | 36 +++ .../rapids_triton/triton/dtype.py | 35 +++ .../rapids_triton/triton/io.py | 145 ++++++++++ .../rapids_triton/triton/message.py | 30 ++ .../rapids_triton/triton/response.py | 35 +++ .../rapids_triton/utils/__init__.py | 0 .../rapids_triton/utils/safe_import.py | 39 +++ 15 files changed, 778 insertions(+), 6 deletions(-) create mode 100644 rapids_triton_python/pyproject.toml create mode 100644 rapids_triton_python/rapids_triton/__init__.py create mode 100644 rapids_triton_python/rapids_triton/client.py create mode 100644 rapids_triton_python/rapids_triton/exceptions.py create mode 100644 rapids_triton_python/rapids_triton/logging.py create mode 100644 rapids_triton_python/rapids_triton/testing.py create mode 100644 rapids_triton_python/rapids_triton/triton/__init__.py create mode 100644 rapids_triton_python/rapids_triton/triton/client.py create mode 100644 rapids_triton_python/rapids_triton/triton/dtype.py create mode 100644 rapids_triton_python/rapids_triton/triton/io.py create mode 100644 rapids_triton_python/rapids_triton/triton/message.py create mode 100644 rapids_triton_python/rapids_triton/triton/response.py create mode 100644 rapids_triton_python/rapids_triton/utils/__init__.py create mode 100644 rapids_triton_python/rapids_triton/utils/safe_import.py diff --git a/ops/Dockerfile b/ops/Dockerfile index 1012b40..5c20636 100644 --- a/ops/Dockerfile +++ b/ops/Dockerfile @@ -64,6 +64,9 @@ RUN conda run --no-capture-output -n triton_test \ && rm -r /sdk_install FROM wheel-install-${USE_CLIENT_WHEEL} as conda-test +COPY ./rapids_triton_python /workspace/rapids_triton_python +RUN conda run --no-capture-output -n triton_test \ + pip install /workspace/rapids_triton_python RUN conda-pack --ignore-missing-files -n triton_test -o /tmp/env.tar \ && mkdir /conda/test/ \ && cd /conda/test/ \ @@ -98,16 +101,16 @@ RUN apt-get update \ # Stage immediately before building; useful for build iteration FROM base as build-prep -RUN mkdir -p /rapids_triton/build /rapids_triton/install +RUN mkdir -p /workspace/build /workspace/install -COPY ./src /rapids_triton/src -COPY ./CMakeLists.txt /rapids_triton -COPY ./cmake /rapids_triton/cmake +COPY ./src /workspace/src +COPY ./CMakeLists.txt /workspace +COPY ./cmake /workspace/cmake ARG BACKEND_NAME=fil ENV BACKEND_NAME=$BACKEND_NAME -WORKDIR /rapids_triton/build +WORKDIR /workspace/build # Remove potentially stale build artifacts RUN if [ -d /opt/tritonserver/backends/${BACKEND_NAME} ]; \ @@ -170,7 +173,7 @@ RUN source /conda/dev/bin/activate \ -DTRITON_ENABLE_STATS="${TRITON_ENABLE_STATS}" \ -DRAPIDS_DEPENDENCIES_VERSION="${RAPIDS_DEPENDENCIES_VERSION}" \ -DTRITON_FIL_USE_TREELITE_STATIC="${TRITON_FIL_USE_TREELITE_STATIC}" \ - -DCMAKE_INSTALL_PREFIX=/rapids_triton/install \ + -DCMAKE_INSTALL_PREFIX=/workspace/install \ ..; ENV CCACHE_DIR=/ccache diff --git a/rapids_triton_python/pyproject.toml b/rapids_triton_python/pyproject.toml new file mode 100644 index 0000000..9787c3b --- /dev/null +++ b/rapids_triton_python/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/rapids_triton_python/rapids_triton/__init__.py b/rapids_triton_python/rapids_triton/__init__.py new file mode 100644 index 0000000..5c99549 --- /dev/null +++ b/rapids_triton_python/rapids_triton/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from rapids_triton.client import Client +from rapids_triton.logging import logger diff --git a/rapids_triton_python/rapids_triton/client.py b/rapids_triton_python/rapids_triton/client.py new file mode 100644 index 0000000..a4b2360 --- /dev/null +++ b/rapids_triton_python/rapids_triton/client.py @@ -0,0 +1,263 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import concurrent.futures +import time +from collections import namedtuple + +from rapids_triton.triton.client import get_triton_client +from rapids_triton.triton.dtype import dtype_to_triton_name +from rapids_triton.triton.io import ( + create_triton_input, + create_triton_output, + destroy_shared_memory_region, +) +from rapids_triton.triton.response import get_response_data +from tritonclient import utils as triton_utils + +# TODO(wphicks): Propagate device ids for cuda shared memory + +MultiModelOutput = namedtuple("MultiModelOutput", ("name", "version", "output")) + + +class Client(object): + def __init__(self, protocol="grpc", host="localhost", port=None, concurrency=4): + self.triton_client = get_triton_client( + protocol=protocol, host=host, port=port, concurrency=concurrency + ) + self._protocol = protocol + + @property + def protocol(self): + return self._protocol + + def create_inputs(self, array_inputs, shared_mem=None): + return [ + create_triton_input( + self.triton_client, + arr, + name, + dtype_to_triton_name(arr.dtype), + protocol=self.protocol, + shared_mem=shared_mem, + ) + for name, arr in array_inputs.items() + ] + + def create_outputs(self, output_sizes, shared_mem=None): + return { + name: create_triton_output( + self.triton_client, + size, + name, + protocol=self.protocol, + shared_mem=shared_mem, + ) + for name, size in output_sizes.items() + } + + def wait_for_server(self, timeout): + server_wait_start = time.time() + while True: + try: + if self.triton_client.is_server_ready(): + break + except triton_utils.InferenceServerException: + pass + if time.time() - server_wait_start > timeout: + raise RuntimeError("Server startup timeout expired") + time.sleep(1) + + def clear_shared_memory(self): + self.triton_client.unregister_cuda_shared_memory() + self.triton_client.unregister_system_shared_memory() + + def release_io(self, io_objs): + for io_ in io_objs: + if io_.name is not None: + self.triton_client.unregister_cuda_shared_memory(name=io_.name) + destroy_shared_memory_region(io_.handle, shared_mem="cuda") + + def get_model_config(self, model_name): + return self.triton_client.get_model_config(model_name).config + + def predict( + self, + model_name, + input_data, + output_sizes, + model_version="1", + shared_mem=None, + attempts=1, + ): + model_version = str(model_version) + + try: + inputs = self.create_inputs(input_data, shared_mem=shared_mem) + outputs = self.create_outputs(output_sizes, shared_mem=shared_mem) + + response = self.triton_client.infer( + model_name, + model_version=model_version, + inputs=[input_.input for input_ in inputs], + outputs=[output_.output for output_ in outputs.values()], + ) + result = { + name: get_response_data(response, handle, name) + for name, (_, handle, _) in outputs.items() + } + self.release_io(inputs) + self.release_io(outputs.values()) + + except triton_utils.InferenceServerException: + if attempts > 1: + return self.predict( + model_name, + input_data, + output_sizes, + model_version=model_version, + shared_mem=shared_mem, + attempts=attempts - 1, + ) + raise + return result + + def predict_async( + self, + model_name, + input_data, + output_sizes, + model_version="1", + shared_mem=None, + attempts=1, + ): + model_version = str(model_version) + + inputs = self.create_inputs(input_data, shared_mem=shared_mem) + outputs = self.create_outputs(output_sizes, shared_mem=shared_mem) + + future_result = concurrent.futures.Future() + + def callback(result, error): + if error is None: + output_arrays = { + name: get_response_data(result, handle, name) + for name, (_, handle, _) in outputs.items() + } + + future_result.set_result(output_arrays) + + self.release_io(outputs.values()) + else: + if isinstance(error, triton_utils.InferenceServerException): + if attempts > 1: + future_result.set_result( + self.predict( + model_name, + input_data, + output_sizes, + model_version=model_version, + shared_mem=shared_mem, + attempts=attempts - 1, + ) + ) + future_result.set_exception(error) + + self.triton_client.async_infer( + model_name, + model_version=model_version, + inputs=[input_.input for input_ in inputs], + outputs=[output_.output for output_ in outputs.values()], + callback=callback, + ) + + if shared_mem is not None: + + def release_callback(fut): + self.release_io(inputs) + + future_result.add_done_callback(release_callback) + return future_result + + def predict_multimodel_async( + self, + model_names, + input_data, + output_sizes, + model_versions=("1",), + shared_mem=None, + executor=None, + attempts=1, + ): + all_models = [ + (name, str(version)) for name in model_names for version in model_versions + ] + + inputs = self.create_inputs(input_data, shared_mem=shared_mem) + + all_future_results = [] + for model_name, version in all_models: + outputs = self.create_outputs(output_sizes, shared_mem=shared_mem) + + def create_callback(future_result, outputs): + def callback(result, error): + if error is None: + output_arrays = { + name: get_response_data(result, handle, name) + for name, (_, handle, _) in outputs.items() + } + + future_result.set_result( + MultiModelOutput( + name=model_name, version=version, output=output_arrays + ) + ) + + self.release_io(outputs.values()) + else: + if isinstance(error, triton_utils.InferenceServerException): + if attempts > 1: + future_result.set_result( + self.predict( + model_name, + input_data, + output_sizes, + model_version=version, + shared_mem=shared_mem, + attempts=attempts - 1, + ) + ) + future_result.set_exception(error) + + return callback + + all_future_results.append(concurrent.futures.Future()) + self.triton_client.async_infer( + model_name, + model_version=version, + inputs=[input_.input for input_ in inputs], + outputs=[output_.output for output_ in outputs.values()], + callback=create_callback(all_future_results[-1], outputs), + ) + + def wait_for_all(future_results, releasable_inputs): + concurrent.futures.wait(future_results) + self.release_io(releasable_inputs) + return [fut.result() for fut in future_results] + + if executor is None: + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(wait_for_all, all_future_results, inputs) + else: + return executor.submit(wait_for_all, all_future_results, inputs) diff --git a/rapids_triton_python/rapids_triton/exceptions.py b/rapids_triton_python/rapids_triton/exceptions.py new file mode 100644 index 0000000..dc7f8c5 --- /dev/null +++ b/rapids_triton_python/rapids_triton/exceptions.py @@ -0,0 +1,18 @@ +# Copyright (c) 2022, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +class IncompatibleSharedMemory(Exception): + """Error thrown if operation cannot be completed with given shared memory + type""" diff --git a/rapids_triton_python/rapids_triton/logging.py b/rapids_triton_python/rapids_triton/logging.py new file mode 100644 index 0000000..2a94de8 --- /dev/null +++ b/rapids_triton_python/rapids_triton/logging.py @@ -0,0 +1,18 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging + +logger = logging.getLogger("rapids_triton") +logger.setLevel(logging.INFO) diff --git a/rapids_triton_python/rapids_triton/testing.py b/rapids_triton_python/rapids_triton/testing.py new file mode 100644 index 0000000..58aa354 --- /dev/null +++ b/rapids_triton_python/rapids_triton/testing.py @@ -0,0 +1,118 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import numpy as np +from rapids_triton.client import Client +from rapids_triton.logging import logger +from rapids_triton.triton.client import STANDARD_PORTS + + +def arrays_close( + a, b, atol=None, rtol=None, total_atol=None, total_rtol=None, assert_close=False +): + """ + Compare numpy arrays for approximate equality + + :param numpy.array a: The array to compare against a reference value + :param numpy.array b: The reference array to compare against + :param float atol: The maximum absolute difference allowed between an + element in a and an element in b before they are considered non-close. + If both atol and rtol are set to None, atol is assumed to be 0. If atol + is set to None and rtol is not None, no absolute threshold is used in + comparisons. + :param float rtol: The maximum relative difference allowed between an + element in a and an element in b before they are considered non-close. + If rtol is set to None, no relative threshold is used in comparisons. + :param int total_atol: The maximum number of elements allowed to be + non-close before the arrays are considered non-close. + :param float total_rtol: The maximum proportion of elements allowed to be + non-close before the arrays are considered non-close. + """ + + if np.any(a.shape != b.shape): + if assert_close: + raise AssertionError( + "Arrays have different shapes:\n{} vs. {}".format(a.shape, b.shape) + ) + return False + + if a.size == 0 and b.size == 0: + return True + + if atol is None and rtol is None: + atol = 0 + if total_atol is None and total_rtol is None: + total_atol = 0 + + diff_mask = np.ones(a.shape, dtype="bool") + + diff = np.abs(a - b) + + if atol is not None: + diff_mask = np.logical_and(diff_mask, diff > atol) + + if rtol is not None: + diff_mask = np.logical_and(diff_mask, diff > rtol * np.abs(b)) + + is_close = True + + mismatch_count = np.sum(diff_mask) + + if total_atol is not None and mismatch_count > total_atol: + is_close = False + + mismatch_proportion = mismatch_count / a.size + if total_rtol is not None and mismatch_proportion > total_rtol: + is_close = False + + if assert_close and not is_close: + total_tol_desc = [] + if total_atol is not None: + total_tol_desc.append(str(int(total_atol))) + if total_rtol is not None: + total_tol_desc.append("{:.2f} %".format(total_rtol * 100)) + total_tol_desc = " or ".join(total_tol_desc) + + msg = """Arrays have more than {} mismatched elements. + +Mismatch in {} ({:.2f} %) elements + a: {} + b: {} + + Mismatched indices: {}""".format( + total_tol_desc, + mismatch_count, + mismatch_proportion * 100, + a, + b, + np.transpose(np.nonzero(diff_mask)), + ) + raise AssertionError(msg) + return is_close + + +def get_random_seed(): + """Provide random seed to allow for easier reproduction of testing failures + + Note: Code taken directly from cuML testing infrastructure""" + current_random_seed = os.getenv("PYTEST_RANDOM_SEED") + if current_random_seed is not None and current_random_seed.isdigit(): + random_seed = int(current_random_seed) + else: + random_seed = np.random.randint(0, 1e6) + os.environ["PYTEST_RANDOM_SEED"] = str(random_seed) + logger.info("Random seed value: %d", random_seed) + return random_seed diff --git a/rapids_triton_python/rapids_triton/triton/__init__.py b/rapids_triton_python/rapids_triton/triton/__init__.py new file mode 100644 index 0000000..3291348 --- /dev/null +++ b/rapids_triton_python/rapids_triton/triton/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/rapids_triton_python/rapids_triton/triton/client.py b/rapids_triton_python/rapids_triton/triton/client.py new file mode 100644 index 0000000..2969b2d --- /dev/null +++ b/rapids_triton_python/rapids_triton/triton/client.py @@ -0,0 +1,36 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import tritonclient.grpc as triton_grpc +import tritonclient.http as triton_http + +STANDARD_PORTS = {"http": 8000, "grpc": 8001} + + +def get_triton_client(protocol="grpc", host="localhost", port=None, concurrency=4): + """Get Triton client instance of desired type""" + + if port is None: + port = STANDARD_PORTS[protocol] + + if protocol == "grpc": + client = triton_grpc.InferenceServerClient(url=f"{host}:{port}", verbose=False) + elif protocol == "http": + client = triton_http.InferenceServerClient( + url=f"{host}:{port}", verbose=False, concurrency=concurrency + ) + else: + raise RuntimeError('Bad protocol: "{}"'.format(protocol)) + + return client diff --git a/rapids_triton_python/rapids_triton/triton/dtype.py b/rapids_triton_python/rapids_triton/triton/dtype.py new file mode 100644 index 0000000..8fa97e6 --- /dev/null +++ b/rapids_triton_python/rapids_triton/triton/dtype.py @@ -0,0 +1,35 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np + +DTYPE_NAMES = { + np.dtype("bool").str: "BOOL", + np.dtype("uint8").str: "UINT8", + np.dtype("uint16").str: "UINT16", + np.dtype("uint32").str: "UINT32", + np.dtype("uint64").str: "UINT64", + np.dtype("int8").str: "INT8", + np.dtype("int16").str: "INT16", + np.dtype("int32").str: "INT32", + np.dtype("int64").str: "INT64", + np.dtype("float16").str: "FP16", + np.dtype("float32").str: "FP32", + np.dtype("float64").str: "FP64", +} + + +def dtype_to_triton_name(dtype): + dtype = np.dtype(dtype).str + return DTYPE_NAMES.get(dtype, "BYTES") diff --git a/rapids_triton_python/rapids_triton/triton/io.py b/rapids_triton_python/rapids_triton/triton/io.py new file mode 100644 index 0000000..d82649d --- /dev/null +++ b/rapids_triton_python/rapids_triton/triton/io.py @@ -0,0 +1,145 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections import namedtuple +from uuid import uuid4 + +import tritonclient.grpc as triton_grpc +import tritonclient.http as triton_http +from rapids_triton.exceptions import IncompatibleSharedMemory +from rapids_triton.utils.safe_import import ImportReplacement +from tritonclient import utils as triton_utils + +try: + import tritonclient.utils.cuda_shared_memory as shm +except OSError: # CUDA libraries not available + shm = ImportReplacement("tritonclient.utils.cuda_shared_memory") + + +TritonInput = namedtuple("TritonInput", ("name", "handle", "input")) +TritonOutput = namedtuple("TritonOutput", ("name", "handle", "output")) + + +def set_unshared_input_data(triton_input, data, protocol="grpc"): + if protocol == "grpc": + triton_input.set_data_from_numpy(data) + else: + triton_input.set_data_from_numpy(data, binary_data=True) + + return TritonInput(None, None, triton_input) + + +def set_shared_input_data(triton_client, triton_input, data, protocol="grpc"): + input_size = data.size * data.itemsize + + input_name = "input_{}".format(uuid4().hex) + + input_handle = shm.create_shared_memory_region(input_name, input_size, 0) + + shm.set_shared_memory_region(input_handle, [data]) + + triton_client.register_cuda_shared_memory( + input_name, shm.get_raw_handle(input_handle), 0, input_size + ) + + triton_input.set_shared_memory(input_name, input_size) + + return TritonInput(input_name, input_handle, triton_input) + + +def set_input_data(triton_client, triton_input, data, protocol="grpc", shared_mem=None): + if shared_mem is None: + return set_unshared_input_data(triton_input, data, protocol=protocol) + if shared_mem == "cuda": + return set_shared_input_data( + triton_client, triton_input, data, protocol=protocol + ) + raise RuntimeError("Unsupported shared memory type") + + +def create_triton_input( + triton_client, data, name, dtype, protocol="grpc", shared_mem=None +): + if protocol == "grpc": + triton_input = triton_grpc.InferInput(name, data.shape, dtype) + else: + triton_input = triton_http.InferInput(name, data.shape, dtype) + + return set_input_data( + triton_client, triton_input, data, protocol=protocol, shared_mem=shared_mem + ) + + +def create_output_handle(triton_client, triton_output, size, shared_mem=None): + if shared_mem is None: + return (None, None) + + output_name = "output_{}".format(uuid4().hex) + output_handle = shm.create_shared_memory_region(output_name, size, 0) + + triton_client.register_cuda_shared_memory( + output_name, shm.get_raw_handle(output_handle), 0, size + ) + + triton_output.set_shared_memory(output_name, size) + + return output_name, output_handle + + +def create_triton_output(triton_client, size, name, protocol="grpc", shared_mem=None): + """Set up output memory in Triton + + Parameters + ---------- + triton_client : Triton client object + The client used to set output parameters + size : int + The size of the output in bytes + name : str + The model-defined name for this output + protocol : 'grpc' or 'http' + The protocol used for communication with the server + """ + if protocol == "grpc": + triton_output = triton_grpc.InferRequestedOutput(name) + else: + triton_output = triton_grpc.InferRequestedOutput(name, binary_data=True) + + output_name, output_handle = create_output_handle( + triton_client, triton_output, size, shared_mem=shared_mem + ) + + return TritonOutput(name=output_name, handle=output_handle, output=triton_output) + + +def destroy_shared_memory_region(handle, shared_mem="cuda"): + """Release memory from a given shared memory handle + + Parameters + ---------- + handle : c_void_p + The handle (as returned by the Triton client) for the region to be + released. + shared_mem : 'cuda' or 'system' or None + The type of shared memory region to release. If None, an exception will + be thrown. + """ + if shared_mem is None: + raise IncompatibleSharedMemory("Attempting to release non-shared memory") + elif shared_mem == "system": + raise NotImplementedError("System shared memory not yet supported") + elif shared_mem == "cuda": + shm.destroy_shared_memory_region(handle) + else: + raise NotImplementedError(f"Unrecognized memory type {shared_mem}") diff --git a/rapids_triton_python/rapids_triton/triton/message.py b/rapids_triton_python/rapids_triton/triton/message.py new file mode 100644 index 0000000..0d522ea --- /dev/null +++ b/rapids_triton_python/rapids_triton/triton/message.py @@ -0,0 +1,30 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +class TritonMessage: + """Adapter to read output from both GRPC and HTTP responses""" + + def __init__(self, message): + self.message = message + + def __getattr__(self, attr): + try: + return getattr(self.message, attr) + except AttributeError: + try: + return self.message[attr] + except Exception: # Re-raise AttributeError + pass + raise diff --git a/rapids_triton_python/rapids_triton/triton/response.py b/rapids_triton_python/rapids_triton/triton/response.py new file mode 100644 index 0000000..4e7235f --- /dev/null +++ b/rapids_triton_python/rapids_triton/triton/response.py @@ -0,0 +1,35 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from rapids_triton.triton.message import TritonMessage +from rapids_triton.utils.safe_import import ImportReplacement +from tritonclient import utils as triton_utils + +try: + import tritonclient.utils.cuda_shared_memory as shm +except OSError: # CUDA libraries not available + shm = ImportReplacement("tritonclient.utils.cuda_shared_memory") + + +def get_response_data(response, output_handle, output_name): + """Convert Triton response to NumPy array""" + if output_handle is None: + return response.as_numpy(output_name) + else: + network_result = TritonMessage(response.get_output(output_name)) + return shm.get_contents_as_numpy( + output_handle, + triton_utils.triton_to_np_dtype(network_result.datatype), + network_result.shape, + ) diff --git a/rapids_triton_python/rapids_triton/utils/__init__.py b/rapids_triton_python/rapids_triton/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rapids_triton_python/rapids_triton/utils/safe_import.py b/rapids_triton_python/rapids_triton/utils/safe_import.py new file mode 100644 index 0000000..1bc93f8 --- /dev/null +++ b/rapids_triton_python/rapids_triton/utils/safe_import.py @@ -0,0 +1,39 @@ +# Copyright (c) 2022, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +class ImportUnavailableError(Exception): + """Error thrown if a symbol is unavailable due to an issue importing it""" + + +class ImportReplacement: + """A class to be used in place of an importable symbol if that symbol + cannot be imported + + Parameters + ---------- + symbol: str + The name or import path to be used in error messages when attempting to + make use of this symbol. E.g. "some_pkg.func" would result in an + exception with message "some_pkg.func could not be imported" + """ + + def __init__(self, symbol): + self._msg = f"{symbol} could not be imported" + + def __getattr__(self, name): + raise ImportUnavailableError(self._msg) + + def __call__(self, *args, **kwargs): + raise ImportUnavailableError(self._msg) From f7ad9dda8a811533f8821b5e21f94129dd31a871 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 6 Aug 2026 17:27:14 -0700 Subject: [PATCH 6/8] Update tests to use nvForest instead of FIL --- conda/environments/triton_test.yml | 1 + conda/environments/triton_test_no_client.yml | 1 + qa/L0_e2e/test_model.py | 20 ++++++++++++++------ 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/conda/environments/triton_test.yml b/conda/environments/triton_test.yml index b1f93d2..e010d45 100644 --- a/conda/environments/triton_test.yml +++ b/conda/environments/triton_test.yml @@ -9,6 +9,7 @@ dependencies: - cuda-version=12.8 - cudf=26.08 - cuml=26.08 + - nvforest=26.08 - flake8 - hypothesis - lightgbm diff --git a/conda/environments/triton_test_no_client.yml b/conda/environments/triton_test_no_client.yml index 79cf498..1b2376e 100644 --- a/conda/environments/triton_test_no_client.yml +++ b/conda/environments/triton_test_no_client.yml @@ -9,6 +9,7 @@ dependencies: - cuda-version=12.8 - cudf=26.08 - cuml=26.08 + - nvforest=26.08 - flake8 - hypothesis - lightgbm diff --git a/qa/L0_e2e/test_model.py b/qa/L0_e2e/test_model.py index 233fdec..3a63b88 100644 --- a/qa/L0_e2e/test_model.py +++ b/qa/L0_e2e/test_model.py @@ -18,9 +18,17 @@ from functools import lru_cache try: - import cuml + import nvforest except Exception: - cuml = None + nvforest = None +try: + from cupy import asnumpy +except Exception: + + def asnumpy(x): + return x + + import numpy as np import pytest import treelite @@ -199,20 +207,20 @@ def __init__( with open(model_path, "rb") as pkl_file: self._base_model = pickle.load(pkl_file) else: - self._base_model = cuml.ForestInference.load( - model_path, is_classifier=is_classifier, model_type=model_format + self._base_model = nvforest.load_model( + model_file=model_path, model_type=model_format ) def predict(self, inputs): if self.predict_proba: - result = self._base_model.predict_proba(inputs["input__0"]) + result = asnumpy(self._base_model.predict_proba(inputs["input__0"])) if len(result.shape) == 1 or ( len(result.shape) == 2 and result.shape[1] == 1 ): result = result.reshape((-1, 1)) result = np.concatenate((1 - result, result), axis=1) else: - result = self._base_model.predict(inputs["input__0"]) + result = asnumpy(self._base_model.predict(inputs["input__0"])) output = {"output__0": result.squeeze()} if self._run_treeshap: treeshap_result = self._xgb_model.predict( From 37359ca79a2bbcd1379bbad45da977d9679d797e Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 6 Aug 2026 17:47:45 -0700 Subject: [PATCH 7/8] Use CUDA 13.3 --- conda/environments/triton_benchmark.yml | 3 +-- conda/environments/triton_test.yml | 3 +-- conda/environments/triton_test_no_client.yml | 2 +- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/conda/environments/triton_benchmark.yml b/conda/environments/triton_benchmark.yml index 59e9bcd..76db0ce 100644 --- a/conda/environments/triton_benchmark.yml +++ b/conda/environments/triton_benchmark.yml @@ -4,7 +4,7 @@ channels: - nvidia - rapidsai dependencies: - - cuda-version=12.8 + - cuda-version=13.3 - cudf=26.08 - libcusolver - libcusparse @@ -15,4 +15,3 @@ dependencies: - pip: - tritonclient[all] - protobuf - - git+https://github.com/rapidsai/rapids-triton.git@branch-25.08#subdirectory=python diff --git a/conda/environments/triton_test.yml b/conda/environments/triton_test.yml index e010d45..0f305e9 100644 --- a/conda/environments/triton_test.yml +++ b/conda/environments/triton_test.yml @@ -6,7 +6,7 @@ channels: dependencies: - aws-sdk-cpp - clang-tools=19.1.7 - - cuda-version=12.8 + - cuda-version=13.3 - cudf=26.08 - cuml=26.08 - nvforest=26.08 @@ -24,4 +24,3 @@ dependencies: - pip: - tritonclient[all] - protobuf - - git+https://github.com/rapidsai/rapids-triton.git@branch-25.08#subdirectory=python diff --git a/conda/environments/triton_test_no_client.yml b/conda/environments/triton_test_no_client.yml index 1b2376e..57d741c 100644 --- a/conda/environments/triton_test_no_client.yml +++ b/conda/environments/triton_test_no_client.yml @@ -6,7 +6,7 @@ channels: dependencies: - aws-sdk-cpp - clang-tools=19.1.7 - - cuda-version=12.8 + - cuda-version=13.3 - cudf=26.08 - cuml=26.08 - nvforest=26.08 From b644609a1b1c48452fda2cb37400529381f9f6d5 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 6 Aug 2026 18:09:29 -0700 Subject: [PATCH 8/8] Install cuda-python --- conda/environments/triton_benchmark.yml | 1 + conda/environments/triton_test.yml | 1 + conda/environments/triton_test_no_client.yml | 1 + 3 files changed, 3 insertions(+) diff --git a/conda/environments/triton_benchmark.yml b/conda/environments/triton_benchmark.yml index 76db0ce..17744c9 100644 --- a/conda/environments/triton_benchmark.yml +++ b/conda/environments/triton_benchmark.yml @@ -5,6 +5,7 @@ channels: - rapidsai dependencies: - cuda-version=13.3 + - cuda-python=13.3 - cudf=26.08 - libcusolver - libcusparse diff --git a/conda/environments/triton_test.yml b/conda/environments/triton_test.yml index 0f305e9..4204734 100644 --- a/conda/environments/triton_test.yml +++ b/conda/environments/triton_test.yml @@ -7,6 +7,7 @@ dependencies: - aws-sdk-cpp - clang-tools=19.1.7 - cuda-version=13.3 + - cuda-python=13.3 - cudf=26.08 - cuml=26.08 - nvforest=26.08 diff --git a/conda/environments/triton_test_no_client.yml b/conda/environments/triton_test_no_client.yml index 57d741c..89557ef 100644 --- a/conda/environments/triton_test_no_client.yml +++ b/conda/environments/triton_test_no_client.yml @@ -7,6 +7,7 @@ dependencies: - aws-sdk-cpp - clang-tools=19.1.7 - cuda-version=13.3 + - cuda-python=13.3 - cudf=26.08 - cuml=26.08 - nvforest=26.08