Function test_data_transfer

verl/tests/ray/test_data_transfer.py:52–109

()

def test_data_transfer():
    ray.init()
    # construct resource pool
    resource_pool = RayResourcePool([8])
    cls_with_init = RayClassWithInitArgs(cls=DummyWorker)
    # construct worker group
    wg = RayWorkerGroup(resource_pool, cls_with_init)

    # this is real dataset size
    batch_size = 4096
    seqlen = 32768

    data_dict = {}

    for i in range(2):
        data_dict[str(i)] = torch.randint(0, 10000, (batch_size, seqlen))

    data = DataProto.from_dict(tensors=data_dict)

    print(data)

    # we manually split data here and send to each worker
    data_list = data.chunk(wg.world_size)

    for i in range(wg.world_size):
        # consolidate is necessary
        if tensordict.__version__ >= '0.5.0':
            data_list[i].batch = data_list[i].batch.consolidate()

    with Timer(name='ray.pickle', initial_text=True):
        for i in range(wg.world_size):
            ray.cloudpickle.pickle.dumps(data_list[i])

    with Timer(name='raw.pickle', initial_text=True):
        import pickle
        for i in range(wg.world_size):
            pickle.dumps(data_list[i])

    # we put in advance
    with Timer(name='put', initial_text=True):
        # takes around 40 seconds
        data_list_ref = parallel_put(data_list)
        # for i in range(wg.world_size):
        #     data_list[i] = ray.put(data_list[i])

    with Timer(name='launch', initial_text=True):
        output_ref = wg.do_nothing(data_list_ref)

    with Timer(name='get', initial_text=True):
        # takes around 40 seconds
        output_lst = ray.get(output_ref)

    for input_data, output_data in zip(data_list, output_lst):
        for key in input_data.batch.keys():
            assert torch.all(torch.eq(input_data.batch[key] + 1,
                                      output_data.batch[key])), (input_data.batch[key], output_data.batch[key], key)

    ray.shutdown()