[GSoC] Parallelisation of AnalysisBase with multiprocessing and dask

package/MDAnalysis/analysis/base.py

@@ -131,8 +131,40 @@
 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)}")
 
 class Results(UserDict):
     r"""Container object for storing results.
@@ -219,6 +251,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,8 +426,8 @@ def _conclude(self):
         """
         pass  # pylint: disable=unnecessary-pass
 
-    def run(self, start=None, stop=None, step=None, frames=None,
-            verbose=None, *, progressbar_kwargs={}):
+    def _compute(self, start=None, stop=None, step=None, 
+        frames=None, verbose=None, *, progressbar_kwargs={}):
         """Perform the calculation
 
         Parameters
@@ -450,6 +487,48 @@ def run(self, start=None, stop=None, step=None, frames=None,
         self._conclude()
         return self
 
+    def _setup_bslices(self, start=None, stop=None, step=None, frames=None):
+        self._bslices = [(start, stop, step, frames)]
+
+    def _setup_scheduler(self, **kwargs):
+        self._scheduler_kwargs = kwargs
+
+    def run(self, start=None, stop=None, step=None, frames=None,
+            verbose=None, scheduler=None, 
+            *, progressbar_kwargs={},
+            ):
+        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()
+            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.run)(start=bstart, stop=bstop, step=bstep, frames=bframes)
+                    for bstart, bstop, bstep, bframes in self._bslices
+                ]
+            )
+            dask_results = computations.compute(**self._scheduler_kwargs)
+            self._remote_results = dask_results
+            self._parallel_conclude()
+
+        return self
+
+    def _parallel_conclude(self):
+        self.results = self._remote_results[0].results
+
 
 class AnalysisFromFunction(AnalysisBase):
     r"""Create an :class:`AnalysisBase` from a function working on AtomGroups

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()

testsuite/MDAnalysisTests/analysis/test_base.py

@@ -205,9 +205,9 @@ def test_start_stop_step(u, run_kwargs, frames):
     ({'frames': [True, True, False, True, False, True, True, False, True,
                  False]}, (0, 1, 3, 5, 6, 8)),
 ])
-def test_frame_slice(run_kwargs, frames):
+def test_frame_slice(run_kwargs, frames, scheduler):
     u = mda.Universe(TPR, XTC)  # dt = 100
-    an = FrameAnalysis(u.trajectory).run(**run_kwargs)
+    an = FrameAnalysis(u.trajectory).run(scheduler=scheduler, **run_kwargs)
     assert an.n_frames == len(frames)
     assert_equal(an.found_frames, frames)
     assert_equal(an.frames, frames, err_msg=FRAMES_ERR)
@@ -222,31 +222,31 @@ def test_frame_slice(run_kwargs, frames):
     ({'start': 4, 'step': 2, 'frames': [4, 6, 8]}),
     ({'start': 0, 'stop': 0, 'step': 0, 'frames': [4, 6, 8]}),
 ])
-def test_frame_fail(u, run_kwargs):
+def test_frame_fail(u, run_kwargs, scheduler):
     an = FrameAnalysis(u.trajectory)
     msg = 'start/stop/step cannot be combined with frames'
     with pytest.raises(ValueError, match=msg):
-        an.run(**run_kwargs)
+        an.run(scheduler=scheduler, **run_kwargs)
 
 
-def test_frame_bool_fail():
+def test_frame_bool_fail(scheduler):
     u = mda.Universe(TPR, XTC)  # dt = 100
     an = FrameAnalysis(u.trajectory)
     frames = [True, True, False]
     msg = 'boolean index did not match indexed array along dimension 0'
     with pytest.raises(IndexError, match=msg):
-        an.run(frames=frames)
+        an.run(scheduler=scheduler, frames=frames)
 
 
-def test_rewind():
+def test_rewind(scheduler):
     u = mda.Universe(TPR, XTC)  # dt = 100
-    an = FrameAnalysis(u.trajectory).run(frames=[0, 2, 3, 5, 9])
+    an = FrameAnalysis(u.trajectory).run(scheduler=scheduler, frames=[0, 2, 3, 5, 9])
     assert_equal(u.trajectory.ts.frame, 0)
 
 
-def test_frames_times():
+def test_frames_times(scheduler):
     u = mda.Universe(TPR, XTC)  # dt = 100
-    an = FrameAnalysis(u.trajectory).run(start=1, stop=8, step=2)
+    an = FrameAnalysis(u.trajectory).run(start=1, stop=8, step=2, scheduler=scheduler)
     frames = np.array([1, 3, 5, 7])
     assert an.n_frames == len(frames)
     assert_equal(an.found_frames, frames)
@@ -259,8 +259,8 @@ def test_verbose(u):
     assert a._verbose
 
 
-def test_verbose_progressbar(u, capsys):
-    an = FrameAnalysis(u.trajectory).run()
+def test_verbose_progressbar(u, capsys, scheduler):
+    an = FrameAnalysis(u.trajectory).run(scheduler=scheduler)
     out, err = capsys.readouterr()
     expected = ''
     actual = err.strip().split('\r')[-1]
@@ -281,9 +281,9 @@ def test_verbose_progressbar_run_with_kwargs(u, capsys):
     actual = err.strip().split('\r')[-1]
     assert actual[:30] == expected[:30]
 
-def test_incomplete_defined_analysis(u):
+def test_incomplete_defined_analysis(u, scheduler):
     with pytest.raises(NotImplementedError):
-        IncompleteAnalysis(u.trajectory).run()
+        IncompleteAnalysis(u.trajectory).run(scheduler=scheduler)
 
 
 def test_old_api(u):
@@ -330,15 +330,15 @@ def test_results_type(u):
     (20, 50, 2, 15),
     (20, 50, None, 30)
 ])
-def test_AnalysisFromFunction(u, start, stop, step, nframes):
+def test_AnalysisFromFunction(u, start, stop, step, nframes, scheduler):
     ana1 = base.AnalysisFromFunction(simple_function, mobile=u.atoms)
-    ana1.run(start=start, stop=stop, step=step)
+    ana1.run(start=start, stop=stop, step=step, scheduler=scheduler)
 
     ana2 = base.AnalysisFromFunction(simple_function, u.atoms)
-    ana2.run(start=start, stop=stop, step=step)
+    ana2.run(start=start, stop=stop, step=step, scheduler=scheduler)
 
     ana3 = base.AnalysisFromFunction(simple_function, u.trajectory, u.atoms)
-    ana3.run(start=start, stop=stop, step=step)
+    ana3.run(start=start, stop=stop, step=step, scheduler=scheduler)
 
     frames = []
     times = []
@@ -365,26 +365,26 @@ def mass_xyz(atomgroup1, atomgroup2, masses):
     return atomgroup1.positions * masses
 
 
-def test_AnalysisFromFunction_args_content(u):
+def test_AnalysisFromFunction_args_content(u, scheduler):
     protein = u.select_atoms('protein')
     masses = protein.masses.reshape(-1, 1)
     another = mda.Universe(TPR, XTC).select_atoms("protein")
     ans = base.AnalysisFromFunction(mass_xyz, protein, another, masses)
     assert len(ans.args) == 3
-    result = np.sum(ans.run().results.timeseries)
+    result = np.sum(ans.run(scheduler=scheduler).results.timeseries)
     assert_almost_equal(result, -317054.67757345125, decimal=6)
     assert (ans.args[0] is protein) and (ans.args[1] is another)
     assert ans._trajectory is protein.universe.trajectory
 
 
-def test_analysis_class():
+def test_analysis_class(scheduler):
     ana_class = base.analysis_class(simple_function)
     assert issubclass(ana_class, base.AnalysisBase)
     assert issubclass(ana_class, base.AnalysisFromFunction)
 
     u = mda.Universe(PSF, DCD)
     step = 2
-    ana = ana_class(u.atoms).run(step=step)
+    ana = ana_class(u.atoms).run(step=step, scheduler=scheduler)
 
     results = []
     for ts in u.trajectory[::step]: