From b88c13487468ac2fbb8429934891f2fa214ca601 Mon Sep 17 00:00:00 2001
From: corey-derochie-amd
 <161367113+corey-derochie-amd@users.noreply.github.com>
Date: Fri, 22 Aug 2025 23:17:34 -0600
Subject: [PATCH] Changed `TestBedChild` to avoid hang if the call fails
 (#1875)

Changed `TestBedChild` protocol to send the result code before the return value to avoid hanging if the call fails. Switched `TestBedChild::GetUniqueId` to use this.
---
 test/common/TestBed.cpp      |  2 +-
 test/common/TestBedChild.cpp | 32 +++++++++++++++++++-------------
 test/common/TestBedChild.hpp |  2 +-
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/test/common/TestBed.cpp b/test/common/TestBed.cpp
index 1481ec52a7..0825a847d1 100644
--- a/test/common/TestBed.cpp
+++ b/test/common/TestBed.cpp
@@ -142,11 +142,11 @@ namespace RcclUnitTesting
     // Tell first rank to get ncclUniqueId
     int getIdCmd = TestBedChild::CHILD_GET_UNIQUE_ID;
     PIPE_WRITE(0, getIdCmd);
+    PIPE_CHECK(0);
 
     // Receive back unique ID from first rank
     ncclUniqueId id;
     PIPE_READ(0, id);
-    PIPE_CHECK(0);
 
     // Send InitComms command to each active child process
     int const cmd = TestBedChild::CHILD_INIT_COMMS;
diff --git a/test/common/TestBedChild.cpp b/test/common/TestBedChild.cpp
index c9b409e6d1..381d8f1dbb 100644
--- a/test/common/TestBedChild.cpp
+++ b/test/common/TestBedChild.cpp
@@ -114,19 +114,20 @@ namespace RcclUnitTesting
     {
       if (verbose) INFO("Child %d received command [%s]:\n", this->childId, ChildCommandNames[command]);;
       ErrCode status = TEST_SUCCESS;
+      std::vector<char> retValBuf;
       switch(command)
       {
-      case CHILD_GET_UNIQUE_ID   : status = GetUniqueId();        break;
-      case CHILD_INIT_COMMS      : status = InitComms();          break;
-      case CHILD_SET_COLL_ARGS   : status = SetCollectiveArgs();  break;
-      case CHILD_ALLOCATE_MEM    : status = AllocateMem();        break;
-      case CHILD_PREPARE_DATA    : status = PrepareData();        break;
-      case CHILD_EXECUTE_COLL    : status = ExecuteCollectives(); break;
-      case CHILD_VALIDATE_RESULTS: status = ValidateResults();    break;
-      case CHILD_LAUNCH_GRAPHS   : status = LaunchGraphs();       break;
-      case CHILD_DEALLOCATE_MEM  : status = DeallocateMem();      break;
-      case CHILD_DESTROY_COMMS   : status = DestroyComms();       break;
-      case CHILD_DESTROY_GRAPHS  : status = DestroyGraphs();      break;
+      case CHILD_GET_UNIQUE_ID   : status = GetUniqueId(retValBuf); break;
+      case CHILD_INIT_COMMS      : status = InitComms();            break;
+      case CHILD_SET_COLL_ARGS   : status = SetCollectiveArgs();    break;
+      case CHILD_ALLOCATE_MEM    : status = AllocateMem();          break;
+      case CHILD_PREPARE_DATA    : status = PrepareData();          break;
+      case CHILD_EXECUTE_COLL    : status = ExecuteCollectives();   break;
+      case CHILD_VALIDATE_RESULTS: status = ValidateResults();      break;
+      case CHILD_LAUNCH_GRAPHS   : status = LaunchGraphs();         break;
+      case CHILD_DEALLOCATE_MEM  : status = DeallocateMem();        break;
+      case CHILD_DESTROY_COMMS   : status = DestroyComms();         break;
+      case CHILD_DESTROY_GRAPHS  : status = DestroyGraphs();        break;
       case CHILD_STOP            : goto stop;
       default: exit(0);
       }
@@ -139,6 +140,10 @@ namespace RcclUnitTesting
         ERROR("Child %d write to parent failed: %s\n", this->childId, strerror(errno));
         break;
       }
+      if (retValBuf.size() > 0 && write(childWriteFd, retValBuf.data(), retValBuf.size()) < 0) {
+        ERROR("Child %d write return value to parent failed: %s\n", this->childId, strerror(errno));
+        break;
+      }
     }
   stop:
     if (verbose) INFO("Child %d exiting execution loop\n", this->childId);
@@ -150,14 +155,15 @@ namespace RcclUnitTesting
     exit(0);
   }
 
-  ErrCode TestBedChild::GetUniqueId()
+  ErrCode TestBedChild::GetUniqueId(std::vector<char>& retValBuf)
   {
     if (this->verbose) INFO("Child %d begins GetUniqueId()\n", this->childId);
 
     // Get a unique ID and pass it back to parent process
     ncclUniqueId id;
     CHILD_NCCL_CALL(ncclGetUniqueId(&id), "ncclGetUniqueId");
-    write(childWriteFd, &id, sizeof(id));
+    retValBuf.resize(sizeof(id));
+    memcpy(retValBuf.data(), &id, sizeof(id));
 
     if (this->verbose) INFO("Child %d finishes GetUniqueId()\n", this->childId);
     return TEST_SUCCESS;
diff --git a/test/common/TestBedChild.hpp b/test/common/TestBedChild.hpp
index 34bef5cebf..44c16de5de 100644
--- a/test/common/TestBedChild.hpp
+++ b/test/common/TestBedChild.hpp
@@ -91,7 +91,7 @@ namespace RcclUnitTesting
 
   protected:
     // Calls ncclGetUniqueId and returns it to parent
-    ErrCode GetUniqueId();
+    ErrCode GetUniqueId(std::vector<char>& retValBuf);
 
     // Initialize RCCL communicators
     ErrCode InitComms();