- cugraph.dask.comms.comms.initialize(comms=None, p2p=False, prows=None, pcols=None, partition_type=1)#
Initialize a communicator for multi-node/multi-gpu communications. It is expected to be called right after client initialization for running multi-GPU algorithms (this wraps raft comms that manages underlying NCCL and UCX comms handles across the workers of a Dask cluster).
It is recommended to also call destroy() when the comms are no longer needed so the underlying resources can be cleaned up.
- commsraft Comms, optional (default=None)
A pre-initialized raft communicator. If provided, this is used for mnmg communications. If not provided, default comms are initialized as per client information.
- p2pbool, optional (default=False)
Initialize UCX endpoints if True.
- prowsint, optional (default=None)
Specifies the number of rows when performing a 2D partitioning of the input graph. If specified, this must be a factor of the total number of parallel processes. When specified with pcols, prows*pcols should be equal to the total number of parallel processes.
- pcolsint, optional (default=None)
Specifies the number of columns when performing a 2D partitioning of the input graph. If specified, this must be a factor of the total number of parallel processes. When specified with prows, prows*pcols should be equal to the total number of parallel processes.
- partition_typeint, optional (default=1)
Valid values are currently 1 or any int other than 1. A value of 1 (the default) represents a partitioning resulting in prows*pcols partitions. A non-1 value currently results in a partitioning of p*pcols partitions, where p is the number of GPUs.
>>> from dask.distributed import Client >>> from dask_cuda import LocalCUDACluster >>> import cugraph.dask.comms as Comms >>> cluster = LocalCUDACluster() >>> client = Client(cluster) >>> Comms.initialize(p2p=True) >>> # DO WORK HERE >>> # All done, clean up >>> Comms.destroy() >>> client.close() >>> cluster.close()