Side effect free & lazy multiprocessing context

Author: m1so

distributed/nanny.py

@@ -36,9 +36,9 @@
 from .utils import (
     TimeoutError,
     get_ip,
+    get_mp_context,
     json_load_robust,
     log_errors,
-    mp_context,
     parse_ports,
     silence_logging,
 )
@@ -662,6 +662,7 @@ async def start(self) -> Status:
             await self.running.wait()
             return self.status
 
+        mp_context = get_mp_context()
         self.init_result_q = init_q = mp_context.Queue()
         self.child_stop_q = mp_context.Queue()
         uid = uuid.uuid4().hex

distributed/process.py

@@ -11,7 +11,7 @@
 
 import dask
 
-from .utils import TimeoutError, mp_context
+from .utils import TimeoutError, get_mp_context
 
 logger = logging.getLogger(__name__)
 
@@ -65,6 +65,7 @@ def __init__(self, loop=None, target=None, name=None, args=(), kwargs={}):
         # monitor from the child and exit when the parent goes away unexpectedly
         # (for example due to SIGKILL). This variable is otherwise unused except
         # for the assignment here.
+        mp_context = get_mp_context()
         parent_alive_pipe, self._keep_child_alive = mp_context.Pipe(duplex=False)
 
         self._process = mp_context.Process(

distributed/tests/test_asyncprocess.py

@@ -16,7 +16,7 @@
 from distributed.compatibility import WINDOWS
 from distributed.metrics import time
 from distributed.process import AsyncProcess
-from distributed.utils import mp_context
+from distributed.utils import get_mp_context
 from distributed.utils_test import gen_test, nodebug, pristine_loop
 
 
@@ -53,6 +53,7 @@ def threads_info(q):
 @nodebug
 @gen_test()
 async def test_simple():
+    mp_context = get_mp_context()
     to_child = mp_context.Queue()
     from_child = mp_context.Queue()
 
@@ -143,6 +144,7 @@ async def test_simple():
 
 @gen_test()
 async def test_exitcode():
+    mp_context = get_mp_context()
     q = mp_context.Queue()
 
     proc = AsyncProcess(target=exit, kwargs={"q": q})
@@ -220,6 +222,7 @@ async def test_close():
 
 @gen_test()
 async def test_exit_callback():
+    mp_context = get_mp_context()
     to_child = mp_context.Queue()
     from_child = mp_context.Queue()
     evt = Event()
@@ -268,6 +271,7 @@ async def test_child_main_thread():
     """
     The main thread in the child should be called "MainThread".
     """
+    mp_context = get_mp_context()
     q = mp_context.Queue()
     proc = AsyncProcess(target=threads_info, args=(q,))
     await proc.start()
@@ -339,6 +343,7 @@ def _parent_process(child_pipe):
     be used to determine if it exited correctly."""
 
     async def parent_process_coroutine():
+        mp_context = get_mp_context()
         worker_ready = mp_context.Event()
 
         worker = AsyncProcess(target=_worker_process, args=(worker_ready, child_pipe))
@@ -377,6 +382,7 @@ def test_asyncprocess_child_teardown_on_parent_exit():
                   \________ <--   child_pipe   <-- ________/
     """
     # When child_pipe is closed, the children_alive pipe unblocks.
+    mp_context = get_mp_context()
     children_alive, child_pipe = mp_context.Pipe(duplex=False)
 
     try:

distributed/tests/test_client.py

@@ -68,7 +68,7 @@
     Scheduler,
 )
 from distributed.sizeof import sizeof
-from distributed.utils import is_valid_xml, mp_context, sync, tmp_text
+from distributed.utils import get_mp_context, is_valid_xml, sync, tmp_text
 from distributed.utils_test import (
     TaskStateMetadataPlugin,
     _UnhashableCallable,
@@ -2201,7 +2201,9 @@ def long_running_client_connection(address):
 
 @gen_cluster()
 async def test_cleanup_after_broken_client_connection(s, a, b):
-    proc = mp_context.Process(target=long_running_client_connection, args=(s.address,))
+    proc = get_mp_context().Process(
+        target=long_running_client_connection, args=(s.address,)
+    )
     proc.daemon = True
     proc.start()
 

distributed/tests/test_diskutils.py

@@ -14,7 +14,7 @@
 
 from distributed.diskutils import WorkSpace
 from distributed.metrics import time
-from distributed.utils import mp_context
+from distributed.utils import get_mp_context
 from distributed.utils_test import captured_logger
 
 
@@ -220,6 +220,7 @@ def test_workspace_concurrency(tmpdir, timeout, max_procs):
     deadlock happens.
     """
     base_dir = str(tmpdir)
+    mp_context = get_mp_context()
 
     err_q = mp_context.Queue()
     purged_q = mp_context.Queue()

distributed/utils.py

@@ -93,7 +93,7 @@ def _initialize_mp_context():
         return ctx
 
 
-mp_context = _initialize_mp_context()
+get_mp_context = toolz.memoize(_initialize_mp_context)
 
 
 def has_arg(func, argname):

distributed/utils_test.py

@@ -64,9 +64,9 @@
     _offload_executor,
     get_ip,
     get_ipv6,
+    get_mp_context,
     iscoroutinefunction,
     log_errors,
-    mp_context,
     reset_logger_locks,
     sync,
     thread_state,
@@ -626,6 +626,8 @@ def cluster(
     enable_proctitle_on_children()
 
     with clean(timeout=active_rpc_timeout, threads=False) as loop:
+        mp_context = get_mp_context()
+
         if nanny:
             _run_worker = run_nanny
         else:
@@ -1528,6 +1530,7 @@ def check_thread_leak():
 
 @contextmanager
 def check_process_leak(check=True):
+    mp_context = get_mp_context()
     for proc in mp_context.active_children():
         proc.terminate()