From 27cd36e6e2e808464c8343066b03db5db2d15413 Mon Sep 17 00:00:00 2001 From: Gene Der Su Date: Fri, 1 Nov 2024 15:08:23 -0700 Subject: [PATCH] [Bugfix] PicklingError on RayTaskError (#9934) Signed-off-by: Gene Su --- vllm/engine/multiprocessing/engine.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/vllm/engine/multiprocessing/engine.py b/vllm/engine/multiprocessing/engine.py index eb1512ca17822..a73b4c825b11c 100644 --- a/vllm/engine/multiprocessing/engine.py +++ b/vllm/engine/multiprocessing/engine.py @@ -5,6 +5,7 @@ import cloudpickle import zmq +from ray.exceptions import RayTaskError from vllm import AsyncEngineArgs, SamplingParams # yapf conflicts with isort for this block @@ -305,6 +306,11 @@ def _health_check(self): def _send_outputs(self, outputs: REQUEST_OUTPUTS_T): """Send List of RequestOutput to RPCClient.""" if outputs: + # RayTaskError might not pickelable here. We need to unpack the + # underlying exception as the real exception in the output. + if (isinstance(outputs, RPCError) + and isinstance(outputs.exception, RayTaskError)): + outputs.exception = outputs.exception.cause output_bytes = pickle.dumps(outputs) self.output_socket.send_multipart((output_bytes, ), copy=False)