xref: /aosp_15_r20/external/pytorch/test/backends/xeon/test_launch.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# Owner(s): ["module: intel"]
2
3import shutil
4import subprocess
5import tempfile
6import unittest
7
8from torch.testing._internal.common_utils import IS_LINUX, run_tests, TestCase
9
10
11@unittest.skipIf(not IS_LINUX, "Only works on linux")
12class TestTorchrun(TestCase):
13    def setUp(self):
14        self._test_dir = tempfile.mkdtemp(prefix=self.__class__.__name__)
15
16    def tearDown(self):
17        shutil.rmtree(self._test_dir)
18
19    def test_cpu_info(self):
20        lscpu_info = """# The following is the parsable format, which can be fed to other
21# programs. Each different item in every column has an unique ID
22# starting from zero.
23# CPU,Core,Socket,Node
240,0,0,0
251,1,0,0
262,2,0,0
273,3,0,0
284,4,1,1
295,5,1,1
306,6,1,1
317,7,1,1
328,0,0,0
339,1,0,0
3410,2,0,0
3511,3,0,0
3612,4,1,1
3713,5,1,1
3814,6,1,1
3915,7,1,1
40"""
41        from torch.backends.xeon.run_cpu import _CPUinfo
42
43        cpuinfo = _CPUinfo(lscpu_info)
44        assert cpuinfo._physical_core_nums() == 8
45        assert cpuinfo._logical_core_nums() == 16
46        assert cpuinfo.get_node_physical_cores(0) == [0, 1, 2, 3]
47        assert cpuinfo.get_node_physical_cores(1) == [4, 5, 6, 7]
48        assert cpuinfo.get_node_logical_cores(0) == [0, 1, 2, 3, 8, 9, 10, 11]
49        assert cpuinfo.get_node_logical_cores(1) == [4, 5, 6, 7, 12, 13, 14, 15]
50        assert cpuinfo.get_all_physical_cores() == [0, 1, 2, 3, 4, 5, 6, 7]
51        assert cpuinfo.get_all_logical_cores() == [
52            0,
53            1,
54            2,
55            3,
56            8,
57            9,
58            10,
59            11,
60            4,
61            5,
62            6,
63            7,
64            12,
65            13,
66            14,
67            15,
68        ]
69        assert cpuinfo.numa_aware_check([0, 1, 2, 3]) == [0]
70        assert cpuinfo.numa_aware_check([4, 5, 6, 7]) == [1]
71        assert cpuinfo.numa_aware_check([2, 3, 4, 5]) == [0, 1]
72
73    def test_multi_threads(self):
74        num = 0
75        with subprocess.Popen(
76            f"python -m torch.backends.xeon.run_cpu --ninstances 4 --use-default-allocator \
77            --disable-iomp --disable-numactl --disable-taskset --log-path {self._test_dir} --no-python pwd",
78            shell=True,
79            stdout=subprocess.PIPE,
80            stderr=subprocess.STDOUT,
81        ) as p:
82            for line in p.stdout.readlines():
83                segs = str(line, "utf-8").strip().split("-")
84                if segs[-1].strip() == "pwd":
85                    num += 1
86        assert num == 4, "Failed to launch multiple instances for inference"
87
88
89if __name__ == "__main__":
90    run_tests()
91