From d654e0b13c4044ceb5fe423d2f4ba30ff601d16e Mon Sep 17 00:00:00 2001 From: Yingge He Date: Thu, 18 Jul 2024 16:17:47 -0700 Subject: [PATCH 1/6] Refactor string input checks --- include/triton/backend/backend_common.h | 20 +++++++ src/backend_common.cc | 70 +++++++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/include/triton/backend/backend_common.h b/include/triton/backend/backend_common.h index e8ba5fa..ef1cc07 100644 --- a/include/triton/backend/backend_common.h +++ b/include/triton/backend/backend_common.h @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -671,4 +672,23 @@ TRITONSERVER_Error* BufferAsTypedString( /// \return a formatted string for logging the request ID. std::string GetRequestId(TRITONBACKEND_Request* request); +/// Validate the contiguous string buffer with correct format +/// .... +/// @param buffer The pointer to the contiguous string buffer. +/// @param buffer_byte_size The size of the buffer in bytes. +/// @param expected_element_cnt The number of expected string elements. +/// @param input_name The name of the input buffer. +/// @param element_idx Returns the number of validated strings. +/// @param set_string_tensor_cb Callback function which sets string input tensor +/// depending on the backend platform. +/// @param onnx_backend Whether the backend platform is ONNX runtime. +/// @return a TRITONSERVER_Error indicating success or failure. +TRITONSERVER_Error* ValidateStringBuffer( + const char* buffer, size_t buffer_byte_size, + const size_t expected_element_cnt, const char* input_name, + size_t* element_idx, + const std::function& + set_string_tensor_cb, + bool onnx_backend = false); + }} // namespace triton::backend diff --git a/src/backend_common.cc b/src/backend_common.cc index 8c8821d..f3d98b3 100644 --- a/src/backend_common.cc +++ b/src/backend_common.cc @@ -1372,4 +1372,74 @@ GetRequestId(TRITONBACKEND_Request* request) return std::string("[request id: ") + request_id + "] "; } +TRITONSERVER_Error* +ValidateStringBuffer( + const char* buffer, size_t buffer_byte_size, + const size_t expected_element_cnt, const char* input_name, + size_t* element_idx, + const std::function& + set_string_tensor_cb, + bool onnx_backend) +{ + *element_idx = 0; + size_t remaining_bytes = buffer_byte_size; + + // Each string in 'buffer' is a 4-byte length followed by the string itself + // with no null-terminator. + while (remaining_bytes >= sizeof(uint32_t)) { + if (*element_idx >= expected_element_cnt) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + std::string( + "unexpected number of string elements " + + std::to_string(*element_idx + 1) + " for inference input '" + + input_name + "', expecting " + + std::to_string(expected_element_cnt)) + .c_str()); + } + + const uint32_t len = *(reinterpret_cast(buffer)); + remaining_bytes -= sizeof(uint32_t); + // Special handling for ONNX runtime backend + if (onnx_backend) { + // Make first byte of size info 0, so that if there is string data + // in front of it, the data becomes valid C string. + *const_cast(buffer) = 0; + } + buffer += sizeof(uint32_t); + + if (remaining_bytes < len) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + std::string( + "incomplete string data for inference input '" + + std::string(input_name) + "', expecting string of length " + + std::to_string(len) + " but only " + + std::to_string(remaining_bytes) + " bytes available") + .c_str()); + } + + set_string_tensor_cb(*element_idx, buffer, len); + buffer += len; + remaining_bytes -= len; + (*element_idx)++; + } + + // Special handling for ONNX runtime backend + if (onnx_backend && remaining_bytes > 0) { + *const_cast(buffer) = 0; + } + + if (*element_idx != expected_element_cnt) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + std::string( + "expected " + std::to_string(expected_element_cnt) + + " strings for inference input '" + input_name + "', got " + + std::to_string(*element_idx)) + .c_str()); + } + return nullptr; +} + }} // namespace triton::backend From 9694903b3b592324e8580c809d84081e5b6541e6 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Fri, 19 Jul 2024 15:58:11 -0700 Subject: [PATCH 2/6] Improve readability --- include/triton/backend/backend_common.h | 12 +++------ src/backend_common.cc | 34 +++++++++---------------- 2 files changed, 15 insertions(+), 31 deletions(-) diff --git a/include/triton/backend/backend_common.h b/include/triton/backend/backend_common.h index ef1cc07..995b671 100644 --- a/include/triton/backend/backend_common.h +++ b/include/triton/backend/backend_common.h @@ -672,23 +672,17 @@ TRITONSERVER_Error* BufferAsTypedString( /// \return a formatted string for logging the request ID. std::string GetRequestId(TRITONBACKEND_Request* request); -/// Validate the contiguous string buffer with correct format +/// Validate the contiguous string buffer with correct format. /// .... /// @param buffer The pointer to the contiguous string buffer. /// @param buffer_byte_size The size of the buffer in bytes. /// @param expected_element_cnt The number of expected string elements. /// @param input_name The name of the input buffer. -/// @param element_idx Returns the number of validated strings. -/// @param set_string_tensor_cb Callback function which sets string input tensor -/// depending on the backend platform. -/// @param onnx_backend Whether the backend platform is ONNX runtime. +/// @param str_list Returns pairs of address and length of parsed strings. /// @return a TRITONSERVER_Error indicating success or failure. TRITONSERVER_Error* ValidateStringBuffer( const char* buffer, size_t buffer_byte_size, const size_t expected_element_cnt, const char* input_name, - size_t* element_idx, - const std::function& - set_string_tensor_cb, - bool onnx_backend = false); + std::vector>* str_list); }} // namespace triton::backend diff --git a/src/backend_common.cc b/src/backend_common.cc index f3d98b3..61e2598 100644 --- a/src/backend_common.cc +++ b/src/backend_common.cc @@ -1376,23 +1376,22 @@ TRITONSERVER_Error* ValidateStringBuffer( const char* buffer, size_t buffer_byte_size, const size_t expected_element_cnt, const char* input_name, - size_t* element_idx, - const std::function& - set_string_tensor_cb, - bool onnx_backend) + std::vector>* str_list) { - *element_idx = 0; + size_t element_idx = 0; size_t remaining_bytes = buffer_byte_size; // Each string in 'buffer' is a 4-byte length followed by the string itself // with no null-terminator. while (remaining_bytes >= sizeof(uint32_t)) { - if (*element_idx >= expected_element_cnt) { + // Do not modify this line. str_list->size() must not exceed + // expected_element_cnt. + if (element_idx >= expected_element_cnt) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INVALID_ARG, std::string( "unexpected number of string elements " + - std::to_string(*element_idx + 1) + " for inference input '" + + std::to_string(element_idx + 1) + " for inference input '" + input_name + "', expecting " + std::to_string(expected_element_cnt)) .c_str()); @@ -1400,12 +1399,6 @@ ValidateStringBuffer( const uint32_t len = *(reinterpret_cast(buffer)); remaining_bytes -= sizeof(uint32_t); - // Special handling for ONNX runtime backend - if (onnx_backend) { - // Make first byte of size info 0, so that if there is string data - // in front of it, the data becomes valid C string. - *const_cast(buffer) = 0; - } buffer += sizeof(uint32_t); if (remaining_bytes < len) { @@ -1419,24 +1412,21 @@ ValidateStringBuffer( .c_str()); } - set_string_tensor_cb(*element_idx, buffer, len); + if (str_list) { + str_list->push_back({buffer, len}); + } buffer += len; remaining_bytes -= len; - (*element_idx)++; - } - - // Special handling for ONNX runtime backend - if (onnx_backend && remaining_bytes > 0) { - *const_cast(buffer) = 0; + element_idx++; } - if (*element_idx != expected_element_cnt) { + if (element_idx != expected_element_cnt) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INTERNAL, std::string( "expected " + std::to_string(expected_element_cnt) + " strings for inference input '" + input_name + "', got " + - std::to_string(*element_idx)) + std::to_string(element_idx)) .c_str()); } return nullptr; From 64be2465d423e81115e6d483ec552d52e30c6939 Mon Sep 17 00:00:00 2001 From: Yingge He <157551214+yinggeh@users.noreply.github.com> Date: Thu, 25 Jul 2024 17:02:08 -0700 Subject: [PATCH 3/6] Update include/triton/backend/backend_common.h Co-authored-by: Tanmay Verma --- include/triton/backend/backend_common.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/triton/backend/backend_common.h b/include/triton/backend/backend_common.h index 995b671..8fafe83 100644 --- a/include/triton/backend/backend_common.h +++ b/include/triton/backend/backend_common.h @@ -672,7 +672,7 @@ TRITONSERVER_Error* BufferAsTypedString( /// \return a formatted string for logging the request ID. std::string GetRequestId(TRITONBACKEND_Request* request); -/// Validate the contiguous string buffer with correct format. +/// Validate the contiguous string buffer with correct format and parse string elements into list of pairs of memory address and length. Note the returned list of pairs points to valid memory as long as memory pointed by buffer remains allocated. /// .... /// @param buffer The pointer to the contiguous string buffer. /// @param buffer_byte_size The size of the buffer in bytes. From ea004a89ec7fdbc305e8e9f4da7f5f75e7ee50f8 Mon Sep 17 00:00:00 2001 From: Yingge He <157551214+yinggeh@users.noreply.github.com> Date: Thu, 25 Jul 2024 17:12:04 -0700 Subject: [PATCH 4/6] Update backend_common.h --- include/triton/backend/backend_common.h | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/include/triton/backend/backend_common.h b/include/triton/backend/backend_common.h index 8fafe83..653f8fa 100644 --- a/include/triton/backend/backend_common.h +++ b/include/triton/backend/backend_common.h @@ -672,14 +672,18 @@ TRITONSERVER_Error* BufferAsTypedString( /// \return a formatted string for logging the request ID. std::string GetRequestId(TRITONBACKEND_Request* request); -/// Validate the contiguous string buffer with correct format and parse string elements into list of pairs of memory address and length. Note the returned list of pairs points to valid memory as long as memory pointed by buffer remains allocated. -/// .... -/// @param buffer The pointer to the contiguous string buffer. -/// @param buffer_byte_size The size of the buffer in bytes. -/// @param expected_element_cnt The number of expected string elements. -/// @param input_name The name of the input buffer. -/// @param str_list Returns pairs of address and length of parsed strings. -/// @return a TRITONSERVER_Error indicating success or failure. +/// Validate the contiguous string buffer with correct format +/// ... and parse string +/// elements into list of pairs of memory address and length. +/// Note the returned list of pairs points to valid memory as long +/// as memory pointed by buffer remains allocated. +/// +/// \param buffer The pointer to the contiguous string buffer. +/// \param buffer_byte_size The size of the buffer in bytes. +/// \param expected_element_cnt The number of expected string elements. +/// \param input_name The name of the input buffer. +/// \param str_list Returns pairs of address and length of parsed strings. +/// \return a TRITONSERVER_Error indicating success or failure. TRITONSERVER_Error* ValidateStringBuffer( const char* buffer, size_t buffer_byte_size, const size_t expected_element_cnt, const char* input_name, From 5d6905d0c347f3ded8df0c6842c04a8847d97bd1 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Sun, 28 Jul 2024 14:16:33 -0700 Subject: [PATCH 5/6] Minor updates --- include/triton/backend/backend_common.h | 1 - 1 file changed, 1 deletion(-) diff --git a/include/triton/backend/backend_common.h b/include/triton/backend/backend_common.h index 653f8fa..bbc80dd 100644 --- a/include/triton/backend/backend_common.h +++ b/include/triton/backend/backend_common.h @@ -28,7 +28,6 @@ #include #include #include -#include #include #include #include From 8364bf4390f23de5097ef1f16db484ac19e795e9 Mon Sep 17 00:00:00 2001 From: Yingge He <157551214+yinggeh@users.noreply.github.com> Date: Tue, 30 Jul 2024 15:40:45 -0700 Subject: [PATCH 6/6] Update copyright --- src/backend_common.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/backend_common.cc b/src/backend_common.cc index 61e2598..fb705aa 100644 --- a/src/backend_common.cc +++ b/src/backend_common.cc @@ -1,4 +1,4 @@ -// Copyright 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions