diff --git a/onnxruntime/core/framework/parallel_executor.cc b/onnxruntime/core/framework/parallel_executor.cc index 72fb4a71a0..9c9091987c 100644 --- a/onnxruntime/core/framework/parallel_executor.cc +++ b/onnxruntime/core/framework/parallel_executor.cc @@ -58,6 +58,26 @@ Status ParallelExecutor::Execute(const SessionState& session_state, const std::v while (out_standings_ > 0) complete_cv_.wait(lock); } + Status status = Status::OK(); + + if (!errors_.empty()) { + if (errors_.size() == 1) + status = errors_.front(); + else { + std::stringstream ss; + ss << "Multiple errors were found."; + for (const auto& s : errors_) { + ss << '\n' + << s; + } + + status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, ss.str()); + } + + LOGS(logger, ERROR) << status; + return status; + } + VLOGS(logger, 1) << "Fetching output."; // ExecutionFrame::Finalize will update 'fetches' with the final output ORT_RETURN_IF_ERROR(root_frame_->GetOutputs(fetches)); @@ -85,31 +105,24 @@ Status ParallelExecutor::Execute(const SessionState& session_state, const std::v if (f_profiler_enabled) { session_state.Profiler().EndTimeAndRecordEvent(profiling::SESSION_EVENT, "ParallelExecutor::Execute", tp); } + return Status::OK(); } -void ParallelExecutor::RunNodeAsync(size_t p_node_index, - const SessionState& session_state, - const logging::Logger& logger) { - try { - RunNodeAsyncInternal(p_node_index, session_state, logger); - } catch (...) { - FinishNodeRun(); - throw; - } -} - -void ParallelExecutor::RunNodeAsyncInternal(size_t p_node_index, - const SessionState& session_state, - const logging::Logger& logger) { +Status ParallelExecutor::RunNodeAsync(size_t p_node_index, + const SessionState& session_state, + const logging::Logger& logger) { LOGS(logger, INFO) << "Begin execution"; + Status status = Status::OK(); + size_t node_index = p_node_index; bool keep_running = true; auto graph_viewer = session_state.GetGraphViewer(); TimePoint sync_time_begin; TimePoint kernel_begin_time; bool f_profiler_enabled = session_state.Profiler().FEnabled(); + // Avoid context switching if possible. while (keep_running) { // TODO: Convert RunNodeAsync return Status. @@ -179,11 +192,14 @@ void ParallelExecutor::RunNodeAsyncInternal(size_t p_node_index, VLOGS(logger, 1) << "Computing kernel: " << p_op_kernel->Node().Name(); // Execute the kernel. - auto status = p_op_kernel->Compute(&op_kernel_context); + status = p_op_kernel->Compute(&op_kernel_context); if (!status.IsOK()) { - ORT_THROW("Compute failed for node: ", graph_viewer->GetNode(node_index)->Name(), - ". Error:", status.ErrorMessage()); + status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Compute failed for node: ", graph_viewer->GetNode(node_index)->Name(), + "\nError:", status); + break; } + if (f_profiler_enabled) { session_state.Profiler().EndTimeAndRecordEvent(profiling::NODE_EVENT, p_op_kernel->Node().Name() + "_kernel_time", @@ -247,21 +263,39 @@ void ParallelExecutor::RunNodeAsyncInternal(size_t p_node_index, } } - FinishNodeRun(); + return status; } void ParallelExecutor::EnqueueNode(size_t p_node_index, const SessionState& session_state, const logging::Logger& logger) { { std::unique_lock lock(complete_mutex_); + // if there are errors there's no point queuing more work + if (!errors_.empty()) + return; + out_standings_++; } executor_pool_->Schedule([this, p_node_index, &session_state, &logger]() { + auto create_exception_message = [p_node_index, &session_state](const std::exception* ex) { + const auto* node = session_state.GetGraphViewer()->GetNode(p_node_index); + + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Exception running nodes starting at ", node->OpType(), + " node '", node->Name(), "'. ", + ex ? ex->what() : "Unknown exception was caught by catch-all handler."); + }; + + Status status; try { - ParallelExecutor::RunNodeAsync(p_node_index, std::cref(session_state), std::cref(logger)); + status = ParallelExecutor::RunNodeAsync(p_node_index, std::cref(session_state), std::cref(logger)); + } catch (const std::exception& ex) { + status = create_exception_message(&ex); } catch (...) { // catch node processing failure exceptions here to prevent app crash. + status = create_exception_message(nullptr); } + + FinishNodeRun(status); }); } } // namespace onnxruntime diff --git a/onnxruntime/core/framework/parallel_executor.h b/onnxruntime/core/framework/parallel_executor.h index c489f5e31c..5f34309937 100644 --- a/onnxruntime/core/framework/parallel_executor.h +++ b/onnxruntime/core/framework/parallel_executor.h @@ -33,18 +33,20 @@ class ParallelExecutor : public IExecutor { private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ParallelExecutor); - void RunNodeAsync(size_t p_node_index, const SessionState& session_state, const logging::Logger& logger); - void RunNodeAsyncInternal(size_t p_node_index, const SessionState& session_state, const logging::Logger& logger); + Status RunNodeAsync(size_t p_node_index, const SessionState& session_state, const logging::Logger& logger); void EnqueueNode(size_t p_node_index, const SessionState& session_state, const logging::Logger& logger); - void FinishNodeRun() { + void FinishNodeRun(const Status& status) { bool finished = false; { //Because we have a mutex here, it's not possible another thread is doing the test("while (out_standings_ > 0)" std::lock_guard lock(complete_mutex_); finished = --out_standings_ == 0; + if (!status.IsOK()) + errors_.push_back(status); } + if (finished) { //std::cout << "all out standing nodes are completed." << std::endl; complete_cv_.notify_all(); @@ -57,6 +59,7 @@ class ParallelExecutor : public IExecutor { int out_standings_; //protected by complete_mutex_ OrtMutex complete_mutex_; OrtCondVar complete_cv_; + std::vector errors_; const bool& terminate_flag_; // TODO: Temporary threadpool for the executor. This is a costly way to handle the problem. diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 871a12b734..af0bbdfe66 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -241,7 +241,8 @@ TensorrtExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, graph_build.AddInitializedTensor(*(tensor.second)); } - ORT_ENFORCE(graph_build.Resolve().IsOK()); + auto status = graph_build.Resolve(); + ORT_ENFORCE(status.IsOK(), status); ONNX_NAMESPACE::ModelProto model_proto = model.ToProto(); model_proto.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); diff --git a/onnxruntime/test/framework/parallel_executor_test.cc b/onnxruntime/test/framework/parallel_executor_test.cc new file mode 100644 index 0000000000..a4ccd7f72d --- /dev/null +++ b/onnxruntime/test/framework/parallel_executor_test.cc @@ -0,0 +1,117 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/framework/data_types.h" +#include "core/framework/op_kernel.h" +#include "test/providers/provider_test_utils.h" +#include "test_utils.h" + +#include "gtest/gtest.h" + +using namespace ONNX_NAMESPACE; +using namespace onnxruntime::common; + +namespace onnxruntime { +namespace test { + +// Test kernel that will return success, or failure, or throw based on the input +struct TestOp { + static constexpr const char* OpName = "TestOp"; + static constexpr const char* OpDomain = "testing"; + + static ONNX_NAMESPACE::OpSchema OpSchema() { + ONNX_NAMESPACE::OpSchema schema; + schema.SetDoc("Return success, error, or throw based on the input.") + .SetName(OpName) + .SetDomain(OpDomain) + .SinceVersion(10) + .Input(0, "action", "Action to take.", "T", OpSchema::Single) + .Output(0, "action_out", "Return input as is", "T", OpSchema::Single) + .TypeConstraint("T", {"tensor(int64)"}, "Type of the action and values component"); + return schema; + } + + class OpKernelImpl final : public OpKernel { + public: + OpKernelImpl(const OpKernelInfo& info) : OpKernel{info} {} + + Status Compute(OpKernelContext* ctx) const override { + const Tensor& action_tensor = *ctx->Input(0); + const int64_t* action = action_tensor.Data(); + + Status status = Status::OK(); + + switch (*action) { + case 0: { + // success + Tensor* Y = ctx->Output(0, action_tensor.Shape()); + void* target = Y->MutableData(); + memcpy(target, action, action_tensor.Size()); + break; + } + case 1: { + // fail + status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Action was ", *action); + break; + } + default: { + ORT_THROW("Throwing as action was ", *action); + } + } + + return status; + } + }; + + static KernelDefBuilder KernelDef() { + KernelDefBuilder def; + def.SetName(OpName) + .SetDomain(OpDomain) + .SinceVersion(10) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .Provider(onnxruntime::kCpuExecutionProvider); + + return def; + } +}; + +// test that the status from TestOp is correctly returned from InferenceSession::Run +TEST(ParallelExecutor, TestStatusPropagation) { + auto registry = std::make_shared(); + std::vector schemas{TestOp::OpSchema()}; + Status status; + ASSERT_TRUE((status = registry->RegisterOpSet(schemas, TestOp::OpDomain, 10, 11)).IsOK()) << status; + KernelCreateFn kernel_create_fn = [](const OpKernelInfo& info) { return new typename TestOp::OpKernelImpl(info); }; + auto kernel_def = TestOp::KernelDef(); + ASSERT_TRUE((status = registry->RegisterCustomKernel(kernel_def, kernel_create_fn)).IsOK()) << status; + + { // test success + OpTester tester{"TestOp", 10, TestOp::OpDomain}; + tester.AddCustomOpRegistry(registry); + + tester.AddInput("action", {1}, {/*success*/ 0}); + tester.AddOutput("action_out", {1}, {0}); + // TensorRT doesn't handle a custom op. Possibly it should, but that would be a separate PR + tester.Run(OpTester::ExpectResult::kExpectSuccess, {}, {kTensorrtExecutionProvider}, nullptr, nullptr, false); + } + + { // test failure + OpTester tester{"TestOp", 10, TestOp::OpDomain}; + tester.AddCustomOpRegistry(registry); + + tester.AddInput("action", {1}, {/*failure*/ 1}); + tester.AddOutput("action_out", {1}, {0}); + tester.Run(OpTester::ExpectResult::kExpectFailure, "Action was 1", {kTensorrtExecutionProvider}, nullptr, nullptr, false); + } + + { // test exception + OpTester tester{"TestOp", 10, TestOp::OpDomain}; + tester.AddCustomOpRegistry(registry); + + tester.AddInput("action", {1}, {/*exception*/ 2}); + tester.AddOutput("action_out", {1}, {0}); + tester.Run(OpTester::ExpectResult::kExpectFailure, "Throwing as action was 2", {kTensorrtExecutionProvider}, nullptr, nullptr, false); + } +} +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/provider_test_utils.cc b/onnxruntime/test/providers/provider_test_utils.cc index e20b6e9630..b274754b0d 100644 --- a/onnxruntime/test/providers/provider_test_utils.cc +++ b/onnxruntime/test/providers/provider_test_utils.cc @@ -308,7 +308,7 @@ void OpTester::ExecuteModel(Model& model, InferenceSession& session_object, Expe std::vector fetches; status = session_object.Run(run_options ? *run_options : default_run_options, feeds, output_names, &fetches); if (status.IsOK()) { - EXPECT_TRUE(expect_result == ExpectResult::kExpectSuccess); + EXPECT_TRUE(expect_result == ExpectResult::kExpectSuccess) << "Expected failure but Run was successful"; if (expect_result == ExpectResult::kExpectFailure) { return; } @@ -364,7 +364,8 @@ void OpTester::Run(ExpectResult expect_result, const std::string& expected_failure_string, const std::unordered_set& excluded_provider_types, const RunOptions* run_options, - std::vector>* execution_providers) { + std::vector>* execution_providers, + bool sequential_execution) { try { #ifndef NDEBUG run_called_ = true; @@ -407,6 +408,7 @@ void OpTester::Run(ExpectResult expect_result, SessionOptions so; so.session_logid = op_; so.session_log_verbosity_level = 1; + so.enable_sequential_execution = sequential_execution; static const std::string all_provider_types[] = { kCpuExecutionProvider, @@ -480,7 +482,15 @@ void OpTester::Run(ExpectResult expect_result, const KernelCreateInfo* kci = reg->TryFindKernel(node, execution_provider->Type()); if (!kci) { valid = false; - break; + for (auto& custom_session_registry : custom_session_registries_) { + if (custom_session_registry->GetKernelRegistry()->TryFindKernel(node, execution_provider->Type())) { + valid = true; + break; + } + } + + if (!valid) + break; } } diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index 492fec9792..5f93eabbcc 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -220,10 +220,8 @@ class OpTester { } void AddCustomOpRegistry(std::shared_ptr registry) { - // need to do some static casting so we can easily use this later - //UTScustom_schema_registries_.push_back(std::static_pointer_cast(registry->GetOpschemaRegistry())); custom_schema_registries_.push_back(registry->GetOpschemaRegistry()); - custom_session_registries_.push_back(std::static_pointer_cast(registry)); + custom_session_registries_.push_back(registry); } void SetOutputAbsErr(const char* name, float v); @@ -244,7 +242,8 @@ class OpTester { void Run(ExpectResult expect_result = ExpectResult::kExpectSuccess, const std::string& expected_failure_string = "", const std::unordered_set& excluded_provider_types = {}, const RunOptions* run_options = nullptr, - std::vector>* execution_providers = nullptr); + std::vector>* execution_providers = nullptr, + bool sequential_execution = true); struct Data { onnxruntime::NodeArg def_;