minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity
Files
65be3fa71a4d35c5166cd1b510e5e9203989455c
sglang/test/registered/unit/platforms/test_platform_interface.py
T
Yuhan ChenYuhan ChenXiaoyu Zhang
 99a44c88d4 Add out-of-tree DFlash extension points (#38740) 
Co-authored-by: Yuhan Chen <yuhanc@fb.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
2026-09-20 14:54:46 +08:00

945 lines
39 KiB
Python
Raw Blame History

"""
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 CpuSRTPlatform
from sglang.srt.platforms.cuda import CudaSRTPlatform
from sglang.srt.platforms.device_mixin import (
CpuArchEnum,
DeviceCapability,
DeviceMixin,
PlatformEnum,
)
from sglang.srt.platforms.interface import SRTPlatform
from sglang.srt.platforms.npu import NPUSRTPlatform
from sglang.srt.platforms.rocm import RocmSRTPlatform
from sglang.srt.platforms.xpu import XpuSRTPlatform
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, 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()
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())
def test_pin_memory_default_is_conservative(self):
mixin = _make_device_mixin(PlatformEnum.OOT, "custom", "custom")
self.assertFalse(mixin.is_pin_memory_available())
self.assertFalse(mixin.is_pin_memory_available(device="cpu"))
@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())
def test_base_pin_memory_default_is_conservative(self):
base = SRTPlatform()
self.assertFalse(base.is_pin_memory_available())
self.assertFalse(base.is_pin_memory_available(device="cpu"))
def test_base_speculative_capability_defaults_are_conservative(self):
base = SRTPlatform()
self.assertFalse(base.supports_speculative_algorithm("DFLASH"))
self.assertFalse(
base.supports_speculative_draft_attention_backend(
"DFLASH", "custom_backend"
)
)
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))
@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)
def test_pin_memory_available_for_cuda_targets(self):
base = CudaSRTPlatform()
self.assertTrue(base.is_pin_memory_available())
self.assertTrue(base.is_pin_memory_available(device="cuda"))
self.assertTrue(base.is_pin_memory_available(device=torch.device("cuda", 0)))
self.assertFalse(base.is_pin_memory_available(device="cpu"))
def test_rocm_inherits_cuda_pin_memory_behavior(self):
base = RocmSRTPlatform()
self.assertTrue(base.is_pin_memory_available())
self.assertTrue(base.is_pin_memory_available(device="cuda"))
self.assertFalse(base.is_pin_memory_available(device="cpu"))
@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 TestXpuDeviceMixin(CustomTestCase):
"""Tests for XPU device operation defaults."""
def test_default_get_device_returns_xpu_device(self):
base = XpuSRTPlatform()
self.assertEqual(base.get_device(2), torch.device("xpu", 2))
# TODO: @patch("torch.xpu.get_device_capability", return_value=(9, 0))
def test_default_get_device_capability_uses_xpu(self):
# torch.ops.sgl_kernel.query_device is only registered by XPU builds
# of sgl-kernel, so patch the op namespace attribute with create=True
# (a dotted @patch target would fail to import on CPU/CUDA machines).
# torch.xpu.current_device() likewise needs an XPU device; mock it.
base = XpuSRTPlatform()
fake_query_device = MagicMock()
fake_query_device.default.return_value = (9, 0)
with (
patch("torch.xpu.current_device", return_value=0),
patch.object(
torch.ops.sgl_kernel, "query_device", fake_query_device, create=True
),
):
self.assertEqual(base.get_device_capability(0), DeviceCapability(9, 0))
fake_query_device.default.assert_called_once_with(0)
def test_pin_memory_available_for_xpu_targets(self):
base = XpuSRTPlatform()
self.assertTrue(base.is_pin_memory_available())
self.assertTrue(base.is_pin_memory_available(device="xpu"))
self.assertTrue(base.is_pin_memory_available(device=torch.device("xpu", 0)))
self.assertFalse(base.is_pin_memory_available(device="cpu"))
@patch("torch.xpu.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_xpu(
self, mock_random_seed, mock_np_seed, mock_torch_seed, mock_xpu_seed
):
XpuSRTPlatform.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_xpu_seed.assert_called_once_with(123)
def test_xpu_srt_platform_capabilities(self):
base = XpuSRTPlatform()
self.assertFalse(base.supports_fp8())
self.assertTrue(base.support_cuda_graph())
self.assertTrue(base.support_piecewise_cuda_graph())
class TestNpuDeviceMixin(CustomTestCase):
"""Tests for NPU device operation defaults."""
def setUp(self):
# torch.device("npu", ...) requires the "npu" device type, which
# torch_npu registers via the privateuse1 backend rename; CPU-only
# builds lack it. Register it per-test so only this suite carries
# the process-wide side effect.
try:
torch.utils.rename_privateuse1_backend("npu")
except Exception:
# Re-registration with a different name raises on some versions;
# real NPU machines may have already renamed the backend.
pass
super().setUp()
def test_default_get_device_returns_npu_device(self):
base = NPUSRTPlatform()
self.assertEqual(base.get_device(2), torch.device("npu", 2))
def test_default_get_device_capability_reports_zero(self):
# torch_npu's get_device_capability is configured via the environment
# variable TORCH_NPU_DEVICE_CAPABILITY purely for native-PyTorch
# compatibility; it does not reflect the real NPU hardware. The
# platform therefore reports (0, 0) without consulting torch.npu.
base = NPUSRTPlatform()
mock_npu = MagicMock()
with patch.object(torch, "npu", mock_npu, create=True):
self.assertEqual(base.get_device_capability(1), DeviceCapability(0, 0))
mock_npu.get_device_capability.assert_not_called()
def test_memory_queries_delegate_to_torch_npu(self):
base = NPUSRTPlatform()
mock_npu = MagicMock()
mock_npu.get_device_properties.return_value.total_memory = 32 * 1024**3
mock_npu.max_memory_allocated.return_value = 5 * 10**8
mock_npu.mem_get_info.return_value = (10**9, 2 * 10**9)
with patch.object(torch, "npu", mock_npu, create=True):
self.assertEqual(base.get_device_total_memory(1), 32 * 1024**3)
mock_npu.get_device_properties.assert_called_once_with(1)
self.assertEqual(base.get_current_memory_usage(), 5 * 10**8)
mock_npu.max_memory_allocated.assert_called_once_with(None)
device = torch.device("npu", 0)
base.get_current_memory_usage(device)
mock_npu.max_memory_allocated.assert_called_with(device)
self.assertEqual(base.get_available_memory(2), (10**9, 2 * 10**9))
mock_npu.mem_get_info.assert_called_once_with(2)
def test_device_info_queries_delegate_to_torch_npu(self):
base = NPUSRTPlatform()
mock_npu = MagicMock()
mock_npu.get_device_name.return_value = "Ascend910B4"
mock_npu.get_device_properties.return_value.uuid = "npu-uuid-0"
with patch.object(torch, "npu", mock_npu, create=True):
self.assertEqual(base.get_device_name(1), "Ascend910B4")
mock_npu.get_device_name.assert_called_once_with(1)
self.assertEqual(base.get_device_uuid(1), "npu-uuid-0")
mock_npu.get_device_properties.assert_called_once_with(1)
def test_device_state_ops_delegate_to_torch_npu(self):
base = NPUSRTPlatform()
mock_npu = MagicMock()
with patch.object(torch, "npu", mock_npu, create=True):
device = torch.device("npu", 3)
base.set_device(device)
mock_npu.set_device.assert_called_once_with(device)
base.empty_cache()
mock_npu.empty_cache.assert_called_once()
base.synchronize()
mock_npu.synchronize.assert_called_once()
def test_pin_memory_available_for_npu_targets(self):
# Pinned memory stays disabled on NPU: torch_npu's pinned-memory +
# non_blocking H2D path is not verified against CANN (see
# NPUSRTPlatform.is_pin_memory_available).
base = NPUSRTPlatform()
self.assertFalse(base.is_pin_memory_available())
self.assertFalse(base.is_pin_memory_available(device="npu"))
self.assertFalse(base.is_pin_memory_available(device=torch.device("npu", 0)))
self.assertFalse(base.is_pin_memory_available(device="cpu"))
self.assertFalse(base.is_pin_memory_available(device=torch.device("cpu")))
def test_default_seed_everything_seeds_npu(self):
mock_npu = MagicMock()
with (
patch.object(torch, "npu", mock_npu, create=True),
patch("torch.manual_seed") as mock_torch_seed,
patch("sglang.srt.platforms.device_mixin.np.random.seed") as mock_np_seed,
patch("sglang.srt.platforms.device_mixin.random.seed") as mock_random_seed,
):
NPUSRTPlatform.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_npu.manual_seed_all.assert_called_once_with(123)
def test_seed_everything_none_seed_is_noop(self):
mock_npu = MagicMock()
with (
patch.object(torch, "npu", mock_npu, create=True),
patch("torch.manual_seed") as mock_torch_seed,
patch("sglang.srt.platforms.device_mixin.np.random.seed") as mock_np_seed,
patch("sglang.srt.platforms.device_mixin.random.seed") as mock_random_seed,
):
NPUSRTPlatform.seed_everything(None)
mock_random_seed.assert_not_called()
mock_np_seed.assert_not_called()
mock_torch_seed.assert_not_called()
mock_npu.manual_seed_all.assert_not_called()
def test_npu_srt_platform_identity(self):
base = NPUSRTPlatform()
self.assertTrue(base.is_npu())
self.assertFalse(base.is_cuda())
self.assertFalse(base.is_cuda_alike())
self.assertEqual(base.device_name, "npu")
self.assertEqual(base.device_type, "npu")
def test_get_default_attention_backend_is_ascend(self):
self.assertEqual(NPUSRTPlatform().get_default_attention_backend(), "ascend")
def test_get_dispatch_key_name_is_npu(self):
self.assertEqual(NPUSRTPlatform().get_dispatch_key_name(), "npu")
def test_npu_srt_platform_capabilities(self):
base = NPUSRTPlatform()
self.assertTrue(base.supports_fp8())
self.assertTrue(base.support_cuda_graph())
self.assertFalse(base.support_piecewise_cuda_graph())
class TestCpuDeviceMixin(CustomTestCase):
"""Tests for CPU device operation defaults (covers both x86 and ARM)."""
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_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)
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("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)
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())
# CPU has no GPU to pin host memory to.
self.assertFalse(base.is_pin_memory_available())
self.assertFalse(base.is_pin_memory_available(device="cpu"))
class TestPinMemoryAvailability(CustomTestCase):
"""Tests for common pin-memory helper dispatch through platforms."""
def test_srt_platform_does_not_shadow_device_mixin_pin_memory_override(self):
class M(DeviceMixin):
def is_pin_memory_available(self, device=None):
return device == "custom"
class P(SRTPlatform, M):
pass
self.assertTrue(P().is_pin_memory_available(device="custom"))
def test_device_mixin_can_precede_srt_platform_for_pin_memory_override(self):
class M(DeviceMixin):
def is_pin_memory_available(self, device=None):
return device == "custom"
class P(M, SRTPlatform):
pass
self.assertTrue(P().is_pin_memory_available(device="custom"))
def test_common_wrapper_dispatches_to_current_platform_with_device(self):
from sglang.srt.utils import common
class P(SRTPlatform):
_enum = PlatformEnum.OOT
device_name = "custom"
device_type = "custom"
def __init__(self):
self.calls = []
def is_pin_memory_available(self, device=None):
self.calls.append(device)
return True
platform = P()
device = torch.device("cuda", 0)
with patch.object(common, "current_platform", platform):
self.assertTrue(common.is_pin_memory_available(device))
self.assertEqual(platform.calls, [device])
def test_common_wrapper_dispatches_to_current_platform_without_device(self):
from sglang.srt.utils import common
class P(SRTPlatform):
_enum = PlatformEnum.OOT
device_name = "custom"
device_type = "custom"
def __init__(self):
self.calls = []
def is_pin_memory_available(self, device=None):
self.calls.append(device)
return True
platform = P()
with patch.object(common, "current_platform", platform):
self.assertTrue(common.is_pin_memory_available())
self.assertEqual(platform.calls, [None])
def test_oot_platform_override_true_is_used(self):
from sglang.srt.utils import common
class P(SRTPlatform):
_enum = PlatformEnum.OOT
device_name = "custom"
device_type = "custom"
def is_pin_memory_available(self, device=None):
return True
with patch.object(common, "current_platform", P()):
self.assertTrue(common.is_pin_memory_available())
def test_oot_platform_override_false_is_used(self):
from sglang.srt.utils import common
class P(SRTPlatform):
_enum = PlatformEnum.OOT
device_name = "custom"
device_type = "custom"
def is_pin_memory_available(self, device=None):
return False
with patch.object(common, "current_platform", P()):
self.assertFalse(common.is_pin_memory_available())
def test_oot_platform_without_override_uses_conservative_default(self):
from sglang.srt.utils import common
class P(SRTPlatform):
_enum = PlatformEnum.OOT
device_name = "custom"
device_type = "custom"
with (
patch.object(common, "current_platform", P()),
patch("torch.cuda.is_available", return_value=True) as mock_cuda_available,
):
self.assertFalse(common.is_pin_memory_available())
mock_cuda_available.assert_not_called()
# ---------------------------------------------------------------------------
# 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()
Reference in New Issue View Git Blame Copy Permalink