Propagate errors from parallel executor (#1262)

* Capture and return any errors from parallel execution so the behavior is equivalent to when the sequential executor is used.

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

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<OrtMutex> 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<Status> errors_;
 
   const bool& terminate_flag_;
   // TODO: Temporary threadpool for the executor.  This is a costly way to handle the problem.

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

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<Tensor>(0);
+      const int64_t* action = action_tensor.Data<int64_t>();
+
+      Status status = Status::OK();
+
+      switch (*action) {
+        case 0: {
+          // success
+          Tensor* Y = ctx->Output(0, action_tensor.Shape());
+          void* target = Y->MutableData<int64_t>();
+          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<int64_t>())
+        .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<CustomRegistry>();
+  std::vector<OpSchema> 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<int64_t>("action", {1}, {/*success*/ 0});
+    tester.AddOutput<int64_t>("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<int64_t>("action", {1}, {/*failure*/ 1});
+    tester.AddOutput<int64_t>("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<int64_t>("action", {1}, {/*exception*/ 2});
+    tester.AddOutput<int64_t>("action_out", {1}, {0});
+    tester.Run(OpTester::ExpectResult::kExpectFailure, "Throwing as action was 2", {kTensorrtExecutionProvider}, nullptr, nullptr, false);
+  }
+}
+}  // namespace test
+}  // namespace onnxruntime

onnxruntime/test/providers/provider_test_utils.cc

@@ -308,7 +308,7 @@ void OpTester::ExecuteModel(Model& model, InferenceSession& session_object, Expe
   std::vector<OrtValue> 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<std::string>& excluded_provider_types,
                    const RunOptions* run_options,
-                   std::vector<std::unique_ptr<IExecutionProvider>>* execution_providers) {
+                   std::vector<std::unique_ptr<IExecutionProvider>>* 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;
           }
         }
 

onnxruntime/test/providers/provider_test_utils.h

@@ -220,10 +220,8 @@ class OpTester {
   }
 
   void AddCustomOpRegistry(std::shared_ptr<CustomRegistry> registry) {
-    // need to do some static casting so we can easily use this later
-    //UTScustom_schema_registries_.push_back(std::static_pointer_cast<IOnnxRuntimeOpSchemaCollection>(registry->GetOpschemaRegistry()));
     custom_schema_registries_.push_back(registry->GetOpschemaRegistry());
-    custom_session_registries_.push_back(std::static_pointer_cast<CustomRegistry>(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<std::string>& excluded_provider_types = {},
            const RunOptions* run_options = nullptr,
-           std::vector<std::unique_ptr<IExecutionProvider>>* execution_providers = nullptr);
+           std::vector<std::unique_ptr<IExecutionProvider>>* execution_providers = nullptr,
+           bool sequential_execution = true);
 
   struct Data {
     onnxruntime::NodeArg def_;