Merge pull request #25 from djarecka/adi611-adi611-patch-updatedask-1

adi611 · web-flow · commit 3d687b06db10 · 2023-08-19T15:14:50.000+05:30
properly closing the dask Client
diff --git a/.github/workflows/testdask.yml b/.github/workflows/testdask.yml
@@ -27,7 +27,7 @@ jobs:
         uses: actions/checkout@v3
         with:
           repository: ${{ github.repository }}
-  
+
       - name: Setup Python version ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:
diff --git a/pydra/engine/workers.py b/pydra/engine/workers.py
@@ -875,18 +875,16 @@ def run_el(self, runnable, rerun=False, **kwargs):
 
     async def exec_dask(self, runnable, rerun=False):
         """Run a task (coroutine wrapper)."""
-        if self.client is None:
-            from dask.distributed import Client
-
-            self.client = await Client(**self.client_args, asynchronous=True)
-
-        if isinstance(runnable, TaskBase):
-            future = self.client.submit(runnable._run, rerun)
-            result = await future
-        else:  # it could be tuple that includes pickle files with tasks and inputs
-            ind, task_main_pkl, task_orig = runnable
-            future = self.client.submit(load_and_run, task_main_pkl, ind, rerun)
-            result = await future
+        from dask.distributed import Client
+
+        async with Client(**self.client_args, asynchronous=True) as client:
+            if isinstance(runnable, TaskBase):
+                future = client.submit(runnable._run, rerun)
+                result = await future
+            else:  # it could be tuple that includes pickle files with tasks and inputs
+                ind, task_main_pkl, task_orig = runnable
+                future = client.submit(load_and_run, task_main_pkl, ind, rerun)
+                result = await future
         return result
 
     def close(self):
