raft_dask.common.comms_utils

raft_dask.common.comms_utils#

4 min read time

Applies to Linux

Functions#

inject_comms_on_handle(handle, nccl_inst, is_ucxx, ...)

inject_comms_on_handle(handle, nccl_inst, is_ucxx, ucp_worker, eps, size, rank, verbose)

inject_comms_on_handle_coll_only(handle, nccl_inst, ...)

inject_comms_on_handle_coll_only(handle, nccl_inst, size, rank, verbose)

perform_test_comm_split(handle, n_colors)

perform_test_comm_split(handle, n_colors)

perform_test_comms_allgather(handle, root)

perform_test_comms_allgather(handle, root)

perform_test_comms_allreduce(handle, root)

perform_test_comms_allreduce(handle, root)

perform_test_comms_bcast(handle, root)

perform_test_comms_bcast(handle, root)

perform_test_comms_device_multicast_sendrecv(handle, ...)

perform_test_comms_device_multicast_sendrecv(handle, n_trials)

perform_test_comms_device_send_or_recv(handle, n_trials)

perform_test_comms_device_send_or_recv(handle, n_trials)

perform_test_comms_device_sendrecv(handle, n_trials)

perform_test_comms_device_sendrecv(handle, n_trials)

perform_test_comms_gather(handle, root)

perform_test_comms_gather(handle, root)

perform_test_comms_gatherv(handle, root)

perform_test_comms_gatherv(handle, root)

perform_test_comms_reduce(handle, root)

perform_test_comms_reduce(handle, root)

perform_test_comms_reducescatter(handle, root)

perform_test_comms_reducescatter(handle, root)

perform_test_comms_send_recv(handle, n_trials)

perform_test_comms_send_recv(handle, n_trials)

Module Contents#

raft_dask.common.comms_utils.inject_comms_on_handle(handle, nccl_inst, is_ucxx, ucp_worker, eps, size, rank, verbose)#

inject_comms_on_handle(handle, nccl_inst, is_ucxx, ucp_worker, eps, size, rank, verbose)

Given a handle and initialized comms, creates a comms_t instance and injects it into the handle.

Parameters:

  • handle (raft.common.Handle) – handle containing comms_t to use

  • nccl_inst (raft.dask.common.nccl) – Initialized nccl comm to use

  • ucp_worker (size_t pointer to initialized ucp_worker_h instance)

  • eps (size_t pointer to array of initialized ucp_ep_h instances)

  • size (int) – Number of workers in cluster

  • rank (int) – Rank of current worker

raft_dask.common.comms_utils.inject_comms_on_handle_coll_only(handle, nccl_inst, size, rank, verbose)#

inject_comms_on_handle_coll_only(handle, nccl_inst, size, rank, verbose)

Given a handle and initialized nccl comm, creates a comms_t instance and injects it into the handle.

Parameters:

  • handle (raft.common.Handle) – handle containing comms_t to use

  • nccl_inst (raft.dask.common.nccl) – Initialized nccl comm to use

  • size (int) – Number of workers in cluster

  • rank (int) – Rank of current worker

raft_dask.common.comms_utils.perform_test_comm_split(handle, n_colors)#

perform_test_comm_split(handle, n_colors)

Performs a p2p send/recv on the current worker

Parameters:

handle (raft.common.Handle) – handle containing comms_t to use

raft_dask.common.comms_utils.perform_test_comms_allgather(handle, root)#

perform_test_comms_allgather(handle, root)

Performs an broadcast on the current worker

Parameters:

handle (raft.common.Handle) – handle containing comms_t to use

raft_dask.common.comms_utils.perform_test_comms_allreduce(handle, root)#

perform_test_comms_allreduce(handle, root)

Performs an allreduce on the current worker

Parameters:

handle (raft.common.Handle) – handle containing comms_t to use

raft_dask.common.comms_utils.perform_test_comms_bcast(handle, root)#

perform_test_comms_bcast(handle, root)

Performs an broadcast on the current worker

Parameters:

handle (raft.common.Handle) – handle containing comms_t to use

raft_dask.common.comms_utils.perform_test_comms_device_multicast_sendrecv(handle, n_trials)#

perform_test_comms_device_multicast_sendrecv(handle, n_trials)

Performs a p2p device concurrent multicast send&recv on the current worker

Parameters:

  • handle (raft.common.Handle) – handle containing comms_t to use

  • n_trilas (int) – Number of test trials

raft_dask.common.comms_utils.perform_test_comms_device_send_or_recv(handle, n_trials)#

perform_test_comms_device_send_or_recv(handle, n_trials)

Performs a p2p device send or recv on the current worker

Parameters:

  • handle (raft.common.Handle) – handle containing comms_t to use

  • n_trilas (int) – Number of test trials

raft_dask.common.comms_utils.perform_test_comms_device_sendrecv(handle, n_trials)#

perform_test_comms_device_sendrecv(handle, n_trials)

Performs a p2p device concurrent send&recv on the current worker

Parameters:

  • handle (raft.common.Handle) – handle containing comms_t to use

  • n_trilas (int) – Number of test trials

raft_dask.common.comms_utils.perform_test_comms_gather(handle, root)#

perform_test_comms_gather(handle, root)

Performs a gather on the current worker

Parameters:

  • handle (raft.common.Handle) – handle containing comms_t to use

  • root (int) – Rank of the root worker

raft_dask.common.comms_utils.perform_test_comms_gatherv(handle, root)#

perform_test_comms_gatherv(handle, root)

Performs a gatherv on the current worker

Parameters:

  • handle (raft.common.Handle) – handle containing comms_t to use

  • root (int) – Rank of the root worker

raft_dask.common.comms_utils.perform_test_comms_reduce(handle, root)#

perform_test_comms_reduce(handle, root)

Performs an allreduce on the current worker

Parameters:

handle (raft.common.Handle) – handle containing comms_t to use

raft_dask.common.comms_utils.perform_test_comms_reducescatter(handle, root)#

perform_test_comms_reducescatter(handle, root)

Performs an allreduce on the current worker

Parameters:

handle (raft.common.Handle) – handle containing comms_t to use

raft_dask.common.comms_utils.perform_test_comms_send_recv(handle, n_trials)#

perform_test_comms_send_recv(handle, n_trials)

Performs a p2p send/recv on the current worker

Parameters:

  • handle (raft.common.Handle) – handle containing comms_t to use

  • n_trilas (int) – Number of test trials