Function test_distributed

tests/data/test_data_iterator.py:80–123

(world_size, num_workers, device_batch_size)

    (3, 6, 4),
])
def test_distributed(world_size, num_workers, device_batch_size):
    start = 0
    global_batch_size = device_batch_size*world_size
    iterators = []
    datasets = [MockDataset("a", 5), MockDataset("b", 11)]
    mixture_rates = [0.8, 0.2]
    device_iterators = []
    bk = torch.utils.data.get_worker_info
    for rank in range(world_size):
        worker_iterators = []
        for worker_id in range(num_workers):
            worker_iterators.append(iter(IterableDatasetMixture(
                datasets, mixture_rates=mixture_rates,
                rank=rank, world_size=world_size,
                worker_info=MockWorkerInfo(worker_id, num_workers),
                global_batch_size=global_batch_size, seed=32, start_index=start)))

        def get_device_batch(_worker_its):
            while True:
                for it in _worker_its:
                    batch = []
                    for _ in range(device_batch_size):
                        batch.append(next(it))
                    yield batch

        device_iterators.append(get_device_batch(worker_iterators))
    torch.utils.data.get_worker_info = bk

    grouped_by_dataset = defaultdict(list)
    for i in range(100):
        global_batch = []
        for it in device_iterators:
            global_batch += next(it)
        global_batch.sort(key=lambda x: x.epoch)
        for ex in global_batch:
            grouped_by_dataset[ex.dataset].append(ex)

    for dataset in datasets:
        items = grouped_by_dataset[dataset.name]
        ds_len = dataset.n
        for epoch in range(len(items)//ds_len):
            epoch_items = items[epoch*ds_len:(epoch+1)*ds_len]
            assert all(x.epoch == epoch for x in epoch_items)
            assert set(x.idx for x in epoch_items) == set(range(ds_len))


@pytest.mark.parametrize("ns,start_index,world_size,rank", [