Use flag instead
@@ -12,6 +12,14 @@ from cupyx.distributed import _store_actions
_DEFAULT_HOST = '127.0.0.1'
_DEFAULT_PORT = 13333
+_exit_mode = False
+
+
+@atexit.register
+def _exit():
+ global _exit_mode
+ _exit_mode = True
+
class ExceptionAwareProcess(multiprocessing.Process):
def __init__(self, *args, **kwargs):
@@ -45,7 +53,6 @@ class TCPStore:
# For implementing a barrier
self._lock = threading.Lock()
self._current_barrier = None
- atexit.register(self.stop)
def _set_process(self, process):
self._process = process
@@ -90,6 +97,8 @@ class TCPStore:
self._process = p
def stop(self):
+ if _exit_mode:
+ return # Prevent shutdown errors
if self._process is not None:
with self._run.get_lock():
self._run.value = 0
