diff --git a/tests/ut/device_allocator/test_cpu_binding.py b/tests/ut/device_allocator/test_cpu_binding.py index ea7bc6cf..b48eb515 100644 --- a/tests/ut/device_allocator/test_cpu_binding.py +++ b/tests/ut/device_allocator/test_cpu_binding.py @@ -1,12 +1,50 @@ +import subprocess import unittest -from unittest.mock import mock_open, patch +from unittest.mock import MagicMock, 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 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): mock_execute_command.side_effect = [ diff --git a/vllm_ascend/cpu_binding.py b/vllm_ascend/cpu_binding.py index abb80aa1..10967e04 100644 --- a/vllm_ascend/cpu_binding.py +++ b/vllm_ascend/cpu_binding.py @@ -37,8 +37,17 @@ def is_arm_cpu() -> bool: def execute_command(cmd: list[str]) -> tuple[str, int]: - with subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as p: - out, _ = p.communicate(timeout=1000) + # Force a C locale so subprocess output remains parseable on localized OSes. + env = os.environ.copy() + env["LC_ALL"] = "C" + env["LANG"] = "C" + env["LC_MESSAGES"] = "C" + with subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env) as p: + try: + out, _ = p.communicate(timeout=1000) + except subprocess.TimeoutExpired: + p.kill() + out, _ = p.communicate() return out.decode(), p.returncode