├── .clang-format ├── .github └── workflows │ └── pre-commit.yml ├── .gitignore ├── .pre-commit-config.yaml ├── CMakeLists.txt ├── LICENSE ├── README.md ├── cmake └── TritonCommonConfig.cmake.in ├── include └── triton │ └── common │ ├── async_work_queue.h │ ├── error.h │ ├── logging.h │ ├── model_config.h │ ├── nvtx.h │ ├── sync_queue.h │ ├── table_printer.h │ ├── thread_pool.h │ └── triton_json.h ├── protobuf ├── CMakeLists.txt ├── grpc_service.proto ├── health.proto └── model_config.proto ├── pyproject.toml └── src ├── CMakeLists.txt ├── async_work_queue.cc ├── error.cc ├── logging.cc ├── model_config.cc ├── table_printer.cc ├── test ├── CMakeLists.txt └── triton_json │ ├── CMakeLists.txt │ └── triton_json_test.cc └── thread_pool.cc /.clang-format: -------------------------------------------------------------------------------- 1 | --- 2 | BasedOnStyle: Google 3 | 4 | IndentWidth: 2 5 | ColumnLimit: 80 6 | ContinuationIndentWidth: 4 7 | UseTab: Never 8 | MaxEmptyLinesToKeep: 2 9 | 10 | SortIncludes: true 11 | CompactNamespaces: true 12 | ReflowComments: true 13 | 14 | DerivePointerAlignment: false 15 | PointerAlignment: Left 16 | 17 | AllowShortIfStatementsOnASingleLine: false 18 | AllowShortBlocksOnASingleLine: false 19 | AllowShortFunctionsOnASingleLine: Inline 20 | 21 | AlwaysBreakAfterReturnType: TopLevelDefinitions 22 | AlignAfterOpenBracket: AlwaysBreak 23 | BreakBeforeBraces: Custom 24 | BraceWrapping: 25 | AfterClass: false 26 | AfterControlStatement: false 27 | AfterEnum: false 28 | AfterFunction: true 29 | AfterNamespace: false 30 | AfterStruct: false 31 | AfterUnion: false 32 | BeforeCatch: true 33 | 34 | BinPackArguments: true 35 | BinPackParameters: true 36 | ConstructorInitializerAllOnOneLineOrOnePerLine: false 37 | 38 | IndentCaseLabels: true 39 | -------------------------------------------------------------------------------- /.github/workflows/pre-commit.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # Redistribution and use in source and binary forms, with or without 4 | # modification, are permitted provided that the following conditions 5 | # are met: 6 | # * Redistributions of source code must retain the above copyright 7 | # notice, this list of conditions and the following disclaimer. 8 | # * Redistributions in binary form must reproduce the above copyright 9 | # notice, this list of conditions and the following disclaimer in the 10 | # documentation and/or other materials provided with the distribution. 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its 12 | # contributors may be used to endorse or promote products derived 13 | # from this software without specific prior written permission. 14 | # 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | name: pre-commit 28 | 29 | on: 30 | pull_request: 31 | 32 | jobs: 33 | pre-commit: 34 | runs-on: ubuntu-22.04 35 | steps: 36 | - uses: actions/checkout@v3 37 | - uses: actions/setup-python@v3 38 | - uses: pre-commit/action@v3.0.0 39 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /build 2 | /.vscode 3 | *.so 4 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # Redistribution and use in source and binary forms, with or without 4 | # modification, are permitted provided that the following conditions 5 | # are met: 6 | # * Redistributions of source code must retain the above copyright 7 | # notice, this list of conditions and the following disclaimer. 8 | # * Redistributions in binary form must reproduce the above copyright 9 | # notice, this list of conditions and the following disclaimer in the 10 | # documentation and/or other materials provided with the distribution. 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its 12 | # contributors may be used to endorse or promote products derived 13 | # from this software without specific prior written permission. 14 | # 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | repos: 28 | - repo: https://github.com/timothycrosley/isort 29 | rev: 5.12.0 30 | hooks: 31 | - id: isort 32 | additional_dependencies: [toml] 33 | - repo: https://github.com/psf/black 34 | rev: 23.1.0 35 | hooks: 36 | - id: black 37 | types_or: [python, cython] 38 | - repo: https://github.com/PyCQA/flake8 39 | rev: 5.0.4 40 | hooks: 41 | - id: flake8 42 | args: [--max-line-length=88, --select=C,E,F,W,B,B950, --extend-ignore = E203,E501] 43 | types_or: [python, cython] 44 | - repo: https://github.com/pre-commit/mirrors-clang-format 45 | rev: v16.0.5 46 | hooks: 47 | - id: clang-format 48 | types_or: [c, c++, cuda, proto, textproto, java] 49 | args: ["-fallback-style=none", "-style=file", "-i"] 50 | - repo: https://github.com/codespell-project/codespell 51 | rev: v2.2.4 52 | hooks: 53 | - id: codespell 54 | additional_dependencies: [tomli] 55 | args: ["--toml", "pyproject.toml"] 56 | exclude: (?x)^(.*stemmer.*|.*stop_words.*|^CHANGELOG.md$) 57 | # More details about these pre-commit hooks here: 58 | # https://pre-commit.com/hooks.html 59 | - repo: https://github.com/pre-commit/pre-commit-hooks 60 | rev: v4.4.0 61 | hooks: 62 | - id: check-case-conflict 63 | - id: check-executables-have-shebangs 64 | - id: check-merge-conflict 65 | - id: check-json 66 | - id: check-toml 67 | - id: check-yaml 68 | - id: check-shebang-scripts-are-executable 69 | - id: end-of-file-fixer 70 | types_or: [c, c++, cuda, proto, textproto, java, python] 71 | - id: mixed-line-ending 72 | - id: requirements-txt-fixer 73 | - id: trailing-whitespace 74 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # Redistribution and use in source and binary forms, with or without 4 | # modification, are permitted provided that the following conditions 5 | # are met: 6 | # * Redistributions of source code must retain the above copyright 7 | # notice, this list of conditions and the following disclaimer. 8 | # * Redistributions in binary form must reproduce the above copyright 9 | # notice, this list of conditions and the following disclaimer in the 10 | # documentation and/or other materials provided with the distribution. 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its 12 | # contributors may be used to endorse or promote products derived 13 | # from this software without specific prior written permission. 14 | # 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | cmake_minimum_required(VERSION 3.17) 28 | 29 | project(tritoncommon LANGUAGES C CXX) 30 | 31 | # 32 | # C++ standard 33 | # 34 | set(TRITON_MIN_CXX_STANDARD 17 CACHE STRING "The minimum C++ standard whose features are requested to build this target.") 35 | 36 | # 37 | # Options 38 | # 39 | # Some components are expensive to build and have extensive 40 | # dependencies, so those parts of the build must be enabled 41 | # explicitly. 42 | option(TRITON_COMMON_ENABLE_PROTOBUF "Build protobuf artifacts" OFF) 43 | option(TRITON_COMMON_ENABLE_PROTOBUF_PYTHON "Build protobuf artifacts for python" ON) 44 | option(TRITON_COMMON_ENABLE_GRPC "Build grpc artifacts" OFF) 45 | option(TRITON_COMMON_ENABLE_JSON "Build json-related libs" ON) 46 | 47 | if(TRITON_COMMON_ENABLE_JSON) 48 | find_package(RapidJSON CONFIG REQUIRED) 49 | message(STATUS "RapidJSON found. Headers: ${RAPIDJSON_INCLUDE_DIRS}") 50 | endif() 51 | 52 | set(THREADS_PREFER_PTHREAD_FLAG TRUE) 53 | find_package(Threads REQUIRED) 54 | 55 | if(CMAKE_CXX_COMPILER_ID STREQUAL "MSVC") 56 | message("Using MSVC as compiler, default target on Windows 10. " 57 | "If the target system is not Windows 10, please update _WIN32_WINNT " 58 | "to corresponding value.") 59 | endif() 60 | 61 | add_library(common-compile-settings INTERFACE) 62 | 63 | target_compile_features(common-compile-settings INTERFACE cxx_std_${TRITON_MIN_CXX_STANDARD}) 64 | 65 | target_compile_options(common-compile-settings INTERFACE 66 | $<$,$,$>: 67 | -Wall -Wextra -Wno-unused-parameter -Wno-type-limits> 68 | $<$:/W0 /D_WIN32_WINNT=0x0A00 /EHsc /Zc:preprocessor> 69 | ) 70 | 71 | # 72 | # SyncQueue 73 | # 74 | add_library( 75 | triton-common-sync-queue INTERFACE 76 | ) 77 | 78 | add_library( 79 | TritonCommon::triton-common-sync-queue ALIAS triton-common-sync-queue 80 | ) 81 | 82 | target_include_directories( 83 | triton-common-sync-queue 84 | INTERFACE 85 | $ 86 | $ 87 | ) 88 | 89 | 90 | # 91 | # JSON utilities 92 | # 93 | if(TRITON_COMMON_ENABLE_JSON) 94 | add_library( 95 | triton-common-json INTERFACE 96 | ) 97 | 98 | add_library( 99 | TritonCommon::triton-common-json ALIAS triton-common-json 100 | ) 101 | 102 | target_include_directories( 103 | triton-common-json 104 | INTERFACE 105 | $ 106 | $ 107 | $ 108 | $ 109 | ) 110 | endif() 111 | 112 | add_subdirectory(src) 113 | 114 | # 115 | # Protobuf and GRPC artifacts 116 | # 117 | if(${TRITON_COMMON_ENABLE_PROTOBUF} OR ${TRITON_COMMON_ENABLE_GRPC}) 118 | add_subdirectory(protobuf) 119 | 120 | set(protobuf_MODULE_COMPATIBLE TRUE CACHE BOOL "protobuf_MODULE_COMPATIBLE" FORCE) 121 | find_package(Protobuf CONFIG REQUIRED) 122 | message(STATUS "Using protobuf ${Protobuf_VERSION}") 123 | 124 | # 125 | # Model Config (depends on protobuf & generated .pb.h file) 126 | # 127 | add_library( 128 | triton-common-model-config 129 | src/model_config.cc 130 | ) 131 | 132 | add_library( 133 | TritonCommon::triton-common-model-config ALIAS triton-common-model-config 134 | ) 135 | 136 | target_include_directories( 137 | triton-common-model-config 138 | PUBLIC 139 | $ 140 | $ 141 | PRIVATE 142 | ${CMAKE_CURRENT_SOURCE_DIR}/src 143 | ${Protobuf_INCLUDE_DIRS} 144 | ) 145 | 146 | target_link_libraries( 147 | triton-common-model-config 148 | PRIVATE 149 | common-compile-settings 150 | protobuf::libprotobuf 151 | proto-library 152 | ) 153 | 154 | set_target_properties( 155 | triton-common-model-config 156 | PROPERTIES 157 | WINDOWS_EXPORT_ALL_SYMBOLS TRUE 158 | POSITION_INDEPENDENT_CODE ON 159 | OUTPUT_NAME tritoncommonmodelconfig 160 | ) 161 | 162 | endif() 163 | 164 | # 165 | # Install 166 | # 167 | include(GNUInstallDirs) 168 | set(INSTALL_CONFIGDIR ${CMAKE_INSTALL_LIBDIR}/cmake/TritonCommon) 169 | 170 | install( 171 | TARGETS 172 | triton-common-async-work-queue 173 | triton-common-error 174 | triton-common-logging 175 | triton-common-sync-queue 176 | triton-common-table-printer 177 | triton-common-thread-pool 178 | common-compile-settings 179 | EXPORT 180 | triton-common-targets 181 | LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} 182 | ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} 183 | ) 184 | 185 | if(TRITON_COMMON_ENABLE_JSON) 186 | install( 187 | TARGETS 188 | triton-common-json 189 | EXPORT 190 | triton-common-targets 191 | LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} 192 | ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} 193 | ) 194 | endif() 195 | 196 | if(${TRITON_COMMON_ENABLE_GRPC} OR ${TRITON_COMMON_ENABLE_PROTOBUF}) 197 | install( 198 | TARGETS 199 | proto-library 200 | triton-common-model-config 201 | # proto-py-library 202 | EXPORT 203 | triton-common-targets 204 | LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} 205 | ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} 206 | ) 207 | endif() 208 | 209 | if(${TRITON_COMMON_ENABLE_GRPC}) 210 | install( 211 | TARGETS 212 | grpc-health-library 213 | grpc-service-library 214 | # grpc-service-py-library 215 | EXPORT 216 | triton-common-targets 217 | LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} 218 | ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} 219 | ) 220 | endif() 221 | 222 | install( 223 | DIRECTORY include/ 224 | DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} 225 | ) 226 | 227 | install( 228 | EXPORT 229 | triton-common-targets 230 | FILE 231 | TritonCommonTargets.cmake 232 | NAMESPACE 233 | TritonCommon:: 234 | DESTINATION 235 | ${INSTALL_CONFIGDIR} 236 | ) 237 | 238 | include(CMakePackageConfigHelpers) 239 | configure_package_config_file( 240 | ${CMAKE_CURRENT_LIST_DIR}/cmake/TritonCommonConfig.cmake.in 241 | ${CMAKE_CURRENT_BINARY_DIR}/TritonCommonConfig.cmake 242 | INSTALL_DESTINATION ${INSTALL_CONFIGDIR} 243 | ) 244 | 245 | install( 246 | FILES 247 | ${CMAKE_CURRENT_BINARY_DIR}/TritonCommonConfig.cmake 248 | DESTINATION 249 | ${INSTALL_CONFIGDIR} 250 | ) 251 | 252 | # 253 | # Export from build tree 254 | # 255 | export( 256 | EXPORT 257 | triton-common-targets 258 | FILE 259 | ${CMAKE_CURRENT_BINARY_DIR}/TritonCommonTargets.cmake 260 | NAMESPACE 261 | TritonCommon:: 262 | ) 263 | 264 | export(PACKAGE TritonCommon) 265 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. 2 | 3 | Redistribution and use in source and binary forms, with or without 4 | modification, are permitted provided that the following conditions 5 | are met: 6 | * Redistributions of source code must retain the above copyright 7 | notice, this list of conditions and the following disclaimer. 8 | * Redistributions in binary form must reproduce the above copyright 9 | notice, this list of conditions and the following disclaimer in the 10 | documentation and/or other materials provided with the distribution. 11 | * Neither the name of NVIDIA CORPORATION nor the names of its 12 | contributors may be used to endorse or promote products derived 13 | from this software without specific prior written permission. 14 | 15 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 28 | 29 | [![License](https://img.shields.io/badge/License-BSD3-lightgrey.svg)](https://opensource.org/licenses/BSD-3-Clause) 30 | 31 | # Triton Inference Server Common 32 | 33 | Common source, scripts and utilities shared across all Triton 34 | repositories. 35 | 36 | This repo is not typically built directly but is instead included in 37 | the build of other repos. To build directly first install the required 38 | dependencies. 39 | 40 | ``` 41 | $ apt-get install rapidjson-dev 42 | ``` 43 | 44 | Use cmake 3.17 or later to build and install in a local directory. 45 | 46 | ``` 47 | $ mkdir build 48 | $ cd build 49 | $ cmake -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install .. 50 | $ make install 51 | ``` 52 | -------------------------------------------------------------------------------- /cmake/TritonCommonConfig.cmake.in: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Redistribution and use in source and binary forms, with or without 4 | # modification, are permitted provided that the following conditions 5 | # are met: 6 | # * Redistributions of source code must retain the above copyright 7 | # notice, this list of conditions and the following disclaimer. 8 | # * Redistributions in binary form must reproduce the above copyright 9 | # notice, this list of conditions and the following disclaimer in the 10 | # documentation and/or other materials provided with the distribution. 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its 12 | # contributors may be used to endorse or promote products derived 13 | # from this software without specific prior written permission. 14 | # 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | @PACKAGE_INIT@ 28 | 29 | set_and_check(TRITONCOMMON_CMAKE_DIR "${CMAKE_CURRENT_LIST_DIR}") 30 | 31 | list(APPEND CMAKE_MODULE_PATH ${TRITONCOMMON_CMAKE_DIR}) 32 | 33 | include(CMakeFindDependencyMacro) 34 | find_dependency(Threads) 35 | 36 | if(NOT TARGET TritonCommon::triton-common-json) 37 | include("${TRITONCOMMON_CMAKE_DIR}/TritonCommonTargets.cmake") 38 | endif() 39 | 40 | check_required_components(triton-common-json 41 | triton-common-sync-queue 42 | triton-common-async-work-queue 43 | triton-common-thread-pool 44 | ) 45 | 46 | set(TRITONCOMMON_LIBRARIES 47 | TritonCommon::triton-common-json 48 | TritonCommon::triton-common-sync-queue 49 | TritonCommon::triton-common-async-work-queue 50 | TritonCommon::triton-common-thread-pool 51 | ) 52 | -------------------------------------------------------------------------------- /include/triton/common/async_work_queue.h: -------------------------------------------------------------------------------- 1 | // Copyright 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | #pragma once 27 | 28 | #include "error.h" 29 | #include "thread_pool.h" 30 | 31 | namespace triton { namespace common { 32 | // Manager for asynchronous worker threads. Use to accelerate copies and 33 | // other such operations by running them in parallel. 34 | // Call Initialize to start the worker threads (once) and AddTask to tasks to 35 | // the queue. 36 | 37 | class AsyncWorkQueue { 38 | public: 39 | // Start 'worker_count' number of worker threads. 40 | static Error Initialize(size_t worker_count); 41 | 42 | // Get the number of worker threads. 43 | static size_t WorkerCount(); 44 | 45 | // Add a 'task' to the queue. The function will take ownership of 'task'. 46 | // Therefore std::move should be used when calling AddTask. 47 | static Error AddTask(std::function&& task); 48 | 49 | protected: 50 | static void Reset(); 51 | 52 | private: 53 | AsyncWorkQueue() = default; 54 | ~AsyncWorkQueue(); 55 | static AsyncWorkQueue* GetSingleton(); 56 | std::unique_ptr thread_pool_; 57 | }; 58 | 59 | }} // namespace triton::common 60 | -------------------------------------------------------------------------------- /include/triton/common/error.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | #pragma once 27 | 28 | #include 29 | #include 30 | 31 | namespace triton { namespace common { 32 | 33 | // 34 | // Error 35 | // 36 | // Error returned by utilities from common repo. 37 | // 38 | class Error { 39 | public: 40 | enum class Code : uint8_t { 41 | SUCCESS, 42 | UNKNOWN, 43 | INTERNAL, 44 | NOT_FOUND, 45 | INVALID_ARG, 46 | UNAVAILABLE, 47 | UNSUPPORTED, 48 | ALREADY_EXISTS, 49 | CANCELLED 50 | }; 51 | 52 | explicit Error(Code code = Code::SUCCESS) : code_(code) {} 53 | explicit Error(Code code, const std::string& msg) : code_(code), msg_(msg) {} 54 | 55 | // Convenience "success" value. Can be used as Error::Success to 56 | // indicate no error. 57 | static const Error Success; 58 | 59 | // Return the code for this status. 60 | Code ErrorCode() const { return code_; } 61 | 62 | // Return the message for this status. 63 | const std::string& Message() const { return msg_; } 64 | 65 | // Return true if this status indicates "ok"/"success", false if 66 | // status indicates some kind of failure. 67 | bool IsOk() const { return code_ == Code::SUCCESS; } 68 | 69 | // Return the status as a string. 70 | std::string AsString() const; 71 | 72 | // Return the constant string name for a code. 73 | static const char* CodeString(const Code code); 74 | 75 | protected: 76 | Code code_; 77 | std::string msg_; 78 | }; 79 | 80 | }} // namespace triton::common 81 | -------------------------------------------------------------------------------- /include/triton/common/logging.h: -------------------------------------------------------------------------------- 1 | // Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | #pragma once 27 | 28 | #include 29 | #include 30 | #include 31 | #include 32 | #include 33 | #include 34 | #include 35 | #include 36 | 37 | #include "table_printer.h" 38 | #ifdef _WIN32 39 | // exclude winsock apis 40 | #define WIN32_LEAN_AND_MEAN 41 | // suppress the min and max definitions in Windef.h. 42 | #define NOMINMAX 43 | #include 44 | #else 45 | #include 46 | #include 47 | #include 48 | #include 49 | #endif 50 | 51 | 52 | namespace triton { namespace common { 53 | 54 | 55 | // Global logger for messages. Controls how log messages are reported. 56 | class Logger { 57 | public: 58 | // Log Formats. 59 | enum class Format { kDEFAULT, kISO8601 }; 60 | 61 | // Log levels. 62 | enum class Level : uint8_t { kERROR = 0, kWARNING = 1, kINFO = 2, kEND }; 63 | 64 | inline static const std::array(Level::kEND)> 65 | LEVEL_NAMES{"E", "W", "I"}; 66 | 67 | Logger(); 68 | 69 | // Is a log level enabled. 70 | bool IsEnabled(Level level) const 71 | { 72 | return enables_[static_cast(level)]; 73 | } 74 | 75 | // Set enable for a log Level. 76 | void SetEnabled(Level level, bool enable) 77 | { 78 | enables_[static_cast(level)] = enable; 79 | } 80 | 81 | // Get the current verbose logging level. 82 | uint32_t VerboseLevel() const { return vlevel_; } 83 | 84 | // Set the current verbose logging level. 85 | void SetVerboseLevel(uint32_t vlevel) { vlevel_ = vlevel; } 86 | 87 | // Whether to escape log messages 88 | // using JSON string escaping rules. 89 | // Default is true but can be disabled by setting 90 | // the following environment variable to '0'. 91 | // If the variable is unset or set to any value !='0' 92 | // log messages will be escaped 93 | // 94 | // TRITON_SERVER_ESCAPE_LOG_MESSAGES=0 95 | bool EscapeLogMessages() const { return escape_log_messages_; }; 96 | 97 | // Get the logging format. 98 | Format LogFormat() { return format_; } 99 | 100 | // Get the logging format as a string. 101 | std::string LogFormatString() 102 | { 103 | switch (format_) { 104 | case Format::kISO8601: 105 | return "ISO8601"; 106 | case Format::kDEFAULT: 107 | return "default"; 108 | default: 109 | return "Invalid format"; 110 | } 111 | } 112 | 113 | // Set the logging format. 114 | void SetLogFormat(Format format) { format_ = format; } 115 | 116 | // Get the log output file name. 117 | const std::string& LogFile() { return filename_; } 118 | 119 | // Set the log output file. Returns an empty string upon 120 | // success, else returns an error string. 121 | const std::string SetLogFile(const std::string& filename) 122 | { 123 | const std::lock_guard lock(mutex_); 124 | file_stream_.close(); 125 | std::string revert_name(filename_); 126 | filename_ = filename; 127 | if (!filename_.empty()) { 128 | file_stream_.open(filename_, std::ios::app); 129 | if (file_stream_.fail()) { 130 | std::stringstream error; 131 | error << __FILE__ << " " << __LINE__ 132 | << ": Failed to open log file: " << std::strerror(errno) 133 | << std::endl; 134 | filename_ = revert_name; 135 | file_stream_.open(filename_, std::ios::app); 136 | return error.str(); 137 | } 138 | } 139 | // will return an empty string 140 | return std::string(); 141 | } 142 | 143 | // Log a message. 144 | void Log(const std::string& msg, const Logger::Level level); 145 | 146 | // Flush the log. 147 | void Flush(); 148 | 149 | private: 150 | inline static const char* ESCAPE_ENVIRONMENT_VARIABLE = 151 | "TRITON_SERVER_ESCAPE_LOG_MESSAGES"; 152 | bool escape_log_messages_; 153 | std::array(Level::kEND)> enables_; 154 | uint32_t vlevel_; 155 | Format format_; 156 | std::mutex mutex_; 157 | std::string filename_; 158 | std::ofstream file_stream_; 159 | }; 160 | 161 | extern Logger gLogger_; 162 | 163 | // A log message. 164 | class LogMessage { 165 | public: 166 | LogMessage( 167 | const char* file, int line, Logger::Level level, 168 | const char* heading = nullptr, 169 | bool escape_log_messages = gLogger_.EscapeLogMessages()) 170 | : path_(file), line_(line), level_(level), pid_(GetProcessId()), 171 | heading_(heading), escape_log_messages_(escape_log_messages) 172 | { 173 | SetTimestamp(); 174 | size_t path_start = path_.rfind('/'); 175 | if (path_start != std::string::npos) { 176 | path_ = path_.substr(path_start + 1, std::string::npos); 177 | } 178 | } 179 | 180 | ~LogMessage(); 181 | 182 | std::stringstream& stream() { return message_; } 183 | 184 | private: 185 | std::string path_; 186 | const int line_; 187 | const Logger::Level level_; 188 | const uint32_t pid_; 189 | void LogPreamble(std::stringstream& stream); 190 | void LogTimestamp(std::stringstream& stream); 191 | 192 | #ifdef _WIN32 193 | SYSTEMTIME timestamp_; 194 | void SetTimestamp() { GetSystemTime(×tamp_); } 195 | static uint32_t GetProcessId() 196 | { 197 | return static_cast(GetCurrentProcessId()); 198 | }; 199 | #else 200 | struct timeval timestamp_; 201 | void SetTimestamp() { gettimeofday(×tamp_, NULL); } 202 | static uint32_t GetProcessId() { return static_cast(getpid()); }; 203 | #endif 204 | std::stringstream message_; 205 | const char* heading_; 206 | bool escape_log_messages_; 207 | }; 208 | 209 | #define LOG_ENABLE_INFO(E) \ 210 | triton::common::gLogger_.SetEnabled(triton::common::Logger::Level::kINFO, (E)) 211 | #define LOG_ENABLE_WARNING(E) \ 212 | triton::common::gLogger_.SetEnabled( \ 213 | triton::common::Logger::Level::kWARNING, (E)) 214 | #define LOG_ENABLE_ERROR(E) \ 215 | triton::common::gLogger_.SetEnabled( \ 216 | triton::common::Logger::Level::kERROR, (E)) 217 | #define LOG_SET_VERBOSE(L) \ 218 | triton::common::gLogger_.SetVerboseLevel( \ 219 | static_cast(std::max(0, (L)))) 220 | #define LOG_SET_OUT_FILE(FN) triton::common::gLogger_.SetLogFile((FN)) 221 | #define LOG_SET_FORMAT(F) triton::common::gLogger_.SetLogFormat((F)) 222 | 223 | #define LOG_VERBOSE_LEVEL triton::common::gLogger_.VerboseLevel() 224 | #define LOG_FORMAT triton::common::gLogger_.LogFormat() 225 | #define LOG_FORMAT_STRING triton::common::gLogger_.LogFormatString() 226 | #define LOG_FILE triton::common::gLogger_.LogFile() 227 | 228 | #ifdef TRITON_ENABLE_LOGGING 229 | 230 | #define LOG_INFO_IS_ON \ 231 | triton::common::gLogger_.IsEnabled(triton::common::Logger::Level::kINFO) 232 | #define LOG_WARNING_IS_ON \ 233 | triton::common::gLogger_.IsEnabled(triton::common::Logger::Level::kWARNING) 234 | #define LOG_ERROR_IS_ON \ 235 | triton::common::gLogger_.IsEnabled(triton::common::Logger::Level::kERROR) 236 | #define LOG_VERBOSE_IS_ON(L) (triton::common::gLogger_.VerboseLevel() >= (L)) 237 | 238 | #else 239 | 240 | // If logging is disabled, define macro to be false to avoid further evaluation 241 | #define LOG_INFO_IS_ON false 242 | #define LOG_WARNING_IS_ON false 243 | #define LOG_ERROR_IS_ON false 244 | #define LOG_VERBOSE_IS_ON(L) false 245 | 246 | #endif // TRITON_ENABLE_LOGGING 247 | 248 | // Macros that use explicitly given filename and line number. 249 | #define LOG_INFO_FL(FN, LN) \ 250 | if (LOG_INFO_IS_ON) \ 251 | triton::common::LogMessage( \ 252 | (char*)(FN), LN, triton::common::Logger::Level::kINFO) \ 253 | .stream() 254 | #define LOG_WARNING_FL(FN, LN) \ 255 | if (LOG_WARNING_IS_ON) \ 256 | triton::common::LogMessage( \ 257 | (char*)(FN), LN, triton::common::Logger::Level::kWARNING) \ 258 | .stream() 259 | #define LOG_ERROR_FL(FN, LN) \ 260 | if (LOG_ERROR_IS_ON) \ 261 | triton::common::LogMessage( \ 262 | (char*)(FN), LN, triton::common::Logger::Level::kERROR) \ 263 | .stream() 264 | #define LOG_VERBOSE_FL(L, FN, LN) \ 265 | if (LOG_VERBOSE_IS_ON(L)) \ 266 | triton::common::LogMessage( \ 267 | (char*)(FN), LN, triton::common::Logger::Level::kINFO) \ 268 | .stream() 269 | 270 | // Macros that use current filename and line number. 271 | #define LOG_INFO LOG_INFO_FL(__FILE__, __LINE__) 272 | #define LOG_WARNING LOG_WARNING_FL(__FILE__, __LINE__) 273 | #define LOG_ERROR LOG_ERROR_FL(__FILE__, __LINE__) 274 | #define LOG_VERBOSE(L) LOG_VERBOSE_FL(L, __FILE__, __LINE__) 275 | 276 | // Macros for use with triton::common::table_printer objects 277 | // 278 | // Data is assumed to be server / backend generated 279 | // and not for use with client input. 280 | // 281 | // Tables are printed without escaping 282 | #define LOG_TABLE_VERBOSE(L, TABLE) \ 283 | \ 284 | do { \ 285 | if (LOG_VERBOSE_IS_ON(L)) \ 286 | triton::common::LogMessage( \ 287 | __FILE__, __LINE__, triton::common::Logger::Level::kINFO, nullptr, \ 288 | false) \ 289 | .stream() \ 290 | << TABLE.PrintTable(); \ 291 | } while (false) 292 | 293 | #define LOG_TABLE_INFO(TABLE) \ 294 | do { \ 295 | if (LOG_INFO_IS_ON) \ 296 | triton::common::LogMessage( \ 297 | __FILE__, __LINE__, triton::common::Logger::Level::kINFO, nullptr, \ 298 | false) \ 299 | .stream() \ 300 | << TABLE.PrintTable(); \ 301 | } while (false) 302 | 303 | 304 | // Macros for use with protobuf messages 305 | // 306 | // Data is serialized via DebugString() 307 | // 308 | // Data is printed without further escaping 309 | #define LOG_PROTOBUF_VERBOSE(L, HEADING, PB_MESSAGE) \ 310 | do { \ 311 | if (LOG_VERBOSE_IS_ON(L)) \ 312 | triton::common::LogMessage( \ 313 | __FILE__, __LINE__, triton::common::Logger::Level::kINFO, HEADING, \ 314 | false) \ 315 | .stream() \ 316 | << PB_MESSAGE.DebugString(); \ 317 | } while (false) 318 | 319 | // Macros for logging errors 320 | #define LOG_STATUS_ERROR(X, MSG) \ 321 | do { \ 322 | const Status& status__ = (X); \ 323 | if (!status__.IsOk()) { \ 324 | LOG_ERROR << (MSG) << ": " << status__.AsString(); \ 325 | } \ 326 | } while (false) 327 | 328 | #define LOG_TRITONSERVER_ERROR(X, MSG) \ 329 | do { \ 330 | TRITONSERVER_Error* err__ = (X); \ 331 | if (err__ != nullptr) { \ 332 | LOG_ERROR << (MSG) << ": " << TRITONSERVER_ErrorCodeString(err__) \ 333 | << " - " << TRITONSERVER_ErrorMessage(err__); \ 334 | TRITONSERVER_ErrorDelete(err__); \ 335 | } \ 336 | } while (false) 337 | 338 | #define LOG_FLUSH triton::common::gLogger_.Flush() 339 | 340 | }} // namespace triton::common 341 | -------------------------------------------------------------------------------- /include/triton/common/model_config.h: -------------------------------------------------------------------------------- 1 | // Copyright 2018-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | #pragma once 27 | 28 | #include 29 | #include 30 | 31 | #include "model_config.pb.h" 32 | 33 | namespace triton { namespace common { 34 | 35 | /// The type for a repeated dims field (used for shape). 36 | using DimsList = ::google::protobuf::RepeatedField<::google::protobuf::int64>; 37 | 38 | /// The type for the metric_tags map. 39 | using MetricTagsMap = ::google::protobuf::Map; 40 | 41 | // Map from a host policy name to map of cmdline 42 | // settings for the host policy. 43 | using HostPolicyCmdlineConfig = std::map; 44 | using HostPolicyCmdlineConfigMap = 45 | std::unordered_map; 46 | 47 | // Map from backend name to list of setting=value pairs of cmdline 48 | // settings for the backend. 49 | using BackendCmdlineConfig = std::vector>; 50 | using BackendCmdlineConfigMap = 51 | std::unordered_map; 52 | 53 | /// The value for a dimension in a shape that indicates that that 54 | /// dimension can take on any size. 55 | constexpr int WILDCARD_DIM = -1; 56 | 57 | constexpr int SCHEDULER_DEFAULT_NICE = 5; 58 | 59 | /// Enumeration for the different platform types. 60 | enum Platform { 61 | PLATFORM_UNKNOWN = 0, 62 | PLATFORM_TENSORRT_PLAN = 1, 63 | PLATFORM_TENSORFLOW_GRAPHDEF = 2, 64 | PLATFORM_TENSORFLOW_SAVEDMODEL = 3, 65 | PLATFORM_ENSEMBLE = 4, 66 | PLATFORM_ONNXRUNTIME_ONNX = 5, 67 | PLATFORM_PYTORCH_LIBTORCH = 6 68 | }; 69 | 70 | /// Get the number of elements in a shape. 71 | /// \param dims The shape. 72 | /// \return The number of elements, or -1 if the number of elements 73 | /// cannot be determined because the shape contains one or more 74 | /// wildcard dimensions. 75 | int64_t GetElementCount(const DimsList& dims); 76 | 77 | /// Get the number of elements in a shape. 78 | /// \param dims The shape. 79 | /// \return The number of elements, or -1 if the number of elements 80 | /// cannot be determined because the shape contains one or more 81 | /// wildcard dimensions. 82 | int64_t GetElementCount(const std::vector& dims); 83 | 84 | /// Get the number of elements in the shape of a model input. 85 | /// \param mio The model input. 86 | /// \return The number of elements, or -1 if the number of elements 87 | /// cannot be determined because the shape contains one or more 88 | /// wildcard dimensions. 89 | int64_t GetElementCount(const inference::ModelInput& mio); 90 | 91 | /// Get the number of elements in the shape of a model output. 92 | /// \param mio The model output. 93 | /// \return The number of elements, or -1 if the number of elements 94 | /// cannot be determined because the shape contains one or more 95 | /// wildcard dimensions. 96 | int64_t GetElementCount(const inference::ModelOutput& mio); 97 | 98 | /// Are values of a datatype fixed-size, or variable-sized. 99 | /// \param dtype The data-type. 100 | /// \return True if datatype values are fixed-sized, false if 101 | /// variable-sized. 102 | bool IsFixedSizeDataType(const inference::DataType dtype); 103 | 104 | /// Get the size of objects of a given datatype in bytes. 105 | /// \param dtype The data-type. 106 | /// \return The size, in bytes, of objects of the datatype, or 0 if 107 | /// size cannot be determine (for example, values of type TYPE_STRING 108 | /// have variable length and so size cannot be determine just from the 109 | /// type). 110 | size_t GetDataTypeByteSize(const inference::DataType dtype); 111 | 112 | /// Get the size, in bytes, of a tensor based on datatype and 113 | /// shape. 114 | /// \param dtype The data-type. 115 | /// \param dims The shape. 116 | /// \return The size, in bytes, of the corresponding tensor, or -1 if 117 | /// unable to determine the size. 118 | int64_t GetByteSize(const inference::DataType& dtype, const DimsList& dims); 119 | 120 | /// Get the size, in bytes, of a tensor based on datatype and 121 | /// shape. 122 | /// \param dtype The data-type. 123 | /// \param dims The shape. 124 | /// \return The size, in bytes, of the corresponding tensor, or -1 if 125 | /// unable to determine the size. 126 | int64_t GetByteSize( 127 | const inference::DataType& dtype, const std::vector& dims); 128 | 129 | /// Get the size, in bytes, of a tensor based on batch-size, datatype 130 | /// and shape. A tensor that has empty shape [] and non-zero 131 | /// batch-size is sized as a tensor with shape [ batch-size ]. 132 | /// \param batch_size The batch-size. May be 0 to indicate no 133 | /// batching. 134 | /// \param dtype The data-type. 135 | /// \param dims The shape. 136 | /// \return The size, in bytes, of the corresponding tensor, or -1 if 137 | /// unable to determine the size. 138 | int64_t GetByteSize( 139 | const int batch_size, const inference::DataType& dtype, 140 | const DimsList& dims); 141 | 142 | /// Get the size, in bytes, of a tensor based on batch-size, datatype 143 | /// and shape. A tensor that has empty shape [] and non-zero 144 | /// batch-size is sized as a tensor with shape [ batch-size ]. 145 | /// \param batch_size The batch-size. May be 0 to indicate no 146 | /// batching. 147 | /// \param dtype The data-type. 148 | /// \param dims The shape. 149 | /// \return The size, in bytes, of the corresponding tensor, or -1 if 150 | /// unable to determine the size. 151 | int64_t GetByteSize( 152 | const int batch_size, const inference::DataType& dtype, 153 | const std::vector& dims); 154 | 155 | /// Get the size, in bytes, of a tensor based on ModelInput. 156 | /// \param mio The ModelInput protobuf. 157 | /// \return The size, in bytes, of the corresponding tensor, or -1 if 158 | /// unable to determine the size. 159 | int64_t GetByteSize(const inference::ModelInput& mio); 160 | 161 | /// Get the size, in bytes, of a tensor based on ModelOutput. 162 | /// \param mio The ModelOutput protobuf. 163 | /// \return The size, in bytes, of the corresponding tensor, or -1 if 164 | /// unable to determine the size. 165 | int64_t GetByteSize(const inference::ModelOutput& mio); 166 | 167 | /// Get the CPU thread nice level associate with a model 168 | /// configuration's priority. 169 | /// \param config The model configuration. 170 | /// \return The nice level. 171 | int GetCpuNiceLevel(const inference::ModelConfig& config); 172 | 173 | /// Compare two model configuration shapes for equality. Wildcard 174 | /// dimensions (that is, dimensions with size WILDCARD_DIM) are 175 | /// compared literally so that to be equal the two shapes must both 176 | /// specify WILDCARD_DIM in the same dimensions. 177 | /// \params dims0 The first shape. 178 | /// \params dims1 The second shape. 179 | /// \return True if the shapes are equal, false if not equal. 180 | bool CompareDims(const DimsList& dims0, const DimsList& dims1); 181 | 182 | /// Compare two model configuration shapes for equality. Wildcard 183 | /// dimensions (that is, dimensions with size WILDCARD_DIM) are 184 | /// compared literally so that to be equal the two shapes must both 185 | /// specify WILDCARD_DIM in the same dimensions. 186 | /// \params dims0 The first shape. 187 | /// \params dims1 The second shape. 188 | /// \return True if the shapes are equal, false if not equal. 189 | bool CompareDims( 190 | const std::vector& dims0, const std::vector& dims1); 191 | 192 | /// Compare two model configuration shapes for equality. Wildcard 193 | /// dimensions (that is, dimensions with size WILDCARD_DIM) are 194 | /// allowed to match with any value. So, a dimension in one shape 195 | /// specified as WILDCARD_DIM will always match the same dimension in 196 | /// the other shape. 197 | /// \params dims0 The first shape. 198 | /// \params dims1 The second shape. 199 | /// \return True if the shapes are equal, false if not equal. 200 | bool CompareDimsWithWildcard(const DimsList& dims0, const DimsList& dims1); 201 | 202 | /// Compare two model configuration shapes for equality. Wildcard 203 | /// dimensions (that is, dimensions with size WILDCARD_DIM) are 204 | /// allowed to match with any value. So, a dimension in one shape 205 | /// specified as WILDCARD_DIM will always match the same dimension in 206 | /// the other shape. 207 | /// \params dims0 The first shape. 208 | /// \params dims1 The second shape. 209 | /// \return True if the shapes are equal, false if not equal. 210 | bool CompareDimsWithWildcard( 211 | const DimsList& dims0, const std::vector& dims1); 212 | 213 | /// Convert a DimsList to string representation. 214 | /// \param dims The DimsList to be converted. 215 | /// \return String representation of the DimsList in pattern 216 | /// "[d0,d1,...,dn]" 217 | std::string DimsListToString(const DimsList& dims); 218 | 219 | /// Convert a vector representing a shape to string representation. 220 | /// \param dims The vector of dimensions to be converted. 221 | /// \return String representation of the vector in pattern 222 | /// "[d0,d1,...,dn]" 223 | std::string DimsListToString( 224 | const std::vector& dims, const int start_idx = 0); 225 | 226 | /// Get the server protocol string representation of a datatype. 227 | /// \param dtype The data type. 228 | /// \return The string representation. 229 | const char* DataTypeToProtocolString(const inference::DataType dtype); 230 | 231 | /// Get the datatype corresponding to a server protocol string 232 | /// representation of a datatype. 233 | /// \param dtype string representation. 234 | /// \return The data type. 235 | inference::DataType ProtocolStringToDataType(const std::string& dtype); 236 | 237 | /// Get the datatype corresponding to a server protocol string 238 | /// representation of a datatype. 239 | /// \param dtype Pointer to string. 240 | /// \param len Length of the string. 241 | /// \return The data type. 242 | inference::DataType ProtocolStringToDataType(const char* dtype, size_t len); 243 | 244 | }} // namespace triton::common 245 | -------------------------------------------------------------------------------- /include/triton/common/nvtx.h: -------------------------------------------------------------------------------- 1 | // Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | #pragma once 27 | 28 | #ifdef TRITON_ENABLE_NVTX 29 | 30 | #include 31 | 32 | namespace triton { namespace common { 33 | 34 | namespace detail { 35 | 36 | class NvtxTritonDomain { 37 | public: 38 | static nvtxDomainHandle_t& GetDomain() 39 | { 40 | static NvtxTritonDomain inst; 41 | return inst.triton_nvtx_domain_; 42 | } 43 | 44 | private: 45 | NvtxTritonDomain() { triton_nvtx_domain_ = nvtxDomainCreateA("Triton"); } 46 | 47 | ~NvtxTritonDomain() { nvtxDomainDestroy(triton_nvtx_domain_); } 48 | 49 | nvtxDomainHandle_t triton_nvtx_domain_; 50 | }; 51 | 52 | } // namespace detail 53 | 54 | // Updates a server stat with duration measured by a C++ scope. 55 | class NvtxRange { 56 | public: 57 | explicit NvtxRange(const char* label, uint32_t rgb = kNvGreen) 58 | { 59 | auto attr = GetAttributes(label, rgb); 60 | nvtxDomainRangePushEx(detail::NvtxTritonDomain::GetDomain(), &attr); 61 | } 62 | 63 | explicit NvtxRange(const std::string& label, uint32_t rgb = kNvGreen) 64 | : NvtxRange(label.c_str(), rgb) 65 | { 66 | } 67 | 68 | ~NvtxRange() { nvtxDomainRangePop(detail::NvtxTritonDomain::GetDomain()); } 69 | 70 | static constexpr uint32_t kNvGreen = 0x76b900; 71 | static constexpr uint32_t kRed = 0xc1121f; 72 | static constexpr uint32_t kGreen = 0x588157; 73 | static constexpr uint32_t kBlue = 0x023047; 74 | static constexpr uint32_t kYellow = 0xffb703; 75 | static constexpr uint32_t kOrange = 0xfb8500; 76 | 77 | private: 78 | nvtxEventAttributes_t GetAttributes(const char* label, uint32_t rgb) 79 | { 80 | nvtxEventAttributes_t attr; 81 | attr.version = NVTX_VERSION; 82 | attr.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE; 83 | attr.colorType = NVTX_COLOR_ARGB; 84 | attr.color = rgb | 0xff000000; 85 | attr.messageType = NVTX_MESSAGE_TYPE_ASCII; 86 | attr.message.ascii = label; 87 | return attr; 88 | } 89 | }; 90 | 91 | }} // namespace triton::common 92 | 93 | #endif // TRITON_ENABLE_NVTX 94 | 95 | // 96 | // Macros to access NVTX functionality. 97 | // For `NVTX_RANGE` macro please refer to the usage below. 98 | // 99 | #ifdef TRITON_ENABLE_NVTX 100 | #define NVTX_INITIALIZE nvtxInitialize(nullptr) 101 | #define NVTX_RANGE1(V, L) triton::common::NvtxRange V(L) 102 | #define NVTX_RANGE2(V, L, RGB) triton::common::NvtxRange V(L, RGB) 103 | #define NVTX_MARKER(L) nvtxMarkA(L) 104 | #else 105 | #define NVTX_INITIALIZE 106 | #define NVTX_RANGE1(V, L) 107 | #define NVTX_RANGE2(V, L, RGB) 108 | #define NVTX_MARKER(L) 109 | #endif // TRITON_ENABLE_NVTX 110 | 111 | // "Overload" for `NVTX_RANGE` macro. 112 | // Usage: 113 | // NVTX_RANGE(nvtx1, "My message") -> Records NVTX marker with kNvGreen color. 114 | // NVTX_RANGE(nvtx1, "My message", NvtxRange::kRed) -> Records NVTX marker with 115 | // kRed color. 116 | #define GET_NVTX_MACRO(_1, _2, _3, NAME, ...) NAME 117 | #define NVTX_RANGE(...) \ 118 | GET_NVTX_MACRO(__VA_ARGS__, NVTX_RANGE2, NVTX_RANGE1)(__VA_ARGS__) 119 | -------------------------------------------------------------------------------- /include/triton/common/sync_queue.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | #pragma once 27 | 28 | #include 29 | #include 30 | #include 31 | 32 | namespace triton { namespace common { 33 | 34 | // 35 | // C++11 doesn't have a sync queue so we implement a simple one. 36 | // 37 | template 38 | class SyncQueue { 39 | public: 40 | SyncQueue() {} 41 | 42 | bool Empty() 43 | { 44 | std::lock_guard lk(mu_); 45 | return queue_.empty(); 46 | } 47 | 48 | Item Get() 49 | { 50 | std::unique_lock lk(mu_); 51 | if (queue_.empty()) { 52 | cv_.wait(lk, [this] { return !queue_.empty(); }); 53 | } 54 | auto res = std::move(queue_.front()); 55 | queue_.pop_front(); 56 | return res; 57 | } 58 | 59 | void Put(const Item& value) 60 | { 61 | { 62 | std::lock_guard lk(mu_); 63 | queue_.push_back(value); 64 | } 65 | cv_.notify_all(); 66 | } 67 | 68 | void Put(Item&& value) 69 | { 70 | { 71 | std::lock_guard lk(mu_); 72 | queue_.push_back(std::move(value)); 73 | } 74 | cv_.notify_all(); 75 | } 76 | 77 | private: 78 | std::mutex mu_; 79 | std::condition_variable cv_; 80 | std::deque queue_; 81 | }; 82 | 83 | }} // namespace triton::common 84 | -------------------------------------------------------------------------------- /include/triton/common/table_printer.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | #pragma once 27 | 28 | #include 29 | #include 30 | #include 31 | #include 32 | 33 | namespace triton { namespace common { 34 | 35 | // 36 | // An ASCII table printer. 37 | // 38 | class TablePrinter { 39 | public: 40 | // Insert a row at the end of the table 41 | void InsertRow(const std::vector& row); 42 | 43 | // Print the table 44 | std::string PrintTable(); 45 | 46 | // TablePrinter will take the ownership of `headers`. 47 | TablePrinter(const std::vector& headers); 48 | 49 | private: 50 | // Update the `shares_` such that all the excess 51 | // amount of space not used a column is fairly allocated 52 | // to the other columns 53 | void FairShare(); 54 | 55 | // Append a row to `table`. This function handles the cases where a wrapping 56 | // occurs. 57 | void AddRow(std::stringstream& table, size_t row_index); 58 | 59 | // Add a row divider 60 | void AddRowDivider(std::stringstream& table); 61 | 62 | // Max row width 63 | std::vector max_widths_; 64 | 65 | // Max row height 66 | std::vector max_heights_; 67 | 68 | // A vector of vectors of vectors containing data items for every column 69 | // The record is stored in a vector of string, where each of the vector items 70 | // contains a single line from the record. For example, ["Item 1", "Item 2", 71 | // "Item 3\n Item 3 line 2"] will be stored as [["Item 1"], ["Item 2"], ["Item 72 | // 3", "Item 3 line 2"]] 73 | std::vector>> data_; 74 | 75 | // Fair share of every column 76 | std::vector shares_; 77 | }; 78 | 79 | }} // namespace triton::common 80 | -------------------------------------------------------------------------------- /include/triton/common/thread_pool.h: -------------------------------------------------------------------------------- 1 | // Copyright 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | #pragma once 27 | 28 | #include 29 | #include 30 | #include 31 | #include 32 | 33 | namespace triton { namespace common { 34 | 35 | // Generic fixed-size Thread Pool to execute tasks asynchronously 36 | 37 | class ThreadPool { 38 | public: 39 | explicit ThreadPool(std::size_t thread_count); 40 | ~ThreadPool(); 41 | ThreadPool(const ThreadPool&) = delete; 42 | ThreadPool& operator=(const ThreadPool&) = delete; 43 | 44 | using Task = std::function; 45 | // Assigns "task" to the task queue for a worker thread to execute when 46 | // available. This will not track the return value of the task. 47 | void Enqueue(Task&& task); 48 | 49 | // Returns the number of tasks waiting in the queue 50 | size_t TaskQueueSize() 51 | { 52 | std::lock_guard lk(queue_mtx_); 53 | return task_queue_.size(); 54 | } 55 | 56 | // Returns the number of threads in thread pool 57 | size_t Size() const { return workers_.size(); } 58 | 59 | private: 60 | std::queue task_queue_; 61 | std::mutex queue_mtx_; 62 | std::condition_variable cv_; 63 | std::vector workers_; 64 | // If true, tells pool to stop accepting work and tells awake worker threads 65 | // to exit when no tasks are left on the queue. 66 | bool stop_ = false; 67 | }; 68 | 69 | }} // namespace triton::common 70 | -------------------------------------------------------------------------------- /include/triton/common/triton_json.h: -------------------------------------------------------------------------------- 1 | // Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | #pragma once 27 | 28 | #ifdef _WIN32 29 | // Remove GetObject definition from windows.h, which prevents calls to 30 | // RapidJSON's GetObject. 31 | // https://github.com/Tencent/rapidjson/issues/1448 32 | #undef GetObject 33 | #include 34 | #else 35 | // Disable class-memaccess warning to facilitate compilation with gcc>7 36 | // https://github.com/Tencent/rapidjson/issues/1700 37 | #pragma GCC diagnostic push 38 | #if defined(__GNUC__) && __GNUC__ >= 8 39 | #pragma GCC diagnostic ignored "-Wclass-memaccess" 40 | #endif 41 | #include 42 | #pragma GCC diagnostic pop 43 | #endif // _WIN32 44 | 45 | #include // CrtAllocator (default) for Writer instantiation 46 | #include // UTF8 (default) for Writer instantiation 47 | #include 48 | #include 49 | #include 50 | #include 51 | #include 52 | 53 | #include 54 | #include 55 | 56 | // This header can be used both within Triton server and externally 57 | // (i.e. in source that interacts only via TRITONSERVER or 58 | // TRITONBACKEND API). Status is handled differently in these cases so 59 | // the following macros must be defined before including this 60 | // header. As an example the defines are shown here as returned by the 61 | // TRITONSERVER API. 62 | // 63 | // #define TRITONJSON_STATUSTYPE TRITONSERVER_Error* 64 | // #define TRITONJSON_STATUSRETURN(M) 65 | // return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL, (M).c_str()) 66 | // #define TRITONJSON_STATUSSUCCESS nullptr 67 | 68 | namespace triton { namespace common { 69 | 70 | // 71 | // A JSON parser/writer. Currently based on rapidjson but the intent 72 | // is to provide an abstraction for JSON functions that make it easy 73 | // to substitute a different JSON parser. Specifically for rapidjson 74 | // the class is also designed to provide safe access and error 75 | // reporting to avoid the cases where rapidjson would just abort the 76 | // entire application (!). 77 | // 78 | class TritonJson { 79 | public: 80 | class Value; 81 | enum class ValueType { 82 | OBJECT = rapidjson::kObjectType, 83 | ARRAY = rapidjson::kArrayType, 84 | }; 85 | 86 | // 87 | // Buffer used when writing JSON representation. 88 | // 89 | class WriteBuffer { 90 | public: 91 | // Get buffer base address. 92 | const char* Base() const { return buffer_.c_str(); } 93 | 94 | // Get a reference to the buffer itself. Useful to efficiently 95 | // move the contents out of the buffer. 96 | std::string& MutableContents() { return buffer_; } 97 | 98 | // Immutable contents. 99 | const std::string& Contents() const { return buffer_; } 100 | 101 | // Interface required by rapidjson::Writer 102 | typedef char Ch; 103 | void Put(char c) { buffer_.push_back(c); } 104 | void Clear() { buffer_.clear(); } 105 | void Flush() { return; } 106 | size_t Size() const { return buffer_.size(); } 107 | 108 | private: 109 | std::string buffer_; 110 | }; 111 | 112 | // 113 | // Utility to serialize input string 114 | // as a JSON string value 115 | 116 | static std::string SerializeString(const std::string& input) 117 | { 118 | WriteBuffer writebuffer; 119 | const unsigned int writeFlags = rapidjson::kWriteNanAndInfFlag; 120 | // Provide default template arguments to pass writeFlags 121 | rapidjson::Writer< 122 | WriteBuffer, rapidjson::UTF8<>, rapidjson::UTF8<>, 123 | rapidjson::CrtAllocator, writeFlags> 124 | writer(writebuffer); 125 | if (RAPIDJSON_UNLIKELY(!writer.String(input.c_str()))) { 126 | return "Error Serializing String"; 127 | } 128 | return writebuffer.Contents(); 129 | } 130 | 131 | // 132 | // Value representing the entire document or an element within a 133 | // document. 134 | // 135 | class Value { 136 | public: 137 | // Empty value. Will become a top-level Document value if 138 | // initialized by parsing or a non-top-level value if initialized 139 | // any other way. 140 | explicit Value() : value_(nullptr), allocator_(nullptr) {} 141 | 142 | // Construct a top-level JSON document. 143 | explicit Value(const ValueType type) 144 | : document_(static_cast(type)), value_(nullptr), 145 | allocator_(&document_.GetAllocator()) 146 | { 147 | } 148 | 149 | // Construct a non-top-level JSON value in a 'document'. 150 | explicit Value(TritonJson::Value& document, const ValueType type) 151 | { 152 | allocator_ = &document.document_.GetAllocator(); 153 | value_ = new (allocator_->Malloc(sizeof(rapidjson::Value))) 154 | rapidjson::Value(static_cast(type)); 155 | } 156 | 157 | // Move constructor. 158 | explicit Value(Value&& other) { *this = std::move(other); } 159 | 160 | // Move assignment operator. 161 | Value& operator=(Value&& other) 162 | { 163 | document_ = std::move(other.document_); 164 | value_ = other.value_; 165 | allocator_ = other.allocator_; 166 | other.value_ = nullptr; 167 | other.allocator_ = nullptr; 168 | return *this; 169 | } 170 | 171 | // Parse JSON into document. Can only be called on top-level 172 | // document value, otherwise error is returned. 173 | TRITONJSON_STATUSTYPE Parse(const char* base, const size_t size) 174 | { 175 | if (value_ != nullptr) { 176 | TRITONJSON_STATUSRETURN( 177 | std::string("JSON parsing only available for top-level document")); 178 | } 179 | const unsigned int parseFlags = rapidjson::kParseNanAndInfFlag; 180 | document_.Parse(base, size); 181 | if (document_.HasParseError()) { 182 | TRITONJSON_STATUSRETURN(std::string( 183 | "failed to parse the request JSON buffer: " + 184 | std::string(GetParseError_En(document_.GetParseError())) + " at " + 185 | std::to_string(document_.GetErrorOffset()))); 186 | } 187 | allocator_ = &document_.GetAllocator(); 188 | return TRITONJSON_STATUSSUCCESS; 189 | } 190 | 191 | // \see Parse(const char* base, const size_t size) 192 | TRITONJSON_STATUSTYPE Parse(const std::string& json) 193 | { 194 | return Parse(json.data(), json.size()); 195 | } 196 | 197 | // Write JSON representation into a 'buffer' in a compact 198 | // format. Can only be called for a top-level document value, 199 | // otherwise error is returned. 200 | TRITONJSON_STATUSTYPE Write(WriteBuffer* buffer) const 201 | { 202 | if (value_ != nullptr) { 203 | TRITONJSON_STATUSRETURN( 204 | std::string("JSON writing only available for top-level document")); 205 | } 206 | const unsigned int writeFlags = rapidjson::kWriteNanAndInfFlag; 207 | // Provide default template arguments to pass writeFlags 208 | rapidjson::Writer< 209 | WriteBuffer, rapidjson::UTF8<>, rapidjson::UTF8<>, 210 | rapidjson::CrtAllocator, writeFlags> 211 | writer(*buffer); 212 | if (!document_.Accept(writer)) { 213 | TRITONJSON_STATUSRETURN( 214 | std::string("Failed to accept document, invalid JSON.")); 215 | } 216 | return TRITONJSON_STATUSSUCCESS; 217 | } 218 | 219 | // Write JSON representation into a 'buffer' in an easy-to-read 220 | // format. Can only be called for a top-level document value, 221 | // otherwise error is returned. 222 | TRITONJSON_STATUSTYPE PrettyWrite(WriteBuffer* buffer) const 223 | { 224 | if (value_ != nullptr) { 225 | TRITONJSON_STATUSRETURN( 226 | std::string("JSON writing only available for top-level document")); 227 | } 228 | 229 | // Can't pass writeFlags with latest release v1.1.0 of rapidjson-dev. 230 | // We would need to build rapidjson from source to capture latest fixes. 231 | // See this issue: 232 | // https://github.com/Tencent/rapidjson/issues/905#issuecomment-370981353 233 | // PrettyWrite is only used for displaying model configs currently, so 234 | // this should not be an issue. 235 | rapidjson::PrettyWriter writer(*buffer); 236 | if (!document_.Accept(writer)) { 237 | TRITONJSON_STATUSRETURN( 238 | std::string("Failed to accept document, invalid JSON.")); 239 | } 240 | return TRITONJSON_STATUSSUCCESS; 241 | } 242 | 243 | // Swap a value with another. 244 | TRITONJSON_STATUSTYPE Swap(TritonJson::Value& other) 245 | { 246 | rapidjson::Value& value = AsMutableValue(); 247 | value.Swap(other.AsMutableValue()); 248 | return TRITONJSON_STATUSSUCCESS; 249 | } 250 | 251 | // Set/overwrite a boolean in a value. This changes the 252 | // type of the value to boolean. 253 | TRITONJSON_STATUSTYPE SetBool(const bool value) 254 | { 255 | rapidjson::Value& v = AsMutableValue(); 256 | v.SetBool(value); 257 | return TRITONJSON_STATUSSUCCESS; 258 | } 259 | 260 | // Set/overwrite a signed integer in a value. This changes the 261 | // type of the value to signed int. 262 | TRITONJSON_STATUSTYPE SetInt(const int64_t value) 263 | { 264 | rapidjson::Value& v = AsMutableValue(); 265 | v.SetInt64(value); 266 | return TRITONJSON_STATUSSUCCESS; 267 | } 268 | 269 | // Set/overwrite an unsigned integer in a value. This changes the 270 | // type of the value to unsigned int. 271 | TRITONJSON_STATUSTYPE SetUInt(const uint64_t value) 272 | { 273 | rapidjson::Value& v = AsMutableValue(); 274 | v.SetUint64(value); 275 | return TRITONJSON_STATUSSUCCESS; 276 | } 277 | 278 | // Set/overwrite a double in a value. This changes the 279 | // type of the value to double. 280 | TRITONJSON_STATUSTYPE SetDouble(const double value) 281 | { 282 | rapidjson::Value& v = AsMutableValue(); 283 | v.SetDouble(value); 284 | return TRITONJSON_STATUSSUCCESS; 285 | } 286 | 287 | // Set/overwrite a string in a value. This changes the 288 | // type of the value to string 289 | TRITONJSON_STATUSTYPE SetString(const std::string& value) 290 | { 291 | rapidjson::Value& v = AsMutableValue(); 292 | v.SetString(value.c_str(), value.length(), *allocator_); 293 | return TRITONJSON_STATUSSUCCESS; 294 | } 295 | 296 | // Set/overwrite a string member with provided name and value in this object 297 | TRITONJSON_STATUSTYPE SetStringObject( 298 | const char* name, const std::string& value) 299 | { 300 | rapidjson::Value& object = AsMutableValue(); 301 | if (!object.IsObject()) { 302 | TRITONJSON_STATUSRETURN( 303 | std::string("attempt to add/replace JSON member '") + name + 304 | "' to non-object"); 305 | } 306 | auto itr = object.FindMember(name); 307 | if (itr == object.MemberEnd()) { 308 | AddString(name, value); 309 | } else { 310 | object.RemoveMember(itr); 311 | object.AddMember( 312 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 313 | rapidjson::Value(value.c_str(), value.size(), *allocator_), 314 | *allocator_); 315 | } 316 | 317 | return TRITONJSON_STATUSSUCCESS; 318 | } 319 | 320 | // Add an array or object as a new member to this value. 'value' 321 | // is moved into this value and so on return 'value' should not be 322 | // used. It is assumed that 'name' can be used by reference, it is 323 | // the caller's responsibility to make sure the lifetime of 'name' 324 | // extends at least as long as the object. 325 | TRITONJSON_STATUSTYPE Add(const char* name, TritonJson::Value&& value) 326 | { 327 | rapidjson::Value& object = AsMutableValue(); 328 | if (!object.IsObject()) { 329 | TRITONJSON_STATUSRETURN( 330 | std::string("attempt to add JSON member '") + name + 331 | "' to non-object"); 332 | } 333 | if (value.value_ == nullptr) { 334 | rapidjson::Value v2; 335 | v2.CopyFrom(value.document_, *allocator_); 336 | object.AddMember( 337 | rapidjson::Value(rapidjson::StringRef(name)).Move(), v2.Move(), 338 | *allocator_); 339 | } else { 340 | object.AddMember( 341 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 342 | value.value_->Move(), *allocator_); 343 | } 344 | value.Release(); 345 | return TRITONJSON_STATUSSUCCESS; 346 | } 347 | 348 | // Add a copy of a string as a new member to this value. It is 349 | // assumed that 'name' can be used by reference, it is the 350 | // caller's responsibility to make sure the lifetime of 'name' 351 | // extends at least as long as the object. 352 | TRITONJSON_STATUSTYPE AddString(const char* name, const std::string& value) 353 | { 354 | rapidjson::Value& object = AsMutableValue(); 355 | if (!object.IsObject()) { 356 | TRITONJSON_STATUSRETURN( 357 | std::string("attempt to add JSON member '") + name + 358 | "' to non-object"); 359 | } 360 | object.AddMember( 361 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 362 | rapidjson::Value(value.c_str(), value.size(), *allocator_).Move(), 363 | *allocator_); 364 | return TRITONJSON_STATUSSUCCESS; 365 | } 366 | 367 | // Add a copy of a explicit-length string as a new member to this 368 | // value. It is assumed that 'name' can be used by reference, it 369 | // is the caller's responsibility to make sure the lifetime of 370 | // 'name' extends at least as long as the object. 371 | TRITONJSON_STATUSTYPE AddString( 372 | const char* name, const char* value, const size_t len) 373 | { 374 | rapidjson::Value& object = AsMutableValue(); 375 | if (!object.IsObject()) { 376 | TRITONJSON_STATUSRETURN( 377 | std::string("attempt to add JSON member '") + name + 378 | "' to non-object"); 379 | } 380 | object.AddMember( 381 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 382 | rapidjson::Value(value, len, *allocator_).Move(), *allocator_); 383 | return TRITONJSON_STATUSSUCCESS; 384 | } 385 | 386 | // Add a reference to a string as a new member to this value. It 387 | // is assumed that 'name' and 'value' can be used by reference, it 388 | // is the caller's responsibility to make sure the lifetime of 389 | // 'name' and 'value' extend at least as long as the object. 390 | TRITONJSON_STATUSTYPE AddStringRef(const char* name, const char* value) 391 | { 392 | rapidjson::Value& object = AsMutableValue(); 393 | if (!object.IsObject()) { 394 | TRITONJSON_STATUSRETURN( 395 | std::string("attempt to add JSON member '") + name + 396 | "' to non-object"); 397 | } 398 | object.AddMember( 399 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 400 | rapidjson::StringRef(value), *allocator_); 401 | return TRITONJSON_STATUSSUCCESS; 402 | } 403 | 404 | // Add a reference to a explicit-length string as a new member to 405 | // this value. It is assumed that 'name' and 'value' can be used 406 | // by reference, it is the caller's responsibility to make sure 407 | // the lifetime of 'name' and 'value' extend at least as long as 408 | // the object. 409 | TRITONJSON_STATUSTYPE AddStringRef( 410 | const char* name, const char* value, const size_t len) 411 | { 412 | rapidjson::Value& object = AsMutableValue(); 413 | if (!object.IsObject()) { 414 | TRITONJSON_STATUSRETURN( 415 | std::string("attempt to add JSON member '") + name + 416 | "' to non-object"); 417 | } 418 | object.AddMember( 419 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 420 | rapidjson::StringRef(value, len), *allocator_); 421 | return TRITONJSON_STATUSSUCCESS; 422 | } 423 | 424 | // Add a boolean new member to this value. It is assumed that 425 | // 'name' can be used by reference, it is the caller's 426 | // responsibility to make sure the lifetime of 'name' extends at 427 | // least as long as the object. 428 | TRITONJSON_STATUSTYPE AddBool(const char* name, const bool value) 429 | { 430 | rapidjson::Value& object = AsMutableValue(); 431 | if (!object.IsObject()) { 432 | TRITONJSON_STATUSRETURN( 433 | std::string("attempt to add JSON member '") + name + 434 | "' to non-object"); 435 | } 436 | object.AddMember( 437 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 438 | rapidjson::Value(value).Move(), *allocator_); 439 | return TRITONJSON_STATUSSUCCESS; 440 | } 441 | 442 | // Add a signed integer as a new member to this value. It is 443 | // assumed that 'name' can be used by reference, it is the 444 | // caller's responsibility to make sure the lifetime of 'name' 445 | // extends at least as long as the object. 446 | TRITONJSON_STATUSTYPE AddInt(const char* name, const int64_t value) 447 | { 448 | rapidjson::Value& object = AsMutableValue(); 449 | if (!object.IsObject()) { 450 | TRITONJSON_STATUSRETURN( 451 | std::string("attempt to add JSON member '") + name + 452 | "' to non-object"); 453 | } 454 | object.AddMember( 455 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 456 | rapidjson::Value(value).Move(), *allocator_); 457 | return TRITONJSON_STATUSSUCCESS; 458 | } 459 | 460 | // Add an unsigned integer as a new member to this value. It is 461 | // assumed that 'name' can be used by reference, it is the 462 | // caller's responsibility to make sure the lifetime of 'name' 463 | // extends at least as long as the object. 464 | TRITONJSON_STATUSTYPE AddUInt(const char* name, const uint64_t value) 465 | { 466 | rapidjson::Value& object = AsMutableValue(); 467 | if (!object.IsObject()) { 468 | TRITONJSON_STATUSRETURN( 469 | std::string("attempt to add JSON member '") + name + 470 | "' to non-object"); 471 | } 472 | object.AddMember( 473 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 474 | rapidjson::Value(value).Move(), *allocator_); 475 | return TRITONJSON_STATUSSUCCESS; 476 | } 477 | 478 | // Add a double as a new member to this value. It is assumed that 479 | // 'name' can be used by reference, it is the caller's 480 | // responsibility to make sure the lifetime of 'name' extends at 481 | // least as long as the object. 482 | TRITONJSON_STATUSTYPE AddDouble(const char* name, const double value) 483 | { 484 | rapidjson::Value& object = AsMutableValue(); 485 | if (!object.IsObject()) { 486 | TRITONJSON_STATUSRETURN( 487 | std::string("attempt to add JSON member '") + name + 488 | "' to non-object"); 489 | } 490 | object.AddMember( 491 | rapidjson::Value(rapidjson::StringRef(name)).Move(), 492 | rapidjson::Value(value).Move(), *allocator_); 493 | return TRITONJSON_STATUSSUCCESS; 494 | } 495 | 496 | // Append an array or object to this value, which must be an 497 | // array. 'value' is moved into this value and so on return 498 | // 'value' should not be used. 499 | TRITONJSON_STATUSTYPE Append(TritonJson::Value&& value) 500 | { 501 | rapidjson::Value& array = AsMutableValue(); 502 | if (!array.IsArray()) { 503 | TRITONJSON_STATUSRETURN( 504 | std::string("attempt to append JSON member to non-array")); 505 | } 506 | if (value.value_ == nullptr) { 507 | rapidjson::Value v2; 508 | v2.CopyFrom(value.document_, *allocator_); 509 | array.PushBack(v2.Move(), *allocator_); 510 | } else { 511 | array.PushBack(value.value_->Move(), *allocator_); 512 | } 513 | 514 | value.Release(); 515 | return TRITONJSON_STATUSSUCCESS; 516 | } 517 | 518 | // Append a copy of a string to this value, which must be an 519 | // array. 520 | TRITONJSON_STATUSTYPE AppendString(const std::string& value) 521 | { 522 | rapidjson::Value& array = AsMutableValue(); 523 | if (!array.IsArray()) { 524 | TRITONJSON_STATUSRETURN( 525 | std::string("attempt to append JSON member to non-array")); 526 | } 527 | array.PushBack( 528 | rapidjson::Value(value.c_str(), value.size(), *allocator_).Move(), 529 | *allocator_); 530 | return TRITONJSON_STATUSSUCCESS; 531 | } 532 | 533 | // Append a copy of an explicit-length string to this value, which 534 | // must be an array. 535 | TRITONJSON_STATUSTYPE AppendString(const char* value, const size_t len) 536 | { 537 | rapidjson::Value& array = AsMutableValue(); 538 | if (!array.IsArray()) { 539 | TRITONJSON_STATUSRETURN( 540 | std::string("attempt to append JSON member to non-array")); 541 | } 542 | array.PushBack( 543 | rapidjson::Value(value, len, *allocator_).Move(), *allocator_); 544 | return TRITONJSON_STATUSSUCCESS; 545 | } 546 | 547 | // Append a reference to a string to this value, which must be an 548 | // array. It is assumed that 'value' can be used by reference, it 549 | // is the caller's responsibility to make sure the lifetime of 550 | // 'value' extends at least as long as the object. 551 | TRITONJSON_STATUSTYPE AppendStringRef(const char* value) 552 | { 553 | rapidjson::Value& array = AsMutableValue(); 554 | if (!array.IsArray()) { 555 | TRITONJSON_STATUSRETURN( 556 | std::string("attempt to append JSON member to non-array")); 557 | } 558 | array.PushBack(rapidjson::StringRef(value), *allocator_); 559 | return TRITONJSON_STATUSSUCCESS; 560 | } 561 | 562 | // Append a reference to a explicit-length string to this value, 563 | // which must be an array. It is assumed that 'value' can be used 564 | // by reference, it is the caller's responsibility to make sure 565 | // the lifetime of 'value' extends at least as long as the object. 566 | TRITONJSON_STATUSTYPE AppendStringRef(const char* value, const size_t len) 567 | { 568 | rapidjson::Value& array = AsMutableValue(); 569 | if (!array.IsArray()) { 570 | TRITONJSON_STATUSRETURN( 571 | std::string("attempt to append JSON member to non-array")); 572 | } 573 | array.PushBack(rapidjson::StringRef(value, len), *allocator_); 574 | return TRITONJSON_STATUSSUCCESS; 575 | } 576 | 577 | // Append a boolean to this value, which must be an array. 578 | TRITONJSON_STATUSTYPE AppendBool(const bool value) 579 | { 580 | rapidjson::Value& array = AsMutableValue(); 581 | if (!array.IsArray()) { 582 | TRITONJSON_STATUSRETURN( 583 | std::string("attempt to append JSON member to non-array")); 584 | } 585 | 586 | array.PushBack(rapidjson::Value(value).Move(), *allocator_); 587 | return TRITONJSON_STATUSSUCCESS; 588 | } 589 | 590 | // Append a signed integer to this value, which must be an array. 591 | TRITONJSON_STATUSTYPE AppendInt(const int64_t value) 592 | { 593 | rapidjson::Value& array = AsMutableValue(); 594 | if (!array.IsArray()) { 595 | TRITONJSON_STATUSRETURN( 596 | std::string("attempt to append JSON member to non-array")); 597 | } 598 | 599 | array.PushBack(rapidjson::Value(value).Move(), *allocator_); 600 | return TRITONJSON_STATUSSUCCESS; 601 | } 602 | 603 | // Append an unsigned integer to this value, which must be an 604 | // array. 605 | TRITONJSON_STATUSTYPE AppendUInt(const uint64_t value) 606 | { 607 | rapidjson::Value& array = AsMutableValue(); 608 | if (!array.IsArray()) { 609 | TRITONJSON_STATUSRETURN( 610 | std::string("attempt to append JSON member to non-array")); 611 | } 612 | 613 | array.PushBack(rapidjson::Value(value).Move(), *allocator_); 614 | return TRITONJSON_STATUSSUCCESS; 615 | } 616 | 617 | // Append a double to this value, which must be an array. 618 | TRITONJSON_STATUSTYPE AppendDouble(const double value) 619 | { 620 | rapidjson::Value& array = AsMutableValue(); 621 | if (!array.IsArray()) { 622 | TRITONJSON_STATUSRETURN( 623 | std::string("attempt to append JSON member to non-array")); 624 | } 625 | 626 | array.PushBack(rapidjson::Value(value).Move(), *allocator_); 627 | return TRITONJSON_STATUSSUCCESS; 628 | } 629 | 630 | // Remove member from this object 631 | TRITONJSON_STATUSTYPE Remove(const char* name) 632 | { 633 | rapidjson::Value& object = AsMutableValue(); 634 | if (!object.IsObject()) { 635 | TRITONJSON_STATUSRETURN( 636 | std::string("attempt to remove JSON member '") + name + 637 | "' to non-object"); 638 | } 639 | auto itr = object.FindMember(name); 640 | if (itr != object.MemberEnd()) { 641 | object.RemoveMember(itr); 642 | } // else report success 643 | 644 | return TRITONJSON_STATUSSUCCESS; 645 | } 646 | 647 | // Check if this value is of the specified type. Return appropriate 648 | // error if not. 649 | TRITONJSON_STATUSTYPE AssertType(TritonJson::ValueType type) const 650 | { 651 | if (static_cast(type) != AsValue().GetType()) { 652 | TRITONJSON_STATUSRETURN(std::string("unexpected type")); 653 | } 654 | return TRITONJSON_STATUSSUCCESS; 655 | } 656 | 657 | // Get the size of an array. If called on non-array returns zero. 658 | size_t ArraySize() const 659 | { 660 | const rapidjson::Value& array = AsValue(); 661 | if (!array.IsArray()) { 662 | return 0; 663 | } 664 | return array.GetArray().Size(); 665 | } 666 | 667 | // Return the specified index contained in this array. 668 | TRITONJSON_STATUSTYPE At( 669 | const size_t idx, TritonJson::Value* value = nullptr) 670 | { 671 | rapidjson::Value& array = AsMutableValue(); 672 | if (!array.IsArray() || (idx >= array.GetArray().Size())) { 673 | TRITONJSON_STATUSRETURN( 674 | std::string("attempt to access non-existing array index '") + 675 | std::to_string(idx) + "'"); 676 | } 677 | *value = TritonJson::Value(array[idx], allocator_); 678 | return TRITONJSON_STATUSSUCCESS; 679 | } 680 | 681 | // Get the names of all members in an object. Error if value is 682 | // not an object. 683 | TRITONJSON_STATUSTYPE Members(std::vector* names) const 684 | { 685 | const rapidjson::Value& object = AsValue(); 686 | if (!object.IsObject()) { 687 | TRITONJSON_STATUSRETURN( 688 | std::string("attempt to get members for non-object")); 689 | } 690 | for (const auto& m : object.GetObject()) { 691 | names->push_back(m.name.GetString()); 692 | } 693 | return TRITONJSON_STATUSSUCCESS; 694 | } 695 | 696 | // Return true if this value is an object and the named member is 697 | // contained in this object. 698 | bool Find(const char* name) const 699 | { 700 | const rapidjson::Value& object = AsValue(); 701 | return object.IsObject() && object.HasMember(name); 702 | } 703 | 704 | // Return true if this value is an object and the named member is 705 | // contained in this object. Return the member in 'value'. 706 | bool Find(const char* name, TritonJson::Value* value) 707 | { 708 | rapidjson::Value& object = AsMutableValue(); 709 | if (object.IsObject() && object.HasMember(name)) { 710 | if (value != nullptr) { 711 | *value = TritonJson::Value(object[name], allocator_); 712 | } 713 | return true; 714 | } 715 | 716 | return false; 717 | } 718 | 719 | // Whether the object is null value. Note that false will also be returned 720 | // if the object is not a JSON value. 721 | bool IsNull() const { return ((value_ != nullptr) && value_->IsNull()); } 722 | 723 | // Return true if the object is an object and it has no members; 724 | // false otherwise. 725 | bool IsEmpty() const 726 | { 727 | const rapidjson::Value& object = AsValue(); 728 | if (object.IsObject() && object.MemberCount() == 0) { 729 | return true; 730 | } 731 | return false; 732 | } 733 | 734 | // Get value as a string. The string may contain null or other 735 | // special characters and so 'len' must be used to determine length. 736 | // Error if value is not a string. 737 | TRITONJSON_STATUSTYPE AsString(const char** value, size_t* len) const 738 | { 739 | if ((value_ == nullptr) || !value_->IsString()) { 740 | TRITONJSON_STATUSRETURN( 741 | std::string("attempt to access JSON non-string as string")); 742 | } 743 | *value = value_->GetString(); 744 | *len = value_->GetStringLength(); 745 | return TRITONJSON_STATUSSUCCESS; 746 | } 747 | 748 | // Get value as a string. The string may contain null or other 749 | // special characters. Error if value is not a string. 750 | TRITONJSON_STATUSTYPE AsString(std::string* str) const 751 | { 752 | if ((value_ == nullptr) || !value_->IsString()) { 753 | TRITONJSON_STATUSRETURN( 754 | std::string("attempt to access JSON non-string as string")); 755 | } 756 | str->assign(value_->GetString(), value_->GetStringLength()); 757 | return TRITONJSON_STATUSSUCCESS; 758 | } 759 | 760 | // Get value as a boolean. Error if value is not a boolean. 761 | TRITONJSON_STATUSTYPE AsBool(bool* value) const 762 | { 763 | if ((value_ == nullptr) || !value_->IsBool()) { 764 | TRITONJSON_STATUSRETURN( 765 | std::string("attempt to access JSON non-boolean as boolean")); 766 | } 767 | *value = value_->GetBool(); 768 | return TRITONJSON_STATUSSUCCESS; 769 | } 770 | 771 | // Get value as a signed integer. Error if value is not a signed 772 | // integer. 773 | TRITONJSON_STATUSTYPE AsInt(int64_t* value) const 774 | { 775 | if ((value_ == nullptr) || !value_->IsInt64()) { 776 | TRITONJSON_STATUSRETURN(std::string( 777 | "attempt to access JSON non-signed-integer as signed-integer")); 778 | } 779 | *value = value_->GetInt64(); 780 | return TRITONJSON_STATUSSUCCESS; 781 | } 782 | 783 | // Get value as an unsigned integer. Error if value is not an 784 | // unsigned integer. 785 | TRITONJSON_STATUSTYPE AsUInt(uint64_t* value) const 786 | { 787 | if ((value_ == nullptr) || !value_->IsUint64()) { 788 | TRITONJSON_STATUSRETURN(std::string( 789 | "attempt to access JSON non-unsigned-integer as unsigned-integer")); 790 | } 791 | *value = value_->GetUint64(); 792 | return TRITONJSON_STATUSSUCCESS; 793 | } 794 | 795 | // Get value as a double. Error if value is not a double. 796 | TRITONJSON_STATUSTYPE AsDouble(double* value) const 797 | { 798 | if ((value_ == nullptr) || !value_->IsNumber()) { 799 | TRITONJSON_STATUSRETURN( 800 | std::string("attempt to access JSON non-number as double")); 801 | } 802 | *value = value_->GetDouble(); 803 | return TRITONJSON_STATUSSUCCESS; 804 | } 805 | 806 | // Get named array member contained in this object. 807 | TRITONJSON_STATUSTYPE MemberAsArray( 808 | const char* name, TritonJson::Value* value) 809 | { 810 | rapidjson::Value& object = AsMutableValue(); 811 | if (!object.IsObject() || !object.HasMember(name)) { 812 | TRITONJSON_STATUSRETURN( 813 | std::string("attempt to access non-existing object member '") + 814 | name + "'"); 815 | } 816 | auto& v = object[name]; 817 | if (!v.IsArray()) { 818 | TRITONJSON_STATUSRETURN( 819 | std::string("attempt to access JSON non-array as array")); 820 | } 821 | *value = TritonJson::Value(v, allocator_); 822 | return TRITONJSON_STATUSSUCCESS; 823 | } 824 | 825 | // Get named object member contained in this object. 826 | TRITONJSON_STATUSTYPE MemberAsObject( 827 | const char* name, TritonJson::Value* value) 828 | { 829 | rapidjson::Value& object = AsMutableValue(); 830 | if (!object.IsObject() || !object.HasMember(name)) { 831 | TRITONJSON_STATUSRETURN( 832 | std::string("attempt to access non-existing object member '") + 833 | name + "'"); 834 | } 835 | auto& v = object[name]; 836 | if (!v.IsObject()) { 837 | TRITONJSON_STATUSRETURN( 838 | std::string("attempt to access JSON non-object as object")); 839 | } 840 | *value = TritonJson::Value(v, allocator_); 841 | return TRITONJSON_STATUSSUCCESS; 842 | } 843 | 844 | // Get object member as a string. The string may contain null or other 845 | // special characters and so 'len' must be used to determine length. 846 | // Error if this is not an object or if the member is not a string. 847 | TRITONJSON_STATUSTYPE MemberAsString( 848 | const char* name, const char** value, size_t* len) const 849 | { 850 | const rapidjson::Value& object = AsValue(); 851 | if (!object.IsObject() || !object.HasMember(name)) { 852 | TRITONJSON_STATUSRETURN( 853 | std::string("attempt to access non-existing object member '") + 854 | name + "'"); 855 | } 856 | const auto& v = object[name]; 857 | if (!v.IsString()) { 858 | TRITONJSON_STATUSRETURN( 859 | std::string("attempt to access JSON non-string as string")); 860 | } 861 | *value = v.GetString(); 862 | *len = v.GetStringLength(); 863 | return TRITONJSON_STATUSSUCCESS; 864 | } 865 | 866 | // Get object member as a string. The string may contain null or 867 | // other special characters. Error if this is not an object or if 868 | // the member is not a string. 869 | TRITONJSON_STATUSTYPE MemberAsString( 870 | const char* name, std::string* str) const 871 | { 872 | const rapidjson::Value& object = AsValue(); 873 | if (!object.IsObject() || !object.HasMember(name)) { 874 | TRITONJSON_STATUSRETURN( 875 | std::string("attempt to access non-existing object member '") + 876 | name + "'"); 877 | } 878 | const auto& v = object[name]; 879 | if (!v.IsString()) { 880 | TRITONJSON_STATUSRETURN( 881 | std::string("attempt to access JSON non-string as string")); 882 | } 883 | str->assign(v.GetString(), v.GetStringLength()); 884 | return TRITONJSON_STATUSSUCCESS; 885 | } 886 | 887 | // Get object member as a boolean. Error if this is not an object 888 | // or if the member is not a boolean. 889 | TRITONJSON_STATUSTYPE MemberAsBool(const char* name, bool* value) const 890 | { 891 | const rapidjson::Value& object = AsValue(); 892 | if (!object.IsObject() || !object.HasMember(name)) { 893 | TRITONJSON_STATUSRETURN( 894 | std::string("attempt to access non-existing object member '") + 895 | name + "'"); 896 | } 897 | const auto& v = object[name]; 898 | if (!v.IsBool()) { 899 | TRITONJSON_STATUSRETURN( 900 | std::string("attempt to access JSON non-boolean as boolean")); 901 | } 902 | *value = v.GetBool(); 903 | return TRITONJSON_STATUSSUCCESS; 904 | } 905 | 906 | // Get object member as a signed integer. Error if this is not an object 907 | // or if the member is not a signed integer. 908 | TRITONJSON_STATUSTYPE MemberAsInt(const char* name, int64_t* value) const 909 | { 910 | const rapidjson::Value& object = AsValue(); 911 | if (!object.IsObject() || !object.HasMember(name)) { 912 | TRITONJSON_STATUSRETURN( 913 | std::string("attempt to access non-existing object member '") + 914 | name + "'"); 915 | } 916 | const auto& v = object[name]; 917 | if (!v.IsInt64()) { 918 | TRITONJSON_STATUSRETURN(std::string( 919 | "attempt to access JSON non-signed-integer as signed-integer")); 920 | } 921 | *value = v.GetInt64(); 922 | return TRITONJSON_STATUSSUCCESS; 923 | } 924 | 925 | // Get object member as an unsigned integer. Error if this is not an object 926 | // or if the member is not an unsigned integer. 927 | TRITONJSON_STATUSTYPE MemberAsUInt(const char* name, uint64_t* value) const 928 | { 929 | const rapidjson::Value& object = AsValue(); 930 | if (!object.IsObject() || !object.HasMember(name)) { 931 | TRITONJSON_STATUSRETURN( 932 | std::string("attempt to access non-existing object member '") + 933 | name + "'"); 934 | } 935 | const auto& v = object[name]; 936 | if (!v.IsUint64()) { 937 | TRITONJSON_STATUSRETURN(std::string( 938 | "attempt to access JSON non-unsigned-integer as unsigned-integer")); 939 | } 940 | *value = v.GetUint64(); 941 | return TRITONJSON_STATUSSUCCESS; 942 | } 943 | 944 | // Get object member as a double. Error if this is not an object 945 | // or if the member is not a double. 946 | TRITONJSON_STATUSTYPE MemberAsDouble(const char* name, double* value) const 947 | { 948 | const rapidjson::Value& object = AsValue(); 949 | if (!object.IsObject() || !object.HasMember(name)) { 950 | TRITONJSON_STATUSRETURN( 951 | std::string("attempt to access non-existing object member '") + 952 | name + "'"); 953 | } 954 | const auto& v = object[name]; 955 | if (!v.IsNumber()) { 956 | TRITONJSON_STATUSRETURN( 957 | std::string("attempt to access JSON non-number as double")); 958 | } 959 | *value = v.GetDouble(); 960 | return TRITONJSON_STATUSSUCCESS; 961 | } 962 | 963 | bool IsString() 964 | { 965 | const rapidjson::Value& object = AsValue(); 966 | return object.IsString(); 967 | } 968 | 969 | bool IsArray() 970 | { 971 | const rapidjson::Value& object = AsValue(); 972 | return object.IsArray(); 973 | } 974 | 975 | bool IsObject() 976 | { 977 | const rapidjson::Value& object = AsValue(); 978 | return object.IsObject(); 979 | } 980 | 981 | bool IsInt() 982 | { 983 | const rapidjson::Value& object = AsValue(); 984 | return object.IsInt64(); 985 | } 986 | 987 | bool IsBool() 988 | { 989 | const rapidjson::Value& object = AsValue(); 990 | return object.IsBool(); 991 | } 992 | 993 | bool IsNumber() 994 | { 995 | const rapidjson::Value& object = AsValue(); 996 | return object.IsNumber(); 997 | } 998 | 999 | // Get array element at a given index within this array. 1000 | TRITONJSON_STATUSTYPE IndexAsArray( 1001 | const size_t idx, TritonJson::Value* value) 1002 | { 1003 | rapidjson::Value& array = AsMutableValue(); 1004 | if (!array.IsArray() || (idx >= array.GetArray().Size())) { 1005 | TRITONJSON_STATUSRETURN( 1006 | std::string("attempt to access non-existing array index '") + 1007 | std::to_string(idx) + "'"); 1008 | } 1009 | auto& v = array[idx]; 1010 | if (!v.IsArray()) { 1011 | TRITONJSON_STATUSRETURN( 1012 | std::string("attempt to access JSON non-array as array")); 1013 | } 1014 | *value = TritonJson::Value(v, allocator_); 1015 | return TRITONJSON_STATUSSUCCESS; 1016 | } 1017 | 1018 | // Get object element at a given index within this array. 1019 | TRITONJSON_STATUSTYPE IndexAsObject( 1020 | const size_t idx, TritonJson::Value* value) 1021 | { 1022 | rapidjson::Value& array = AsMutableValue(); 1023 | if (!array.IsArray() || (idx >= array.GetArray().Size())) { 1024 | TRITONJSON_STATUSRETURN( 1025 | std::string("attempt to access non-existing array index '") + 1026 | std::to_string(idx) + "'"); 1027 | } 1028 | auto& v = array[idx]; 1029 | if (!v.IsObject()) { 1030 | TRITONJSON_STATUSRETURN( 1031 | std::string("attempt to access JSON non-object as object")); 1032 | } 1033 | *value = TritonJson::Value(v, allocator_); 1034 | return TRITONJSON_STATUSSUCCESS; 1035 | } 1036 | 1037 | // Get array index as a string. The string may contain null or 1038 | // other special characters and so 'len' must be used to determine 1039 | // length. Error if this is not an array or if the index element 1040 | // is not a string. 1041 | TRITONJSON_STATUSTYPE IndexAsString( 1042 | const size_t idx, const char** value, size_t* len) const 1043 | { 1044 | const rapidjson::Value& array = AsValue(); 1045 | if (!array.IsArray() || (idx >= array.GetArray().Size())) { 1046 | TRITONJSON_STATUSRETURN( 1047 | std::string("attempt to access non-existing array index '") + 1048 | std::to_string(idx) + "'"); 1049 | } 1050 | const auto& v = array[idx]; 1051 | if (!v.IsString()) { 1052 | TRITONJSON_STATUSRETURN( 1053 | std::string("attempt to access JSON non-string as string")); 1054 | } 1055 | *value = v.GetString(); 1056 | *len = v.GetStringLength(); 1057 | return TRITONJSON_STATUSSUCCESS; 1058 | } 1059 | 1060 | // Get array index as a string. The string may contain null or 1061 | // other special characters. Error if this is not an array or if 1062 | // the index element is not a string. 1063 | TRITONJSON_STATUSTYPE IndexAsString( 1064 | const size_t idx, std::string* str) const 1065 | { 1066 | const rapidjson::Value& array = AsValue(); 1067 | if (!array.IsArray() || (idx >= array.GetArray().Size())) { 1068 | TRITONJSON_STATUSRETURN( 1069 | std::string("attempt to access non-existing array index '") + 1070 | std::to_string(idx) + "'"); 1071 | } 1072 | const auto& v = array[idx]; 1073 | if (!v.IsString()) { 1074 | TRITONJSON_STATUSRETURN( 1075 | std::string("attempt to access JSON non-string as string")); 1076 | } 1077 | str->assign(v.GetString(), v.GetStringLength()); 1078 | return TRITONJSON_STATUSSUCCESS; 1079 | } 1080 | 1081 | // Get array index as a boolean. Error if this is not an array or 1082 | // if the index element is not a boolean. 1083 | TRITONJSON_STATUSTYPE IndexAsBool(const size_t idx, bool* value) const 1084 | { 1085 | const rapidjson::Value& array = AsValue(); 1086 | if (!array.IsArray() || (idx >= array.GetArray().Size())) { 1087 | TRITONJSON_STATUSRETURN( 1088 | std::string("attempt to access non-existing array index '") + 1089 | std::to_string(idx) + "'"); 1090 | } 1091 | const auto& v = array[idx]; 1092 | if (!v.IsBool()) { 1093 | TRITONJSON_STATUSRETURN( 1094 | std::string("attempt to access JSON non-boolean as boolean")); 1095 | } 1096 | *value = v.GetBool(); 1097 | return TRITONJSON_STATUSSUCCESS; 1098 | } 1099 | 1100 | // Get array index as a signed integer. Error if this is not an array or 1101 | // if the index element is not a signed integer. 1102 | TRITONJSON_STATUSTYPE IndexAsInt(const size_t idx, int64_t* value) const 1103 | { 1104 | const rapidjson::Value& array = AsValue(); 1105 | if (!array.IsArray() || (idx >= array.GetArray().Size())) { 1106 | TRITONJSON_STATUSRETURN( 1107 | std::string("attempt to access non-existing array index '") + 1108 | std::to_string(idx) + "'"); 1109 | } 1110 | const auto& v = array[idx]; 1111 | if (!v.IsInt64()) { 1112 | TRITONJSON_STATUSRETURN(std::string( 1113 | "attempt to access JSON non-signed-integer as signed-integer")); 1114 | } 1115 | *value = v.GetInt64(); 1116 | return TRITONJSON_STATUSSUCCESS; 1117 | } 1118 | 1119 | // Get array index as an unsigned integer. Error if this is not an array or 1120 | // if the index element is not an unsigned integer. 1121 | TRITONJSON_STATUSTYPE IndexAsUInt(const size_t idx, uint64_t* value) const 1122 | { 1123 | const rapidjson::Value& array = AsValue(); 1124 | if (!array.IsArray() || (idx >= array.GetArray().Size())) { 1125 | TRITONJSON_STATUSRETURN( 1126 | std::string("attempt to access non-existing array index '") + 1127 | std::to_string(idx) + "'"); 1128 | } 1129 | const auto& v = array[idx]; 1130 | if (!v.IsUint64()) { 1131 | TRITONJSON_STATUSRETURN(std::string( 1132 | "attempt to access JSON non-unsigned-integer as unsigned-integer")); 1133 | } 1134 | *value = v.GetUint64(); 1135 | return TRITONJSON_STATUSSUCCESS; 1136 | } 1137 | 1138 | // Get array index as a double. Error if this is not an array or 1139 | // if the index element is not a double. 1140 | TRITONJSON_STATUSTYPE IndexAsDouble(const size_t idx, double* value) const 1141 | { 1142 | const rapidjson::Value& array = AsValue(); 1143 | if (!array.IsArray() || (idx >= array.GetArray().Size())) { 1144 | TRITONJSON_STATUSRETURN( 1145 | std::string("attempt to access non-existing array index '") + 1146 | std::to_string(idx) + "'"); 1147 | } 1148 | const auto& v = array[idx]; 1149 | if (!v.IsNumber()) { 1150 | TRITONJSON_STATUSRETURN( 1151 | std::string("attempt to access JSON non-number as double")); 1152 | } 1153 | *value = v.GetDouble(); 1154 | return TRITONJSON_STATUSSUCCESS; 1155 | } 1156 | 1157 | // Release/clear a value. 1158 | void Release() 1159 | { 1160 | if (value_ != nullptr) { 1161 | allocator_->Free(value_); 1162 | } 1163 | } 1164 | 1165 | private: 1166 | // Construct a non-top-level JSON value that references an 1167 | // existing element in a document. 1168 | explicit Value( 1169 | rapidjson::Value& v, rapidjson::Document::AllocatorType* allocator) 1170 | : value_(&v), allocator_(allocator) 1171 | { 1172 | } 1173 | 1174 | // Return a value object that can be used for both a top-level 1175 | // document as well as an element within a document. 1176 | const rapidjson::Value& AsValue() const 1177 | { 1178 | if (value_ == nullptr) { 1179 | return document_; 1180 | } 1181 | return *value_; 1182 | } 1183 | 1184 | rapidjson::Value& AsMutableValue() 1185 | { 1186 | if (value_ == nullptr) { 1187 | return document_; 1188 | } 1189 | return *value_; 1190 | } 1191 | 1192 | // If this object a document or value. Based on this only one or 1193 | // document_ or value_ is valid. 1194 | rapidjson::Document document_; 1195 | rapidjson::Value* value_; 1196 | rapidjson::Document::AllocatorType* allocator_; 1197 | }; 1198 | }; 1199 | 1200 | }} // namespace triton::common 1201 | -------------------------------------------------------------------------------- /protobuf/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Redistribution and use in source and binary forms, with or without 4 | # modification, are permitted provided that the following conditions 5 | # are met: 6 | # * Redistributions of source code must retain the above copyright 7 | # notice, this list of conditions and the following disclaimer. 8 | # * Redistributions in binary form must reproduce the above copyright 9 | # notice, this list of conditions and the following disclaimer in the 10 | # documentation and/or other materials provided with the distribution. 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its 12 | # contributors may be used to endorse or promote products derived 13 | # from this software without specific prior written permission. 14 | # 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | cmake_minimum_required (VERSION 3.18) 28 | 29 | set(protobuf_MODULE_COMPATIBLE TRUE CACHE BOOL "protobuf_MODULE_COMPATIBLE" FORCE) 30 | find_package(Protobuf CONFIG REQUIRED) 31 | message(STATUS "Using protobuf ${Protobuf_VERSION}") 32 | include_directories(${Protobuf_INCLUDE_DIRS}) 33 | 34 | if(${TRITON_COMMON_ENABLE_GRPC}) 35 | find_package(gRPC CONFIG REQUIRED) 36 | message(STATUS "Using gRPC ${gRPC_VERSION}") 37 | include_directories($) 38 | endif() 39 | 40 | # 41 | # Protobuf 42 | # 43 | if(${TRITON_COMMON_ENABLE_PROTOBUF}) 44 | file(GLOB proto-srcs *.proto) 45 | protobuf_generate_cpp(PROTO_SRCS PROTO_HDRS ${proto-srcs}) 46 | if(TRITON_COMMON_ENABLE_PROTOBUF_PYTHON) 47 | protobuf_generate_python(PROTO_PY ${proto-srcs}) 48 | endif() 49 | 50 | set(_PROTOBUF_PROTOC $) 51 | if(${TRITON_COMMON_ENABLE_GRPC}) 52 | set(_GRPC_CPP_PLUGIN_EXECUTABLE $) 53 | endif() 54 | 55 | add_library( 56 | proto-library EXCLUDE_FROM_ALL OBJECT 57 | ${PROTO_SRCS} ${PROTO_HDRS} 58 | ) 59 | 60 | target_include_directories( 61 | proto-library 62 | PUBLIC 63 | $ 64 | ) 65 | 66 | target_link_libraries( 67 | proto-library 68 | PRIVATE 69 | common-compile-settings 70 | ) 71 | 72 | set_target_properties( 73 | proto-library 74 | PROPERTIES 75 | POSITION_INDEPENDENT_CODE ON 76 | ) 77 | 78 | if(TRITON_COMMON_ENABLE_PROTOBUF_PYTHON) 79 | add_custom_target(proto-py-library DEPENDS ${PROTO_PY}) 80 | endif() 81 | 82 | install( 83 | FILES 84 | ${PROTO_HDRS} 85 | DESTINATION include 86 | ) 87 | endif() 88 | 89 | # 90 | # GRPC 91 | # 92 | if(${TRITON_COMMON_ENABLE_GRPC}) 93 | get_filename_component(grpc_service_proto_abspath "grpc_service.proto" ABSOLUTE) 94 | get_filename_component(grpc_service_proto_dir "${grpc_service_proto_abspath}" PATH) 95 | set(GRPC_SRCS "grpc_service.grpc.pb.cc") 96 | set(GRPC_HDRS "grpc_service.grpc.pb.h") 97 | if(TRITON_COMMON_ENABLE_PROTOBUF_PYTHON) 98 | set(GRPC_PY "grpc_service.grpc.py") 99 | endif() 100 | 101 | add_custom_command( 102 | OUTPUT "${GRPC_SRCS}" "${GRPC_HDRS}" 103 | COMMAND ${_PROTOBUF_PROTOC} 104 | ARGS 105 | --grpc_out "${CMAKE_CURRENT_BINARY_DIR}" 106 | --cpp_out "${CMAKE_CURRENT_BINARY_DIR}" 107 | -I "${grpc_service_proto_dir}" 108 | --plugin=protoc-gen-grpc="${_GRPC_CPP_PLUGIN_EXECUTABLE}" 109 | "grpc_service.proto" 110 | DEPENDS "grpc_service.proto" proto-library 111 | ) 112 | 113 | if(TRITON_COMMON_ENABLE_PROTOBUF_PYTHON) 114 | find_package(Python REQUIRED COMPONENTS Interpreter) 115 | add_custom_command( 116 | OUTPUT "${GRPC_PY}" 117 | COMMAND ${Python_EXECUTABLE} 118 | ARGS 119 | -m grpc_tools.protoc 120 | -I "${grpc_service_proto_dir}" 121 | --grpc_python_out "${CMAKE_CURRENT_BINARY_DIR}" 122 | "grpc_service.proto" 123 | DEPENDS "grpc_service.proto" proto-library 124 | ) 125 | endif() 126 | 127 | add_library( 128 | grpc-service-library EXCLUDE_FROM_ALL OBJECT 129 | ${GRPC_SRCS} ${GRPC_HDRS} 130 | ) 131 | 132 | target_include_directories( 133 | grpc-service-library 134 | PUBLIC 135 | $ 136 | ) 137 | 138 | target_link_libraries( 139 | grpc-service-library 140 | PRIVATE 141 | common-compile-settings 142 | ) 143 | 144 | set_target_properties( 145 | grpc-service-library 146 | PROPERTIES 147 | POSITION_INDEPENDENT_CODE ON 148 | ) 149 | 150 | if(TRITON_COMMON_ENABLE_PROTOBUF_PYTHON) 151 | add_custom_target(grpc-service-py-library DEPENDS ${GRPC_PY}) 152 | endif() 153 | 154 | install( 155 | FILES 156 | ${CMAKE_CURRENT_BINARY_DIR}/grpc_service.grpc.pb.h 157 | DESTINATION include 158 | OPTIONAL 159 | ) 160 | endif() 161 | 162 | # 163 | # GRPC Health Service 164 | # 165 | if(${TRITON_COMMON_ENABLE_GRPC}) 166 | get_filename_component(health_proto_abspath "health.proto" ABSOLUTE) 167 | get_filename_component(health_proto_dir "${health_proto_abspath}" PATH) 168 | set(HEALTH_SRCS "health.grpc.pb.cc") 169 | set(HEALTH_HDRS "health.grpc.pb.h") 170 | 171 | add_custom_command( 172 | OUTPUT "${HEALTH_SRCS}" "${HEALTH_HDRS}" 173 | COMMAND ${_PROTOBUF_PROTOC} 174 | ARGS 175 | --grpc_out "${CMAKE_CURRENT_BINARY_DIR}" 176 | --cpp_out "${CMAKE_CURRENT_BINARY_DIR}" 177 | -I "${health_proto_dir}" 178 | --plugin=protoc-gen-grpc="${_GRPC_CPP_PLUGIN_EXECUTABLE}" 179 | "health.proto" 180 | DEPENDS "health.proto" proto-library 181 | ) 182 | 183 | add_library( 184 | grpc-health-library EXCLUDE_FROM_ALL OBJECT 185 | ${HEALTH_SRCS} ${HEALTH_HDRS} 186 | ) 187 | 188 | target_include_directories( 189 | grpc-health-library 190 | PUBLIC 191 | $ 192 | ) 193 | 194 | target_link_libraries( 195 | grpc-health-library 196 | PRIVATE 197 | common-compile-settings 198 | ) 199 | 200 | set_target_properties( 201 | grpc-health-library 202 | PROPERTIES 203 | POSITION_INDEPENDENT_CODE ON 204 | ) 205 | 206 | install( 207 | FILES 208 | ${CMAKE_CURRENT_BINARY_DIR}/health.grpc.pb.h 209 | DESTINATION include 210 | OPTIONAL 211 | ) 212 | endif() 213 | -------------------------------------------------------------------------------- /protobuf/health.proto: -------------------------------------------------------------------------------- 1 | // Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | syntax = "proto3"; 28 | 29 | package grpc.health.v1; 30 | 31 | //@@.. cpp:namespace:: grpc.health.v1 32 | 33 | //@@ 34 | //@@.. cpp:var:: message HealthCheckRequest 35 | //@@ 36 | //@@ Request message for HealthCheck 37 | //@@ 38 | message HealthCheckRequest 39 | { 40 | string service = 1; 41 | } 42 | 43 | //@@ 44 | //@@.. cpp:var:: message HealthCheckResponse 45 | //@@ 46 | //@@ Response message for HealthCheck 47 | //@@ 48 | message HealthCheckResponse 49 | { 50 | //@@ 51 | //@@.. cpp:enum:: ServingStatus 52 | //@@ 53 | //@@ Statuses supported by GRPC's health check. 54 | //@@ 55 | enum ServingStatus { 56 | UNKNOWN = 0; 57 | SERVING = 1; 58 | NOT_SERVING = 2; 59 | SERVICE_UNKNOWN = 3; 60 | } 61 | ServingStatus status = 1; 62 | } 63 | 64 | //@@ 65 | //@@.. cpp:var:: service Health 66 | //@@ 67 | //@@ Health service for GRPC endpoints. 68 | //@@ 69 | service Health 70 | { 71 | //@@ .. cpp:var:: rpc Check(HealthCheckRequest) returns 72 | //@@ (HealthCheckResponse) 73 | //@@ 74 | //@@ Get serving status of the inference server. 75 | //@@ 76 | rpc Check(HealthCheckRequest) returns (HealthCheckResponse); 77 | } 78 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | # Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # Redistribution and use in source and binary forms, with or without 4 | # modification, are permitted provided that the following conditions 5 | # are met: 6 | # * Redistributions of source code must retain the above copyright 7 | # notice, this list of conditions and the following disclaimer. 8 | # * Redistributions in binary form must reproduce the above copyright 9 | # notice, this list of conditions and the following disclaimer in the 10 | # documentation and/or other materials provided with the distribution. 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its 12 | # contributors may be used to endorse or promote products derived 13 | # from this software without specific prior written permission. 14 | # 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | [tool.codespell] 28 | # note: pre-commit passes explicit lists of files here, which this skip file list doesn't override - 29 | # this is only to allow you to run codespell interactively 30 | skip = "./.git,./.github" 31 | # ignore short words, and typename parameters like OffsetT 32 | ignore-regex = "\\b(.{1,4}|[A-Z]\\w*T)\\b" 33 | # use the 'clear' dictionary for unambiguous spelling mistakes 34 | builtin = "clear" 35 | # disable warnings about binary files and wrong encoding 36 | quiet-level = 3 37 | 38 | [tool.isort] 39 | profile = "black" 40 | use_parentheses = true 41 | multi_line_output = 3 42 | include_trailing_comma = true 43 | force_grid_wrap = 0 44 | ensure_newline_before_comments = true 45 | line_length = 88 46 | balanced_wrapping = true 47 | indent = " " 48 | skip = ["build"] 49 | -------------------------------------------------------------------------------- /src/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # Redistribution and use in source and binary forms, with or without 4 | # modification, are permitted provided that the following conditions 5 | # are met: 6 | # * Redistributions of source code must retain the above copyright 7 | # notice, this list of conditions and the following disclaimer. 8 | # * Redistributions in binary form must reproduce the above copyright 9 | # notice, this list of conditions and the following disclaimer in the 10 | # documentation and/or other materials provided with the distribution. 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its 12 | # contributors may be used to endorse or promote products derived 13 | # from this software without specific prior written permission. 14 | # 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | 28 | # 29 | # Error 30 | # 31 | add_library( 32 | triton-common-error 33 | error.cc 34 | ) 35 | 36 | add_library( 37 | TritonCommon::triton-common-error ALIAS triton-common-error 38 | ) 39 | 40 | target_include_directories( 41 | triton-common-error 42 | PUBLIC 43 | $ 44 | $ 45 | PRIVATE 46 | ${CMAKE_CURRENT_SOURCE_DIR} 47 | ) 48 | 49 | target_link_libraries(triton-common-error PRIVATE common-compile-settings) 50 | 51 | # 52 | # Logging 53 | # 54 | add_library( 55 | triton-common-logging 56 | logging.cc 57 | ) 58 | 59 | add_library( 60 | TritonCommon::triton-common-logging ALIAS triton-common-logging 61 | ) 62 | 63 | target_include_directories( 64 | triton-common-logging 65 | PUBLIC 66 | $ 67 | $ 68 | PRIVATE 69 | ${RAPIDJSON_INCLUDE_DIRS} 70 | ${CMAKE_CURRENT_SOURCE_DIR} 71 | ) 72 | 73 | if(${TRITON_ENABLE_LOGGING}) 74 | target_compile_definitions( 75 | triton-common-logging 76 | PRIVATE TRITON_ENABLE_LOGGING=1 77 | ) 78 | endif() # TRITON_ENABLE_LOGGING 79 | 80 | target_link_libraries(triton-common-logging PRIVATE common-compile-settings) 81 | 82 | # 83 | # Async Work Queue 84 | # 85 | add_library( 86 | triton-common-async-work-queue 87 | async_work_queue.cc 88 | error.cc 89 | thread_pool.cc 90 | ) 91 | 92 | add_library( 93 | TritonCommon::triton-common-async-work-queue ALIAS triton-common-async-work-queue 94 | ) 95 | 96 | target_include_directories( 97 | triton-common-async-work-queue 98 | PUBLIC 99 | $ 100 | $ 101 | PRIVATE 102 | ${CMAKE_CURRENT_SOURCE_DIR} 103 | ) 104 | 105 | target_link_libraries(triton-common-async-work-queue 106 | PUBLIC 107 | Threads::Threads 108 | PRIVATE 109 | common-compile-settings 110 | ) 111 | 112 | # 113 | # Thread Pool 114 | # 115 | add_library( 116 | triton-common-thread-pool 117 | thread_pool.cc 118 | ) 119 | 120 | add_library( 121 | TritonCommon::triton-common-thread-pool ALIAS triton-common-thread-pool 122 | ) 123 | 124 | target_include_directories( 125 | triton-common-thread-pool 126 | PUBLIC 127 | $ 128 | $ 129 | PRIVATE 130 | ${CMAKE_CURRENT_SOURCE_DIR} 131 | ) 132 | 133 | target_link_libraries(triton-common-thread-pool 134 | PUBLIC 135 | Threads::Threads 136 | PRIVATE 137 | common-compile-settings 138 | ) 139 | 140 | # 141 | # Table Printer 142 | # 143 | add_library( 144 | triton-common-table-printer 145 | table_printer.cc 146 | ) 147 | 148 | add_library( 149 | TritonBackend::triton-common-table-printer ALIAS triton-common-table-printer 150 | ) 151 | 152 | target_include_directories( 153 | triton-common-table-printer 154 | PUBLIC 155 | $ 156 | $ 157 | PRIVATE 158 | ${CMAKE_CURRENT_SOURCE_DIR}/src 159 | ) 160 | 161 | target_link_libraries(triton-common-table-printer PRIVATE common-compile-settings) 162 | 163 | set_target_properties( 164 | triton-common-async-work-queue 165 | triton-common-error 166 | triton-common-logging 167 | triton-common-table-printer 168 | triton-common-thread-pool 169 | PROPERTIES 170 | WINDOWS_EXPORT_ALL_SYMBOLS TRUE 171 | POSITION_INDEPENDENT_CODE ON 172 | ) 173 | 174 | set_target_properties( 175 | triton-common-async-work-queue 176 | PROPERTIES 177 | OUTPUT_NAME tritonasyncworkqueue 178 | ) 179 | 180 | set_target_properties( 181 | triton-common-thread-pool 182 | PROPERTIES 183 | OUTPUT_NAME tritonthreadpool 184 | ) 185 | 186 | set_target_properties( 187 | triton-common-error 188 | PROPERTIES 189 | OUTPUT_NAME tritoncommonerror 190 | ) 191 | 192 | set_target_properties( 193 | triton-common-logging 194 | PROPERTIES 195 | OUTPUT_NAME tritoncommonlogging 196 | ) 197 | 198 | set_target_properties( 199 | triton-common-table-printer 200 | PROPERTIES 201 | OUTPUT_NAME tritontableprinter 202 | ) 203 | 204 | add_subdirectory(test) 205 | -------------------------------------------------------------------------------- /src/async_work_queue.cc: -------------------------------------------------------------------------------- 1 | // Copyright 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | #include "triton/common/async_work_queue.h" 28 | 29 | namespace triton { namespace common { 30 | 31 | AsyncWorkQueue::~AsyncWorkQueue() 32 | { 33 | GetSingleton()->thread_pool_.reset(); 34 | } 35 | 36 | AsyncWorkQueue* 37 | AsyncWorkQueue::GetSingleton() 38 | { 39 | static AsyncWorkQueue singleton; 40 | return &singleton; 41 | } 42 | 43 | Error 44 | AsyncWorkQueue::Initialize(size_t worker_count) 45 | { 46 | if (worker_count < 1) { 47 | return Error( 48 | Error::Code::INVALID_ARG, 49 | "Async work queue must be initialized with positive 'worker_count'"); 50 | } 51 | 52 | static std::mutex init_mtx; 53 | std::lock_guard lk(init_mtx); 54 | 55 | if (GetSingleton()->thread_pool_) { 56 | return Error( 57 | Error::Code::ALREADY_EXISTS, 58 | "Async work queue has been initialized with " + 59 | std::to_string(GetSingleton()->thread_pool_->Size()) + 60 | " 'worker_count'"); 61 | } 62 | 63 | GetSingleton()->thread_pool_.reset(new ThreadPool(worker_count)); 64 | return Error::Success; 65 | } 66 | 67 | size_t 68 | AsyncWorkQueue::WorkerCount() 69 | { 70 | if (!GetSingleton()->thread_pool_) { 71 | return 0; 72 | } 73 | return GetSingleton()->thread_pool_->Size(); 74 | } 75 | 76 | Error 77 | AsyncWorkQueue::AddTask(std::function&& task) 78 | { 79 | if (!GetSingleton()->thread_pool_) { 80 | return Error( 81 | Error::Code::UNAVAILABLE, 82 | "Async work queue must be initialized before adding task"); 83 | } 84 | GetSingleton()->thread_pool_->Enqueue(std::move(task)); 85 | 86 | return Error::Success; 87 | } 88 | 89 | void 90 | AsyncWorkQueue::Reset() 91 | { 92 | // Reconstruct the singleton to reset it 93 | GetSingleton()->~AsyncWorkQueue(); 94 | new (GetSingleton()) AsyncWorkQueue(); 95 | } 96 | 97 | }} // namespace triton::common 98 | -------------------------------------------------------------------------------- /src/error.cc: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | #include "triton/common/error.h" 28 | 29 | namespace triton { namespace common { 30 | 31 | const Error Error::Success(Error::Code::SUCCESS); 32 | 33 | std::string 34 | Error::AsString() const 35 | { 36 | std::string str(CodeString(code_)); 37 | str += ": " + msg_; 38 | return str; 39 | } 40 | 41 | const char* 42 | Error::CodeString(const Code code) 43 | { 44 | switch (code) { 45 | case Error::Code::SUCCESS: 46 | return "OK"; 47 | case Error::Code::UNKNOWN: 48 | return "Unknown"; 49 | case Error::Code::INTERNAL: 50 | return "Internal"; 51 | case Error::Code::NOT_FOUND: 52 | return "Not found"; 53 | case Error::Code::INVALID_ARG: 54 | return "Invalid argument"; 55 | case Error::Code::UNAVAILABLE: 56 | return "Unavailable"; 57 | case Error::Code::UNSUPPORTED: 58 | return "Unsupported"; 59 | case Error::Code::ALREADY_EXISTS: 60 | return "Already exists"; 61 | default: 62 | break; 63 | } 64 | 65 | return ""; 66 | } 67 | 68 | }} // namespace triton::common 69 | -------------------------------------------------------------------------------- /src/logging.cc: -------------------------------------------------------------------------------- 1 | // Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | #include 28 | #include 29 | #include 30 | 31 | // Defined but not used 32 | #define TRITONJSON_STATUSTYPE uint8_t 33 | #define TRITONJSON_STATUSRETURN(M) 34 | #define TRITONJSON_STATUSSUCCESS 0 35 | 36 | #include "triton/common/logging.h" 37 | #include "triton/common/triton_json.h" 38 | 39 | namespace triton { namespace common { 40 | 41 | Logger gLogger_; 42 | 43 | Logger::Logger() 44 | : enables_{true, true, true}, vlevel_(0), format_(Format::kDEFAULT) 45 | { 46 | const char* value = std::getenv(Logger::ESCAPE_ENVIRONMENT_VARIABLE); 47 | escape_log_messages_ = (value && std::strcmp(value, "0") == 0) ? false : true; 48 | } 49 | 50 | void 51 | Logger::Log(const std::string& msg, const Level level) 52 | { 53 | const std::lock_guard lock(mutex_); 54 | if (file_stream_.is_open()) { 55 | file_stream_ << msg << std::endl; 56 | } else if (level == Level::kINFO) { 57 | std::cout << msg << std::endl; 58 | } else { // kWARNING or kERROR 59 | std::cerr << msg << std::endl; 60 | } 61 | } 62 | 63 | void 64 | Logger::Flush() 65 | { 66 | std::cerr << std::flush; 67 | } 68 | 69 | #ifdef _WIN32 70 | 71 | void 72 | LogMessage::LogTimestamp(std::stringstream& stream) 73 | { 74 | switch (gLogger_.LogFormat()) { 75 | case Logger::Format::kDEFAULT: { 76 | stream << std::setfill('0') << std::setw(2) << timestamp_.wMonth 77 | << std::setw(2) << timestamp_.wDay << ' ' << std::setw(2) 78 | << timestamp_.wHour << ':' << std::setw(2) << timestamp_.wMinute 79 | << ':' << std::setw(2) << timestamp_.wSecond << '.' << std::setw(6) 80 | << timestamp_.wMilliseconds * 1000; 81 | break; 82 | } 83 | case Logger::Format::kISO8601: { 84 | stream << timestamp_.wYear << '-' << std::setfill('0') << std::setw(2) 85 | << timestamp_.wMonth << '-' << std::setw(2) << timestamp_.wDay 86 | << 'T' << std::setw(2) << timestamp_.wHour << ':' << std::setw(2) 87 | << timestamp_.wMinute << ':' << std::setw(2) << timestamp_.wSecond 88 | << "Z"; 89 | break; 90 | } 91 | } 92 | } 93 | #else 94 | void 95 | LogMessage::LogTimestamp(std::stringstream& stream) 96 | { 97 | struct tm tm_time; 98 | gmtime_r(((time_t*)&(timestamp_.tv_sec)), &tm_time); 99 | 100 | switch (gLogger_.LogFormat()) { 101 | case Logger::Format::kDEFAULT: { 102 | stream << std::setfill('0') << std::setw(2) << (tm_time.tm_mon + 1) 103 | << std::setw(2) << tm_time.tm_mday << ' ' << std::setw(2) 104 | << tm_time.tm_hour << ':' << std::setw(2) << tm_time.tm_min << ':' 105 | << std::setw(2) << tm_time.tm_sec << '.' << std::setw(6) 106 | << timestamp_.tv_usec; 107 | break; 108 | } 109 | case Logger::Format::kISO8601: { 110 | stream << (tm_time.tm_year + 1900) << '-' << std::setfill('0') 111 | << std::setw(2) << (tm_time.tm_mon + 1) << '-' << std::setw(2) 112 | << tm_time.tm_mday << 'T' << std::setw(2) << tm_time.tm_hour << ':' 113 | << std::setw(2) << tm_time.tm_min << ':' << std::setw(2) 114 | << tm_time.tm_sec << "Z"; 115 | break; 116 | } 117 | } 118 | } 119 | 120 | #endif 121 | 122 | void 123 | LogMessage::LogPreamble(std::stringstream& stream) 124 | { 125 | switch (gLogger_.LogFormat()) { 126 | case Logger::Format::kDEFAULT: { 127 | stream << Logger::LEVEL_NAMES[static_cast(level_)]; 128 | LogTimestamp(stream); 129 | stream << ' ' << pid_ << ' ' << path_ << ':' << line_ << "] "; 130 | 131 | break; 132 | } 133 | case Logger::Format::kISO8601: { 134 | LogTimestamp(stream); 135 | stream << " " << Logger::LEVEL_NAMES[static_cast(level_)] << ' ' 136 | << pid_ << ' ' << path_ << ':' << line_ << "] "; 137 | break; 138 | } 139 | } 140 | } 141 | 142 | 143 | LogMessage::~LogMessage() 144 | { 145 | std::stringstream log_record; 146 | LogPreamble(log_record); 147 | std::string escaped_message = 148 | escape_log_messages_ ? TritonJson::SerializeString(message_.str()) 149 | : message_.str(); 150 | if (heading_ != nullptr) { 151 | std::string escaped_heading = gLogger_.EscapeLogMessages() 152 | ? TritonJson::SerializeString(heading_) 153 | : heading_; 154 | log_record << escaped_heading << '\n'; 155 | } 156 | log_record << escaped_message; 157 | gLogger_.Log(log_record.str(), level_); 158 | } 159 | 160 | }} // namespace triton::common 161 | -------------------------------------------------------------------------------- /src/model_config.cc: -------------------------------------------------------------------------------- 1 | // Copyright 2018-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | #include "triton/common/model_config.h" 28 | 29 | namespace triton { namespace common { 30 | 31 | bool 32 | IsFixedSizeDataType(const inference::DataType dtype) 33 | { 34 | return dtype != inference::DataType::TYPE_STRING; 35 | } 36 | 37 | size_t 38 | GetDataTypeByteSize(const inference::DataType dtype) 39 | { 40 | switch (dtype) { 41 | case inference::DataType::TYPE_BOOL: 42 | return 1; 43 | case inference::DataType::TYPE_UINT8: 44 | return 1; 45 | case inference::DataType::TYPE_UINT16: 46 | return 2; 47 | case inference::DataType::TYPE_UINT32: 48 | return 4; 49 | case inference::DataType::TYPE_UINT64: 50 | return 8; 51 | case inference::DataType::TYPE_INT8: 52 | return 1; 53 | case inference::DataType::TYPE_INT16: 54 | return 2; 55 | case inference::DataType::TYPE_INT32: 56 | return 4; 57 | case inference::DataType::TYPE_INT64: 58 | return 8; 59 | case inference::DataType::TYPE_FP16: 60 | return 2; 61 | case inference::DataType::TYPE_FP32: 62 | return 4; 63 | case inference::DataType::TYPE_FP64: 64 | return 8; 65 | case inference::DataType::TYPE_STRING: 66 | return 0; 67 | case inference::DataType::TYPE_BF16: 68 | return 2; 69 | default: 70 | break; 71 | } 72 | 73 | return 0; 74 | } 75 | 76 | int64_t 77 | GetElementCount(const DimsList& dims) 78 | { 79 | bool first = true; 80 | int64_t cnt = 0; 81 | for (auto dim : dims) { 82 | if (dim == WILDCARD_DIM) { 83 | return -1; 84 | } 85 | 86 | if (first) { 87 | cnt = dim; 88 | first = false; 89 | } else { 90 | cnt *= dim; 91 | } 92 | } 93 | 94 | return cnt; 95 | } 96 | 97 | int64_t 98 | GetElementCount(const std::vector& dims) 99 | { 100 | bool first = true; 101 | int64_t cnt = 0; 102 | for (auto dim : dims) { 103 | if (dim == WILDCARD_DIM) { 104 | return -1; 105 | } 106 | 107 | if (first) { 108 | cnt = dim; 109 | first = false; 110 | } else { 111 | cnt *= dim; 112 | } 113 | } 114 | 115 | return cnt; 116 | } 117 | 118 | int64_t 119 | GetElementCount(const inference::ModelInput& mio) 120 | { 121 | return GetElementCount(mio.dims()); 122 | } 123 | 124 | int64_t 125 | GetElementCount(const inference::ModelOutput& mio) 126 | { 127 | return GetElementCount(mio.dims()); 128 | } 129 | 130 | int64_t 131 | GetByteSize(const inference::DataType& dtype, const DimsList& dims) 132 | { 133 | size_t dt_size = GetDataTypeByteSize(dtype); 134 | if (dt_size == 0) { 135 | return -1; 136 | } 137 | 138 | int64_t cnt = GetElementCount(dims); 139 | if (cnt == -1) { 140 | return -1; 141 | } 142 | 143 | return cnt * dt_size; 144 | } 145 | 146 | int64_t 147 | GetByteSize(const inference::DataType& dtype, const std::vector& dims) 148 | { 149 | size_t dt_size = GetDataTypeByteSize(dtype); 150 | if (dt_size == 0) { 151 | return -1; 152 | } 153 | 154 | int64_t cnt = GetElementCount(dims); 155 | if (cnt == -1) { 156 | return -1; 157 | } 158 | 159 | return cnt * dt_size; 160 | } 161 | 162 | int64_t 163 | GetByteSize( 164 | const int batch_size, const inference::DataType& dtype, 165 | const DimsList& dims) 166 | { 167 | if (dims.size() == 0) { 168 | return batch_size * GetDataTypeByteSize(dtype); 169 | } 170 | 171 | int64_t bs = GetByteSize(dtype, dims); 172 | if (bs == -1) { 173 | return -1; 174 | } 175 | 176 | return std::max(1, batch_size) * bs; 177 | } 178 | 179 | int64_t 180 | GetByteSize( 181 | const int batch_size, const inference::DataType& dtype, 182 | const std::vector& dims) 183 | { 184 | if (dims.size() == 0) { 185 | return batch_size * GetDataTypeByteSize(dtype); 186 | } 187 | 188 | int64_t bs = GetByteSize(dtype, dims); 189 | if (bs == -1) { 190 | return -1; 191 | } 192 | 193 | return std::max(1, batch_size) * bs; 194 | } 195 | 196 | int64_t 197 | GetByteSize(const inference::ModelInput& mio) 198 | { 199 | return GetByteSize(mio.data_type(), mio.dims()); 200 | } 201 | 202 | int64_t 203 | GetByteSize(const inference::ModelOutput& mio) 204 | { 205 | return GetByteSize(mio.data_type(), mio.dims()); 206 | } 207 | 208 | int 209 | GetCpuNiceLevel(const inference::ModelConfig& config) 210 | { 211 | int nice = SCHEDULER_DEFAULT_NICE; 212 | if (config.has_optimization()) { 213 | switch (config.optimization().priority()) { 214 | case inference::ModelOptimizationPolicy::PRIORITY_MAX: 215 | nice = 0; 216 | break; 217 | case inference::ModelOptimizationPolicy::PRIORITY_MIN: 218 | nice = 19; 219 | break; 220 | default: 221 | nice = SCHEDULER_DEFAULT_NICE; 222 | break; 223 | } 224 | } 225 | 226 | return nice; 227 | } 228 | 229 | bool 230 | CompareDims(const DimsList& dims0, const DimsList& dims1) 231 | { 232 | if (dims0.size() != dims1.size()) { 233 | return false; 234 | } 235 | 236 | for (int i = 0; i < dims0.size(); ++i) { 237 | if (dims0[i] != dims1[i]) { 238 | return false; 239 | } 240 | } 241 | 242 | return true; 243 | } 244 | 245 | bool 246 | CompareDims( 247 | const std::vector& dims0, const std::vector& dims1) 248 | { 249 | if (dims0.size() != dims1.size()) { 250 | return false; 251 | } 252 | 253 | for (size_t i = 0; i < dims0.size(); ++i) { 254 | if (dims0[i] != dims1[i]) { 255 | return false; 256 | } 257 | } 258 | 259 | return true; 260 | } 261 | 262 | bool 263 | CompareDimsWithWildcard(const DimsList& dims0, const DimsList& dims1) 264 | { 265 | if (dims0.size() != dims1.size()) { 266 | return false; 267 | } 268 | 269 | for (int i = 0; i < dims0.size(); ++i) { 270 | if ((dims0[i] != WILDCARD_DIM) && (dims1[i] != WILDCARD_DIM) && 271 | (dims0[i] != dims1[i])) { 272 | return false; 273 | } 274 | } 275 | 276 | return true; 277 | } 278 | 279 | bool 280 | CompareDimsWithWildcard( 281 | const DimsList& dims0, const std::vector& dims1) 282 | { 283 | if (dims0.size() != (int64_t)dims1.size()) { 284 | return false; 285 | } 286 | 287 | for (int i = 0; i < dims0.size(); ++i) { 288 | if ((dims0[i] != WILDCARD_DIM) && (dims1[i] != WILDCARD_DIM) && 289 | (dims0[i] != dims1[i])) { 290 | return false; 291 | } 292 | } 293 | 294 | return true; 295 | } 296 | 297 | std::string 298 | DimsListToString(const DimsList& dims) 299 | { 300 | bool first = true; 301 | 302 | std::string str("["); 303 | for (const auto& dim : dims) { 304 | if (!first) { 305 | str += ","; 306 | } 307 | str += std::to_string(dim); 308 | first = false; 309 | } 310 | 311 | str += "]"; 312 | return str; 313 | } 314 | 315 | std::string 316 | DimsListToString(const std::vector& dims, const int start_idx) 317 | { 318 | int idx = 0; 319 | 320 | std::string str("["); 321 | for (const auto& dim : dims) { 322 | if (idx >= start_idx) { 323 | if (idx > start_idx) { 324 | str += ","; 325 | } 326 | str += std::to_string(dim); 327 | } 328 | 329 | idx++; 330 | } 331 | 332 | str += "]"; 333 | return str; 334 | } 335 | 336 | const char* 337 | DataTypeToProtocolString(const inference::DataType dtype) 338 | { 339 | switch (dtype) { 340 | case inference::DataType::TYPE_BOOL: 341 | return "BOOL"; 342 | case inference::DataType::TYPE_UINT8: 343 | return "UINT8"; 344 | case inference::DataType::TYPE_UINT16: 345 | return "UINT16"; 346 | case inference::DataType::TYPE_UINT32: 347 | return "UINT32"; 348 | case inference::DataType::TYPE_UINT64: 349 | return "UINT64"; 350 | case inference::DataType::TYPE_INT8: 351 | return "INT8"; 352 | case inference::DataType::TYPE_INT16: 353 | return "INT16"; 354 | case inference::DataType::TYPE_INT32: 355 | return "INT32"; 356 | case inference::DataType::TYPE_INT64: 357 | return "INT64"; 358 | case inference::DataType::TYPE_FP16: 359 | return "FP16"; 360 | case inference::DataType::TYPE_FP32: 361 | return "FP32"; 362 | case inference::DataType::TYPE_FP64: 363 | return "FP64"; 364 | case inference::DataType::TYPE_STRING: 365 | return "BYTES"; 366 | case inference::DataType::TYPE_BF16: 367 | return "BF16"; 368 | default: 369 | break; 370 | } 371 | 372 | return ""; 373 | } 374 | 375 | inference::DataType 376 | ProtocolStringToDataType(const std::string& dtype) 377 | { 378 | return ProtocolStringToDataType(dtype.c_str(), dtype.size()); 379 | } 380 | 381 | inference::DataType 382 | ProtocolStringToDataType(const char* dtype, size_t len) 383 | { 384 | if (len < 4 || len > 6) { 385 | return inference::DataType::TYPE_INVALID; 386 | } 387 | 388 | if ((*dtype == 'I') && (len != 6)) { 389 | if ((dtype[1] == 'N') && (dtype[2] == 'T')) { 390 | if ((dtype[3] == '8') && (len == 4)) { 391 | return inference::DataType::TYPE_INT8; 392 | } else if ((dtype[3] == '1') && (dtype[4] == '6')) { 393 | return inference::DataType::TYPE_INT16; 394 | } else if ((dtype[3] == '3') && (dtype[4] == '2')) { 395 | return inference::DataType::TYPE_INT32; 396 | } else if ((dtype[3] == '6') && (dtype[4] == '4')) { 397 | return inference::DataType::TYPE_INT64; 398 | } 399 | } 400 | } else if ((*dtype == 'U') && (len != 4)) { 401 | if ((dtype[1] == 'I') && (dtype[2] == 'N') && (dtype[3] == 'T')) { 402 | if ((dtype[4] == '8') && (len == 5)) { 403 | return inference::DataType::TYPE_UINT8; 404 | } else if ((dtype[4] == '1') && (dtype[5] == '6')) { 405 | return inference::DataType::TYPE_UINT16; 406 | } else if ((dtype[4] == '3') && (dtype[5] == '2')) { 407 | return inference::DataType::TYPE_UINT32; 408 | } else if ((dtype[4] == '6') && (dtype[5] == '4')) { 409 | return inference::DataType::TYPE_UINT64; 410 | } 411 | } 412 | } else if ((*dtype == 'F') && (dtype[1] == 'P') && (len == 4)) { 413 | if ((dtype[2] == '1') && (dtype[3] == '6')) { 414 | return inference::DataType::TYPE_FP16; 415 | } else if ((dtype[2] == '3') && (dtype[3] == '2')) { 416 | return inference::DataType::TYPE_FP32; 417 | } else if ((dtype[2] == '6') && (dtype[3] == '4')) { 418 | return inference::DataType::TYPE_FP64; 419 | } 420 | } else if (*dtype == 'B') { 421 | switch (dtype[1]) { 422 | case 'Y': 423 | if (!strcmp(dtype + 2, "TES")) { 424 | return inference::DataType::TYPE_STRING; 425 | } 426 | break; 427 | case 'O': 428 | if (!strcmp(dtype + 2, "OL")) { 429 | return inference::DataType::TYPE_BOOL; 430 | } 431 | break; 432 | case 'F': 433 | if (!strcmp(dtype + 2, "16")) { 434 | return inference::DataType::TYPE_BF16; 435 | } 436 | break; 437 | } 438 | } 439 | 440 | return inference::DataType::TYPE_INVALID; 441 | } 442 | 443 | }} // namespace triton::common 444 | -------------------------------------------------------------------------------- /src/table_printer.cc: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | #include "triton/common/table_printer.h" 28 | 29 | #ifdef _WIN32 30 | // suppress the min and max definitions in Windef.h. 31 | #define NOMINMAX 32 | #include 33 | #else 34 | #include 35 | #include 36 | #endif 37 | #include 38 | #include 39 | #include 40 | #include 41 | #include 42 | #include 43 | #include 44 | #include 45 | 46 | namespace triton { namespace common { 47 | 48 | // 49 | // ASCII table printer. 50 | // 51 | void 52 | TablePrinter::InsertRow(const std::vector& row) 53 | { 54 | std::vector> table_row; 55 | 56 | // Number of lines in each field in the record 57 | size_t max_height = 0; 58 | 59 | // Update max length of data items in each row 60 | for (size_t i = 0; i < row.size(); ++i) { 61 | table_row.push_back(std::vector{}); 62 | std::stringstream ss(row[i]); 63 | std::string line; 64 | 65 | size_t max_width = 0; 66 | while (std::getline(ss, line, '\n')) { 67 | table_row[i].push_back(line); 68 | if (line.size() > max_width) 69 | max_width = line.size(); 70 | } 71 | 72 | if (max_width > max_widths_[i]) 73 | max_widths_[i] = max_width; 74 | 75 | size_t number_of_lines = table_row[i].size(); 76 | if (max_height < number_of_lines) 77 | max_height = number_of_lines; 78 | } 79 | 80 | max_heights_.push_back(max_height); 81 | data_.emplace_back(table_row); 82 | } 83 | 84 | void 85 | TablePrinter::FairShare() 86 | { 87 | // initialize original index locations 88 | size_t array_size = max_widths_.size(); 89 | std::vector idx(array_size); 90 | iota(idx.begin(), idx.end(), 0); 91 | 92 | stable_sort(idx.begin(), idx.end(), [this](size_t i1, size_t i2) { 93 | return this->max_widths_[i1] < this->max_widths_[i2]; 94 | }); 95 | 96 | size_t loop_index = 1; 97 | for (auto itr = idx.begin(); itr != idx.end(); ++itr) { 98 | // If a column is not using all the space allocated to it 99 | if (max_widths_[*itr] < shares_[*itr]) { 100 | float excess = shares_[*itr] - max_widths_[*itr]; 101 | shares_[*itr] -= excess; 102 | 103 | if (itr == idx.end() - 1) 104 | break; 105 | auto update_itr = idx.begin() + (itr - idx.begin() + 1); 106 | 107 | // excess amount of unused space that must be distributed evenly to the 108 | // next columns 109 | float excess_per_column = excess / (array_size - loop_index); 110 | 111 | for (; update_itr != idx.end(); ++update_itr) { 112 | shares_[*update_itr] += excess_per_column; 113 | excess -= excess_per_column; 114 | } 115 | } 116 | ++loop_index; 117 | } 118 | 119 | // Remove any decimal shares 120 | for (auto itr = idx.begin(); itr != idx.end(); ++itr) { 121 | shares_[*itr] = (size_t)shares_[*itr]; 122 | } 123 | 124 | // For each record 125 | for (size_t i = 0; i < data_.size(); i++) { 126 | auto current_row = data_[i]; 127 | 128 | // For each field in the record 129 | for (size_t j = 0; j < current_row.size(); j++) { 130 | // For each line in the record 131 | for (size_t line_index = 0; line_index < current_row[j].size(); 132 | line_index++) { 133 | std::string line = current_row[j][line_index]; 134 | size_t num_rows = (line.size() + shares_[j] - 1) / shares_[j]; 135 | 136 | // If the number of rows required for this record is larger than 1, we 137 | // will break that line and put it in multiple lines 138 | if (num_rows > 1) { 139 | // Remove the multi-line field, it will be replaced by the line 140 | // that can fits the column size 141 | data_[i][j].erase(data_[i][j].begin() + line_index); 142 | for (size_t k = 0; k < num_rows; k++) { 143 | size_t start_index = 144 | std::min((size_t)(k * shares_[j]), line.size()); 145 | size_t end_index = 146 | std::min((size_t)((k + 1) * shares_[j]), line.size()); 147 | data_[i][j].insert( 148 | data_[i][j].begin() + line_index + k, 149 | line.substr(start_index, end_index - start_index)); 150 | } 151 | 152 | // We need to advance the index for the split lines. 153 | line_index += num_rows - 1; 154 | } 155 | 156 | if (max_heights_[i] < (num_rows - 1 + current_row[j].size())) 157 | max_heights_[i] += num_rows - 1; 158 | } 159 | } 160 | } 161 | } 162 | 163 | void 164 | TablePrinter::AddRow(std::stringstream& table, size_t row_index) 165 | { 166 | auto row = data_[row_index]; 167 | size_t max_height = max_heights_[row_index]; 168 | 169 | for (size_t j = 0; j < max_height; j++) { 170 | table << "|" << std::left; 171 | 172 | for (size_t i = 0; i < row.size(); i++) { 173 | if (j < row[i].size()) 174 | table << " " << std::setw(shares_[i]) << row[i][j] << " |"; 175 | else 176 | table << " " << std::setw(shares_[i]) << " " 177 | << " |"; 178 | } 179 | 180 | // Do not add new line if this is the last row of this record 181 | if (j != max_height - 1) 182 | table << "\n"; 183 | } 184 | table << "\n"; 185 | } 186 | 187 | void 188 | TablePrinter::AddRowDivider(std::stringstream& table) 189 | { 190 | table << "+"; 191 | for (const auto& share : shares_) { 192 | for (size_t i = 0; i < share + 2; i++) table << "-"; 193 | table << "+"; 194 | } 195 | table << "\n"; 196 | } 197 | 198 | std::string 199 | TablePrinter::PrintTable() 200 | { 201 | std::stringstream table; 202 | table << "\n"; 203 | 204 | FairShare(); 205 | 206 | AddRowDivider(table); 207 | // Add table headers 208 | AddRow(table, 0); 209 | AddRowDivider(table); 210 | 211 | for (size_t j = 1; j < data_.size(); j++) { 212 | AddRow(table, j); 213 | } 214 | 215 | AddRowDivider(table); 216 | 217 | return table.str(); 218 | } 219 | 220 | // TablePrinter will take the ownership of `headers`. 221 | TablePrinter::TablePrinter(const std::vector& headers) 222 | { 223 | // terminal size 224 | size_t column_size = 500; 225 | #ifdef _WIN32 226 | CONSOLE_SCREEN_BUFFER_INFO csbi; 227 | int ret = GetConsoleScreenBufferInfo(GetStdHandle(STD_OUTPUT_HANDLE), &csbi); 228 | if (ret && (csbi.dwSize.X != 0)) { 229 | column_size = csbi.dwSize.X; 230 | } 231 | #else 232 | struct winsize terminal_size; 233 | int status = ioctl(STDOUT_FILENO, TIOCGWINSZ, &terminal_size); 234 | if ((status == 0) && (terminal_size.ws_col != 0)) { 235 | column_size = terminal_size.ws_col; 236 | } 237 | #endif 238 | 239 | for (size_t i = 0; i < headers.size(); ++i) { 240 | max_widths_.emplace_back(0); 241 | } 242 | 243 | // Calculate fair share of every column 244 | size_t number_of_columns = headers.size(); 245 | 246 | // Terminal width is the actual terminal width minus two times spaces 247 | // required before and after each column and number of columns plus 1 for 248 | // the pipes between the columns 249 | size_t min_required = (2 * number_of_columns) + (number_of_columns + 1); 250 | // Ensure terminal_width is never negative 251 | size_t terminal_width = 0; 252 | if (column_size > min_required) { 253 | terminal_width = column_size - min_required; 254 | } else { 255 | // Ensure each column gets at least 1 unit width 256 | terminal_width = headers.size(); 257 | } 258 | size_t equal_share = terminal_width / headers.size(); 259 | 260 | for (size_t i = 0; i < headers.size(); ++i) { 261 | shares_.emplace_back(equal_share); 262 | terminal_width -= equal_share; 263 | } 264 | 265 | InsertRow(headers); 266 | } 267 | 268 | }} // namespace triton::common 269 | -------------------------------------------------------------------------------- /src/test/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # Redistribution and use in source and binary forms, with or without 4 | # modification, are permitted provided that the following conditions 5 | # are met: 6 | # * Redistributions of source code must retain the above copyright 7 | # notice, this list of conditions and the following disclaimer. 8 | # * Redistributions in binary form must reproduce the above copyright 9 | # notice, this list of conditions and the following disclaimer in the 10 | # documentation and/or other materials provided with the distribution. 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its 12 | # contributors may be used to endorse or promote products derived 13 | # from this software without specific prior written permission. 14 | # 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | cmake_minimum_required (VERSION 3.18) 28 | 29 | # 30 | # Unit tests 31 | # 32 | include(FetchContent) 33 | 34 | FetchContent_Declare( 35 | googletest 36 | URL https://github.com/google/googletest/archive/9406a60c7839052e4944ea4dbc8344762a89f9bd.zip 37 | ) 38 | FetchContent_MakeAvailable(googletest) 39 | 40 | if (TRITON_COMMON_ENABLE_JSON) 41 | add_subdirectory(triton_json triton_json) 42 | endif() 43 | 44 | -------------------------------------------------------------------------------- /src/test/triton_json/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # Redistribution and use in source and binary forms, with or without 4 | # modification, are permitted provided that the following conditions 5 | # are met: 6 | # * Redistributions of source code must retain the above copyright 7 | # notice, this list of conditions and the following disclaimer. 8 | # * Redistributions in binary form must reproduce the above copyright 9 | # notice, this list of conditions and the following disclaimer in the 10 | # documentation and/or other materials provided with the distribution. 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its 12 | # contributors may be used to endorse or promote products derived 13 | # from this software without specific prior written permission. 14 | # 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | cmake_minimum_required (VERSION 3.18) 28 | add_executable(triton-json-test triton_json_test.cc ../../error.cc) 29 | target_link_libraries( 30 | triton-json-test 31 | GTest::gtest 32 | GTest::gtest_main) 33 | 34 | target_include_directories( 35 | triton-json-test 36 | PUBLIC 37 | $ 38 | $ 39 | $ 40 | PRIVATE 41 | ${CMAKE_CURRENT_SOURCE_DIR} 42 | ) 43 | 44 | set_target_properties( 45 | triton-json-test 46 | PROPERTIES 47 | OUTPUT_NAME triton_json_test 48 | ) 49 | 50 | install( 51 | TARGETS triton-json-test 52 | RUNTIME DESTINATION bin 53 | ) 54 | -------------------------------------------------------------------------------- /src/test/triton_json/triton_json_test.cc: -------------------------------------------------------------------------------- 1 | // Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | #include "triton/common/error.h" 28 | 29 | #define TRITONJSON_STATUSTYPE Error 30 | #define TRITONJSON_STATUSRETURN(M) \ 31 | return Error(Error::Code::INTERNAL, (M).c_str()) 32 | #define TRITONJSON_STATUSSUCCESS Error() 33 | 34 | #include "gtest/gtest.h" 35 | #include "triton/common/triton_json.h" 36 | 37 | namespace { 38 | 39 | TEST(JsonTypeCheck, TestIsBool) 40 | { 41 | triton::common::TritonJson::Value value; 42 | value.Parse("{\"x\": true}"); 43 | 44 | triton::common::TritonJson::Value x; 45 | ASSERT_TRUE(value.Find("x", &x)); 46 | ASSERT_TRUE(x.IsBool()); 47 | 48 | value.Parse("{\"x\": false}"); 49 | ASSERT_TRUE(value.Find("x", &x)); 50 | ASSERT_TRUE(x.IsBool()); 51 | 52 | value.Parse("{\"x\": \"x\"}"); 53 | ASSERT_TRUE(value.Find("x", &x)); 54 | ASSERT_FALSE(x.IsBool()); 55 | } 56 | 57 | TEST(JsonTypeCheck, TestIsObject) 58 | { 59 | triton::common::TritonJson::Value value; 60 | value.Parse("{\"x\": {}}"); 61 | 62 | triton::common::TritonJson::Value x; 63 | ASSERT_TRUE(value.Find("x", &x)); 64 | ASSERT_TRUE(x.IsObject()); 65 | 66 | value.Parse("{\"x\": 2}"); 67 | ASSERT_TRUE(value.Find("x", &x)); 68 | ASSERT_FALSE(x.IsObject()); 69 | } 70 | 71 | TEST(JsonTypeCheck, TestIsString) 72 | { 73 | triton::common::TritonJson::Value value; 74 | value.Parse("{\"x\": \"123\"}"); 75 | 76 | triton::common::TritonJson::Value x; 77 | ASSERT_TRUE(value.Find("x", &x)); 78 | ASSERT_TRUE(x.IsString()); 79 | 80 | value.Parse("{\"x\": 2}"); 81 | ASSERT_TRUE(value.Find("x", &x)); 82 | ASSERT_FALSE(x.IsString()); 83 | } 84 | 85 | TEST(JsonTypeCheck, TestIsArray) 86 | { 87 | triton::common::TritonJson::Value value; 88 | value.Parse("{\"x\": []}"); 89 | 90 | triton::common::TritonJson::Value x; 91 | ASSERT_TRUE(value.Find("x", &x)); 92 | ASSERT_TRUE(x.IsArray()); 93 | 94 | value.Parse("{\"x\": 2}"); 95 | ASSERT_TRUE(value.Find("x", &x)); 96 | ASSERT_FALSE(x.IsArray()); 97 | } 98 | 99 | TEST(JsonTypeCheck, TestIsNumber) 100 | { 101 | triton::common::TritonJson::Value value; 102 | value.Parse("{\"x\": 2.0}"); 103 | 104 | triton::common::TritonJson::Value x; 105 | ASSERT_TRUE(value.Find("x", &x)); 106 | ASSERT_TRUE(x.IsNumber()); 107 | 108 | value.Parse("{\"x\": 2.001}"); 109 | ASSERT_TRUE(value.Find("x", &x)); 110 | ASSERT_TRUE(x.IsNumber()); 111 | 112 | value.Parse("{\"x\": 2}"); 113 | ASSERT_TRUE(value.Find("x", &x)); 114 | ASSERT_TRUE(x.IsNumber()); 115 | 116 | value.Parse("{\"x\": \"a\"}"); 117 | ASSERT_TRUE(value.Find("x", &x)); 118 | ASSERT_FALSE(x.IsNumber()); 119 | } 120 | 121 | TEST(JsonTypeCheck, TestIsInt) 122 | { 123 | triton::common::TritonJson::Value value; 124 | value.Parse("{\"x\": 2.0}"); 125 | 126 | triton::common::TritonJson::Value x; 127 | ASSERT_TRUE(value.Find("x", &x)); 128 | ASSERT_FALSE(x.IsInt()); 129 | 130 | value.Parse("{\"x\": 2}"); 131 | ASSERT_TRUE(value.Find("x", &x)); 132 | ASSERT_TRUE(x.IsInt()); 133 | 134 | value.Parse("{\"x\": -2}"); 135 | ASSERT_TRUE(value.Find("x", &x)); 136 | ASSERT_TRUE(x.IsInt()); 137 | 138 | value.Parse("{\"x\": \"a\"}"); 139 | ASSERT_TRUE(value.Find("x", &x)); 140 | ASSERT_FALSE(x.IsInt()); 141 | } 142 | 143 | } // namespace 144 | 145 | int 146 | main(int argc, char** argv) 147 | { 148 | ::testing::InitGoogleTest(&argc, argv); 149 | return RUN_ALL_TESTS(); 150 | } 151 | -------------------------------------------------------------------------------- /src/thread_pool.cc: -------------------------------------------------------------------------------- 1 | // Copyright 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | // 3 | // Redistribution and use in source and binary forms, with or without 4 | // modification, are permitted provided that the following conditions 5 | // are met: 6 | // * Redistributions of source code must retain the above copyright 7 | // notice, this list of conditions and the following disclaimer. 8 | // * Redistributions in binary form must reproduce the above copyright 9 | // notice, this list of conditions and the following disclaimer in the 10 | // documentation and/or other materials provided with the distribution. 11 | // * Neither the name of NVIDIA CORPORATION nor the names of its 12 | // contributors may be used to endorse or promote products derived 13 | // from this software without specific prior written permission. 14 | // 15 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | #include "triton/common/thread_pool.h" 28 | 29 | #include 30 | 31 | namespace triton { namespace common { 32 | 33 | ThreadPool::ThreadPool(size_t thread_count) 34 | { 35 | if (!thread_count) { 36 | throw std::invalid_argument("Thread count must be greater than zero."); 37 | } 38 | 39 | // Define infinite loop for each thread to wait for a task to complete 40 | const auto worker_loop = [this]() { 41 | while (true) { 42 | Task task; 43 | { 44 | std::unique_lock lk(queue_mtx_); 45 | // Wake if there's a task to do, or the pool has been stopped. 46 | cv_.wait(lk, [&]() { return !task_queue_.empty() || stop_; }); 47 | // Exit condition 48 | if (stop_ && task_queue_.empty()) { 49 | break; 50 | } 51 | task = std::move(task_queue_.front()); 52 | task_queue_.pop(); 53 | } 54 | 55 | // Execute task - ensure function has a valid target 56 | if (task) { 57 | task(); 58 | } 59 | } 60 | }; 61 | 62 | workers_.reserve(thread_count); 63 | for (size_t i = 0; i < thread_count; ++i) { 64 | workers_.emplace_back(worker_loop); 65 | } 66 | } 67 | 68 | ThreadPool::~ThreadPool() 69 | { 70 | { 71 | std::lock_guard lk(queue_mtx_); 72 | // Signal to each worker that it should exit loop when tasks are finished 73 | stop_ = true; 74 | } 75 | // Wake all threads to clean up 76 | cv_.notify_all(); 77 | for (auto& t : workers_) { 78 | t.join(); 79 | } 80 | } 81 | 82 | void 83 | ThreadPool::Enqueue(Task&& task) 84 | { 85 | { 86 | std::lock_guard lk(queue_mtx_); 87 | // Don't accept more work if pool is shutting down 88 | if (stop_) { 89 | return; 90 | } 91 | task_queue_.push(std::move(task)); 92 | } 93 | // Only wake one thread per task 94 | // Todo: DLIS-3859 if ThreadPool gets used more. 95 | cv_.notify_one(); 96 | } 97 | 98 | }} // namespace triton::common 99 | --------------------------------------------------------------------------------