[BugFix] Lazy import ray (#10021)

This commit is contained in:
Gene Der Su 2024-11-05 02:04:10 -08:00 committed by GitHub
parent ad23318928
commit 7a83b1aec0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -5,7 +5,6 @@ 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
@ -306,11 +305,17 @@ 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
try:
from ray.exceptions import RayTaskError
# 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
except ImportError:
pass
output_bytes = pickle.dumps(outputs)
self.output_socket.send_multipart((output_bytes, ), copy=False)