Извлечение данных с удаленного сервера в загрузчике данных Pytorch

Question

У меня большой файл hd5 (~ 100 ГБ), содержащий функции изображения из re snet. Этот файл находится на моей локальной машине (ноутбуке). Моя модель обучена на узле кластера, который имеет ограничение хранения 25 ГБ. Сейчас я использую torch.distributed.rp c для передачи данных с локального компьютера в кластер. Я выставляю сервер на моей локальной машине следующим образом:

num_worker = 4

utils.WORLD_SIZE = num_worker + 1

import os
import torch
import utils
import torch.distributed.rpc as rpc

def run_worker(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '8888'

    rpc.init_rpc(utils.SERVER_NAME,
                 rank = rank, 
                 world_size = world_size)
    print("Server Initialized", flush=True)
    rpc.shutdown()

if __name__ == "__main__":
    rank = 0
    world_size = utils.WORLD_SIZE
    run_worker(rank, world_size)

Этот сервер отправляет данные с локального компьютера в кластер. (другие классы опущены)

Теперь для запроса данных из кластера я инициализирую rp c для каждого работника, используя worker_init_fn для загрузчика данных,

def worker_init_fn(worker_id):
    rpc.init_rpc(utils.CLIENT_NAME.format(worker_id+1),
                rank=worker_id+1, world_size=utils.WORLD_SIZE)

    worker_info = torch.utils.data.get_worker_info()
    dataset = worker_info.dataset
    worker_id = worker_info.id
    server_info = rpc.get_worker_info(utils.SERVER_NAME)
    dataset.server_ref = rpc.remote(server_info, utils.Server)

Сейчас , когда я запускаю свой код, тренировка l oop завершает одну итерацию набора данных и после этого зависает, и я получаю следующую ошибку на стороне кластера:

Traceback (most recent call last):
  File "custom_datasets.py", line 134, in <module>
    main()
  File "custom_datasets.py", line 110, in main
    for i, (images, labels) in enumerate(mn_dataset_loader):
  File "/home/kanishk/.local/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 345, in __next__
    data = self._next_data()
  File "/home/kanishk/.local/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 856, in _next_data
    return self._process_data(data)
  File "/home/kanishk/.local/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 881, in _process_data
    data.reraise()
  File "/home/kanishk/.local/lib/python3.7/site-packages/torch/_utils.py", line 394, in reraise
    raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/home/kanishk/.local/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py", line 135, in _worker_loop
    init_fn(worker_id)
  File "custom_datasets.py", line 78, in worker_init_fn
    rank=worker_id+1, world_size=utils.WORLD_SIZE)
  File "/home/kanishk/.local/lib/python3.7/site-packages/torch/distributed/rpc/__init__.py", line 67, in init_rpc
    store, _, _ = next(rendezvous_iterator)
  File "/home/kanishk/.local/lib/python3.7/site-packages/torch/distributed/rendezvous.py", line 168, in _env_rendezvous_handler
    store = TCPStore(master_addr, master_port, world_size, start_daemon)
RuntimeError: connect() timed out.

Выше проблемы не возникает, когда я установите num_worker = 0, но код кластера очень медленный. Я думаю, что ошибка из-за многопоточности, но я не уверен, как решить эту проблему. Пожалуйста, помогите мне решить проблему.