diff --git a/onnxruntime/core/graph/model.cc b/onnxruntime/core/graph/model.cc index f88427b5e9..e57e8f9b48 100644 --- a/onnxruntime/core/graph/model.cc +++ b/onnxruntime/core/graph/model.cc @@ -26,6 +26,8 @@ using namespace ONNX_NAMESPACE; using namespace onnxruntime; using namespace onnxruntime::common; +static constexpr int DEFAULT_PROTOBUF_BLOCK_SIZE = 4 * 1024 * 1024; + namespace onnxruntime { Model::Model(const std::string& graph_name, bool is_onnx_domain_only, @@ -237,6 +239,7 @@ Status Model::Load(std::istream& model_istream, ModelProto* p_model_proto) { if (!p_model_proto) { return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Null model_proto ptr."); } + google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream); const bool result = p_model_proto->ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof(); if (!result) { @@ -447,7 +450,13 @@ Status Model::Load(int fd, ONNX_NAMESPACE::ModelProto& model_proto) { } #if GOOGLE_PROTOBUF_VERSION >= 3002000 - FileInputStream input(fd); + size_t file_size = 0; + int block_size = -1; + Status st = Env::Default().GetFileLength(fd, file_size); + if (st.IsOK()) { + block_size = std::min(DEFAULT_PROTOBUF_BLOCK_SIZE, static_cast(file_size)); + } + FileInputStream input(fd, block_size); const bool result = model_proto.ParseFromZeroCopyStream(&input) && input.GetErrno() == 0; if (!result) { return Status(ONNXRUNTIME, INVALID_PROTOBUF, "Protobuf parsing failed."); diff --git a/onnxruntime/core/platform/env.h b/onnxruntime/core/platform/env.h index 221c194282..2959ab7349 100644 --- a/onnxruntime/core/platform/env.h +++ b/onnxruntime/core/platform/env.h @@ -134,6 +134,7 @@ class Env { * Gets the length of the specified file. */ virtual common::Status GetFileLength(_In_z_ const ORTCHAR_T* file_path, size_t& length) const = 0; + virtual common::Status GetFileLength(int fd, /*out*/ size_t& file_size) const = 0; /** * Copies the content of the file into the provided buffer. diff --git a/onnxruntime/core/platform/posix/env.cc b/onnxruntime/core/platform/posix/env.cc index d6b9119773..35f7b1db40 100644 --- a/onnxruntime/core/platform/posix/env.cc +++ b/onnxruntime/core/platform/posix/env.cc @@ -231,21 +231,30 @@ class PosixEnv : public Env { Status GetFileLength(const PathChar* file_path, size_t& length) const override { ScopedFileDescriptor file_descriptor{open(file_path, O_RDONLY)}; - if (file_descriptor.Get() < 0) { - return ReportSystemError("open", file_path); + return GetFileLength(file_descriptor.Get(), length); + } + + common::Status GetFileLength(int fd, /*out*/ size_t& file_size) const override { + using namespace common; + if (fd < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid fd was supplied: ", fd); } - struct stat stbuf; - if (fstat(file_descriptor.Get(), &stbuf) != 0) { - return ReportSystemError("fstat", file_path); + struct stat buf; + int rc = fstat(fd, &buf); + if (rc < 0) { + return ReportSystemError("fstat", ""); } - if (!S_ISREG(stbuf.st_mode)) { - return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, - "GetFileLength: input is not a regular file"); + if (buf.st_size < 0) { + return ORT_MAKE_STATUS(SYSTEM, FAIL, "Received negative size from stat call"); } - length = static_cast(stbuf.st_size); + if (static_cast(buf.st_size) > std::numeric_limits::max()) { + return ORT_MAKE_STATUS(SYSTEM, FAIL, "File is too large."); + } + + file_size = static_cast(buf.st_size); return Status::OK(); } diff --git a/onnxruntime/core/platform/windows/env.cc b/onnxruntime/core/platform/windows/env.cc index e00ad5ddbe..6ea293e858 100644 --- a/onnxruntime/core/platform/windows/env.cc +++ b/onnxruntime/core/platform/windows/env.cc @@ -193,6 +193,30 @@ class WindowsEnv : public Env { return Status::OK(); } + common::Status GetFileLength(int fd, /*out*/ size_t& file_size) const override { + using namespace common; + if (fd < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid fd was supplied: ", fd); + } + + struct _stat buf; + int rc = _fstat(fd, &buf); + if (rc < 0) { + return Status(SYSTEM, errno); + } + + if (buf.st_size < 0) { + return ORT_MAKE_STATUS(SYSTEM, FAIL, "Received negative size from stat call"); + } + + if (static_cast(buf.st_size) > std::numeric_limits::max()) { + return ORT_MAKE_STATUS(SYSTEM, FAIL, "File is too large."); + } + + file_size = static_cast(buf.st_size); + return Status::OK(); + } + Status ReadFileIntoBuffer(_In_z_ const ORTCHAR_T* const file_path, const FileOffsetType offset, const size_t length, const gsl::span buffer) const override { ORT_RETURN_IF_NOT(file_path); diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 2a64562282..702d2cd0ba 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -263,9 +263,8 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, const : graph_transformation_mgr_(session_options.max_num_graph_transformation_steps), logging_manager_(session_env.GetLoggingManager()), insert_cast_transformer_("CastFloat16Transformer") { - google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream); - const bool result = model_proto_.ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof(); - ORT_ENFORCE(result, "Could not parse model successfully while constructing the inference session"); + Status st = Model::Load(model_istream, &model_proto_); + ORT_ENFORCE(st.IsOK(), "Could not parse model successfully while constructing the inference session"); is_model_proto_parsed_ = true; // Finalize session options and initialize assets of this session instance ConstructorCommon(session_options, session_env); @@ -531,12 +530,9 @@ common::Status InferenceSession::Load(std::istream& model_istream) { auto loader = [this, &model_istream](std::shared_ptr& model) { ModelProto model_proto; - - google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream); - const bool result = model_proto.ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof(); - if (!result) { - return Status(common::ONNXRUNTIME, common::INVALID_PROTOBUF, - "Failed to load model because protobuf parsing failed."); + Status st = Model::Load(model_istream, &model_proto); + if (!st.IsOK()) { + return st; } #ifdef ENABLE_LANGUAGE_INTEROP_OPS LoadInterOp(model_proto, interop_domains_, [&](const char* msg) { LOGS(*session_logger_, WARNING) << msg; }); diff --git a/onnxruntime/test/providers/memcpy_test.cc b/onnxruntime/test/providers/memcpy_test.cc index 53fa70c2d4..db2480ca22 100644 --- a/onnxruntime/test/providers/memcpy_test.cc +++ b/onnxruntime/test/providers/memcpy_test.cc @@ -39,9 +39,8 @@ TEST(MemcpyTest, copy1) { ONNX_NAMESPACE::ModelProto mp; std::ifstream model_istream("testdata/matmul_1.onnx", std::ifstream::in | std::ifstream::binary); - google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream); - const bool result = mp.ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof(); - ASSERT_TRUE(result); + st = Model::Load(model_istream, &mp); + ASSERT_STATUS_OK(st); Model model(mp, nullptr, DefaultLoggingManager().DefaultLogger()); ASSERT_STATUS_OK(model.MainGraph().Resolve());