diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index ae902322e5..7369925e9d 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -38,7 +38,7 @@ #include "core/platform/threadpool.h" #include "core/providers/cpu/controlflow/utils.h" #include "core/providers/cpu/cpu_execution_provider.h" -#include "core/flatbuffers/ort.fbs.h" +#include "core/flatbuffers/flatbuffers_utils.h" #ifdef USE_DML // TODO: This is necessary for the workaround in TransformGraph #include "core/providers/dml/DmlExecutionProvider/src/GraphTransformer.h" #endif @@ -486,7 +486,7 @@ common::Status InferenceSession::SaveToOrtFormat(const std::basic_string load_ort_format_mo } ORT_RETURN_IF_ERROR(load_ort_format_model_bytes()); + + // Verify the ort_format_model_bytes_ is a valid InferenceSessionBuffer before we access the data + flatbuffers::Verifier verifier(ort_format_model_bytes_.data(), ort_format_model_bytes_.size()); + ORT_RETURN_IF_NOT(fbs::VerifyInferenceSessionBuffer(verifier)); + const auto* fbs_session = fbs::GetInferenceSession(ort_format_model_bytes_.data()); ORT_RETURN_IF(nullptr == fbs_session, "InferenceSession is null. Invalid ORT format model."); @@ -1146,7 +1151,7 @@ common::Status InferenceSession::Initialize() { if ((has_explicit_type && model_type == "ORT") || (!has_explicit_type && - inference_session_utils::IsOrtFormatModel(session_options_.optimized_model_filepath))) { + experimental::utils::IsOrtFormatModel(session_options_.optimized_model_filepath))) { ORT_RETURN_IF_ERROR_SESSIONID_(SaveToOrtFormat(session_options_.optimized_model_filepath)); } else { ORT_RETURN_IF_ERROR_SESSIONID_(Model::Save(*model_, session_options_.optimized_model_filepath)); diff --git a/onnxruntime/test/framework/ort_model_only_test.cc b/onnxruntime/test/framework/ort_model_only_test.cc index c7000cd046..be81f8e9a7 100644 --- a/onnxruntime/test/framework/ort_model_only_test.cc +++ b/onnxruntime/test/framework/ort_model_only_test.cc @@ -50,6 +50,7 @@ struct OrtModelTestInfo { std::vector output_names; std::function&)> output_verifier; std::vector> configs; + bool run_use_buffer{false}; }; static void RunOrtModel(const OrtModelTestInfo& test_info) { @@ -58,8 +59,21 @@ static void RunOrtModel(const OrtModelTestInfo& test_info) { for (const auto& config : test_info.configs) so.AddConfigEntry(config.first.c_str(), config.second.c_str()); + std::vector model_data; InferenceSessionGetGraphWrapper session_object{so, GetEnvironment()}; - ASSERT_STATUS_OK(session_object.Load(test_info.model_filename)); // infer type from filename + if (test_info.run_use_buffer) { + // Load the file into a buffer and use the buffer to create inference session + size_t num_bytes = 0; + ASSERT_STATUS_OK(Env::Default().GetFileLength(test_info.model_filename.c_str(), num_bytes)); + model_data.resize(num_bytes); + std::ifstream bytes_stream(test_info.model_filename, std::ifstream::in | std::ifstream::binary); + bytes_stream.read(model_data.data(), num_bytes); + bytes_stream.close(); + ASSERT_STATUS_OK(session_object.Load(model_data.data(), static_cast(num_bytes))); + } else { + ASSERT_STATUS_OK(session_object.Load(test_info.model_filename)); // infer type from filename + } + ASSERT_STATUS_OK(session_object.Initialize()); std::vector fetches; @@ -304,8 +318,7 @@ TEST(OrtModelOnlyTests, SerializeToOrtFormatMLOps) { #endif // #if !defined(DISABLE_ML_OPS) #endif // #if !defined(ORT_MINIMAL_BUILD) -// test that we can deserialize and run a previously saved ORT format model -TEST(OrtModelOnlyTests, LoadOrtFormatModel) { +OrtModelTestInfo GetTestInfoForLoadOrtFormatModel() { OrtModelTestInfo test_info; test_info.model_filename = ORT_TSTR("testdata/ort_github_issue_4031.onnx.ort"); test_info.logid = "LoadOrtFormatModel"; @@ -323,13 +336,26 @@ TEST(OrtModelOnlyTests, LoadOrtFormatModel) { ASSERT_TRUE(output.Data()[0] == 125.f); }; + return test_info; +} + +// test that we can deserialize and run a previously saved ORT format model +TEST(OrtModelOnlyTests, LoadOrtFormatModel) { + OrtModelTestInfo test_info = GetTestInfoForLoadOrtFormatModel(); + RunOrtModel(test_info); +} + +// Load the model from a buffer instead of a file path +TEST(OrtModelOnlyTests, LoadOrtFormatModelFromBuffer) { + OrtModelTestInfo test_info = GetTestInfoForLoadOrtFormatModel(); + test_info.run_use_buffer = true; RunOrtModel(test_info); } #if !defined(DISABLE_ML_OPS) // test that we can deserialize and run a previously saved ORT format model // for a model with sequence and map outputs -TEST(OrtModelOnlyTests, LoadOrtFormatModelMLOps) { +OrtModelTestInfo GetTestInfoForLoadOrtFormatModelMLOps() { OrtModelTestInfo test_info; test_info.model_filename = ORT_TSTR("testdata/sklearn_bin_voting_classifier_soft.ort"); test_info.logid = "LoadOrtFormatModelMLOps"; @@ -363,8 +389,23 @@ TEST(OrtModelOnlyTests, LoadOrtFormatModelMLOps) { } }; + return test_info; +} + +// test that we can deserialize and run a previously saved ORT format model +// for a model with sequence and map outputs +TEST(OrtModelOnlyTests, LoadOrtFormatModelMLOps) { + OrtModelTestInfo test_info = GetTestInfoForLoadOrtFormatModelMLOps(); RunOrtModel(test_info); } + +// Load the model from a buffer instead of a file path +TEST(OrtModelOnlyTests, LoadOrtFormatModelMLOpsFromBuffer) { + OrtModelTestInfo test_info = GetTestInfoForLoadOrtFormatModelMLOps(); + test_info.run_use_buffer = true; + RunOrtModel(test_info); +} + #endif // !defined(DISABLE_ML_OPS) } // namespace test diff --git a/onnxruntime/test/testdata/ort_github_issue_4031.onnx.ort b/onnxruntime/test/testdata/ort_github_issue_4031.onnx.ort index 0fe711ee71..6d10929180 100644 Binary files a/onnxruntime/test/testdata/ort_github_issue_4031.onnx.ort and b/onnxruntime/test/testdata/ort_github_issue_4031.onnx.ort differ diff --git a/onnxruntime/test/testdata/sklearn_bin_voting_classifier_soft.ort b/onnxruntime/test/testdata/sklearn_bin_voting_classifier_soft.ort index 7eae709f43..eaef3784e7 100644 Binary files a/onnxruntime/test/testdata/sklearn_bin_voting_classifier_soft.ort and b/onnxruntime/test/testdata/sklearn_bin_voting_classifier_soft.ort differ