fix missing import making distributed optional

MDAnalysis · Oct 31, 2018 · 114a2b0 · 114a2b0
1 parent dd88867
commit 114a2b0
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 10 deletions.
diff --git a/pmda/leaflet.py b/pmda/leaflet.py
@@ -268,9 +268,26 @@ def run(self,
         if scheduler is None and n_jobs == 1:
             scheduler = 'single-threaded'
 
-        # fall back to multiprocessing, we tried everything
-        if scheduler is None:
-            scheduler = 'multiprocessing'
+        if n_blocks is None:
+            if scheduler == 'multiprocessing':
+                n_blocks = n_jobs
+            else:
+                try:
+                    from dask import distributed
+                    if isinstance(scheduler, distributed.Client):
+                        n_blocks = len(scheduler.ncores())
+                    else:
+                        n_blocks = 1
+                        warnings.warn(
+                            "Couldn't guess ideal number of blocks from scheduler. "
+                            "Setting n_blocks=1. "
+                            "Please provide `n_blocks` in call to method.")
+                except ImportError:
+                    n_blocks = 1
+                    warnings.warn(
+                        "Couldn't guess ideal number of blocks from scheduler. "
+                        "Setting n_blocks=1. "
+                        "Please provide `n_blocks` in call to method.")
 
         scheduler_kwargs = {'scheduler': scheduler}
         if scheduler == 'multiprocessing':

diff --git a/pmda/parallel.py b/pmda/parallel.py
@@ -315,14 +315,23 @@ def run(self,
         if n_blocks is None:
             if scheduler == 'multiprocessing':
                 n_blocks = n_jobs
-            elif isinstance(scheduler, dask.distributed.Client):
-                n_blocks = len(scheduler.ncores())
             else:
-                n_blocks = 1
-                warnings.warn(
-                    "Couldn't guess ideal number of blocks from scheduler. "
-                    "Setting n_blocks=1. "
-                    "Please provide `n_blocks` in call to method.")
+                try:
+                    from dask import distributed
+                    if isinstance(scheduler, distributed.Client):
+                        n_blocks = len(scheduler.ncores())
+                    else:
+                        n_blocks = 1
+                        warnings.warn(
+                            "Couldn't guess ideal number of blocks from scheduler. "
+                            "Setting n_blocks=1. "
+                            "Please provide `n_blocks` in call to method.")
+                except ImportError:
+                    n_blocks = 1
+                    warnings.warn(
+                        "Couldn't guess ideal number of blocks from scheduler. "
+                        "Setting n_blocks=1. "
+                        "Please provide `n_blocks` in call to method.")
 
         scheduler_kwargs = {'scheduler': scheduler}
         if scheduler == 'multiprocessing':