Repository: tomMoral/loky

Issue:
initializer for get_reusable_executor
It is possible to set `initializer` for `ProcessPoolExecutor` but I could not find how to pass an initializer to `get_reusable_executor`. 

Discussion:

Patch:
diff --git a/CHANGES.md b/CHANGES.md
index 033eccb..e03a520 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,6 +1,7 @@
 ### 2.2.0.dev0 (in development)
 
 - Fix bad interaction between `max_workers=None` and `reuse='auto'` (#132).
+- Add initializer for `get_reusable_executor` (#134)
 
 ### 2.1.2 - 2018-06-04 - Release Highligths
 
diff --git a/examples/deadlock_results.py b/examples/deadlocks/deadlock_results.py
similarity index 95%
rename from examples/deadlock_results.py
rename to examples/deadlocks/deadlock_results.py
index dc11ba6..d1de4cb 100644
--- a/examples/deadlock_results.py
+++ b/examples/deadlocks/deadlock_results.py
@@ -1,4 +1,5 @@
 """Deadlock with unpickling error for the result
+=================================================
 
 This example highlights the fact that the ProcessPoolExecutor implementation
 from concurrent.futures is not robust to pickling error (at least in versions
diff --git a/examples/deadlock_timeout_shutdown.py b/examples/deadlocks/deadlock_timeout_shutdown.py
similarity index 95%
rename from examples/deadlock_timeout_shutdown.py
rename to examples/deadlocks/deadlock_timeout_shutdown.py
index 419467b..ba5a7e0 100644
--- a/examples/deadlock_timeout_shutdown.py
+++ b/examples/deadlocks/deadlock_timeout_shutdown.py
@@ -1,4 +1,5 @@
 """Deadlock with pickling errors
+================================
 
 This example highlights the fact that the ProcessPoolExecutor implementation
 from concurrent.futures is not robust to pickling error (at least in versions
diff --git a/examples/reuseable_executor.py b/examples/reuseable_executor.py
new file mode 100644
index 0000000..36833c3
--- /dev/null
+++ b/examples/reuseable_executor.py
@@ -0,0 +1,42 @@
+"""
+Advanced Executor Setup
+=======================
+
+It is possible to reuse an executor even if it has a complicated setup. When
+using the parameter `reuse=True`, the executor is resized if needed but the
+arguments stay the same.
+"""
+from time import sleep
+from loky import get_reusable_executor
+
+INITIALIZER_STATUS = "uninitialized"
+
+
+def initializer(self, x):
+    global INITIALIZER_STATUS
+    INITIALIZER_STATUS = x
+
+
+def test_initializer(self, delay=0):
+    print(delay)
+    sleep(delay)
+
+    global INITIALIZER_STATUS
+    return INITIALIZER_STATUS
+
+
+executor = get_reusable_executor(
+    max_workers=2, initializer=initializer, initargs=('initialized',))
+
+assert executor.submit(test_initializer).result() == 'initialized'
+
+# With reuse=True, the executor use the same initializer
+executor = get_reusable_executor(max_workers=4, reuse=True)
+for x in executor.map(test_initializer, delay=.5):
+    assert x == 'initialized'
+
+# With reuse='auto', the initializer is not used anymore as a new executor is
+# created.
+executor = get_reusable_executor(max_workers=4)
+for x in executor.map(test_initializer, delay=.1):
+    assert x == 'uninitialized'
diff --git a/loky/reusable_executor.py b/loky/reusable_executor.py
index 881ec83..30b217f 100644
--- a/loky/reusable_executor.py
+++ b/loky/reusable_executor.py
@@ -38,8 +38,9 @@ def _get_next_executor_id():
 
 
 def get_reusable_executor(max_workers=None, context=None, timeout=10,
-                          kill_workers=False, job_reducers=None,
-                          result_reducers=None, reuse="auto"):
+                          kill_workers=False, reuse="auto",
+                          job_reducers=None, result_reducers=None,
+                          initializer=None, initargs=()):
     """Return the current ReusableExectutor instance.
 
     Start a new instance if it has not been started already or if the previous
@@ -69,9 +70,12 @@ def get_reusable_executor(max_workers=None, context=None, timeout=10,
 
     The ``job_reducers`` and ``result_reducers`` are used to customize the
     pickling of tasks and results send to the executor.
+
+    When provided, the ``initializer`` is run first in newly spawned
+    processes with argument ``initargs``.
     """
     with _executor_lock:
-        global _executor, _executor_args
+        global _executor, _executor_kwargs
         executor = _executor
 
         if max_workers is None:
@@ -84,25 +88,27 @@ def get_reusable_executor(max_workers=None, context=None, timeout=10,
                 "max_workers must be greater than 0, got {}."
                 .format(max_workers))
 
-        args = dict(context=context, timeout=timeout,
-                    job_reducers=job_reducers, result_reducers=result_reducers)
         if isinstance(context, STRING_TYPE):
             context = get_context(context)
         if context is not None and context.get_start_method() == "fork":
             raise ValueError("Cannot use reusable executor with the 'fork' "
                              "context")
+
+        kwargs = dict(context=context, timeout=timeout,
+                      job_reducers=job_reducers,
+                      result_reducers=result_reducers,
+                      initializer=initializer, initargs=initargs)
         if executor is None:
             mp.util.debug("Create a executor with max_workers={}."
                           .format(max_workers))
             executor_id = _get_next_executor_id()
-            _executor_args = args
+            _executor_kwargs = kwargs
             _executor = executor = _ReusablePoolExecutor(
-                _executor_lock, max_workers=max_workers, context=context,
-                timeout=timeout, executor_id=executor_id,
-                job_reducers=job_reducers, result_reducers=result_reducers)
+                _executor_lock, max_workers=max_workers,
+                executor_id=executor_id, **kwargs)
         else:
             if reuse == 'auto':
-                reuse = args == _executor_args
+                reuse = kwargs == _executor_kwargs
             if (executor._flags.broken or executor._flags.shutdown
                     or not reuse):
                 if executor._flags.broken:
@@ -116,10 +122,10 @@ def get_reusable_executor(max_workers=None, context=None, timeout=10,
                     "previous instance cannot be reused ({})."
                     .format(max_workers, reason))
                 executor.shutdown(wait=True, kill_workers=kill_workers)
-                _executor = executor = _executor_args = None
+                _executor = executor = _executor_kwargs = None
                 # Recursive call to build a new instance
                 return get_reusable_executor(max_workers=max_workers,
-                                             **args)
+                                             **kwargs)
             else:
                 mp.util.debug("Reusing existing executor with max_workers={}."
                               .format(executor._max_workers))
@@ -131,11 +137,11 @@ def get_reusable_executor(max_workers=None, context=None, timeout=10,
 class _ReusablePoolExecutor(ProcessPoolExecutor):
     def __init__(self, submit_resize_lock, max_workers=None, context=None,
                  timeout=None, executor_id=0, job_reducers=None,
-                 result_reducers=None):
+                 result_reducers=None, initializer=None, initargs=()):
         super(_ReusablePoolExecutor, self).__init__(
             max_workers=max_workers, context=context, timeout=timeout,
-            job_reducers=job_reducers,
-            result_reducers=result_reducers)
+            job_reducers=job_reducers, result_reducers=result_reducers,
+            initializer=initializer, initargs=initargs)
         self.executor_id = executor_id
         self._submit_resize_lock = submit_resize_lock
 


Test patch:
diff --git a/tests/test_reusable_executor.py b/tests/test_reusable_executor.py
index 76e1aef..85385bf 100644
--- a/tests/test_reusable_executor.py
+++ b/tests/test_reusable_executor.py
@@ -43,6 +43,9 @@ except ImportError:
     pass
 
 
+INITIALIZER_STATUS = 'uninitialized'
+
+
 def clean_warning_registry():
     """Safe way to reset warnings."""
     warnings.resetwarnings()
@@ -651,3 +654,37 @@ def test_reusable_executor_reuse_true():
 
     executor4 = get_reusable_executor()
     assert executor4 is executor3
+
+
+class TestExecutorInitializer(ReusableExecutorMixin):
+    def _initializer(self, x):
+        global INITIALIZER_STATUS
+        INITIALIZER_STATUS = x
+
+    def _test_initializer(self, delay=0):
+        sleep(delay)
+
+        global INITIALIZER_STATUS
+        return INITIALIZER_STATUS
+
+    def test_reusable_initializer(self):
+        executor = get_reusable_executor(
+            max_workers=2, initializer=self._initializer, initargs=('done',))
+
+        assert executor.submit(self._test_initializer).result() == 'done'
+
+        # when the initializer change, the executor is re-spawned
+        executor = get_reusable_executor(
+            max_workers=2, initializer=self._initializer, initargs=(42,))
+
+        assert executor.submit(self._test_initializer).result() == 42
+
+        # With reuse=True, the executor use the same initializer
+        executor = get_reusable_executor(max_workers=4, reuse=True)
+        for x in executor.map(self._test_initializer, delay=.1):
+            assert x == 42
+
+        # With reuse='auto', the initializer is not used anymore
+        executor = get_reusable_executor(max_workers=4)
+        for x in executor.map(self._test_initializer, delay=.1):
+            assert x == 'uninitialized'