dgl.distributed.node_split

dgl.distributed.node_split(nodes, partition_book=None, ntype='_N', rank=None, force_even=True, node_trainer_ids=None)[source]

Split nodes and return a subset for the local rank.

This function splits the input nodes based on the partition book and returns a subset of nodes for the local rank. This method is used for dividing workloads for distributed training.

The input nodes are stored as a vector of masks. The length of the vector is the same as the number of nodes in a graph; 1 indicates that the vertex in the corresponding location exists.

There are two strategies to split the nodes. By default, it splits the nodes in a way to maximize data locality. That is, all nodes that belong to a process are returned. If force_even is set to true, the nodes are split evenly so that each process gets almost the same number of nodes.

When force_even is True, the data locality is still preserved if a graph is partitioned with Metis and the node/edge IDs are shuffled. In this case, majority of the nodes returned for a process are the ones that belong to the process. If node/edge IDs are not shuffled, data locality is not guaranteed.

Parameters

  • nodes (1D tensor or DistTensor) – A boolean mask vector that indicates input nodes.

  • partition_book (GraphPartitionBook, optional) – The graph partition book

  • ntype (str, optional) – The node type of the input nodes.

  • rank (int, optional) – The rank of a process. If not given, the rank of the current process is used.

  • force_even (bool, optional) – Force the nodes are split evenly.

  • node_trainer_ids (1D tensor or DistTensor, optional) – If not None, split the nodes to the trainers on the same machine according to trainer IDs assigned to each node. Otherwise, split randomly.

Returns

The vector of node IDs that belong to the rank.

Return type

1D-tensor