From 918e3fb13aee02bb75e22f1c2a892cb2cbadc1cb Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Thu, 21 Oct 2021 21:34:54 +0200
Subject: [PATCH] Handle `UCXNotConnected` error (#5449)

Handle UCXNotConnected error
---
 distributed/comm/ucx.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/distributed/comm/ucx.py b/distributed/comm/ucx.py
index ecd6eaef72..992367bdf8 100644
--- 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,