Repository: dask/dask

Issue:
Support configuring the `multiprocessing` context
`multiprocessing` supports different context, with different tradeoffs. Currently Dask only allows setting a custom *global* pool. This causes some problems:

1. It's not possible to use a different context with multiple pools: users are forced to use a single global pool if they want to use a different context.
2. It impedes changing the default context (see followup issue #3759).

Solution: support setting a multiprocessing context in the configuration, instead of just defaulting to the default `fork` context.



Discussion:
Going to start on this during the sprint.
+1!

On Sat, Jul 14, 2018 at 11:38 AM, Itamar Turner-Trauring <
notifications@github.com> wrote:

> Going to start on this during the sprint.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/3758#issuecomment-405034813>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszOn0huXwJtzYi2BANEDv5dmvi8t9ks5uGh5ogaJpZM4VP5Cn>
> .
>

Also cc @jcrist

On Sat, Jul 14, 2018 at 12:00 PM, Matthew Rocklin <mrocklin@anaconda.com>
wrote:

> +1!
>
> On Sat, Jul 14, 2018 at 11:38 AM, Itamar Turner-Trauring <
> notifications@github.com> wrote:
>
>> Going to start on this during the sprint.
>>
>> —
>> You are receiving this because you are subscribed to this thread.
>> Reply to this email directly, view it on GitHub
>> <https://github.com/dask/dask/issues/3758#issuecomment-405034813>, or mute
>> the thread
>> <https://github.com/notifications/unsubscribe-auth/AASszOn0huXwJtzYi2BANEDv5dmvi8t9ks5uGh5ogaJpZM4VP5Cn>
>> .
>>
>
>

Bit tricky in that one needs to create compatible Lock/Queue/etc. from the same context object, so need to figure out how to ensure that happens.

Resolving patch:
diff --git a/dask/multiprocessing.py b/dask/multiprocessing.py
--- a/dask/multiprocessing.py
+++ b/dask/multiprocessing.py
@@ -4,6 +4,7 @@
 import traceback
 import pickle
 import sys
+from warnings import warn
 
 import cloudpickle
 
@@ -120,6 +121,23 @@ def pack_exception(e, dumps):
     return result
 
 
+_CONTEXT_UNSUPPORTED = """\
+The 'multiprocessing.context' configuration option will be ignored on Python 2
+and on Windows, because they each only support a single context.
+"""
+
+
+def get_context():
+    """ Return the current multiprocessing context."""
+    if sys.platform == "win32" or sys.version_info.major == 2:
+        # Just do the default, since we can't change it:
+        if config.get("multiprocessing.context", None) is not None:
+            warn(_CONTEXT_UNSUPPORTED, UserWarning)
+        return multiprocessing
+    context_name = config.get("multiprocessing.context", None)
+    return multiprocessing.get_context(context_name)
+
+
 def get(dsk, keys, num_workers=None, func_loads=None, func_dumps=None,
         optimize_graph=True, **kwargs):
     """ Multiprocessed get function appropriate for Bags
@@ -143,8 +161,9 @@ def get(dsk, keys, num_workers=None, func_loads=None, func_dumps=None,
     """
     pool = config.get('pool', None)
     if pool is None:
-        pool = multiprocessing.Pool(num_workers,
-                                    initializer=initialize_worker_process)
+        context = get_context()
+        pool = context.Pool(num_workers,
+                            initializer=initialize_worker_process)
         cleanup = True
     else:
         cleanup = False
diff --git a/dask/utils.py b/dask/utils.py
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -12,7 +12,6 @@
 from importlib import import_module
 from numbers import Integral
 from threading import Lock
-import multiprocessing as mp
 import uuid
 import warnings
 from weakref import WeakValueDictionary
@@ -809,7 +808,7 @@ def get_scheduler_lock(get=None, collection=None, scheduler=None):
                                scheduler=scheduler)
 
     if actual_get == multiprocessing.get:
-        return mp.Manager().Lock()
+        return multiprocessing.get_context().Manager().Lock()
 
     return SerializableLock()
 


Test patch:
diff --git a/dask/tests/test_multiprocessing.py b/dask/tests/test_multiprocessing.py
--- a/dask/tests/test_multiprocessing.py
+++ b/dask/tests/test_multiprocessing.py
@@ -1,3 +1,6 @@
+from __future__ import absolute_import, division, print_function
+
+import sys
 import multiprocessing
 from operator import add
 import pickle
@@ -6,10 +9,11 @@
 import numpy as np
 
 import pytest
-
-from dask import compute, delayed
+from dask import compute, config, delayed
 from dask.context import set_options
-from dask.multiprocessing import get, _dumps, remote_exception
+from dask.multiprocessing import (
+    get, _dumps, get_context, remote_exception
+)
 from dask.utils_test import inc
 
 
@@ -134,3 +138,70 @@ def f():
         results, = compute([delayed(f, pure=False)() for i in range(N)])
 
     assert len(set(results)) == N
+
+
+@pytest.mark.skipif(sys.platform == 'win32',
+                    reason="Windows doesn't support different contexts")
+@pytest.mark.skipif(sys.version_info.major == 2,
+                    reason="Python 2 doesn't support different contexts")
+def test_custom_context_used_python3_posix():
+    """ The 'multiprocessing.context' config is used to create the pool.
+
+    We assume default is 'fork', and therefore test for 'spawn'.  If default
+    context is changed this test will need to be modified to be different than
+    that.
+    """
+    # We check for spawn by ensuring subprocess doesn't have modules only
+    # parent process should have:
+    def check_for_pytest():
+        import sys
+        return "FAKE_MODULE_FOR_TEST" in sys.modules
+
+    import sys
+    sys.modules["FAKE_MODULE_FOR_TEST"] = 1
+    try:
+        with config.set({"multiprocessing.context": "spawn"}):
+            result = get({"x": (check_for_pytest,)}, "x")
+        assert not result
+    finally:
+        del sys.modules["FAKE_MODULE_FOR_TEST"]
+
+
+@pytest.mark.skipif(sys.platform == 'win32',
+                    reason="Windows doesn't support different contexts")
+@pytest.mark.skipif(sys.version_info.major == 2,
+                    reason="Python 2 doesn't support different contexts")
+def test_get_context_using_python3_posix():
+    """ get_context() respects configuration.
+
+    If default context is changed this test will need to change too.
+    """
+    assert get_context() is multiprocessing.get_context(None)
+    with config.set({"multiprocessing.context": "forkserver"}):
+        assert get_context() is multiprocessing.get_context("forkserver")
+    with config.set({"multiprocessing.context": "spawn"}):
+        assert get_context() is multiprocessing.get_context("spawn")
+
+
+@pytest.mark.skipif(sys.platform != 'win32' and sys.version_info.major > 2,
+                    reason="Python 3 POSIX supports different contexts")
+def test_custom_context_ignored_elsewhere():
+    """ On Python 2/Windows, setting 'multiprocessing.context' doesn't explode.
+
+    Presumption is it's not used since unsupported, but mostly we care about
+    not breaking anything.
+    """
+    assert get({'x': (inc, 1)}, 'x') == 2
+    with pytest.warns(UserWarning):
+        with config.set({"multiprocessing.context": "forkserver"}):
+            assert get({'x': (inc, 1)}, 'x') == 2
+
+
+@pytest.mark.skipif(sys.platform != 'win32' and sys.version_info.major > 2,
+                    reason="Python 3 POSIX supports different contexts")
+def test_get_context_always_default():
+    """ On Python 2/Windows, get_context() always returns same context."""
+    assert get_context() is multiprocessing
+    with pytest.warns(UserWarning):
+        with config.set({"multiprocessing.context": "forkserver"}):
+            assert get_context() is multiprocessing