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

[ROCm][AMD] unify CUDA_VISIBLE_DEVICES usage in cuda/rocm (#6352)

Browse Source
This commit is contained in:
Hongxia Yang
2024-07-12 00:30:46 -04:00
committed by GitHub
parent d26a8b3f1f
commit b6c16cf8ff
5 changed files with 10 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/distributed/test_utils.py
View File

@@ -1,7 +1,7 @@
import ray
import vllm.envs as envs
from vllm.utils import (cuda_device_count_stateless, is_hip,
from vllm.utils import (cuda_device_count_stateless,
update_environment_variables)
@@ -22,11 +22,6 @@ class _CUDADeviceCountStatelessTestActor:
def test_cuda_device_count_stateless():
"""Test that cuda_device_count_stateless changes return value if
CUDA_VISIBLE_DEVICES is changed."""
if is_hip():
# Set HIP_VISIBLE_DEVICES == CUDA_VISIBLE_DEVICES. Conversion
# is handled by `update_environment_variables`
update_environment_variables(
{"CUDA_VISIBLE_DEVICES": envs.CUDA_VISIBLE_DEVICES})
actor = _CUDADeviceCountStatelessTestActor.options( # type: ignore
num_gpus=2).remote()
assert sorted(ray.get(