1255 lines
56 KiB
Python
1255 lines
56 KiB
Python
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
|
|
import subprocess
|
|
import unittest
|
|
from collections import defaultdict
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, call, mock_open, patch
|
|
|
|
import vllm_ascend.cpu_binding as cpu_binding_module
|
|
from vllm_ascend.cpu_binding import CpuAlloc, DeviceInfo, bind_cpus, is_arm_cpu
|
|
from vllm_ascend.utils import AscendDeviceType
|
|
|
|
|
|
def make_cpu_alloc(rank_id=0):
|
|
cpu_alloc = object.__new__(CpuAlloc)
|
|
cpu_alloc.rank_id = rank_id
|
|
cpu_alloc.device_info = SimpleNamespace(
|
|
running_npu_list=[0],
|
|
all_logic_npus=[0],
|
|
allowed_cpus=[],
|
|
npu_affinity={},
|
|
total_logic_npus=0,
|
|
)
|
|
cpu_alloc.cpu_node = {}
|
|
cpu_alloc.numa_to_cpu_map = defaultdict(list)
|
|
cpu_alloc.npu_cpu_pool = {}
|
|
cpu_alloc.assign_main = {}
|
|
cpu_alloc.assign_acl = {}
|
|
cpu_alloc.assign_rel = {}
|
|
cpu_alloc.uvb_cpu_pool = []
|
|
return cpu_alloc
|
|
|
|
|
|
class TestDeviceInfo(unittest.TestCase):
|
|
@patch("vllm_ascend.cpu_binding.subprocess.Popen")
|
|
def test_execute_command(self, mock_popen):
|
|
process = MagicMock()
|
|
process.communicate.return_value = (b"command-output", b"")
|
|
process.returncode = 7
|
|
mock_popen.return_value.__enter__.return_value = process
|
|
|
|
output, return_code = cpu_binding_module.execute_command(["dummy", "cmd"])
|
|
|
|
self.assertEqual(output, "command-output")
|
|
self.assertEqual(return_code, 7)
|
|
mock_popen.assert_called_once()
|
|
args, kwargs = mock_popen.call_args
|
|
self.assertEqual(args[0], ["dummy", "cmd"])
|
|
self.assertEqual(kwargs["shell"], False)
|
|
self.assertEqual(kwargs["stdout"], subprocess.PIPE)
|
|
self.assertEqual(kwargs["stderr"], subprocess.PIPE)
|
|
self.assertEqual(kwargs["env"]["LC_ALL"], "C")
|
|
self.assertEqual(kwargs["env"]["LANG"], "C")
|
|
self.assertEqual(kwargs["env"]["LC_MESSAGES"], "C")
|
|
|
|
@patch("vllm_ascend.cpu_binding.subprocess.Popen")
|
|
def test_execute_command_kills_timed_out_process(self, mock_popen):
|
|
process = MagicMock()
|
|
process.communicate.side_effect = [
|
|
subprocess.TimeoutExpired(cmd=["dummy", "cmd"], timeout=1000),
|
|
(b"command-output", b""),
|
|
]
|
|
process.returncode = -9
|
|
mock_popen.return_value.__enter__.return_value = process
|
|
|
|
output, return_code = cpu_binding_module.execute_command(["dummy", "cmd"])
|
|
|
|
self.assertEqual(output, "command-output")
|
|
self.assertEqual(return_code, -9)
|
|
process.kill.assert_called_once_with()
|
|
self.assertEqual(process.communicate.call_count, 2)
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def setUp(self, mock_execute_command):
|
|
visible_devices_patcher = patch("vllm_ascend.cpu_binding.ASCEND_RT_VISIBLE_DEVICES", None)
|
|
visible_devices_patcher.start()
|
|
self.addCleanup(visible_devices_patcher.stop)
|
|
|
|
mock_execute_command.side_effect = [
|
|
("NPU ID Chip ID Chip Logic ID Chip Name\n0 0 0 Ascend\n0 1 - Mcu\n1 0 1 Ascend", 0),
|
|
("| NPU Chip | Process id |\n| 0 0 | 1234 | vllm | 56000 |\n| 1 0 | 1235 | vllm | 56000 |", 0),
|
|
("", 0),
|
|
]
|
|
self.device_info = DeviceInfo()
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_get_npu_map_info(self, mock_execute_command):
|
|
execute_result_list = [
|
|
("NPU ID Chip ID Chip Logic ID Chip Phy-ID Chip Name\n0 0 0 0 Ascend\n0 1 1 1 Ascend\n0 2 - - Mcu", 0),
|
|
("NPU ID Chip ID Chip Logic ID Chip Name\n8 0 0 Ascend\n8 1 - Mcu\n9 0 1 Ascend", 0),
|
|
]
|
|
result_list = [
|
|
{"0": {"0": "0", "1": "1"}},
|
|
{"8": {"0": "0"}, "9": {"0": "1"}},
|
|
]
|
|
for result in execute_result_list:
|
|
mock_execute_command.return_value = result
|
|
npu_map_info = self.device_info.get_npu_map_info()
|
|
expected = result_list.pop(0)
|
|
self.assertEqual(npu_map_info, expected)
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_get_npu_map_info_without_chip_logic_id_uses_npu_id(self, mock_execute_command):
|
|
mock_execute_command.return_value = (
|
|
"NPU ID Slot ID Chip ID Chip Phy-ID Chip Name\n"
|
|
"0 0 0 0 Ascend950DT\n"
|
|
"1 1 0 1 Ascend950DT\n"
|
|
"2 2 0 2 Ascend950DT",
|
|
0,
|
|
)
|
|
|
|
self.assertEqual(
|
|
self.device_info.get_npu_map_info(),
|
|
{
|
|
"0": {"0": "0"},
|
|
"1": {"0": "1"},
|
|
"2": {"0": "2"},
|
|
},
|
|
)
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_get_running_npus(self, mock_execute_command):
|
|
mock_execute_command.side_effect = [
|
|
("| NPU Chip | Process id |\n| 0 1 | 1236 | vllm | 56000 |", 0),
|
|
("", 0),
|
|
("| NPU Chip | Process id |\n| 1 0 | 1236 | vllm | 56000 |", 0),
|
|
]
|
|
with self.assertRaises(RuntimeError):
|
|
self.device_info.get_running_npus()
|
|
with self.assertRaises(RuntimeError):
|
|
self.device_info.get_running_npus()
|
|
running_npus = self.device_info.get_running_npus()
|
|
self.assertEqual(len(running_npus), 1)
|
|
|
|
@patch("vllm_ascend.cpu_binding.ASCEND_RT_VISIBLE_DEVICES", "1,5")
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_get_running_npus_filters_invalid_rows_and_visible_devices(self, mock_execute_command):
|
|
device_info = object.__new__(DeviceInfo)
|
|
device_info.npu_map_info = {"0": {"0": "0", "1": "1"}}
|
|
mock_execute_command.return_value = (
|
|
"ignored before header\n"
|
|
"| NPU Chip | Process id |\n"
|
|
"| malformed |\n"
|
|
"| xx yy | 1001 | vllm |\n"
|
|
"| 0 0 | 1234 | vllm |\n"
|
|
"| 0 1 | 2345 | vllm |",
|
|
0,
|
|
)
|
|
|
|
self.assertEqual(device_info.get_running_npus(), [1])
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_get_running_npus_skips_non_pipe_rows_inside_process_section(self, mock_execute_command):
|
|
device_info = object.__new__(DeviceInfo)
|
|
device_info.npu_map_info = {"0": {"0": "0"}}
|
|
mock_execute_command.return_value = (
|
|
"| NPU Chip | Process id |\nseparator row\n| 0 0 | 1234 | vllm |",
|
|
0,
|
|
)
|
|
|
|
self.assertEqual(device_info.get_running_npus(), [0])
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_get_running_npus_from_npu_only_process_table(self, mock_execute_command):
|
|
device_info = object.__new__(DeviceInfo)
|
|
device_info.npu_map_info = {"0": {"0": "0"}, "1": {"0": "1"}}
|
|
mock_execute_command.return_value = (
|
|
"| NPU ID | Process id | Process name | Process memory(MB) |\n"
|
|
"| 0 | 2018733 | VLLMWorker_TP | 30212 |\n"
|
|
"| 1 | 2018734 | VLLMWorker_TP | 30212 |",
|
|
0,
|
|
)
|
|
|
|
self.assertEqual(device_info.get_running_npus(), [0, 1])
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_get_running_npus_from_npu_chip_process_table_with_extra_spaces(self, mock_execute_command):
|
|
device_info = object.__new__(DeviceInfo)
|
|
device_info.npu_map_info = {"0": {"0": "0"}, "1": {"0": "1"}}
|
|
mock_execute_command.return_value = (
|
|
"| NPU Chip | Process id | Process name | Process memory(MB) |\n"
|
|
"| 0 0 | 3428811 | python3.10 | 56600 |\n"
|
|
"| 1 0 | 3428818 | python3.10 | 56474 |",
|
|
0,
|
|
)
|
|
|
|
self.assertEqual(device_info.get_running_npus(), [0, 1])
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_get_running_npus_raises_for_ambiguous_npu_only_map(self, mock_execute_command):
|
|
device_info = object.__new__(DeviceInfo)
|
|
device_info.npu_map_info = {"0": {"0": "0", "1": "1"}}
|
|
mock_execute_command.return_value = (
|
|
"| NPU ID | Process id | Process name | Process memory(MB) |\n"
|
|
"| 0 | 1234 | vllm | 56000 |",
|
|
0,
|
|
)
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
device_info.get_running_npus()
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_parse_topo_affinity(self, mock_execute_command):
|
|
mock_execute_command.return_value = ("NPU0 X HCCS HCCS HCCS HCCS HCCS HCCS HCCS 0-3", 0)
|
|
affinity = self.device_info.parse_topo_affinity()
|
|
expected = {0: [0, 1, 2, 3]}
|
|
self.assertEqual(affinity, expected)
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_parse_topo_affinity_skips_affinity_header_and_non_npu_rows(self, mock_execute_command):
|
|
device_info = object.__new__(DeviceInfo)
|
|
mock_execute_command.return_value = (
|
|
"HEADER\nNPU Chip Affinity\nnot-an-npu row\nNPU0 x x x 2-3",
|
|
0,
|
|
)
|
|
|
|
self.assertEqual(device_info.parse_topo_affinity(), {0: [2, 3]})
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_parse_topo_affinity_skips_topo_matrix_without_cpu_affinity(self, mock_execute_command):
|
|
device_info = object.__new__(DeviceInfo)
|
|
mock_execute_command.return_value = (
|
|
" NPU0 NPU1 NIC0\n"
|
|
"NPU0 X UB NA\n"
|
|
"NPU1 UB X NA\n"
|
|
"NIC0 NA NA X",
|
|
0,
|
|
)
|
|
|
|
self.assertEqual(device_info.parse_topo_affinity(), {})
|
|
|
|
def test_resolve_logic_id_from_npu_only_single_chip_map(self):
|
|
device_info = object.__new__(DeviceInfo)
|
|
device_info.npu_map_info = {"7": {"0": "7"}}
|
|
|
|
self.assertEqual(device_info.resolve_logic_id("7", None), 7)
|
|
|
|
def test_resolve_logic_id_from_chip_aware_map(self):
|
|
device_info = object.__new__(DeviceInfo)
|
|
device_info.npu_map_info = {"0": {"0": "0", "1": "1"}}
|
|
|
|
self.assertEqual(device_info.resolve_logic_id("0", "1"), 1)
|
|
|
|
def test_expand_cpu_list(self):
|
|
result = self.device_info.expand_cpu_list("0-2, 4, 6-8")
|
|
self.assertEqual(result, [0, 1, 2, 4, 6, 7, 8])
|
|
|
|
def test_get_all_logic_npus(self):
|
|
self.assertEqual(self.device_info.all_logic_npus, [0, 1])
|
|
self.assertEqual(self.device_info.total_logic_npus, 2)
|
|
|
|
def test_get_all_logic_npus_filters_invalid_values(self):
|
|
device_info = object.__new__(DeviceInfo)
|
|
device_info.npu_map_info = {
|
|
"0": {"0": "0", "1": "", "2": "abc"},
|
|
"1": {"0": "2"},
|
|
}
|
|
|
|
self.assertEqual(device_info.get_all_logic_npus(), [0, 2])
|
|
|
|
@patch("vllm_ascend.cpu_binding.os.path.exists", return_value=False)
|
|
def test_parse_allowed_cpus_returns_empty_when_status_file_missing(self, _mock_exists):
|
|
device_info = object.__new__(DeviceInfo)
|
|
|
|
self.assertEqual(device_info.parse_allowed_cpus(), [])
|
|
|
|
@patch("vllm_ascend.cpu_binding.os.path.exists", return_value=True)
|
|
@patch("builtins.open", new_callable=mock_open, read_data="Name:\tpython\nState:\tR\n")
|
|
def test_parse_allowed_cpus_raises_when_field_missing(self, _mock_open, _mock_exists):
|
|
device_info = object.__new__(DeviceInfo)
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
device_info.parse_allowed_cpus()
|
|
|
|
|
|
class TestCpuAlloc(unittest.TestCase):
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def setUp(self, mock_execute_command):
|
|
visible_devices_patcher = patch("vllm_ascend.cpu_binding.ASCEND_RT_VISIBLE_DEVICES", None)
|
|
visible_devices_patcher.start()
|
|
self.addCleanup(visible_devices_patcher.stop)
|
|
|
|
device_type_patcher = patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
device_type_patcher.start()
|
|
self.addCleanup(device_type_patcher.stop)
|
|
|
|
mock_execute_command.side_effect = [
|
|
("NPU ID Chip ID Chip Logic ID Chip Name\n0 0 0 Ascend\n0 1 - Mcu\n1 0 1 Ascend", 0),
|
|
("| NPU Chip | Process id |\n| 0 0 | 1234 | vllm | 56000 |\n| 1 0 | 1235 | vllm | 56000 |", 0),
|
|
("", 0),
|
|
]
|
|
self.cpu_alloc = CpuAlloc(0)
|
|
|
|
def test_average_distribute(self):
|
|
self.cpu_alloc.npu_cpu_pool = {0: [10, 11, 12, 13], 1: [10, 11, 12, 13]}
|
|
groups = {"[10, 11, 12, 13]": [0, 1]}
|
|
result = self.cpu_alloc.average_distribute(groups)
|
|
self.assertEqual(result, {0: [10, 11], 1: [12, 13]})
|
|
|
|
self.cpu_alloc.npu_cpu_pool = {
|
|
0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
|
|
1: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
|
|
2: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
|
|
}
|
|
groups = {"[0, 1, 2, 3, 4, 5]": [0, 1, 2]}
|
|
result = self.cpu_alloc.average_distribute(groups)
|
|
self.assertEqual(
|
|
result,
|
|
{
|
|
0: [0, 1, 2, 3],
|
|
1: [4, 5, 6, 7],
|
|
2: [8, 9, 10, 11, 12, 13],
|
|
},
|
|
)
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type")
|
|
def test_binding_mode_table(self, mock_get_device_type):
|
|
mock_get_device_type.return_value = AscendDeviceType.A2
|
|
self.assertEqual(self.cpu_alloc._binding_mode(), "topo_affinity")
|
|
mock_get_device_type.return_value = AscendDeviceType.A3
|
|
self.assertEqual(self.cpu_alloc._binding_mode(), "global_slice")
|
|
mock_get_device_type.return_value = AscendDeviceType.A5
|
|
self.assertEqual(self.cpu_alloc._binding_mode(), "topo_affinity")
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type")
|
|
def test_build_cpu_pools_fallback_to_global_slice(self, mock_get_device_type):
|
|
mock_get_device_type.return_value = AscendDeviceType.A2
|
|
self.cpu_alloc.device_info.npu_affinity = {}
|
|
with (
|
|
patch.object(self.cpu_alloc, "build_cpu_node_map") as mock_build_cpu_node_map,
|
|
patch.object(self.cpu_alloc, "build_global_slice_cpu_pool") as mock_build_global_slice_cpu_pool,
|
|
):
|
|
self.cpu_alloc.build_cpu_pools()
|
|
mock_build_cpu_node_map.assert_called_once()
|
|
mock_build_global_slice_cpu_pool.assert_called_once()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type")
|
|
def test_build_cpu_pools_global_slice_mode(self, mock_get_device_type):
|
|
mock_get_device_type.return_value = AscendDeviceType.A3
|
|
with (
|
|
patch.object(self.cpu_alloc, "build_cpu_node_map") as mock_build_cpu_node_map,
|
|
patch.object(self.cpu_alloc, "build_global_slice_cpu_pool") as mock_build_global_slice_cpu_pool,
|
|
):
|
|
self.cpu_alloc.build_cpu_pools()
|
|
mock_build_cpu_node_map.assert_called_once()
|
|
mock_build_global_slice_cpu_pool.assert_called_once()
|
|
|
|
def test_extend_numa(self):
|
|
result = self.cpu_alloc.extend_numa([])
|
|
self.assertEqual(result, [])
|
|
self.cpu_alloc.cpu_node = {0: 0, 1: 0, 2: 1, 3: 1}
|
|
self.cpu_alloc.numa_to_cpu_map = {0: [0, 1], 1: [2, 3]}
|
|
self.cpu_alloc.device_info.allowed_cpus = [0, 1, 2, 3]
|
|
result = self.cpu_alloc.extend_numa([0, 1])
|
|
self.assertEqual(result, [0, 1, 2, 3])
|
|
self.cpu_alloc.device_info.allowed_cpus = [0, 1, 3]
|
|
result = self.cpu_alloc.extend_numa([0, 1])
|
|
self.assertEqual(result, [0, 1, 3])
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_build_cpu_node_map(self, mock_execute_command):
|
|
mock_execute_command.return_value = ("", 0)
|
|
with self.assertRaises(RuntimeError):
|
|
self.cpu_alloc.build_cpu_node_map()
|
|
mock_execute_command.return_value = ("0 0\n1 1\n2 0\n3 1", 0)
|
|
self.cpu_alloc.build_cpu_node_map()
|
|
expected_cpu_node = {0: 0, 1: 1, 2: 0, 3: 1}
|
|
expected_numa_to_cpu_map = {0: [0, 2], 1: [1, 3]}
|
|
self.assertEqual(self.cpu_alloc.cpu_node, expected_cpu_node)
|
|
self.assertEqual(self.cpu_alloc.numa_to_cpu_map, expected_numa_to_cpu_map)
|
|
|
|
def test_build_global_slice_cpu_pool_uses_total_logic_npus(self):
|
|
self.cpu_alloc.device_info.running_npu_list = [1]
|
|
self.cpu_alloc.device_info.allowed_cpus = list(range(20))
|
|
self.cpu_alloc.device_info.total_logic_npus = 2
|
|
self.cpu_alloc.device_info.npu_affinity = {0: [0, 1], 1: [2, 3]}
|
|
|
|
self.cpu_alloc.build_global_slice_cpu_pool()
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[1], list(range(10, 20)))
|
|
|
|
def test_build_global_slice_cpu_pool_splits_same_cpuset_across_processes(self):
|
|
def build_single_card_process(visible_npu):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [visible_npu]
|
|
cpu_alloc.device_info.allowed_cpus = list(range(20))
|
|
cpu_alloc.device_info.total_logic_npus = 4
|
|
cpu_alloc.device_info.npu_affinity = {}
|
|
|
|
cpu_alloc.build_global_slice_cpu_pool()
|
|
cpu_alloc.allocate()
|
|
return cpu_alloc
|
|
|
|
npu0_process = build_single_card_process(0)
|
|
npu3_process = build_single_card_process(3)
|
|
|
|
self.assertEqual(npu0_process.npu_cpu_pool, {0: [0, 1, 2, 3, 4]})
|
|
self.assertEqual(npu3_process.npu_cpu_pool, {3: [15, 16, 17, 18, 19]})
|
|
self.assertFalse(set(npu0_process.npu_cpu_pool[0]) & set(npu3_process.npu_cpu_pool[3]))
|
|
self.assertFalse(set(npu0_process.assign_main[0]) & set(npu3_process.assign_main[3]))
|
|
|
|
def test_build_global_slice_cpu_pool_distributes_remainder_by_npu_id(self):
|
|
self.cpu_alloc.device_info.running_npu_list = [0, 1, 2]
|
|
self.cpu_alloc.device_info.allowed_cpus = list(range(17))
|
|
self.cpu_alloc.device_info.total_logic_npus = 3
|
|
self.cpu_alloc.device_info.npu_affinity = {}
|
|
|
|
self.cpu_alloc.build_global_slice_cpu_pool()
|
|
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[0], [0, 1, 2, 3, 4, 5])
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[1], [6, 7, 8, 9, 10, 11])
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[2], [12, 13, 14, 15, 16])
|
|
|
|
def test_build_global_slice_cpu_pool_fallback_to_affinity_len(self):
|
|
self.cpu_alloc.device_info.running_npu_list = [0, 1]
|
|
self.cpu_alloc.device_info.allowed_cpus = list(range(12))
|
|
self.cpu_alloc.device_info.total_logic_npus = 0
|
|
self.cpu_alloc.device_info.npu_affinity = {0: [0, 1], 1: [2, 3]}
|
|
|
|
self.cpu_alloc.build_global_slice_cpu_pool()
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[0], [0, 1, 2, 3, 4, 5])
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[1], [6, 7, 8, 9, 10, 11])
|
|
|
|
def test_build_global_slice_cpu_pool_fallback_to_running_len(self):
|
|
self.cpu_alloc.device_info.running_npu_list = [0, 1]
|
|
self.cpu_alloc.device_info.allowed_cpus = list(range(12))
|
|
self.cpu_alloc.device_info.total_logic_npus = 0
|
|
self.cpu_alloc.device_info.npu_affinity = {}
|
|
|
|
self.cpu_alloc.build_global_slice_cpu_pool()
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[0], [0, 1, 2, 3, 4, 5])
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[1], [6, 7, 8, 9, 10, 11])
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
def test_build_global_slice_cpu_pool_raises_when_cpu_insufficient(self, _mock_get_device_type):
|
|
self.cpu_alloc.device_info.running_npu_list = [0, 1]
|
|
self.cpu_alloc.device_info.allowed_cpus = list(range(8))
|
|
self.cpu_alloc.device_info.total_logic_npus = 2
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
self.cpu_alloc.build_global_slice_cpu_pool()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A5)
|
|
def test_build_global_slice_cpu_pool_allows_ascend_950_without_irq_reservation(self, _mock_get_device_type):
|
|
self.cpu_alloc.device_info.running_npu_list = [0, 1]
|
|
self.cpu_alloc.device_info.allowed_cpus = list(range(6))
|
|
self.cpu_alloc.device_info.total_logic_npus = 2
|
|
|
|
self.cpu_alloc.build_global_slice_cpu_pool()
|
|
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[0], [0, 1, 2])
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool[1], [3, 4, 5])
|
|
|
|
def test_build_global_slice_cpu_pool_raises_invalid_npu_id(self):
|
|
self.cpu_alloc.device_info.running_npu_list = [2]
|
|
self.cpu_alloc.device_info.allowed_cpus = list(range(12))
|
|
self.cpu_alloc.device_info.total_logic_npus = 2
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
self.cpu_alloc.build_global_slice_cpu_pool()
|
|
|
|
def test_build_global_slice_cpu_pool_returns_when_running_or_allowed_empty(self):
|
|
self.cpu_alloc.device_info.running_npu_list = []
|
|
self.cpu_alloc.device_info.allowed_cpus = list(range(12))
|
|
self.cpu_alloc.build_global_slice_cpu_pool()
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool, {})
|
|
|
|
self.cpu_alloc.device_info.running_npu_list = [0]
|
|
self.cpu_alloc.device_info.allowed_cpus = []
|
|
self.cpu_alloc.build_global_slice_cpu_pool()
|
|
self.assertEqual(self.cpu_alloc.npu_cpu_pool, {})
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_allocate(self, _mock_execute_command, _mock_get_device_type):
|
|
self.cpu_alloc.device_info.running_npu_list = [0]
|
|
self.cpu_alloc.npu_cpu_pool = {0: [0, 1, 2, 3, 4]}
|
|
self.cpu_alloc.allocate()
|
|
self.assertEqual(self.cpu_alloc.assign_main[0], [2])
|
|
self.assertEqual(self.cpu_alloc.assign_acl[0], [3])
|
|
self.assertEqual(self.cpu_alloc.assign_rel[0], [4])
|
|
self.cpu_alloc.npu_cpu_pool = {0: [0, 1]}
|
|
with self.assertRaises(RuntimeError):
|
|
self.cpu_alloc.allocate()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A5)
|
|
def test_allocate_ascend_950_assigns_cluster_to_main_only(self, _mock_get_device_type):
|
|
self.cpu_alloc.device_info.running_npu_list = [0]
|
|
self.cpu_alloc.npu_cpu_pool = {0: [0, 1, 2, 3, 4]}
|
|
|
|
self.cpu_alloc.allocate()
|
|
|
|
self.assertEqual(self.cpu_alloc.assign_main[0], [0, 1, 2, 3, 4])
|
|
self.assertEqual(self.cpu_alloc.assign_acl[0], [])
|
|
self.assertEqual(self.cpu_alloc.assign_rel[0], [])
|
|
|
|
self.cpu_alloc.assign_main = {}
|
|
self.cpu_alloc.assign_acl = {}
|
|
self.cpu_alloc.assign_rel = {}
|
|
self.cpu_alloc.npu_cpu_pool = {0: [0, 1]}
|
|
self.cpu_alloc.allocate()
|
|
self.assertEqual(self.cpu_alloc.assign_main[0], [0, 1])
|
|
self.assertEqual(self.cpu_alloc.assign_acl[0], [])
|
|
self.assertEqual(self.cpu_alloc.assign_rel[0], [])
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_threads(self, mock_execute_command):
|
|
thread_message = "1234 1234 ? 00:00:03 acl_thread\n4567 4567 ? 00:00:03 release_thread"
|
|
mock_execute_command.return_value = (thread_message, 0)
|
|
self.cpu_alloc.device_info.running_npu_list = [0]
|
|
self.cpu_alloc.assign_main = {0: [0, 1]}
|
|
self.cpu_alloc.assign_acl = {0: [2]}
|
|
self.cpu_alloc.assign_rel = {0: [3]}
|
|
self.cpu_alloc.bind_threads()
|
|
mock_execute_command.assert_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type")
|
|
@patch("vllm_ascend.cpu_binding.os.listdir")
|
|
@patch("builtins.open", new_callable=mock_open, read_data="123: 0 0 0 0 sq_send_trigger_irq\n")
|
|
@patch("vllm_ascend.cpu_binding.shutil.which")
|
|
@patch("vllm_ascend.cpu_binding.os.access")
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_npu_irq_a3_uses_card_chip_mapping(
|
|
self, mock_execute_command, mock_access, mock_which, _mock_open, mock_listdir, mock_get_device_type
|
|
):
|
|
mock_access.return_value = True
|
|
mock_which.return_value = None
|
|
mock_listdir.side_effect = FileNotFoundError
|
|
mock_get_device_type.return_value = AscendDeviceType.A3
|
|
mock_execute_command.return_value = ("PCIe Bus Info 0000:03:00.0", 0)
|
|
self.cpu_alloc.rank_id = 0
|
|
self.cpu_alloc.device_info.running_npu_list = [3]
|
|
self.cpu_alloc.npu_cpu_pool = {3: [0, 1, 2, 3, 4]}
|
|
|
|
self.cpu_alloc.bind_npu_irq()
|
|
|
|
mock_execute_command.assert_any_call(["npu-smi", "info", "-t", "board", "-i", "1", "-c", "1"])
|
|
|
|
|
|
class TestCpuBindingSupplemental(unittest.TestCase):
|
|
def setUp(self):
|
|
visible_devices_patcher = patch("vllm_ascend.cpu_binding.ASCEND_RT_VISIBLE_DEVICES", None)
|
|
visible_devices_patcher.start()
|
|
self.addCleanup(visible_devices_patcher.stop)
|
|
|
|
device_type_patcher = patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
device_type_patcher.start()
|
|
self.addCleanup(device_type_patcher.stop)
|
|
|
|
def test_cpu_to_mask_handles_single_and_multi_group_masks(self):
|
|
self.assertEqual(CpuAlloc.cpu_to_mask(3), "00000008")
|
|
self.assertEqual(CpuAlloc.cpu_to_mask(35), "00000008,00000000")
|
|
|
|
def test_get_threads_map_skips_irrelevant_lines(self):
|
|
thread_message = (
|
|
"bad-line\n"
|
|
"123 456 ? 00:00:01 acl_thread\n"
|
|
"123 789 ? 00:00:01 release_thread\n"
|
|
"123 999 ? 00:00:01 worker_thread\n"
|
|
"555 666 ? 00:00:01 acl_thread"
|
|
)
|
|
|
|
self.assertEqual(
|
|
CpuAlloc.get_threads_map(thread_message),
|
|
{
|
|
"123": {"acl_thread": ["456"], "release_thread": ["789"]},
|
|
"555": {"acl_thread": ["666"], "release_thread": []},
|
|
},
|
|
)
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_skips_empty_cpu_list(self, mock_execute_command):
|
|
CpuAlloc.bind("123", [], False)
|
|
|
|
mock_execute_command.assert_not_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command", return_value=("ok", 0))
|
|
def test_bind_uses_sub_thread_flag(self, mock_execute_command):
|
|
CpuAlloc.bind("123", [1, 2], True)
|
|
|
|
mock_execute_command.assert_called_once_with(["taskset", "-acp", "1,2", "123"])
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command", return_value=("failed", 1))
|
|
def test_bind_raises_for_failed_taskset(self, mock_execute_command):
|
|
with self.assertRaises(RuntimeError):
|
|
CpuAlloc.bind("123", [1, 2], False)
|
|
|
|
mock_execute_command.assert_called_once_with(["taskset", "-cp", "1,2", "123"])
|
|
|
|
def test_extend_numa_returns_original_list_when_multiple_nodes_present(self):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.cpu_node = {0: 0, 1: 1}
|
|
|
|
self.assertEqual(cpu_alloc.extend_numa([0, 1]), [0, 1])
|
|
|
|
def test_parse_threads_per_core(self):
|
|
self.assertEqual(CpuAlloc.parse_threads_per_core("CPU(s): 384\nThread(s) per core: 1\n"), 1)
|
|
self.assertEqual(CpuAlloc.parse_threads_per_core("Thread(s) per core: 2\nCore(s) per socket: 96\n"), 2)
|
|
self.assertIsNone(CpuAlloc.parse_threads_per_core("CPU(s): 384\n"))
|
|
|
|
def test_get_uvb_poll_window_threads(self):
|
|
thread_message = (
|
|
"100 101 ? 00:00:01 uvb_poll_window_thread\n"
|
|
"200 201 ? 00:00:01 worker_thread\n"
|
|
"bad-line uvb_poll_window_thread\n"
|
|
"300 301 ? 00:00:01 uvb_poll_window_thread"
|
|
)
|
|
|
|
self.assertEqual(CpuAlloc.get_uvb_poll_window_threads(thread_message), ["101", "301"])
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_uvb_poll_window_threads(self, mock_execute_command):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.numa_to_cpu_map = {0: [0, 1, 2, 3], 1: [4, 5]}
|
|
cpu_alloc.device_info.allowed_cpus = [0, 1, 2, 3, 4, 5]
|
|
mock_execute_command.return_value = (
|
|
"100 101 ? 00:00:01 uvb_poll_window_thread\n200 201 ? 00:00:01 uvb_poll_window_thread",
|
|
0,
|
|
)
|
|
|
|
with (
|
|
patch.object(cpu_alloc, "bind") as mock_bind,
|
|
patch("vllm_ascend.cpu_binding.logger.info") as mock_logger_info,
|
|
):
|
|
cpu_alloc.bind_uvb_poll_window_threads()
|
|
|
|
self.assertEqual(cpu_alloc.uvb_cpu_pool, [1, 2, 3])
|
|
self.assertEqual(
|
|
mock_bind.call_args_list,
|
|
[
|
|
call("101", [1, 2, 3], False),
|
|
call("201", [1, 2, 3], False),
|
|
],
|
|
)
|
|
mock_logger_info.assert_called_once_with(
|
|
"[cpu_bind_ascend_950] uvb_poll_window_thread tids=[%s] cpus=[%s]",
|
|
"101 201",
|
|
"1 2 3",
|
|
)
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_uvb_poll_window_threads_skips_empty_inputs(self, mock_execute_command):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.numa_to_cpu_map = {0: [0, 1]}
|
|
cpu_alloc.device_info.allowed_cpus = [0]
|
|
|
|
with patch.object(cpu_alloc, "bind") as mock_bind:
|
|
cpu_alloc.bind_uvb_poll_window_threads()
|
|
mock_execute_command.assert_not_called()
|
|
mock_bind.assert_not_called()
|
|
|
|
cpu_alloc.device_info.allowed_cpus = [0, 1]
|
|
mock_execute_command.return_value = ("100 101 ? 00:00:01 worker_thread", 0)
|
|
with patch.object(cpu_alloc, "bind") as mock_bind:
|
|
cpu_alloc.bind_uvb_poll_window_threads()
|
|
mock_bind.assert_not_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.logger.warning")
|
|
@patch("vllm_ascend.cpu_binding.execute_command", return_value=("100 101 ? 00:00:01 worker_thread", 0))
|
|
def test_bind_uvb_poll_window_threads_logs_pid_host_when_thread_missing(
|
|
self, _mock_execute_command, mock_logger_warning
|
|
):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.numa_to_cpu_map = {0: [0, 1]}
|
|
cpu_alloc.device_info.allowed_cpus = [0, 1]
|
|
|
|
cpu_alloc.bind_uvb_poll_window_threads()
|
|
|
|
self.assertIn("--pid=host", mock_logger_warning.call_args[0][0])
|
|
|
|
@patch("vllm_ascend.cpu_binding.logger.warning")
|
|
@patch(
|
|
"vllm_ascend.cpu_binding.execute_command",
|
|
return_value=("100 101 ? 00:00:01 uvb_poll_window_thread", 0),
|
|
)
|
|
def test_bind_uvb_poll_window_threads_logs_pid_host_when_bind_fails(
|
|
self, _mock_execute_command, mock_logger_warning
|
|
):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.numa_to_cpu_map = {0: [0, 1]}
|
|
cpu_alloc.device_info.allowed_cpus = [0, 1]
|
|
|
|
with patch.object(cpu_alloc, "bind", side_effect=RuntimeError("failed")):
|
|
cpu_alloc.bind_uvb_poll_window_threads()
|
|
|
|
self.assertIn("--pid=host", mock_logger_warning.call_args[0][0])
|
|
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_build_cpu_node_map_skips_blank_and_header_rows(self, mock_execute_command):
|
|
cpu_alloc = make_cpu_alloc()
|
|
mock_execute_command.return_value = ("CPU NODE\n\n0 0\n1 1", 0)
|
|
|
|
cpu_alloc.build_cpu_node_map()
|
|
|
|
self.assertEqual(cpu_alloc.cpu_node, {0: 0, 1: 1})
|
|
self.assertEqual(cpu_alloc.numa_to_cpu_map, {0: [0], 1: [1]})
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value="unknown")
|
|
def test_binding_mode_defaults_to_topo_affinity_for_unknown_device(self, _mock_get_device_type):
|
|
self.assertEqual(CpuAlloc._binding_mode(), "topo_affinity")
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
def test_build_cpu_pools_raises_on_affinity_conflict(self, _mock_get_device_type):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.device_info.allowed_cpus = [8, 9]
|
|
cpu_alloc.device_info.npu_affinity = {0: [0, 1]}
|
|
|
|
with patch.object(cpu_alloc, "build_cpu_node_map"), self.assertRaises(RuntimeError):
|
|
cpu_alloc.build_cpu_pools()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
def test_build_cpu_pools_topo_mode_builds_and_splits_duplicate_groups(self, _mock_get_device_type):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.all_logic_npus = [0, 1, 2]
|
|
cpu_alloc.device_info.running_npu_list = [0, 1, 2]
|
|
cpu_alloc.device_info.allowed_cpus = [0, 1, 2, 3]
|
|
cpu_alloc.device_info.npu_affinity = {0: [0, 1], 1: [2, 3], 2: [2, 3]}
|
|
|
|
with (
|
|
patch.object(cpu_alloc, "build_cpu_node_map"),
|
|
patch.object(cpu_alloc, "extend_numa", side_effect=lambda cpus: cpus),
|
|
):
|
|
cpu_alloc.build_cpu_pools()
|
|
|
|
self.assertEqual(cpu_alloc.npu_cpu_pool, {0: [0, 1], 1: [2], 2: [3]})
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
def test_build_cpu_pools_topo_mode_skips_non_running_npu_without_cpuset_overlap(self, _mock_get_device_type):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.all_logic_npus = [0, 1]
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.device_info.allowed_cpus = [192, 193]
|
|
cpu_alloc.device_info.npu_affinity = {0: [192, 193], 1: [128, 129]}
|
|
|
|
with (
|
|
patch.object(cpu_alloc, "build_cpu_node_map"),
|
|
patch.object(cpu_alloc, "extend_numa", side_effect=lambda cpus: cpus),
|
|
):
|
|
cpu_alloc.build_cpu_pools()
|
|
|
|
self.assertEqual(cpu_alloc.npu_cpu_pool, {0: [192, 193]})
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
def test_build_cpu_pools_topo_mode_excludes_non_running_npu_from_final_pool(self, _mock_get_device_type):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.all_logic_npus = [0, 1]
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.device_info.allowed_cpus = list(range(192, 202))
|
|
cpu_alloc.device_info.npu_affinity = {
|
|
0: list(range(192, 202)),
|
|
1: list(range(192, 202)),
|
|
}
|
|
|
|
with (
|
|
patch.object(cpu_alloc, "build_cpu_node_map"),
|
|
patch.object(cpu_alloc, "extend_numa", side_effect=lambda cpus: cpus),
|
|
):
|
|
cpu_alloc.build_cpu_pools()
|
|
cpu_alloc.allocate()
|
|
|
|
self.assertEqual(cpu_alloc.npu_cpu_pool, {0: [192, 193, 194, 195, 196]})
|
|
self.assertEqual(cpu_alloc.assign_main, {0: [194]})
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
def test_build_cpu_pools_topo_mode_splits_hidden_same_affinity_npus_across_processes(self, _mock_get_device_type):
|
|
def build_single_card_process(visible_npu):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.all_logic_npus = list(range(8))
|
|
cpu_alloc.device_info.running_npu_list = [visible_npu]
|
|
cpu_alloc.device_info.allowed_cpus = list(range(144, 192))
|
|
cpu_alloc.device_info.npu_affinity = {
|
|
0: list(range(144, 168)),
|
|
1: list(range(0, 24)),
|
|
2: list(range(144, 168)),
|
|
3: list(range(0, 24)),
|
|
4: list(range(96, 120)),
|
|
5: list(range(48, 72)),
|
|
6: list(range(96, 120)),
|
|
7: list(range(48, 72)),
|
|
}
|
|
|
|
with (
|
|
patch.object(cpu_alloc, "build_cpu_node_map"),
|
|
patch.object(cpu_alloc, "extend_numa", return_value=list(range(144, 192))),
|
|
):
|
|
cpu_alloc.build_cpu_pools()
|
|
cpu_alloc.allocate()
|
|
return cpu_alloc
|
|
|
|
npu0_process = build_single_card_process(0)
|
|
npu2_process = build_single_card_process(2)
|
|
|
|
self.assertEqual(npu0_process.npu_cpu_pool, {0: list(range(144, 168))})
|
|
self.assertEqual(npu2_process.npu_cpu_pool, {2: list(range(168, 192))})
|
|
self.assertFalse(set(npu0_process.npu_cpu_pool[0]) & set(npu2_process.npu_cpu_pool[2]))
|
|
self.assertFalse(set(npu0_process.assign_main[0]) & set(npu2_process.assign_main[2]))
|
|
self.assertFalse(set(npu0_process.assign_acl[0]) & set(npu2_process.assign_acl[2]))
|
|
self.assertFalse(set(npu0_process.assign_rel[0]) & set(npu2_process.assign_rel[2]))
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A5)
|
|
def test_build_ascend_950_cpu_pools_assigns_hidden_global_clusters(self, _mock_get_device_type):
|
|
def build_dp_process(visible_npus):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.all_logic_npus = list(range(8))
|
|
cpu_alloc.device_info.running_npu_list = visible_npus
|
|
cpu_alloc.device_info.allowed_cpus = list(range(384))
|
|
cpu_alloc.device_info.npu_affinity = {
|
|
0: list(range(288, 384)),
|
|
1: list(range(288, 384)),
|
|
2: list(range(96, 192)),
|
|
3: list(range(96, 192)),
|
|
4: list(range(96, 192)),
|
|
5: list(range(96, 192)),
|
|
6: list(range(288, 384)),
|
|
7: list(range(288, 384)),
|
|
}
|
|
cpu_alloc.cpu_node = {cpu: cpu // 96 for cpu in range(384)}
|
|
cpu_alloc.numa_to_cpu_map = {
|
|
0: list(range(0, 96)),
|
|
1: list(range(96, 192)),
|
|
2: list(range(192, 288)),
|
|
3: list(range(288, 384)),
|
|
}
|
|
|
|
with patch.object(cpu_alloc, "get_ascend_950_cluster_size", return_value=16):
|
|
self.assertTrue(cpu_alloc.build_ascend_950_cpu_pools())
|
|
return cpu_alloc
|
|
|
|
dp0_process = build_dp_process([0, 1, 2, 3])
|
|
dp1_process = build_dp_process([4, 5, 6, 7])
|
|
|
|
self.assertEqual(dp0_process.npu_cpu_pool[0], list(range(288, 304)))
|
|
self.assertEqual(dp0_process.npu_cpu_pool[1], list(range(304, 320)))
|
|
self.assertEqual(dp0_process.npu_cpu_pool[2], list(range(96, 112)))
|
|
self.assertEqual(dp0_process.npu_cpu_pool[3], list(range(112, 128)))
|
|
self.assertEqual(dp1_process.npu_cpu_pool[4], list(range(128, 144)))
|
|
self.assertEqual(dp1_process.npu_cpu_pool[5], list(range(144, 160)))
|
|
self.assertEqual(dp1_process.npu_cpu_pool[6], list(range(320, 336)))
|
|
self.assertEqual(dp1_process.npu_cpu_pool[7], list(range(336, 352)))
|
|
dp0_cpus = set().union(*(set(cpus) for cpus in dp0_process.npu_cpu_pool.values()))
|
|
dp1_cpus = set().union(*(set(cpus) for cpus in dp1_process.npu_cpu_pool.values()))
|
|
self.assertFalse(dp0_cpus & dp1_cpus)
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A5)
|
|
def test_build_ascend_950_cpu_pools_skips_invalid_inputs(self, _mock_get_device_type):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.all_logic_npus = [0]
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.device_info.allowed_cpus = list(range(32))
|
|
cpu_alloc.cpu_node = {cpu: 0 for cpu in range(32)}
|
|
cpu_alloc.numa_to_cpu_map = {0: list(range(32))}
|
|
|
|
cpu_alloc.device_info.npu_affinity = {}
|
|
self.assertFalse(cpu_alloc.build_ascend_950_cpu_pools())
|
|
|
|
cpu_alloc.device_info.npu_affinity = {0: list(range(8))}
|
|
with patch.object(cpu_alloc, "get_ascend_950_cluster_size", return_value=None):
|
|
self.assertFalse(cpu_alloc.build_ascend_950_cpu_pools())
|
|
|
|
cpu_alloc.device_info.allowed_cpus = [0, 100]
|
|
cpu_alloc.device_info.npu_affinity = {0: [0, 100]}
|
|
cpu_alloc.cpu_node[100] = 1
|
|
with patch.object(cpu_alloc, "get_ascend_950_cluster_size", return_value=8):
|
|
self.assertFalse(cpu_alloc.build_ascend_950_cpu_pools())
|
|
|
|
cpu_alloc.device_info.allowed_cpus = list(range(32))
|
|
cpu_alloc.device_info.all_logic_npus = [0, 1, 2]
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.device_info.npu_affinity = {0: list(range(8)), 1: list(range(8)), 2: list(range(8))}
|
|
with patch.object(cpu_alloc, "get_ascend_950_cluster_size", return_value=16):
|
|
self.assertFalse(cpu_alloc.build_ascend_950_cpu_pools())
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.logger.info")
|
|
def test_print_plan_handles_empty_release_assignment(self, mock_logger_info, _mock_get_device_type):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [1]
|
|
cpu_alloc.rank_id = 0
|
|
cpu_alloc.assign_main = {1: [2, 3]}
|
|
cpu_alloc.assign_acl = {1: [4]}
|
|
cpu_alloc.assign_rel = {1: []}
|
|
|
|
cpu_alloc.print_plan()
|
|
|
|
self.assertEqual(mock_logger_info.call_count, 2)
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A5)
|
|
@patch("vllm_ascend.cpu_binding.logger.info")
|
|
def test_print_plan_uses_ascend_950_worker_log(self, mock_logger_info, _mock_get_device_type):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [1]
|
|
cpu_alloc.rank_id = 0
|
|
cpu_alloc.assign_main = {1: [2, 3]}
|
|
cpu_alloc.assign_acl = {1: []}
|
|
cpu_alloc.assign_rel = {1: []}
|
|
|
|
cpu_alloc.print_plan()
|
|
|
|
self.assertEqual(
|
|
mock_logger_info.call_args_list,
|
|
[
|
|
call("The CPU allocation plan is as follows:"),
|
|
call("Ascend 950 NPU%s: worker=[%s]", 1, "2 3"),
|
|
],
|
|
)
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.logger.info")
|
|
def test_print_plan_handles_non_empty_release_assignment(self, mock_logger_info, _mock_get_device_type):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [1]
|
|
cpu_alloc.rank_id = 0
|
|
cpu_alloc.assign_main = {1: [2, 3]}
|
|
cpu_alloc.assign_acl = {1: [4]}
|
|
cpu_alloc.assign_rel = {1: [5]}
|
|
|
|
cpu_alloc.print_plan()
|
|
|
|
self.assertEqual(mock_logger_info.call_count, 2)
|
|
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value=None)
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_memory_skips_when_migratepages_missing(self, mock_execute_command, _mock_which):
|
|
cpu_alloc = make_cpu_alloc()
|
|
|
|
cpu_alloc.bind_memory("999", 0)
|
|
|
|
mock_execute_command.assert_not_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value="/usr/bin/migratepages")
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_memory_skips_when_cpu_pool_or_numa_invalid(self, mock_execute_command, _mock_which):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.numa_to_cpu_map = {0: [0], 1: [1]}
|
|
cpu_alloc.bind_memory("1000", 0)
|
|
mock_execute_command.assert_not_called()
|
|
|
|
cpu_alloc.npu_cpu_pool = {0: [8]}
|
|
cpu_alloc.cpu_node = {8: 3}
|
|
cpu_alloc.bind_memory("1000", 0)
|
|
mock_execute_command.assert_not_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value="/usr/bin/migratepages")
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_memory_executes_on_valid_numa_target(self, mock_execute_command, _mock_which):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.npu_cpu_pool = {0: [8, 9]}
|
|
cpu_alloc.cpu_node = {8: 1}
|
|
cpu_alloc.numa_to_cpu_map = {0: [0], 1: [8, 9]}
|
|
|
|
cpu_alloc.bind_memory("1000", 0)
|
|
|
|
mock_execute_command.assert_called_once_with(["migratepages", "1000", "0,1", "1"])
|
|
|
|
@patch("vllm_ascend.cpu_binding.psutil.Process")
|
|
@patch(
|
|
"vllm_ascend.cpu_binding.execute_command",
|
|
return_value=(
|
|
"1000 2000 ? 00:00:01 acl_thread\n1000 3000 ? 00:00:01 release_thread",
|
|
0,
|
|
),
|
|
)
|
|
def test_bind_threads_binds_main_acl_and_release_threads(self, _mock_execute_command, mock_process):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.assign_main = {0: [1, 2]}
|
|
cpu_alloc.assign_acl = {0: [3]}
|
|
cpu_alloc.assign_rel = {0: [4]}
|
|
mock_process.return_value.pid = 1000
|
|
|
|
with patch.object(cpu_alloc, "bind") as mock_bind, patch.object(cpu_alloc, "bind_memory") as mock_bind_memory:
|
|
cpu_alloc.bind_threads()
|
|
|
|
self.assertEqual(
|
|
mock_bind.call_args_list,
|
|
[
|
|
call("1000", [1, 2], True),
|
|
call("2000", [3], False),
|
|
call("3000", [4], False),
|
|
],
|
|
)
|
|
mock_bind_memory.assert_called_once_with("1000", 0)
|
|
|
|
@patch("vllm_ascend.cpu_binding.psutil.Process")
|
|
def test_bind_ascend_950_threads_binds_only_main_and_memory(self, mock_process):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.assign_main = {0: [1, 2, 3]}
|
|
mock_process.return_value.pid = 1000
|
|
|
|
with patch.object(cpu_alloc, "bind") as mock_bind, patch.object(cpu_alloc, "bind_memory") as mock_bind_memory:
|
|
cpu_alloc.bind_ascend_950_threads()
|
|
|
|
mock_bind.assert_called_once_with("1000", [1, 2, 3], True)
|
|
mock_bind_memory.assert_called_once_with("1000", 0)
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.os.access", return_value=False)
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_npu_irq_returns_when_irq_path_not_writable(
|
|
self, mock_execute_command, _mock_access, _mock_get_device_type
|
|
):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
mock_execute_command.assert_not_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A5)
|
|
@patch("vllm_ascend.cpu_binding.os.access")
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_npu_irq_skips_on_ascend_950(self, mock_execute_command, mock_access, _mock_get_device_type):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.npu_cpu_pool = {0: [8, 9, 10]}
|
|
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
mock_access.assert_not_called()
|
|
mock_execute_command.assert_not_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.os.access", return_value=True)
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_npu_irq_returns_when_current_npu_has_no_cpu_pool(
|
|
self, mock_execute_command, _mock_access, _mock_get_device_type
|
|
):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.npu_cpu_pool = {}
|
|
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
mock_execute_command.assert_not_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("builtins.open", new_callable=mock_open, read_data="123: 0 0 0 sq_send_trigger_irq\n")
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value=None)
|
|
@patch("vllm_ascend.cpu_binding.os.access", return_value=True)
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_npu_irq_skips_when_cpu_pool_too_small(
|
|
self, mock_execute_command, _mock_access, _mock_which, _mock_open, _mock_get_device_type
|
|
):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.npu_cpu_pool = {0: [7]}
|
|
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
mock_execute_command.assert_not_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("builtins.open", new_callable=mock_open, read_data="123: 0 0 0 sq_send_trigger_irq\n")
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value=None)
|
|
@patch("vllm_ascend.cpu_binding.os.access", return_value=True)
|
|
@patch("vllm_ascend.cpu_binding.execute_command", return_value=("board info without pci", 0))
|
|
def test_bind_npu_irq_skips_when_pci_address_missing(
|
|
self, mock_execute_command, _mock_access, _mock_which, _mock_open, _mock_get_device_type
|
|
):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.npu_cpu_pool = {0: [7, 8]}
|
|
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
mock_execute_command.assert_called_once_with(["npu-smi", "info", "-t", "board", "-i", "0"])
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.os.listdir", return_value=["456", "457"])
|
|
@patch("builtins.open", new_callable=mock_open, read_data="123: 0 0 0 sq_send_trigger_irq\n")
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value=None)
|
|
@patch("vllm_ascend.cpu_binding.os.access", return_value=True)
|
|
@patch("vllm_ascend.cpu_binding.execute_command", return_value=("prefix\nPCIe Bus Info 0000:03:00.0", 0))
|
|
def test_bind_npu_irq_skips_when_sq_irq_not_found(
|
|
self, _mock_execute_command, _mock_access, _mock_which, _mock_open, _mock_listdir, _mock_get_device_type
|
|
):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.npu_cpu_pool = {0: [7, 8]}
|
|
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.os.listdir", return_value=["123", "124"])
|
|
@patch("builtins.open", new_callable=mock_open, read_data="123: 0 0 0 sq_send_trigger_irq\n")
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value="/bin/systemctl")
|
|
@patch("vllm_ascend.cpu_binding.os.access", return_value=True)
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_npu_irq_stops_irqbalance_and_writes_affinity_masks(
|
|
self, mock_execute_command, _mock_access, _mock_which, mock_file, _mock_listdir, _mock_get_device_type
|
|
):
|
|
mock_execute_command.side_effect = [
|
|
("irqbalance.service enabled\n", 0),
|
|
("", 0),
|
|
("stopped", 0),
|
|
("prefix\nPCIe Bus Info 0000:03:00.0", 0),
|
|
]
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.npu_cpu_pool = {0: [8, 9, 10]}
|
|
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
self.assertIn(call(["systemctl", "stop", "irqbalance"]), mock_execute_command.call_args_list)
|
|
self.assertIn(call(["npu-smi", "info", "-t", "board", "-i", "0"]), mock_execute_command.call_args_list)
|
|
handle = mock_file()
|
|
self.assertEqual(handle.write.call_args_list, [call("00000100"), call("00000200")])
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.os.listdir", return_value=["123", "124"])
|
|
@patch("builtins.open", new_callable=mock_open, read_data="123: 0 0 0 sq_send_trigger_irq\n")
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value="/bin/systemctl")
|
|
@patch("vllm_ascend.cpu_binding.os.access", return_value=True)
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_npu_irq_keeps_irqbalance_when_inactive(
|
|
self, mock_execute_command, _mock_access, _mock_which, _mock_open, _mock_listdir, _mock_get_device_type
|
|
):
|
|
mock_execute_command.side_effect = [
|
|
("irqbalance.service enabled\n", 0),
|
|
("", 3),
|
|
("prefix\nPCIe Bus Info 0000:03:00.0", 0),
|
|
]
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.npu_cpu_pool = {0: [8, 9, 10]}
|
|
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
self.assertNotIn(call(["systemctl", "stop", "irqbalance"]), mock_execute_command.call_args_list)
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.os.listdir", return_value=["123", "124"])
|
|
@patch("builtins.open", new_callable=mock_open, read_data="123: 0 0 0 sq_send_trigger_irq\n")
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value="/bin/systemctl")
|
|
@patch("vllm_ascend.cpu_binding.os.access", return_value=True)
|
|
@patch("vllm_ascend.cpu_binding.execute_command")
|
|
def test_bind_npu_irq_skips_irqbalance_handling_when_service_absent(
|
|
self, mock_execute_command, _mock_access, _mock_which, _mock_open, _mock_listdir, _mock_get_device_type
|
|
):
|
|
mock_execute_command.side_effect = [
|
|
("another.service enabled\n", 0),
|
|
("prefix\nPCIe Bus Info 0000:03:00.0", 0),
|
|
]
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.npu_cpu_pool = {0: [8, 9, 10]}
|
|
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
self.assertNotIn(call(["systemctl", "is-active", "--quiet", "irqbalance"]), mock_execute_command.call_args_list)
|
|
|
|
@patch("vllm_ascend.cpu_binding.get_ascend_device_type", return_value=AscendDeviceType.A2)
|
|
@patch("vllm_ascend.cpu_binding.os.listdir", return_value=["123", "124"])
|
|
@patch(
|
|
"builtins.open",
|
|
new_callable=mock_open,
|
|
read_data="100: 0 0 0 other_irq\n123: 0 0 0 sq_send_trigger_irq\n",
|
|
)
|
|
@patch("vllm_ascend.cpu_binding.shutil.which", return_value=None)
|
|
@patch("vllm_ascend.cpu_binding.os.access", return_value=True)
|
|
@patch("vllm_ascend.cpu_binding.execute_command", return_value=("prefix\nPCIe Bus Info 0000:03:00.0", 0))
|
|
def test_bind_npu_irq_scans_multiple_interrupt_lines(
|
|
self, _mock_execute_command, _mock_access, _mock_which, mock_file, _mock_listdir, _mock_get_device_type
|
|
):
|
|
cpu_alloc = make_cpu_alloc()
|
|
cpu_alloc.device_info.running_npu_list = [0]
|
|
cpu_alloc.npu_cpu_pool = {0: [8, 9, 10]}
|
|
|
|
cpu_alloc.bind_npu_irq()
|
|
|
|
handle = mock_file()
|
|
self.assertEqual(handle.write.call_args_list, [call("00000100"), call("00000200")])
|
|
|
|
def test_run_all_invokes_steps_in_order(self):
|
|
cpu_alloc = make_cpu_alloc()
|
|
calls = []
|
|
|
|
def build_cpu_pools() -> bool:
|
|
calls.append("build_cpu_pools")
|
|
return True
|
|
|
|
with (
|
|
patch.object(cpu_alloc, "build_cpu_pools", side_effect=build_cpu_pools),
|
|
patch.object(cpu_alloc, "allocate", side_effect=lambda: calls.append("allocate")),
|
|
patch.object(cpu_alloc, "print_plan", side_effect=lambda: calls.append("print_plan")),
|
|
patch.object(cpu_alloc, "bind_threads", side_effect=lambda: calls.append("bind_threads")),
|
|
patch.object(cpu_alloc, "bind_npu_irq", side_effect=lambda: calls.append("bind_npu_irq")),
|
|
):
|
|
cpu_alloc.run_all()
|
|
|
|
self.assertEqual(calls, ["build_cpu_pools", "allocate", "print_plan", "bind_threads", "bind_npu_irq"])
|
|
|
|
def test_run_all_returns_when_cpu_pool_build_is_skipped(self):
|
|
cpu_alloc = make_cpu_alloc()
|
|
calls = []
|
|
|
|
def build_cpu_pools() -> bool:
|
|
calls.append("build_cpu_pools")
|
|
return False
|
|
|
|
with (
|
|
patch.object(cpu_alloc, "build_cpu_pools", side_effect=build_cpu_pools),
|
|
patch.object(cpu_alloc, "allocate", side_effect=lambda: calls.append("allocate")),
|
|
patch.object(cpu_alloc, "print_plan", side_effect=lambda: calls.append("print_plan")),
|
|
patch.object(cpu_alloc, "bind_threads", side_effect=lambda: calls.append("bind_threads")),
|
|
patch.object(cpu_alloc, "bind_npu_irq", side_effect=lambda: calls.append("bind_npu_irq")),
|
|
):
|
|
cpu_alloc.run_all()
|
|
|
|
self.assertEqual(calls, ["build_cpu_pools"])
|
|
|
|
|
|
class TestBindingSwitch(unittest.TestCase):
|
|
@patch("vllm_ascend.cpu_binding.platform.machine")
|
|
def test_is_arm_cpu(self, mock_machine):
|
|
mock_machine.return_value = "x86_64"
|
|
self.assertFalse(is_arm_cpu())
|
|
mock_machine.return_value = "aarch64"
|
|
self.assertTrue(is_arm_cpu())
|
|
mock_machine.return_value = "armv8"
|
|
self.assertTrue(is_arm_cpu())
|
|
mock_machine.return_value = "mips64"
|
|
self.assertFalse(is_arm_cpu())
|
|
|
|
@patch("vllm_ascend.cpu_binding.CpuAlloc")
|
|
@patch("vllm_ascend.cpu_binding.is_arm_cpu")
|
|
def test_bind_cpus_skip_non_arm(self, mock_is_arm_cpu, mock_cpu_alloc):
|
|
mock_is_arm_cpu.return_value = False
|
|
bind_cpus(0)
|
|
mock_cpu_alloc.assert_not_called()
|
|
|
|
@patch("vllm_ascend.cpu_binding.CpuAlloc")
|
|
@patch("vllm_ascend.cpu_binding.is_arm_cpu", return_value=True)
|
|
def test_bind_cpus_runs_allocator_on_arm(self, _mock_is_arm_cpu, mock_cpu_alloc):
|
|
bind_cpus(1)
|
|
|
|
mock_cpu_alloc.assert_called_once_with(1)
|
|
mock_cpu_alloc.return_value.run_all.assert_called_once_with()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|