move ProcessGroupGlooTest to gtest (pytorch#32133)

Summary: Pull Request resolved: pytorch#32133 We should do this to better debug the test. Differential Revision: D19375479 fbshipit-source-id: 8c2bf61bae605a38252bb793b091ade479bea11a
wuhuikx · Jan 15, 2020 · 2928105 · 2928105
1 parent bc0a628
commit 2928105
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 23 deletions.
diff --git a/torch/lib/c10d/test/CMakeLists.txt b/torch/lib/c10d/test/CMakeLists.txt
@@ -19,7 +19,7 @@ c10d_add_test(TCPStoreTest.cpp c10d gtest_main)
 
 if(USE_CUDA)
   if(USE_C10D_GLOO)
-    c10d_add_test(ProcessGroupGlooTest.cpp c10d c10d_cuda_test)
+    c10d_add_test(ProcessGroupGlooTest.cpp c10d c10d_cuda_test gtest_main)
     c10d_add_test(ProcessGroupGlooAsyncTest.cpp c10d c10d_cuda_test)
   endif()
   if(USE_C10D_NCCL)
@@ -29,7 +29,7 @@ if(USE_CUDA)
   endif()
 else()
   if(USE_C10D_GLOO)
-    c10d_add_test(ProcessGroupGlooTest.cpp c10d c10d)
+    c10d_add_test(ProcessGroupGlooTest.cpp c10d c10d gtest_main)
   endif()
 endif()
 

diff --git a/torch/lib/c10d/test/ProcessGroupGlooTest.cpp b/torch/lib/c10d/test/ProcessGroupGlooTest.cpp
@@ -9,6 +9,7 @@
 #include <sstream>
 #include <thread>
 
+#include <gtest/gtest.h>
 #include <torch/cuda.h>
 
 #include <c10d/FileStore.hpp>
@@ -370,7 +371,8 @@ void testRecv(const std::string& path) {
   senderThread.join();
 }
 
-int main(int argc, char** argv) {
+TEST(ProcessGroupGlooTest, testExceptionsThrown) {
+  // test SIGSTOP
   {
     TemporaryFile file;
     auto work = testSignal(file.path, SIGSTOP);
@@ -381,6 +383,7 @@ int main(int argc, char** argv) {
     }
   }
 
+  // test SIGKILL
   {
     TemporaryFile file;
     auto work = testSignal(file.path, SIGKILL);
@@ -390,50 +393,61 @@ int main(int argc, char** argv) {
       std::cout << "SIGKILL test got: " << ex.what() << std::endl;
     }
   }
+}
 
+TEST(ProcessGroupGlooTest, testAllReduceCPU) {
   {
     TemporaryFile file;
     testAllreduce(file.path, at::DeviceType::CPU);
   }
+}
 
-#ifdef USE_CUDA
-  {
-    if (torch::cuda::is_available()) {
-      TemporaryFile file;
-      testAllreduce(file.path, at::DeviceType::CUDA);
-    }
-  }
-#endif
-
+TEST(ProcessGroupGlooTest, testBroadcastCPU) {
   {
     TemporaryFile file;
     testBroadcast(file.path, at::DeviceType::CPU);
   }
+}
 
-#ifdef USE_CUDA
-  {
-    if (torch::cuda::is_available()) {
-      TemporaryFile file;
-      testBroadcast(file.path, at::DeviceType::CUDA);
-    }
-  }
-#endif
-
+TEST(ProcessGroupGlooTest, testBarrier) {
   {
     TemporaryFile file;
     testBarrier(file.path);
   }
+}
 
+TEST(ProcessGroupGlooTest, testSend) {
   {
     TemporaryFile file;
     testSend(file.path);
   }
+}
 
+TEST(ProcessGroupGlooTest, testRecv) {
   {
     TemporaryFile file;
     testRecv(file.path);
   }
+}
+
+#ifdef USE_CUDA
+// CUDA-only tests
+TEST(ProcessGroupGlooTest, testAllReduceCUDA) {
+  {
+    if (torch::cuda::is_available()) {
+      TemporaryFile file;
+      testAllreduce(file.path, at::DeviceType::CUDA);
+    }
+  }
+}
 
-  std::cout << "Test successful" << std::endl;
-  return 0;
+TEST(ProcessGroupGlooTest, testBroadcastCUDA) {
+  {
+    if (torch::cuda::is_available()) {
+      TemporaryFile file;
+      testBroadcast(file.path, at::DeviceType::CUDA);
+    }
+  }
 }
+
+#endif