# Owner(s): ["module: intel"] import shutil import subprocess import tempfile import unittest from torch.testing._internal.common_utils import IS_LINUX, run_tests, TestCase @unittest.skipIf(not IS_LINUX, "Only works on linux") class TestTorchrun(TestCase): def setUp(self): self._test_dir = tempfile.mkdtemp(prefix=self.__class__.__name__) def tearDown(self): shutil.rmtree(self._test_dir) def test_cpu_info(self): lscpu_info = """# The following is the parsable format, which can be fed to other # programs. Each different item in every column has an unique ID # starting from zero. # CPU,Core,Socket,Node 0,0,0,0 1,1,0,0 2,2,0,0 3,3,0,0 4,4,1,1 5,5,1,1 6,6,1,1 7,7,1,1 8,0,0,0 9,1,0,0 10,2,0,0 11,3,0,0 12,4,1,1 13,5,1,1 14,6,1,1 15,7,1,1 """ from torch.backends.xeon.run_cpu import _CPUinfo cpuinfo = _CPUinfo(lscpu_info) assert cpuinfo._physical_core_nums() == 8 assert cpuinfo._logical_core_nums() == 16 assert cpuinfo.get_node_physical_cores(0) == [0, 1, 2, 3] assert cpuinfo.get_node_physical_cores(1) == [4, 5, 6, 7] assert cpuinfo.get_node_logical_cores(0) == [0, 1, 2, 3, 8, 9, 10, 11] assert cpuinfo.get_node_logical_cores(1) == [4, 5, 6, 7, 12, 13, 14, 15] assert cpuinfo.get_all_physical_cores() == [0, 1, 2, 3, 4, 5, 6, 7] assert cpuinfo.get_all_logical_cores() == [ 0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 6, 7, 12, 13, 14, 15, ] assert cpuinfo.numa_aware_check([0, 1, 2, 3]) == [0] assert cpuinfo.numa_aware_check([4, 5, 6, 7]) == [1] assert cpuinfo.numa_aware_check([2, 3, 4, 5]) == [0, 1] def test_multi_threads(self): num = 0 with subprocess.Popen( f"python -m torch.backends.xeon.run_cpu --ninstances 4 --use-default-allocator \ --disable-iomp --disable-numactl --disable-taskset --log-path {self._test_dir} --no-python pwd", shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, ) as p: for line in p.stdout.readlines(): segs = str(line, "utf-8").strip().split("-") if segs[-1].strip() == "pwd": num += 1 assert num == 4, "Failed to launch multiple instances for inference" if __name__ == "__main__": run_tests()