fix: fix async algorithm aborting (#78)

* tmp * . * update * format * remove unused
BaguaSys · Aug 30, 2021 · 6eaeb80 · 6eaeb80
1 parent dbf56fb
commit 6eaeb80
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 7 deletions.
diff --git a/bagua-core/bagua-core-internal/src/comm_ops/decentralized_full_precision_asynchronous.rs b/bagua-core/bagua-core-internal/src/comm_ops/decentralized_full_precision_asynchronous.rs
@@ -44,6 +44,11 @@ impl CommOpTrait for DecentralizedFullPrecisionAsynchronous {
                 false,
                 &mut |c, t| {
 
+             if c.check_abort() {
+                 tracing::debug!("process {} exits due to previous abortion", c.rank);
+                 return
+             }
+
              let start_time = std::time::Instant::now();
              tracing::debug!("async model average start");
 
@@ -103,22 +108,32 @@ impl CommOpTrait for DecentralizedFullPrecisionAsynchronous {
 
              let comm_ready_event = CUDA_EVENT_POOL.take().event;
 
-             unsafe {
+             let ret = unsafe {
                  cpp::cpp!([
                      comm_ready_event as "cudaEvent_t",
-                     comm_stream as "cudaStream_t"]
+                     comm_stream as "cudaStream_t"] -> bool as "bool"
                  {
                      CUDACHECK(cudaEventRecord(comm_ready_event, comm_stream));
-                     CUDACHECK(cudaEventSynchronize(comm_ready_event));
-                 });
+                     cudaError_t err = cudaEventSynchronize(comm_ready_event);
+                     if (err != cudaSuccess) {
+                         printf("Warning: Cuda error %s:%d '%s'\n", __FILE__,__LINE__,cudaGetErrorString(err));
+                         return false;
+                     }
+                     return true;
+                 })
+             };
+
+             if !ret {
+                 tracing::debug!("process {} early stopped due to communication failure", c.rank);
+                 c.set_abort();
+                 return
              }
 
              if c.check_abort() {
-                 tracing::debug!("async model average on process {} early stopped due to communicator abortion", c.rank);
+                 tracing::debug!("process {} early stopped due to communicator abortion", c.rank);
                  return
              }
 
-
              // do we need to wait default stream?
              unsafe {
                  cpp::cpp!([

diff --git a/bagua-core/bagua-core-internal/src/communicators/mod.rs b/bagua-core/bagua-core-internal/src/communicators/mod.rs
@@ -456,7 +456,7 @@ impl BaguaCommunicatorInner {
     pub fn abort(&self) {
         let communicator_ptr = self.comm_ptr;
 
-        self.aborted.store(true, Ordering::Relaxed);
+        self.set_abort();
 
         unsafe {
             cpp::cpp!([communicator_ptr as "Al::NCCLCommunicator*"]
@@ -466,6 +466,10 @@ impl BaguaCommunicatorInner {
         }
     }
 
+    pub fn set_abort(&self) {
+        self.aborted.store(true, Ordering::Relaxed);
+    }
+
     pub fn check_abort(&self) -> bool {
         self.aborted.load(Ordering::Relaxed)
     }