795 lines
32 KiB
Python
795 lines
32 KiB
Python
"""
|
|
Unit tests for SGLang platform abstraction layer.
|
|
|
|
Tests DeviceMixin, SRTPlatform, PlatformEnum, CpuArchEnum, DeviceCapability,
|
|
and the platform discovery / lazy initialization mechanism.
|
|
"""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import torch
|
|
|
|
from sglang.srt.platforms import _load_platform_class, _resolve_platform
|
|
from sglang.srt.platforms.cpu import CpuDeviceMixin, CpuSRTPlatform
|
|
from sglang.srt.platforms.cuda import CudaDeviceMixin, CudaSRTPlatform
|
|
from sglang.srt.platforms.device_mixin import (
|
|
CpuArchEnum,
|
|
DeviceCapability,
|
|
DeviceMixin,
|
|
PlatformEnum,
|
|
)
|
|
from sglang.srt.platforms.interface import SRTPlatform
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cpu_ci(est_time=7, suite="base-a-test-cpu")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers: factory functions to reduce boilerplate
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_device_mixin(enum, name, dtype):
|
|
"""Create a concrete DeviceMixin subclass for testing."""
|
|
|
|
class M(DeviceMixin):
|
|
_enum = enum
|
|
device_name = name
|
|
device_type = dtype
|
|
|
|
def get_device_total_memory(self, device_id=0):
|
|
return 10**9
|
|
|
|
def get_current_memory_usage(self, device=None):
|
|
return 5 * 10**8
|
|
|
|
return M()
|
|
|
|
|
|
class _StubPlatform(SRTPlatform):
|
|
"""Concrete SRTPlatform with minimal defaults for testing overrides."""
|
|
|
|
_enum = PlatformEnum.CUDA
|
|
device_name = "cuda"
|
|
device_type = "cuda"
|
|
|
|
def get_device_total_memory(self, device_id=0):
|
|
return 10**9
|
|
|
|
def get_current_memory_usage(self, device=None):
|
|
return 5 * 10**8
|
|
|
|
def get_default_attention_backend(self):
|
|
return "flashinfer"
|
|
|
|
def get_graph_runner_cls(self):
|
|
return object
|
|
|
|
def get_mha_kv_pool_cls(self):
|
|
return object
|
|
|
|
def get_mla_kv_pool_cls(self):
|
|
return object
|
|
|
|
def get_dsa_kv_pool_cls(self):
|
|
return object
|
|
|
|
def get_paged_allocator_cls(self):
|
|
return object
|
|
|
|
def get_piecewise_backend_cls(self):
|
|
return object
|
|
|
|
|
|
def _make_platform_ep(name, load_fn=None):
|
|
"""Create a mock entry point for platform plugins."""
|
|
ep = MagicMock()
|
|
ep.name = name
|
|
if load_fn is not None:
|
|
ep.load.return_value = load_fn
|
|
else:
|
|
ep.load.return_value = MagicMock()
|
|
return ep
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# PlatformEnum & CpuArchEnum
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPlatformEnum(CustomTestCase):
|
|
"""Tests for PlatformEnum enumeration."""
|
|
|
|
def test_all_expected_values_exist(self):
|
|
expected = {
|
|
"CUDA",
|
|
"ROCM",
|
|
"CPU",
|
|
"XPU",
|
|
"MUSA",
|
|
"NPU",
|
|
"TPU",
|
|
"MPS",
|
|
"OOT",
|
|
"UNSPECIFIED",
|
|
}
|
|
actual = {member.name for member in PlatformEnum}
|
|
self.assertEqual(actual, expected)
|
|
|
|
|
|
class TestCpuArchEnum(CustomTestCase):
|
|
"""Tests for CpuArchEnum enumeration."""
|
|
|
|
def test_all_expected_values_exist(self):
|
|
expected = {"X86", "ARM", "UNSPECIFIED"}
|
|
actual = {member.name for member in CpuArchEnum}
|
|
self.assertEqual(actual, expected)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DeviceCapability
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDeviceCapability(CustomTestCase):
|
|
"""Tests for DeviceCapability custom logic (formatting, conversion)."""
|
|
|
|
def test_as_version_str(self):
|
|
self.assertEqual(DeviceCapability(major=9, minor=0).as_version_str(), "9.0")
|
|
self.assertEqual(DeviceCapability(major=8, minor=9).as_version_str(), "8.9")
|
|
|
|
def test_to_int(self):
|
|
self.assertEqual(DeviceCapability(major=9, minor=0).to_int(), 90)
|
|
self.assertEqual(DeviceCapability(major=8, minor=9).to_int(), 89)
|
|
self.assertEqual(DeviceCapability(major=0, minor=0).to_int(), 0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DeviceMixin
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Platform identity test data: (enum, name, dtype, true_method)
|
|
_PLATFORM_IDENTITY = [
|
|
(PlatformEnum.CUDA, "cuda", "cuda", "is_cuda"),
|
|
(PlatformEnum.ROCM, "rocm", "hip", "is_rocm"),
|
|
(PlatformEnum.CPU, "cpu", "cpu", "is_cpu"),
|
|
(PlatformEnum.XPU, "xpu", "xpu", "is_xpu"),
|
|
(PlatformEnum.MUSA, "musa", "musa", "is_musa"),
|
|
(PlatformEnum.NPU, "npu", "npu", "is_npu"),
|
|
(PlatformEnum.TPU, "tpu", "tpu", "is_tpu"),
|
|
(PlatformEnum.MPS, "mps", "mps", "is_mps"),
|
|
]
|
|
|
|
# is_cuda_alike test data: (enum, name, dtype, expected)
|
|
_CUDA_ALIKE = [
|
|
(PlatformEnum.CUDA, "cuda", "cuda", True),
|
|
(PlatformEnum.ROCM, "rocm", "hip", True),
|
|
(PlatformEnum.MUSA, "musa", "musa", True),
|
|
(PlatformEnum.CPU, "cpu", "cpu", False),
|
|
(PlatformEnum.NPU, "npu", "npu", False),
|
|
]
|
|
|
|
|
|
class TestDeviceMixin(CustomTestCase):
|
|
"""Tests for DeviceMixin base class."""
|
|
|
|
def test_platform_identity_methods(self):
|
|
"""Each platform type returns True for its identity method."""
|
|
for enum_val, name, dtype, method in _PLATFORM_IDENTITY:
|
|
with self.subTest(method=method, enum=enum_val.name):
|
|
mixin = _make_device_mixin(enum_val, name, dtype)
|
|
self.assertTrue(getattr(mixin, method)())
|
|
|
|
def test_is_cuda_alike(self):
|
|
"""is_cuda_alike is True for CUDA/ROCM/MUSA, False otherwise."""
|
|
for enum_val, name, dtype, expected in _CUDA_ALIKE:
|
|
with self.subTest(enum=enum_val.name):
|
|
mixin = _make_device_mixin(enum_val, name, dtype)
|
|
self.assertEqual(mixin.is_cuda_alike(), expected)
|
|
|
|
def test_is_out_of_tree(self):
|
|
oot = _make_device_mixin(PlatformEnum.OOT, "custom", "custom")
|
|
self.assertTrue(oot.is_out_of_tree())
|
|
cuda = _make_device_mixin(PlatformEnum.CUDA, "cuda", "cuda")
|
|
self.assertFalse(cuda.is_out_of_tree())
|
|
|
|
@patch("platform.machine")
|
|
def test_get_cpu_architecture(self, mock_machine):
|
|
"""get_cpu_architecture maps common strings to CpuArchEnum."""
|
|
cases = [
|
|
("x86_64", CpuArchEnum.X86),
|
|
("amd64", CpuArchEnum.X86),
|
|
("i386", CpuArchEnum.X86),
|
|
("i686", CpuArchEnum.X86),
|
|
("X86_64", CpuArchEnum.X86), # case insensitive
|
|
("arm64", CpuArchEnum.ARM),
|
|
("aarch64", CpuArchEnum.ARM),
|
|
("unknown_arch", CpuArchEnum.UNSPECIFIED),
|
|
]
|
|
for machine_str, expected in cases:
|
|
with self.subTest(machine=machine_str):
|
|
mock_machine.return_value = machine_str
|
|
self.assertEqual(DeviceMixin.get_cpu_architecture(), expected)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SRTPlatform
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSRTPlatform(CustomTestCase):
|
|
"""Tests for SRTPlatform base class and default behaviors."""
|
|
|
|
def test_compile_backend_signature_compatibility(self):
|
|
"""get_compile_backend accepts mode keyword arg without error."""
|
|
base = SRTPlatform()
|
|
self.assertEqual(base.get_compile_backend(mode="npugraph_ex"), "inductor")
|
|
|
|
def test_base_device_identity_stays_unspecified(self):
|
|
"""The abstract SRT base should not claim any concrete in-tree device."""
|
|
base = SRTPlatform()
|
|
self.assertFalse(base.is_cuda())
|
|
self.assertFalse(base.is_cuda_alike())
|
|
|
|
|
|
class TestCudaDeviceMixin(CustomTestCase):
|
|
"""Tests for CUDA device operation defaults."""
|
|
|
|
def test_default_get_device_returns_cuda_device(self):
|
|
base = CudaSRTPlatform()
|
|
self.assertEqual(base.get_device(2), torch.device("cuda", 2))
|
|
|
|
def test_cuda_platform_identity(self):
|
|
base = CudaSRTPlatform()
|
|
self.assertTrue(base.is_cuda())
|
|
self.assertTrue(base.is_cuda_alike())
|
|
self.assertIsInstance(base, CudaDeviceMixin)
|
|
|
|
@patch("torch.cuda.get_device_properties")
|
|
def test_default_get_device_total_memory_uses_cuda(
|
|
self, mock_get_device_properties
|
|
):
|
|
mock_get_device_properties.return_value.total_memory = 123
|
|
base = CudaSRTPlatform()
|
|
self.assertEqual(base.get_device_total_memory(1), 123)
|
|
mock_get_device_properties.assert_called_once_with(1)
|
|
|
|
@patch("torch.cuda.max_memory_allocated", return_value=456)
|
|
def test_default_get_current_memory_usage_uses_cuda(
|
|
self, mock_max_memory_allocated
|
|
):
|
|
base = CudaSRTPlatform()
|
|
device = torch.device("cuda", 1)
|
|
self.assertEqual(base.get_current_memory_usage(device), 456.0)
|
|
mock_max_memory_allocated.assert_called_once_with(device)
|
|
|
|
@patch("torch.cuda.set_device")
|
|
def test_default_set_device_uses_cuda(self, mock_set_device):
|
|
base = CudaSRTPlatform()
|
|
device = torch.device("cuda", 1)
|
|
base.set_device(device)
|
|
mock_set_device.assert_called_once_with(device)
|
|
|
|
@patch("torch.cuda.get_device_name", return_value="NVIDIA H100")
|
|
def test_default_get_device_name_uses_cuda(self, mock_get_device_name):
|
|
base = CudaSRTPlatform()
|
|
self.assertEqual(base.get_device_name(1), "NVIDIA H100")
|
|
mock_get_device_name.assert_called_once_with(1)
|
|
|
|
@patch("torch.cuda.get_device_properties")
|
|
def test_default_get_device_uuid_uses_cuda(self, mock_get_device_properties):
|
|
mock_get_device_properties.return_value.uuid = "1234"
|
|
base = CudaSRTPlatform()
|
|
self.assertEqual(base.get_device_uuid(1), "1234")
|
|
mock_get_device_properties.assert_called_once_with(1)
|
|
|
|
@patch("torch.cuda.get_device_capability", return_value=(9, 0))
|
|
def test_default_get_device_capability_uses_cuda(self, mock_get_device_capability):
|
|
base = CudaSRTPlatform()
|
|
self.assertEqual(base.get_device_capability(1), DeviceCapability(9, 0))
|
|
mock_get_device_capability.assert_called_once_with(1)
|
|
|
|
@patch("torch.cuda.empty_cache")
|
|
def test_default_empty_cache_uses_cuda(self, mock_empty_cache):
|
|
base = CudaSRTPlatform()
|
|
base.empty_cache()
|
|
mock_empty_cache.assert_called_once_with()
|
|
|
|
@patch("torch.cuda.synchronize")
|
|
def test_default_synchronize_uses_cuda(self, mock_synchronize):
|
|
base = CudaSRTPlatform()
|
|
base.synchronize()
|
|
mock_synchronize.assert_called_once_with()
|
|
|
|
@patch("torch.cuda.mem_get_info", return_value=(123, 456), create=True)
|
|
def test_default_get_available_memory_uses_cuda(self, mock_mem_get_info):
|
|
base = CudaSRTPlatform()
|
|
self.assertEqual(base.get_available_memory(1), (123, 456))
|
|
mock_mem_get_info.assert_called_once_with(1)
|
|
|
|
def test_default_distributed_backend_is_nccl(self):
|
|
base = CudaSRTPlatform()
|
|
self.assertEqual(base.get_torch_distributed_backend_str(), "nccl")
|
|
|
|
@patch("torch.cuda.manual_seed_all")
|
|
@patch("torch.manual_seed")
|
|
@patch("sglang.srt.platforms.device_mixin.np.random.seed")
|
|
@patch("sglang.srt.platforms.device_mixin.random.seed")
|
|
def test_default_seed_everything_seeds_cuda(
|
|
self, mock_random_seed, mock_np_seed, mock_torch_seed, mock_cuda_seed
|
|
):
|
|
CudaSRTPlatform.seed_everything(123)
|
|
mock_random_seed.assert_called_once_with(123)
|
|
mock_np_seed.assert_called_once_with(123)
|
|
mock_torch_seed.assert_called_once_with(123)
|
|
mock_cuda_seed.assert_called_once_with(123)
|
|
|
|
def test_cuda_srt_platform_capabilities(self):
|
|
base = CudaSRTPlatform()
|
|
self.assertTrue(base.supports_fp8())
|
|
self.assertTrue(base.support_cuda_graph())
|
|
self.assertTrue(base.support_piecewise_cuda_graph())
|
|
|
|
|
|
class TestCpuDeviceMixin(CustomTestCase):
|
|
"""Tests for CPU device operation defaults (covers both x86 and ARM)."""
|
|
|
|
def test_cpu_platform_identity(self):
|
|
base = CpuSRTPlatform()
|
|
self.assertTrue(base.is_cpu())
|
|
self.assertFalse(base.is_cuda())
|
|
self.assertFalse(base.is_cuda_alike())
|
|
self.assertIsInstance(base, CpuDeviceMixin)
|
|
|
|
def test_default_get_device_returns_cpu_device(self):
|
|
base = CpuSRTPlatform()
|
|
# ``local_rank`` is ignored — CPU has no per-rank device.
|
|
self.assertEqual(base.get_device(0), torch.device("cpu"))
|
|
self.assertEqual(base.get_device(7), torch.device("cpu"))
|
|
|
|
@patch("sglang.srt.platforms.cpu.psutil.virtual_memory")
|
|
def test_default_get_device_total_memory_uses_psutil(self, mock_vm):
|
|
mock_vm.return_value.total = 12345
|
|
base = CpuSRTPlatform()
|
|
self.assertEqual(base.get_device_total_memory(), 12345)
|
|
|
|
@patch("sglang.srt.platforms.cpu.psutil.virtual_memory")
|
|
def test_default_get_available_memory_uses_psutil(self, mock_vm):
|
|
mock_vm.return_value.available = 100
|
|
mock_vm.return_value.total = 200
|
|
base = CpuSRTPlatform()
|
|
self.assertEqual(base.get_available_memory(), (100, 200))
|
|
|
|
@patch("sglang.srt.platforms.cpu.psutil.virtual_memory")
|
|
def test_default_get_current_memory_usage_is_system_used(self, mock_vm):
|
|
mock_vm.return_value.total = 1000
|
|
mock_vm.return_value.available = 300
|
|
base = CpuSRTPlatform()
|
|
# system-used == total - available (not per-process RSS)
|
|
self.assertEqual(base.get_current_memory_usage(), 700.0)
|
|
|
|
@patch("sglang.srt.platforms.cpu.psutil.virtual_memory")
|
|
def test_memory_free_contract_yields_available(self, mock_vm):
|
|
# The [Active] contract free = total - used must yield psutil.available.
|
|
mock_vm.return_value.total = 1000
|
|
mock_vm.return_value.available = 300
|
|
base = CpuSRTPlatform()
|
|
free = base.get_device_total_memory() - base.get_current_memory_usage()
|
|
self.assertEqual(free, 300)
|
|
|
|
@patch("torch.cpu.set_device")
|
|
def test_default_set_device_uses_torch_cpu(self, mock_set_device):
|
|
base = CpuSRTPlatform()
|
|
device = torch.device("cpu")
|
|
base.set_device(device)
|
|
# Documented CPU no-op, but called for symmetry with CudaDeviceMixin.
|
|
mock_set_device.assert_called_once_with(device)
|
|
|
|
def test_default_set_device_does_not_flip_default(self):
|
|
base = CpuSRTPlatform()
|
|
# Must not call torch.set_default_device — process-wide default stays put.
|
|
before = torch.empty(0).device
|
|
base.set_device(torch.device("cpu"))
|
|
after = torch.empty(0).device
|
|
self.assertEqual(before, after)
|
|
|
|
@patch("sglang.srt.platforms.cpu.gc.collect")
|
|
def test_default_empty_cache_calls_gc_collect(self, mock_collect):
|
|
base = CpuSRTPlatform()
|
|
base.empty_cache()
|
|
mock_collect.assert_called_once_with()
|
|
|
|
@patch("torch.cpu.synchronize")
|
|
def test_default_synchronize_uses_torch_cpu(self, mock_synchronize):
|
|
base = CpuSRTPlatform()
|
|
base.synchronize()
|
|
mock_synchronize.assert_called_once_with()
|
|
|
|
def test_default_distributed_backend_is_gloo(self):
|
|
base = CpuSRTPlatform()
|
|
self.assertEqual(base.get_torch_distributed_backend_str(), "gloo")
|
|
|
|
@patch("platform.machine", return_value="aarch64")
|
|
def test_cpu_arch_property_resolves_and_caches(self, mock_machine):
|
|
base = CpuSRTPlatform()
|
|
self.assertEqual(base.cpu_arch, CpuArchEnum.ARM)
|
|
# cached_property: second access must not re-query platform.machine
|
|
call_count = mock_machine.call_count
|
|
self.assertEqual(base.cpu_arch, CpuArchEnum.ARM)
|
|
self.assertEqual(mock_machine.call_count, call_count)
|
|
|
|
@patch("platform.machine", return_value="aarch64")
|
|
def test_get_device_name_arm_branch(self, _mock_machine):
|
|
base = CpuSRTPlatform()
|
|
name = base.get_device_name()
|
|
self.assertIn("aarch64", name)
|
|
|
|
@patch("platform.machine", return_value="x86_64")
|
|
def test_get_device_name_x86_branch(self, _mock_machine):
|
|
base = CpuSRTPlatform()
|
|
name = base.get_device_name()
|
|
self.assertIn("x86_64", name)
|
|
|
|
@patch("platform.machine", return_value="aarch64")
|
|
def test_get_device_uuid_returns_machine(self, _mock_machine):
|
|
base = CpuSRTPlatform()
|
|
self.assertEqual(base.get_device_uuid(), "aarch64")
|
|
|
|
def test_get_device_capability_returns_none(self):
|
|
base = CpuSRTPlatform()
|
|
self.assertIsNone(base.get_device_capability())
|
|
|
|
def test_cpu_srt_platform_capabilities(self):
|
|
base = CpuSRTPlatform()
|
|
self.assertFalse(base.supports_fp8())
|
|
self.assertFalse(base.support_cuda_graph())
|
|
self.assertFalse(base.support_piecewise_cuda_graph())
|
|
# Override of the SRTPlatform default (True) — no GPU to pin to.
|
|
self.assertFalse(base.is_pin_memory_available())
|
|
|
|
|
|
class TestSRTPlatformOverrides(CustomTestCase):
|
|
"""Tests for SRTPlatform method overrides via plugins."""
|
|
|
|
def test_custom_get_dispatch_key_name(self):
|
|
class P(_StubPlatform):
|
|
_enum = PlatformEnum.NPU
|
|
device_name = "npu"
|
|
device_type = "npu"
|
|
|
|
def get_dispatch_key_name(self):
|
|
return "npu"
|
|
|
|
self.assertEqual(P().get_dispatch_key_name(), "npu")
|
|
|
|
def test_custom_get_compile_backend(self):
|
|
class P(_StubPlatform):
|
|
_enum = PlatformEnum.NPU
|
|
device_name = "npu"
|
|
device_type = "npu"
|
|
|
|
def get_compile_backend(self, mode=None):
|
|
return "inductor"
|
|
|
|
self.assertEqual(P().get_compile_backend(mode="npugraph_ex"), "inductor")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Platform Discovery: _resolve_platform
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestResolvePlatformWithEnv(CustomTestCase):
|
|
"""Tests for _resolve_platform when SGLANG_PLATFORM is set."""
|
|
|
|
@patch("sglang.srt.platforms.entry_points")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_selected_plugin_activates(self, mock_envs, mock_ep):
|
|
"""When SGLANG_PLATFORM matches an entry point, it activates that plugin."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = "my_hardware"
|
|
plugin_fn = MagicMock(return_value="pkg.Mod:MyPlatform")
|
|
mock_ep.return_value = [_make_platform_ep("my_hardware", plugin_fn)]
|
|
with patch("sglang.srt.platforms._load_platform_class") as mock_load:
|
|
mock_instance = MagicMock()
|
|
mock_load.return_value = MagicMock(return_value=mock_instance)
|
|
result = _resolve_platform()
|
|
mock_load.assert_called_once_with("pkg.Mod:MyPlatform")
|
|
self.assertEqual(result, mock_instance)
|
|
|
|
@patch("sglang.srt.platforms.entry_points")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_selected_plugin_not_found(self, mock_envs, mock_ep):
|
|
"""When SGLANG_PLATFORM names a nonexistent plugin, raise RuntimeError."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = "nonexistent"
|
|
mock_ep.return_value = []
|
|
with self.assertRaises(RuntimeError):
|
|
_resolve_platform()
|
|
|
|
@patch("sglang.srt.platforms.entry_points")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_selected_plugin_hardware_unavailable(self, mock_envs, mock_ep):
|
|
"""When activate() returns None, hardware is not available."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = "my_hardware"
|
|
plugin_fn = MagicMock(return_value=None)
|
|
mock_ep.return_value = [_make_platform_ep("my_hardware", plugin_fn)]
|
|
with self.assertRaises(RuntimeError):
|
|
_resolve_platform()
|
|
|
|
@patch("sglang.srt.platforms.entry_points")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_selected_plugin_load_exception(self, mock_envs, mock_ep):
|
|
"""When ep.load() or activate() throws, exception is re-raised."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = "my_hardware"
|
|
plugin_fn = MagicMock(side_effect=ImportError("missing dep"))
|
|
mock_ep.return_value = [_make_platform_ep("my_hardware", plugin_fn)]
|
|
with self.assertRaises(ImportError):
|
|
_resolve_platform()
|
|
|
|
@patch("sglang.srt.platforms.entry_points")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_other_plugins_not_loaded(self, mock_envs, mock_ep):
|
|
"""When SGLANG_PLATFORM is set, other plugins are not imported."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = "target_hw"
|
|
target_fn = MagicMock(return_value="pkg.Mod:TargetPlatform")
|
|
other_ep = _make_platform_ep("other_hw") # default load returns MagicMock
|
|
target_ep = _make_platform_ep("target_hw", target_fn)
|
|
mock_ep.return_value = [other_ep, target_ep]
|
|
with patch("sglang.srt.platforms._load_platform_class") as mock_load:
|
|
mock_load.return_value = MagicMock(return_value=MagicMock())
|
|
_resolve_platform()
|
|
# Only the target entry point should be loaded
|
|
target_ep.load.assert_called_once()
|
|
other_ep.load.assert_not_called()
|
|
|
|
|
|
class TestResolvePlatformAutoDiscover(CustomTestCase):
|
|
"""Tests for _resolve_platform auto-discovery when SGLANG_PLATFORM is not set."""
|
|
|
|
@patch("sglang.srt.platforms.torch")
|
|
def test_is_cuda_available_excludes_rocm(self, mock_torch):
|
|
"""ROCm exposes torch.cuda, but should not use the CUDA platform identity."""
|
|
mock_torch.cuda.is_available.return_value = True
|
|
mock_torch.version.hip = "6.0"
|
|
|
|
import sglang.srt.platforms as plat_mod
|
|
|
|
self.assertFalse(plat_mod._is_cuda_available())
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_single_plugin_activates(self, mock_envs, mock_load):
|
|
"""When exactly one plugin activates, return its platform instance."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
plugin_fn = MagicMock(return_value="pkg.Mod:MyPlatform")
|
|
mock_load.return_value = {"my_hw": (plugin_fn, "my-hw-dist")}
|
|
with patch("sglang.srt.platforms._load_platform_class") as mock_resolve:
|
|
mock_instance = MagicMock()
|
|
mock_resolve.return_value = MagicMock(return_value=mock_instance)
|
|
result = _resolve_platform()
|
|
mock_resolve.assert_called_once_with("pkg.Mod:MyPlatform")
|
|
self.assertEqual(result, mock_instance)
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms._is_cuda_available")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_no_plugin_activates_cuda_fallback(
|
|
self, mock_envs, mock_is_cuda_available, mock_load
|
|
):
|
|
"""When CUDA is available and no plugin activates, return CUDA defaults."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
mock_is_cuda_available.return_value = True
|
|
mock_load.return_value = {}
|
|
result = _resolve_platform()
|
|
self.assertIsInstance(result, CudaSRTPlatform)
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms._is_cuda_available")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_no_plugin_no_cuda_activates_base_fallback(
|
|
self, mock_envs, mock_is_cuda_available, mock_load
|
|
):
|
|
"""When no plugin or CUDA is available, return the abstract base platform."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
mock_is_cuda_available.return_value = False
|
|
mock_load.return_value = {}
|
|
result = _resolve_platform()
|
|
self.assertIsInstance(result, SRTPlatform)
|
|
self.assertNotIsInstance(result, CudaSRTPlatform)
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms.torch")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_no_plugin_rocm_does_not_activate_cuda_fallback(
|
|
self, mock_envs, mock_torch, mock_load
|
|
):
|
|
"""ROCm exposes torch.cuda but must not use the CUDA fallback platform."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
mock_torch.cuda.is_available.return_value = True
|
|
mock_torch.version.hip = "6.0"
|
|
mock_load.return_value = {}
|
|
|
|
result = _resolve_platform()
|
|
|
|
self.assertIsInstance(result, SRTPlatform)
|
|
self.assertNotIsInstance(result, CudaSRTPlatform)
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms._is_cuda_available")
|
|
@patch("sglang.srt.platforms._is_cpu_available")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_no_plugin_cpu_engine_enabled_activates_cpu_fallback(
|
|
self, mock_envs, mock_is_cpu, mock_is_cuda, mock_load
|
|
):
|
|
"""SGLANG_USE_CPU_ENGINE=1 + no plugins → CpuSRTPlatform."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
mock_is_cpu.return_value = True
|
|
mock_is_cuda.return_value = False
|
|
mock_load.return_value = {}
|
|
result = _resolve_platform()
|
|
self.assertIsInstance(result, CpuSRTPlatform)
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms._is_cuda_available")
|
|
@patch("sglang.srt.platforms._is_cpu_available")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_cpu_engine_wins_over_cuda(
|
|
self, mock_envs, mock_is_cpu, mock_is_cuda, mock_load
|
|
):
|
|
"""When both CPU engine and CUDA are available, explicit opt-in wins."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
mock_is_cpu.return_value = True
|
|
mock_is_cuda.return_value = True
|
|
mock_load.return_value = {}
|
|
result = _resolve_platform()
|
|
self.assertIsInstance(result, CpuSRTPlatform)
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms._is_cuda_available")
|
|
@patch("sglang.srt.platforms._is_cpu_available")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_no_plugin_cpu_engine_disabled_prefers_cuda(
|
|
self, mock_envs, mock_is_cpu, mock_is_cuda, mock_load
|
|
):
|
|
"""Regression: CPU opt-out leaves the existing CUDA fallback path intact."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
mock_is_cpu.return_value = False
|
|
mock_is_cuda.return_value = True
|
|
mock_load.return_value = {}
|
|
result = _resolve_platform()
|
|
self.assertIsInstance(result, CudaSRTPlatform)
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_multiple_plugins_activate_raises(self, mock_envs, mock_load):
|
|
"""When multiple plugins activate, raise RuntimeError."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
fn1 = MagicMock(return_value="pkg1.Mod:Platform1")
|
|
fn2 = MagicMock(return_value="pkg2.Mod:Platform2")
|
|
mock_load.return_value = {"hw1": (fn1, "hw1-dist"), "hw2": (fn2, "hw2-dist")}
|
|
with self.assertRaises(RuntimeError):
|
|
_resolve_platform()
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_plugin_exception_does_not_crash(self, mock_envs, mock_load):
|
|
"""When a plugin's activate() throws, it is skipped, others continue."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
bad_fn = MagicMock(side_effect=RuntimeError("broken"))
|
|
good_fn = MagicMock(return_value="pkg.Mod:GoodPlatform")
|
|
mock_load.return_value = {
|
|
"bad": (bad_fn, "bad-dist"),
|
|
"good": (good_fn, "good-dist"),
|
|
}
|
|
with patch("sglang.srt.platforms._load_platform_class") as mock_resolve:
|
|
mock_instance = MagicMock()
|
|
mock_resolve.return_value = MagicMock(return_value=mock_instance)
|
|
result = _resolve_platform()
|
|
mock_resolve.assert_called_once_with("pkg.Mod:GoodPlatform")
|
|
self.assertEqual(result, mock_instance)
|
|
|
|
@patch("sglang.srt.platforms.load_plugins_by_group")
|
|
@patch("sglang.srt.platforms.envs")
|
|
def test_plugin_returns_none_is_skipped(self, mock_envs, mock_load):
|
|
"""When a plugin's activate() returns None, it is skipped (hardware unavailable)."""
|
|
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
|
none_fn = MagicMock(return_value=None)
|
|
good_fn = MagicMock(return_value="pkg.Mod:GoodPlatform")
|
|
mock_load.return_value = {
|
|
"unavailable": (none_fn, "unavail-dist"),
|
|
"good": (good_fn, "good-dist"),
|
|
}
|
|
with patch("sglang.srt.platforms._load_platform_class") as mock_resolve:
|
|
mock_instance = MagicMock()
|
|
mock_resolve.return_value = MagicMock(return_value=mock_instance)
|
|
result = _resolve_platform()
|
|
# Only the good plugin activated; single activation succeeds
|
|
mock_resolve.assert_called_once_with("pkg.Mod:GoodPlatform")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Platform Discovery: _load_platform_class
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLoadPlatformClass(CustomTestCase):
|
|
"""Tests for _load_platform_class qualname resolution."""
|
|
|
|
@patch("sglang.srt.platforms.pkgutil.resolve_name")
|
|
def test_valid_subclass(self, mock_resolve):
|
|
"""Valid SRTPlatform subclass resolves successfully."""
|
|
mock_resolve.return_value = type("MyPlatform", (SRTPlatform,), {})
|
|
result = _load_platform_class("pkg.Mod:MyPlatform")
|
|
self.assertTrue(issubclass(result, SRTPlatform))
|
|
|
|
@patch("sglang.srt.platforms.pkgutil.resolve_name")
|
|
def test_non_subclass_raises_type_error(self, mock_resolve):
|
|
"""Non-SRTPlatform class raises TypeError."""
|
|
mock_resolve.return_value = str
|
|
with self.assertRaises(TypeError):
|
|
_load_platform_class("builtins.str")
|
|
|
|
@patch("sglang.srt.platforms.pkgutil.resolve_name")
|
|
def test_non_type_raises_type_error(self, mock_resolve):
|
|
"""Non-type object raises TypeError."""
|
|
mock_resolve.return_value = "not a class"
|
|
with self.assertRaises(TypeError):
|
|
_load_platform_class("something")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Platform Discovery: current_platform lazy init
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCurrentPlatformLazyInit(CustomTestCase):
|
|
"""Tests for current_platform lazy initialization via module __getattr__."""
|
|
|
|
def setUp(self):
|
|
"""Reset module-level cache before each test."""
|
|
import sglang.srt.platforms as plat_mod
|
|
|
|
self._saved_platform = plat_mod._current_platform
|
|
plat_mod._current_platform = None
|
|
|
|
def tearDown(self):
|
|
"""Restore original _current_platform after each test."""
|
|
import sglang.srt.platforms as plat_mod
|
|
|
|
plat_mod._current_platform = self._saved_platform
|
|
|
|
@patch("sglang.srt.platforms._resolve_platform")
|
|
def test_first_access_triggers_resolve(self, mock_resolve):
|
|
"""First access to current_platform calls _resolve_platform."""
|
|
mock_instance = MagicMock(spec=SRTPlatform)
|
|
mock_resolve.return_value = mock_instance
|
|
import sglang.srt.platforms as plat_mod
|
|
|
|
result = plat_mod.current_platform
|
|
mock_resolve.assert_called_once()
|
|
self.assertEqual(result, mock_instance)
|
|
|
|
@patch("sglang.srt.platforms._resolve_platform")
|
|
def test_subsequent_access_uses_cache(self, mock_resolve):
|
|
"""Subsequent accesses return cached instance without re-resolving."""
|
|
mock_instance = MagicMock(spec=SRTPlatform)
|
|
mock_resolve.return_value = mock_instance
|
|
import sglang.srt.platforms as plat_mod
|
|
|
|
_ = plat_mod.current_platform
|
|
_ = plat_mod.current_platform
|
|
mock_resolve.assert_called_once()
|
|
|
|
def test_other_attribute_raises_error(self):
|
|
"""Accessing non-existent module attribute raises AttributeError."""
|
|
import sglang.srt.platforms as plat_mod
|
|
|
|
with self.assertRaises(AttributeError):
|
|
_ = plat_mod.nonexistent_attribute
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import unittest
|
|
|
|
unittest.main()
|