biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

replace cuda_device_count_stateless() to current_platform.device_count() (#37841)

Browse Source 
Signed-off-by: Liao, Wei <wei.liao@intel.com>
Signed-off-by: wliao2 <wei.liao@intel.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
wliao2
2026-03-31 07:32:54 -07:00
committed by GitHub
parent e8057c00bc
commit 4dfad17ed1
20 changed files with 96 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/utils.py
View File

@@ -58,7 +58,6 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.utils.mem_constants import GB_bytes
from vllm.utils.network_utils import get_open_port
from vllm.utils.torch_utils import (
cuda_device_count_stateless,
set_random_seed, # noqa: F401 - re-exported for use in test files
)
@@ -384,7 +383,7 @@ class RemoteVLLMServer:
elif current_platform.is_cuda():
with _nvml():
total_used = 0
device_count = cuda_device_count_stateless()
device_count = current_platform.device_count()
for i in range(device_count):
handle = nvmlDeviceGetHandleByIndex(i)
mem_info = nvmlDeviceGetMemoryInfo(handle)
@@ -1497,7 +1496,7 @@ def multi_gpu_marks(*, num_gpus: int):
"""Get a collection of pytest marks to apply for `@multi_gpu_test`."""
test_selector = pytest.mark.distributed(num_gpus=num_gpus)
test_skipif = pytest.mark.skipif(
cuda_device_count_stateless() < num_gpus,
current_platform.device_count() < num_gpus,
reason=f"Need at least {num_gpus} GPUs to run the test.",
)
@@ -1529,7 +1528,7 @@ def gpu_tier_mark(*, min_gpus: int = 1, max_gpus: int | None = None):
@gpu_tier_mark(max_gpus=1) # only on single-GPU
@gpu_tier_mark(min_gpus=2, max_gpus=4) # 2-4 GPUs only
"""
gpu_count = cuda_device_count_stateless()
gpu_count = current_platform.device_count()
marks = []
if min_gpus > 1: