raft_dask.common

7 min read time

Applies to Linux

Submodules

• raft_dask.common.comms
• raft_dask.common.comms_utils
• raft_dask.common.nccl
• raft_dask.common.ucx
• raft_dask.common.utils

Classes

Comms	Initializes and manages underlying RCCL and UCX comms handles across
UCX	Singleton UCX context to encapsulate all interactions with the

Functions

local_handle(sessionId[, dask_worker])	Simple helper function for retrieving the local handle_t instance
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)

Package Contents

class raft_dask.common.Comms(comms_p2p=False, client=None, verbose=False, streams_per_handle=0, nccl_root_location='scheduler')

Initializes and manages underlying RCCL and UCX comms handles across the workers of a Dask cluster. It is expected that init() will be called explicitly. It is recommended to also call destroy() when the comms are no longer needed so the underlying resources can be cleaned up. This class is not meant to be thread-safe.

Examples

# The following code block assumes we have wrapped a C++
# function in a Python function called `run_algorithm`,
# which takes a `raft::handle_t` as a single argument.
# Once the `Comms` instance is successfully initialized,
# the underlying `raft::handle_t` will contain an instance
# of `raft::comms::comms_t`

from dask_cuda import LocalCUDACluster
from dask.distributed import Client

from raft.dask.common import Comms, local_handle

cluster = LocalCUDACluster()
client = Client(cluster)

def _use_comms(sessionId):
    return run_algorithm(local_handle(sessionId))

comms = Comms(client=client)
comms.init()

futures = [client.submit(_use_comms,
                         comms.sessionId,
                         workers=[w],
                         pure=False) # Don't memoize
               for w in cb.worker_addresses]
wait(dfs, timeout=5)

comms.destroy()
client.close()
cluster.close()

valid_nccl_placements: ClassVar[tuple] = ('client', 'worker', 'scheduler')

__del__()

create_nccl_uniqueid()

destroy()

Shuts down initialized comms and cleans up resources. This will be called automatically by the Comms destructor, but may be called earlier to save resources.

init(workers=None)

Initializes the underlying comms. NCCL is required but UCX is only initialized if comms_p2p == True

Parameters:

workers (Sequence) – Unique collection of workers for initializing comms.

worker_info(workers)

Builds a dictionary of { (worker_address, worker_port) :

(worker_rank, worker_port ) }

raft_dask.common.local_handle(sessionId, dask_worker=None)

Simple helper function for retrieving the local handle_t instance for a comms session on a worker.

Parameters:

• sessionId (str) – session identifier from an initialized comms instance

• dask_worker (dask_worker object) –

  (Note: if called by client.run(), this is supplied by Dask

  and not the client)

Returns:

handle

Return type:

raft.Handle or None

raft_dask.common.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.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.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.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.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.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.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.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.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.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.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.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.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.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

class raft_dask.common.UCX(listener_callback)

Singleton UCX context to encapsulate all interactions with the UCXX API and guarantee only a single listener & endpoints are created by RAFT Comms on a single process.

_UCX__instance = None

static get(listener_callback=_connection_func)

__del__()

_create_endpoint(ip, port)

_create_listener()

add_server_endpoint(ep)

close_endpoints()

get_endpoint(ip, port)

get_worker()

listener_port()