[GSoC] Parallelisation of AnalysisBase with multiprocessing and dask

package/MDAnalysis/analysis/base.py

@@ -131,8 +131,49 @@
 from MDAnalysis.core.groups import AtomGroup
 from MDAnalysis.lib.log import ProgressBar
 
+from functools import partial
+
 logger = logging.getLogger(__name__)
 
+def localdelayed(obj):
+    """
+    Mock implementation of `dask.delayed.delayed` function
+    with the same semantics
+    """
+    if isinstance(obj, Iterable):
+
+        class inner:
+            def __init__(self, iterable):
+                self._computations = iterable
+
+            def compute(self):
+                return [f.compute() for f in self._computations]
+
+        return inner(obj)
+
+    elif isinstance(obj, Callable):
+
+        class inner:
+            def __init__(self, *a, **kwa):
+                self._a = a
+                self._kwa = kwa
+                self._func = obj
+
+            def compute(self):
+                return self._func(*self._a, **self._kwa)
+
+        return inner
+    else:
+        raise ValueError(f"Argument should be Iterable or Callable, got {type(obj)}")
+
+from itertools import islice
+
+def split_every(n, iterable):
+    i = iter(iterable)
+    piece = list(islice(i, n))
+    while piece:
+        yield piece
+        piece = list(islice(i, n))
 
 class Results(UserDict):
     r"""Container object for storing results.
@@ -219,6 +260,11 @@ def __getstate__(self):
     def __setstate__(self, state):
         self.data = state
 
+from typing import Callable, Iterable
+
+
+# from dask import delayed
+
 
 class AnalysisBase(object):
     r"""Base class for defining multi-frame analysis
@@ -389,9 +435,10 @@ def _conclude(self):
         """
         pass  # pylint: disable=unnecessary-pass
 
-    def run(self, start=None, stop=None, step=None, frames=None,
-            verbose=None, *, progressbar_kwargs={}):
-        """Perform the calculation
+    def _compute(self, start=None, stop=None, step=None, 
+        frames=None, verbose=None, *, progressbar_kwargs={}):
+        """Perform the calculation on frames that have been setup prior to that
+        using _setup_frames()
 
         Parameters
         ----------
@@ -447,8 +494,177 @@ def run(self, start=None, stop=None, step=None, frames=None,
             self.times[i] = ts.time
             self._single_frame()
         logger.info("Finishing up")
+        return self
+
+    def _setup_bslices(self, start=None, stop=None, step=None, frames=None, 
+                       n_bslices=None): 
+        """
+        Set the self._bslices (the workload distribution scheme) for the future delayed
+        computations.
+
+        Parameters
+        ----------
+        start : int, optional
+            start frame of analysis
+        stop : int, optional
+            stop frame of analysis
+        step : int, optional
+            number of frames to skip between each analysed frame
+        frames : array_like, optional
+            array of integers or booleans to slice trajectory; `frames` can
+            only be used *instead* of `start`, `stop`, and `step`. Setting
+            *both* `frames` and at least one of `start`, `stop`, `step` to a
+            non-default value will raise a :exc:`ValueError`.
+
+        Raises
+        ------
+        ValueError
+            if *both* `frames` and at least one of `start`, `stop`, or `frames`
+            is provided (i.e., set to another value than ``None``)
+
+        Returns
+        -------
+        bslices : list of (bstart, bstop, bstep, bframes) tuples.
+            Iterator will have size of self._n_bslices.
+        """
+        if frames is not None:
+            if not all(opt is None for opt in [start, stop, step]):
+                raise ValueError("start/stop/step cannot be combined with "
+                                 "frames")
+            self.n_frames_total = len(frames)
+            n_bslices = self._n_bslices
+            slices = [(None, None, None, frames[i::n_bslices]) for i in range(n_bslices)]
+
+        else:  # frames is None
+            start, stop, step = self._trajectory.check_slice_indices(start, stop, step)
+            n_frames = len(range(start, stop, step))
+            self.start, self.stop, self.step = start, stop, step
+            self.n_frames_total = n_frames
+
+            ## this part was taken from pmda with slight modifications
+            n_bslices = self._n_bslices
+            bsizes = np.ones(n_bslices, dtype=np.int64) * n_frames // n_bslices
+            bsizes += np.arange(n_bslices, dtype=np.int64) < n_frames % n_bslices
+            # This can give a last index that is larger than the real last index;
+            # this is not a problem for slicing but it's not pretty.
+            # Example: original [0:20:3] -> n_frames=7, start=0, step=3:
+            #          last frame 21 instead of 20
+            bsizes *= step
+            idx = np.cumsum(np.concatenate(([start], bsizes)))
+            slices = [(bstart, bstop, step, None) 
+                       for bstart, bstop in zip(idx[:-1], idx[1:])]
+
+            # fix very last stop index: make sure it's within trajectory range or None
+            # (no really critical because the slices will work regardless, but neater)
+            last = slices[-1]
+            last_stop = min(last[1], stop) if stop is not None else stop
+            slices[-1] = (last[0], last_stop, last[2], None)
+
+        self._bslices = slices
+        return self._bslices
+
+
+    def _setup_scheduler(self, scheduler, n_workers):
+        """
+        Configure parameters necessary for running a distributed workload.
+
+        Parameters
+        ----------
+        scheduler : dask.distributed.Client
+            dask scheduler object
+        n_workers : int, optional
+            number of workers (local or remote processes).
+        """
+        if scheduler == 'localdask':
+            n_workers = 1
+        else:
+            from dask.distributed import Client
+            if isinstance(scheduler, Client):
+                # TODO: add assertions that make sure that 
+                # you don't set 
+                # both `scheduler` and `n_workers/threads_per_worker`
+                n_workers = len(scheduler.ncores())
+            else:
+                kwargs = {'n_workers':n_workers}
+            self._scheduler_kwargs = kwargs
+        self._n_bslices = self._n_workers = n_workers
+
+    def run(self, start=None, stop=None, step=None, frames=None,
+            verbose=None, n_workers=1, 
+            *, scheduler=None,
+            progressbar_kwargs={},
+            ):
+        """Perform the calculation
+
+        Parameters
+        ----------
+        start : int, optional
+            start frame of analysis
+        stop : int, optional
+            stop frame of analysis
+        step : int, optional
+            number of frames to skip between each analysed frame
+        frames : array_like, optional
+            array of integers or booleans to slice trajectory; `frames` can
+            only be used *instead* of `start`, `stop`, and `step`. Setting
+            *both* `frames` and at least one of `start`, `stop`, `step` to a
+            non-default value will raise a :exc:`ValueError`.
+
+            .. versionadded:: 2.2.0
+
+        verbose : bool, optional
+            Turn on verbosity
+
+        scheduler : str, optional
+            Enables running with different schedulers.
+
+        progressbar_kwargs : dict, optional
+            ProgressBar keywords with custom parameters regarding progress bar position, etc; 
+            see :class:`MDAnalysis.lib.log.ProgressBar` for full list.
+
+
+        .. versionchanged:: 2.2.0
+            Added ability to analyze arbitrary frames by passing a list of
+            frame indices in the `frames` keyword argument.
+
+        .. versionchanged:: 2.5.0
+            Add `progressbar_kwargs` parameter, 
+            allowing to modify description, position etc of tqdm progressbars
+        """
+        if scheduler is None: # fallback to the local scheduler
+            self._compute(start=start, stop=stop, step=step, frames=frames, 
+            verbose=verbose, progressbar_kwargs=progressbar_kwargs)
+        else:
+            self._setup_scheduler(scheduler=scheduler, n_workers=n_workers)
+            if scheduler == 'localdask': # imitation of dask mainly for testing purposes
+                delayed = localdelayed
+            else: # elif scheduler is a type of dask scheduler
+                try:
+                    from dask.delayed import delayed
+                except ImportError:
+                    def delayed(*args, **kwargs): # need implementation to avoid syntax warnings below
+                        pass
+                    # raising exception here allows me to skip potentionally time-consuming setup steps
+                    raise ImportError('Please install dask for this functionality')
+            self._setup_bslices(start=start, stop=stop, step=step,
+                frames=frames)
+
+            computations = delayed(
+                [
+                    delayed(self._compute)(start=bstart, stop=bstop, step=bstep, frames=bframes)
+                    for bstart, bstop, bstep, bframes in self._bslices
+                ]
+            )
+            dask_results = computations.compute()
+            self._remote_results = dask_results
+            self._parallel_conclude()
+
         self._conclude()
+
         return self
+
+    def _parallel_conclude(self):
+        self.results = self._remote_results[0].results
 
 
 class AnalysisFromFunction(AnalysisBase):

testsuite/MDAnalysisTests/analysis/conftest.py

@@ -0,0 +1,17 @@
+import pytest
+
+@pytest.fixture(params=[None, 'localdask',], scope='module')
+def scheduler(request):
+    """
+    Fixture for testing all possible schedulers.
+    If scheduler raises an exception, you should explicitly check it
+    using `with pytest.raises(NotImplementedError)`
+    """
+    return request.param
+
+@pytest.fixture(params=[None])
+def localscheduler(request):
+    """
+    Fixture for testing only local, i.e. before-dask, codepath.
+    """
+    return request.param

testsuite/MDAnalysisTests/analysis/test_align.py

@@ -417,9 +417,9 @@ def universe(self):
     def reference(self):
         return mda.Universe(PSF, CRD)
 
-    def test_average_structure_deprecated_attrs(self, universe, reference):
+    def test_average_structure_deprecated_attrs(self, universe, reference, scheduler):
         # Issue #3278 - remove in MDAnalysis 3.0.0
-        avg = align.AverageStructure(universe, reference).run(stop=2)
+        avg = align.AverageStructure(universe, reference).run(stop=2, scheduler=scheduler)
 
         wmsg = "The `universe` attribute was deprecated in MDAnalysis 2.0.0"
         with pytest.warns(DeprecationWarning, match=wmsg):
@@ -434,45 +434,45 @@ def test_average_structure_deprecated_attrs(self, universe, reference):
         with pytest.warns(DeprecationWarning, match=wmsg):
             assert avg.rmsd == avg.results.rmsd
 
-    def test_average_structure(self, universe, reference):
+    def test_average_structure(self, universe, reference, scheduler):
         ref, rmsd = _get_aligned_average_positions(self.ref_files, reference)
-        avg = align.AverageStructure(universe, reference).run()
+        avg = align.AverageStructure(universe, reference).run(scheduler=scheduler)
         assert_almost_equal(avg.results.universe.atoms.positions, ref,
                             decimal=4)
         assert_almost_equal(avg.results.rmsd, rmsd)
 
-    def test_average_structure_mass_weighted(self, universe, reference):
+    def test_average_structure_mass_weighted(self, universe, reference, scheduler):
         ref, rmsd = _get_aligned_average_positions(self.ref_files, reference, weights='mass')
-        avg = align.AverageStructure(universe, reference, weights='mass').run()
+        avg = align.AverageStructure(universe, reference, weights='mass').run(scheduler=scheduler)
         assert_almost_equal(avg.results.universe.atoms.positions, ref,
                             decimal=4)
         assert_almost_equal(avg.results.rmsd, rmsd)
 
-    def test_average_structure_select(self, universe, reference):
+    def test_average_structure_select(self, universe, reference, scheduler):
         select = 'protein and name CA and resid 3-5'
         ref, rmsd = _get_aligned_average_positions(self.ref_files, reference, select=select)
-        avg = align.AverageStructure(universe, reference, select=select).run()
+        avg = align.AverageStructure(universe, reference, select=select).run(scheduler=scheduler)
         assert_almost_equal(avg.results.universe.atoms.positions, ref,
                             decimal=4)
         assert_almost_equal(avg.results.rmsd, rmsd)
 
-    def test_average_structure_no_ref(self, universe):
+    def test_average_structure_no_ref(self, universe, scheduler):
         ref, rmsd = _get_aligned_average_positions(self.ref_files, universe)
-        avg = align.AverageStructure(universe).run()
+        avg = align.AverageStructure(universe).run(scheduler=scheduler)
         assert_almost_equal(avg.results.universe.atoms.positions, ref,
                             decimal=4)
         assert_almost_equal(avg.results.rmsd, rmsd)
 
-    def test_average_structure_no_msf(self, universe):
-        avg = align.AverageStructure(universe).run()
+    def test_average_structure_no_msf(self, universe, scheduler):
+        avg = align.AverageStructure(universe).run(scheduler=scheduler)
         assert not hasattr(avg, 'msf')
 
     def test_mismatch_atoms(self, universe):
         u = mda.Merge(universe.atoms[:10])
         with pytest.raises(SelectionError):
             align.AverageStructure(universe, u)
 
-    def test_average_structure_ref_frame(self, universe):
+    def test_average_structure_ref_frame(self, universe, scheduler):
         ref_frame = 3
         u = mda.Merge(universe.atoms)
 
@@ -483,13 +483,13 @@ def test_average_structure_ref_frame(self, universe):
         # back to start
         universe.trajectory[0]
         ref, rmsd = _get_aligned_average_positions(self.ref_files, u)
-        avg = align.AverageStructure(universe, ref_frame=ref_frame).run()
+        avg = align.AverageStructure(universe, ref_frame=ref_frame).run(scheduler=scheduler)
         assert_almost_equal(avg.results.universe.atoms.positions, ref,
                             decimal=4)
         assert_almost_equal(avg.results.rmsd, rmsd)
 
-    def test_average_structure_in_memory(self, universe):
-        avg = align.AverageStructure(universe, in_memory=True).run()
+    def test_average_structure_in_memory(self, universe, scheduler):
+        avg = align.AverageStructure(universe, in_memory=True).run(scheduler=scheduler)
         reference_coordinates = universe.trajectory.timeseries().mean(axis=1)
         assert_almost_equal(avg.results.universe.atoms.positions,
                             reference_coordinates, decimal=4)

testsuite/MDAnalysisTests/analysis/test_atomicdistances.py

@@ -115,21 +115,23 @@ def test_ad_exceptions(self, ad_ag1, ad_ag3, ad_ag4):
     # only need to test that this class correctly applies distance calcs
     # calc_bonds() is tested elsewhere
     def test_ad_pairwise_dist(self, ad_ag1, ad_ag2,
-                              expected_dist):
+                              expected_dist,
+                              scheduler):
         '''Ensure that pairwise distances between atoms are
         correctly calculated without PBCs.'''
         pairwise_no_pbc = (ad.AtomicDistances(ad_ag1, ad_ag2,
-                                              pbc=False).run())
+                                              pbc=False).run(scheduler=scheduler))
         actual = pairwise_no_pbc.results
 
         # compare with expected values from dist()
         assert_allclose(actual, expected_dist)
 
     def test_ad_pairwise_dist_pbc(self, ad_ag1, ad_ag2,
-                                  expected_pbc_dist):
+                                  expected_pbc_dist,
+                                  scheduler):
         '''Ensure that pairwise distances between atoms are
         correctly calculated with PBCs.'''
-        pairwise_pbc = (ad.AtomicDistances(ad_ag1, ad_ag2).run())
+        pairwise_pbc = (ad.AtomicDistances(ad_ag1, ad_ag2).run(scheduler=scheduler))
         actual = pairwise_pbc.results
 
         # compare with expected values from dist()