def _split_local(partition_book, rank, elements, local_eles):
    """Split the input element list with respect to data locality."""
    num_clients = role.get_num_trainers()
    num_client_per_part = num_clients // partition_book.num_partitions()
    if rank is None:
        rank = role.get_trainer_rank()
    assert (
        rank < num_clients
    ), "The input rank ({}) is incorrect. #Trainers: {}".format(
        rank, num_clients
    )
    # all ranks of the clients in the same machine are in a contiguous range.
    client_id_in_part = rank % num_client_per_part
    local_eles = _get_overlap(elements, local_eles)

    # get a subset for the local client.
    size = len(local_eles) // num_client_per_part
    # if this isn't the last client in the partition.
    if client_id_in_part + 1 < num_client_per_part:
        return local_eles[
            (size * client_id_in_part) : (size * (client_id_in_part + 1))
        ]
    else:
        return local_eles[(size * client_id_in_part) :]
Split the input element list with respect to data locality.
python/dgl/distributed/dist_graph.py