From 9b663c867bc7b279797db4e7bbdbd98494454a28 Mon Sep 17 00:00:00 2001 From: Casey Jao Date: Wed, 7 Jun 2023 20:54:47 -0400 Subject: [PATCH] Runner-Ex: fix Dask poll --- covalent/executor/executor_plugins/dask.py | 10 +++++----- tests/functional_tests/workflow_cancellation_test.py | 3 --- 2 files changed, 5 insertions(+), 8 deletions(-) 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)