diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 5499698ce7..9b70f8154d 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -164,7 +164,7 @@ set(onnxruntime_test_framework_libs ) set(onnxruntime_test_server_libs - onnxruntime_test_utils + onnxruntime_test_utils_for_framework onnxruntime_test_utils_for_server ) @@ -581,14 +581,15 @@ if (onnxruntime_BUILD_SERVER) set_source_files_properties("${TEST_SRC_DIR}/server/unit_tests/json_handling_tests.cc" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter) set_source_files_properties("${TEST_SRC_DIR}/server/unit_tests/converter_tests.cc" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter) set_source_files_properties("${TEST_SRC_DIR}/server/unit_tests/util_tests.cc" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter) + set_source_files_properties("${TEST_SRC_DIR}/server/unit_tests/executor_test.cc" PROPERTIES COMPILE_FLAGS -Wno-unused-parameter) endif() endif() add_library(onnxruntime_test_utils_for_server ${onnxruntime_test_server_src}) - onnxruntime_add_include_to_target(onnxruntime_test_utils_for_server onnxruntime_test_utils gtest gmock gsl onnx onnx_proto server_proto) + onnxruntime_add_include_to_target(onnxruntime_test_utils_for_server onnxruntime_test_utils_for_framework gtest gmock gsl onnx onnx_proto server_proto) add_dependencies(onnxruntime_test_utils_for_server onnxruntime_server_lib onnxruntime_server_http_core_lib Boost ${onnxruntime_EXTERNAL_DEPENDENCIES}) target_include_directories(onnxruntime_test_utils_for_server PUBLIC ${Boost_INCLUDE_DIR} ${REPO_ROOT}/cmake/external/re2 ${CMAKE_CURRENT_BINARY_DIR}/onnx ${ONNXRUNTIME_ROOT}/server/http ${ONNXRUNTIME_ROOT}/server/http/core PRIVATE ${ONNXRUNTIME_ROOT} ) - target_link_libraries(onnxruntime_test_utils_for_server ${Boost_LIBRARIES} ${onnx_test_libs}) + target_link_libraries(onnxruntime_test_utils_for_server ${Boost_LIBRARIES}) AddTest( TARGET onnxruntime_server_tests diff --git a/onnxruntime/server/environment.cc b/onnxruntime/server/environment.cc index 5235cad266..2faed1f537 100644 --- a/onnxruntime/server/environment.cc +++ b/onnxruntime/server/environment.cc @@ -10,15 +10,17 @@ namespace onnxruntime { namespace server { -ServerEnvironment::ServerEnvironment(logging::Severity severity) : severity_(severity), +ServerEnvironment::ServerEnvironment(logging::Severity severity, logging::LoggingManager::InstanceType instance_type, bool env_init) : severity_(severity), logger_id_("ServerApp"), default_logging_manager_( std::unique_ptr{new LogSink{}}, severity, /* default_filter_user_data */ false, - logging::LoggingManager::InstanceType::Default, + instance_type, &logger_id_) { - auto status = onnxruntime::Environment::Create(runtime_environment_); + if (env_init) { + auto status = onnxruntime::Environment::Create(runtime_environment_); + } // The session initialization MUST BE AFTER environment creation session = std::make_unique(options_, &default_logging_manager_); diff --git a/onnxruntime/server/environment.h b/onnxruntime/server/environment.h index 4e0e408a7f..b0a983814d 100644 --- a/onnxruntime/server/environment.h +++ b/onnxruntime/server/environment.h @@ -17,7 +17,7 @@ namespace logging = logging; class ServerEnvironment { public: - explicit ServerEnvironment(logging::Severity severity); + explicit ServerEnvironment(logging::Severity severity, logging::LoggingManager::InstanceType instance_type = logging::LoggingManager::Default, bool env_init = true); ~ServerEnvironment() = default; ServerEnvironment(const ServerEnvironment&) = delete; diff --git a/onnxruntime/server/executor.cc b/onnxruntime/server/executor.cc index 48535129b9..734a6b6ac6 100644 --- a/onnxruntime/server/executor.cc +++ b/onnxruntime/server/executor.cc @@ -25,6 +25,7 @@ namespace server { namespace protobufutil = google::protobuf::util; protobufutil::Status Executor::SetMLValue(const onnx::TensorProto& input_tensor, + MemBufferArray& buffers, OrtAllocatorInfo* cpu_allocator_info, /* out */ MLValue& ml_value) { auto logger = env_->GetLogger(request_id_); @@ -36,12 +37,10 @@ protobufutil::Status Executor::SetMLValue(const onnx::TensorProto& input_tensor, return GenerateProtobufStatus(status, "GetSizeInBytesFromTensorProto() failed: " + status.ToString()); } - std::unique_ptr data(new char[cpu_tensor_length]); - memset(data.get(), 0, cpu_tensor_length); - OrtCallback deleter; + auto* buf = buffers.AllocNewBuffer(cpu_tensor_length); status = onnxruntime::utils::TensorProtoToMLValue(onnxruntime::Env::Default(), nullptr, input_tensor, - onnxruntime::MemBuffer(data.release(), cpu_tensor_length, *cpu_allocator_info), + onnxruntime::MemBuffer(buf, cpu_tensor_length, *cpu_allocator_info), ml_value, deleter); if (!status.IsOK()) { LOGS(*logger, ERROR) << "TensorProtoToMLValue() failed. Message: " << status.ToString(); @@ -51,12 +50,15 @@ protobufutil::Status Executor::SetMLValue(const onnx::TensorProto& input_tensor, return protobufutil::Status::OK; } -protobufutil::Status Executor::SetNameMLValueMap(onnxruntime::NameMLValMap& name_value_map, const onnxruntime::server::PredictRequest& request) { +protobufutil::Status Executor::SetNameMLValueMap(onnxruntime::NameMLValMap& name_value_map, + const onnxruntime::server::PredictRequest& request, + MemBufferArray& buffers) { auto logger = env_->GetLogger(request_id_); - OrtAllocatorInfo* cpu_allocator_info = nullptr; - auto ort_status = OrtCreateAllocatorInfo("Cpu", OrtArenaAllocator, 0, OrtMemTypeDefault, &cpu_allocator_info); - if (ort_status != nullptr || cpu_allocator_info == nullptr) { + OrtAllocatorInfo* allocator_info = nullptr; + auto ort_status = OrtCreateCpuAllocatorInfo(OrtArenaAllocator, OrtMemTypeDefault, &allocator_info); + + if (ort_status != nullptr || allocator_info == nullptr) { LOGS(*logger, ERROR) << "OrtCreateAllocatorInfo failed"; return protobufutil::Status(protobufutil::error::Code::RESOURCE_EXHAUSTED, "OrtCreateAllocatorInfo() failed"); } @@ -66,19 +68,22 @@ protobufutil::Status Executor::SetNameMLValueMap(onnxruntime::NameMLValMap& name using_raw_data_ = using_raw_data_ && input.second.has_raw_data(); MLValue ml_value; - auto status = SetMLValue(input.second, cpu_allocator_info, ml_value); + auto status = SetMLValue(input.second, buffers, allocator_info, ml_value); if (status != protobufutil::Status::OK) { + OrtReleaseAllocatorInfo(allocator_info); LOGS(*logger, ERROR) << "SetMLValue() failed! Input name: " << input.first; return status; } auto insertion_result = name_value_map.insert(std::make_pair(input.first, ml_value)); if (!insertion_result.second) { + OrtReleaseAllocatorInfo(allocator_info); LOGS(*logger, ERROR) << "SetNameMLValueMap() failed! Input name: " << input.first << " Trying to overwrite existing input value"; return protobufutil::Status(protobufutil::error::Code::INVALID_ARGUMENT, "SetNameMLValueMap() failed: Cannot have two inputs with the same name"); } } + OrtReleaseAllocatorInfo(allocator_info); return protobufutil::Status::OK; } @@ -89,8 +94,9 @@ protobufutil::Status Executor::Predict(const std::string& model_name, auto logger = env_->GetLogger(request_id_); // Convert PredictRequest to NameMLValMap + MemBufferArray buffer_array; onnxruntime::NameMLValMap name_ml_value_map{}; - auto conversion_status = SetNameMLValueMap(name_ml_value_map, request); + auto conversion_status = SetNameMLValueMap(name_ml_value_map, request, buffer_array); if (conversion_status != protobufutil::Status::OK) { return conversion_status; } diff --git a/onnxruntime/server/executor.h b/onnxruntime/server/executor.h index 47aaa49b84..179d8915fa 100644 --- a/onnxruntime/server/executor.h +++ b/onnxruntime/server/executor.h @@ -7,6 +7,7 @@ #include "environment.h" #include "predict.pb.h" +#include "util.h" namespace onnxruntime { namespace server { @@ -14,8 +15,8 @@ namespace server { class Executor { public: Executor(ServerEnvironment* server_env, std::string request_id) : env_(server_env), - request_id_(std::move(request_id)), - using_raw_data_(true) {} + request_id_(std::move(request_id)), + using_raw_data_(true) {} // Prediction method google::protobuf::util::Status Predict(const std::string& model_name, @@ -29,10 +30,13 @@ class Executor { bool using_raw_data_; google::protobuf::util::Status SetMLValue(const onnx::TensorProto& input_tensor, + MemBufferArray& buffers, OrtAllocatorInfo* cpu_allocator_info, /* out */ MLValue& ml_value); - google::protobuf::util::Status SetNameMLValueMap(onnxruntime::NameMLValMap& name_value_map, const onnxruntime::server::PredictRequest& request); + google::protobuf::util::Status SetNameMLValueMap(onnxruntime::NameMLValMap& name_value_map, + const onnxruntime::server::PredictRequest& request, + MemBufferArray& buffers); }; } // namespace server diff --git a/onnxruntime/server/http/predict_request_handler.cc b/onnxruntime/server/http/predict_request_handler.cc index cde66d32f8..30e259a77c 100644 --- a/onnxruntime/server/http/predict_request_handler.cc +++ b/onnxruntime/server/http/predict_request_handler.cc @@ -62,10 +62,9 @@ void Predict(const std::string& name, } // Run Prediction - protobufutil::Status status; Executor executor(env.get(), context.request_id); PredictResponse predict_response{}; - status = executor.Predict(name, version, predict_request, predict_response); + auto status = executor.Predict(name, version, predict_request, predict_response); if (!status.ok()) { GenerateErrorResponse(logger, GetHttpStatusCode((status)), status.error_message(), context); return; diff --git a/onnxruntime/server/util.h b/onnxruntime/server/util.h index e46a3a13fa..10b7aeb158 100644 --- a/onnxruntime/server/util.h +++ b/onnxruntime/server/util.h @@ -10,6 +10,34 @@ namespace onnxruntime { namespace server { +/** + * A RAII container for MLValue buffers + */ +class MemBufferArray { + public: + MemBufferArray() = default; + + uint8_t* AllocNewBuffer(size_t tensor_length) { + auto* data = new uint8_t[tensor_length]; + memset(data, 0, tensor_length); + buffers_.push_back(data); + return data; + } + + ~MemBufferArray() { + FreeBuffers(); + } + + private: + std::vector buffers_; + + void FreeBuffers() { + for (auto* buf : buffers_) { + delete[] buf; + } + } +}; + // Generate protobuf status from ONNX Runtime status google::protobuf::util::Status GenerateProtobufStatus(const onnxruntime::common::Status& onnx_status, const std::string& message); diff --git a/onnxruntime/test/server/unit_tests/executor_test.cc b/onnxruntime/test/server/unit_tests/executor_test.cc new file mode 100644 index 0000000000..75ab5d6817 --- /dev/null +++ b/onnxruntime/test/server/unit_tests/executor_test.cc @@ -0,0 +1,46 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include "gtest/gtest.h" + +#include "server/executor.h" +#include "server/http/json_handling.h" +#include "test/test_environment.h" + +namespace onnxruntime { +namespace server { +namespace test { + +TEST(ExecutorTests, TestMul_1) { + const static auto model_file = "testdata/mul_1.pb"; + const static auto input_json = R"({"inputs":{"X":{"dims":[3,2],"dataType":1,"floatData":[1,2,3,4,5,6]}},"outputFilter":["Y"]})"; + const static auto expected = R"({"outputs":{"Y":{"dims":["3","2"],"dataType":1,"floatData":[1,4,9,16,25,36]}}})"; + + onnxruntime::server::ServerEnvironment env(logging::Severity::kWARNING, logging::LoggingManager::InstanceType::Temporal, false); + + auto status = env.InitializeModel(model_file); + EXPECT_TRUE(status.IsOK()); + + status = env.GetSession()->Initialize(); + EXPECT_TRUE(status.IsOK()); + + onnxruntime::server::Executor executor(&env, "RequestId"); + onnxruntime::server::PredictRequest request{}; + onnxruntime::server::PredictResponse response{}; + + auto protostatus = onnxruntime::server::GetRequestFromJson(input_json, request); + EXPECT_TRUE(protostatus.ok()); + + auto prediction_res = executor.Predict("Name", "version", request, response); + EXPECT_TRUE(prediction_res.ok()); + + std::string body; + protostatus = GenerateResponseInJson(response, body); + EXPECT_EQ(expected, body); +} + +} // namespace test +} // namespace server +} // namespace onnxruntime