Refine code

paddle/fluid/framework/details/threaded_ssa_graph_executor.cc

@@ -78,6 +78,10 @@ FeedFetchList ThreadedSSAGraphExecutor::Run(
     set.clear();
   };
 
+  // Clean run context
+  run_op_futures_.clear();
+  exception_.reset();
+
   // Step 3. Execution
   while (!pending_vars.empty()) {
     // 1. Run All Ready ops
@@ -98,14 +102,15 @@ FeedFetchList ThreadedSSAGraphExecutor::Run(
     if (timeout) {
       std::lock_guard<std::mutex> l(exception_mu_);
       if (exception_) {
+        for (auto &run_op_future : run_op_futures_) {
+          run_op_future.wait();
+        }
         std::exception *exp = exception_.get();
         if (dynamic_cast<platform::EOFException *>(exp)) {
           auto e = *static_cast<platform::EOFException *>(exp);
-          exception_.reset();
           throw e;
         } else if (dynamic_cast<platform::EnforceNotMet *>(exp)) {
           auto e = *static_cast<platform::EnforceNotMet *>(exp);
-          exception_.reset();
           throw e;
         } else {
           LOG(FATAL) << "Unknown exception.";
@@ -222,7 +227,7 @@ void ThreadedSSAGraphExecutor::RunOp(
     }
   };
   if (pool_) {
-    pool_->enqueue(op_run);
+    run_op_futures_.emplace_back(pool_->enqueue(op_run));
   } else {
     op_run();
   }

paddle/fluid/framework/details/threaded_ssa_graph_executor.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include <deque>
+#include <list>
 #include <string>
 #include <unordered_set>
 #include <utility>
@@ -77,6 +78,8 @@ class ThreadedSSAGraphExecutor : public SSAGraphExecutor {
 
  private:
   ExecutionStrategy strategy_;
+  // use std::list because clear(), push_back, and for_each are O(1)
+  std::list<std::future<void>> run_op_futures_;
 };
 
 }  // namespace details