Adding timeout functionality/EnvVar to TestBed (#1044)

* Adding timeout functionality/EnvVar to TestBed
* updating timeout unit to microseconds

Signed-off-by: Tim Hu <timhu102@amd.com>

test/ReduceTests.cpp

@@ -7,7 +7,7 @@
 
 namespace RcclUnitTesting
 {
-    TEST(Reduce, OutOfPlace)
+  TEST(Reduce, OutOfPlace)
   {
     TestBed testBed;
 

test/common/EnvVars.cpp

@@ -63,6 +63,7 @@ namespace RcclUnitTesting
     maxRanksPerGpu = GetEnvVar("UT_MAX_RANKS_PER_GPU", 1);
     showTiming     = GetEnvVar("UT_SHOW_TIMING",  1);
     useInteractive = GetEnvVar("UT_INTERACTIVE",  0);
+    timeoutUs      = GetEnvVar("UT_TIMEOUT_US" ,  200000);
 
     // Limit number of supported reduction operators to just ncclSum if only allReduce is built
 #ifdef BUILD_ALLREDUCE_ONLY
@@ -195,7 +196,8 @@ namespace RcclUnitTesting
         std::make_tuple("UT_MAX_RANKS_PER_GPU", maxRanksPerGpu, "Maximum number of ranks using the same GPU"),
         std::make_tuple("UT_PRINT_VALUES"     , printValues   , "Print array values (-1 for all)"),
         std::make_tuple("UT_SHOW_TIMING"      , showTiming    , "Show timing table"),
-        std::make_tuple("UT_INTERACTIVE"      , useInteractive, "Run in interactive mode")
+        std::make_tuple("UT_INTERACTIVE"      , useInteractive, "Run in interactive mode"),
+        std::make_tuple("UT_TIMEOUT_US"       , timeoutUs     , "Timeout limit for collective calls in us")
       };
 
     printf("================================================================================\n");

test/common/EnvVars.hpp

@@ -28,6 +28,7 @@ namespace RcclUnitTesting
     int  maxRanksPerGpu; // Number of ranks using the same GPU     [UT_MAX_RANKS_PER_GPU]
     bool showTiming;     // Show timing per case at end            [UT_SHOW_TIMING]
     bool useInteractive; // Run in interactive mode                [UT_INTERACTIVE]
+    int  timeoutUs;      // Set timeout for child in microseconds  [UT_TIMEOUT_US]
 
     // Constructor that parses and collects environment variables
     EnvVars();

test/common/ErrCode.hpp

@@ -11,7 +11,8 @@ namespace RcclUnitTesting
   typedef enum
   {
     TEST_SUCCESS = 0,
-    TEST_FAIL    = 1
+    TEST_FAIL    = 1,
+    TEST_TIMEOUT = 2
   } ErrCode;
 
 #define ERROR(...) printf("\033[0;31m" "[ ERROR    ] " "\033[0m" __VA_ARGS__)

test/common/TestBed.cpp

@@ -296,6 +296,7 @@ namespace RcclUnitTesting
       {
         InteractiveWait("Starting ExecuteCollectives for child " + std::to_string(childId));
         PIPE_WRITE(childId, cmd);
+        PIPE_WRITE(childId, ev.timeoutUs);
         PIPE_WRITE(childId, useHipGraph);
         int tempCurrentRanks = currentRanks.size();
         PIPE_WRITE(childId, tempCurrentRanks);

test/common/TestBedChild.cpp

@@ -393,6 +393,9 @@ namespace RcclUnitTesting
 
   ErrCode TestBedChild::ExecuteCollectives()
   {
+    int timeoutUs = 0;
+    PIPE_READ(timeoutUs);
+
     bool useHipGraph = false;
     PIPE_READ(useHipGraph);
 
@@ -432,7 +435,6 @@ namespace RcclUnitTesting
     {
       for (int localRank : localRanksToExecute)
       {
-        CHECK_HIP(hipSetDevice(this->deviceIds[localRank]));
         if (this->verbose) INFO("Capturing stream for rank %d\n", localRank);
         CHECK_HIP(hipSetDevice(this->deviceIds[localRank]));
         for (int i = 0; i < this->numStreamsPerGroup; i++)
@@ -659,10 +661,47 @@ namespace RcclUnitTesting
     }
 
     // Synchronize
+    std::vector<hipStream_t> streamsToComplete;
+    for (int localRank : localRanksToExecute)
+    {
+      for (int i = 0; i < this->numStreamsPerGroup; i++)
+        streamsToComplete.push_back(this->streams[localRank][i]);
+    }
+    int usElapsed = 0;
+    using namespace std::chrono;
+    using Clock = std::chrono::high_resolution_clock;
+    if (this->verbose) INFO("Starting sychronization and timing\n");
+    const auto start = Clock::now();
+    while (!streamsToComplete.empty() && usElapsed < timeoutUs)
+    {
+      for (int i = 0; i < streamsToComplete.size(); i++)
+      {
+        if (hipStreamQuery(streamsToComplete[i]) == hipSuccess)
+        {
+          streamsToComplete.erase(streamsToComplete.begin() + i);
+          i--;
+        }  
+      }
+      usElapsed = duration_cast<microseconds>(Clock::now() - start).count();
+    }
+
+    // timed out
+    if (!streamsToComplete.empty())
+    {
+      if (this->verbose) INFO("Collective timed out, aborting\n");
+      for (int localRank : localRanksToExecute)
+      {
+        ncclCommAbort(this->comms[localRank]); 
+        timeoutUs = -1;
+      }
+    }
+
+    // extra sync to flush GPU cache for validation later
+    // TODO: remove this after figuring out & fixing the exact behavior 
+    // of fencing between kernels and at hipStreamQuery
     for (int localRank : localRanksToExecute)
     {
       if (this->verbose) INFO("Starting synchronization for rank %d\n", localRank);
-      CHECK_HIP(hipSetDevice(this->deviceIds[localRank]));
       for (int i = 0; i < this->numStreamsPerGroup; i++)
         CHECK_HIP(hipStreamSynchronize(this->streams[localRank][i]));
     }
@@ -699,6 +738,10 @@ namespace RcclUnitTesting
                  collArg.expected.ToString(collArg.dataType, numOutputElementsToPrint).c_str());
         }
     }
+
+    if (timeoutUs == -1)
+      return TEST_TIMEOUT;
+
     if (this->verbose) INFO("Child %d finishes ExecuteCollectives()\n", this->childId);
     return TEST_SUCCESS;
   }