Merge pull request #263 from python-adaptive/support-loky

add support for loky

Author: basnijholt

adaptive/runner.py

@@ -12,9 +12,13 @@
 from adaptive.notebook_integration import in_ipynb, live_info, live_plot
 
 try:
-    import ipyparallel
+    if sys.version_info < (3, 8):
+        # XXX: remove when ipyparallel 6.2.5 is released
+        import ipyparallel
 
-    with_ipyparallel = True
+        with_ipyparallel = True
+    else:
+        with_ipyparallel = False
 except ModuleNotFoundError:
     with_ipyparallel = False
 
@@ -32,6 +36,13 @@
 except ModuleNotFoundError:
     with_mpi4py = False
 
+try:
+    import loky
+
+    with_loky = True
+except ModuleNotFoundError:
+    with_loky = False
+
 with suppress(ModuleNotFoundError):
     import uvloop
 
@@ -232,10 +243,13 @@ def _remove_unfinished(self):
 
     def _cleanup(self):
         if self.shutdown_executor:
-            # XXX: temporary set wait=True for Python 3.7
+            # XXX: temporary set wait=True because of a bug with Python ≥3.7
+            # and loky in any Python version.
             # see https://github.com/python-adaptive/adaptive/issues/156
             # and https://github.com/python-adaptive/adaptive/pull/164
-            self.executor.shutdown(wait=True if sys.version_info >= (3, 7) else False)
+            # and https://bugs.python.org/issue36281
+            # and https://github.com/joblib/loky/issues/241
+            self.executor.shutdown(wait=True)
         self.end_time = time.time()
 
     @property
@@ -269,7 +283,8 @@ class BlockingRunner(BaseRunner):
         the learner as its sole argument, and return True when we should
         stop requesting more points.
     executor : `concurrent.futures.Executor`, `distributed.Client`,\
-               `mpi4py.futures.MPIPoolExecutor`, or `ipyparallel.Client`, optional
+               `mpi4py.futures.MPIPoolExecutor`, `ipyparallel.Client` or\
+               `loky.get_reusable_executor`, optional
         The executor in which to evaluate the function to be learned.
         If not provided, a new `~concurrent.futures.ProcessPoolExecutor`.
     ntasks : int, optional
@@ -386,7 +401,8 @@ class AsyncRunner(BaseRunner):
         stop requesting more points. If not provided, the runner will run
         forever, or until ``self.task.cancel()`` is called.
     executor : `concurrent.futures.Executor`, `distributed.Client`,\
-               `mpi4py.futures.MPIPoolExecutor`, or `ipyparallel.Client`, optional
+               `mpi4py.futures.MPIPoolExecutor`, `ipyparallel.Client` or\
+               `loky.get_reusable_executor`, optional
         The executor in which to evaluate the function to be learned.
         If not provided, a new `~concurrent.futures.ProcessPoolExecutor`.
     ntasks : int, optional
@@ -740,9 +756,16 @@ def shutdown(self, wait=True):
         pass
 
 
+def _default_executor():
+    if with_loky:
+        return loky.get_reusable_executor()
+    else:
+        return concurrent.ProcessPoolExecutor()
+
+
 def _ensure_executor(executor):
     if executor is None:
-        executor = concurrent.ProcessPoolExecutor()
+        executor = _default_executor()
 
     if isinstance(executor, concurrent.Executor):
         return executor
@@ -765,6 +788,8 @@ def _get_ncores(ex):
         ex, (concurrent.ProcessPoolExecutor, concurrent.ThreadPoolExecutor)
     ):
         return ex._max_workers  # not public API!
+    elif with_loky and isinstance(ex, loky.reusable_executor._ReusablePoolExecutor):
+        return ex._max_workers  # not public API!
     elif isinstance(ex, SequentialExecutor):
         return 1
     elif with_distributed and isinstance(ex, distributed.cfexecutor.ClientExecutor):

adaptive/tests/test_runner.py

@@ -15,6 +15,7 @@
     stop_after,
     with_distributed,
     with_ipyparallel,
+    with_loky,
 )
 
 
@@ -91,6 +92,13 @@ def ipyparallel_executor():
         raise RuntimeError("Could not stop ipcluster")
 
 
+@pytest.fixture(scope="session")
+def loky_executor():
+    import loky
+
+    return loky.get_reusable_executor()
+
+
 def linear(x):
     return x
 
@@ -134,3 +142,12 @@ def test_distributed_executor():
     BlockingRunner(learner, trivial_goal, executor=client)
     client.shutdown()
     assert learner.npoints > 0
+
+
+@pytest.mark.skipif(not with_loky, reason="loky not installed")
+def test_loky_executor(loky_executor):
+    learner = Learner1D(lambda x: x, (-1, 1))
+    BlockingRunner(
+        learner, trivial_goal, executor=loky_executor, shutdown_executor=True
+    )
+    assert learner.npoints > 0

docs/source/conf.py

@@ -139,6 +139,7 @@
     "holoviews": ("https://holoviews.org/", None),
     "ipyparallel": ("https://ipyparallel.readthedocs.io/en/stable/", None),
     "scipy": ("https://docs.scipy.org/doc/scipy/reference", None),
+    "loky": ("https://loky.readthedocs.io/en/stable/", None),
 }
 
 

docs/source/tutorial/tutorial.parallelism.rst

@@ -116,3 +116,21 @@ How you call MPI might depend on your specific queuing system, with SLURM for ex
     #SBATCH --ntasks 100
 
     srun -n $SLURM_NTASKS --mpi=pmi2 ~/miniconda3/envs/py37_min/bin/python -m mpi4py.futures run_learner.py
+
+`loky.get_reusable_executor`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This executor is basically a powered-up version of `~concurrent.futures.ProcessPoolExecutor`, check its `documentation <https://loky.readthedocs.io/>`_.
+Among other things, it allows to *reuse* the executor and uses ``cloudpickle`` for serialization.
+This means you can even learn closures, lambdas, or other functions that are not picklable with `pickle`.
+
+.. code:: python
+
+    from loky import get_reusable_executor
+    ex = get_reusable_executor()
+
+    f = lambda x: x
+    learner = adaptive.Learner1D(f, bounds=(-1, 1))
+
+    runner = adaptive.Runner(learner, goal=lambda l: l.loss() < 0.01, executor=ex)
+    runner.live_info()

environment.yml

@@ -12,6 +12,7 @@ dependencies:
   - ipyparallel
   - distributed
   - ipykernel>=4.8*
+  - loky
   - jupyter_client>=5.2.2
   - ipywidgets
   - scikit-optimize

setup.py

@@ -53,6 +53,7 @@ def get_version_and_cmdclass(package_name):
     "other": [
         "ipyparallel",
         "distributed",
+        "loky",
         "scikit-optimize",
         "wexpect" if os.name == "nt" else "pexpect",
     ],