[GSoC] Parallelisation of AnalysisBase with multiprocessing and dask

.github/actions/setup-deps/action.yaml

@@ -62,6 +62,8 @@ inputs:
     default: 'chemfiles-python>=0.9'
   clustalw:
     default: 'clustalw=2.1'
+  dask:
+    default: 'dask'
   distopia:
     default: 'distopia>=0.2.0'
   h5py:
@@ -134,6 +136,7 @@ runs:
           ${{ inputs.biopython }}
           ${{ inputs.chemfiles-python }}
           ${{ inputs.clustalw }}
+          ${{ inputs.dask }}
           ${{ inputs.distopia }}
           ${{ inputs.gsd }}
           ${{ inputs.h5py }}

package/CHANGELOG

@@ -52,6 +52,8 @@ Fixes
  * Fix groups.py doctests using sphinx directives (Issue #3925, PR #4374)
 
 Enhancements
+ * Introduce parallelization API to `AnalysisBase` and to `analysis.rms.RMSD` class
+   (Issue #4158, PR #4304)
  * Improve error message for `AtomGroup.unwrap()` when bonds are not present.(Issue #4436, PR #4642)
  * Add `analysis.DSSP` module for protein secondary structure assignment, based on [pydssp](https://github.com/ShintaroMinami/PyDSSP)
  * Added a tqdm progress bar for `MDAnalysis.analysis.pca.PCA.transform()`

package/MDAnalysis/analysis/__init__.py

@@ -23,6 +23,7 @@
 
 __all__ = [
     'align',
+    'backends',
     'base',
     'contacts',
     'density',
@@ -45,6 +46,7 @@
     'pca',
     'psa',
     'rdf',
+    'results',
     'rms',
     'waterdynamics',
 ]

package/MDAnalysis/analysis/backends.py

@@ -0,0 +1,333 @@
+"""Analysis backends --- :mod:`MDAnalysis.analysis.backends`
+============================================================
+
+.. versionadded:: 2.8.0
+
+
+The :mod:`backends` module provides :class:`BackendBase` base class to
+implement custom execution backends for
+:meth:`MDAnalysis.analysis.base.AnalysisBase.run` and its
+subclasses.
+
+.. SeeAlso:: :ref:`parallel-analysis`
+
+.. _backends:
+
+Backends
+--------
+
+Three built-in backend classes are provided:
+
+* *serial*: :class:`BackendSerial`, that is equivalent to using no
+  parallelization and is the default
+
+* *multiprocessing*: :class:`BackendMultiprocessing` that supports
+  parallelization via standard Python :mod:`multiprocessing` module
+  and uses default :mod:`pickle` serialization
+
+* *dask*: :class:`BackendDask`, that uses the same process-based
+  parallelization as :class:`BackendMultiprocessing`, but different
+  serialization algorithm via `dask <https://dask.org/>`_ (see `dask
+  serialization algorithms
+  <https://distributed.dask.org/en/latest/serialization.html>`_ for details)
+
+Classes
+-------
+
+"""
+import warnings
+from typing import Callable
+from MDAnalysis.lib.util import is_installed
+
+
+class BackendBase:
+    """Base class for backend implementation.
+
+    Initializes an instance and performs checks for its validity, such as
+    ``n_workers`` and possibly other ones.
+
+    Parameters
+    ----------
+    n_workers : int
+        number of workers (usually, processes) over which the work is split
+
+    Examples
+    --------
+    .. code-block:: python
+
+        from MDAnalysis.analysis.backends import BackendBase
+
+        class ThreadsBackend(BackendBase):
+            def apply(self, func, computations):
+                from multiprocessing.dummy import Pool
+
+                with Pool(processes=self.n_workers) as pool:
+                    results = pool.map(func, computations)
+                return results
+
+        import MDAnalysis as mda
+        from MDAnalysis.tests.datafiles import PSF, DCD
+        from MDAnalysis.analysis.rms import RMSD
+
+        u = mda.Universe(PSF, DCD)
+        ref = mda.Universe(PSF, DCD)
+
+        R = RMSD(u, ref)
+
+        n_workers = 2
+        backend = ThreadsBackend(n_workers=n_workers)
+        R.run(backend=backend, unsupported_backend=True)
+
+    .. warning::
+        Using `ThreadsBackend` above will lead to erroneous results, since it
+        is an educational example. Do not use it for real analysis.
+
+
+    .. versionadded:: 2.8.0
+
+    """
+
+    def __init__(self, n_workers: int):
+        self.n_workers = n_workers
+        self._validate()
+
+    def _get_checks(self):
+        """Get dictionary with ``condition: error_message`` pairs that ensure the
+        validity of the backend instance
+
+        Returns
+        -------
+        dict
+            dictionary with ``condition: error_message`` pairs that will get
+            checked during ``_validate()`` run
+        """
+        return {
+            isinstance(self.n_workers, int) and self.n_workers > 0:
+            f"n_workers should be positive integer, got {self.n_workers=}",
+        }
+
+    def _get_warnings(self):
+        """Get dictionary with ``condition: warning_message`` pairs that ensure
+        the good usage of the backend instance
+
+        Returns
+        -------
+        dict
+            dictionary with ``condition: warning_message`` pairs that will get
+            checked during ``_validate()`` run
+        """
+        return dict()
+
+    def _validate(self):
+        """Check correctness (e.g. ``dask`` is installed if using ``backend='dask'``)
+        and good usage (e.g. ``n_workers=1`` if backend is serial) of the backend
+
+        Raises
+        ------
+        ValueError
+            if one of the conditions in :meth:`_get_checks` is ``True``
+        """
+        for check, msg in self._get_checks().items():
+            if not check:
+                raise ValueError(msg)
+        for check, msg in self._get_warnings().items():
+            if not check:
+                warnings.warn(msg)
+
+    def apply(self, func: Callable, computations: list) -> list:
+        """map function `func` to all tasks in the `computations` list
+
+        Main method that will get called when using an instance of
+        ``BackendBase``.  It is equivalent to running ``[func(item) for item in
+        computations]`` while using the parallel backend capabilities.
+
+        Parameters
+        ----------
+        func : Callable
+            function to be called on each of the tasks in computations list
+        computations : list
+            computation tasks to apply function to
+
+        Returns
+        -------
+        list
+            list of results of the function
+
+        """
+        raise NotImplementedError
+
+
+class BackendSerial(BackendBase):
+    """A built-in backend that does serial execution of the function, without any
+    parallelization.
+
+    Parameters
+    ----------
+    n_workers : int
+        Is ignored in this class, and if ``n_workers`` > 1, a warning will be
+        given.
+
+
+    .. versionadded:: 2.8.0
+    """
+
+    def _get_warnings(self):
+        """Get dictionary with ``condition: warning_message`` pairs that ensure
+        the good usage of the backend instance. Here, it checks if the number
+        of workers is not 1, otherwise gives warning.
+
+        Returns
+        -------
+        dict
+            dictionary with ``condition: warning_message`` pairs that will get
+            checked during ``_validate()`` run
+        """
+        return {
+            self.n_workers == 1:
+            "n_workers is ignored when executing with backend='serial'"
+        }
+
+    def apply(self, func: Callable, computations: list) -> list:
+        """
+        Serially applies `func` to each task object in ``computations``.
+
+        Parameters
+        ----------
+        func : Callable
+            function to be called on each of the tasks in computations list
+        computations : list
+            computation tasks to apply function to
+
+        Returns
+        -------
+        list
+            list of results of the function
+        """
+        return [func(task) for task in computations]
+
+
+class BackendMultiprocessing(BackendBase):
+    """A built-in backend that executes a given function using the
+    :meth:`multiprocessing.Pool.map <multiprocessing.pool.Pool.map>` method.
+
+    Parameters
+    ----------
+    n_workers : int
+        number of processes in :class:`multiprocessing.Pool
+        <multiprocessing.pool.Pool>` to distribute the workload
+        between. Must be a positive integer.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        from MDAnalysis.analysis.backends import BackendMultiprocessing
+        import multiprocessing as mp
+
+        backend_obj = BackendMultiprocessing(n_workers=mp.cpu_count())
+
+
+    .. versionadded:: 2.8.0
+
+    """
+
+    def apply(self, func: Callable, computations: list) -> list:
+        """Applies `func` to each object in ``computations`` using `multiprocessing`'s `Pool.map`.
+
+        Parameters
+        ----------
+        func : Callable
+            function to be called on each of the tasks in computations list
+        computations : list
+            computation tasks to apply function to
+
+        Returns
+        -------
+        list
+            list of results of the function
+        """
+        from multiprocessing import Pool
+
+        with Pool(processes=self.n_workers) as pool:
+            results = pool.map(func, computations)
+        return results
+
+
+class BackendDask(BackendBase):
+    """A built-in backend that executes a given function with *dask*.
+
+    Execution is performed with the :func:`dask.compute` function of
+    :class:`dask.delayed.Delayed` object (created with
+    :func:`dask.delayed.delayed`) with ``scheduler='processes'`` and
+    ``chunksize=1`` (this ensures uniform distribution of tasks among
+    processes). Requires the `dask package <https://docs.dask.org/en/stable/>`_
+    to be `installed <https://docs.dask.org/en/stable/install.html>`_.
+
+    Parameters
+    ----------
+    n_workers : int
+        number of processes in to distribute the workload
+        between. Must be a positive integer. Workers are actually
+        :class:`multiprocessing.pool.Pool` processes, but they use a different and
+        more flexible `serialization protocol
+        <https://docs.dask.org/en/stable/phases-of-computation.html#graph-serialization>`_.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        from MDAnalysis.analysis.backends import BackendDask
+        import multiprocessing as mp
+
+        backend_obj = BackendDask(n_workers=mp.cpu_count())
+
+
+    .. versionadded:: 2.8.0
+
+    """
+
+    def apply(self, func: Callable, computations: list) -> list:
+        """Applies `func` to each object in ``computations``.
+
+        Parameters
+        ----------
+        func : Callable
+            function to be called on each of the tasks in computations list
+        computations : list
+            computation tasks to apply function to
+
+        Returns
+        -------
+        list
+            list of results of the function
+        """
+        from dask.delayed import delayed
+        import dask
+
+        computations = [delayed(func)(task) for task in computations]
+        results = dask.compute(computations,
+                               scheduler="processes",
+                               chunksize=1,
+                               num_workers=self.n_workers)[0]
+        return results
+
+    def _get_checks(self):
+        """Get dictionary with ``condition: error_message`` pairs that ensure the
+        validity of the backend instance. Here checks if ``dask`` module is
+        installed in the environment.
+
+        Returns
+        -------
+        dict
+            dictionary with ``condition: error_message`` pairs that will get
+            checked during ``_validate()`` run
+        """
+        base_checks = super()._get_checks()
+        checks = {
+            is_installed("dask"):
+            ("module 'dask' is missing. Please install 'dask': "
+             "https://docs.dask.org/en/stable/install.html")
+        }
+        return base_checks | checks