A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 05:15 UTC
[Core] Move ray-specific WorkerWrapperBase methods to RayWorkerWrapper (#35328)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 4, 2026 at 20:11 UTC
16c472abe7e0e77e7924080bd4ed55bdceb86c53
2 files changed
+24
-29
vllm/v1/executor/ray_utils.py
+24
@@ -16,6 +16,7 @@ from vllm.platforms import current_platform
16
from vllm.sequence import IntermediateTensors
17
from vllm.utils.network_utils import get_ip
18
from vllm.v1.outputs import AsyncModelRunnerOutput
19
+from vllm.v1.serial_utils import run_method
20
from vllm.v1.worker.worker_base import WorkerWrapperBase
21
22
if TYPE_CHECKING:
@@ -50,6 +51,29 @@ try:
51
# that thread.
52
self.compiled_dag_cuda_device_set = False
53
54
+ def adjust_rank(self, rank_mapping: dict[int, int]) -> None:
55
+ """
56
+ Adjust the rpc_rank based on the given mapping.
57
+ It is only used during the initialization of the executor,
58
+ to adjust the rpc_rank of workers after we create all workers.
59
+ """
60
+ if self.rpc_rank in rank_mapping:
61
+ self.rpc_rank = rank_mapping[self.rpc_rank]
62
+
63
+ def execute_method(self, method: str | bytes, *args, **kwargs):
64
+ try:
65
+ return run_method(self, method, args, kwargs)
66
+ except Exception as e:
67
+ # if the driver worker also execute methods,
68
+ # exceptions in the rest worker may cause deadlock in rpc
69
+ # see https://github.com/vllm-project/vllm/issues/3455
70
+ msg = (
71
+ f"Error executing method {method!r}. "
72
+ "This might cause deadlock in distributed execution."
73
+ )
74
+ logger.exception(msg)
75
+ raise e
76
+
77
def get_node_ip(self) -> str:
78
return get_ip()
79
vllm/v1/worker/worker_base.py
-29
@@ -15,7 +15,6 @@ from vllm.tracing import instrument
15
from vllm.utils.import_utils import resolve_obj_by_qualname
16
from vllm.utils.system_utils import update_environment_variables
17
from vllm.v1.kv_cache_interface import KVCacheSpec
18
-from vllm.v1.serial_utils import run_method
18
19
if TYPE_CHECKING:
20
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -211,15 +210,6 @@ class WorkerWrapperBase:
210
if self.worker is not None:
211
self.worker.shutdown()
212
214
- def adjust_rank(self, rank_mapping: dict[int, int]) -> None:
215
- """
216
- Adjust the rpc_rank based on the given mapping.
217
- It is only used during the initialization of the executor,
218
- to adjust the rpc_rank of workers after we create all workers.
219
- """
220
- if self.rpc_rank in rank_mapping:
221
- self.rpc_rank = rank_mapping[self.rpc_rank]
222
-
213
def update_environment_variables(
214
self,
215
envs_list: list[dict[str, str]],
@@ -325,25 +315,6 @@ class WorkerWrapperBase:
315
# To make vLLM config available during device initialization
316
self.worker.init_device() # type: ignore
317
328
- def execute_method(self, method: str | bytes, *args, **kwargs):
329
- try:
330
- # method resolution order:
331
- # if a method is defined in this class, it will be called directly.
332
- # otherwise, since we define `__getattr__` and redirect attribute
333
- # query to `self.worker`, the method will be called on the worker.
334
- return run_method(self, method, args, kwargs)
335
- except Exception as e:
336
- # if the driver worker also execute methods,
337
- # exceptions in the rest worker may cause deadlock in rpc like ray
338
- # see https://github.com/vllm-project/vllm/issues/3455
339
- # print the error and inform the user to solve the error
340
- msg = (
341
- f"Error executing method {method!r}. "
342
- "This might cause deadlock in distributed execution."
343
- )
344
- logger.exception(msg)
345
- raise e
346
-
318
def __getattr__(self, attr: str):
319
return getattr(self.worker, attr)
320