Skip to content

Commit 9f0a0b4

Browse files
jingxu10EikanWang
authored andcommitted
fixed an affinity bug (#363)
* fixed an affinity bug * bug fix in test_launcher.py
1 parent 197ac93 commit 9f0a0b4

File tree

2 files changed

+4
-1
lines changed

2 files changed

+4
-1
lines changed

intel_extension_for_pytorch/cpu/launch.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -457,6 +457,7 @@ def get_ccl_worker_affinity(self, nproc_per_node, ccl_worker_count, total_cores)
457457
for proc in range(ppn):
458458
for ccl_worker in range(ccl_worker_count):
459459
affinity += str(proc * cores_per_rank + ccl_worker) + ","
460+
affinity = affinity[:-1]
460461
return affinity
461462

462463
def launch(self, args):

tests/cpu/test_launcher.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ def find_lib(self, lib_type):
1616
library_paths = []
1717
if "CONDA_PREFIX" in os.environ:
1818
library_paths.append(os.environ["CONDA_PREFIX"] + "/lib/")
19+
elif "VIRTUAL_ENV" in os.environ:
20+
library_paths.append(os.environ["VIRTUAL_ENV"] + "/lib/")
1921

2022
library_paths += ["{}/.local/lib/".format(expanduser("~")), "/usr/local/lib/",
2123
"/usr/local/lib64/", "/usr/lib/", "/usr/lib64/"]
@@ -63,7 +65,7 @@ def test_mpi_pin_domain_and_ccl_worker_affinity(self):
6365
expect_pin_domain = "[0xffffff0,0xffffff00000000,]"
6466
self.assertEqual(pin_doamin, expect_pin_domain)
6567
ccl_worker_affinity = launcher.get_ccl_worker_affinity(proc_per_node, ccl_worker_count, total_cores)
66-
expect_ccl_worker_affinity = "0,1,2,3,28,29,30,31,"
68+
expect_ccl_worker_affinity = "0,1,2,3,28,29,30,31"
6769
self.assertEqual(ccl_worker_affinity, expect_ccl_worker_affinity)
6870

6971

0 commit comments

Comments
 (0)