Execution on a cluster with Dask

Author: w8jcik

.gitignore

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

docs/source/CHANGELOG.md

@@ -9,9 +9,49 @@ 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.
+
+### 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`.
+  ```
+  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/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,84 @@
+import multiprocessing as mp
+from enum import Enum
+from typing import Callable, List, Optional
+
+import tqdm
+
+try:
+    from dask.distributed import Client
+    from dask_jobqueue import JobQueueCluster
+
+    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, process_count: int, process_threads_count: int):
+        client = Client(n_workers=process_count, threads_per_worker=process_threads_count)
+
+        return client.gather(client.map(func, entries))
+
+
+    def parallel_execution_dask_cluster(func: Callable, entries: List, cluster: JobQueueCluster):
+        client = Client(cluster)
+
+        return client.gather(client.map(func, entries))
+
+
+class RunnerMode(Enum):
+    SERIAL = "serial"
+    MULTIPROCESSING = "multiprocessing"
+    DASK_LOCAL = "dask_local"
+    DASK_JOB_QUEUE_CLUSTER = "dask_slurm"
+
+
+class Runner:
+    mode: RunnerMode
+    thread_count: Optional[int]
+    process_count: Optional[int]
+    dask_cluster = 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_cluster is not None, "Execution using Dask requires installation of optional dependencies. The optional pip package group is called 'dask'"
+
+        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 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
+
+    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, process_count=self.process_count, process_threads_count=self.thread_count)
+        if self.mode == RunnerMode.DASK_JOB_QUEUE_CLUSTER:
+            return parallel_execution_dask_cluster(func, entries, cluster=self.cluster)
+        else:
+            assert False

src/moldrug/utils.py

@@ -3,7 +3,6 @@
 import bz2
 import collections.abc
 import datetime
-import multiprocessing as mp
 import os
 import random
 import shutil
@@ -12,20 +11,20 @@
 import time
 from copy import deepcopy
 from inspect import signature
-from typing import Callable, Dict, Iterable, List, Union
+from typing import Callable, Dict, Iterable, List, Optional, Union
 from warnings import warn
 
 import dill as pickle
 import numpy as np
 import pandas as pd
-import tqdm
 from crem.crem import grow_mol, mutate_mol
 from meeko import (MoleculePreparation, PDBQTMolecule, PDBQTWriterLegacy,
                    RDKitMolCreate)
 from rdkit import Chem, RDLogger
 from rdkit.Chem import AllChem, DataStructs, Descriptors, Lipinski, rdFMCS
 
 from moldrug import __version__
+from moldrug.runner import Runner, RunnerMode, parallel_execution_multiprocessing
 
 RDLogger.DisableLog('rdApp.*')
 # # in order to pickle the isotope properties of the molecule
@@ -1164,9 +1163,7 @@ def __call__(self, njobs: int = 1, pick: int = None):
             args_list.append((individual, kwargs_copy))
 
         print('Calculating cost function...')
-        pool = mp.Pool(njobs)
-        self.pop = [individual for individual in tqdm.tqdm(pool.imap(self.__costfunc__, args_list), total=len(args_list))]
-        pool.close()
+        self.pop = parallel_execution_multiprocessing(self.__costfunc__, args_list, njobs)
 
         # Clean directory
         costfunc_jobs_tmp_dir.cleanup()
@@ -1408,7 +1405,7 @@ def __init__(self, seed_mol: Union[Chem.rdchem.Mol, Iterable[Chem.rdchem.Mol]],
         self.InitIndividual = Individual(self._seed_mol[0], idx=0, randomseed=self.randomseed)
         self.pop = []
 
-    def __call__(self, njobs: int = 1):
+    def __call__(self, njobs: int = 1, runner: Optional[Runner] = None):
         """Call definition
 
         Parameters
@@ -1421,6 +1418,12 @@ def __call__(self, njobs: int = 1):
         RuntimeError
             Error during the initialization of the population.
         """
+        if njobs > 1:
+            assert runner is None, "Both njobs > 1 and runner have been specified. Please use only one of the parameters."
+
+        if runner is None:
+            runner = Runner(RunnerMode.MULTIPROCESSING, process_count=njobs)
+
         ts = time.time()
         # Counting the calls
         self.NumCalls += 1
@@ -1492,18 +1495,8 @@ def __call__(self, njobs: int = 1):
                 args_list.append((individual, kwargs_copy))
 
             print(f'\n\nCreating the first population with {len(self.pop)} members:')
-            try:
-                pool = mp.Pool(njobs)
-                self.pop = [individual for individual in tqdm.tqdm(pool.imap(self.__costfunc__, args_list), total=len(args_list))]
-                pool.close()
-            except Exception as e1:
-                warn("Parallelization did not work. Trying with serial...")
-                try:
-                    self.pop = [self.__costfunc__(args) for args in tqdm.tqdm(args_list, total=len(args_list))]
-                except Exception as e2:
-                    raise RuntimeError("Serial did not work either. Here are the ucurred exceptions:\n"
-                                       f"=========Parellel=========:\n {e1}\n"
-                                       f"==========Serial==========:\n {e2}")
+            self.pop = runner.run(self.__costfunc__, entries=args_list)
+
             # Clean directory
             costfunc_jobs_tmp_dir.cleanup()
 
@@ -1603,19 +1596,8 @@ def __call__(self, njobs: int = 1):
                     args_list.append((individual, kwargs_copy))
                 print(f'Evaluating generation {self.NumGens} / {self.maxiter + number_of_previous_generations}:')
 
-                # Calculating cost fucntion in parallel
-                try:
-                    pool = mp.Pool(njobs)
-                    popc = [individual for individual in tqdm.tqdm(pool.imap(self.__costfunc__, args_list), total=len(args_list))]
-                    pool.close()
-                except Exception as e1:
-                    warn("Parallelization did not work. Trying with serial...")
-                    try:
-                        popc = [self.__costfunc__(args) for args in tqdm.tqdm(args_list, total=len(args_list))]
-                    except Exception as e2:
-                        raise RuntimeError("Serial did not work either. Here are the ucurred exceptions:\n"
-                                           f"=========Parellel=========:\n {e1}\n"
-                                           f"==========Serial==========:\n {e2}")
+                # Calculating cost function in parallel
+                popc = runner.run(self.__costfunc__, args_list)
 
                 # Clean directory
                 costfunc_jobs_tmp_dir.cleanup()