saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-26 19:52:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix server's executor memory leak and add tests (#984)

Browse source 
* Fixes memory leaks in executor and adds executor tests

* Remove logging test from executor_test.cc

* Uses RAII to allocate and freee MLValue buffers

* Change private member buffer to buffer_
This commit is contained in:
tmccrmck 2019-05-09 22:51:57 -07:00 committed by Pranav Sharma
parent 3fb24f6409
commit af7090c25f
8 changed files with 108 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
cmake/onnxruntime_unittests.cmake
View file

@ -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

8
onnxruntime/server/environment.cc
View file

@ -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<logging::ISink>{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<onnxruntime::InferenceSession>(options_, &default_logging_manager_);

2
onnxruntime/server/environment.h
View file

@ -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;

26
onnxruntime/server/executor.cc
View file

@ -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<char[]> 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;
}

10
onnxruntime/server/executor.h
View file

@ -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

3
onnxruntime/server/http/predict_request_handler.cc
View file

@ -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;

28
onnxruntime/server/util.h
View file

@ -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<uint8_t*> 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);

46
onnxruntime/test/server/unit_tests/executor_test.cc Normal file
View file

@ -0,0 +1,46 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <iostream>
#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