Close client during wait_until_finish; rm async.

alxmrs · Sep 22, 2022 · 09365f6 · 09365f6
1 parent 676d752
commit 09365f6
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/sdks/python/apache_beam/runners/dask/dask_runner.py b/sdks/python/apache_beam/runners/dask/dask_runner.py
@@ -88,13 +88,14 @@ def wait_until_finish(self, duration=None) -> PipelineState:
         # Convert milliseconds to seconds
         duration /= 1000
       self.client.wait_for_workers(timeout=duration)
-      self.client.gather(self.futures, errors='raise', asynchronous=True)
+      self.client.gather(self.futures, errors='raise')
       self._state = PipelineState.DONE
     except:  # pylint: disable=broad-except
       self._state = PipelineState.FAILED
       raise
-    # finally:
-    #   self.client.close(timeout=duration)
+    finally:
+      self.client.close()
+      self.client.cluster.close()
     return self._state
 
   def cancel(self) -> PipelineState: