107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from sglang.srt.arg_groups.overrides import resolution_result
|
|
from sglang.srt.arg_groups.platform_hook import handle_cpu_backends
|
|
from sglang.srt.arg_groups.validation_hook import validate_ib_devices
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=6, suite="base-b-test-cpu")
|
|
register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
|
|
|
|
|
|
class TestServerArgsCPUBackend(unittest.TestCase):
|
|
def _make_server_args(self, attention_backend=None):
|
|
server_args = ServerArgs.__new__(ServerArgs)
|
|
server_args.device = "cpu"
|
|
server_args.attention_backend = attention_backend
|
|
server_args.sampling_backend = None
|
|
return server_args
|
|
|
|
@patch("sglang.srt.arg_groups.platform_hook.is_host_cpu_arm64", return_value=True)
|
|
def test_arm_cpu_defaults_to_torch_native(self, _mock_is_arm64):
|
|
server_args = self._make_server_args()
|
|
|
|
handle_cpu_backends(server_args)
|
|
|
|
self.assertEqual(
|
|
resolution_result(server_args, "attention_backend"), "torch_native"
|
|
)
|
|
self.assertEqual(resolution_result(server_args, "sampling_backend"), "pytorch")
|
|
|
|
@patch("sglang.srt.arg_groups.platform_hook.is_host_cpu_arm64", return_value=False)
|
|
def test_x86_cpu_defaults_to_intel_amx(self, _mock_is_arm64):
|
|
server_args = self._make_server_args()
|
|
|
|
handle_cpu_backends(server_args)
|
|
|
|
self.assertEqual(
|
|
resolution_result(server_args, "attention_backend"), "intel_amx"
|
|
)
|
|
self.assertEqual(resolution_result(server_args, "sampling_backend"), "pytorch")
|
|
|
|
|
|
class TestServerArgsIBDeviceValidation(unittest.TestCase):
|
|
def _validate_ib_devices(self, device_str, available_devices=None):
|
|
available_devices = available_devices or [
|
|
"mlx5_0",
|
|
"mlx5_1",
|
|
"mlx5_2",
|
|
"mlx5_3",
|
|
]
|
|
real_isdir = os.path.isdir
|
|
real_listdir = os.listdir
|
|
|
|
with patch(
|
|
"sglang.srt.arg_groups.validation_hook.os.path.isdir",
|
|
side_effect=lambda path: (
|
|
True if path == "/sys/class/infiniband" else real_isdir(path)
|
|
),
|
|
), patch(
|
|
"sglang.srt.arg_groups.validation_hook.os.listdir",
|
|
side_effect=lambda path: (
|
|
available_devices
|
|
if path == "/sys/class/infiniband"
|
|
else real_listdir(path)
|
|
),
|
|
):
|
|
return validate_ib_devices(device_str)
|
|
|
|
def test_validate_ib_devices_accepts_comma_separated(self):
|
|
self.assertEqual(
|
|
self._validate_ib_devices("mlx5_0, mlx5_1"),
|
|
"mlx5_0,mlx5_1",
|
|
)
|
|
|
|
def test_validate_ib_devices_accepts_json_object(self):
|
|
result = self._validate_ib_devices(
|
|
'{"0": "mlx5_0, mlx5_1", "1": "mlx5_2, mlx5_3"}'
|
|
)
|
|
self.assertEqual(
|
|
json.loads(result),
|
|
{"0": "mlx5_0,mlx5_1", "1": "mlx5_2,mlx5_3"},
|
|
)
|
|
|
|
def test_validate_ib_devices_accepts_json_file(self):
|
|
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as file:
|
|
json.dump({"0": "mlx5_0, mlx5_1", "1": "mlx5_2"}, file)
|
|
json_file = file.name
|
|
|
|
try:
|
|
result = self._validate_ib_devices(json_file)
|
|
finally:
|
|
os.unlink(json_file)
|
|
|
|
self.assertEqual(
|
|
json.loads(result),
|
|
{"0": "mlx5_0,mlx5_1", "1": "mlx5_2"},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|