Improve factory interfaces

Change-Id: I012e2de0635b3d6701ac441520ae96aba6be8899

Author: adrien-berchet

.gitignore

@@ -81,6 +81,7 @@ coverage.xml
 *.cover
 .hypothesis/
 reports
+dask-worker-space
 
 # Translations
 *.mo

bluepyparallel/evaluator.py

@@ -73,14 +73,15 @@ def evaluate(
     if func_kwargs is None:
         func_kwargs = {}
 
+    # Drop exception column if present
+    if "exception" in df.columns:
+        logger.warning("The 'exception' column is going to be replaced")
+        df = df.drop(columns=["exception"])
+
     # Shallow copy the given DataFrame to add internal rows
     to_evaluate = df.copy()
     task_ids = to_evaluate.index
 
-    if "exception" in to_evaluate.columns:
-        logger.warning("The exception column is going to be replaced")
-        to_evaluate = to_evaluate.drop(columns=["exception"])
-
     # Set default new columns
     if new_columns is None:
         new_columns = [["data", ""]]

bluepyparallel/parallel.py

@@ -2,7 +2,6 @@
 import logging
 import multiprocessing
 import os
-import time
 from abc import abstractmethod
 from collections.abc import Iterator
 from functools import partial
@@ -36,14 +35,19 @@ class ParallelFactory:
     _CHUNK_SIZE = "PARALLEL_CHUNK_SIZE"
 
     # pylint: disable=unused-argument
-    def __init__(self, batch_size=None, chunk_size=None, **kwargs):
+    def __init__(self, batch_size=None, chunk_size=None):
         self.batch_size = batch_size or int(os.getenv(self._BATCH_SIZE, "0")) or None
         L.info("Using %s=%s", self._BATCH_SIZE, self.batch_size)
 
         self.chunk_size = batch_size or int(os.getenv(self._CHUNK_SIZE, "0")) or None
         L.info("Using %s=%s", self._CHUNK_SIZE, self.chunk_size)
 
-        self.nb_processes = 1
+        if not hasattr(self, "nb_processes"):
+            self.nb_processes = 1
+
+    def __del__(self):
+        """Call the shutdown method."""
+        self.shutdown()
 
     @abstractmethod
     def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
@@ -62,7 +66,12 @@ def _with_batches(self, mapper, func, iterable, batch_size=None):
 
         batch_size = batch_size or self.batch_size
         if batch_size is not None:
-            iterables = np.array_split(iterable, len(iterable) // min(batch_size, len(iterable)))
+            iterables = [
+                _iterable.tolist()
+                for _iterable in np.array_split(
+                    iterable, len(iterable) // min(batch_size, len(iterable))
+                )
+            ]
         else:
             iterables = [iterable]
 
@@ -113,17 +122,17 @@ class MultiprocessingFactory(ParallelFactory):
 
     _CHUNKSIZE = "PARALLEL_CHUNKSIZE"
 
-    def __init__(self, processes=None, **kwargs):
+    def __init__(self, batch_size=None, chunk_size=None, processes=None, **kwargs):
         """Initialize multiprocessing factory."""
 
-        super().__init__(**kwargs)
+        super().__init__(batch_size, chunk_size)
 
-        self.pool = NestedPool(processes=processes)
         self.nb_processes = processes or os.cpu_count()
+        self.pool = NestedPool(processes=self.nb_processes, **kwargs)
 
     def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
         """Get a NestedPool."""
-        self._chunksize_to_kwargs(chunk_size, kwargs)
+        self._chunksize_to_kwargs(chunk_size, kwargs, label="chunksize")
 
         def _mapper(func, iterable):
             return self._with_batches(
@@ -144,29 +153,25 @@ class IPyParallelFactory(ParallelFactory):
 
     _IPYTHON_PROFILE = "IPYTHON_PROFILE"
 
-    def __init__(self, **kwargs):
+    def __init__(self, batch_size=None, chunk_size=None, profile=None, **kwargs):
         """Initialize the ipyparallel factory."""
-
-        super().__init__(**kwargs)
-        self.rc = None
-        self.nb_processes = 1
-
-    def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
-        """Get an ipyparallel mapper using the profile name provided."""
-        profile = os.getenv(self._IPYTHON_PROFILE, None)
+        profile = profile or os.getenv(self._IPYTHON_PROFILE, None)
         L.debug("Using %s=%s", self._IPYTHON_PROFILE, profile)
-        self.rc = ipyparallel.Client(profile=profile)
+        self.rc = ipyparallel.Client(profile=profile, **kwargs)
         self.nb_processes = len(self.rc.ids)
-        lview = self.rc.load_balanced_view()
+        self.lview = self.rc.load_balanced_view()
+        super().__init__(batch_size, chunk_size)
 
+    def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
+        """Get an ipyparallel mapper using the profile name provided."""
         if "ordered" not in kwargs:
             kwargs["ordered"] = False
 
         self._chunksize_to_kwargs(chunk_size, kwargs)
 
         def _mapper(func, iterable):
             return self._with_batches(
-                partial(lview.imap, **kwargs), func, iterable, batch_size=batch_size
+                partial(self.lview.imap, **kwargs), func, iterable, batch_size=batch_size
             )
 
         return _mapper
@@ -182,27 +187,34 @@ class DaskFactory(ParallelFactory):
 
     _SCHEDULER_PATH = "PARALLEL_DASK_SCHEDULER_PATH"
 
-    def __init__(self, **kwargs):
+    def __init__(
+        self, batch_size=None, chunk_size=None, scheduler_file=None, address=None, **kwargs
+    ):
         """Initialize the dask factory."""
-        dask_scheduler_path = os.getenv(self._SCHEDULER_PATH)
+        dask_scheduler_path = scheduler_file or os.getenv(self._SCHEDULER_PATH)
+        self.interactive = True
         if dask_scheduler_path:
-            self.interactive = True
             L.info("Connecting dask_mpi with scheduler %s", dask_scheduler_path)
-            self.client = dask.distributed.Client(scheduler_file=dask_scheduler_path)
-        else:
+        if address:
+            L.info("Connecting dask_mpi with address %s", address)
+        if not dask_scheduler_path and not address:
             self.interactive = False
-            dask_mpi.initialize()
             L.info("Starting dask_mpi...")
-            self.client = dask.distributed.Client()
+            dask_mpi.initialize()
+        self.client = dask.distributed.Client(
+            address=address,
+            scheduler_file=dask_scheduler_path,
+            **kwargs,
+        )
         self.nb_processes = len(self.client.scheduler_info()["workers"])
-        super().__init__(**kwargs)
+        super().__init__(batch_size, chunk_size)
 
     def shutdown(self):
-        """Retire the workers on the scheduler."""
+        """Close the scheduler and the cluster if it was created by the factory."""
+        cluster = self.client.cluster
+        self.client.close()
         if not self.interactive:
-            time.sleep(1)
-            self.client.retire_workers()
-            self.client = None
+            cluster.close()
 
     def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
         """Get a Dask mapper."""

examples/large_computation.py

@@ -4,11 +4,12 @@
 import time
 from bluepyparallel import evaluate
 from bluepyparallel import init_parallel_factory
-from data_validation_framework.util import apply_to_df
 
 
 def func(row):
     """Trivial computation"""
+    time.sleep(5)
+
     if row["data"] in [1, 3]:
         raise ValueError(f"The value {row['data']} is forbidden")
     else:
@@ -20,7 +21,7 @@ def func(row):
     batch_size = int(sys.argv[2]) if len(sys.argv) >= 3 else None
     chunk_size = int(sys.argv[3]) if len(sys.argv) >= 4 else None
     df = pd.DataFrame()
-    df["data"] = np.arange(1e6)
+    df["data"] = np.arange(200)
 
     parallel_factory = init_parallel_factory(parallel_lib, batch_size=batch_size)
     df = evaluate(

tests/__init__.py

@@ -1,7 +1,46 @@
 """Prepare the tests."""
+# pylint: disable=redefined-outer-name
+import copy
+import os
+
+import dask.distributed
 import pytest
 
+from bluepyparallel import init_parallel_factory
+
 
 @pytest.fixture
 def db_url(tmpdir):
     return tmpdir / "db.sql"
+
+
+@pytest.fixture(params=[None, "multiprocessing", "ipyparallel", "dask"])
+def factory_type(request):
+    return request.param
+
+
+@pytest.fixture(scope="session")
+def dask_cluster():
+    cluster = dask.distributed.LocalCluster()
+    yield cluster
+    cluster.close()
+
+
+@pytest.fixture(
+    params=[
+        {},
+        {"chunk_size": 2},
+        {"batch_size": 2},
+        {"chunk_size": 2, "batch_size": 2},
+        {"chunk_size": 999, "batch_size": 999},
+    ]
+)
+def parallel_factory(factory_type, dask_cluster, request):
+    factory_kwargs = copy.deepcopy(request.param)
+    if factory_type == "dask":
+        factory_kwargs["address"] = dask_cluster
+    elif factory_type == "ipyparallel":
+        tox_name = os.environ.get("TOX_ENV_NAME")
+        if tox_name:
+            factory_kwargs["cluster_id"] = f"bluepyparallel_{tox_name}"
+    return init_parallel_factory(factory_type, **factory_kwargs)

tests/conftest.py

@@ -61,11 +61,8 @@ class TestEvaluate:
     """Test the bluepyparallel.evaluator.evaluate function."""
 
     @pytest.mark.parametrize("with_sql", [True, False])
-    @pytest.mark.parametrize("factory_type", [None, "multiprocessing"])
-    def test_evaluate(self, input_df, new_columns, expected_df, db_url, with_sql, factory_type):
+    def test_evaluate(self, input_df, new_columns, expected_df, db_url, with_sql, parallel_factory):
         """Test evaluator on a trivial example."""
-        parallel_factory = init_parallel_factory(factory_type)
-
         result_df = evaluate(
             input_df,
             _evaluation_function,
@@ -88,7 +85,6 @@ def test_evaluate(self, input_df, new_columns, expected_df, db_url, with_sql, fa
         ],
     )
     @pytest.mark.parametrize("with_sql", [True, False])
-    @pytest.mark.parametrize("factory_type", [None, "multiprocessing"])
     def test_evaluate_args_kwargs(
         self,
         input_df,
@@ -97,10 +93,9 @@ def test_evaluate_args_kwargs(
         db_url,
         func_args_kwargs,
         with_sql,
-        factory_type,
+        parallel_factory,
     ):
         """Test evaluator on a trivial example with passing args or kwargs."""
-        parallel_factory = init_parallel_factory(factory_type)
         args, kwargs = deepcopy(func_args_kwargs)
 
         result_df = evaluate(
@@ -124,11 +119,8 @@ def test_evaluate_args_kwargs(
 
         assert_frame_equal(result_df, expected_df, check_like=True)
 
-    @pytest.mark.parametrize("factory_type", [None, "multiprocessing"])
-    def test_evaluate_resume(self, input_df, new_columns, expected_df, db_url, factory_type):
+    def test_evaluate_resume(self, input_df, new_columns, expected_df, db_url, parallel_factory):
         """Test evaluator on a trivial example."""
-        parallel_factory = init_parallel_factory(factory_type)
-
         # Compute some values
         tmp_df = evaluate(
             input_df.loc[[0, 2]],
@@ -193,11 +185,10 @@ def test_evaluate_resume_bad_cols(self, input_df, new_columns, db_url):
                 db_url=db_url,
             )
 
-    @pytest.mark.parametrize("factory_type", [None, "multiprocessing"])
-    def test_evaluate_overwrite_db(self, input_df, new_columns, expected_df, db_url, factory_type):
+    def test_evaluate_overwrite_db(
+        self, input_df, new_columns, expected_df, db_url, parallel_factory
+    ):
         """Test evaluator on a trivial example."""
-        parallel_factory = init_parallel_factory(factory_type)
-
         # Compute once
         previous_df = input_df.copy(deep=True)
         previous_df["name"] += "_previous"
@@ -228,7 +219,6 @@ class TestBenchmark:
         @pytest.mark.parametrize("df_size", ["small", "big"])
         @pytest.mark.parametrize("function_type", ["fast", "slow"])
         @pytest.mark.parametrize("with_sql", [True, False])
-        @pytest.mark.parametrize("factory_type", [None, "multiprocessing"])
         def test_evaluate(
             self,
             input_df,
@@ -238,12 +228,10 @@ def test_evaluate(
             df_size,
             function_type,
             with_sql,
-            factory_type,
+            parallel_factory,
             benchmark,
         ):
             """Test evaluator on a trivial example."""
-            parallel_factory = init_parallel_factory(factory_type, processes=None)
-
             if df_size == "big":
                 input_df = input_df.loc[np.repeat(input_df.index.values, 50)].reset_index(drop=True)
                 expected_df = expected_df.loc[np.repeat(expected_df.index.values, 50)].reset_index(

tests/test_evaluator.py

@@ -6,7 +6,6 @@ testdeps =
     pytest-benchmark
     pytest-cov
     pytest-html
-    pytest-xdist
 
 [tox]
 envlist =
@@ -19,7 +18,21 @@ minversion = 3.1.0
 
 [testenv]
 deps = {[base]testdeps}
-commands = pytest -n 2 --basetemp={envtmpdir} --cov={envsitepackagesdir}/{[base]name} --cov-branch --no-cov-on-fail --html reports/pytest-{envname}.html --self-contained-html --benchmark-skip {posargs}
+commands_pre =
+    - ipcluster stop --cluster-id={[base]name}_{envname} --debug
+    ipcluster start -n 2 --daemonize --log-to-file --cluster-id={[base]name}_{envname} --debug
+commands =
+    pytest \
+    --basetemp={envtmpdir} \
+    --cov={[base]name} \
+    --cov-branch \
+    --no-cov-on-fail \
+    --html reports/pytest-{envname}.html \
+    --self-contained-html \
+    --benchmark-skip \
+    {posargs}
+commands_post =
+    - ipcluster stop --cluster-id={[base]name}_{envname} --debug
 
 [testenv:check-version]
 skip_install = true