diff --git a/covalent/executor/executor_plugins/dask.py b/covalent/executor/executor_plugins/dask.py index d5f22fbd3..96df57e1f 100644 --- 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 index 9d1eaee4d..76ae2fa11 100644 --- 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)