Changing to async produces a timeout error instead of stuck in infini…

…te loop.
alxmrs · Sep 21, 2022 · 676d752 · 676d752
1 parent 41623ec
commit 676d752
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/sdks/python/apache_beam/runners/dask/dask_runner.py b/sdks/python/apache_beam/runners/dask/dask_runner.py
@@ -88,11 +88,13 @@ 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')
+      self.client.gather(self.futures, errors='raise', asynchronous=True)
       self._state = PipelineState.DONE
     except:  # pylint: disable=broad-except
       self._state = PipelineState.FAILED
       raise
+    # finally:
+    #   self.client.close(timeout=duration)
     return self._state
 
   def cancel(self) -> PipelineState: