Merge pull request #13 from w8jcik/add/dask

This looks great!
Distributed parallelization will enable moldrug to scale more effectively on HPC systems, making it practical to use more computationally demanding fitness functions.

Thanks again, @w8jcik!

Author: ale94mleon

.gitignore

@@ -19,6 +19,7 @@ ideas
 projects
 pkl
 tmp
+tmp*
 tox.ini
 
 build

docs/source/CHANGELOG.md

@@ -9,9 +9,50 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
-- Move from .rst to .md on the documentation.
+- Move from `.rst` to `.md` on the documentation.
 - Update installation instructions.
 
+### Fixed
+
+- The tests no longer rely on `/tmp`; they now create temporary files in the current directory instead. This approach was already used throughout the package, except in the tests. The change makes the behavior consistent and avoids issues that occurred on certain clusters when using `/tmp`.
+- Prevent race condition when creating `error` directory during parallel execution.
+
+### Added
+
+- Cluster configuration can be provided to calculate cost function on multiple nodes of a cluster.  
+  Cluster support requires installation of optional packages related to Dask, a package set called `cluster`.
+
+- ```sh
+  pip install moldrug[cluster]
+  ```
+
+  An example of local parallel execution (same as before, for reference)
+
+  ```python
+  for _ in range(NumbCalls):
+      ga(njobs=njobs)
+  ```
+
+  An example of parallel execution on a SLURM cluster
+
+  ```python
+  cluster = SLURMCluster(
+      queue='short',
+      cores=12,
+      processes=1,
+      memory='8GB',
+      walltime='00:30:00',
+      job_extra_directives=[]
+  )
+  
+  cluster.scale(4)
+
+  runner = Runner(RunnerMode.DASK_JOB_QUEUE_CLUSTER, dask_cluster=cluster)
+
+  for _ in range(NumbCalls):
+      ga(runner=runner)
+  ```
+
 ## [3.7.3] - 2024.07.05
 
 ### Changed

docs/source/installation.md

@@ -10,6 +10,14 @@ Strongly recommended; although this project is still in beta state.
 pip install moldrug
 ```
 
+````{note}
+To enable calculation of a cost function on [a cluster](https://jobqueue.dask.org/en/latest/clusters-api.html) install optional set of dependencies
+
+```sh
+pip install moldrug[cluster]
+```
+````
+
 ````{admonition} Linux 🐧
 :class: Tip
 
@@ -119,3 +127,5 @@ Finally, `pip install moldrug` inside the container.
 pip install -e .[dev]
 pytest tests
 ```
+
+To test execution on a cluster, change the `execution_mode` variable in `test_moldrug.py`. The tests you can run on the cluster afterward are called `test_single_receptor_command_line` and `test_multi_receptor`. Each cluster is different, and the current configuration is tailored to our two internal clusters. To test on your own cluster, you’ll also need to modify the `SLURMCluster` configuration in `test_moldrug.py`. During the execution of each test, multiple entries should appear in the Slurm queue.

docs/source/notebooks/quickstart.ipynb

@@ -1133,6 +1133,34 @@
    "source": [
     "plot_dist(out.SawIndividuals, properties=['qed', 'sa_score', 'vina_score', 'cost'])"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "By the default the calculations are running locally, but it is also possible to evaluate new generations on a cluster.\n",
+    "\n",
+    "First a cluster needs to be configured\n",
+    "\n",
+    "```python\n",
+    "cluster = SLURMCluster(\n",
+    "    queue='short',\n",
+    "    cores=12,\n",
+    "    processes=1,\n",
+    "    memory='8GB',\n",
+    "    walltime='00:30:00',\n",
+    "    job_extra_directives=[]\n",
+    ")\n",
+    "\n",
+    "cluster.scale(4)\n",
+    "\n",
+    "runner = Runner(RunnerMode.DASK_JOB_QUEUE_CLUSTER, dask_cluster=cluster)\n",
+    "```\n",
+    "\n",
+    "and then each iteration takes the `runner` as an argument. That is one needs to call `out(runner=runner)` instead of `out(njobs=njobs)`.\n",
+    "\n",
+    "An example demonstrates use of SLURM, the other supported cluster types can be found here https://jobqueue.dask.org/en/latest/clusters-api.html."
+   ]
   }
  ],
  "metadata": {

pyproject.toml

@@ -63,6 +63,7 @@ file = "LICENSE"
 
 [project.optional-dependencies]
 dev = ["requests", "pytest"]
+cluster = ["dask[distributed]", "dask-jobqueue"]
 
 [tool.versioningit]
 default-version = "1+unknown"

src/moldrug/cli.py

@@ -8,6 +8,7 @@
 import argparse
 import datetime
 import inspect
+import importlib
 import os
 import sys
 from typing import Union
@@ -114,7 +115,8 @@ def _translate_config(self):
         InitArgs = MainConfig.copy()
 
         # Modifying InitArgs
-        _ = [InitArgs.pop(key, None) for key in ['type', 'njobs', 'pick']]
+        for key in ['type', 'njobs', 'cluster', 'pick']:
+            InitArgs.pop(key, None)
         InitArgs['costfunc'] = self.costfunc
 
         # Getting call arguments
@@ -125,6 +127,23 @@ def _translate_config(self):
             except KeyError:
                 pass
 
+        if 'cluster' in MainConfig:
+            if 'type' not in MainConfig['cluster'] or 'kwargs' not in MainConfig['cluster']:
+                raise ValueError("The cluster configuration must contain 'type' and 'kwargs'.")
+
+            from moldrug.runner import Runner, RunnerMode  # terminates with a message if dask is not installed
+
+            try:
+                cluster_class = getattr(importlib.import_module("dask_jobqueue"), MainConfig['cluster']['type'])
+            except ImportError:
+                raise ImportError(f"Unable to import {MainConfig['cluster']['type']} from dask_jobqueue module.")
+
+            cluster = cluster_class(**MainConfig['cluster']['kwargs'])
+            cluster.scale(MainConfig.get('njobs', 1))
+
+            CallArgs['runner'] = Runner(RunnerMode.DASK_JOB_QUEUE_CLUSTER, dask_cluster=cluster)
+            del CallArgs['njobs']
+
         # Checking for follow jobs and sanity check on the arguments
         if FollowConfig:
             # Defining the possible mutable arguments with its default values depending on the type of run
@@ -145,7 +164,7 @@ def _translate_config(self):
                         InitArgs[param.name] = param.default
 
                 MutableArgs = {
-                    'njobs': CallArgs['njobs'],
+                    'njobs': MainConfig['njobs'],
                     'crem_db_path': InitArgs['crem_db_path'],
                     'maxiter': InitArgs['maxiter'],
                     'popsize': InitArgs['popsize'],

src/moldrug/constraintconf.py

@@ -162,7 +162,7 @@ def generate_conformers(mol: Chem.rdchem.Mol,
     """
     # Creating the error directory if needed
     if not os.path.isdir('error'):
-        os.makedirs('error')
+        os.makedirs('error', exist_ok=True)
     # if SMILES to be fixed are not given, assume to the MCS
     if ref_smi:
         if not Chem.MolFromSmiles(ref_smi):

src/moldrug/fitness.py

@@ -311,7 +311,7 @@ def _vinadock(
 
     # Creating the error directory if needed
     if not os.path.isdir('error'):
-        os.makedirs('error')
+        os.makedirs('error', exist_ok=True)
     # Creating the working directory if needed
     if not os.path.exists(wd):
         os.makedirs(wd)

src/moldrug/runner.py

@@ -0,0 +1,85 @@
+import multiprocessing as mp
+from enum import Enum
+from typing import Callable, List, Optional
+
+import tqdm
+
+try:
+    from dask.distributed import Client
+
+    dask_available = True
+except ImportError:
+    dask_available = False
+
+
+def serial_execution(func: Callable, entries: List):
+    return [func(args) for args in tqdm.tqdm(entries, total=len(entries))]
+
+
+def parallel_execution_multiprocessing(func: Callable, entries: List, process_count: int):
+    pool = mp.Pool(process_count)
+    results = [entry for entry in tqdm.tqdm(pool.imap(func, entries), total=len(entries))]
+    pool.close()
+
+    return results
+
+
+if dask_available:
+    def parallel_execution_dask_local(func: Callable, entries: List, client: Client):
+        return client.gather(client.map(func, entries))
+
+    def parallel_execution_dask_cluster(func: Callable, entries: List, client: Client):
+        return client.gather(client.map(func, entries))
+
+
+class RunnerMode(Enum):
+    SERIAL = "serial"
+    MULTIPROCESSING = "multiprocessing"
+    DASK_LOCAL = "dask_local"
+    DASK_JOB_QUEUE_CLUSTER = "dask_job_queue_cluster"
+
+
+class Runner:
+    mode: RunnerMode
+    thread_count: Optional[int]
+    process_count: Optional[int]
+    dask_cluster = None
+    client = None
+
+    def __init__(self, mode: RunnerMode, thread_count: Optional[int] = None, process_count: Optional[int] = None, dask_cluster = None):
+        if mode in [RunnerMode.DASK_LOCAL, RunnerMode.DASK_JOB_QUEUE_CLUSTER]:
+            assert dask_available is not None, "Execution using Dask requires installation of optional dependencies. The optional pip package group is called 'cluster'"
+
+        if mode is RunnerMode.SERIAL:
+            assert thread_count is None and process_count is None and dask_cluster is None, "Serial execution doesn't take any parameters."
+        elif mode in [RunnerMode.MULTIPROCESSING]:
+            assert thread_count is None and process_count is not None and dask_cluster is None, "Only process count is needed."
+        elif mode in [RunnerMode.DASK_LOCAL]:
+            assert (thread_count is not None or process_count is not None) and dask_cluster is None, "Only process count and/or thread count are needed."
+        elif mode in [RunnerMode.DASK_JOB_QUEUE_CLUSTER]:
+            assert thread_count is None and process_count is None and dask_cluster is not None, "Dask execution takes only a Dask cluster object."
+        else:
+            assert False
+
+        self.mode = mode
+        self.thread_count = thread_count
+        self.process_count = process_count
+        self.cluster = dask_cluster
+
+        if mode in [RunnerMode.DASK_LOCAL]:
+            self.client = Client(n_workers=process_count, threads_per_worker=thread_count)
+
+        if mode in [RunnerMode.DASK_JOB_QUEUE_CLUSTER]:
+            self.client = Client(dask_cluster)
+
+    def run(self, func: Callable, entries: List):
+        if self.mode == RunnerMode.SERIAL:
+            return serial_execution(func, entries)
+        elif self.mode == RunnerMode.MULTIPROCESSING:
+            return parallel_execution_multiprocessing(func, entries, process_count=self.process_count)
+        elif self.mode == RunnerMode.DASK_LOCAL:
+            return parallel_execution_dask_local(func, entries, self.client)
+        if self.mode == RunnerMode.DASK_JOB_QUEUE_CLUSTER:
+            return parallel_execution_dask_cluster(func, entries, self.client)
+        else:
+            assert False