Merge pull request #7952 from kerrylu/main

Provide .stop() method for cupyx.distributed._Backend
cupy · Oct 23, 2023 · 2b86893 · 2b86893
2 parents 6bd7021 + 6b40d0f
commit 2b86893
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/cupyx/distributed/_comm.py b/cupyx/distributed/_comm.py
@@ -61,3 +61,7 @@ def all_to_all(self, in_array, out_array, stream=None):
     @abc.abstractmethod
     def barrier(self):
         pass
+
+    def stop(self):
+        if self.rank == 0:
+            self._store.stop()
diff --git a/cupyx/distributed/_store.py b/cupyx/distributed/_store.py
@@ -106,7 +106,8 @@ def stop(self):
         if self._process is not None:
             with self._run.get_lock():
                 self._run.value = 0
-            self._process.join()
+            if self._process.is_alive():
+                self._process.join()
 
 
 class TCPStoreProxy: