[Bugfix] PicklingError on RayTaskError (#9934)

Signed-off-by: Gene Su <e870252314@gmail.com>
This commit is contained in:
Gene Der Su 2024-11-01 15:08:23 -07:00 committed by GitHub
parent 18bd7587b7
commit 27cd36e6e2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -5,6 +5,7 @@ from typing import Iterator, List, Optional, Union
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 @@ class MQLLMEngine:
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)