diff --git a/qa/L0_http/http_request_many_chunks.py b/qa/L0_http/http_request_many_chunks.py new file mode 100755 index 0000000000..a67d9a82c5 --- /dev/null +++ b/qa/L0_http/http_request_many_chunks.py @@ -0,0 +1,153 @@ +#!/usr/bin/python +# Copyright 2025, 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 +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import socket +import unittest + + +class HTTPRequestManyChunksTest(unittest.TestCase): + def setUp(self): + self._model_name = "simple" + self._local_host = "localhost" + self._http_port = 8000 + self._malicious_chunk_count = ( + 1000000 # large enough to cause a stack overflow if using alloca() + ) + self._parse_error = ( + "failed to parse the request JSON buffer: Invalid value. at 0" + ) + + def send_chunked_request( + self, header: str, chunk_count: int, expected_response: str + ): + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + header = ( + f"{header}" + f"Host: {self._local_host}:{self._http_port}\r\n" + f"Content-Type: application/octet-stream\r\n" + f"Transfer-Encoding: chunked\r\n" + f"Connection: close\r\n" + f"\r\n" + ) + try: + s.connect((self._local_host, self._http_port)) + # HTTP request with chunked encoding + s.sendall((header.encode())) + + # Send chunked payload + for _ in range(chunk_count): + s.send(b"1\r\nA\r\n") + # End chunked encoding + s.sendall(b"0\r\n\r\n") + + # Receive response + response = b"" + while True: + try: + chunk = s.recv(4096) + if not chunk: + break + response += chunk + except socket.timeout: + break + self.assertIn(expected_response, response.decode()) + except Exception as e: + raise (e) + finally: + s.close() + + def test_infer(self): + request_header = ( + f"POST /v2/models/{self._model_name}/infer HTTP/1.1\r\n" + f"Inference-Header-Content-Length: 0\r\n" + ) + + self.send_chunked_request( + request_header, + self._malicious_chunk_count, + "Raw request must only have 1 input (found 1) to be deduced but got 2 inputs in 'simple' model configuration", + ) + + def test_registry_index(self): + request_header = f"POST /v2/repository/index HTTP/1.1\r\n" + + self.send_chunked_request( + request_header, self._malicious_chunk_count, self._parse_error + ) + + def test_model_control(self): + load_request_header = ( + f"POST /v2/repository/models/{self._model_name}/load HTTP/1.1\r\n" + ) + unload_request_header = load_request_header.replace("/load", "/unload") + + self.send_chunked_request( + load_request_header, self._malicious_chunk_count, self._parse_error + ) + self.send_chunked_request( + unload_request_header, self._malicious_chunk_count, self._parse_error + ) + + def test_trace(self): + request_header = ( + f"POST /v2/models/{self._model_name}/trace/setting HTTP/1.1\r\n" + ) + + self.send_chunked_request( + request_header, self._malicious_chunk_count, self._parse_error + ) + + def test_logging(self): + request_header = f"POST /v2/logging HTTP/1.1\r\n" + + self.send_chunked_request( + request_header, self._malicious_chunk_count, self._parse_error + ) + + def test_system_shm_register(self): + request_header = f"POST /v2/systemsharedmemory/region/test_system_shm_register/register HTTP/1.1\r\n" + + self.send_chunked_request( + request_header, self._malicious_chunk_count, self._parse_error + ) + + def test_cuda_shm_register(self): + request_header = f"POST /v2/cudasharedmemory/region/test_cuda_shm_register/register HTTP/1.1\r\n" + + self.send_chunked_request( + request_header, self._malicious_chunk_count, self._parse_error + ) + + def test_generate(self): + request_header = f"POST /v2/models/{self._model_name}/generate HTTP/1.1\r\n" + self.send_chunked_request( + request_header, self._malicious_chunk_count, self._parse_error + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/qa/L0_http/test.sh b/qa/L0_http/test.sh index a96bd6a1c6..0fc25501a3 100755 --- a/qa/L0_http/test.sh +++ b/qa/L0_http/test.sh @@ -40,6 +40,7 @@ fi export CUDA_VISIBLE_DEVICES=0 +source ../common/util.sh RET=0 CLIENT_PLUGIN_TEST="./http_client_plugin_test.py" @@ -129,7 +130,6 @@ set -e CLIENT_LOG=`pwd`/client.log SERVER_ARGS="--backend-directory=${BACKEND_DIR} --model-repository=${MODELDIR}" -source ../common/util.sh run_server if [ "$SERVER_PID" == "0" ]; then @@ -855,7 +855,32 @@ elif [ `grep -c "Error: --http-max-input-size must be greater than 0." ${SERVER_ RET=1 fi -### +### Test HTTP Requests Containing Many Chunks ### +MODELDIR="`pwd`/models" +REQUEST_MANY_CHUNKS_PY="http_request_many_chunks.py" +CLIENT_LOG="./client.http_request_many_chunks.log" +SERVER_ARGS="--model-repository=${MODELDIR} --log-verbose=1 --model-control-mode=explicit --load-model=simple" +SERVER_LOG="./inference_server_request_many_chunks.log" + +run_server +if [ "$SERVER_PID" == "0" ]; then + echo -e "\n***\n*** Failed to start $SERVER\n***" + cat $SERVER_LOG + exit 1 +fi + +set +e +python $REQUEST_MANY_CHUNKS_PY -v >> ${CLIENT_LOG} 2>&1 +if [ $? -ne 0 ]; then + echo -e "\n***\n*** HTTP Request Many Chunks Test Failed\n***" + cat $SERVER_LOG + cat $CLIENT_LOG + RET=1 +fi +set -e + +kill $SERVER_PID +wait $SERVER_PID if [ $RET -eq 0 ]; then echo -e "\n***\n*** Test Passed\n***" diff --git a/qa/L0_sagemaker/sagemaker_request_many_chunks.py b/qa/L0_sagemaker/sagemaker_request_many_chunks.py new file mode 100755 index 0000000000..415ddcba28 --- /dev/null +++ b/qa/L0_sagemaker/sagemaker_request_many_chunks.py @@ -0,0 +1,91 @@ +#!/usr/bin/python +# Copyright 2025, 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 +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import socket +import unittest + + +class SagemakerRequestManyChunksTest(unittest.TestCase): + def setUp(self): + self._local_host = "localhost" + self._sagemaker_port = 8080 + self._malicious_chunk_count = ( + 1000000 # large enough to cause a stack overflow if using alloca() + ) + + def send_chunked_request( + self, header: str, chunk_count: int, expected_response: str + ): + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + header = ( + f"{header}" + f"Host: {self._local_host}:{self._sagemaker_port}\r\n" + f"Content-Type: application/octet-stream\r\n" + f"Transfer-Encoding: chunked\r\n" + f"Connection: close\r\n" + f"\r\n" + ) + try: + s.connect((self._local_host, self._sagemaker_port)) + # HTTP request with chunked encoding + s.sendall((header.encode())) + + # Send chunked payload + for _ in range(chunk_count): + s.send(b"1\r\nA\r\n") + # End chunked encoding + s.sendall(b"0\r\n\r\n") + + # Receive response + response = b"" + while True: + try: + chunk = s.recv(4096) + if not chunk: + break + response += chunk + except socket.timeout: + break + self.assertIn(expected_response, response.decode()) + except Exception as e: + raise (e) + finally: + s.close() + + def test_load_model(self): + request_header = ( + f"POST /models HTTP/1.1\r\n" f"X-Amzn-SageMaker-Target-Model: ZZZZZZZ\r\n" + ) + self.send_chunked_request( + request_header, + self._malicious_chunk_count, + "failed to parse the request JSON buffer: Invalid value. at 0", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/qa/L0_sagemaker/test.sh b/qa/L0_sagemaker/test.sh index dde4c794d8..c86b278e38 100755 --- a/qa/L0_sagemaker/test.sh +++ b/qa/L0_sagemaker/test.sh @@ -565,6 +565,53 @@ kill $SERVER_PID wait $SERVE_PID # MME end +### Test Sagemaker Requests Containing Many Chunks ### +rm -rf models && mkdir models && \ + cp -r $DATADIR/qa_model_repository/onnx_int32_int32_int32 models/sm_model && \ + rm -r models/sm_model/2 && rm -r models/sm_model/3 && \ + sed -i "s/onnx_int32_int32_int32/sm_model/" models/sm_model/config.pbtxt + +export SAGEMAKER_TRITON_DEFAULT_MODEL_NAME=sm_model +REQUEST_MANY_CHUNKS_PY="sagemaker_request_many_chunks.py" +CLIENT_LOG="./client.sagemaker_request_many_chunks.log" +SERVER_LOG="./server.sagemaker_request_many_chunks.log" + +serve > $SERVER_LOG 2>&1 & +SERVE_PID=$! +# Obtain Triton PID in such way as $! will return the script PID +sleep 1 +SERVER_PID=`ps | grep tritonserver | awk '{ printf $1 }'` +sagemaker_wait_for_server_ready $SERVER_PID 10 +if [ "$WAIT_RET" != "0" ]; then + echo -e "\n***\n*** Failed to start $SERVER\n***" + kill $SERVER_PID || true + cat $SERVER_LOG + exit 1 +fi + +# Ping +set +e +code=`curl -s -w %{http_code} -o ./ping.out localhost:8080/ping` +set -e +if [ "$code" != "200" ]; then + cat ./ping.out + echo -e "\n***\n*** Test Failed\n***" + RET=1 +fi + +set +e +python $REQUEST_MANY_CHUNKS_PY >>$CLIENT_LOG 2>&1 +if [ $? -ne 0 ]; then + echo -e "\n***\n*** Sagemaker Request Many Chunks Test Failed\n***" + cat $SERVER_LOG + cat $CLIENT_LOG + RET=1 +fi +set -e + +kill $SERVER_PID +wait $SERVE_PID + unlink /opt/ml/model rm -rf /opt/ml/model diff --git a/src/http_server.cc b/src/http_server.cc index e3ee3efe17..f62f73c948 100644 --- a/src/http_server.cc +++ b/src/http_server.cc @@ -1053,62 +1053,6 @@ ValidateOutputParameter(triton::common::TritonJson::Value& io) return nullptr; // success } -TRITONSERVER_Error* -EVBufferToJson( - triton::common::TritonJson::Value* document, evbuffer_iovec* v, int* v_idx, - const size_t length, int n) -{ - size_t offset = 0, remaining_length = length; - char* json_base; - std::vector json_buffer; - - // No need to memcpy when number of iovecs is 1 - if ((n > 0) && (v[0].iov_len >= remaining_length)) { - json_base = static_cast(v[0].iov_base); - if (v[0].iov_len > remaining_length) { - v[0].iov_base = static_cast(json_base + remaining_length); - v[0].iov_len -= remaining_length; - remaining_length = 0; - } else if (v[0].iov_len == remaining_length) { - remaining_length = 0; - *v_idx += 1; - } - } else { - json_buffer.resize(length); - json_base = json_buffer.data(); - while ((remaining_length > 0) && (*v_idx < n)) { - char* base = static_cast(v[*v_idx].iov_base); - size_t base_size; - if (v[*v_idx].iov_len > remaining_length) { - base_size = remaining_length; - v[*v_idx].iov_base = static_cast(base + remaining_length); - v[*v_idx].iov_len -= remaining_length; - remaining_length = 0; - } else { - base_size = v[*v_idx].iov_len; - remaining_length -= v[*v_idx].iov_len; - *v_idx += 1; - } - - memcpy(json_base + offset, base, base_size); - offset += base_size; - } - } - - if (remaining_length != 0) { - return TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - std::string( - "unexpected size for request JSON, expecting " + - std::to_string(remaining_length) + " more bytes") - .c_str()); - } - - RETURN_IF_ERR(document->Parse(json_base, length)); - - return nullptr; // success -} - std::string CompressionTypeUsed(const std::string accept_encoding) { @@ -1432,34 +1376,16 @@ HTTPAPIServer::HandleRepositoryIndex( } TRITONSERVER_Error* err = nullptr; - - struct evbuffer_iovec* v = nullptr; - int v_idx = 0; - int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0); - if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); - if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) { - err = TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - "unexpected error getting registry index request body"); - } - } - + triton::common::TritonJson::Value index_request; bool ready = false; + size_t buffer_len = 0; + RETURN_AND_RESPOND_IF_ERR( + req, EVRequestToJson(req, "registry index", &index_request, &buffer_len)); - if (err == nullptr) { - // If no request json then just use all default values. - size_t buffer_len = evbuffer_get_length(req->buffer_in); - if (buffer_len > 0) { - triton::common::TritonJson::Value index_request; - err = EVBufferToJson(&index_request, v, &v_idx, buffer_len, n); - if (err == nullptr) { - triton::common::TritonJson::Value ready_json; - if (index_request.Find("ready", &ready_json)) { - err = ready_json.AsBool(&ready); - } - } + if (buffer_len > 0) { + triton::common::TritonJson::Value ready_json; + if (index_request.Find("ready", &ready_json)) { + err = ready_json.AsBool(&ready); } } @@ -1508,19 +1434,6 @@ HTTPAPIServer::HandleRepositoryControl( "'repository_name' specification is not supported"); } else { if (action == "load") { - struct evbuffer_iovec* v = nullptr; - int v_idx = 0; - int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0); - if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); - if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) { - RETURN_AND_RESPOND_IF_ERR( - req, TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - "unexpected error getting load model request buffers")); - } - } static auto param_deleter = [](std::vector* params) { if (params != nullptr) { @@ -1538,15 +1451,15 @@ HTTPAPIServer::HandleRepositoryControl( std::list> binary_files; // WAR for the const-ness check std::vector const_params; - size_t buffer_len = evbuffer_get_length(req->buffer_in); - if (buffer_len > 0) { - triton::common::TritonJson::Value request; - RETURN_AND_RESPOND_IF_ERR( - req, EVBufferToJson(&request, v, &v_idx, buffer_len, n)); + triton::common::TritonJson::Value load_request; + size_t buffer_len = 0; + RETURN_AND_RESPOND_IF_ERR( + req, EVRequestToJson(req, "load model", &load_request, &buffer_len)); + if (buffer_len > 0) { // Parse request body for parameters triton::common::TritonJson::Value param_json; - if (request.Find("parameters", ¶m_json)) { + if (load_request.Find("parameters", ¶m_json)) { // Iterate over each member in 'param_json' std::vector members; RETURN_AND_RESPOND_IF_ERR(req, param_json.Members(&members)); @@ -1592,37 +1505,25 @@ HTTPAPIServer::HandleRepositoryControl( // Check if the dependent models should be removed bool unload_dependents = false; { - struct evbuffer_iovec* v = nullptr; - int v_idx = 0; - int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0); - if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); - if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) { - err = TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - "unexpected error getting model control request body"); - } - } + triton::common::TritonJson::Value control_request; + size_t buffer_len = 0; + RETURN_AND_RESPOND_IF_ERR( + req, EVRequestToJson( + req, "unload model", &control_request, &buffer_len)); - size_t buffer_len = evbuffer_get_length(req->buffer_in); if (buffer_len > 0) { - triton::common::TritonJson::Value control_request; - err = EVBufferToJson(&control_request, v, &v_idx, buffer_len, n); - if (err == nullptr) { - triton::common::TritonJson::Value params_json; - if (control_request.Find("parameters", ¶ms_json)) { - triton::common::TritonJson::Value ud_json; - if (params_json.Find("unload_dependents", &ud_json)) { - auto parse_err = ud_json.AsBool(&unload_dependents); - if (parse_err != nullptr) { - err = TRITONSERVER_ErrorNew( - TRITONSERVER_ErrorCode(parse_err), - (std::string("Unable to parse 'unload_dependents': ") + - TRITONSERVER_ErrorMessage(parse_err)) - .c_str()); - TRITONSERVER_ErrorDelete(parse_err); - } + triton::common::TritonJson::Value params_json; + if (control_request.Find("parameters", ¶ms_json)) { + triton::common::TritonJson::Value ud_json; + if (params_json.Find("unload_dependents", &ud_json)) { + auto parse_err = ud_json.AsBool(&unload_dependents); + if (parse_err != nullptr) { + err = TRITONSERVER_ErrorNew( + TRITONSERVER_ErrorCode(parse_err), + (std::string("Unable to parse 'unload_dependents': ") + + TRITONSERVER_ErrorMessage(parse_err)) + .c_str()); + TRITONSERVER_ErrorDelete(parse_err); } } } @@ -1868,42 +1769,27 @@ HTTPAPIServer::HandleTrace(evhtp_request_t* req, const std::string& model_name) // Perform trace setting update if requested if (req->method == htp_method_POST) { - struct evbuffer_iovec* v = nullptr; - int v_idx = 0; - int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0); - if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); - if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) { - RETURN_AND_RESPOND_IF_ERR( - req, TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - "unexpected error getting trace request buffers")); - } - } - - triton::common::TritonJson::Value request; - size_t buffer_len = evbuffer_get_length(req->buffer_in); + triton::common::TritonJson::Value trace_request; RETURN_AND_RESPOND_IF_ERR( - req, EVBufferToJson(&request, v, &v_idx, buffer_len, n)); + req, EVRequestToJsonAllowsEmpty(req, "trace", &trace_request)); TraceManager::NewSetting new_setting; triton::common::TritonJson::Value setting_json; - if (request.Find("trace_file", &setting_json)) { + if (trace_request.Find("trace_file", &setting_json)) { RETURN_AND_RESPOND_IF_ERR( req, TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_UNSUPPORTED, "trace file location can not be updated through network " "protocol")); } - if (request.Find("trace_level", &setting_json)) { + if (trace_request.Find("trace_level", &setting_json)) { if (setting_json.IsNull()) { new_setting.clear_level_ = true; } else { triton::common::TritonJson::Value level_array; RETURN_AND_RESPOND_IF_ERR( - req, request.MemberAsArray("trace_level", &level_array)); + req, trace_request.MemberAsArray("trace_level", &level_array)); for (size_t i = 0; i < level_array.ArraySize(); ++i) { std::string level_str; RETURN_AND_RESPOND_IF_ERR( @@ -1930,7 +1816,7 @@ HTTPAPIServer::HandleTrace(evhtp_request_t* req, const std::string& model_name) } } } - if (request.Find("trace_rate", &setting_json)) { + if (trace_request.Find("trace_rate", &setting_json)) { if (setting_json.IsNull()) { new_setting.clear_rate_ = true; } else { @@ -1963,7 +1849,7 @@ HTTPAPIServer::HandleTrace(evhtp_request_t* req, const std::string& model_name) } } } - if (request.Find("trace_count", &setting_json)) { + if (trace_request.Find("trace_count", &setting_json)) { if (setting_json.IsNull()) { new_setting.clear_count_ = true; } else { @@ -2005,7 +1891,7 @@ HTTPAPIServer::HandleTrace(evhtp_request_t* req, const std::string& model_name) } } } - if (request.Find("log_frequency", &setting_json)) { + if (trace_request.Find("log_frequency", &setting_json)) { if (setting_json.IsNull()) { new_setting.clear_log_frequency_ = true; } else { @@ -2128,28 +2014,13 @@ HTTPAPIServer::HandleLogging(evhtp_request_t* req) #ifdef TRITON_ENABLE_LOGGING // Perform log setting update if requested if (req->method == htp_method_POST) { - struct evbuffer_iovec* v = nullptr; - int v_idx = 0; - int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0); - if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); - if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) { - RETURN_AND_RESPOND_IF_ERR( - req, - TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - "unexpected error getting dynamic logging request buffers")); - } - } - triton::common::TritonJson::Value request; - size_t buffer_len = evbuffer_get_length(req->buffer_in); + triton::common::TritonJson::Value log_request; RETURN_AND_RESPOND_IF_ERR( - req, EVBufferToJson(&request, v, &v_idx, buffer_len, n)); + req, EVRequestToJsonAllowsEmpty(req, "dynamic logging", &log_request)); // Server and Core repos do not have the same Logger object // Each update must be applied to both server and core repo versions triton::common::TritonJson::Value setting_json; - if (request.Find("log_file", &setting_json)) { + if (log_request.Find("log_file", &setting_json)) { if (!setting_json.IsNull()) { RETURN_AND_RESPOND_IF_ERR( req, TRITONSERVER_ErrorNew( @@ -2158,7 +2029,7 @@ HTTPAPIServer::HandleLogging(evhtp_request_t* req) "protocol")); } } - if (request.Find("log_info", &setting_json)) { + if (log_request.Find("log_info", &setting_json)) { if (!setting_json.IsNull()) { bool log_info_status; RETURN_AND_RESPOND_IF_ERR(req, setting_json.AsBool(&log_info_status)); @@ -2166,7 +2037,7 @@ HTTPAPIServer::HandleLogging(evhtp_request_t* req) TRITONSERVER_ServerOptionsSetLogInfo(nullptr, log_info_status); } } - if (request.Find("log_warning", &setting_json)) { + if (log_request.Find("log_warning", &setting_json)) { if (!setting_json.IsNull()) { bool log_warn_status; RETURN_AND_RESPOND_IF_ERR(req, setting_json.AsBool(&log_warn_status)); @@ -2174,7 +2045,7 @@ HTTPAPIServer::HandleLogging(evhtp_request_t* req) TRITONSERVER_ServerOptionsSetLogWarn(nullptr, log_warn_status); } } - if (request.Find("log_error", &setting_json)) { + if (log_request.Find("log_error", &setting_json)) { if (!setting_json.IsNull()) { bool log_error_status; RETURN_AND_RESPOND_IF_ERR(req, setting_json.AsBool(&log_error_status)); @@ -2182,7 +2053,7 @@ HTTPAPIServer::HandleLogging(evhtp_request_t* req) TRITONSERVER_ServerOptionsSetLogError(nullptr, log_error_status); } } - if (request.Find("log_verbose_level", &setting_json)) { + if (log_request.Find("log_verbose_level", &setting_json)) { if (!setting_json.IsNull()) { uint64_t verbose_level; RETURN_AND_RESPOND_IF_ERR(req, setting_json.AsUInt(&verbose_level)); @@ -2191,7 +2062,7 @@ HTTPAPIServer::HandleLogging(evhtp_request_t* req) nullptr, static_cast(verbose_level)); } } - if (request.Find("log_format", &setting_json)) { + if (log_request.Find("log_format", &setting_json)) { if (!setting_json.IsNull()) { std::string log_format_parse; RETURN_AND_RESPOND_IF_ERR( @@ -2310,63 +2181,46 @@ HTTPAPIServer::HandleSystemSharedMemory( TRITONSERVER_ERROR_INVALID_ARG, "'region name' is necessary to register system shared memory region"); } else { - struct evbuffer_iovec* v = nullptr; - int v_idx = 0; - int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0); - if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); - if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) { - err = TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - "unexpected error getting register request buffers"); - } + triton::common::TritonJson::Value register_request; + triton::common::TritonJson::Value key_json; + RETURN_AND_RESPOND_IF_ERR( + req, EVRequestToJsonAllowsEmpty(req, action, ®ister_request)); + if (!register_request.Find("key", &key_json)) { + err = TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "Shared memory register request has no 'key' field"); } + const char* shm_key = nullptr; if (err == nullptr) { - triton::common::TritonJson::Value register_request; - size_t buffer_len = evbuffer_get_length(req->buffer_in); - err = EVBufferToJson(®ister_request, v, &v_idx, buffer_len, n); - if (err == nullptr) { - triton::common::TritonJson::Value key_json; - if (!register_request.Find("key", &key_json)) { - err = TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - "Shared memory register request has no 'key' field"); - } - - const char* shm_key = nullptr; - if (err == nullptr) { - size_t shm_key_len; - err = key_json.AsString(&shm_key, &shm_key_len); - } - - uint64_t offset = 0; - if (err == nullptr) { - triton::common::TritonJson::Value offset_json; - if (register_request.Find("offset", &offset_json)) { - err = offset_json.AsUInt(&offset); - } - } + size_t shm_key_len; + err = key_json.AsString(&shm_key, &shm_key_len); + } - uint64_t byte_size = 0; - if (err == nullptr) { - triton::common::TritonJson::Value byte_size_json; - if (!register_request.Find("byte_size", &byte_size_json)) { - err = TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - "Shared memory register request has no 'byte_size' field"); - } else { - err = byte_size_json.AsUInt(&byte_size); - } - } + uint64_t offset = 0; + if (err == nullptr) { + triton::common::TritonJson::Value offset_json; + if (register_request.Find("offset", &offset_json)) { + err = offset_json.AsUInt(&offset); + } + } - if (err == nullptr) { - err = shm_manager_->RegisterSystemSharedMemory( - region_name, shm_key, offset, byte_size); - } + uint64_t byte_size = 0; + if (err == nullptr) { + triton::common::TritonJson::Value byte_size_json; + if (!register_request.Find("byte_size", &byte_size_json)) { + err = TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "Shared memory register request has no 'byte_size' field"); + } else { + err = byte_size_json.AsUInt(&byte_size); } } + + if (err == nullptr) { + err = shm_manager_->RegisterSystemSharedMemory( + region_name, shm_key, offset, byte_size); + } } } else if (action == "unregister") { if (region_name.empty()) { @@ -2417,67 +2271,50 @@ HTTPAPIServer::HandleCudaSharedMemory( "'region name' is necessary to register cuda shared memory region"); } else { #ifdef TRITON_ENABLE_GPU - struct evbuffer_iovec* v = nullptr; - int v_idx = 0; - int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0); - if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); - if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) { - err = TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - "unexpected error getting register request buffers"); - } + triton::common::TritonJson::Value register_request; + RETURN_AND_RESPOND_IF_ERR( + req, EVRequestToJsonAllowsEmpty(req, action, ®ister_request)); + const char* b64_handle = nullptr; + size_t b64_handle_len = 0; + triton::common::TritonJson::Value raw_handle_json; + if (!register_request.Find("raw_handle", &raw_handle_json)) { + err = TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "Shared memory register request has no 'raw_handle' field"); + } else { + err = + raw_handle_json.MemberAsString("b64", &b64_handle, &b64_handle_len); } - if (err == nullptr) { - triton::common::TritonJson::Value register_request; - size_t buffer_len = evbuffer_get_length(req->buffer_in); - err = EVBufferToJson(®ister_request, v, &v_idx, buffer_len, n); - if (err == nullptr) { - const char* b64_handle = nullptr; - size_t b64_handle_len = 0; - triton::common::TritonJson::Value raw_handle_json; - if (!register_request.Find("raw_handle", &raw_handle_json)) { - err = TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - "Shared memory register request has no 'raw_handle' field"); - } else { - err = raw_handle_json.MemberAsString( - "b64", &b64_handle, &b64_handle_len); - } - uint64_t byte_size = 0; - if (err == nullptr) { - err = register_request.MemberAsUInt("byte_size", &byte_size); - } + uint64_t byte_size = 0; + if (err == nullptr) { + err = register_request.MemberAsUInt("byte_size", &byte_size); + } - uint64_t device_id = 0; - if (err == nullptr) { - err = register_request.MemberAsUInt("device_id", &device_id); - } + uint64_t device_id = 0; + if (err == nullptr) { + err = register_request.MemberAsUInt("device_id", &device_id); + } - if (err == nullptr) { - size_t decoded_size; - std::vector raw_handle; - RETURN_AND_RESPOND_IF_ERR( - req, DecodeBase64( - b64_handle, b64_handle_len, raw_handle, decoded_size, - "raw_handle")); + if (err == nullptr) { + size_t decoded_size; + std::vector raw_handle; + RETURN_AND_RESPOND_IF_ERR( + req, DecodeBase64( + b64_handle, b64_handle_len, raw_handle, decoded_size, + "raw_handle")); - if (decoded_size != sizeof(cudaIpcMemHandle_t)) { - err = TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - "'raw_handle' must be a valid base64 encoded " - "cudaIpcMemHandle_t"); - } else { - raw_handle.resize(sizeof(cudaIpcMemHandle_t)); - err = shm_manager_->RegisterCUDASharedMemory( - region_name.c_str(), - reinterpret_cast( - raw_handle.data()), - byte_size, device_id); - } - } + if (decoded_size != sizeof(cudaIpcMemHandle_t)) { + err = TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "'raw_handle' must be a valid base64 encoded " + "cudaIpcMemHandle_t"); + } else { + raw_handle.resize(sizeof(cudaIpcMemHandle_t)); + err = shm_manager_->RegisterCUDASharedMemory( + region_name.c_str(), + reinterpret_cast(raw_handle.data()), + byte_size, device_id); } } #else @@ -3028,25 +2865,50 @@ HTTPAPIServer::ParseJsonTritonRequestID( return nullptr; // Success } -// TODO: Can refactor other non-inference routes to re-use this helper instead. TRITONSERVER_Error* -HTTPAPIServer::EVRequestToJson( - evhtp_request_t* req, triton::common::TritonJson::Value* request_json_ptr) +HTTPAPIServer::EVRequestToJsonImpl( + evhtp_request_t* req, std::string_view request_kind, bool allows_empty_body, + triton::common::TritonJson::Value* request_json, size_t* buffer_len) { struct evbuffer_iovec* v = nullptr; int v_idx = 0; + std::vector v_vec; + int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0); if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); + try { + v_vec = std::vector(n); + } + catch (const std::bad_alloc& e) { + // Handle memory allocation failure + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + (std::string("Memory allocation failed for evbuffer: ") + e.what()) + .c_str()); + } + catch (const std::exception& e) { + // Catch any other std exceptions + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("Exception while creating evbuffer vector: ") + e.what()) + .c_str()); + } + + v = v_vec.data(); if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INTERNAL, - "Unexpected error getting request buffers"); + ("Unexpected error getting " + std::string(request_kind) + + " request buffers") + .c_str()); } } - size_t buffer_len = evbuffer_get_length(req->buffer_in); - RETURN_IF_ERR(EVBufferToJson(request_json_ptr, v, &v_idx, buffer_len, n)); + + *buffer_len = evbuffer_get_length(req->buffer_in); + if (allows_empty_body || *buffer_len > 0) { + RETURN_IF_ERR(EVBufferToJson(request_json, v, &v_idx, *buffer_len, n)); + } + return nullptr; // success } @@ -3063,11 +2925,29 @@ HTTPAPIServer::EVBufferToInput( // body. struct evbuffer_iovec* v = nullptr; int v_idx = 0; + std::vector v_vec; int n = evbuffer_peek(input_buffer, -1, NULL, NULL, 0); if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); + try { + v_vec = std::vector(n); + } + catch (const std::bad_alloc& e) { + // Handle memory allocation failure + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + (std::string("Memory allocation failed for evbuffer: ") + e.what()) + .c_str()); + } + catch (const std::exception& e) { + // Catch any other std exceptions + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("Exception while creating evbuffer vector: ") + e.what()) + .c_str()); + } + + v = v_vec.data(); if (evbuffer_peek(input_buffer, -1, NULL, v, n) != n) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INTERNAL, @@ -3120,10 +3000,30 @@ HTTPAPIServer::EVBufferToRawInput( } else { struct evbuffer_iovec* v = nullptr; int v_idx = 0; + std::vector v_vec; + int n = evbuffer_peek(input_buffer, -1, NULL, NULL, 0); if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); + try { + v_vec = std::vector(n); + } + catch (const std::bad_alloc& e) { + // Handle memory allocation failure + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + (std::string("Memory allocation failed for evbuffer: ") + e.what()) + .c_str()); + } + catch (const std::exception& e) { + // Catch any other std exceptions + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("Exception while creating evbuffer vector: ") + + e.what()) + .c_str()); + } + + v = v_vec.data(); if (evbuffer_peek(input_buffer, -1, NULL, v, n) != n) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INTERNAL, @@ -3155,6 +3055,62 @@ HTTPAPIServer::EVBufferToRawInput( return nullptr; // success } +TRITONSERVER_Error* +HTTPAPIServer::EVBufferToJson( + triton::common::TritonJson::Value* document, evbuffer_iovec* v, int* v_idx, + const size_t length, int n) +{ + size_t offset = 0, remaining_length = length; + char* json_base; + std::vector json_buffer; + + // No need to memcpy when number of iovecs is 1 + if ((n > 0) && (v[0].iov_len >= remaining_length)) { + json_base = static_cast(v[0].iov_base); + if (v[0].iov_len > remaining_length) { + v[0].iov_base = static_cast(json_base + remaining_length); + v[0].iov_len -= remaining_length; + remaining_length = 0; + } else if (v[0].iov_len == remaining_length) { + remaining_length = 0; + *v_idx += 1; + } + } else { + json_buffer.resize(length); + json_base = json_buffer.data(); + while ((remaining_length > 0) && (*v_idx < n)) { + char* base = static_cast(v[*v_idx].iov_base); + size_t base_size; + if (v[*v_idx].iov_len > remaining_length) { + base_size = remaining_length; + v[*v_idx].iov_base = static_cast(base + remaining_length); + v[*v_idx].iov_len -= remaining_length; + remaining_length = 0; + } else { + base_size = v[*v_idx].iov_len; + remaining_length -= v[*v_idx].iov_len; + *v_idx += 1; + } + + memcpy(json_base + offset, base, base_size); + offset += base_size; + } + } + + if (remaining_length != 0) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + std::string( + "unexpected size for request JSON, expecting " + + std::to_string(remaining_length) + " more bytes") + .c_str()); + } + + RETURN_IF_ERR(document->Parse(json_base, length)); + + return nullptr; // success +} + struct HeaderSearchPayload { HeaderSearchPayload( const re2::RE2& regex, TRITONSERVER_InferenceRequest* request) @@ -3429,7 +3385,8 @@ HTTPAPIServer::HandleGenerate( // arbitrary JSON message convoluted (added key is reference to a string and // thus the string must live as long as the JSON message). triton::common::TritonJson::Value request; - RETURN_AND_CALLBACK_IF_ERR(EVRequestToJson(req, &request), error_callback); + RETURN_AND_CALLBACK_IF_ERR( + EVRequestToJsonAllowsEmpty(req, "generate", &request), error_callback); RETURN_AND_CALLBACK_IF_ERR( ParseJsonTritonRequestID(request, irequest), error_callback); diff --git a/src/http_server.h b/src/http_server.h index 877327b4d5..c5ed0eb6de 100644 --- a/src/http_server.h +++ b/src/http_server.h @@ -607,9 +607,34 @@ class HTTPAPIServer : public HTTPServer { std::map* input_metadata, triton::common::TritonJson::Value* meta_data_root); + // Internal utility method for parsing evhtp request to JSON + // Should not be called directly - use EVRequestToJson or + // EVRequestToJsonAllowsEmpty instead + TRITONSERVER_Error* EVRequestToJsonImpl( + evhtp_request_t* req, std::string_view request_kind, + bool allows_empty_body, triton::common::TritonJson::Value* request_json, + size_t* buffer_len); + // Parses full evhtp request and its evbuffers into JSON. + TRITONSERVER_Error* EVRequestToJsonAllowsEmpty( + evhtp_request_t* req, std::string_view request_kind, + triton::common::TritonJson::Value* request_json) + { + size_t buffer_len = 0; + TRITONSERVER_Error* err = + EVRequestToJsonImpl(req, request_kind, true, request_json, &buffer_len); + return err; + } + TRITONSERVER_Error* EVRequestToJson( - evhtp_request_t* req, triton::common::TritonJson::Value* request_json); + evhtp_request_t* req, std::string_view request_kind, + triton::common::TritonJson::Value* request_json, size_t* buffer_len) + { + TRITONSERVER_Error* err = + EVRequestToJsonImpl(req, request_kind, false, request_json, buffer_len); + return err; + } + // Parses evhtp request buffers into Triton Inference Request. TRITONSERVER_Error* EVRequestToTritonRequest( evhtp_request_t* req, const std::string& model_name, @@ -622,6 +647,9 @@ class HTTPAPIServer : public HTTPServer { TRITONSERVER_Error* EVBufferToRawInput( const std::string& model_name, TRITONSERVER_InferenceRequest* irequest, evbuffer* input_buffer, InferRequestClass* infer_req); + TRITONSERVER_Error* EVBufferToJson( + triton::common::TritonJson::Value* document, evbuffer_iovec* v, + int* v_idx, const size_t length, int n); // Helpers for parsing JSON requests for Triton-specific fields diff --git a/src/sagemaker_server.cc b/src/sagemaker_server.cc index 52074f2b9d..49b8a3ba06 100644 --- a/src/sagemaker_server.cc +++ b/src/sagemaker_server.cc @@ -54,63 +54,6 @@ EVBufferAddErrorJson(evbuffer* buffer, TRITONSERVER_Error* err) evbuffer_add(buffer, buffer_json.Base(), buffer_json.Size()); } - -TRITONSERVER_Error* -EVBufferToJson( - triton::common::TritonJson::Value* document, evbuffer_iovec* v, int* v_idx, - const size_t length, int n) -{ - size_t offset = 0, remaining_length = length; - char* json_base; - std::vector json_buffer; - - // No need to memcpy when number of iovecs is 1 - if ((n > 0) && (v[0].iov_len >= remaining_length)) { - json_base = static_cast(v[0].iov_base); - if (v[0].iov_len > remaining_length) { - v[0].iov_base = static_cast(json_base + remaining_length); - v[0].iov_len -= remaining_length; - remaining_length = 0; - } else if (v[0].iov_len == remaining_length) { - remaining_length = 0; - *v_idx += 1; - } - } else { - json_buffer.resize(length); - json_base = json_buffer.data(); - while ((remaining_length > 0) && (*v_idx < n)) { - char* base = static_cast(v[*v_idx].iov_base); - size_t base_size; - if (v[*v_idx].iov_len > remaining_length) { - base_size = remaining_length; - v[*v_idx].iov_base = static_cast(base + remaining_length); - v[*v_idx].iov_len -= remaining_length; - remaining_length = 0; - } else { - base_size = v[*v_idx].iov_len; - remaining_length -= v[*v_idx].iov_len; - *v_idx += 1; - } - - memcpy(json_base + offset, base, base_size); - offset += base_size; - } - } - - if (remaining_length != 0) { - return TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - std::string( - "unexpected size for request JSON, expecting " + - std::to_string(remaining_length) + " more bytes") - .c_str()); - } - - RETURN_IF_ERR(document->Parse(json_base, length)); - - return nullptr; // success -} - } // namespace @@ -266,7 +209,9 @@ SagemakerAPIServer::Handle(evhtp_request_t* req) std::unordered_map parse_load_map; ParseSageMakerRequest(req, &parse_load_map, "load"); - SageMakerMMELoadModel(req, parse_load_map); + if (!parse_load_map.empty()) { + SageMakerMMELoadModel(req, parse_load_map); + } return; } break; @@ -320,29 +265,15 @@ SagemakerAPIServer::ParseSageMakerRequest( std::unordered_map* parse_map, const std::string& action) { - struct evbuffer_iovec* v = nullptr; - int v_idx = 0; - int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0); - if (n > 0) { - v = static_cast( - alloca(sizeof(struct evbuffer_iovec) * n)); - if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) { - HTTP_RESPOND_IF_ERR( - req, TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - "unexpected error getting load model request buffers")); - } - } + size_t buffer_len; + triton::common::TritonJson::Value request; + HTTP_RESPOND_IF_ERR( + req, EVRequestToJson(req, "load model", &request, &buffer_len)); std::string model_name_string; std::string url_string; - size_t buffer_len = evbuffer_get_length(req->buffer_in); if (buffer_len > 0) { - triton::common::TritonJson::Value request; - HTTP_RESPOND_IF_ERR( - req, EVBufferToJson(&request, v, &v_idx, buffer_len, n)); - triton::common::TritonJson::Value url; triton::common::TritonJson::Value model_name;