fix: propagate remote exception traceback to parameter server

Author: SongXiaoXi

checkpoint_engine/ps.py

@@ -1284,9 +1284,9 @@ def _update_per_bucket(
                     dist.broadcast(buffer_b, src=brank)
                     resp = socket.recv()
                     if resp != b"":
-                        exception_obj = pickle.loads(resp)
+                        msg = pickle.loads(resp)
                         logger.error(
-                            f"[rank{self._rank}] receive error response '{type(exception_obj).__name__}: {exception_obj}' from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}"
+                            f"[rank{self._rank}] receive error response from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}: {msg}"
                         )
                         ret_code.fill_(1)
                     dist.all_reduce(ret_code, op=dist.ReduceOp.SUM)

checkpoint_engine/worker.py

@@ -4,6 +4,7 @@
 
 import torch
 import zmq
+import traceback
 
 from checkpoint_engine.device_utils import DeviceManager, npu_generate_uuid
 
@@ -63,7 +64,10 @@ def update_weights_from_ipc(
         assert buffer.dtype == torch.uint8
         socket.send(b"")
     except Exception as e:
-        socket.send_pyobj(e)
+        msg = "".join(
+            traceback.format_exception(type(e), e, e.__traceback__)
+        )
+        socket.send_pyobj(msg)
         socket.recv()  # wait for ack
         raise
     try:
@@ -83,7 +87,10 @@ def update_weights_from_ipc(
                 except Exception as e:  # noqa: BLE001
                     # Send exception back to Parameter Server.
                     # Don't raise here. Because all workers should quit in the same way by receiving the exception from PS
-                    socket.send_pyobj(e)
+                    msg = "".join(
+                        traceback.format_exception(type(e), e, e.__traceback__)
+                    )
+                    socket.send_pyobj(msg)
             elif isinstance(
                 payload, Exception
             ):  # error occurred, got force quit signal from Parameter Server