Split the input element list with respect to data locality.
(partition_book, rank, elements, local_eles)
| 1699 | |
| 1700 | |
| 1701 | def _split_local(partition_book, rank, elements, local_eles): |
| 1702 | """Split the input element list with respect to data locality.""" |
| 1703 | num_clients = role.get_num_trainers() |
| 1704 | num_client_per_part = num_clients // partition_book.num_partitions() |
| 1705 | if rank is None: |
| 1706 | rank = role.get_trainer_rank() |
| 1707 | assert ( |
| 1708 | rank < num_clients |
| 1709 | ), "The input rank ({}) is incorrect. #Trainers: {}".format( |
| 1710 | rank, num_clients |
| 1711 | ) |
| 1712 | # all ranks of the clients in the same machine are in a contiguous range. |
| 1713 | client_id_in_part = rank % num_client_per_part |
| 1714 | local_eles = _get_overlap(elements, local_eles) |
| 1715 | |
| 1716 | # get a subset for the local client. |
| 1717 | size = len(local_eles) // num_client_per_part |
| 1718 | # if this isn't the last client in the partition. |
| 1719 | if client_id_in_part + 1 < num_client_per_part: |
| 1720 | return local_eles[ |
| 1721 | (size * client_id_in_part) : (size * (client_id_in_part + 1)) |
| 1722 | ] |
| 1723 | else: |
| 1724 | return local_eles[(size * client_id_in_part) :] |
| 1725 | |
| 1726 | |
| 1727 | def _even_offset(n, k): |
no test coverage detected