Stop setting CUDA_VISIBLE_DEVICES internally in vLLM, add device_ids arg (#45026)

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Codex <codex@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: kourosh hakhamaneshi <kouroshHakha@users.noreply.github.com>
This commit is contained in:
Tyler Michael Smith
2026-06-20 13:38:10 -07:00
committed by GitHub
co-authored by Claude Codex mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> kourosh hakhamaneshi
parent e9de72fe6c
commit ebfbcfe46a
24 changed files with 722 additions and 185 deletions
@@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
using TileShapeD = typename MlaSm100Type::TileShapeD;
arguments.problem_shape =
cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
// Assumes device 0 when getting sm_count.
arguments.hw_info.sm_count =
sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
if (sm_count <= 0) {
int current_device = 0;
cudaGetDevice(&current_device);
arguments.hw_info.sm_count =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
} else {
arguments.hw_info.sm_count = sm_count;
}
arguments.split_kv = static_cast<int>(num_kv_splits);
MlaSm100Type::Fmha::set_split_kv(arguments);
+193
View File
@@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):
args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer")
assert args.model == "s3://bucket/model"
assert args.tokenizer == "s3://bucket/tokenizer"
class TestDeviceIds:
def test_device_ids_with_cvd_out_of_range(self, monkeypatch):
"""--device-ids index beyond the CVD set raises ValueError."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "4,5")
args = EngineArgs(model="m", device_ids=[0, 2])
with pytest.raises(ValueError, match="out of range"):
args._resolve_device_ids()
def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch):
"""--device-ids are CVD-local indices resolved to physical ids."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "4,5")
args = EngineArgs(model="m", device_ids=[0, 1])
assert args._resolve_device_ids() == [4, 5]
def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch):
"""--device-ids support UUID CVD values resolved by the platform."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890")
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
),
)
args = EngineArgs(model="m", device_ids=[0, 1])
assert args._resolve_device_ids() == [4, 5]
def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch):
"""UUID --device-ids are resolved to physical IDs immediately."""
from vllm.platforms import current_platform
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]),
)
args = EngineArgs(model="m", device_ids=["GPU-abcd1234"])
assert args._resolve_device_ids() == [4]
def test_device_ids_reject_mixed_integer_and_uuid_args(self):
"""--device-ids must not mix CVD indices and UUIDs."""
args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"])
with pytest.raises(ValueError, match="must not mix"):
args._resolve_device_ids()
def test_no_device_ids(self):
"""No --device-ids returns None."""
args = EngineArgs(model="m")
assert args._resolve_device_ids() is None
def test_cli_parsing(self):
"""--device-ids parses comma-separated string from CLI."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"])
assert parsed.device_ids == [0, 2, 4]
def test_cli_parsing_uuid(self):
"""--device-ids parses comma-separated UUID strings from CLI."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(
["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"]
)
assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"]
def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch):
"""assigned_physical_gpu_ids are already physical and not composed with CVD."""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5])
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
assert current_platform.device_id_to_physical_device_id(0) == 4
assert current_platform.device_id_to_physical_device_id(1) == 5
assert current_platform.logical_device_id_to_visible_device_id(0) == 0
assert current_platform.logical_device_id_to_visible_device_id(1) == 1
def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch):
"""Physical IDs map back to visible ordinals when CVD uses UUIDs."""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5])
monkeypatch.setenv(
current_platform.device_control_env_var,
"GPU-abcd1234,GPU-ef567890",
)
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
),
)
assert current_platform.logical_device_id_to_visible_device_id(0) == 1
def test_device_ids_reject_duplicates(self):
"""--device-ids must not contain duplicate entries."""
args = EngineArgs(model="m", device_ids=[2, 2])
with pytest.raises(ValueError, match="duplicates"):
args._resolve_device_ids()
def test_cli_parsing_strips_whitespace(self):
"""--device-ids tolerates whitespace around commas."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"])
assert parsed.device_ids == [0, 2, 4]
def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch):
"""visible_device_id_to_physical_device_id maps torch device ordinals,
independent of the logical-to-physical mapping.
Regression test: CustomAllreduce passes device.index (a visible
ordinal) and must not index into assigned_physical_gpu_ids, which
raised IndexError for non-identity --device-ids like [2, 3].
"""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3])
monkeypatch.delenv(current_platform.device_control_env_var, raising=False)
# CVD unset: visible ordinal == physical ID, even beyond the
# assigned list's length.
assert current_platform.visible_device_id_to_physical_device_id(2) == 2
assert current_platform.visible_device_id_to_physical_device_id(3) == 3
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
assert current_platform.visible_device_id_to_physical_device_id(1) == 5
with pytest.raises(IndexError, match="out of range"):
current_platform.visible_device_id_to_physical_device_id(2)
class TestDpDeviceIdSharding:
def test_dp_supervisor_device_ids_stay_env_relative(self):
"""Regression test: the DP supervisor must pass env-relative indices,
not physical IDs, because each child re-resolves --device-ids
against its inherited device-control env var."""
import argparse
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
args = argparse.Namespace(
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None
)
assert _build_device_ids(args, local_rank=0) == [0, 1]
assert _build_device_ids(args, local_rank=1) == [2, 3]
def test_dp_supervisor_shards_user_device_ids(self):
"""User-provided --device-ids are sharded across DP children."""
import argparse
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
args = argparse.Namespace(
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7]
)
assert _build_device_ids(args, local_rank=0) == [4, 5]
assert _build_device_ids(args, local_rank=1) == [6, 7]
with pytest.raises(ValueError, match="needs devices"):
_build_device_ids(args, local_rank=2)
def test_dp_rank_shards_user_assigned_gpu_ids(self):
"""get_physical_gpu_ids_for_local_dp_rank slices the user-provided
--device-ids list instead of recomputing from the env var."""
from vllm.platforms import current_platform
from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank
evar = current_platform.device_control_env_var
assert get_physical_gpu_ids_for_local_dp_rank(
evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
) == [6, 7]
with pytest.raises(ValueError, match="needs devices"):
get_physical_gpu_ids_for_local_dp_rank(
evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
)
@@ -364,7 +364,7 @@ class MockVLLMServer:
await self._serve_task
def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):
def launch_mock_vllm(child_args: argparse.Namespace):
logger.info("Launching mock vLLM on port %s", child_args.port)
mock_vllm = MockVLLMServer(
port=child_args.port,
@@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]
def launch_mock_vllm_with_drain(
child_args: argparse.Namespace, env_updates: dict[str, str]
child_args: argparse.Namespace,
):
logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
mock_vllm = MockVLLMServer(
+9
View File
@@ -302,6 +302,14 @@ class ParallelConfig:
Each entry must use `numactl --physcpubind` CPU-list syntax, for example
`"0-3"` or `"0,2,4-7"`.
"""
assigned_physical_gpu_ids: list[int] | None = None
"""Mapping from vLLM-local logical GPU IDs to physical GPU IDs.
For example, ``[2, 3]`` means logical GPU 0 maps to physical GPU 2,
and logical GPU 1 maps to physical GPU 3. Physical IDs are used only
at platform/topology boundaries such as NVML, NIC affinity, P2P
checks, and final CUDA device selection when needed. When None,
logical IDs map to visible device IDs in order."""
distributed_timeout_seconds: int | None = None
"""Timeout in seconds for distributed operations (e.g., init_process_group).
@@ -772,6 +780,7 @@ class ParallelConfig:
"numa_bind",
"numa_bind_nodes",
"numa_bind_cpus",
"assigned_physical_gpu_ids",
}
from vllm.config.utils import get_hash_factors, hash_factors
@@ -704,7 +704,14 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
self.num_experts = num_experts
self.cleanup()
gpus_per_node = torch.accelerator.device_count()
from vllm.platforms.interface import get_assigned_physical_gpu_ids
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
gpus_per_node = (
len(assigned_physical_gpu_ids)
if assigned_physical_gpu_ids is not None
else torch.accelerator.device_count()
)
logger.debug(
"Making One-sided NVLink mapping: rank=%d, world size=%d",
self.rank,
@@ -320,13 +320,21 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
is_distributed = dist.is_initialized()
num_dev = current_platform.device_count()
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
if cuda_visible_devices is None:
cuda_visible_devices = ",".join(str(i) for i in range(num_dev))
from vllm.platforms.interface import get_assigned_physical_gpu_ids
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
if assigned_physical_gpu_ids is not None:
# Key by the ordered list: the cache stores directed local-index
# pairs, so permutations of the same set are distinct mappings.
cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids)
num_dev = len(assigned_physical_gpu_ids)
else:
num_dev = current_platform.device_count()
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev))
path = os.path.join(
envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json"
envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json"
)
os.makedirs(os.path.dirname(path), exist_ok=True)
from vllm.distributed.parallel_state import get_world_group
@@ -338,7 +346,15 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
# enter this block to calculate the cache
logger.info("generating GPU P2P access cache in %s", path)
cache: dict[str, bool] = {}
ids = list(range(num_dev))
# The probe subprocesses inherit this process's device-control env
# var, so they must be given visible ordinals, not physical IDs.
if assigned_physical_gpu_ids is not None:
ids = [
current_platform.logical_device_id_to_visible_device_id(local)
for local in range(num_dev)
]
else:
ids = list(range(num_dev))
# batch of all pairs of GPUs
batch_src, batch_tgt = zip(*list(product(ids, ids)))
# NOTE: we use `subprocess` rather than `multiprocessing` here
@@ -368,8 +384,11 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
) from e
with open(output_file.name, "rb") as f:
result = pickle.load(f)
# Cache entries must be keyed by local indices (0..N-1) because
# gpu_p2p_access_check() is called with local ranks.
id_to_local = {device_id: local for local, device_id in enumerate(ids)}
for _i, _j, r in zip(batch_src, batch_tgt, result):
cache[f"{_i}->{_j}"] = r
cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r
with open(path, "w") as f:
json.dump(cache, f, indent=4)
if is_distributed:
@@ -34,7 +34,12 @@ def _can_p2p(rank: int, world_size: int) -> bool:
continue
if envs.VLLM_SKIP_P2P_CHECK:
logger.debug("Skipping P2P check and trusting the driver's P2P report.")
return torch.cuda.can_device_access_peer(rank, i)
# can_device_access_peer takes visible device ordinals, while
# rank and i are logical local IDs.
return torch.cuda.can_device_access_peer(
current_platform.logical_device_id_to_visible_device_id(rank),
current_platform.logical_device_id_to_visible_device_id(i),
)
if not gpu_p2p_access_check(rank, i):
return False
return True
@@ -126,13 +131,10 @@ class CustomAllreduce:
CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],
max_size,
)
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
if cuda_visible_devices:
device_ids = list(map(int, cuda_visible_devices.split(",")))
else:
device_ids = list(range(current_platform.device_count()))
physical_device_id = device_ids[device.index]
# device.index is a visible ordinal, not a logical local ID.
physical_device_id = current_platform.visible_device_id_to_physical_device_id(
device.index
)
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
gather_list = [
torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
@@ -129,12 +129,10 @@ class QuickAllReduce:
assert isinstance(device, torch.device)
self.device = device
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
if cuda_visible_devices:
device_ids = list(map(int, cuda_visible_devices.split(",")))
else:
device_ids = list(range(current_platform.device_count()))
physical_device_id = device_ids[device.index]
# device.index is a visible ordinal, not a logical local ID.
physical_device_id = current_platform.visible_device_id_to_physical_device_id(
device.index
)
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
gather_list = [
torch.tensor([0], dtype=torch.int, device="cpu")
@@ -840,7 +840,13 @@ class MessageQueue:
The MessageQueue instance for the calling process,
and a list of handles (only non-empty for the reader process).
"""
local_size = current_platform.device_count()
from vllm.platforms.interface import get_assigned_physical_gpu_ids
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
if assigned_physical_gpu_ids is not None:
local_size = len(assigned_physical_gpu_ids)
else:
local_size = current_platform.device_count()
rank = dist.get_rank()
same_node = rank // local_size == reader_rank // local_size
buffer_io = MessageQueue(
@@ -482,10 +482,11 @@ def _init_lmcache_engine(
)
# Change current device.
num_gpus = torch.accelerator.device_count()
local_rank = parallel_config.rank % num_gpus
torch.accelerator.set_device_index(local_rank)
device = torch.device(f"cuda:{local_rank}")
from vllm.distributed.parallel_state import get_world_group
device_index = get_world_group().device_index
torch.accelerator.set_device_index(device_index)
device = torch.device(f"cuda:{device_index}")
metadata = LMCacheEngineMetadata(
model_config.model,
parallel_config.world_size,
+24 -4
View File
@@ -392,6 +392,14 @@ class GroupCoordinator:
self.rank = torch.distributed.get_rank()
self.local_rank = local_rank
self.device_index: int
if _WORLD is not None:
self.device_index = _WORLD.device_index
else:
assert local_rank >= 0, (
"local_rank must be provided when creating the world group"
)
self.device_index = local_rank
self_device_group = None
self_cpu_group = None
@@ -442,11 +450,18 @@ class GroupCoordinator:
from vllm.platforms import current_platform
if current_platform.is_cuda_alike():
self.device = torch.device(f"cuda:{local_rank}")
visible_device_index = (
current_platform.logical_device_id_to_visible_device_id(
self.device_index
)
)
self.device = torch.device(f"cuda:{visible_device_index}")
elif current_platform.is_xpu():
self.device = torch.device(f"xpu:{local_rank}")
self.device = torch.device(f"xpu:{self.device_index}")
elif current_platform.is_out_of_tree():
self.device = torch.device(f"{current_platform.device_name}:{local_rank}")
self.device = torch.device(
f"{current_platform.device_name}:{self.device_index}"
)
else:
self.device = torch.device("cpu")
@@ -1438,7 +1453,12 @@ def _init_process_group_for_split_group(
"""
if torch.accelerator.is_available() and backend != "gloo":
init_backend = "cpu:gloo,cuda:nccl"
device_id: torch.device | None = torch.device(f"cuda:{local_rank}")
from vllm.platforms import current_platform
visible_device_index = current_platform.logical_device_id_to_visible_device_id(
local_rank
)
device_id: torch.device | None = torch.device(f"cuda:{visible_device_index}")
else:
init_backend = "gloo"
device_id = None
+19 -3
View File
@@ -86,6 +86,15 @@ class StatelessGroupCoordinator(GroupCoordinator):
self.rank = global_rank
self.local_rank = local_rank
from vllm.distributed.parallel_state import _WORLD
if _WORLD is not None:
self.device_index = _WORLD.device_index
else:
assert local_rank >= 0, (
"local_rank must be provided when creating the world group"
)
self.device_index = local_rank
self_device_group = None
self_cpu_group = None
@@ -152,11 +161,18 @@ class StatelessGroupCoordinator(GroupCoordinator):
self.tcp_store_group = self_tcp_store_group
if current_platform.is_cuda_alike():
self.device = torch.device(f"cuda:{local_rank}")
visible_device_index = (
current_platform.logical_device_id_to_visible_device_id(
self.device_index
)
)
self.device = torch.device(f"cuda:{visible_device_index}")
elif current_platform.is_xpu():
self.device = torch.device(f"xpu:{local_rank}")
self.device = torch.device(f"xpu:{self.device_index}")
elif current_platform.is_out_of_tree():
self.device = torch.device(f"{current_platform.device_name}:{local_rank}")
self.device = torch.device(
f"{current_platform.device_name}:{self.device_index}"
)
else:
self.device = torch.device("cpu")
+58
View File
@@ -6,6 +6,7 @@ import copy
import dataclasses
import functools
import json
import os
import sys
from collections.abc import Callable
from dataclasses import MISSING, asdict, dataclass, fields, is_dataclass
@@ -465,6 +466,7 @@ class EngineArgs:
numa_bind: bool = ParallelConfig.numa_bind
numa_bind_nodes: list[int] | None = ParallelConfig.numa_bind_nodes
numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus
device_ids: list[int | str] | None = None
tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
@@ -979,6 +981,20 @@ class EngineArgs:
parallel_group.add_argument(
"--numa-bind-cpus", **parallel_kwargs["numa_bind_cpus"]
)
parallel_group.add_argument(
"--device-ids",
type=lambda s: [
int(device_id) if device_id.isdigit() else device_id
for device_id in (part.strip() for part in s.split(","))
],
default=None,
help="Comma-separated physical GPU device IDs or UUIDs to use "
'(e.g. --device-ids "2,3,5,7"). Avoids setting '
"CUDA_VISIBLE_DEVICES, preserving full GPU topology "
"visibility for GPU-NIC affinity and DeepGEMM. "
"Note: has no effect with Ray executors; use Ray "
"placement groups for GPU selection instead.",
)
parallel_group.add_argument(
"--tensor-parallel-size", "-tp", **parallel_kwargs["tensor_parallel_size"]
)
@@ -1716,6 +1732,47 @@ class EngineArgs:
)
return SpeculativeConfig(**self.speculative_config)
def _resolve_device_ids(self) -> list[int] | None:
if not self.device_ids:
return None
if self.distributed_executor_backend == "ray":
logger.warning(
"--device-ids has no effect when using the Ray executor. "
"Use Ray placement groups for GPU selection instead."
)
ids = self.device_ids
if len(set(ids)) != len(ids):
raise ValueError(f"--device-ids must not contain duplicates: {ids}")
if all(isinstance(i, str) for i in ids):
return [
current_platform.device_control_id_to_physical_device_id(i)
for i in cast(list[str], ids)
]
if any(isinstance(i, str) for i in ids):
raise ValueError("--device-ids must not mix integer IDs and UUIDs")
int_ids = cast(list[int], ids)
# Compose with CUDA_VISIBLE_DEVICES: if CVD is set, treat
# --device-ids values as indices into the CVD-visible set.
cvd = getattr(
envs,
current_platform.device_control_env_var,
os.environ.get(current_platform.device_control_env_var),
)
if cvd:
cvd_ids = [
current_platform.device_control_id_to_physical_device_id(x)
for x in cvd.split(",")
]
for i in int_ids:
if i >= len(cvd_ids):
raise ValueError(
f"--device-ids index {i} is out of range for "
f"{current_platform.device_control_env_var}"
f"={cvd} ({len(cvd_ids)} devices visible)"
)
return [cvd_ids[i] for i in int_ids]
return int_ids
def create_diffusion_config(self) -> DiffusionConfig | None:
if self.diffusion_config is None:
return None
@@ -2029,6 +2086,7 @@ class EngineArgs:
cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
_api_process_count=self._api_process_count,
_api_process_rank=self._api_process_rank,
assigned_physical_gpu_ids=self._resolve_device_ids(),
numa_bind=self.numa_bind,
numa_bind_nodes=self.numa_bind_nodes,
numa_bind_cpus=self.numa_bind_cpus,
+19 -18
View File
@@ -23,12 +23,10 @@ import uvloop
from fastapi import FastAPI, Response
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.system_utils import (
decorate_logs,
kill_process_tree,
set_process_title,
update_environment_variables,
)
logger = init_logger(__name__)
@@ -127,22 +125,29 @@ def _build_vllm_dp_server_args(
child_args.data_parallel_multi_port_external_lb = False
child_args.data_parallel_supervisor_port = None
child_args.api_server_count = 1
child_args.device_ids = _build_device_ids(args, local_rank)
return child_args
def _build_vllm_dp_server_env(
args: argparse.Namespace, local_rank: int
) -> dict[str, str]:
# set visible devices for the child process
def _build_device_ids(args: argparse.Namespace, local_rank: int) -> list[int | str]:
"""Build the --device-ids value for a DP child process.
The child resolves these against its own inherited device-control env
var (e.g. CUDA_VISIBLE_DEVICES), so integer IDs must stay env-relative
here rather than being translated to physical IDs.
"""
devices_per_rank = args.tensor_parallel_size * args.pipeline_parallel_size
start = local_rank * devices_per_rank
stop = start + devices_per_rank
device_env = current_platform.device_control_env_var
visible_devices = ",".join(
str(current_platform.device_id_to_physical_device_id(idx))
for idx in range(start, stop)
)
return {device_env: visible_devices}
device_ids = getattr(args, "device_ids", None)
if device_ids is not None:
if stop > len(device_ids):
raise ValueError(
f"--device-ids has {len(device_ids)} entries, but DP rank "
f"{local_rank} needs devices [{start}, {stop})"
)
return device_ids[start:stop]
return list(range(start, stop))
def _child_base_url(args: argparse.Namespace, port: int) -> str:
@@ -228,9 +233,7 @@ def _build_dp_supervisor_app(supervisor: DPSupervisor) -> FastAPI:
return app
def _run_vllm_dp_server(
child_args: argparse.Namespace, env_updates: dict[str, str]
) -> None:
def _run_vllm_dp_server(child_args: argparse.Namespace) -> None:
"""
Entrypoint function for the vLLM DP Server.
"""
@@ -241,7 +244,6 @@ def _run_vllm_dp_server(
os.setpgrp()
name = f"APIServer_DP{child_args.data_parallel_rank}"
update_environment_variables(env_updates)
set_process_title(name)
decorate_logs(name)
uvloop.run(run_server(child_args))
@@ -345,11 +347,10 @@ class DPSupervisor:
context = multiprocessing.get_context("spawn")
for local_rank in range(self.args.data_parallel_size_local):
child_args = _build_vllm_dp_server_args(self.args, local_rank)
child_env = _build_vllm_dp_server_env(self.args, local_rank)
process = context.Process(
target=_run_vllm_dp_server,
name=f"APIServer_DPRank_{child_args.data_parallel_rank}",
args=(child_args, child_env),
args=(child_args,),
)
process.start()
self._processes.append(process)
+9
View File
@@ -685,6 +685,15 @@ class CudaPlatformBase(Platform):
# all the related functions work on real physical device ids.
# the major benefit of using NVML is that it will not initialize CUDA
class NvmlCudaPlatform(CudaPlatformBase):
@classmethod
@with_nvml_context
def device_control_id_to_physical_device_id(cls, device_id: str) -> int:
try:
return int(device_id)
except ValueError:
handle = pynvml.nvmlDeviceGetHandleByUUID(device_id)
return pynvml.nvmlDeviceGetIndex(handle)
@classmethod
@cache
@with_nvml_context
+102 -1
View File
@@ -30,6 +30,33 @@ else:
logger = init_logger(__name__)
_assigned_physical_gpu_ids: list[int] | None = None
def set_assigned_physical_gpu_ids(ids: list[int]) -> None:
"""Set the physical GPU IDs assigned to this worker process.
Called during worker init so that device_id_to_physical_device_id()
can map local_rank to the correct physical device without relying
on CUDA_VISIBLE_DEVICES.
Idempotent: a second call with the same value is a no-op.
Raises RuntimeError if called again with a different value.
This is expected to run during single-threaded worker initialization."""
global _assigned_physical_gpu_ids
if _assigned_physical_gpu_ids is not None:
if _assigned_physical_gpu_ids != ids:
raise RuntimeError(
f"set_assigned_physical_gpu_ids called with conflicting values: "
f"existing={_assigned_physical_gpu_ids}, new={ids}"
)
return
_assigned_physical_gpu_ids = ids
def get_assigned_physical_gpu_ids() -> list[int] | None:
return _assigned_physical_gpu_ids
@functools.cache
def in_wsl() -> bool:
@@ -233,8 +260,34 @@ class Platform:
"""
import vllm.kernels # noqa: F401
@classmethod
def device_control_id_to_physical_device_id(cls, device_id: str) -> int:
"""Map one device-control env entry to an integer physical device ID."""
try:
return int(device_id)
except ValueError as e:
raise ValueError(
f"Non-integer device ID {device_id!r} is not supported by "
f"{cls.device_name}."
) from e
@classmethod
def device_id_to_physical_device_id(cls, device_id: int):
"""Map a vLLM-local logical device ID to a physical device ID.
The input is a logical local ID (e.g. a local rank), NOT a visible
device ordinal; for the latter use
visible_device_id_to_physical_device_id(). The two coincide only
when no logical-to-physical mapping is in effect.
"""
if _assigned_physical_gpu_ids is not None:
if device_id >= len(_assigned_physical_gpu_ids):
raise IndexError(
f"device_id {device_id} is out of range for "
f"assigned_physical_gpu_ids {_assigned_physical_gpu_ids} "
f"({len(_assigned_physical_gpu_ids)} devices assigned)"
)
return _assigned_physical_gpu_ids[device_id]
# Treat empty device control env var as unset. This is a valid
# configuration in Ray setups where the engine is launched in
# a CPU-only placement group located on a GPU node.
@@ -244,10 +297,58 @@ class Platform:
):
device_ids = os.environ[cls.device_control_env_var].split(",")
physical_device_id = device_ids[device_id]
return int(physical_device_id)
return cls.device_control_id_to_physical_device_id(physical_device_id)
else:
return device_id
@classmethod
def logical_device_id_to_visible_device_id(cls, device_id: int) -> int:
"""Map a vLLM-local logical device ID to the current process's
visible accelerator ordinal.
vLLM internals use logical local IDs. Physical IDs are used only
at platform/topology boundaries. This helper performs the final
translation needed by APIs such as ``torch.device("cuda:N")``.
"""
physical_device_id = cls.device_id_to_physical_device_id(device_id)
device_control_env = os.environ.get(cls.device_control_env_var, "")
if not device_control_env:
return physical_device_id
visible_physical_device_ids = [
cls.device_control_id_to_physical_device_id(physical_id)
for physical_id in device_control_env.split(",")
]
if physical_device_id not in visible_physical_device_ids:
raise RuntimeError(
f"Physical device {physical_device_id} for logical device "
f"{device_id} is not visible in {cls.device_control_env_var}="
f"{device_control_env}"
)
return visible_physical_device_ids.index(physical_device_id)
@classmethod
def visible_device_id_to_physical_device_id(cls, device_id: int) -> int:
"""Map a visible accelerator ordinal (e.g. ``torch.device.index``)
to a physical device ID.
This is the inverse of the env-var translation performed by
logical_device_id_to_visible_device_id() and is independent of any
logical-to-physical mapping set via set_assigned_physical_gpu_ids().
"""
device_control_env = os.environ.get(cls.device_control_env_var, "")
if not device_control_env:
return device_id
visible_device_ids = device_control_env.split(",")
if device_id >= len(visible_device_ids):
raise IndexError(
f"visible device ordinal {device_id} is out of range for "
f"{cls.device_control_env_var}={device_control_env}"
)
return cls.device_control_id_to_physical_device_id(
visible_device_ids[device_id]
)
@classmethod
def import_kernels(cls) -> None:
"""Import any platform-specific C kernels."""
+16 -9
View File
@@ -74,7 +74,7 @@ from vllm.v1.engine.utils import (
EngineHandshakeMetadata,
EngineZmqAddresses,
SignalCallback,
get_device_indices,
get_physical_gpu_ids_for_local_dp_rank,
)
from vllm.v1.executor import Executor
from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind
@@ -2175,23 +2175,30 @@ class EngineCoreActorMixin:
pass
else:
device_control_env_var = current_platform.device_control_env_var
self._set_cuda_visible_devices(
self._set_assigned_physical_gpu_ids(
vllm_config, local_dp_rank, device_control_env_var
)
def _set_cuda_visible_devices(
self, vllm_config: VllmConfig, local_dp_rank: int, device_control_env_var: str
def _set_assigned_physical_gpu_ids(
self,
vllm_config: VllmConfig,
local_dp_rank: int,
device_control_env_var: str,
):
world_size = vllm_config.parallel_config.world_size
# Set CUDA_VISIBLE_DEVICES or equivalent.
try:
value = get_device_indices(
device_control_env_var, local_dp_rank, world_size
physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(
device_control_env_var,
local_dp_rank,
world_size,
user_assigned_gpu_ids=(
vllm_config.parallel_config.assigned_physical_gpu_ids
),
)
os.environ[device_control_env_var] = value
vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids
except IndexError as e:
raise Exception(
f"Error setting {device_control_env_var}: "
f"Error computing assigned_physical_gpu_ids: "
f"local range: [{local_dp_rank * world_size}, "
f"{(local_dp_rank + 1) * world_size}) "
f'base value: "{os.getenv(device_control_env_var)}"'
+66 -43
View File
@@ -12,7 +12,6 @@ from multiprocessing import Process, connection
from multiprocessing.process import BaseProcess
from multiprocessing.queues import Queue
from typing import TYPE_CHECKING, cast
from unittest.mock import patch
import msgspec
import zmq
@@ -175,38 +174,38 @@ class CoreEngineProcManager:
self.manager_stopped = threading.Event()
self.failed_proc_name: str | None = None
# All ranks share this config object: capture the user-provided
# --device-ids list before the per-rank shard overwrites it. Mutating
# the config before each proc.start() works because the spawn method
# pickles process args at start() time, sequentially per rank.
user_assigned_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids
try:
for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
# Adjust device control in DP for platforms that cannot rely
# on torch.accelerator.set_device_index(), and for Ray launchers.
device_control_context: contextlib.AbstractContextManager[None] = (
contextlib.nullcontext()
)
# Populate the logical-to-physical GPU mapping in DP for
# platforms that cannot rely on
# torch.accelerator.set_device_index(), and for Ray.
needs_device_env_isolation = not (
current_platform.is_cuda_alike() or current_platform.is_xpu()
)
if is_dp and (
needs_device_env_isolation or vllm_config.parallel_config.use_ray
):
device_control_context = set_device_control_env_var(
vllm_config, local_dp_rank
set_assigned_physical_gpu_ids_for_dp_rank(
vllm_config, local_dp_rank, user_assigned_gpu_ids
)
with (
device_control_context,
numa_utils.configure_subprocess(
# EngineCore itself does not have a TP/PP-local rank.
# When DP is enabled, set_device_control_env_var()
# narrows visible devices to this DP shard first, so
# local_rank=0 means "the first local GPU in this
# shard". The actual TP/PP worker processes spawned by
# the executor are bound separately with their own
# local_rank values.
vllm_config,
local_rank=0,
dp_local_rank=local_dp_rank,
process_kind="EngineCore",
),
with numa_utils.configure_subprocess(
# EngineCore itself does not have a TP/PP-local rank.
# When DP is enabled, set_assigned_physical_gpu_ids_for_dp_rank()
# populates the logical-to-physical mapping for this DP
# shard, so local_rank=0 means "the first local GPU in
# this shard". The actual TP/PP worker processes spawned
# by the executor are bound separately with their own
# local_rank values.
vllm_config,
local_rank=0,
dp_local_rank=local_dp_rank,
process_kind="EngineCore",
):
proc.start()
finally:
@@ -281,55 +280,79 @@ class SignalCallback:
self._event.set()
@contextlib.contextmanager
def set_device_control_env_var(
vllm_config: VllmConfig, local_dp_rank: int
) -> Iterator[None]:
def set_assigned_physical_gpu_ids_for_dp_rank(
vllm_config: VllmConfig,
local_dp_rank: int,
user_assigned_gpu_ids: list[int] | None = None,
) -> None:
"""
Temporarily set CUDA_VISIBLE_DEVICES or equivalent
for engine subprocess.
Populate assigned_physical_gpu_ids on the config for the given DP rank.
user_assigned_gpu_ids is the full (un-sharded) --device-ids list, if the
user provided one; this DP rank's shard is sliced from it. It is passed
explicitly rather than read from the config because callers may reuse
one config object across DP ranks, overwriting the field each time.
"""
world_size = vllm_config.parallel_config.world_size
local_world_size = vllm_config.parallel_config.local_world_size
evar = current_platform.device_control_env_var
value = get_device_indices(evar, local_dp_rank, world_size, local_world_size)
with patch.dict(os.environ, values=((evar, value),)):
yield
physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(
evar,
local_dp_rank,
world_size,
local_world_size,
user_assigned_gpu_ids=user_assigned_gpu_ids,
)
vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids
def get_device_indices(
def get_physical_gpu_ids_for_local_dp_rank(
device_control_env_var: str,
local_dp_rank: int,
world_size: int,
local_world_size: int | None = None,
):
user_assigned_gpu_ids: list[int] | None = None,
) -> list[int]:
"""
Returns a comma-separated string of device indices for the specified
Returns list of physical GPU IDs for the specified
data parallel rank.
For example, if world_size=2 and local_dp_rank=1, and there are 4 devices,
this will select devices 2 and 3 for local_dp_rank=1.
this will return [2, 3] for local_dp_rank=1.
If user_assigned_gpu_ids is provided (e.g. from --device-ids), this DP
rank's shard is sliced from it instead of being derived from the
device-control env var.
"""
if local_world_size is None:
local_world_size = world_size
if user_assigned_gpu_ids is not None:
start = local_dp_rank * world_size
stop = start + local_world_size
if stop > len(user_assigned_gpu_ids):
raise ValueError(
f"--device-ids provides {len(user_assigned_gpu_ids)} devices, "
f"but DP rank {local_dp_rank} needs devices [{start}, {stop})"
)
return user_assigned_gpu_ids[start:stop]
try:
value = ",".join(
str(current_platform.device_id_to_physical_device_id(i))
return [
current_platform.device_id_to_physical_device_id(i)
for i in range(
local_dp_rank * world_size,
local_dp_rank * world_size + local_world_size,
)
)
]
except IndexError as e:
raise Exception(
f"Error setting {device_control_env_var}: "
f"Error computing device indices for "
f"{device_control_env_var}: "
f"local range: [{local_dp_rank * world_size}, "
f"{(local_dp_rank + 1) * world_size}) "
"base value: "
f'"{os.getenv(device_control_env_var)}"'
) from e
return value
def _apply_dp_identity_suffix(dp_vllm_config, dp_rank: int) -> None:
@@ -453,11 +476,11 @@ class CoreEngineActorManager:
# https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/intel_gpu.py#L56 # noqa: E501
if current_platform.is_xpu():
device_evar = current_platform.device_control_env_var
device_indices = get_device_indices(
physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(
device_evar, local_index, world_size
)
actor_env_vars = self.env_vars_dict.copy()
actor_env_vars[device_evar] = device_indices
actor_env_vars[device_evar] = ",".join(str(d) for d in physical_gpu_ids)
runtime_env = RuntimeEnv(env_vars=actor_env_vars)
actor = (
+10
View File
@@ -826,6 +826,16 @@ class WorkerProc:
signal.signal(signal.SIGTERM, signal_handler)
signal.signal(signal.SIGINT, signal_handler)
# Publish the logical-to-physical mapping early so topology helpers
# work before init_device (needed by set_worker_net_device below).
assigned_physical_gpu_ids = kwargs[
"vllm_config"
].parallel_config.assigned_physical_gpu_ids
if assigned_physical_gpu_ids is not None:
from vllm.platforms.interface import set_assigned_physical_gpu_ids
set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
# Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set
set_worker_net_device(kwargs.get("local_rank", 0), kwargs["vllm_config"])
+22 -31
View File
@@ -258,30 +258,35 @@ class RayDistributedExecutor(Executor):
}
self.collective_rpc("adjust_rank", args=(rerank_mapping,))
# Get the set of GPU IDs used on each node.
worker_node_and_gpu_ids = []
# Get the set of physical GPU IDs used on each node.
worker_node_and_physical_gpu_ids = []
for worker in [self.driver_dummy_worker] + self.workers:
if worker is None:
# driver_dummy_worker can be None when using ray spmd worker.
continue
worker_node_and_gpu_ids.append(
ray.get(worker.get_node_and_gpu_ids.remote()) # type: ignore[attr-defined]
worker_node_and_physical_gpu_ids.append(
ray.get(worker.get_node_and_physical_gpu_ids.remote()) # type: ignore[attr-defined]
)
node_workers = defaultdict(list) # node id -> list of worker ranks
node_gpus = defaultdict(list) # node id -> list of gpu ids
node_physical_gpu_ids = defaultdict(list) # node id -> physical GPU IDs
for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):
for i, (node_id, physical_gpu_ids) in enumerate(
worker_node_and_physical_gpu_ids
):
node_workers[node_id].append(i)
# `gpu_ids` can be a list of strings or integers.
# `physical_gpu_ids` can be a list of strings or integers.
# convert them to integers for consistency.
# NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs),
# NOTE: physical GPU IDs can be larger than 9 (e.g. 16 GPUs),
# string sorting is not sufficient.
# see https://github.com/vllm-project/vllm/issues/5590
gpu_ids = [int(x) for x in gpu_ids]
node_gpus[node_id].extend(gpu_ids)
for node_id, gpu_ids in node_gpus.items():
node_gpus[node_id] = sorted(gpu_ids)
physical_gpu_ids = [
current_platform.device_control_id_to_physical_device_id(str(x))
for x in physical_gpu_ids
]
node_physical_gpu_ids[node_id].extend(physical_gpu_ids)
for node_id, physical_gpu_ids in node_physical_gpu_ids.items():
node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids)
all_ips = set(worker_ips + [driver_ip])
n_ips = len(all_ips)
@@ -297,23 +302,8 @@ class RayDistributedExecutor(Executor):
" each node."
)
# Set environment variables for the driver and workers.
# We set CUDA_VISIBLE_DEVICES to ALL GPUs on the node for each worker.
# This is needed because:
# 1. Ray's compiled DAG needs to find the allocated GPU in
# CUDA_VISIBLE_DEVICES.
# 2. vLLM's communication layer (NCCL, CustomAllreduce) needs to see
# all GPUs for P2P checks and communication setup. Though if it was
# just this reason, we could have also just kept the visible devices
# unset.
# Each worker will use local_rank to index into the visible devices.
all_args_to_update_environment_variables = [
{
current_platform.device_control_env_var: ",".join(
map(str, node_gpus[node_id])
),
}
for (node_id, _) in worker_node_and_gpu_ids
all_args_to_update_environment_variables: list[dict[str, str]] = [
{} for _ in worker_node_and_physical_gpu_ids
]
# Environment variables to copy from driver to workers
@@ -336,7 +326,7 @@ class RayDistributedExecutor(Executor):
"update_environment_variables", args=(self._get_env_vars_to_be_updated(),)
)
if len(node_gpus) == 1:
if len(node_physical_gpu_ids) == 1:
# in single node case, we don't need to get the IP address.
# the loopback address is sufficient
# NOTE: a node may have several IP addresses, one for each
@@ -352,10 +342,11 @@ class RayDistributedExecutor(Executor):
# Initialize the actual workers inside worker wrapper.
all_kwargs = []
for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids):
for rank, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids):
local_rank = node_workers[node_id].index(rank)
kwargs = dict(
vllm_config=self.vllm_config,
assigned_physical_gpu_ids=sorted(node_physical_gpu_ids[node_id]),
local_rank=local_rank,
rank=rank,
distributed_init_method=distributed_init_method,
+56 -31
View File
@@ -79,24 +79,25 @@ class RayWorkerProc(WorkerProc):
1. __init__: lightweight setup, stores init args (no device/model init)
2. initialize_worker: called after GPU IDs are discovered, completes
the full WorkerProc initialization with the correct local_rank and
CUDA_VISIBLE_DEVICES.
logical-to-physical GPU mapping.
CUDA_VISIBLE_DEVICES setup flow:
GPU assignment flow:
1. RayExecutorV2 enables RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES so Ray does
not set CUDA_VISIBLE_DEVICES on RayWorkerProc actors at creation time.
2. Each actor is scheduled with a placement group and bundle index; Ray resolves
the physical GPU ID for that bundle at placement time.
3. After placement, the worker discovers that GPU ID and sets
CUDA_VISIBLE_DEVICES before finishing WorkerProc initialization.
3. After placement, the executor discovers each worker's GPU ID and passes the
node's logical-to-physical mapping (assigned_physical_gpu_ids) to
initialize_worker(); CUDA_VISIBLE_DEVICES is never modified.
There is no workaround for this unset-and-reset sequence when the placement group
is externally managed: scheduling must complete before CUDA_VISIBLE_DEVICES can
match the GPU tied to the worker's bundle.
Scheduling must complete before the mapping is known when the placement
group is externally managed: only then is the GPU tied to the worker's
bundle resolved.
This sequence allows multiple vLLM instances to coexist on the same node:
each instance is unaware which physical devices others hold, and the
externally managed placement group avoids CUDA_VISIBLE_DEVICES conflicts
externally managed placement group avoids device assignment conflicts
by binding workers to specific placement group bundles.
"""
@@ -120,28 +121,33 @@ class RayWorkerProc(WorkerProc):
is_driver_worker=is_driver_worker,
)
def get_node_and_gpu_ids(self) -> tuple[str, list[int]]:
"""Return (node_id, gpu_ids) assigned to this actor by Ray."""
def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]:
"""Return (node_id, physical_gpu_ids) assigned to this actor by Ray."""
node_id = ray.get_runtime_context().get_node_id()
device_key = current_platform.ray_device_key
if not device_key:
raise RuntimeError(
f"current platform {current_platform.device_name} does not support ray."
)
gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]
return node_id, [int(x) for x in gpu_ids]
physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]
return node_id, [
current_platform.device_control_id_to_physical_device_id(str(x))
for x in physical_gpu_ids
]
def initialize_worker(
self,
local_rank: int,
env_vars: dict[str, str],
driver_env_vars: dict[str, str] | None = None,
assigned_physical_gpu_ids: list[int] | None = None,
) -> None:
"""Complete initialization after GPU assignment is known.
*driver_env_vars* are applied with ``setdefault`` they fill
in missing vars but never overwrite node-local values.
*env_vars* (e.g. CUDA_VISIBLE_DEVICES) always overwrite.
*env_vars* always overwrite.
*assigned_physical_gpu_ids* maps local_rank to physical CUDA device ID.
"""
if driver_env_vars:
for key, value in driver_env_vars.items():
@@ -149,6 +155,13 @@ class RayWorkerProc(WorkerProc):
for key, value in env_vars.items():
os.environ[key] = value
if assigned_physical_gpu_ids is not None:
vllm_config = self._init_kwargs["vllm_config"]
assert isinstance(vllm_config, VllmConfig)
vllm_config.parallel_config.assigned_physical_gpu_ids = (
assigned_physical_gpu_ids
)
self.local_rank = local_rank
super().__init__(
local_rank=local_rank,
@@ -365,36 +378,48 @@ class RayExecutorV2(MultiprocExecutor):
)
self.ray_worker_handles.append(handle)
# Step 6: Discover GPU IDs assigned to each worker via Ray runtime context.
worker_node_and_gpu_ids = ray.get(
[h.actor.get_node_and_gpu_ids.remote() for h in self.ray_worker_handles]
# Step 6: Discover physical GPU IDs assigned to each worker via Ray
# runtime context.
worker_node_and_physical_gpu_ids = ray.get(
[
h.actor.get_node_and_physical_gpu_ids.remote()
for h in self.ray_worker_handles
]
)
node_workers: dict[str, list[int]] = defaultdict(list)
node_gpus: dict[str, list[int]] = defaultdict(list)
for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):
node_physical_gpu_ids: dict[str, list[int]] = defaultdict(list)
for i, (node_id, physical_gpu_ids) in enumerate(
worker_node_and_physical_gpu_ids
):
node_workers[node_id].append(i)
node_gpus[node_id].extend(gpu_ids)
for node_id, gpu_ids in node_gpus.items():
node_gpus[node_id] = sorted(gpu_ids)
node_physical_gpu_ids[node_id].extend(physical_gpu_ids)
for node_id, physical_gpu_ids in node_physical_gpu_ids.items():
node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids)
# Step 7: Initialize workers with correct local_rank and
# CUDA_VISIBLE_DEVICES. Each worker sees all GPUs assigned to
# this executor on its node; local_rank indexes into that set.
# Step 7: Initialize workers with local logical ranks and the
# logical-to-physical GPU mapping discovered from Ray placement.
init_worker_refs = []
for i, (node_id, _) in enumerate(worker_node_and_gpu_ids):
for i, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids):
local_rank = node_workers[node_id].index(i)
worker_env_vars = {
current_platform.device_control_env_var: ",".join(
map(str, node_gpus[node_id])
),
}
assigned_physical_gpu_ids = sorted(node_physical_gpu_ids[node_id])
worker_env_vars: dict[str, str] = {}
self.ray_worker_handles[i].local_rank = local_rank
init_worker_refs.append(
self.ray_worker_handles[i].actor.initialize_worker.remote(
local_rank, worker_env_vars, self.driver_env_vars
local_rank,
worker_env_vars,
self.driver_env_vars,
assigned_physical_gpu_ids=assigned_physical_gpu_ids,
)
)
# Also set on the executor-side config for consistency. The mapping
# is per-node, so only do this when all workers share one node.
if len(node_physical_gpu_ids) == 1:
node_id_0 = worker_node_and_physical_gpu_ids[0][0]
self.vllm_config.parallel_config.assigned_physical_gpu_ids = sorted(
node_physical_gpu_ids[node_id_0]
)
ray.get(init_worker_refs)
# Step 8: Collect response MQ handles
+5 -3
View File
@@ -93,7 +93,7 @@ try:
def get_node_ip(self) -> str:
return get_ip()
def get_node_and_gpu_ids(self) -> tuple[str, list[int]]:
def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]:
node_id = ray.get_runtime_context().get_node_id()
device_key = vllm.platforms.current_platform.ray_device_key
if not device_key:
@@ -101,8 +101,10 @@ try:
"current platform %s does not support ray.",
vllm.platforms.current_platform.device_name,
)
gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]
return node_id, gpu_ids
physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[
device_key
]
return node_id, physical_gpu_ids
def setup_device_if_necessary(self):
# TODO(swang): This is needed right now because Ray CG executes
+38 -10
View File
@@ -270,19 +270,47 @@ class Worker(WorkerBase):
# DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
self.local_rank += dp_local_rank * tp_pp_world_size
# Publish the logical-to-physical mapping for topology queries
# such as NIC affinity and P2P checks.
assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids
if assigned_physical_gpu_ids is not None:
from vllm.platforms.interface import set_assigned_physical_gpu_ids
set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
assert self.local_rank < len(assigned_physical_gpu_ids), (
f"local_rank {self.local_rank} is out of bounds for "
f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}"
)
# NOTE(patch pr45026): local_world_size is derived from
# parallel_config.nnodes, which is only set for the "mp"
# multi-node backend. With the "ray"/"external_launcher"
# backends nnodes stays 1, so local_world_size collapses to
# the full world_size and this check wrongly fires on
# cross-node deployments. assigned_physical_gpu_ids is already
# per-node and the local_rank bound above fully validates the
# mapping for these backends, so skip the check for them.
if parallel_config.distributed_executor_backend not in (
"ray",
"external_launcher",
):
assert self.parallel_config.local_world_size <= len(
assigned_physical_gpu_ids
), (
f"local_world_size ({self.parallel_config.local_world_size})"
" exceeds assigned_physical_gpu_ids count "
f"({len(assigned_physical_gpu_ids)})"
)
else:
assert self.local_rank < torch.accelerator.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
visible_device_count = (
torch.accelerator.device_count() if torch.cuda.is_available() else 0
)
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must "
f"be less than or equal to the number of visible devices "
f"({visible_device_count})."
f"DP adjusted local rank {self.local_rank} is out of "
f"bounds for {torch.accelerator.device_count()} devices."
)
self.device = torch.device(f"cuda:{self.local_rank}")
visible_device_index = (
current_platform.logical_device_id_to_visible_device_id(self.local_rank)
)
self.device = torch.device(f"cuda:{visible_device_index}")
torch.accelerator.set_device_index(self.device)
current_platform.check_if_supports_dtype(self.model_config.dtype)
+6
View File
@@ -286,6 +286,12 @@ class WorkerWrapperBase:
extended_calls,
)
assigned_physical_gpu_ids = kwargs.pop("assigned_physical_gpu_ids", None)
if assigned_physical_gpu_ids is not None:
vllm_config.parallel_config.assigned_physical_gpu_ids = (
assigned_physical_gpu_ids
)
shared_worker_lock = kwargs.pop("shared_worker_lock", None)
if shared_worker_lock is None:
msg = (