Runner-Ex: fix Dask poll

AgnostiqHQ · Jun 8, 2023 · 9b663c8 · 9b663c8
1 parent 4fc0fa0
commit 9b663c8
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 8 deletions.
diff --git a/covalent/executor/executor_plugins/dask.py b/covalent/executor/executor_plugins/dask.py
@@ -273,14 +273,14 @@ async def send(
     async def poll(self, task_group_metadata: Dict, poll_data: Any):
         fut = _futures.pop(poll_data)
         app_log.debug(f"Future {fut}")
-        await fut
+        try:
+            await fut
+        except CancelledError:
+            raise TaskCancelledError()
 
         _clients.pop(poll_data)
 
-        if fut.cancelled():
-            raise TaskCancelledError()
-        else:
-            return {"status": StatusEnum.READY.value}
+        return {"status": StatusEnum.READY.value}
 
         # raise NotImplementedError
 

diff --git a/tests/functional_tests/workflow_cancellation_test.py b/tests/functional_tests/workflow_cancellation_test.py
@@ -52,9 +52,6 @@ def workflow(x):
     ct.cancel(dispatch_id)
 
     result = ct.get_result(dispatch_id, wait=True)
-    print(result.error)
-    for node_id in result.lattice.transport_graph._graph.nodes:
-        print(result.get_node_result(node_id))
 
     assert result.status == ct.status.CANCELLED
     rm._delete_result(dispatch_id)