gigl.distributed.utils.partition_book#

Functions#

build_partition_book(num_entities, rank, world_size)

Builds a range-based partition book for a given number of entities, rank, and world size.

get_ids_on_rank(partition_book, rank)

Provided a tensor-based partition book or a range-based bartition book and a rank, returns all the ids that are stored on that rank.

get_total_ids(partition_book)

Returns the total number of ids (e.g. the total number of nodes) from a partition book.

Module Contents#

gigl.distributed.utils.partition_book.build_partition_book(num_entities, rank, world_size)[source]#

Builds a range-based partition book for a given number of entities, rank, and world size.

The partition book is balanced, i.e. the difference between the number of entities in any two partitions is at most 1.

Examples

num_entities = 10, world_size = 2, rank = 0 -> RangePartitionBook(partition_ranges=[5, 10], partition_idx=0)

num_entities = 7, world_size = 3, rank = 0 -> RangePartitionBook(partition_ranges=[2, 4, 7], partition_idx=0)

Parameters:

  • num_entities (int) – Number of entities

  • rank (int) – Rank of current machine

  • world_size (int) – Total number of machines

Returns:

Range-based partition book

Return type:

RangePartitionBook

gigl.distributed.utils.partition_book.get_ids_on_rank(partition_book, rank)[source]#

Provided a tensor-based partition book or a range-based bartition book and a rank, returns all the ids that are stored on that rank. :param partition_book: Tensor or range-based partition book :type partition_book: Union[torch.Tensor, PartitionBook] :param rank: Rank of current machine :type rank: int

Parameters:

  • partition_book (Union[torch.Tensor, graphlearn_torch.partition.PartitionBook])

  • rank (int)

Return type:

torch.Tensor

gigl.distributed.utils.partition_book.get_total_ids(partition_book)[source]#

Returns the total number of ids (e.g. the total number of nodes) from a partition book. :param partition_book: Tensor or range-based partition book :type partition_book: Union[torch.Tensor, PartitionBook]

Returns:

Total number of ids in the partition book

Return type:

int

Parameters:

partition_book (Union[torch.Tensor, graphlearn_torch.partition.PartitionBook])