File size: 577 Bytes
d13981b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | from eval.distributed_eval_utils import distributed_batch_plan
def test_distributed_batch_plan_keeps_uneven_ranks_in_final_sync() -> None:
plan = distributed_batch_plan(total_items=2500, num_processes=8, batch_size=1)
assert plan.local_item_counts == (313, 313, 313, 313, 312, 312, 312, 312)
assert plan.sync_batches == 313
def test_distributed_batch_plan_handles_batched_uneven_shards() -> None:
plan = distributed_batch_plan(total_items=10, num_processes=3, batch_size=2)
assert plan.local_item_counts == (4, 3, 3)
assert plan.sync_batches == 2
|