From 5a1ed16ed529ef5e041518203ccd6b7d56c77bca Mon Sep 17 00:00:00 2001 From: Yingge He Date: Thu, 18 Jul 2024 15:52:12 -0700 Subject: [PATCH 1/4] Refactor string input checks --- src/tensorflow.cc | 77 +++++++++++------------------------------------ 1 file changed, 17 insertions(+), 60 deletions(-) diff --git a/src/tensorflow.cc b/src/tensorflow.cc index b512996..6e3eeee 100644 --- a/src/tensorflow.cc +++ b/src/tensorflow.cc @@ -582,8 +582,11 @@ SetStringInputTensor( &contiguous_buffer, stream, &cuda_copy); if (err != nullptr) { RESPOND_AND_SET_NULL_IF_ERROR(response, err); - FillStringTensor( - tensor, tensor_offset + element_idx, request_element_cnt - element_idx); + if (element_idx < request_element_cnt) { + FillStringTensor( + tensor, tensor_offset + element_idx, + request_element_cnt - element_idx); + } free(contiguous_buffer); return cuda_copy; } @@ -595,68 +598,22 @@ SetStringInputTensor( } #endif // TRITON_ENABLE_GPU - // Parse content and assign to 'tensor'. Each string in 'content' - // is a 4-byte length followed by the string itself with no - // null-terminator. - while (content_byte_size >= sizeof(uint32_t)) { - if (element_idx >= request_element_cnt) { - RESPOND_AND_SET_NULL_IF_ERROR( - response, - TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - std::string( - "unexpected number of string elements " + - std::to_string(element_idx + 1) + " for inference input '" + - name + "', expecting " + std::to_string(request_element_cnt)) - .c_str())); - FillStringTensor( - tensor, tensor_offset + element_idx, - request_element_cnt - element_idx); - free(contiguous_buffer); - return cuda_copy; - } - - const uint32_t len = *(reinterpret_cast(content)); - content += sizeof(uint32_t); - content_byte_size -= sizeof(uint32_t); - - if (content_byte_size < len) { - RESPOND_AND_SET_NULL_IF_ERROR( - response, - TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - std::string( - "incomplete string data for inference input '" + - std::string(name) + "', expecting string of length " + - std::to_string(len) + " but only " + - std::to_string(content_byte_size) + " bytes available") - .c_str())); - FillStringTensor( - tensor, tensor_offset + element_idx, - request_element_cnt - element_idx); - free(contiguous_buffer); - return cuda_copy; - } - + auto callback = [](TRITONTF_Tensor* tensor, const size_t tensor_offset, + const size_t element_idx, const char* content, + const uint32_t len) { TRITONTF_TensorSetString(tensor, tensor_offset + element_idx, content, len); - content += len; - content_byte_size -= len; - element_idx++; - } - - if ((*response != nullptr) && (element_idx != request_element_cnt)) { - RESPOND_AND_SET_NULL_IF_ERROR( - response, TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - std::string( - "expected " + std::to_string(request_element_cnt) + - " strings for inference input '" + name + "', got " + - std::to_string(element_idx)) - .c_str())); + }; + auto fn = std::bind( + callback, tensor, tensor_offset, std::placeholders::_1, + std::placeholders::_2, std::placeholders::_3); + + err = ValidateStringBuffer( + content, content_byte_size, request_element_cnt, name, &element_idx, fn); + if (err != nullptr) { + RESPOND_AND_SET_NULL_IF_ERROR(response, err); FillStringTensor( tensor, tensor_offset + element_idx, request_element_cnt - element_idx); } - free(contiguous_buffer); return cuda_copy; } From 9f2b66790eeae5cb062f29b681450db7b7cd4127 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Fri, 19 Jul 2024 15:56:27 -0700 Subject: [PATCH 2/4] Improve readability --- src/tensorflow.cc | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/tensorflow.cc b/src/tensorflow.cc index 6e3eeee..b913de8 100644 --- a/src/tensorflow.cc +++ b/src/tensorflow.cc @@ -598,17 +598,15 @@ SetStringInputTensor( } #endif // TRITON_ENABLE_GPU - auto callback = [](TRITONTF_Tensor* tensor, const size_t tensor_offset, - const size_t element_idx, const char* content, - const uint32_t len) { - TRITONTF_TensorSetString(tensor, tensor_offset + element_idx, content, len); - }; - auto fn = std::bind( - callback, tensor, tensor_offset, std::placeholders::_1, - std::placeholders::_2, std::placeholders::_3); - + std::vector> str_list; err = ValidateStringBuffer( - content, content_byte_size, request_element_cnt, name, &element_idx, fn); + content, content_byte_size, request_element_cnt, name, &str_list); + // Set string values. + for (; element_idx < str_list.size(); ++element_idx) { + const auto& [addr, len] = str_list[element_idx]; + TRITONTF_TensorSetString(tensor, tensor_offset + element_idx, addr, len); + } + if (err != nullptr) { RESPOND_AND_SET_NULL_IF_ERROR(response, err); FillStringTensor( From 42b8a0e5b48385f017ac2aa9cd2623a3a5153d78 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Fri, 26 Jul 2024 11:02:59 -0700 Subject: [PATCH 3/4] Minor fix --- src/tensorflow.cc | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/tensorflow.cc b/src/tensorflow.cc index b913de8..885384d 100644 --- a/src/tensorflow.cc +++ b/src/tensorflow.cc @@ -567,7 +567,6 @@ SetStringInputTensor( cudaStream_t stream, const char* host_policy_name) { bool cuda_copy = false; - size_t element_idx = 0; // For string data type, we always need to have the data on CPU so // that we can read string length and construct the string @@ -582,11 +581,7 @@ SetStringInputTensor( &contiguous_buffer, stream, &cuda_copy); if (err != nullptr) { RESPOND_AND_SET_NULL_IF_ERROR(response, err); - if (element_idx < request_element_cnt) { - FillStringTensor( - tensor, tensor_offset + element_idx, - request_element_cnt - element_idx); - } + FillStringTensor(tensor, tensor_offset, request_element_cnt); free(contiguous_buffer); return cuda_copy; } @@ -602,15 +597,16 @@ SetStringInputTensor( err = ValidateStringBuffer( content, content_byte_size, request_element_cnt, name, &str_list); // Set string values. - for (; element_idx < str_list.size(); ++element_idx) { + for (size_t element_idx = 0; element_idx < str_list.size(); ++element_idx) { const auto& [addr, len] = str_list[element_idx]; TRITONTF_TensorSetString(tensor, tensor_offset + element_idx, addr, len); } + size_t element_cnt = str_list.size(); if (err != nullptr) { RESPOND_AND_SET_NULL_IF_ERROR(response, err); FillStringTensor( - tensor, tensor_offset + element_idx, request_element_cnt - element_idx); + tensor, tensor_offset + element_cnt, request_element_cnt - element_cnt); } free(contiguous_buffer); return cuda_copy; From 44c3c845f9a89e12b4fde55d0a8588ed00fa70b4 Mon Sep 17 00:00:00 2001 From: Yingge He <157551214+yinggeh@users.noreply.github.com> Date: Tue, 30 Jul 2024 15:42:00 -0700 Subject: [PATCH 4/4] Update tensorflow.cc --- src/tensorflow.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tensorflow.cc b/src/tensorflow.cc index 885384d..a297ea0 100644 --- a/src/tensorflow.cc +++ b/src/tensorflow.cc @@ -1,4 +1,4 @@ -// Copyright 2020-2023, 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