Repository: tomMoral/loky

Issue:
BUG set_loky_pickler does not change the behavior in workers
With the current behavior, `loky` is always using `cloudpickle` to serialize the results from the computation in the worker. This should be changed by the `set_loky_pickler`. We have 3 options:

 - We can set it when we launch each worker. But this can lead to discrepancy between the workers and race conditions.
- We can set it with each task, bu this can bring a large overhead.
- We can also set it we create the `ReusableProcessExecutor` and kill the executor if `set_loky_pickler` has been called between two `get_reusable_executor`.

Discussion:

Patch:
diff --git a/loky/backend/reduction.py b/loky/backend/reduction.py
index 68cf141..0d40c5e 100644
--- a/loky/backend/reduction.py
+++ b/loky/backend/reduction.py
@@ -129,18 +129,26 @@ except ImportError:
     DEFAULT_ENV = "pickle"
 
 ENV_LOKY_PICKLER = os.environ.get("LOKY_PICKLER", DEFAULT_ENV)
-_LokyPickler = ENV_LOKY_PICKLER
+_LokyPickler = None
+_loky_pickler_name = None
 
 
 def set_loky_pickler(loky_pickler=None):
-    global _LokyPickler
+    global _LokyPickler, _loky_pickler_name
 
     if loky_pickler is None:
         loky_pickler = ENV_LOKY_PICKLER
 
     loky_pickler_cls = None
 
-    if loky_pickler in ["cloudpickle", "", None]:
+    # The default loky_pickler is cloudpickle
+    if loky_pickler in ["", None]:
+        loky_pickler = "cloudpickle"
+
+    if loky_pickler == _loky_pickler_name:
+        return
+
+    if loky_pickler == "cloudpickle":
         from cloudpickle import CloudPickler as loky_pickler_cls
     else:
         try:
@@ -198,11 +206,17 @@ def set_loky_pickler(loky_pickler=None):
                 self.dispatch_table[type] = reduce_func
 
     _LokyPickler = CustomizablePickler
+    _loky_pickler_name = loky_pickler
+
+
+def get_loky_pickler_name():
+    global _loky_pickler_name
+    return _loky_pickler_name
 
 
 def get_loky_pickler():
     global _LokyPickler
-    return _LokyPickler._loky_pickler_cls.__name__
+    return _LokyPickler
 
 
 # Set it to its default value
diff --git a/loky/backend/semaphore_tracker.py b/loky/backend/semaphore_tracker.py
index f494237..7d3f23e 100644
--- a/loky/backend/semaphore_tracker.py
+++ b/loky/backend/semaphore_tracker.py
@@ -23,10 +23,10 @@
 #
 
 import os
-import signal
 import sys
-import threading
+import signal
 import warnings
+import threading
 
 from . import spawn
 from multiprocessing import util
@@ -36,6 +36,9 @@ try:
 except ImportError:
     from .semlock import sem_unlink
 
+if sys.version_info < (3,):
+    BrokenPipeError = IOError
+
 __all__ = ['ensure_running', 'register', 'unregister']
 
 VERBOSE = False
diff --git a/loky/process_executor.py b/loky/process_executor.py
index 69d62a4..73672a8 100644
--- a/loky/process_executor.py
+++ b/loky/process_executor.py
@@ -80,6 +80,7 @@ from .backend.compat import wait
 from .backend.compat import set_cause
 from .backend.context import cpu_count
 from .backend.queues import Queue, SimpleQueue, Full
+from .backend.reduction import set_loky_pickler, get_loky_pickler_name
 from .backend.utils import recursive_terminate, get_exitcodes_terminated_worker
 
 try:
@@ -263,6 +264,13 @@ class _CallItem(object):
         self.args = args
         self.kwargs = kwargs
 
+        # Store the current loky_pickler so it is correctly set in the worker
+        self.loky_pickler = get_loky_pickler_name()
+
+    def __call__(self):
+        set_loky_pickler(self.loky_pickler)
+        return self.fn(*self.args, **self.kwargs)
+
     def __repr__(self):
         return "CallItem({}, {}, {}, {})".format(
             self.work_id, self.fn, self.args, self.kwargs)
@@ -407,7 +415,7 @@ def _process_worker(call_queue, result_queue, initializer, initargs,
             with worker_exit_lock:
                 return
         try:
-            r = call_item.fn(*call_item.args, **call_item.kwargs)
+            r = call_item()
         except BaseException as e:
             exc = _ExceptionWithTraceback(e)
             result_queue.put(_ResultItem(call_item.work_id, exception=exc))


Test patch:
diff --git a/tests/test_cloudpickle_wrapper.py b/tests/test_cloudpickle_wrapper.py
index 5e5cb14..ca8a774 100644
--- a/tests/test_cloudpickle_wrapper.py
+++ b/tests/test_cloudpickle_wrapper.py
@@ -136,48 +136,75 @@ class TestCloudpickleWrapper:
         finally:
             os.unlink(filename)
 
-    @pytest.mark.parametrize('loky_pickler, should_fail', [
-        (None, False), ("''", False), ("'cloudpickle'", False),
-        ("'pickle'", True)
-    ])
-    def test_set_loky_pickler(self, loky_pickler, should_fail):
+    @pytest.mark.parametrize('loky_pickler',
+                             [None, "''", "'cloudpickle'", "'pickle'"])
+    def test_set_loky_pickler(self, loky_pickler):
         # Test that the function set_loky_pickler correctly changes the pickler
         # used in loky.
         code = """if True:
-            from pickle import Pickler
             from loky import set_loky_pickler
-            from cloudpickle import CloudPickler
             from loky import get_reusable_executor
+            from loky import wrap_non_picklable_objects
             from loky.backend.reduction import get_loky_pickler
+            from loky.backend.reduction import get_loky_pickler_name
+
+            from pickle import Pickler
+            from cloudpickle import CloudPickler
 
+            # Check the default loky_pickler is cloudpickle
+            current_loky_pickler_name = get_loky_pickler_name()
+            assert current_loky_pickler_name == 'cloudpickle', (
+                "default got loky_pickler={{}}"
+                .format(current_loky_pickler_name))
+            assert issubclass(get_loky_pickler(), CloudPickler)
+
+            # Check that setting loky pickler to a value is working
             loky_pickler = {loky_pickler}
             set_loky_pickler(loky_pickler)
 
-            def test_func(x):
-                pass
-
-            current_loky_pickler = get_loky_pickler()
-            if loky_pickler in [None, "", "cloudpickle"]:
-                assert current_loky_pickler == 'CloudPickler', (
-                    "Expected CloudPickler and got {{}}"
-                    .format(current_loky_pickler))
+            if loky_pickler in [None, '', 'cloudpickle']:
+                expected_loky_pickler = CloudPickler
+                expected_loky_pickler_name = 'cloudpickle'
+            elif loky_pickler == 'pickle':
+                expected_loky_pickler = Pickler
+                expected_loky_pickler_name = 'pickle'
             else:
-                assert current_loky_pickler == Pickler.__name__, (
-                    "Expected {{}} and got {{}}"
-                    .format(Pickler.__name__, current_loky_pickler))
+                raise RuntimeError("unexpected value {{}} for loky_pickler"
+                                   .format(loky_pickler))
+
+
+            current_loky_pickler_name = get_loky_pickler_name()
+            assert current_loky_pickler_name == expected_loky_pickler_name, (
+                    "Expected 'pickle' and got {{}}"
+                    .format(current_loky_pickler_name))
+            assert issubclass(get_loky_pickler(), expected_loky_pickler)
 
             # Make sure that the default behavior is restored when
             # set_loky_pickler is used without arguments
             set_loky_pickler()
-            current_loky_pickler = get_loky_pickler()
-            assert current_loky_pickler == 'CloudPickler', (
-                "default got loky_pickler={{}}".format(current_loky_pickler))
+            current_loky_pickler_name = get_loky_pickler_name()
+            assert current_loky_pickler_name == 'cloudpickle', (
+                "default got loky_pickler={{}}"
+                .format(current_loky_pickler_name))
+            assert issubclass(get_loky_pickler(), CloudPickler)
 
-            # Test that the behavior expected. This should only fail when
-            # using the default pickle without the cloudpickle_wrapper.
+            # Check that the loky pickler in the workers is the correct one.
             set_loky_pickler(loky_pickler)
             e = get_reusable_executor()
-            e.submit(test_func, 42).result()
+            worker_loky_pickler = e.submit(get_loky_pickler_name).result()
+            assert worker_loky_pickler == expected_loky_pickler_name, (
+                "expected {{}} but got {{}} for the worker loky pickler"
+                .format(loky_pickler, worker_loky_pickler)
+            )
+
+
+            # Check that for cloudpickle, this does not fail and for pickle, it
+            # fails with the correct Error.
+            def test_func(x):
+                return x
+
+            e = get_reusable_executor()
+            assert e.submit(test_func, 42).result() == 42
             print("ok")
         """.format(loky_pickler=loky_pickler)
         cmd = [sys.executable]
@@ -187,8 +214,9 @@ class TestCloudpickleWrapper:
             with open(filename, mode='wb') as f:
                 f.write(code.encode('ascii'))
             cmd += [filename]
-            if should_fail:
-                with pytest.raises(ValueError, match=r'task has failed to un'):
+            if loky_pickler == "'pickle'":
+                match = r"(Can't get|has no) attribute 'test_func'"
+                with pytest.raises(ValueError, match=match):
                     check_subprocess_call(cmd, timeout=10)
             else:
                 check_subprocess_call(cmd, stdout_regex=r'ok', timeout=10)