gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fixed p2p tests failing when there are no supported p2p devices (#386)

Browse Source
This commit is contained in:
Victor
2021-12-06 18:14:03 -08:00
committed by GitHub
parent b7f0e87dc2
commit 0ff1a26b70
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/test/unit/runtime/test_comm.py
View File

@@ -23,13 +23,13 @@ def get_p2p_matrix():
def get_p2p_devices(): def get_p2p_devices():
matrix = get_p2p_matrix() matrix = get_p2p_matrix()
idx = np.where(matrix == "OK") idx = np.where(matrix == "OK")
return f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}" return [f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}"] if len(idx[0]) > 0 else []
def get_non_p2p_devices(): def get_non_p2p_devices():
matrix = get_p2p_matrix() matrix = get_p2p_matrix()
idx = np.where(matrix == "NS") idx = np.where(matrix == "NS")
return f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}" return [f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}"] if len(idx[0]) > 0 else []
p2p_devices = get_p2p_devices() p2p_devices = get_p2p_devices()