Replace AsyncProcess exit handler by weakref.finalize

distributed/process.py

@@ -81,15 +81,26 @@ def __init__(self, loop=None, target=None, name=None, args=(), kwargs={}):
                 dask.config.global_config,
             ),
         )
-        _dangling.add(self._process)
         self._name = self._process.name
+        self._proc_finalizer = weakref.finalize(self, self._get_finalizer())
         self._watch_q = PyQueue()
         self._exit_future = Future()
         self._exit_callback = None
         self._closed = False
 
         self._start_threads()
 
+    def _get_finalizer(self):
+        def finalize(proc=self._process, r=repr(self)):
+            if proc.is_alive():
+                try:
+                    logger.info("reaping stray process %s" % (proc,))
+                    proc.terminate()
+                except OSError:
+                    pass
+
+        return finalize
+
     def __repr__(self):
         return "<%s %s>" % (self.__class__.__name__, self._name)
 
@@ -118,8 +129,8 @@ def stop_thread(q):
             # We don't join the thread here as a finalizer can be called
             # asynchronously from anywhere
 
-        self._finalizer = weakref.finalize(self, stop_thread, q=self._watch_q)
-        self._finalizer.atexit = False
+        self._thread_finalizer = weakref.finalize(self, stop_thread, q=self._watch_q)
+        self._thread_finalizer.atexit = False
 
     def _on_exit(self, exitcode):
         # Called from the event loop when the child process exited
@@ -292,7 +303,7 @@ def close(self):
         immediately and does not ensure the child process has exited.
         """
         if not self._closed:
-            self._finalizer()
+            self._thread_finalizer()
             self._process = None
             self._closed = True
 
@@ -332,17 +343,3 @@ def daemon(self):
     @daemon.setter
     def daemon(self, value):
         self._process.daemon = value
-
-
-_dangling = weakref.WeakSet()
-
-
-@atexit.register
-def _cleanup_dangling():
-    for proc in list(_dangling):
-        if proc.is_alive():
-            try:
-                logger.info("reaping stray process %s" % (proc,))
-                proc.terminate()
-            except OSError:
-                pass

distributed/utils_test.py

@@ -44,7 +44,6 @@
 from .core import connect, rpc, CommClosedError, Status
 from .deploy import SpecCluster
 from .metrics import time
-from .process import _cleanup_dangling
 from .proctitle import enable_proctitle_on_children
 from .security import Security
 from .utils import (
@@ -1451,7 +1450,6 @@ def check_process_leak(check=True):
         else:
             assert not mp_context.active_children()
 
-    _cleanup_dangling()
     for proc in mp_context.active_children():
         proc.terminate()