use condition_variable and wait_until in nccl dump on timeout (pytorch#120544) (#1365)

Fixes test_c10d_nccl.py -k test_timeout_dumps_timing_enabled_True.

Pull Request resolved: pytorch#120544
Approved by: https://github.com/atalman

Co-authored-by: Jeff Daily <[email protected]>

Author: BLOrange-AMD

c10/util/signal_handler.cpp

@@ -10,10 +10,14 @@
 #include <sys/syscall.h>
 #include <unistd.h>
 
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
 #include <cstdint>
 #include <cstdio>
 #include <cstdlib>
 #include <iostream>
+#include <mutex>
 
 #ifdef C10_ANDROID
 #ifndef SYS_gettid
@@ -109,8 +113,9 @@ FatalSignalHandler::FatalSignalHandler()
     : fatalSignalHandlersInstalled(false),
       fatalSignalReceived(false),
       fatalSignalName("<UNKNOWN>"),
-      writingCond(PTHREAD_COND_INITIALIZER),
-      writingMutex(PTHREAD_MUTEX_INITIALIZER) {}
+      writingCond(),
+      writingMutex(),
+      signalReceived(false) {}
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
 FatalSignalHandler::signal_handler FatalSignalHandler::kSignalHandlers[] = {
@@ -157,8 +162,10 @@ void FatalSignalHandler::callPreviousSignalHandler(
 
 // needsLock signals whether we need to lock our writing mutex.
 void FatalSignalHandler::stacktraceSignalHandler(bool needsLock) {
+  std::unique_lock<std::mutex> ul(writingMutex, std::defer_lock);
   if (needsLock) {
-    pthread_mutex_lock(&writingMutex);
+    ul.lock();
+    signalReceived = true;
   }
   pid_t tid = static_cast<pid_t>(syscall(SYS_gettid));
   std::string backtrace = fmt::format(
@@ -170,8 +177,8 @@ void FatalSignalHandler::stacktraceSignalHandler(bool needsLock) {
       c10::get_backtrace());
   std::cerr << backtrace << std::endl;
   if (needsLock) {
-    pthread_mutex_unlock(&writingMutex);
-    pthread_cond_signal(&writingCond);
+    ul.unlock();
+    writingCond.notify_all();
   }
 }
 
@@ -204,23 +211,32 @@ void FatalSignalHandler::fatalSignalHandler(int signum) {
     pid_t pid = getpid();
     pid_t currentTid = static_cast<pid_t>(syscall(SYS_gettid));
     struct dirent* entry = nullptr;
-    pthread_mutex_lock(&writingMutex);
+    std::unique_lock<std::mutex> ul(writingMutex);
     while ((entry = readdir(procDir)) != nullptr) {
       if (entry->d_name[0] == '.') {
         continue;
       }
       pid_t tid = atoi(entry->d_name);
       // If we've found the current thread then we'll jump into the SIGUSR2
-      // handler before calling pthread_cond_wait thus deadlocking, so branch
-      // our directly to the backtrace handler instead of signaling it.
+      // handler instead of signaling to avoid deadlocking.
       if (tid != currentTid) {
+        signalReceived = false;
         syscall(SYS_tgkill, pid, tid, SIGUSR2);
-        pthread_cond_wait(&writingCond, &writingMutex);
+        auto now = std::chrono::system_clock::now();
+        using namespace std::chrono_literals;
+        // we use wait_until instead of wait because on ROCm there was
+        // a single thread that wouldn't receive the SIGUSR2
+        if (std::cv_status::timeout == writingCond.wait_until(ul, now + 2s)) {
+          if (!signalReceived) {
+            std::cerr << "signal lost waiting for stacktrace " << pid << ":"
+                      << tid << std::endl;
+            break;
+          }
+        }
       } else {
         stacktraceSignalHandler(false);
       }
     }
-    pthread_mutex_unlock(&writingMutex);
   } else {
     perror("Failed to open /proc/self/task");
   }

c10/util/signal_handler.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <atomic>
+#include <condition_variable>
 #include <csignal>
 #include <cstdint>
 #include <mutex>
@@ -89,8 +90,10 @@ class C10_API FatalSignalHandler {
   // This wait condition is used to wait for other threads to finish writing
   // their stack trace when in fatal sig handler (we can't use pthread_join
   // because there's no way to convert from a tid to a pthread_t).
-  pthread_cond_t writingCond;
-  pthread_mutex_t writingMutex;
+  std::condition_variable writingCond;
+  std::mutex writingMutex;
+  // used to indicate if the other thread responded to the signal
+  bool signalReceived;
 
   struct signal_handler {
     const char* name;

test/distributed/test_c10d_nccl.py

@@ -3955,6 +3955,10 @@ def _check_return_codes(self, elapsed_time):
     @requires_nccl()
     @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
     def test_timeout_dumps_on_stuck_ranks(self):
+        # need rank0 to crash quicker after detecting timeout
+        os.environ['TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC'] = '1'
+        # restore this env var to its prior default in case another test changed it
+        os.environ['TORCH_NCCL_COORD_CHECK_MILSEC'] = '1000'
 
         if self.rank == self.MAIN_PROCESS_RANK:
             # wait for both rank0 and 1 to crash before looking for both ranks' output