intel-extension-for-pytorch

1
import unittest
2
from common_utils import TestCase
3
from utils.cpuinfo import construct_numa_config
4
from intel_extension_for_pytorch.cpu.launch import (
5
    CPUPoolList,
6
    Launcher,
7
    DistributedTrainingLauncher,
8
)
9
import os
10
from os.path import expanduser
11
import glob
12
import subprocess
13

14

15
class TestLauncher(TestCase):
16
    launch_scripts = [
17
        ["python", "-m", "intel_extension_for_pytorch.cpu.launch"],
18
        ["ipexrun"],
19
    ]
20

21
    def find_lib(self, lib_type):
22
        library_paths = []
23
        if "CONDA_PREFIX" in os.environ:
24
            library_paths.append(f'{os.environ["CONDA_PREFIX"]}/lib/')
25
        elif "VIRTUAL_ENV" in os.environ:
26
            library_paths.append(f'{os.environ["VIRTUAL_ENV"]}/lib/')
27

28
        library_paths += [
29
            f'{expanduser("~")}/.local/lib/',
30
            "/usr/local/lib/",
31
            "/usr/local/lib64/",
32
            "/usr/lib/",
33
            "/usr/lib64/",
34
        ]
35
        lib_find = False
36
        for lib_path in library_paths:
37
            library_file = f"{lib_path}/lib{lib_type}.so"
38
            matches = glob.glob(library_file)
39
            if len(matches) > 0:
40
                lib_find = True
41
                break
42
        return lib_find
43

44
    def del_env(self, env_name):
45
        if env_name in os.environ:
46
            del os.environ[env_name]
47

48
    def test_memory_allocator_setup(self):
49
        launcher = Launcher()
50

51
        # tcmalloc
52
        find_tcmalloc = self.find_lib("tcmalloc")
53
        launcher.set_memory_allocator(memory_allocator="tcmalloc")
54
        ld_preload = (
55
            ":".join(launcher.ld_preload) if len(launcher.ld_preload) > 0 else ""
56
        )
57
        tcmalloc_enabled = "libtcmalloc.so" in ld_preload
58
        self.assertEqual(find_tcmalloc, tcmalloc_enabled)
59

60
        # jemalloc
61
        find_jemalloc = self.find_lib("jemalloc")
62
        launcher.set_memory_allocator(memory_allocator="jemalloc")
63
        ld_preload = (
64
            ":".join(launcher.ld_preload) if len(launcher.ld_preload) > 0 else ""
65
        )
66
        jemalloc_enabled = "libjemalloc.so" in ld_preload
67
        self.assertEqual(find_jemalloc, jemalloc_enabled)
68
        if jemalloc_enabled:
69
            self.assertTrue("MALLOC_CONF" in launcher.environ_set)
70
            self.assertTrue(
71
                launcher.environ_set["MALLOC_CONF"]
72
                == "oversize_threshold:1,background_thread:true,metadata_thp:auto"
73
            )
74

75
        self.del_env("MALLOC_CONF")
76
        launcher.set_memory_allocator(memory_allocator="jemalloc", benchmark=True)
77
        if jemalloc_enabled:
78
            self.assertTrue("MALLOC_CONF" in launcher.environ_set)
79
            self.assertTrue(
80
                launcher.environ_set["MALLOC_CONF"]
81
                == "oversize_threshold:1,background_thread:false,metadata_thp:always,dirty_decay_ms:-1,muzzy_decay_ms:-1"
82
            )
83

84
    def test_mpi_pin_domain_and_ccl_worker_affinity(self):
85
        # HT ON, use_logical_cores ON
86
        nprocs_per_node = 2
87
        ccl_worker_count = 4
88
        lscpu_txt = construct_numa_config(
89
            nprocs_per_node, 28, enable_ht=True, numa_mode=1
90
        )
91
        launcher = DistributedTrainingLauncher(lscpu_txt=lscpu_txt)
92

93
        launcher.cpuinfo.gen_pools_ondemand(
94
            ninstances=nprocs_per_node, use_logical_cores=True
95
        )
96
        pin_domain_affinity = launcher.get_pin_domain_affinity(
97
            launcher.cpuinfo.pools_ondemand, ccl_worker_count
98
        )
99
        expect_pin_domain = "[0xffffff0,0xffffff00000000]"
100
        self.assertEqual(pin_domain_affinity["pin_domain"], expect_pin_domain)
101
        expected_ccl_worker_affinity = "0,1,2,3,28,29,30,31"
102
        self.assertEqual(pin_domain_affinity["affinity"], expected_ccl_worker_affinity)
103

104
        # HT ON, use_logical_cores OFF
105
        nprocs_per_node = 2
106
        ccl_worker_count = 4
107
        lscpu_txt = construct_numa_config(
108
            nprocs_per_node, 28, enable_ht=True, numa_mode=1
109
        )
110
        launcher = DistributedTrainingLauncher(lscpu_txt=lscpu_txt)
111

112
        launcher.cpuinfo.gen_pools_ondemand(
113
            ninstances=nprocs_per_node, use_logical_cores=True
114
        )
115
        pin_domain_affinity = launcher.get_pin_domain_affinity(
116
            launcher.cpuinfo.pools_ondemand,
117
            ccl_worker_count,
118
            logical_cores_for_ccl=True,
119
        )
120
        expect_pin_domain = "[0xfffffff,0xfffffff0000000]"
121
        self.assertEqual(pin_domain_affinity["pin_domain"], expect_pin_domain)
122
        expected_ccl_worker_affinity = "56,57,58,59,84,85,86,87"
123
        self.assertEqual(pin_domain_affinity["affinity"], expected_ccl_worker_affinity)
124

125
        # HT OFF, use_logical_cores ON
126
        nprocs_per_node = 2
127
        ccl_worker_count = 4
128
        lscpu_txt = construct_numa_config(
129
            nprocs_per_node, 28, enable_ht=False, numa_mode=1
130
        )
131
        launcher = DistributedTrainingLauncher(lscpu_txt=lscpu_txt)
132

133
        launcher.cpuinfo.gen_pools_ondemand(
134
            ninstances=nprocs_per_node, use_logical_cores=True
135
        )
136
        pin_domain_affinity = launcher.get_pin_domain_affinity(
137
            launcher.cpuinfo.pools_ondemand,
138
            ccl_worker_count,
139
            logical_cores_for_ccl=True,
140
        )
141
        expect_pin_domain = "[0xffffff0,0xffffff00000000]"
142
        self.assertEqual(pin_domain_affinity["pin_domain"], expect_pin_domain)
143
        expected_ccl_worker_affinity = "0,1,2,3,28,29,30,31"
144
        self.assertEqual(pin_domain_affinity["affinity"], expected_ccl_worker_affinity)
145

146
        # nodes_list
147
        nprocs_per_node = 2
148
        ccl_worker_count = 2
149
        lscpu_txt = construct_numa_config(4, 14, enable_ht=True, numa_mode=1)
150
        launcher = DistributedTrainingLauncher(lscpu_txt=lscpu_txt)
151

152
        launcher.cpuinfo.gen_pools_ondemand(
153
            ninstances=nprocs_per_node, nodes_list=[1, 2], use_logical_cores=True
154
        )
155
        pin_domain_affinity = launcher.get_pin_domain_affinity(
156
            launcher.cpuinfo.pools_ondemand, ccl_worker_count
157
        )
158
        expect_pin_domain = "[0xfff0000,0x3ffc0000000]"
159
        self.assertEqual(pin_domain_affinity["pin_domain"], expect_pin_domain)
160
        expected_ccl_worker_affinity = "14,15,28,29"
161
        self.assertEqual(pin_domain_affinity["affinity"], expected_ccl_worker_affinity)
162

163
        # ncores_per_instance
164
        nprocs_per_node = 2
165
        ccl_worker_count = 4
166
        lscpu_txt = construct_numa_config(
167
            nprocs_per_node, 28, enable_ht=True, numa_mode=1
168
        )
169
        launcher = DistributedTrainingLauncher(lscpu_txt=lscpu_txt)
170

171
        launcher.cpuinfo.gen_pools_ondemand(
172
            ninstances=nprocs_per_node,
173
            ncores_per_instance=(8 + ccl_worker_count) * nprocs_per_node,
174
            use_logical_cores=True,
175
        )
176
        pin_domain_affinity = launcher.get_pin_domain_affinity(
177
            launcher.cpuinfo.pools_ondemand, ccl_worker_count
178
        )
179
        expect_pin_domain = "[0xff0,0xff0000]"
180
        self.assertEqual(pin_domain_affinity["pin_domain"], expect_pin_domain)
181
        expected_ccl_worker_affinity = "0,1,2,3,12,13,14,15"
182
        self.assertEqual(pin_domain_affinity["affinity"], expected_ccl_worker_affinity)
183

184
        # e-cores
185
        nprocs_per_node = 2
186
        ccl_worker_count = 4
187
        lscpu_txt = construct_numa_config(
188
            nprocs_per_node, 28, enable_ht=True, n_e_cores=4, numa_mode=0
189
        )
190
        launcher = DistributedTrainingLauncher(lscpu_txt=lscpu_txt)
191

192
        launcher.cpuinfo.gen_pools_ondemand(
193
            ninstances=nprocs_per_node, use_logical_cores=True
194
        )
195
        pin_domain_affinity = launcher.get_pin_domain_affinity(
196
            launcher.cpuinfo.pools_ondemand,
197
            ccl_worker_count,
198
            logical_cores_for_ccl=True,
199
        )
200
        expect_pin_domain = "[0xfffffff,0xfffffff000000000000000]"
201
        self.assertEqual(pin_domain_affinity["pin_domain"], expect_pin_domain)
202
        expected_ccl_worker_affinity = "28,29,30,31,88,89,90,91"
203
        self.assertEqual(pin_domain_affinity["affinity"], expected_ccl_worker_affinity)
204

205
    def test_launcher_scripts(self):
206
        for launch_script in self.launch_scripts:
207
            cmd = launch_script + ["--help"]
208
            r = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
209
            self.assertEqual(r.returncode, 0)
210

211
    def verify_affinity(self, pools, ground_truth):
212
        self.assertEqual(len(pools), ground_truth["ninstances"])
213
        self.assertEqual(len(pools[0]), ground_truth["ncores_per_instance"])
214
        self.assertEqual(
215
            len(set([c.cpu for p in pools for c in p])), ground_truth["num_cores_sum"]
216
        )
217
        self.assertEqual(
218
            len(set([c.node for p in pools for c in p])), ground_truth["num_nodes_sum"]
219
        )
220
        for i in range(ground_truth["ninstances"]):
221
            self.assertEqual(
222
                len(set([c.cpu for c in pools[i]])), ground_truth["num_cores"][i]
223
            )
224
            self.assertEqual(
225
                len(set([c.node for c in pools[i]])), ground_truth["num_nodes"][i]
226
            )
227
            pool_txt = pools[i].get_pool_txt()
228
            self.assertEqual(pool_txt["cores"], ground_truth["pools_cores"][i])
229
            self.assertEqual(pool_txt["nodes"], ground_truth["pools_nodes"][i])
230

231
    def test_core_affinity(self):
232
        # mode 0
233
        num_nodes = 2
234
        n_phycores_per_node = 28
235
        lscpu_txt = construct_numa_config(
236
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=0
237
        )
238
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
239
        ground_truth = {
240
            "ninstances": 1,
241
            "ncores_per_instance": 112,
242
            "num_cores_sum": 112,
243
            "num_nodes_sum": 2,
244
            "num_cores": [112],
245
            "num_nodes": [2],
246
            "pools_cores": ["0-111"],
247
            "pools_nodes": ["0,1"],
248
        }
249
        self.verify_affinity([cpuinfo.pool_all], ground_truth)
250

251
        cpuinfo.gen_pools_ondemand(ninstances=2)
252
        ground_truth = {
253
            "ninstances": 2,
254
            "ncores_per_instance": 28,
255
            "num_cores_sum": 56,
256
            "num_nodes_sum": 2,
257
            "num_cores": [28, 28],
258
            "num_nodes": [1, 1],
259
            "pools_cores": ["0-27", "56-83"],
260
            "pools_nodes": ["0", "1"],
261
        }
262
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
263

264
        cpuinfo.gen_pools_ondemand(ninstances=4)
265
        ground_truth = {
266
            "ninstances": 4,
267
            "ncores_per_instance": 14,
268
            "num_cores_sum": 56,
269
            "num_nodes_sum": 2,
270
            "num_cores": [14, 14, 14, 14],
271
            "num_nodes": [1, 1, 1, 1],
272
            "pools_cores": ["0-13", "14-27", "56-69", "70-83"],
273
            "pools_nodes": ["0", "0", "1", "1"],
274
        }
275
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
276

277
        cpuinfo.gen_pools_ondemand(ncores_per_instance=28)
278
        ground_truth = {
279
            "ninstances": 2,
280
            "ncores_per_instance": 28,
281
            "num_cores_sum": 56,
282
            "num_nodes_sum": 2,
283
            "num_cores": [28, 28],
284
            "num_nodes": [1, 1],
285
            "pools_cores": ["0-27", "56-83"],
286
            "pools_nodes": ["0", "1"],
287
        }
288
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
289

290
        cpuinfo.gen_pools_ondemand(ncores_per_instance=14)
291
        ground_truth = {
292
            "ninstances": 4,
293
            "ncores_per_instance": 14,
294
            "num_cores_sum": 56,
295
            "num_nodes_sum": 2,
296
            "num_cores": [14, 14, 14, 14],
297
            "num_nodes": [1, 1, 1, 1],
298
            "pools_cores": ["0-13", "14-27", "56-69", "70-83"],
299
            "pools_nodes": ["0", "0", "1", "1"],
300
        }
301
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
302

303
        cores_list_local = []
304
        cores_list_local.extend(list(i for i in range(14, 28)))
305
        cores_list_local.extend(list(i for i in range(42, 56)))
306
        cpuinfo.gen_pools_ondemand(cores_list=cores_list_local)
307
        ground_truth = {
308
            "ninstances": 1,
309
            "ncores_per_instance": 28,
310
            "num_cores_sum": 28,
311
            "num_nodes_sum": 1,
312
            "num_cores": [28],
313
            "num_nodes": [1],
314
            "pools_cores": ["14-27,42-55"],
315
            "pools_nodes": ["0"],
316
        }
317
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
318

319
        num_nodes = 4
320
        n_phycores_per_node = 14
321
        lscpu_txt = construct_numa_config(
322
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=0
323
        )
324
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
325
        ground_truth = {
326
            "ninstances": 1,
327
            "ncores_per_instance": 112,
328
            "num_cores_sum": 112,
329
            "num_nodes_sum": 4,
330
            "num_cores": [112],
331
            "num_nodes": [4],
332
            "pools_cores": ["0-111"],
333
            "pools_nodes": ["0,1,2,3"],
334
        }
335
        self.verify_affinity([cpuinfo.pool_all], ground_truth)
336

337
        cpuinfo.gen_pools_ondemand(nodes_list=[1, 2])
338
        ground_truth = {
339
            "ninstances": 1,
340
            "ncores_per_instance": 28,
341
            "num_cores_sum": 28,
342
            "num_nodes_sum": 2,
343
            "num_cores": [28],
344
            "num_nodes": [2],
345
            "pools_cores": ["28-41,56-69"],
346
            "pools_nodes": ["1,2"],
347
        }
348
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
349

350
        num_nodes = 2
351
        n_phycores_per_node = 28
352
        lscpu_txt = construct_numa_config(
353
            num_nodes, n_phycores_per_node, enable_ht=True, n_e_cores=4, numa_mode=0
354
        )
355
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
356
        cpuinfo.gen_pools_ondemand(ninstances=2)
357
        ground_truth = {
358
            "ninstances": 2,
359
            "ncores_per_instance": 28,
360
            "num_cores_sum": 56,
361
            "num_nodes_sum": 2,
362
            "num_cores": [28, 28],
363
            "num_nodes": [1, 1],
364
            "pools_cores": ["0-27", "60-87"],
365
            "pools_nodes": ["0", "1"],
366
        }
367
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
368

369
        # mode 1
370
        num_nodes = 2
371
        n_phycores_per_node = 28
372
        lscpu_txt = construct_numa_config(
373
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=1
374
        )
375
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
376
        cpuinfo.gen_pools_ondemand(ninstances=1)
377
        ground_truth = {
378
            "ninstances": 1,
379
            "ncores_per_instance": 56,
380
            "num_cores_sum": 56,
381
            "num_nodes_sum": 2,
382
            "num_cores": [56],
383
            "num_nodes": [2],
384
            "pools_cores": ["0-55"],
385
            "pools_nodes": ["0,1"],
386
        }
387
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
388

389
        cpuinfo.gen_pools_ondemand(ninstances=2)
390
        ground_truth = {
391
            "ninstances": 2,
392
            "ncores_per_instance": 28,
393
            "num_cores_sum": 56,
394
            "num_nodes_sum": 2,
395
            "num_cores": [28, 28],
396
            "num_nodes": [1, 1],
397
            "pools_cores": ["0-27", "28-55"],
398
            "pools_nodes": ["0", "1"],
399
        }
400
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
401

402
        cpuinfo.gen_pools_ondemand(ninstances=4)
403
        ground_truth = {
404
            "ninstances": 4,
405
            "ncores_per_instance": 14,
406
            "num_cores_sum": 56,
407
            "num_nodes_sum": 2,
408
            "num_cores": [14, 14, 14, 14],
409
            "num_nodes": [1, 1, 1, 1],
410
            "pools_cores": ["0-13", "14-27", "28-41", "42-55"],
411
            "pools_nodes": ["0", "0", "1", "1"],
412
        }
413
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
414

415
        cpuinfo.gen_pools_ondemand(ncores_per_instance=28)
416
        ground_truth = {
417
            "ninstances": 2,
418
            "ncores_per_instance": 28,
419
            "num_cores_sum": 56,
420
            "num_nodes_sum": 2,
421
            "num_cores": [28, 28],
422
            "num_nodes": [1, 1],
423
            "pools_cores": ["0-27", "28-55"],
424
            "pools_nodes": ["0", "1"],
425
        }
426
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
427

428
        cpuinfo.gen_pools_ondemand(ncores_per_instance=14)
429
        ground_truth = {
430
            "ninstances": 4,
431
            "ncores_per_instance": 14,
432
            "num_cores_sum": 56,
433
            "num_nodes_sum": 2,
434
            "num_cores": [14, 14, 14, 14],
435
            "num_nodes": [1, 1, 1, 1],
436
            "pools_cores": ["0-13", "14-27", "28-41", "42-55"],
437
            "pools_nodes": ["0", "0", "1", "1"],
438
        }
439
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
440

441
        cores_list_local = []
442
        cores_list_local.extend(list(i for i in range(14, 28)))
443
        cores_list_local.extend(list(i for i in range(42, 56)))
444
        cpuinfo.gen_pools_ondemand(ninstances=2, cores_list=cores_list_local)
445
        ground_truth = {
446
            "ninstances": 2,
447
            "ncores_per_instance": 14,
448
            "num_cores_sum": 28,
449
            "num_nodes_sum": 2,
450
            "num_cores": [14, 14],
451
            "num_nodes": [1, 1],
452
            "pools_cores": ["14-27", "42-55"],
453
            "pools_nodes": ["0", "1"],
454
        }
455
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
456

457
        num_nodes = 4
458
        n_phycores_per_node = 14
459
        lscpu_txt = construct_numa_config(
460
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=1
461
        )
462
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
463
        cpuinfo.gen_pools_ondemand(ninstances=2, nodes_list=[1, 2])
464
        ground_truth = {
465
            "ninstances": 2,
466
            "ncores_per_instance": 14,
467
            "num_cores_sum": 28,
468
            "num_nodes_sum": 2,
469
            "num_cores": [14, 14],
470
            "num_nodes": [1, 1],
471
            "pools_cores": ["14-27", "28-41"],
472
            "pools_nodes": ["1", "2"],
473
        }
474
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
475

476
        num_nodes = 2
477
        n_phycores_per_node = 28
478
        lscpu_txt = construct_numa_config(
479
            num_nodes, n_phycores_per_node, enable_ht=True, n_e_cores=4, numa_mode=1
480
        )
481
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
482
        cpuinfo.gen_pools_ondemand(ninstances=2)
483
        ground_truth = {
484
            "ninstances": 2,
485
            "ncores_per_instance": 28,
486
            "num_cores_sum": 56,
487
            "num_nodes_sum": 2,
488
            "num_cores": [28, 28],
489
            "num_nodes": [1, 1],
490
            "pools_cores": ["0-27", "28-55"],
491
            "pools_nodes": ["0", "1"],
492
        }
493
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
494

495
        # mode 2
496
        num_nodes = 2
497
        n_phycores_per_node = 28
498
        lscpu_txt = construct_numa_config(
499
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=2
500
        )
501
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
502
        cpuinfo.gen_pools_ondemand(ninstances=2)
503
        ground_truth = {
504
            "ninstances": 2,
505
            "ncores_per_instance": 28,
506
            "num_cores_sum": 56,
507
            "num_nodes_sum": 2,
508
            "num_cores": [28, 28],
509
            "num_nodes": [1, 1],
510
            "pools_cores": [
511
                "0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,32,34,36,38,40,42,44,46,48,50,52,54",
512
                "56,58,60,62,64,66,68,70,72,74,76,78,80,82,84,86,88,90,92,94,96,98,100,102,104,106,108,110",
513
            ],
514
            "pools_nodes": ["0", "1"],
515
        }
516
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
517

518
        cpuinfo.gen_pools_ondemand(ninstances=4)
519
        ground_truth = {
520
            "ninstances": 4,
521
            "ncores_per_instance": 14,
522
            "num_cores_sum": 56,
523
            "num_nodes_sum": 2,
524
            "num_cores": [14, 14, 14, 14],
525
            "num_nodes": [1, 1, 1, 1],
526
            "pools_cores": [
527
                "0,2,4,6,8,10,12,14,16,18,20,22,24,26",
528
                "28,30,32,34,36,38,40,42,44,46,48,50,52,54",
529
                "56,58,60,62,64,66,68,70,72,74,76,78,80,82",
530
                "84,86,88,90,92,94,96,98,100,102,104,106,108,110",
531
            ],
532
            "pools_nodes": ["0", "0", "1", "1"],
533
        }
534
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
535

536
        cpuinfo.gen_pools_ondemand(ncores_per_instance=28)
537
        ground_truth = {
538
            "ninstances": 2,
539
            "ncores_per_instance": 28,
540
            "num_cores_sum": 56,
541
            "num_nodes_sum": 2,
542
            "num_cores": [28, 28],
543
            "num_nodes": [1, 1],
544
            "pools_cores": [
545
                "0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,32,34,36,38,40,42,44,46,48,50,52,54",
546
                "56,58,60,62,64,66,68,70,72,74,76,78,80,82,84,86,88,90,92,94,96,98,100,102,104,106,108,110",
547
            ],
548
            "pools_nodes": ["0", "1"],
549
        }
550
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
551

552
        cpuinfo.gen_pools_ondemand(ncores_per_instance=14)
553
        ground_truth = {
554
            "ninstances": 4,
555
            "ncores_per_instance": 14,
556
            "num_cores_sum": 56,
557
            "num_nodes_sum": 2,
558
            "num_cores": [14, 14, 14, 14],
559
            "num_nodes": [1, 1, 1, 1],
560
            "pools_cores": [
561
                "0,2,4,6,8,10,12,14,16,18,20,22,24,26",
562
                "28,30,32,34,36,38,40,42,44,46,48,50,52,54",
563
                "56,58,60,62,64,66,68,70,72,74,76,78,80,82",
564
                "84,86,88,90,92,94,96,98,100,102,104,106,108,110",
565
            ],
566
            "pools_nodes": ["0", "0", "1", "1"],
567
        }
568
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
569

570
        cpuinfo.gen_pools_ondemand(ninstances=3)
571
        ground_truth = {
572
            "ninstances": 3,
573
            "ncores_per_instance": 18,
574
            "num_cores_sum": 54,
575
            "num_nodes_sum": 2,
576
            "num_cores": [18, 18, 18],
577
            "num_nodes": [1, 2, 1],
578
            "pools_cores": [
579
                "0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,32,34",
580
                "36,38,40,42,44,46,48,50,52,54,56,58,60,62,64,66,68,70",
581
                "72,74,76,78,80,82,84,86,88,90,92,94,96,98,100,102,104,106",
582
            ],
583
            "pools_nodes": ["0", "0,1", "1"],
584
        }
585
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
586

587
        cores_list_local = []
588
        cores_list_local.extend(list(i for i in range(14, 28)))
589
        cores_list_local.extend(list(i for i in range(98, 112)))
590
        cpuinfo.gen_pools_ondemand(ninstances=2, cores_list=cores_list_local)
591
        ground_truth = {
592
            "ninstances": 2,
593
            "ncores_per_instance": 14,
594
            "num_cores_sum": 28,
595
            "num_nodes_sum": 2,
596
            "num_cores": [14, 14],
597
            "num_nodes": [1, 1],
598
            "pools_cores": ["14-27", "98-111"],
599
            "pools_nodes": ["0", "1"],
600
        }
601
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
602

603
        num_nodes = 4
604
        n_phycores_per_node = 14
605
        lscpu_txt = construct_numa_config(
606
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=2
607
        )
608
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
609
        cpuinfo.gen_pools_ondemand(nodes_list=[1, 2])
610
        ground_truth = {
611
            "ninstances": 1,
612
            "ncores_per_instance": 28,
613
            "num_cores_sum": 28,
614
            "num_nodes_sum": 2,
615
            "num_cores": [28],
616
            "num_nodes": [2],
617
            "pools_cores": [
618
                "28,30,32,34,36,38,40,42,44,46,48,50,52,54,56,58,60,62,64,66,68,70,72,74,76,78,80,82"
619
            ],
620
            "pools_nodes": ["1,2"],
621
        }
622
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
623

624
        num_nodes = 2
625
        n_phycores_per_node = 28
626
        lscpu_txt = construct_numa_config(
627
            num_nodes, n_phycores_per_node, enable_ht=True, n_e_cores=4, numa_mode=2
628
        )
629
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
630
        cpuinfo.gen_pools_ondemand(ninstances=2)
631
        ground_truth = {
632
            "ninstances": 2,
633
            "ncores_per_instance": 28,
634
            "num_cores_sum": 56,
635
            "num_nodes_sum": 2,
636
            "num_cores": [28, 28],
637
            "num_nodes": [1, 1],
638
            "pools_cores": [
639
                "0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,32,34,36,38,40,42,44,46,48,50,52,54",
640
                "60,62,64,66,68,70,72,74,76,78,80,82,84,86,88,90,92,94,96,98,100,102,104,106,108,110,112,114",
641
            ],
642
            "pools_nodes": ["0", "1"],
643
        }
644
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
645

646
    def test_core_affinity_with_logical_cores(self):
647
        num_nodes = 2
648
        n_phycores_per_node = 28
649
        lscpu_txt = construct_numa_config(
650
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=1
651
        )
652
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
653
        cpuinfo.gen_pools_ondemand(ninstances=2, use_logical_cores=True)
654
        ground_truth = {
655
            "ninstances": 2,
656
            "ncores_per_instance": 56,
657
            "num_cores_sum": 112,
658
            "num_nodes_sum": 2,
659
            "num_cores": [56, 56],
660
            "num_nodes": [1, 1],
661
            "pools_cores": ["0-27,56-83", "28-55,84-111"],
662
            "pools_nodes": ["0", "1"],
663
        }
664
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
665

666
    def test_core_affinity_with_skip_cross_node_cores(self):
667
        num_nodes = 2
668
        n_phycores_per_node = 28
669
        lscpu_txt = construct_numa_config(
670
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=1
671
        )
672
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
673
        cpuinfo.gen_pools_ondemand(ninstances=3, skip_cross_node_cores=True)
674
        ground_truth = {
675
            "ninstances": 3,
676
            "ncores_per_instance": 14,
677
            "num_cores_sum": 42,
678
            "num_nodes_sum": 2,
679
            "num_cores": [14, 14, 14],
680
            "num_nodes": [1, 1, 1],
681
            "pools_cores": ["0-13", "14-27", "28-41"],
682
            "pools_nodes": ["0", "0", "1"],
683
        }
684
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
685

686
    def test_core_affinity_with_skip_cross_node_cores_and_use_logical_core(self):
687
        num_nodes = 2
688
        n_phycores_per_node = 28
689
        lscpu_txt = construct_numa_config(
690
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=1
691
        )
692
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
693
        cpuinfo.gen_pools_ondemand(
694
            ninstances=7, use_logical_cores=True, skip_cross_node_cores=True
695
        )
696
        ground_truth = {
697
            "ninstances": 7,
698
            "ncores_per_instance": 14,
699
            "num_cores_sum": 98,
700
            "num_nodes_sum": 2,
701
            "num_cores": [14, 14, 14, 14, 14, 14, 14],
702
            "num_nodes": [1, 1, 1, 1, 1, 1, 1],
703
            "pools_cores": [
704
                "0-6,56-62",
705
                "7-13,63-69",
706
                "14-20,70-76",
707
                "21-27,77-83",
708
                "28-34,84-90",
709
                "35-41,91-97",
710
                "42-48,98-104",
711
            ],
712
            "pools_nodes": ["0", "0", "0", "0", "1", "1", "1"],
713
        }
714
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
715

716
    def test_core_affinity_with_skip_cross_node_cores_and_node_id_use_logical_core(
717
        self,
718
    ):
719
        num_nodes = 4
720
        n_phycores_per_node = 14
721
        lscpu_txt = construct_numa_config(
722
            num_nodes, n_phycores_per_node, enable_ht=True, numa_mode=1
723
        )
724
        cpuinfo = CPUPoolList(lscpu_txt=lscpu_txt)
725
        cpuinfo.gen_pools_ondemand(
726
            ninstances=3,
727
            nodes_list=[1, 2],
728
            use_logical_cores=True,
729
            skip_cross_node_cores=True,
730
        )
731
        ground_truth = {
732
            "ninstances": 3,
733
            "ncores_per_instance": 14,
734
            "num_cores_sum": 42,
735
            "num_nodes_sum": 2,
736
            "num_cores": [14, 14, 14],
737
            "num_nodes": [1, 1, 1],
738
            "pools_cores": ["14-20,70-76", "21-27,77-83", "28-34,84-90"],
739
            "pools_nodes": ["1", "1", "2"],
740
        }
741
        self.verify_affinity(cpuinfo.pools_ondemand, ground_truth)
742

743

744
if __name__ == "__main__":
745
    test = unittest.main()
746