Deadlock with torch.distributed.rpc with num_workers > 1

Kanishk_Jain: <p>I have a large (93GB) .h5 file containing image features on my local system and my model is trained on SLURM ADA cluster which has a storage limit of 25GB.<br>
I am trying to use <strong>torch.distributed.rpc</strong> framework for requesting image features in <strong>Dataset.getitem</strong> using remote call to rpc server on my local system.</p>
<p>Code for initializing RPC server (local system):</p>
<pre><code class="lang-auto">import os
import torch.distributed.rpc as rpc

def run_worker(rank, world_size):
    os.environ['MASTER_ADDR'] = 'XX.X.XX.XX'
    os.environ['MASTER_PORT'] = 'XXXX'
    
    rpc.init_rpc(utils.SERVER_NAME,
                 rank=rank, 
                 world_size=world_size)
    print("Server Initialized", flush=True)

    rpc.shutdown()


if __name__ == "__main__":
    world_size = 2
    rank = 0
    
    run_worker(rank, world_size)
</code></pre>
<p>Code for RPC server for requesting data from local system (On ADA),</p>
<pre><code class="lang-auto">import os
import torch.distributed.rpc as rpc

def run_worker(rank, world_size):
    os.environ['MASTER_ADDR'] = 'XX.X.XX.XX'
    os.environ['MASTER_PORT'] = 'XXXX'

    rpc.init_rpc(utils.CLIENT_NAME.format(rank), rank=rank, world_size=world_size)
    print("Client Initialized", flush=True)

    main()

    rpc.shutdown()

if __name__ == '__main__':
    world_size = 2
    rank = 1

    run_worker(rank, world_size)
</code></pre>
<p>In my <strong>data_loader</strong> I have specified <strong>num_worker=8</strong>,<br>
Simplified code for dataset.getitem is (On ADA),</p>
<pre><code class="lang-auto">def __getitem__(self, index):
        ....
        ....

        print("fetching image for image_id {}, item {}".format(image_id, item), flush=True)
        v = utils._remote_method(utils.Server._load_image, self.server_ref, [index, self.coco_id_to_index])

        return v, ......
</code></pre>
<p>Now in my training loop when I call <strong>enumerate(data_loader)</strong>, multi-process data loading is enabled and <strong>getitem</strong> function is called <strong>num_worker</strong> times and a deadlock is reached.<br>
I am not sure why this deadlock is occuring because whenver <strong>getitem</strong> is called a remote call should be made to RPC server on my local system to request for data.</p>
<p>How can I resolve the deadlock? Is there any other way to solve my problem if large file doesn’t fit in my ADA system, I don’t want to compromise on my latency.</p>
<p>Edit: When I set num_worker=0 my code is working, but it is very slow 20 sec/iterations.</p>

mrshenli: <p>Hey <a class="mention" href="/u/kanishk_jain">@Kanishk_Jain</a></p>
<p>Thanks for trying out RPC.</p>
<blockquote>
<p>multi-process data loading is enabled and  <strong>getitem</strong>  function is called  <strong>num_worker</strong>  times and a deadlock is reached.</p>
</blockquote>
<p>It could be because it depleted the RPC threads in the thread pool. <code>num_send_recv_threads</code> by default is 4. Does it work if you bump up the number of threads? Sth like:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch.distributed.rpc import ProcessGroupRpcBackendOptions
from datetime import timedelta
import os
import torch.distributed.rpc as rpc

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'

options = ProcessGroupRpcBackendOptions()
options.rpc_timeout = timedelta(seconds=60)
options.init_method = "env://"
options.num_send_recv_threads = 32

rpc.init_rpc("client", rank=0, world_size=2, rpc_backend_options=options)
rpc.shutdown()
</code></pre>
<p>The current <code>rpc_backend_options</code> API is too verbose and not well documented. We will improve that in the next release.</p>

mrshenli: <p>Regarding the concern on speed, hopefully, using more workers in the data loader will help to boost the throughput. We are also working on making the RPC comm layer more efficient by adding <a href="https://github.com/pytorch/tensorpipe/tree/master/tensorpipe" rel="nofollow noopener">TensorPipe</a> as a new backend, so that RPC does not have to do two round trips for each message as with <code>ProcessGroup</code> and TensorPipe would also allow using multiple comm media.</p>