Skip to content

Commit

Permalink
TST: Add a new test (untested, sorry) and a brief comment
Browse files Browse the repository at this point in the history
  • Loading branch information
seberg committed Feb 13, 2024
1 parent 0b69050 commit dce6ba3
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 0 deletions.
2 changes: 2 additions & 0 deletions python/raft-dask/raft_dask/common/comms.py
Original file line number Diff line number Diff line change
Expand Up @@ -693,6 +693,8 @@ def _func_worker_ranks(client, workers):
For each worker connected to the client, compute a global rank which takes
into account the NVML device index and the worker IP
(group workers on same host and order by NVML device).
Note that the reason for sorting was nvbug 4149999 and is presumably
fixed afterNCCL 2.19.3.
Parameters
----------
Expand Down
11 changes: 11 additions & 0 deletions python/raft-dask/raft_dask/test/test_comms.py
Original file line number Diff line number Diff line change
Expand Up @@ -246,6 +246,17 @@ def test_collectives(client, func, root_location):
cb.destroy()


@pytest.mark.nccl
@pytest.mark.parametrize("subset", [-1, 1, slice(None, None, -2)])
def test_comm_init_worker_subset(client, subset):
# Basic test that initializing a subset of workers is fine
cb = Comms(comms_p2p=True, verbose=True)

workers = list(client.scheduler_info()["workers"].keys())
workers = workers[subset]
cb.init(workers=workers)


@pytest.mark.nccl
def test_comm_split(client):

Expand Down

0 comments on commit dce6ba3

Please sign in to comment.