Add exception handling support for waitall

docs/architecture/exception_handling.md

@@ -123,6 +123,3 @@ except mx.base.MXNetError as ex:
 d.asnumpy()
 ```
 
-### Limitation
-
-Rethrowing exceptions as part of `mx.nd.waitall` is not supported. So if your code executes a few operators and then calls `waitall` instead of `wait_to_read`/`asnumpy`, the exception will disappear. Please avoid waitalls in your code unless you are confident about your code not throwing exception in any scenario.

python/mxnet/ndarray/ndarray.py

@@ -157,13 +157,6 @@ def waitall():
     """Wait for all async operations to finish in MXNet.
 
     This function is used for benchmarking only.
-
-    .. warning::
-
-       If your code has exceptions, `waitall` can cause silent failures.
-       For this reason you should avoid `waitall` in your code.
-       Use it only if you are confident that your code is error free.
-       Then make sure you call `wait_to_read` on all outputs after `waitall`.
     """
     check_call(_LIB.MXNDArrayWaitAll())
 

src/engine/threaded_engine.cc

@@ -415,6 +415,25 @@ void ThreadedEngine::WaitForAll() {
   finished_cv_.wait(lock, [this]() {
       return pending_.load() == 0 || kill_.load();
     });
+  std::exception_ptr tmp;
+  if (!global_exception_refs_.empty()) {
+    // iterate through all exception refs
+    for (auto itr = global_exception_refs_.begin();
+         itr != global_exception_refs_.end(); ++itr) {
+      const std::shared_ptr<std::exception_ptr>& ptr = *itr;
+      // the first exception will be saved to be rethrown later
+      if (*ptr != nullptr && !tmp) {
+        tmp = *ptr;
+      }
+      // clear exceptions, WaitToRead following WaitForAll shouldn't throw
+      *ptr = nullptr;
+    }
+    // A waitall following a waitall shouldn't throw any exceptions
+    global_exception_refs_.clear();
+    if (tmp != nullptr) {
+      std::rethrow_exception(tmp);
+    }
+  }
 }
 
 inline void ThreadedEngine::OnComplete(ThreadedOpr* threaded_opr) {
@@ -428,6 +447,14 @@ inline void ThreadedEngine::OnComplete(ThreadedOpr* threaded_opr) {
   for (auto&& i : threaded_opr->mutable_vars) {
     if (threaded_opr->opr_exception && *threaded_opr->opr_exception) {
       i->var_exception = threaded_opr->opr_exception;
+      // add current operator exceptions to global exceptions if not already
+      // added
+      auto it = std::find(global_exception_refs_.begin(),
+                          global_exception_refs_.end(),
+                          threaded_opr->opr_exception);
+      if (it == global_exception_refs_.end()) {
+        global_exception_refs_.push_back(threaded_opr->opr_exception);
+      }
     }
     const bool debug_info = (engine_info_ && debug_wait_var_ == i);
     if (debug_info) {

src/engine/threaded_engine.h

@@ -177,7 +177,11 @@ class ThreadedVar final
   static std::atomic<std::size_t> counter;
   ~ThreadedVar() { LOG(INFO) << __func__ << " " << --counter; }
 #endif  // ENGINE_DEBUG
-  /*! \brief exception_ptr associated with the ThreadedVar */
+  /*!
+   * \brief exception_ptr associated with the ThreadedOpr
+   * cannot modify state of exception object since dereferencing
+   * exception_ptr is undefined behavior. Using shared_ptr to hold
+   * exception_ptr and overcome this limitation */
   std::shared_ptr<std::exception_ptr> var_exception;
 
  private:
@@ -254,7 +258,11 @@ struct ThreadedOpr final : public Opr,
   }
   // define possible debug information
   DEFINE_ENGINE_DEBUG_INFO(ThreadedOpr);
-  /*! \brief exception_ptr associated with the ThreadedOpr */
+  /*!
+   * \brief exception_ptr associated with the ThreadedOpr
+   * cannot modify state of exception object since dereferencing
+   * exception_ptr is undefined behavior. Using shared_ptr to hold
+   * exception_ptr and overcome this limitation */
   std::shared_ptr<std::exception_ptr> opr_exception;
 };  // struct ThreadedOpr
 
@@ -432,6 +440,9 @@ class ThreadedEngine : public Engine {
   };
   /*! thread local store for bulk */
   typedef dmlc::ThreadLocalStore<BulkStatus> BulkStatusStore;
+  /*! shared_ptr to exception_ptr, used for exception handling */
+  typedef std::shared_ptr<std::exception_ptr> ExceptionRef;
+
   /*!
    * \brief check if thee is duplication in const_vars and mutable_vars.
    * \param const_vars the variables to read from.
@@ -460,13 +471,25 @@ class ThreadedEngine : public Engine {
     for (auto&& i : threaded_opr->const_vars) {
       if (i->var_exception && *i->var_exception) {
         threaded_opr->opr_exception = i->var_exception;
+        auto it = std::find(global_exception_refs_.begin(),
+                            global_exception_refs_.end(),
+                            threaded_opr->opr_exception);
+        if (it == global_exception_refs_.end()) {
+            global_exception_refs_.push_back(threaded_opr->opr_exception);
+        }
         break;
       }
     }
     if (!(threaded_opr->opr_exception && *threaded_opr->opr_exception)) {
       for (auto&& i : threaded_opr->mutable_vars) {
         if (i->var_exception && *i->var_exception) {
           threaded_opr->opr_exception = i->var_exception;
+          auto it = std::find(global_exception_refs_.begin(),
+                              global_exception_refs_.end(),
+                              threaded_opr->opr_exception);
+          if (it == global_exception_refs_.end()) {
+            global_exception_refs_.push_back(threaded_opr->opr_exception);
+          }
           break;
         }
       }
@@ -542,6 +565,8 @@ class ThreadedEngine : public Engine {
    */
   std::mutex finished_m_;
   std::condition_variable finished_cv_;
+  /*! \brief global exception refs, which are rethrown when WaitForAll is called */
+  std::vector<ExceptionRef> global_exception_refs_;
 
   /*!
    * \brief Holding a shared_ptr to the object pool to prevent it from being destructed too early

tests/python/unittest/test_exc_handling.py

@@ -34,11 +34,11 @@ def imperative(exec_numpy=True):
             c.asnumpy()
 
     imperative(exec_numpy=False)
-    assert_raises(MXNetError, imperative, True)
+    assert_raises(MXNetError, imperative, exec_numpy=True)
 
 @with_seed()
 def test_exc_symbolic():
-    def symbolic(exec_backward=True):
+    def symbolic(exec_backward=True, waitall=True):
         x = mx.sym.Variable('x')
         y = mx.sym.Variable('y')
         z = mx.sym.Variable('z')
@@ -58,16 +58,25 @@ def symbolic(exec_backward=True):
         outputs = exec1.forward()
         if exec_backward:
             exec1.backward()
-            exec1.grad_arrays[0].asnumpy()
+            if waitall:
+                mx.nd.waitall()
+            else:
+                exec1.grad_arrays[0].asnumpy()
         else:
-            outputs[0].asnumpy()
+            if waitall:
+                mx.nd.waitall()
+            else:
+                outputs[0].asnumpy()
 
-    assert_raises(MXNetError, symbolic, False)
-    assert_raises(MXNetError, symbolic, True)
+    assert_raises(MXNetError, symbolic, exec_backward=False)
+    assert_raises(MXNetError, symbolic, exec_backward=True)
+
+    assert_raises(MXNetError, symbolic, exec_backward=False, waitall=True)
+    assert_raises(MXNetError, symbolic, exec_backward=True, waitall=True)
 
 @with_seed()
 def test_exc_gluon():
-    def gluon(exec_wait=True):
+    def gluon(exec_wait=True, waitall=False):
         model = nn.Sequential()
         model.add(nn.Dense(128, activation='tanh', in_units=10, flatten=False))
         model.add(nn.Dropout(1))
@@ -77,46 +86,83 @@ def gluon(exec_wait=True):
         y = model(x)
         model.collect_params().initialize(ctx=[default_context()])
         z = model(mx.nd.random.normal(10, -10, (32, 2, 10), ctx=default_context()))
-        if exec_wait:
+        if waitall:
+            mx.nd.waitall()
+        elif exec_wait:
             z.wait_to_read()
 
     gluon(exec_wait=False)
-    assert_raises(MXNetError, gluon, True)
+    assert_raises(MXNetError, gluon, exec_wait=True)
+
+    assert_raises(MXNetError, gluon, waitall=True)
 
 @with_seed()
 def test_exc_multiple_waits():
-    caught = False
-    try:
-        a = mx.nd.random.normal(0, -1, (2, 2)).copyto(default_context())
-        a.wait_to_read()
-    except MXNetError:
-        caught = True
-    assert caught, "No exception thrown"
-    try:
-        b = mx.nd.random.normal(0, -1, (2, 2)).copyto(default_context())
-        b.wait_to_read()
-    except MXNetError:
-        caught = True
-    assert caught, "No exception thrown"
+    def multiple_waits(waitall=False):
+        caught = False
+        try:
+            a = mx.nd.random.normal(0, -1, (2, 2)).copyto(default_context())
+            if waitall:
+                mx.nd.waitall()
+            else:
+                a.wait_to_read()
+        except MXNetError:
+            caught = True
+        assert caught, "No exception thrown"
+        try:
+            b = mx.nd.random.normal(0, -1, (2, 2)).copyto(default_context())
+            if waitall:
+                mx.nd.waitall()
+            else:
+                b.wait_to_read()
+        except MXNetError:
+            caught = True
+        assert caught, "No exception thrown"
+
+    multiple_waits(waitall=False)
+    multiple_waits(waitall=True)
 
 @with_seed()
 def test_exc_post_fail():
+    def post_fail(waitall=False):
+        caught = False
+        try:
+            a, b = mx.nd.random_normal(0, -1, (2, 2)).copyto(default_context())
+            if waitall:
+                mx.nd.waitall()
+            else:
+                a.asnumpy()
+        except MXNetError:
+            caught = True
+        assert caught, "No exception thrown"
+        b.asnumpy()
+    post_fail(waitall=False)
+    post_fail(waitall=True)
+
+@with_seed()
+def test_exc_mutable_var_fail():
+    def mutable_var_check(waitall=False):
+        a, b = mx.nd.random_normal(0, -1, (2, 2)).copyto(default_context())
+        a = mx.nd.dot(a, a)
+        if waitall:
+            mx.nd.waitall()
+        else:
+            a.asnumpy()
+    assert_raises(MXNetError, mutable_var_check, waitall=False)
+    assert_raises(MXNetError, mutable_var_check, waitall=True)
+
+@with_seed()
+def test_multiple_waitalls():
     caught = False
     try:
-        a, b = mx.nd.random_normal(0, -1, (2, 2)).copyto(default_context())
-        a.asnumpy()
+        a = mx.nd.random.normal(0, -1, (2, 2)).copyto(default_context())
+        mx.nd.waitall()
     except MXNetError:
         caught = True
     assert caught, "No exception thrown"
-    b.asnumpy()
+    mx.nd.waitall()
+
 
-@with_seed()
-def test_exc_mutable_var_fail():
-    def mutable_var_check():
-        a, b = mx.nd.random_normal(0, -1, (2, 2)).copyto(default_context())
-        a = mx.nd.dot(a, a)
-        a.asnumpy()
-    assert_raises(MXNetError, mutable_var_check)
 
 if __name__ == '__main__':
     import nose

tests/python/unittest/test_operator.py

@@ -6961,7 +6961,12 @@ def get_output_names_callback(name, arr):
 
         op_exe = op_sym.simple_bind(ctx=mx.current_context(), grad_req='null')
         op_exe.set_monitor_callback(get_output_names_callback, monitor_all=False)
-        op_exe.forward()
+        try:
+            op_exe.forward()
+            mx.nd.waitall()
+        except mx.base.MXNetError:
+            # skip errors since test is to check output names
+            pass
         for output_name, expected_name in zip(output_names, expected_names):
             assert output_name == expected_name