Repository: dask/distributed

Issue:
Catch exceptions on both send and recv in UCX comm
Following on from a review of #6996, there are a few other places in the UCX communicator backend where we should be catching exceptions on message transfer.

In addition, some of the try catch blocks can have their scope tightened (to avoid hiding other errors).

cc @pentschev, I will self-assign, but don't have privileges to do so.


Discussion:

Resolving patch:
diff --git a/distributed/comm/ucx.py b/distributed/comm/ucx.py
--- a/distributed/comm/ucx.py
+++ b/distributed/comm/ucx.py
@@ -254,27 +254,28 @@ async def write(
     ) -> int:
         if self.closed():
             raise CommClosedError("Endpoint is closed -- unable to send message")
-        try:
-            if serializers is None:
-                serializers = ("cuda", "dask", "pickle", "error")
-            # msg can also be a list of dicts when sending batched messages
-            frames = await to_frames(
-                msg,
-                serializers=serializers,
-                on_error=on_error,
-                allow_offload=self.allow_offload,
-            )
-            nframes = len(frames)
-            cuda_frames = tuple(hasattr(f, "__cuda_array_interface__") for f in frames)
-            sizes = tuple(nbytes(f) for f in frames)
-            cuda_send_frames, send_frames = zip(
-                *(
-                    (is_cuda, each_frame)
-                    for is_cuda, each_frame in zip(cuda_frames, frames)
-                    if nbytes(each_frame) > 0
-                )
+
+        if serializers is None:
+            serializers = ("cuda", "dask", "pickle", "error")
+        # msg can also be a list of dicts when sending batched messages
+        frames = await to_frames(
+            msg,
+            serializers=serializers,
+            on_error=on_error,
+            allow_offload=self.allow_offload,
+        )
+        nframes = len(frames)
+        cuda_frames = tuple(hasattr(f, "__cuda_array_interface__") for f in frames)
+        sizes = tuple(nbytes(f) for f in frames)
+        cuda_send_frames, send_frames = zip(
+            *(
+                (is_cuda, each_frame)
+                for is_cuda, each_frame in zip(cuda_frames, frames)
+                if nbytes(each_frame) > 0
             )
+        )
 
+        try:
             # Send meta data
 
             # Send close flag and number of frames (_Bool, int64)
@@ -326,7 +327,7 @@ async def read(self, deserializers=("cuda", "dask", "pickle", "error")):
             header = struct.unpack(header_fmt, header)
             cuda_frames, sizes = header[:nframes], header[nframes:]
         except BaseException as e:
-            # In addition to UCX exceptions, may be CancelledError or a another
+            # In addition to UCX exceptions, may be CancelledError or another
             # "low-level" exception. The only safe thing to do is to abort.
             # (See also https://github.com/dask/distributed/pull/6574).
             self.abort()
@@ -352,14 +353,29 @@ async def read(self, deserializers=("cuda", "dask", "pickle", "error")):
             if any(cuda_recv_frames):
                 synchronize_stream(0)
 
-            for each_frame in recv_frames:
-                await self.ep.recv(each_frame)
-            msg = await from_frames(
-                frames,
-                deserialize=self.deserialize,
-                deserializers=deserializers,
-                allow_offload=self.allow_offload,
-            )
+            try:
+                for each_frame in recv_frames:
+                    await self.ep.recv(each_frame)
+            except BaseException as e:
+                # In addition to UCX exceptions, may be CancelledError or another
+                # "low-level" exception. The only safe thing to do is to abort.
+                # (See also https://github.com/dask/distributed/pull/6574).
+                self.abort()
+                raise CommClosedError(
+                    f"Connection closed by writer.\nInner exception: {e!r}"
+                )
+
+            try:
+                msg = await from_frames(
+                    frames,
+                    deserialize=self.deserialize,
+                    deserializers=deserializers,
+                    allow_offload=self.allow_offload,
+                )
+            except EOFError:
+                # Frames possibly garbled or truncated by communication error
+                self.abort()
+                raise CommClosedError("Aborted stream on truncated data")
             return msg
 
     async def close(self):


Test patch: