Single Dask client is enough

Author: w8jcik

src/moldrug/runner.py

@@ -6,7 +6,6 @@
 
 try:
     from dask.distributed import Client
-    from dask_jobqueue import JobQueueCluster
 
     dask_available = True
 except ImportError:
@@ -26,15 +25,10 @@ def parallel_execution_multiprocessing(func: Callable, entries: List, process_co
 
 
 if dask_available:
-    def parallel_execution_dask_local(func: Callable, entries: List, process_count: int, process_threads_count: int):
-        client = Client(n_workers=process_count, threads_per_worker=process_threads_count)
-
+    def parallel_execution_dask_local(func: Callable, entries: List, client: Client):
         return client.gather(client.map(func, entries))
 
-
-    def parallel_execution_dask_cluster(func: Callable, entries: List, cluster: JobQueueCluster):
-        client = Client(cluster)
-
+    def parallel_execution_dask_cluster(func: Callable, entries: List, client: Client):
         return client.gather(client.map(func, entries))
 
 
@@ -50,10 +44,11 @@ class Runner:
     thread_count: Optional[int]
     process_count: Optional[int]
     dask_cluster = None
+    client = None
 
     def __init__(self, mode: RunnerMode, thread_count: Optional[int] = None, process_count: Optional[int] = None, dask_cluster = None):
         if mode in [RunnerMode.DASK_LOCAL, RunnerMode.DASK_JOB_QUEUE_CLUSTER]:
-            assert dask_cluster is not None, "Execution using Dask requires installation of optional dependencies. The optional pip package group is called 'cluster'"
+            assert dask_available is not None, "Execution using Dask requires installation of optional dependencies. The optional pip package group is called 'cluster'"
 
         if mode is RunnerMode.SERIAL:
             assert thread_count is None and process_count is None and dask_cluster is None, "Serial execution doesn't take any parameters."
@@ -62,7 +57,7 @@ def __init__(self, mode: RunnerMode, thread_count: Optional[int] = None, process
         elif mode in [RunnerMode.DASK_LOCAL]:
             assert (thread_count is not None or process_count is not None) and dask_cluster is None, "Only process count and/or thread count are needed."
         elif mode in [RunnerMode.DASK_JOB_QUEUE_CLUSTER]:
-            assert thread_count is None and process_count is None and dask_cluster is not None, "Dask cluster execution takes only a Dask cluster object."
+            assert thread_count is None and process_count is None and dask_cluster is not None, "Dask execution takes only a Dask cluster object."
         else:
             assert False
 
@@ -71,14 +66,20 @@ def __init__(self, mode: RunnerMode, thread_count: Optional[int] = None, process
         self.process_count = process_count
         self.cluster = dask_cluster
 
+        if mode in [RunnerMode.DASK_LOCAL]:
+            self.client = Client(n_workers=process_count, threads_per_worker=thread_count)
+
+        if mode in [RunnerMode.DASK_JOB_QUEUE_CLUSTER]:
+            self.client = Client(dask_cluster)
+
     def run(self, func: Callable, entries: List):
         if self.mode == RunnerMode.SERIAL:
             return serial_execution(func, entries)
         elif self.mode == RunnerMode.MULTIPROCESSING:
             return parallel_execution_multiprocessing(func, entries, process_count=self.process_count)
         elif self.mode == RunnerMode.DASK_LOCAL:
-            return parallel_execution_dask_local(func, entries, process_count=self.process_count, process_threads_count=self.thread_count)
+            return parallel_execution_dask_local(func, entries, self.client)
         if self.mode == RunnerMode.DASK_JOB_QUEUE_CLUSTER:
-            return parallel_execution_dask_cluster(func, entries, cluster=self.cluster)
+            return parallel_execution_dask_cluster(func, entries, self.client)
         else:
             assert False