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, 23:45 UTC
[BugFix] Set CUDA_VISIBLE_DEVICES before spawning the subprocesses (#21211)
Signed-off-by: Yinghai Lu <yinghai@thinkingmachines.ai> Signed-off-by: Nick Hill <nhill@redhat.com> Signed-off-by: Rui Qiao <ruisearch42@gmail.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: Rui Qiao <ruisearch42@gmail.com>
Yinghai Lu committed
Jul 23, 2025 at 21:44 UTC
11ef7a611ec015523301930a25422cf68216b5c4
2 files changed
+69
-26
vllm/v1/engine/core.py
+31
-20
@@ -910,22 +910,6 @@ class DPEngineCoreProc(EngineCoreProc):
910
logger.debug("Setting kv_transfer_config.engine_id to %s",
911
vllm_config.kv_transfer_config.engine_id)
912
913
- from vllm.platforms import current_platform
914
- device_control_env_var = current_platform.device_control_env_var
915
- world_size = vllm_config.parallel_config.world_size
916
- # Set CUDA_VISIBLE_DEVICES or equivalent.
917
- try:
918
- os.environ[device_control_env_var] = ",".join(
919
- str(current_platform.device_id_to_physical_device_id(i))
920
- for i in range(local_dp_rank *
921
- world_size, (local_dp_rank + 1) * world_size))
922
- except IndexError as e:
923
- raise Exception(
924
- f"Error setting {device_control_env_var}: "
925
- f"local range: [{local_dp_rank * world_size}, "
926
- f"{(local_dp_rank + 1) * world_size}) "
927
- f"base value: \"{os.getenv(device_control_env_var)}\"") from e
928
-
913
self.dp_rank = dp_rank
914
self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()
915
@@ -1088,14 +1072,41 @@ class DPEngineCoreActor(DPEngineCoreProc):
1072
vllm_config.parallel_config.data_parallel_rank_local = \
1073
local_dp_rank
1074
1091
- # Ray sets CUDA_VISIBLE_DEVICES to empty string,
1092
- # we clean this up to be able to properly initialize
1093
- # data parallel groups.
1094
- del os.environ['CUDA_VISIBLE_DEVICES']
1075
+ # Set CUDA_VISIBLE_DEVICES as early as possible in actor life cycle
1076
+ # NOTE: in MP we set CUDA_VISIBLE_DEVICES at process creation time,
1077
+ # and this cannot be done in the same way for Ray because:
1078
+ # 1) Ray manages life cycle of all ray workers (including
1079
+ # DPEngineCoreActor)
1080
+ # 2) Ray sets CUDA_VISIBLE_DEVICES based on num_gpus configuration
1081
+ # To bypass 2, we need to also set
1082
+ # RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES, but vLLM workers created
1083
+ # thereafter would have CUDA_VISIBLE_DEVICES set, which is sticky:
1084
+ # https://github.com/ray-project/ray/blob/e752fc319ddedd9779a0989b6d3613909bad75c9/python/ray/_private/worker.py#L456 # noqa: E501
1085
+ # But vLLM worker assumes visibility into all local GPUs, therefore
1086
+ # this results in incorrect indexing into the GPU ID list.
1087
+ self._set_cuda_visible_devices(vllm_config, local_dp_rank)
1088
1089
super().__init__(vllm_config, local_client, "", executor_class,
1090
log_stats)
1091
1092
+ def _set_cuda_visible_devices(self, vllm_config: VllmConfig,
1093
+ local_dp_rank: int):
1094
+ from vllm.platforms import current_platform
1095
+ device_control_env_var = current_platform.device_control_env_var
1096
+ world_size = vllm_config.parallel_config.world_size
1097
+ # Set CUDA_VISIBLE_DEVICES or equivalent.
1098
+ try:
1099
+ os.environ[device_control_env_var] = ",".join(
1100
+ str(current_platform.device_id_to_physical_device_id(i))
1101
+ for i in range(local_dp_rank *
1102
+ world_size, (local_dp_rank + 1) * world_size))
1103
+ except IndexError as e:
1104
+ raise Exception(
1105
+ f"Error setting {device_control_env_var}: "
1106
+ f"local range: [{local_dp_rank * world_size}, "
1107
+ f"{(local_dp_rank + 1) * world_size}) "
1108
+ f"base value: \"{os.getenv(device_control_env_var)}\"") from e
1109
+
1110
def _decorate_logs(self):
1111
pass
1112
vllm/v1/engine/utils.py
+38
-6
@@ -10,12 +10,14 @@ from enum import Enum, auto
10
from multiprocessing import Process, connection
11
from multiprocessing.process import BaseProcess
12
from typing import TYPE_CHECKING, Callable, Optional, Union
13
+from unittest.mock import patch
14
15
import msgspec
16
import zmq
17
18
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
19
from vllm.logger import init_logger
20
+from vllm.platforms import current_platform
21
from vllm.ray.ray_env import get_env_vars_to_copy
22
from vllm.utils import get_mp_context, get_open_zmq_ipc_path, zmq_socket_ctx
23
from vllm.v1.engine.coordinator import DPCoordinator
@@ -105,10 +107,13 @@ class CoreEngineProcManager:
107
"client_handshake_address"] = client_handshake_address
108
109
self.processes: list[BaseProcess] = []
110
+ local_dp_ranks = []
111
for index in range(local_engine_count):
112
local_index = local_start_index + index
113
global_index = start_index + index
114
+
115
# Start EngineCore in background process.
116
+ local_dp_ranks.append(local_index)
117
self.processes.append(
118
context.Process(target=target_fn,
119
name=f"EngineCore_{global_index}",
@@ -118,9 +123,14 @@ class CoreEngineProcManager:
123
}))
124
125
self._finalizer = weakref.finalize(self, shutdown, self.processes)
126
+
127
+ data_parallel = vllm_config.parallel_config.data_parallel_size > 1
128
try:
122
- for proc in self.processes:
123
- proc.start()
129
+ for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
130
+ with set_device_control_env_var(
131
+ vllm_config, local_dp_rank) if (
132
+ data_parallel) else contextlib.nullcontext():
133
+ proc.start()
134
finally:
135
# Kill other procs if not all are running.
136
if self.finished_procs():
@@ -145,6 +155,30 @@ class CoreEngineProcManager:
155
}
156
157
158
+@contextlib.contextmanager
159
+def set_device_control_env_var(vllm_config: VllmConfig,
160
+ local_dp_rank: int) -> Iterator[None]:
161
+ """
162
+ Temporarily set CUDA_VISIBLE_DEVICES or equivalent
163
+ for engine subprocess.
164
+ """
165
+ world_size = vllm_config.parallel_config.world_size
166
+ evar = current_platform.device_control_env_var
167
+ try:
168
+ value = ",".join(
169
+ str(current_platform.device_id_to_physical_device_id(i))
170
+ for i in range(local_dp_rank * world_size, (local_dp_rank + 1) *
171
+ world_size))
172
+ except IndexError as e:
173
+ raise Exception(f"Error setting {evar}: "
174
+ f"local range: [{local_dp_rank * world_size}, "
175
+ f"{(local_dp_rank + 1) * world_size}) "
176
+ "base value: "
177
+ f"\"{os.getenv(evar)}\"") from e
178
+ with patch.dict(os.environ, values=((evar, value), )):
179
+ yield
180
+
181
+
182
class CoreEngineActorManager:
183
"""
184
Utility class to handle creation, readiness, and shutdown
@@ -215,10 +249,9 @@ class CoreEngineActorManager:
249
250
self.placement_group_is_local = []
251
refs = []
218
- for index in range(dp_size):
219
- local_index = local_dp_ranks[index]
252
+ for index, local_index, pg in zip(range(dp_size), local_dp_ranks,
253
+ placement_groups):
254
dp_vllm_config = copy.deepcopy(vllm_config)
221
- pg = placement_groups[index]
255
dp_vllm_config.parallel_config.placement_group = pg
256
local_client = index < local_engine_count
257
actor = ray.remote(DPEngineCoreActor).options(
@@ -264,7 +297,6 @@ class CoreEngineActorManager:
297
local_engine_count = \
298
vllm_config.parallel_config.data_parallel_size_local
299
267
- nodes = list_nodes()
300
nodes = sorted(list_nodes(),
301
key=lambda node: node.node_ip != dp_master_ip)
302
assert nodes[0].node_ip == dp_master_ip, (