Handle UCXNotConnected error (#5449)

Handle UCXNotConnected error
dask · Oct 21, 2021 · 918e3fb · 918e3fb
1 parent 3afc670
commit 918e3fb
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/distributed/comm/ucx.py b/distributed/comm/ucx.py
@@ -383,10 +383,10 @@ async def connect(self, address: str, deserialize=True, **connection_args) -> UC
         init_once()
         try:
             ep = await ucx_create_endpoint(ip, port)
-        except (
-            ucp.exceptions.UCXCloseError,
-            ucp.exceptions.UCXCanceled,
-        ) + (getattr(ucp.exceptions, "UCXConnectionReset", ()),):
+        except (ucp.exceptions.UCXCloseError, ucp.exceptions.UCXCanceled,) + (
+            getattr(ucp.exceptions, "UCXConnectionReset", ()),
+            getattr(ucp.exceptions, "UCXNotConnected", ()),
+        ):
             raise CommClosedError("Connection closed before handshake completed")
         return self.comm_class(
             ep,