[Hardware] Support platforms and plugin system (#774)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
This commit is contained in:
1 parent
b11d436e19
commit
c9392dff01
103 files changed
+1770
-1218
No files matched your search
@@ -1,5 +1,6 @@
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -42,6 +43,14 @@ def _make_dummy_torch(call_log):
|
||||
return _Torch
|
||||
|
||||
|
||||
def _make_mock_platform(device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"):
|
||||
"""Create a mock platform for testing."""
|
||||
mock_platform = MagicMock()
|
||||
mock_platform.device_type = device_type
|
||||
mock_platform.device_control_env_var = env_var
|
||||
return mock_platform
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
|
||||
def test_set_stage_devices_respects_logical_ids(monkeypatch):
|
||||
# Preserve an existing logical mapping and ensure devices "0,1" map through it.
|
||||
@@ -49,9 +58,51 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
|
||||
call_log: list[int] = []
|
||||
dummy_torch = _make_dummy_torch(call_log)
|
||||
monkeypatch.setitem(sys.modules, "torch", dummy_torch)
|
||||
monkeypatch.setattr("vllm_omni.utils.detect_device_type", lambda: "cuda")
|
||||
monkeypatch.setattr("vllm_omni.utils.get_device_control_env_var", lambda: "CUDA_VISIBLE_DEVICES")
|
||||
|
||||
# Mock the platform at the source module where it's defined
|
||||
mock_platform = _make_mock_platform(device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.platforms.current_omni_platform",
|
||||
mock_platform,
|
||||
)
|
||||
|
||||
set_stage_devices(stage_id=0, devices="0,1")
|
||||
|
||||
assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
|
||||
def test_set_stage_devices_npu_platform(monkeypatch):
|
||||
"""Test that set_stage_devices works correctly for NPU platform."""
|
||||
monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,5")
|
||||
call_log: list[int] = []
|
||||
|
||||
# Create NPU mock torch
|
||||
class _Npu:
|
||||
@staticmethod
|
||||
def is_available():
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def set_device(idx):
|
||||
call_log.append(idx)
|
||||
|
||||
@staticmethod
|
||||
def device_count():
|
||||
return 2
|
||||
|
||||
class _NpuTorch:
|
||||
npu = _Npu
|
||||
|
||||
monkeypatch.setitem(sys.modules, "torch", _NpuTorch)
|
||||
|
||||
# Mock NPU platform at the source module where it's defined
|
||||
mock_platform = _make_mock_platform(device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES")
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.platforms.current_omni_platform",
|
||||
mock_platform,
|
||||
)
|
||||
|
||||
set_stage_devices(stage_id=0, devices="0,1")
|
||||
|
||||
assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "4,5"
|
||||
Reference in new issue
Block a user