gigl.distributed.dist_range_partitioner#

Attributes#

logger

Classes#

DistRangePartitioner

This class is responsible for implementing range-based partitioning. Rather than using a tensor-based partition

Module Contents#

class gigl.distributed.dist_range_partitioner.DistRangePartitioner(should_assign_edges_by_src_node=False, node_ids=None, node_features=None, edge_index=None, edge_features=None, positive_labels=None, negative_labels=None)[source]#

Bases: gigl.distributed.dist_partitioner.DistPartitioner

This class is responsible for implementing range-based partitioning. Rather than using a tensor-based partition book, this approach stores the upper bound of ids for each rank. For example, a range partition book [4, 8, 12] stores edge ids 0-3 on the 0th rank, 4-7 on the 1st rank, and 8-11 on the 2nd rank. While keeping the same id-indexing pattern for rank lookup as the tensor-based partitioning, this partition book does a search through these partition bounds to fetch the ranks, rather than using a direct index lookup. For example, to get the rank of node ids 1 and 6 by doing node_pb[[1, 6]], the range partition book uses torch.searchsorted on the partition bounds to return [0, 1], the ranks of each of these ids. As a result, the range-based partition book trades off more efficient memory storage for a slower lookup time for indices.

Initializes the parameters of the partitioner. Also optionally takes in node and edge tensors as arguments and registers them to the partitioner. Registered entities should be a dictionary of Dict[[NodeType or EdgeType], torch.Tensor] if heterogeneous or a torch.Tensor if homogeneous. This class assumes the distributed context has already been initialized outside of this class with the glt.distributed.init_worker_group() function and that rpc has been initialized with glt_distributed.init_rpc(). :param should_assign_edges_by_src_node: Whether edges should be assigned to the machine of the source nodes during partitioning :type should_assign_edges_by_src_node: bool :param node_ids: Optionally registered node ids from input. Tensors should be of shape [num_nodes_on_current_rank] :type node_ids: Optional[Union[torch.Tensor, Dict[NodeType, torch.Tensor]]] :param node_features: Optionally registered node feats from input. Tensors should be of shope [num_nodes_on_current_rank, node_feat_dim] :type node_features: Optional[Union[torch.Tensor, Dict[NodeType, torch.Tensor]]] :param edge_index: Optionally registered edge indexes from input. Tensors should be of shape [2, num_edges_on_current_rank] :type edge_index: Optional[Union[torch.Tensor, Dict[EdgeType, torch.Tensor]]] :param edge_features: Optionally registered edge features from input. Tensors should be of shape [num_edges_on_current_rank, edge_feat_dim] :type edge_features: Optional[Union[torch.Tensor, Dict[EdgeType, torch.Tensor]]] :param positive_labels: Optionally registered positive labels from input. Tensors should be of shape [2, num_pos_labels_on_current_rank] :type positive_labels: Optional[Union[torch.Tensor, Dict[EdgeType, torch.Tensor]]] :param negative_labels: Optionally registered negative labels from input. Tensors should be of shape [2, num_neg_labels_on_current_rank] :type negative_labels: Optional[Union[torch.Tensor, Dict[EdgeType, torch.Tensor]]]

Parameters:
partition_edge_index_and_edge_features(node_partition_book)[source]#

Partitions edges of a graph, including edge indices and edge features. If heterogeneous, partitions edges for all edge types. You must call partition_node first to get the node partition book as input. The difference between this function and its parent is that we no longer need to check that the edge_ids have been pre-computed as a prerequisite for partitioning edges and edge features.

Parameters:

node_partition_book (Union[PartitionBook, Dict[NodeType, PartitionBook]]) – The computed Node Partition Book

Returns:

Union[

Tuple[GraphPartitionData, FeaturePartitionData, PartitionBook], Tuple[Dict[EdgeType, GraphPartitionData], Dict[EdgeType, FeaturePartitionData], Dict[EdgeType, PartitionBook]],

]: Partitioned Graph Data, Feature Data, and corresponding edge partition book, is a dictionary if heterogeneous

Return type:

Union[tuple[gigl.types.graph.GraphPartitionData, Optional[gigl.types.graph.FeaturePartitionData], graphlearn_torch.partition.PartitionBook], tuple[Dict[gigl.src.common.types.graph_data.EdgeType, gigl.types.graph.GraphPartitionData], Optional[Dict[gigl.src.common.types.graph_data.EdgeType, gigl.types.graph.FeaturePartitionData]], Dict[gigl.src.common.types.graph_data.EdgeType, graphlearn_torch.partition.PartitionBook]]]

register_edge_index(edge_index)[source]#

Registers the edge_index to the partitioner. Unlike the tensor-based partitioner, this register pattern does not automatically infer edge ids,as they are not needed for partitioning.

For optimal memory management, it is recommended that the reference to edge_index tensor be deleted after calling this function using del <tensor>, as maintaining both original and intermediate tensors can cause OOM concerns.

Parameters:

edge_index (Union[torch.Tensor, Dict[EdgeType, torch.Tensor]]) – Input edge index which is either a torch.Tensor if homogeneous or a Dict if heterogeneous

Return type:

None

gigl.distributed.dist_range_partitioner.logger[source]#