Common

Common#

2026-03-13

8 min read time

Applies to Linux

This page provides raft_dask class and function references for the publicly-exposed elements of the raft_dask.common package, which implements Multi-Node Multi-GPU (MNMG) communicator support via NCCL and UCX/UCXX.

Communicator#

class raft_dask.common.comms.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()
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.

valid_nccl_placements: ClassVar[tuple] = ('client', 'worker', 'scheduler')
worker_info(workers)
Builds a dictionary of { (worker_address, worker_port) :

(worker_rank, worker_port ) }

raft_dask.common.comms.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.comms.get_raft_comm_state(sessionId, state_object=None, dask_worker=None)

Retrieves RAFT comms state on the scheduler node, for the given sessionId, creating a new session if it does not exist. If no session id is given, returns the state dict for all sessions.

Parameters:

  • sessionId (SessionId value to retrieve from the dask_scheduler instances)

  • state_object (Object (either Worker, or Scheduler) on which the raft) – comm state will retrieved (or created)

  • dask_worker (dask_worker object) –

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

    and not the client)

Returns:

session state – session state associated with sessionId

Return type:

str

raft_dask.common.comms.get_ucx(dask_worker=None)

A simple convenience wrapper to make sure UCP listener and endpoints are only ever assigned once per worker.

Parameters:

dask_worker (dask_worker object) –

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

and not the client)

UCX#

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

add_server_endpoint(ep)
close_endpoints()
static get(listener_callback=_connection_func)
get_endpoint(ip, port)
get_worker()
listener_port()

NCCL#

class raft_dask.common.nccl.nccl

A NCCL/RCCL wrapper for initializing and closing NCCL/RCCL comms in Python. RCCL is the AMD implementation of NCCL.

abort()

nccl.abort(self)

Call abort on the underlying nccl comm

cu_device()

nccl.cu_device(self)

Get the device backing the underlying comm

Returns:

device id

Return type:

int

destroy()

nccl.destroy(self)

Call destroy on the underlying NCCL comm

get_comm()

nccl.get_comm(self)

Returns the underlying nccl comm in a size_t (similar to void*). This can be safely typecasted from size_t into ncclComm_t*

Returns:

ncclComm_t instance pointer

Return type:

size_t

static get_unique_id()

nccl.get_unique_id()

Returns a new nccl unique id

Returns:

nccl unique id

Return type:

str

init(nranks, commId, rank)

nccl.init(self, nranks, commId, rank)

Construct a nccl-py object

Parameters:

  • nranks (int size of clique)

  • commId (string unique id from client)

  • rank (int rank of current worker)

user_rank()

nccl.user_rank(self)

Get the rank id of the current comm

Returns:

rank

Return type:

int

raft_dask.common.nccl.unique_id()

unique_id()

Returns a new ncclUniqueId converted to a character array that can be safely serialized and shared to a remote worker.

Returns:

128-byte unique id

Return type:

str

Comms Utilities#

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_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_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_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_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_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_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

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_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_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

Utilities#

raft_dask.common.utils.parse_host_port(address)

Given a string address with host/port, build a tuple(host, port)

Parameters:

address (string address to parse)

Returns:

tuple with host and port info

Return type:

tuple(host, port)