diff --git a/pmda/leaflet.py b/pmda/leaflet.py index ec648b5e..dbd55b91 100644 --- 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 index f595740c..81322a19 100644 --- 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':