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 14, 05:30 UTC
[V0 Deprecation] Remove `vllm.worker` and update according imports (#25901)
Aaron Pham committed
Sep 29, 2025 at 19:26 UTC
6a113d9aed8221a9c234535958e70e34ab6cac5b
11 files changed
+276
-327
tests/model_executor/model_loader/tensorizer_loader/conftest.py
+1
-1
@@ -10,7 +10,7 @@ from vllm.model_executor.model_loader import tensorizer as tensorizer_mod
10
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
11
from vllm.utils import get_distributed_init_method, get_ip, get_open_port
12
from vllm.v1.executor.abstract import UniProcExecutor
13
-from vllm.worker.worker_base import WorkerWrapperBase
13
+from vllm.v1.worker.worker_base import WorkerWrapperBase
14
15
MODEL_REF = "facebook/opt-125m"
16
tools/pre_commit/check_pickle_imports.py
-1
@@ -36,7 +36,6 @@ ALLOWED_FILES = {
36
'benchmarks/cutlass_benchmarks/w8a8_benchmarks.py',
37
'benchmarks/cutlass_benchmarks/sparse_benchmarks.py',
38
# cloudpickle
39
- 'vllm/worker/worker_base.py',
39
'vllm/executor/mp_distributed_executor.py',
40
'vllm/executor/ray_distributed_executor.py',
41
'vllm/entrypoints/llm.py',
vllm/executor/executor_base.py
+5
-5
@@ -19,7 +19,7 @@ from vllm.sequence import ExecuteModelRequest
19
from vllm.tasks import SupportedTask
20
from vllm.utils import make_async
21
from vllm.v1.outputs import PoolerOutput, SamplerOutput
22
-from vllm.worker.worker_base import WorkerBase
22
+from vllm.v1.worker.worker_base import WorkerBase
23
24
logger = init_logger(__name__)
25
@@ -30,7 +30,7 @@ class ExecutorBase(ABC):
30
"""Base class for all executors.
31
32
An executor is responsible for executing the model on one device,
33
- or it can be a distributed executor
33
+ or it can be a distributed executor
34
that can execute the model on multiple devices.
35
"""
36
@@ -83,7 +83,7 @@ class ExecutorBase(ABC):
83
84
Returns:
85
A list containing the results from each worker.
86
-
86
+
87
Note:
88
It is recommended to use this API to only pass control messages,
89
and set up data-plane communication to pass data.
@@ -100,7 +100,7 @@ class ExecutorBase(ABC):
100
101
Returns a tuple `(num_gpu_blocks, num_cpu_blocks)`, where
102
`num_gpu_blocks` are blocks that are "active" on the device and can be
103
- appended to.
103
+ appended to.
104
`num_cpu_blocks` refers to "swapped" blocks in CPU memory and cannot be
105
appended to.
106
"""
@@ -327,7 +327,7 @@ class DistributedExecutorBase(ExecutorBase):
327
run only in the remote TP workers, not the driver worker.
328
It will also be run asynchronously and return a list of futures
329
rather than blocking on the results.
330
-
330
+
331
# TODO: simplify and merge with collective_rpc
332
"""
333
raise NotImplementedError
vllm/executor/ray_utils.py
+1
-1
@@ -16,7 +16,7 @@ from vllm.logger import init_logger
16
from vllm.platforms import current_platform
17
from vllm.sequence import ExecuteModelRequest, IntermediateTensors
18
from vllm.utils import get_ip
19
-from vllm.worker.worker_base import WorkerWrapperBase
19
+from vllm.v1.worker.worker_base import WorkerWrapperBase
20
21
if TYPE_CHECKING:
22
from vllm.v1.core.sched.output import SchedulerOutput
vllm/executor/uniproc_executor.py
+5
-5
@@ -19,7 +19,7 @@ from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,
19
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
20
from vllm.v1.executor.utils import get_and_update_mm_cache
21
from vllm.v1.outputs import AsyncModelRunnerOutput
22
-from vllm.worker.worker_base import WorkerWrapperBase
22
+from vllm.v1.worker.worker_base import WorkerWrapperBase
23
24
logger = init_logger(__name__)
25
@@ -160,10 +160,10 @@ class ExecutorWithExternalLauncher(UniProcExecutor):
160
"""
161
Determine the number of available KV blocks.
162
Add an additional all_reduce to get the min across all ranks.
163
- Note that even if we have the same `gpu_memory_utilization` and
164
- `swap_space`, the available memory in every rank might still
165
- differ because NCCL can take different amounts of memory in
166
- different ranks. Therefore, it is necessary to test if all ranks
163
+ Note that even if we have the same `gpu_memory_utilization` and
164
+ `swap_space`, the available memory in every rank might still
165
+ differ because NCCL can take different amounts of memory in
166
+ different ranks. Therefore, it is necessary to test if all ranks
167
agree on the same KV cache configuration.
168
"""
169
a, b = super().determine_num_available_blocks()
vllm/platforms/cuda.py
+1
-11
@@ -110,17 +110,7 @@ class CudaPlatformBase(Platform):
110
model_config = vllm_config.model_config
111
112
if parallel_config.worker_cls == "auto":
113
- if vllm_config.speculative_config:
114
- if not envs.VLLM_USE_V1:
115
- raise NotImplementedError(
116
- "Speculative decoding is not supported on vLLM V0.")
117
- parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
118
- else:
119
- if envs.VLLM_USE_V1:
120
- parallel_config.worker_cls = \
121
- "vllm.v1.worker.gpu_worker.Worker"
122
- else:
123
- parallel_config.worker_cls = "vllm.worker.worker.Worker"
113
+ parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
114
115
cache_config = vllm_config.cache_config
116
if cache_config and cache_config.block_size is None:
vllm/platforms/rocm.py
+1
-11
@@ -327,17 +327,7 @@ class RocmPlatform(Platform):
327
cache_config.block_size = 16
328
329
if parallel_config.worker_cls == "auto":
330
- if vllm_config.speculative_config:
331
- if not use_v1:
332
- raise NotImplementedError(
333
- "Speculative decoding is not supported on vLLM V0.")
334
- parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
335
- else:
336
- if use_v1:
337
- parallel_config.worker_cls = \
338
- "vllm.v1.worker.gpu_worker.Worker"
339
- else:
340
- parallel_config.worker_cls = "vllm.worker.worker.Worker"
330
+ parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
331
# Aiter rms norm perform best when CUDA Graph capture is enabled.
332
if (use_v1 and use_aiter_rms_norm and not is_eager_execution
333
and "-rms_norm" not in compilation_config.custom_ops):
vllm/v1/executor/multiproc_executor.py
+2
-2
@@ -41,7 +41,7 @@ from vllm.v1.executor.abstract import Executor, FailureCallback
41
from vllm.v1.executor.utils import get_and_update_mm_cache
42
from vllm.v1.outputs import (AsyncModelRunnerOutput, DraftTokenIds,
43
ModelRunnerOutput)
44
-from vllm.worker.worker_base import WorkerWrapperBase
44
+from vllm.v1.worker.worker_base import WorkerWrapperBase
45
46
logger = init_logger(__name__)
47
@@ -702,7 +702,7 @@ class WorkerProc:
702
703
def set_multiprocessing_worker_envs():
704
""" Set up environment variables that should be used when there are workers
705
- in a multiprocessing environment. This should be called by the parent
705
+ in a multiprocessing environment. This should be called by the parent
706
process before worker processes are created"""
707
708
_maybe_force_spawn()
vllm/v1/worker/worker_base.py
+260
-11
@@ -1,23 +1,35 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Optional
4
+from __future__ import annotations
5
+
6
+import os
7
+from typing import Any, Callable, Optional, TypeVar, Union
8
9
import torch
10
import torch.nn as nn
11
9
-from vllm.config import VllmConfig
12
+from vllm.config import VllmConfig, set_current_vllm_config
13
from vllm.logger import init_logger
14
+from vllm.lora.request import LoRARequest
15
+from vllm.sequence import ExecuteModelRequest
16
+from vllm.utils import (enable_trace_function_call_for_thread,
17
+ resolve_obj_by_qualname, run_method,
18
+ update_environment_variables,
19
+ warn_for_unimplemented_methods)
20
from vllm.v1.kv_cache_interface import KVCacheSpec
12
-from vllm.worker.worker_base import WorkerBase as WorkerBaseV0
21
+from vllm.v1.outputs import SamplerOutput
22
23
logger = init_logger(__name__)
24
25
+_R = TypeVar("_R")
26
17
-class WorkerBase(WorkerBaseV0):
18
- """
19
- Abstract class for v1 worker, mainly define some methods for v1.
20
- For methods shared by v0 and v1, define them in v0 WorkerBase
27
+
28
+@warn_for_unimplemented_methods
29
+class WorkerBase:
30
+ """Worker interface that allows vLLM to cleanly separate implementations for
31
+ different hardware. Also abstracts control plane communication, e.g., to
32
+ communicate request metadata to other workers.
33
"""
34
35
def __init__(
@@ -27,10 +39,10 @@ class WorkerBase(WorkerBaseV0):
39
rank: int,
40
distributed_init_method: str,
41
is_driver_worker: bool = False,
30
- ):
42
+ ) -> None:
43
"""
44
Initialize common worker components.
33
-
45
+
46
Args:
47
vllm_config: Complete vLLM configuration
48
local_rank: Local device index
@@ -39,8 +51,21 @@ class WorkerBase(WorkerBaseV0):
51
is_driver_worker: Whether this worker handles driver
52
responsibilities
53
"""
42
- # Configuration storage
43
- super().__init__(vllm_config=vllm_config)
54
+ self.vllm_config = vllm_config
55
+ self.model_config = vllm_config.model_config
56
+ self.cache_config = vllm_config.cache_config
57
+ self.lora_config = vllm_config.lora_config
58
+ self.load_config = vllm_config.load_config
59
+ self.parallel_config = vllm_config.parallel_config
60
+ self.scheduler_config = vllm_config.scheduler_config
61
+ self.device_config = vllm_config.device_config
62
+ self.speculative_config = vllm_config.speculative_config
63
+ self.observability_config = vllm_config.observability_config
64
+ self.kv_transfer_config = vllm_config.kv_transfer_config
65
+ self.compilation_config = vllm_config.compilation_config
66
+
67
+ from vllm.platforms import current_platform
68
+ self.current_platform = current_platform
69
70
self.parallel_config.rank = rank
71
self.local_rank = local_rank
@@ -63,3 +88,227 @@ class WorkerBase(WorkerBaseV0):
88
def check_health(self) -> None:
89
"""Basic health check (override for device-specific checks)."""
90
return
91
+
92
+ def init_device(self) -> None:
93
+ """Initialize device state, such as loading the model or other on-device
94
+ memory allocations.
95
+ """
96
+ raise NotImplementedError
97
+
98
+ def initialize_cache(self, num_gpu_blocks: int,
99
+ num_cpu_blocks: int) -> None:
100
+ """Initialize the KV cache with the given size in blocks.
101
+ """
102
+ raise NotImplementedError
103
+
104
+ def get_model(self) -> nn.Module:
105
+ raise NotImplementedError
106
+
107
+ def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R:
108
+ """Apply a function on the model inside this worker."""
109
+ return fn(self.get_model())
110
+
111
+ def load_model(self) -> None:
112
+ """Load model onto target device."""
113
+ raise NotImplementedError
114
+
115
+ def execute_model(
116
+ self,
117
+ execute_model_req: Optional[ExecuteModelRequest] = None
118
+ ) -> Optional[list[SamplerOutput]]:
119
+ raise NotImplementedError
120
+
121
+ def start_worker_execution_loop(self) -> None:
122
+ """Execute model loop in parallel worker.
123
+
124
+ You can stop the loop by executing a driver worker with an empty output.
125
+ See `stop_remote_worker_execution_loop` for more details.
126
+ """
127
+ with self.current_platform.inference_mode():
128
+ while True:
129
+ output = self.execute_model(execute_model_req=None)
130
+ if output is None:
131
+ return None
132
+
133
+ def determine_num_available_blocks(self) -> tuple[int, int]:
134
+ """Determine the number of available blocks for the GPU KV cache and
135
+ swappable CPU KV cache.
136
+
137
+ The implementation may run profiling or other heuristics to determine
138
+ the size of caches.
139
+
140
+ Returns a tuple[num_gpu_blocks, num_cpu_blocks], where num_gpu_blocks
141
+ are blocks that are "active" on the device and can be appended to.
142
+ num_cpu_blocks refers to "swapped" blocks in CPU memory and cannot be
143
+ appended to.
144
+ """
145
+ raise NotImplementedError
146
+
147
+ def get_cache_block_size_bytes(self) -> int:
148
+ """Return the size of a single cache block, in bytes. Used in
149
+ speculative decoding.
150
+ """
151
+ raise NotImplementedError
152
+
153
+ def add_lora(self, lora_request: LoRARequest) -> bool:
154
+ raise NotImplementedError
155
+
156
+ def remove_lora(self, lora_id: int) -> bool:
157
+ raise NotImplementedError
158
+
159
+ def pin_lora(self, lora_id: int) -> bool:
160
+ raise NotImplementedError
161
+
162
+ def list_loras(self) -> set[int]:
163
+ raise NotImplementedError
164
+
165
+ @property
166
+ def vocab_size(self) -> int:
167
+ """Get vocabulary size from model configuration."""
168
+ return self.model_config.get_vocab_size()
169
+
170
+ def shutdown(self) -> None:
171
+ """Clean up resources held by the worker."""
172
+ return
173
+
174
+
175
+class WorkerWrapperBase:
176
+ """
177
+ This class represents one process in an executor/engine. It is responsible
178
+ for lazily initializing the worker and handling the worker's lifecycle.
179
+ We first instantiate the WorkerWrapper, which remembers the worker module
180
+ and class name. Then, when we call `update_environment_variables`, and the
181
+ real initialization happens in `init_worker`.
182
+ """
183
+
184
+ def __init__(
185
+ self,
186
+ vllm_config: VllmConfig,
187
+ rpc_rank: int = 0,
188
+ ) -> None:
189
+ """
190
+ Initialize the worker wrapper with the given vllm_config and rpc_rank.
191
+ Note: rpc_rank is the rank of the worker in the executor. In most cases,
192
+ it is also the rank of the worker in the distributed group. However,
193
+ when multiple executors work together, they can be different.
194
+ e.g. in the case of SPMD-style offline inference with TP=2,
195
+ users can launch 2 engines/executors, each with only 1 worker.
196
+ All workers have rpc_rank=0, but they have different ranks in the TP
197
+ group.
198
+ """
199
+ self.rpc_rank = rpc_rank
200
+ self.worker: Optional[WorkerBase] = None
201
+ self.vllm_config: Optional[VllmConfig] = None
202
+ # do not store this `vllm_config`, `init_worker` will set the final
203
+ # one. TODO: investigate if we can remove this field in
204
+ # `WorkerWrapperBase`, `init_cached_hf_modules` should be
205
+ # unnecessary now.
206
+ if vllm_config.model_config is not None:
207
+ # it can be None in tests
208
+ trust_remote_code = vllm_config.model_config.trust_remote_code
209
+ if trust_remote_code:
210
+ # note: lazy import to avoid importing torch before initializing
211
+ from vllm.utils import init_cached_hf_modules
212
+ init_cached_hf_modules()
213
+
214
+ def shutdown(self) -> None:
215
+ if self.worker is not None:
216
+ self.worker.shutdown()
217
+
218
+ def adjust_rank(self, rank_mapping: dict[int, int]) -> None:
219
+ """
220
+ Adjust the rpc_rank based on the given mapping.
221
+ It is only used during the initialization of the executor,
222
+ to adjust the rpc_rank of workers after we create all workers.
223
+ """
224
+ if self.rpc_rank in rank_mapping:
225
+ self.rpc_rank = rank_mapping[self.rpc_rank]
226
+
227
+ def update_environment_variables(
228
+ self,
229
+ envs_list: list[dict[str, str]],
230
+ ) -> None:
231
+ envs = envs_list[self.rpc_rank]
232
+ key = 'CUDA_VISIBLE_DEVICES'
233
+ if key in envs and key in os.environ:
234
+ # overwriting CUDA_VISIBLE_DEVICES is desired behavior
235
+ # suppress the warning in `update_environment_variables`
236
+ del os.environ[key]
237
+ update_environment_variables(envs)
238
+
239
+ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None:
240
+ """
241
+ Here we inject some common logic before initializing the worker.
242
+ Arguments are passed to the worker class constructor.
243
+ """
244
+ kwargs = all_kwargs[self.rpc_rank]
245
+ self.vllm_config = kwargs.get("vllm_config")
246
+ assert self.vllm_config is not None, (
247
+ "vllm_config is required to initialize the worker")
248
+ enable_trace_function_call_for_thread(self.vllm_config)
249
+
250
+ from vllm.plugins import load_general_plugins
251
+ load_general_plugins()
252
+
253
+ if isinstance(self.vllm_config.parallel_config.worker_cls, str):
254
+ worker_class = resolve_obj_by_qualname(
255
+ self.vllm_config.parallel_config.worker_cls)
256
+ else:
257
+ raise ValueError(
258
+ "passing worker_cls is no longer supported. Please pass keep the class in a separate module and pass the qualified name of the class as a string." # noqa: E501
259
+ )
260
+ if self.vllm_config.parallel_config.worker_extension_cls:
261
+ worker_extension_cls = resolve_obj_by_qualname(
262
+ self.vllm_config.parallel_config.worker_extension_cls)
263
+ extended_calls = []
264
+ if worker_extension_cls not in worker_class.__bases__:
265
+ # check any conflicts between worker and worker_extension_cls
266
+ for attr in dir(worker_extension_cls):
267
+ if attr.startswith("__"):
268
+ continue
269
+ assert not hasattr(worker_class, attr), (
270
+ f"Worker class {worker_class} already has an attribute"
271
+ f" {attr}, which conflicts with the worker"
272
+ f" extension class {worker_extension_cls}.")
273
+ if callable(getattr(worker_extension_cls, attr)):
274
+ extended_calls.append(attr)
275
+ # dynamically inherit the worker extension class
276
+ worker_class.__bases__ = worker_class.__bases__ + (
277
+ worker_extension_cls, )
278
+ logger.info(
279
+ "Injected %s into %s for extended collective_rpc calls %s",
280
+ worker_extension_cls, worker_class, extended_calls)
281
+ with set_current_vllm_config(self.vllm_config):
282
+ # To make vLLM config available during worker initialization
283
+ self.worker = worker_class(**kwargs)
284
+ assert self.worker is not None
285
+
286
+ def initialize_from_config(self, kv_cache_configs: list[Any]) -> None:
287
+ kv_cache_config = kv_cache_configs[self.rpc_rank]
288
+ with set_current_vllm_config(self.vllm_config):
289
+ self.worker.initialize_from_config(kv_cache_config) # type: ignore
290
+
291
+ def init_device(self):
292
+ with set_current_vllm_config(self.vllm_config):
293
+ # To make vLLM config available during device initialization
294
+ self.worker.init_device() # type: ignore
295
+
296
+ def execute_method(self, method: Union[str, bytes], *args, **kwargs):
297
+ try:
298
+ # method resolution order:
299
+ # if a method is defined in this class, it will be called directly.
300
+ # otherwise, since we define `__getattr__` and redirect attribute
301
+ # query to `self.worker`, the method will be called on the worker.
302
+ return run_method(self, method, args, kwargs)
303
+ except Exception as e:
304
+ # if the driver worker also execute methods,
305
+ # exceptions in the rest worker may cause deadlock in rpc like ray
306
+ # see https://github.com/vllm-project/vllm/issues/3455
307
+ # print the error and inform the user to solve the error
308
+ msg = (f"Error executing method {method!r}. "
309
+ "This might cause deadlock in distributed execution.")
310
+ logger.exception(msg)
311
+ raise e
312
+
313
+ def __getattr__(self, attr):
314
+ return getattr(self.worker, attr)
vllm/worker/__init__.py
vllm/worker/worker_base.py
deleted
-279
@@ -1,279 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import os
5
-from typing import (Any, Callable, Dict, List, Optional, Set, Tuple, TypeVar,
6
- Union)
7
-
8
-import cloudpickle
9
-import torch.nn as nn
10
-
11
-from vllm.config import VllmConfig, set_current_vllm_config
12
-from vllm.logger import init_logger
13
-from vllm.lora.request import LoRARequest
14
-from vllm.sequence import ExecuteModelRequest
15
-from vllm.utils import (enable_trace_function_call_for_thread,
16
- resolve_obj_by_qualname, run_method,
17
- update_environment_variables,
18
- warn_for_unimplemented_methods)
19
-from vllm.v1.outputs import SamplerOutput
20
-
21
-logger = init_logger(__name__)
22
-
23
-_R = TypeVar("_R")
24
-
25
-
26
-@warn_for_unimplemented_methods
27
-class WorkerBase:
28
- """Worker interface that allows vLLM to cleanly separate implementations for
29
- different hardware. Also abstracts control plane communication, e.g., to
30
- communicate request metadata to other workers.
31
- """
32
-
33
- def __init__(
34
- self,
35
- vllm_config: VllmConfig,
36
- ) -> None:
37
- self.vllm_config = vllm_config
38
- self.model_config = vllm_config.model_config
39
- self.cache_config = vllm_config.cache_config
40
- self.lora_config = vllm_config.lora_config
41
- self.load_config = vllm_config.load_config
42
- self.parallel_config = vllm_config.parallel_config
43
- self.scheduler_config = vllm_config.scheduler_config
44
- self.device_config = vllm_config.device_config
45
- self.speculative_config = vllm_config.speculative_config
46
- self.observability_config = vllm_config.observability_config
47
- self.kv_transfer_config = vllm_config.kv_transfer_config
48
- self.compilation_config = vllm_config.compilation_config
49
- from vllm.platforms import current_platform
50
- self.current_platform = current_platform
51
-
52
- def init_device(self) -> None:
53
- """Initialize device state, such as loading the model or other on-device
54
- memory allocations.
55
- """
56
- raise NotImplementedError
57
-
58
- def initialize_cache(self, num_gpu_blocks: int,
59
- num_cpu_blocks: int) -> None:
60
- """Initialize the KV cache with the given size in blocks.
61
- """
62
- raise NotImplementedError
63
-
64
- def get_model(self) -> nn.Module:
65
- raise NotImplementedError
66
-
67
- def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R:
68
- """Apply a function on the model inside this worker."""
69
- return fn(self.get_model())
70
-
71
- def load_model(self) -> None:
72
- """Load model onto target device."""
73
- raise NotImplementedError
74
-
75
- def execute_model(
76
- self,
77
- execute_model_req: Optional[ExecuteModelRequest] = None
78
- ) -> Optional[List[SamplerOutput]]:
79
- raise NotImplementedError
80
-
81
- def start_worker_execution_loop(self) -> None:
82
- """Execute model loop in parallel worker.
83
-
84
- You can stop the loop by executing a driver worker with an empty output.
85
- See `stop_remote_worker_execution_loop` for more details.
86
- """
87
- with self.current_platform.inference_mode():
88
- while True:
89
- output = self.execute_model(execute_model_req=None)
90
- if output is None:
91
- return None
92
-
93
- def determine_num_available_blocks(self) -> Tuple[int, int]:
94
- """Determine the number of available blocks for the GPU KV cache and
95
- swappable CPU KV cache.
96
-
97
- The implementation may run profiling or other heuristics to determine
98
- the size of caches.
99
-
100
- Returns a Tuple[num_gpu_blocks, num_cpu_blocks], where num_gpu_blocks
101
- are blocks that are "active" on the device and can be appended to.
102
- num_cpu_blocks refers to "swapped" blocks in CPU memory and cannot be
103
- appended to.
104
- """
105
- raise NotImplementedError
106
-
107
- def get_cache_block_size_bytes(self) -> int:
108
- """Return the size of a single cache block, in bytes. Used in
109
- speculative decoding.
110
- """
111
- raise NotImplementedError
112
-
113
- def add_lora(self, lora_request: LoRARequest) -> bool:
114
- raise NotImplementedError
115
-
116
- def remove_lora(self, lora_id: int) -> bool:
117
- raise NotImplementedError
118
-
119
- def pin_lora(self, lora_id: int) -> bool:
120
- raise NotImplementedError
121
-
122
- def list_loras(self) -> Set[int]:
123
- raise NotImplementedError
124
-
125
- @property
126
- def vocab_size(self) -> int:
127
- """Get vocabulary size from model configuration."""
128
- return self.model_config.get_vocab_size()
129
-
130
- def shutdown(self) -> None:
131
- """Clean up resources held by the worker."""
132
- return
133
-
134
-
135
-class WorkerWrapperBase:
136
- """
137
- This class represents one process in an executor/engine. It is responsible
138
- for lazily initializing the worker and handling the worker's lifecycle.
139
- We first instantiate the WorkerWrapper, which remembers the worker module
140
- and class name. Then, when we call `update_environment_variables`, and the
141
- real initialization happens in `init_worker`.
142
- """
143
-
144
- def __init__(
145
- self,
146
- vllm_config: VllmConfig,
147
- rpc_rank: int = 0,
148
- ) -> None:
149
- """
150
- Initialize the worker wrapper with the given vllm_config and rpc_rank.
151
- Note: rpc_rank is the rank of the worker in the executor. In most cases,
152
- it is also the rank of the worker in the distributed group. However,
153
- when multiple executors work together, they can be different.
154
- e.g. in the case of SPMD-style offline inference with TP=2,
155
- users can launch 2 engines/executors, each with only 1 worker.
156
- All workers have rpc_rank=0, but they have different ranks in the TP
157
- group.
158
- """
159
- self.rpc_rank = rpc_rank
160
- self.worker: Optional[WorkerBase] = None
161
- self.vllm_config: Optional[VllmConfig] = None
162
- # do not store this `vllm_config`, `init_worker` will set the final
163
- # one. TODO: investigate if we can remove this field in
164
- # `WorkerWrapperBase`, `init_cached_hf_modules` should be
165
- # unnecessary now.
166
- if vllm_config.model_config is not None:
167
- # it can be None in tests
168
- trust_remote_code = vllm_config.model_config.trust_remote_code
169
- if trust_remote_code:
170
- # note: lazy import to avoid importing torch before initializing
171
- from vllm.utils import init_cached_hf_modules
172
- init_cached_hf_modules()
173
-
174
- def shutdown(self) -> None:
175
- if self.worker is not None:
176
- self.worker.shutdown()
177
-
178
- def adjust_rank(self, rank_mapping: Dict[int, int]) -> None:
179
- """
180
- Adjust the rpc_rank based on the given mapping.
181
- It is only used during the initialization of the executor,
182
- to adjust the rpc_rank of workers after we create all workers.
183
- """
184
- if self.rpc_rank in rank_mapping:
185
- self.rpc_rank = rank_mapping[self.rpc_rank]
186
-
187
- def update_environment_variables(self, envs_list: List[Dict[str,
188
- str]]) -> None:
189
- envs = envs_list[self.rpc_rank]
190
- key = 'CUDA_VISIBLE_DEVICES'
191
- if key in envs and key in os.environ:
192
- # overwriting CUDA_VISIBLE_DEVICES is desired behavior
193
- # suppress the warning in `update_environment_variables`
194
- del os.environ[key]
195
- update_environment_variables(envs)
196
-
197
- def init_worker(self, all_kwargs: List[Dict[str, Any]]) -> None:
198
- """
199
- Here we inject some common logic before initializing the worker.
200
- Arguments are passed to the worker class constructor.
201
- """
202
- kwargs = all_kwargs[self.rpc_rank]
203
- self.vllm_config = kwargs.get("vllm_config")
204
- assert self.vllm_config is not None, (
205
- "vllm_config is required to initialize the worker")
206
- enable_trace_function_call_for_thread(self.vllm_config)
207
-
208
- from vllm.plugins import load_general_plugins
209
- load_general_plugins()
210
-
211
- if isinstance(self.vllm_config.parallel_config.worker_cls, str):
212
- worker_class = resolve_obj_by_qualname(
213
- self.vllm_config.parallel_config.worker_cls)
214
- else:
215
- logger.warning(
216
- "passing worker_cls as a class object is strongly deprecated,"
217
- " as the serialization of class objects can be tricky and"
218
- " error-prone. To be safe, please keep the class in a separate"
219
- " module and pass the qualified name of the class as a string."
220
- )
221
- assert isinstance(self.vllm_config.parallel_config.worker_cls,
222
- bytes)
223
- worker_class = cloudpickle.loads(
224
- self.vllm_config.parallel_config.worker_cls)
225
- if self.vllm_config.parallel_config.worker_extension_cls:
226
- worker_extension_cls = resolve_obj_by_qualname(
227
- self.vllm_config.parallel_config.worker_extension_cls)
228
- extended_calls = []
229
- if worker_extension_cls not in worker_class.__bases__:
230
- # check any conflicts between worker and worker_extension_cls
231
- for attr in dir(worker_extension_cls):
232
- if attr.startswith("__"):
233
- continue
234
- assert not hasattr(worker_class, attr), (
235
- f"Worker class {worker_class} already has an attribute"
236
- f" {attr}, which conflicts with the worker"
237
- f" extension class {worker_extension_cls}.")
238
- if callable(getattr(worker_extension_cls, attr)):
239
- extended_calls.append(attr)
240
- # dynamically inherit the worker extension class
241
- worker_class.__bases__ = worker_class.__bases__ + (
242
- worker_extension_cls, )
243
- logger.info(
244
- "Injected %s into %s for extended collective_rpc calls %s",
245
- worker_extension_cls, worker_class, extended_calls)
246
- with set_current_vllm_config(self.vllm_config):
247
- # To make vLLM config available during worker initialization
248
- self.worker = worker_class(**kwargs)
249
- assert self.worker is not None
250
-
251
- def initialize_from_config(self, kv_cache_configs: List[Any]) -> None:
252
- kv_cache_config = kv_cache_configs[self.rpc_rank]
253
- with set_current_vllm_config(self.vllm_config):
254
- self.worker.initialize_from_config(kv_cache_config) # type: ignore
255
-
256
- def init_device(self):
257
- with set_current_vllm_config(self.vllm_config):
258
- # To make vLLM config available during device initialization
259
- self.worker.init_device() # type: ignore
260
-
261
- def execute_method(self, method: Union[str, bytes], *args, **kwargs):
262
- try:
263
- # method resolution order:
264
- # if a method is defined in this class, it will be called directly.
265
- # otherwise, since we define `__getattr__` and redirect attribute
266
- # query to `self.worker`, the method will be called on the worker.
267
- return run_method(self, method, args, kwargs)
268
- except Exception as e:
269
- # if the driver worker also execute methods,
270
- # exceptions in the rest worker may cause deadlock in rpc like ray
271
- # see https://github.com/vllm-project/vllm/issues/3455
272
- # print the error and inform the user to solve the error
273
- msg = (f"Error executing method {method!r}. "
274
- "This might cause deadlock in distributed execution.")
275
- logger.exception(msg)
276
- raise e
277
-
278
- def __getattr__(self, attr):
279
- return getattr(self.worker, attr)