Back out "Revert D20449887: [dt][caffe2] enable using smart exceptions in async nets" (pytorch#36172)

Summary:
Pull Request resolved: pytorch#36172

Original commit changeset: 3d7801613f86

D20449887 broke some OSS tests as the OSS export sync wasn't working correctly.

Test Plan:
Manually export latest version to OSS to trigger the tests

+ test plan in D20449887

verified onnx tests are passing in pytorch#36172

Reviewed By: andrewwdye

Differential Revision: D20902279

fbshipit-source-id: bc30fcc9f5cc8076f69a5d92675fd27455948372

Author: d4l3k

caffe2/core/event.h

@@ -107,8 +107,8 @@ class CAFFE2_API Event {
     event_resetter_[type_](this);
 #ifdef CAFFE2_USE_EXCEPTION_PTR
     caught_exception_ = nullptr;
-    exception_timestamp_ = 0;
 #endif // CAFFE2_USE_EXCEPTION_PTR
+    error_timestamp_ = 0;
   }
 
   const DeviceOption& GetDeviceOption() const {
@@ -126,6 +126,11 @@ class CAFFE2_API Event {
   }
 
   void SetFinished(const char* err_msg = nullptr) {
+    typedef std::chrono::high_resolution_clock clock;
+    error_timestamp_ = std::chrono::duration_cast<std::chrono::nanoseconds>(
+                           clock::now().time_since_epoch())
+                           .count();
+
     CAFFE_ENFORCE(event_finished_setter_[type_]);
     return event_finished_setter_[type_](this, err_msg);
   }
@@ -175,9 +180,6 @@ class CAFFE2_API Event {
 #ifdef CAFFE2_USE_EXCEPTION_PTR
     if (!caught_exception_) {
       caught_exception_ = std::current_exception();
-      typedef std::chrono::high_resolution_clock clock;
-      exception_timestamp_ =
-          clock::now().time_since_epoch() / std::chrono::milliseconds(1);
     }
     CAFFE_ENFORCE(caught_exception_, "No exception found");
 #else
@@ -199,13 +201,8 @@ class CAFFE2_API Event {
 #endif // CAFFE2_USE_EXCEPTION_PTR
   }
 
-  int64_t ExceptionTimestamp() const {
-#ifdef CAFFE2_USE_EXCEPTION_PTR
-    return exception_timestamp_;
-#else
-    VLOG(1) << "No support for exceptions in Event";
-    return 0;
-#endif // CAFFE2_USE_EXCEPTION_PTR
+  int64_t ErrorTimestamp() const {
+    return error_timestamp_;
   }
 
   void RethrowException() const {
@@ -229,20 +226,17 @@ class CAFFE2_API Event {
 
 #ifdef CAFFE2_USE_EXCEPTION_PTR
   std::exception_ptr caught_exception_;
-  int64_t exception_timestamp_{};
 #endif // CAFFE2_USE_EXCEPTION_PTR
+  int64_t error_timestamp_{};
 
   static EventCreateFunction event_creator_[MaxDeviceTypes];
   static EventRecordFunction event_recorder_[MaxDeviceTypes];
-  static EventWaitFunction event_waiter_[MaxDeviceTypes]
-                                        [MaxDeviceTypes];
+  static EventWaitFunction event_waiter_[MaxDeviceTypes][MaxDeviceTypes];
   static EventFinishFunction event_finisher_[MaxDeviceTypes];
 
   static EventQueryFunction event_querier_[MaxDeviceTypes];
-  static EventErrorMessageFunction
-      event_err_msg_getter_[MaxDeviceTypes];
-  static EventSetFinishedFunction
-      event_finished_setter_[MaxDeviceTypes];
+  static EventErrorMessageFunction event_err_msg_getter_[MaxDeviceTypes];
+  static EventSetFinishedFunction event_finished_setter_[MaxDeviceTypes];
   static EventResetFunction event_resetter_[MaxDeviceTypes];
 
   static EventSetCallbackFunction event_callback_setter_[MaxDeviceTypes];

caffe2/core/net_async_base.cc

@@ -119,14 +119,14 @@ bool AsyncNetBase::handleRunError() {
   for (int task_id = 0; task_id < tasksNum(); ++task_id) {
     if (event(task_id).HasException()) {
       if (first_exc_task_id >= 0) {
-        auto exc_ts = event(task_id).ExceptionTimestamp();
+        auto exc_ts = event(task_id).ErrorTimestamp();
         if (exc_ts < first_exc_ts) {
           first_exc_task_id = task_id;
           first_exc_ts = exc_ts;
         }
       } else {
         first_exc_task_id = task_id;
-        first_exc_ts = event(task_id).ExceptionTimestamp();
+        first_exc_ts = event(task_id).ErrorTimestamp();
       }
     }
   }
@@ -497,7 +497,13 @@ void AsyncNetBase::finalizeEvents() {
                 if (op != pending_op) {
                   try {
                     op->CancelAsyncCallback();
-                    op->event().SetFinished("Cancelled");
+
+                    // throw and catch exception to preserve stack trace
+                    try {
+                      throw AsyncNetCancelled();
+                    } catch (const AsyncNetCancelled& e) {
+                      op->event().SetFinishedWithException(e.what());
+                    }
                   } catch (const EnforceNotMet&) {
                     // ignore
                   }

caffe2/core/net_async_base.h

@@ -57,6 +57,12 @@ struct ExecutionOptions {
   bool run_root_tasks_inline_ = false;
 };
 
+struct CAFFE2_API AsyncNetCancelled : public std::exception {
+  const char* what() const noexcept override {
+    return "Cancelled";
+  }
+};
+
 class CAFFE2_API AsyncNetBase : public NetBase {
  public:
   AsyncNetBase(const std::shared_ptr<const NetDef>& net_def, Workspace* ws);

caffe2/core/net_async_scheduling.cc

@@ -280,7 +280,13 @@ void AsyncSchedulingNet::CancelAndFinishAsyncTasks() {
       // event, and in some other cases (CUDA)
       try {
         lastTaskOp(tid)->CancelAsyncCallback();
-        event(tid).SetFinished("Cancelled");
+
+        // throw and catch exception to preserve stack trace
+        try {
+          throw AsyncNetCancelled();
+        } catch (const AsyncNetCancelled& e) {
+          event(tid).SetFinishedWithException(e.what());
+        }
       } catch (const EnforceNotMet&) {
         // ignore
       }

caffe2/core/net_test.cc

@@ -787,8 +787,13 @@ TEST(NetTest, PendingOpsAndNetFailure) {
   Workspace ws;
   std::unique_ptr<NetBase> net(CreateNet(net_def, &ws));
 
-  // net is not stuck and returns false
-  ASSERT_FALSE(net->Run());
+  try {
+    // net is not stuck and returns false
+    ASSERT_FALSE(net->Run());
+  } catch (const caffe2::AsyncNetCancelled&) {
+    // Cancellation exception is fine since if the ops run concurrently the
+    // NotFinishingOp may be cancelled with an exception.
+  }
 }
 
 class AsyncErrorOp final : public Operator<CPUContext> {