[Hardware] Support platforms and plugin system (#774)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
This commit is contained in:
Canlin Guo authored and GitHub committed 2026-01-28 22:12:52 +08:00
1 parent b11d436e19
commit c9392dff01
103 files changed
+1770 -1218

No files matched your search

+53 -2
View File
@@ -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"