Feat: Improve how dask can be configured

Author: adrien-berchet

bluepyparallel/parallel.py

@@ -1,9 +1,11 @@
 """Parallel helper."""
+import json
 import logging
 import multiprocessing
 import os
 from abc import abstractmethod
 from collections.abc import Iterator
+from copy import deepcopy
 from functools import partial
 from multiprocessing.pool import Pool
 
@@ -33,6 +35,7 @@
 except ImportError:  # pragma: no cover
     ipyparallel_available = False
 
+from bluepyparallel.utils import replace_values_in_docstring
 
 L = logging.getLogger(__name__)
 
@@ -210,15 +213,78 @@ def shutdown(self):
             pass
 
 
+_DEFAULT_DASK_CONFIG = {
+    "temporary-directory": None,
+    "distributed": {
+        "worker": {
+            "use_file_locking": False,
+            "memory": {
+                "target": False,
+                "spill": False,
+                "pause": 0.8,
+                "terminate": 0.95,
+            },
+            "profile": {
+                "enabled": False,
+            },
+        },
+        "admin": {
+            "tick": {
+                "limit": "1h",
+            },
+        },
+    },
+}
+
+_DASK_CONFIG_DOCSTRING = """
+It is possible to pass a custom dask configuration in several ways.
+The simplest way is to pass a dictionary to the `dask_config` argument.
+Another way is to create a YAML file containing the configuration and then set the `DASK_CONFIG`
+environment variable to its path. Note that this environment variable must be set before `dask`
+is imported and can not be updated afterwards.
+Also, it is possible to use the `TMPDIR` environment variable to specify the directory in which
+the dask internals will be created. Note that this value will be overridden if a dask configuration
+is given.
+If no config is provided, the following is used:
+
+.. code-block:: JSON
+
+    <>
+""".replace(
+    "<>",
+    json.dumps(
+        _DEFAULT_DASK_CONFIG,
+        sort_keys=True,
+        indent=4,
+    ).replace("\n", "\n" + " " * 4),
+)
+
+
+@replace_values_in_docstring(external_config_block=_DASK_CONFIG_DOCSTRING)
 class DaskFactory(ParallelFactory):
-    """Parallel helper class using dask."""
+    """Parallel helper class using dask.
+
+    <external_config_block>
+    """
 
     _SCHEDULER_PATH = "PARALLEL_DASK_SCHEDULER_PATH"
 
     def __init__(
-        self, batch_size=None, chunk_size=None, scheduler_file=None, address=None, **kwargs
+        self,
+        batch_size=None,
+        chunk_size=None,
+        scheduler_file=None,
+        address=None,
+        dask_config=None,
+        **kwargs,
     ):
         """Initialize the dask factory."""
+        _default_config = deepcopy(_DEFAULT_DASK_CONFIG)
+        _default_config["temporary-directory"] = os.environ.get("TMPDIR", None)
+        dask.config.update_defaults(_default_config)
+        if dask_config is not None:  # pragma: no cover
+            dask.config.set(dask_config)
+
         dask_scheduler_path = scheduler_file or os.getenv(self._SCHEDULER_PATH)
         self.interactive = True
         if dask_scheduler_path:  # pragma: no cover
@@ -269,36 +335,14 @@ def _dask_mapper(in_dask_func, iterable):
         return _mapper
 
 
+@replace_values_in_docstring(external_config_block=_DASK_CONFIG_DOCSTRING)
 class DaskDataFrameFactory(DaskFactory):
-    """Parallel helper class using dask.dataframe."""
+    """Parallel helper class using `dask.dataframe`.
 
-    _SCHEDULER_PATH = "PARALLEL_DASK_SCHEDULER_PATH"
+    <external_config_block>
+    """
 
-    def __init__(
-        self,
-        batch_size=None,
-        chunk_size=None,
-        scheduler_file=None,
-        address=None,
-        dask_config=None,
-        **kwargs,
-    ):
-        super().__init__(
-            batch_size, chunk_size, scheduler_file=scheduler_file, address=address, **kwargs
-        )
-        if dask_config is None:  # pragma: no cover
-            dask_config = {
-                "distributed.worker.use_file_locking": False,
-                "distributed.worker.memory.target": False,
-                "distributed.worker.memory.spill": False,
-                "distributed.worker.memory.pause": 0.8,
-                "distributed.worker.memory.terminate": 95,
-                "distributed.worker.profile.interval": "10000ms",
-                "distributed.worker.profile.cycle": "1000000ms",
-                "distributed.admin.tick.limit": "1h",
-            }
-
-        dask.config.set(dask_config)
+    _SCHEDULER_PATH = "PARALLEL_DASK_SCHEDULER_PATH"
 
     def _with_batches(self, *args, **kwargs):
         """Specific process for batches."""

bluepyparallel/utils.py

@@ -0,0 +1,16 @@
+"""Some utils for the BluePyParallel package."""
+
+
+def replace_values_in_docstring(**kwargs):
+    """Decorator to replace keywords in docstrings by the actual value of a variable.
+
+    .. Note::
+        The keyword must be enclose by <> in the docstring, like <MyKeyword>.
+    """
+
+    def inner(func):
+        for k, v in kwargs.items():
+            func.__doc__ = func.__doc__.replace(f"<{k}>", str(v))
+        return func
+
+    return inner

setup.py

@@ -33,6 +33,7 @@
 
 test_reqs = [
     "mpi4py>=3.0.1",
+    "packaging>=20",
     "pytest>=6.1",
     "pytest-benchmark>=3.4",
     "pytest-cov>=3",

tests/test_parallel.py

@@ -1,15 +1,23 @@
 """Test the ``bluepyparallel.parallel`` module."""
 # pylint: disable=missing-function-docstring
 # pylint: disable=redefined-outer-name
+import importlib.metadata
+import json
+import subprocess
+import sys
 from collections.abc import Iterator
 from copy import deepcopy
 
 import pandas as pd
 import pytest
+import yaml
+from packaging.version import Version
 
 from bluepyparallel import init_parallel_factory
 from bluepyparallel.parallel import DaskDataFrameFactory
 
+dask_version = Version(importlib.metadata.version("dask"))
+
 
 def _evaluation_function_range(element, coeff_a=1.0, coeff_b=1.0):
     """Mock evaluation function."""
@@ -98,3 +106,123 @@ def test_bad_factory_name(self):
         """Test a factory with a wrong name."""
         with pytest.raises(KeyError):
             init_parallel_factory("UNKNOWN FACTORY")
+
+
+@pytest.fixture(params=[True, False])
+def env_tmpdir(tmpdir, request):
+    if request.param:
+        dask_tmpdir = str(tmpdir / "tmpdir_from_env")
+        yield dask_tmpdir
+    else:
+        yield None
+
+
+@pytest.fixture(params=[True, False])
+def env_daskconfig(tmpdir, request):
+    if request.param:
+        # Create dask config file
+        dask_config = {"distributed": {"worker": {"memory": {"pause": 0.123456}}}}
+        filepath = str(tmpdir / "dask_config.yml")
+        with open(filepath, "w", encoding="utf-8") as file:
+            yaml.dump(dask_config, file)
+
+        yield filepath
+    else:
+        yield None
+
+
+@pytest.mark.parametrize(
+    "dask_config",
+    [
+        pytest.param(None, id="No Dask config"),
+        pytest.param({"temporary-directory": "tmpdir"}, id="Dask config with tmp dir"),
+    ],
+)
+@pytest.mark.parametrize(
+    "factory_type",
+    [
+        pytest.param("dask", id="Dask"),
+        pytest.param("dask_dataframe", id="Dask-dataframe"),
+    ],
+)
+@pytest.mark.skipif(dask_version < Version("2023.4"), reason="Requires dask >= 2023.4")
+def test_dask_config(tmpdir, dask_config, factory_type, env_tmpdir, env_daskconfig):
+    """Test the methods to update dask configuration."""
+    dask_config_tmpdir = str(tmpdir / "tmpdir_from_config")
+    has_tmpdir = (
+        dask_config is not None and dask_config.get("temporary-directory", None) is not None
+    )
+    if has_tmpdir:
+        dask_config["temporary-directory"] = dask_config_tmpdir
+    dask_config_str = "None" if dask_config is None else json.dumps(dask_config)
+
+    # Must test using a subprocess because the DASK_CONFIG environment variable is only considered
+    # when dask is imported
+    code = """if True:  # This is just to avoid indentation issue
+    import os
+
+    import dask.config
+    import dask.distributed
+
+    from bluepyparallel import init_parallel_factory
+    from bluepyparallel.parallel import DaskDataFrameFactory
+    from bluepyparallel.parallel import DaskFactory
+
+
+    dask_cluster = dask.distributed.LocalCluster(dashboard_address=None)
+
+    has_tmpdir = {has_tmpdir}
+    dask_config_tmpdir = "{tmpdir}"
+    env_tmpdir = os.getenv("TMPDIR", None)
+    dask_config = {dask_config}
+    env_daskconfig = os.getenv("DASK_CONFIG", None)
+
+    print("Values in subprocess:")
+    print("tmpdir:", dask_config_tmpdir)
+    print("has_tmpdir:", has_tmpdir)
+    print("env_tmpdir:", env_tmpdir)
+    print("env_daskconfig:", env_daskconfig)
+    print("dask_config:", dask_config)
+
+    factory_kwargs = {{
+        "address": dask_cluster,
+        "dask_config": dask_config,
+    }}
+
+    factory = init_parallel_factory("{factory_type}", **factory_kwargs)
+
+    print("tmpdir in dask.config:", dask.config.get("temporary-directory"))
+    print(
+        "distributed.worker.memory.pause in dask.config:",
+        dask.config.get("distributed.worker.memory.pause"),
+    )
+
+    if "{factory_type}" == "dask":
+        assert isinstance(factory, DaskFactory)
+    else:
+        assert isinstance(factory, DaskDataFrameFactory)
+
+    if env_daskconfig is not None:
+        assert dask.config.get("distributed.worker.memory.pause") == 0.123456
+    else:
+        assert dask.config.get("distributed.worker.memory.pause") == 0.8
+
+    if has_tmpdir:
+        assert dask.config.get("temporary-directory") == dask_config_tmpdir
+    else:
+        if env_tmpdir is not None:
+            assert dask.config.get("temporary-directory") == env_tmpdir
+        else:
+            assert dask.config.get("temporary-directory") is None
+    """.format(
+        dask_config=dask_config_str,
+        factory_type=factory_type,
+        tmpdir=dask_config_tmpdir,
+        has_tmpdir=has_tmpdir,
+    )
+    envs = {}
+    if env_daskconfig is not None:
+        envs["DASK_CONFIG"] = env_daskconfig
+    if env_tmpdir is not None:
+        envs["TMPDIR"] = env_tmpdir
+    subprocess.check_call([sys.executable, "-c", code], env=envs)