How to timeout all_reduce or prevent it from hangs

username: Yi_Zhang
cooked: <p>Hi,</p>
<p>Anyone knows how to stop / terminate an all-reduce call properly when it doesn’t get reply from other processes?</p>
<p>Just to explain my question, please see the sample code below. I have 4 processes divided into two sub-groups (group1 and group2), and a shared Queue with 5 elements. Each process will try to get one element from the queue in the while loop until the queue becomes empty. And inside the while loop, each process will do all_reduce with its “neighbor” in the same sub-group.<br>
The problem is that when one of the process get the last element, and at the same time, the shared queue is empty and its “neighbor” process already exit the while loop, it will get hanged and waiting forever for the all_reduce reply. Is there any way to set timeout for all_reduce call? Or some other ways to solve this situation?</p>
<p>Thanks. Please see codes attached below.</p>
<pre><code class="lang-auto">import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

def run(rank, a, q):
    dist_init_method = 'tcp://{master_ip}:{master_port}'.format(
                master_ip='127.0.0.1', master_port='12346')
    world_size = 4
    torch.distributed.init_process_group(backend="nccl",
                                        init_method=dist_init_method,
                                        world_size=world_size,
                                        rank=rank)
    group1 = dist.new_group([0, 1])
    group2 = dist.new_group([2, 3])
    tensor = torch.ones(1)
    device = torch.device('cuda', rank)
    tensor = tensor.to(device)
    
    while not q.empty():
        current_index = q.get()
        print(f'Process {rank} current index is: {current_index}')
        
        if rank == 0 or rank == 1:
            dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group1)
        else:
            dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group2)
    
    print('Rank ', rank, ' has data ', tensor[0])

if __name__ == "__main__":
    a = 1
    ctx = mp.get_context('spawn')
    q = ctx.Queue()
    for index in range(5):
        q.put(index)
    
    mp.spawn(run, args=(a, q), nprocs=4)
</code></pre>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/yi_zhang">@Yi_Zhang</a>, you can set a timeout in <code>init_process_group</code>. For NCCL backend, it also requires setting <code>NCCL_BLOCKING_WAIT</code> env var to 1.</p>
<p>More explanation can be found here <a href="https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group" rel="nofollow noopener">https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group</a>. (search for <code>NCCL_BLOCKING_WAIT </code>)</p>

username: Yi_Zhang
cooked: <p>Hi <a class="mention" href="/u/mrshenli">@mrshenli</a>, thanks for your reply.<br>
I tried with  “gloo”, it does terminate the process and throws Exceptions. But with “nccl”, I add the following line in the bash script to run the .py file,</p>
<pre><code class="lang-auto">export NCCL_BLOCKING_WAIT=1
</code></pre>
<p>However it doesn’t work.<br>
Also tried with adding the following line in the .py file itself, but doesn’t work either.</p>
<pre><code class="lang-auto">os.environ["NCCL_BLOCKING_WAIT"] = "1"
</code></pre>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/yi_zhang">@Yi_Zhang</a>, did you set the env var within each spawned process (i.e., in <code>run</code> function) and before calling <code>init_process_group</code>?</p>

username: Yi_Zhang
cooked: <p>Hi <a class="mention" href="/u/mrshenli">@mrshenli</a>, yes, I did that for each process. Don’t understand where is the mistake. Here is the sample code:</p>
<pre><code class="lang-auto">import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import time
import datetime

def run(rank, a, q):
   os.environ["NCCL_BLOCKING_WAIT"] = "1"
   print('Rank ', rank, 'NCCL_BLOCKING_WAIT is: ', os.environ["NCCL_BLOCKING_WAIT"])
   dist_init_method = 'tcp://{master_ip}:{master_port}'.format(
               master_ip='127.0.0.1', master_port='12346')
   torch.distributed.init_process_group(backend="nccl",
                                       init_method=dist_init_method,
                                       timeout=datetime.timedelta(seconds=5),
                                       world_size=4,
                                       rank=rank)
   group1 = dist.new_group([0, 1])
   group2 = dist.new_group([2, 3])
   tensor = torch.ones(1)

   device = torch.device('cuda', rank)
   tensor = tensor.to(device)

   while not q.empty():
       print('Rank ', rank, ' in the loop ')
       current_index = q.get()
       print(f'Process {rank} current index is: {current_index}')
       try:
           if rank == 0 or rank == 1:
               dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group1)
               print(f'Process {rank} all_reduce tensor is: {tensor}')
           else:
               dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group2)
               print(f'Process {rank} all_reduce tensor is: {tensor}')
       except Exception:
           pass
   print('Rank ', rank, ' has data ', tensor[0])

if __name__ == "__main__":
   a = 1
   ctx = mp.get_context('spawn')
   q = ctx.Queue()
   flag = ctx.Queue()
   for index in range(5):
       q.put(index)
   
   mp.spawn(run, args=(a, q), nprocs=4)
</code></pre>
<p>An update, I tried “gloo” without setting timeout, and it can terminate properly. I’m wondering maybe “gloo” takes care of this situation by itself? It doesn’t have something to do with the “timeout”?</p>

username: mrshenli
cooked: <p>I confirm that I can reproduce this locally.</p>
<p>Hey <a class="mention" href="/u/osalpekar">@osalpekar</a>, do you know if we miss anything here? Would I be correct if I assume the following code is expected to abort the op in this case?</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/ecb88c5d11895a68e5f20917d27a0debbc0f0697/torch/lib/c10d/ProcessGroupNCCL.cpp#L301-L335" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/ecb88c5d11895a68e5f20917d27a0debbc0f0697/torch/lib/c10d/ProcessGroupNCCL.cpp#L301-L335" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/ecb88c5d11895a68e5f20917d27a0debbc0f0697/torch/lib/c10d/ProcessGroupNCCL.cpp#L301-L335</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="301" style="counter-reset: li-counter 300 ;">
<li>if (blockingWait_) {</li>
<li>  // Use the passed in timeout if provided, otherwise use the default</li>
<li>  // opTimeout for each WorkNCCL object.</li>
<li>  std::chrono::milliseconds workTimeout =</li>
<li>      timeout == kNoTimeout ? opTimeout_ : timeout;</li>
<li>  // Wait for the operation to complete.</li>
<li>  while (!isCompleted()) {</li>
<li>    auto currentTimepoint = std::chrono::steady_clock::now();</li>
<li>    if (std::chrono::duration_cast&lt;std::chrono::milliseconds&gt;(</li>
<li>            currentTimepoint - workStartTime_) &gt; workTimeout) {</li>
<li>      // When operation times out due to some errors that are not</li>
<li>      // detected by nccl communicators, ncclCommWatchdog can not check this</li>
<li>      // time out error and thus can not abort ncclComms accordingly.</li>
<li>      // So explicitly abort ncclComms here before throwing this timed out</li>
<li>      // exception to users, after this, ncclCommWatchdog can detect nccl</li>
<li>      // communicators are aborted and clean up devNCCLCommMap_ accordingly.</li>
<li>      // if throwing timed out excepiton without aborting nccl communicators</li>
<li>      // here, it was observed that CUDA GPU will have 100% utilization and</li>
<li>      // can not run new events successfully.</li>
<li>      for (const auto&amp; ncclComm : ncclComms_) {</li>
</ol></code></pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/ecb88c5d11895a68e5f20917d27a0debbc0f0697/torch/lib/c10d/ProcessGroupNCCL.cpp#L301-L335" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


username: osalpekar
cooked: <p><a class="mention" href="/u/mrshenli">@mrshenli</a> - Yes that is the code block that should abort the op if it times out.</p>
<p><a class="mention" href="/u/yi_zhang">@Yi_Zhang</a> - There is a workaround. The all_reduce call actually returns an async work handle. You can capture that handle and wait on it as such:</p>
<pre><code class="lang-auto">work = dist.all_reduce(..., async_op=True)
work.wait(SOME_TIMEOUT)
</code></pre>
<p>If the all_reduce call times out, then the wait call will throw an exception.</p>
<p>In the meantime, let me try to repro from your most recent code snippet.</p>

username: Yi_Zhang
cooked: <p><a class="mention" href="/u/mrshenli">@mrshenli</a> <a class="mention" href="/u/osalpekar">@osalpekar</a>,  thanks for your reply. I find another way to avoid this situation without using timeout. I just add some checks to make sure the pairs of processes terminate after same number of rounds. But I’m still curious to know if you have any answer for the  timeout issue. Thanks</p>