Merge remote-tracking branch 'nccl/master' into develop

src/ras/collectives.cc

@@ -606,6 +606,10 @@ static ncclResult_t rasCollCommsInit(struct rasCollRequest** pReq, size_t* pReqL
   for (int commIdx = 0; commIdx < nNcclComms; commIdx++) {
     if (ncclComms[commIdx] == nullptr) // nullptr's are always at the end after sorting.
       break;
+    if (!__atomic_load_n(&ncclComms[commIdx]->peerInfoValid, __ATOMIC_ACQUIRE)) {
+      // Critical data is not yet initialized -- ignore the communicator.
+      continue;
+    }
     // A process may manage multiple GPUs and thus have multiple communicators with the same commHash.
     // Comparing just the commHash is OK though within communicators that are part of the same process.
     if (commIdx == 0 || ncclComms[commIdx]->commHash != ncclComms[commIdx-1]->commHash) {
@@ -651,6 +655,8 @@ static ncclResult_t rasCollCommsInit(struct rasCollRequest** pReq, size_t* pReqL
   // collCommIdx counts rasCollComms::comm (comm); commIdx indexes ncclComms.
   for (int collCommIdx = 0, commIdx = 0; collCommIdx < nComms; collCommIdx++) {
     struct ncclComm* ncclComm = ncclComms[commIdx];
+    if (!__atomic_load_n(&ncclComm->peerInfoValid, __ATOMIC_ACQUIRE))
+      continue;
 
     comm->commId.commHash = ncclComm->commHash;
     comm->commId.hostHash = ncclComm->peerInfo->hostHash;
@@ -663,15 +669,15 @@ static ncclResult_t rasCollCommsInit(struct rasCollRequest** pReq, size_t* pReqL
          commIdx++) {
       ncclComm = ncclComms[commIdx];
       struct rasCollComms::comm::rank* rank = comm->ranks+comm->nRanks;
-      ncclResult_t asyncError;
       rank->commRank = ncclComm->rank;
       // rasNetSendCollReq initializes coll->peers[0] to our rasNetListeningSocket.addr, so peerIdx is initially
       // always 0.  It will increase after we send this response back to the peer we got the request from.
       rank->peerIdx = 0;
       memcpy(rank->collOpCounts, ncclComm->seqNumber, sizeof(rank->collOpCounts));
       rank->status.initState = ncclComm->initState;
-      if (ncclCommGetAsyncError(ncclComm, &asyncError) == ncclSuccess)
-        rank->status.asyncError = asyncError;
+      rank->status.asyncError = __atomic_load_n(&ncclComm->asyncResult, __ATOMIC_ACQUIRE);
+      if (rank->status.asyncError == ncclSuccess && ncclComm->proxyState)
+        rank->status.asyncError = __atomic_load_n(&ncclComm->proxyState->asyncResult, __ATOMIC_ACQUIRE);
       rank->status.finalizeCalled = (ncclComm->finalizeCalled != 0);
       rank->status.destroyFlag = (ncclComm->destroyFlag != 0);
       rank->status.abortFlag = (__atomic_load_n(ncclComm->abortFlag, __ATOMIC_ACQUIRE) != 0);
@@ -680,7 +686,7 @@ static ncclResult_t rasCollCommsInit(struct rasCollRequest** pReq, size_t* pReqL
       comm->nRanks++;
     } // for (commIdx)
 
-    if (firstNewSkipMissingIdx != -1 &&
+    if (__atomic_load_n(&ncclComm->peerInfoValid, __ATOMIC_ACQUIRE) && firstNewSkipMissingIdx != -1 &&
         memcmp(req->comms.skipMissingRanksComms+firstNewSkipMissingIdx, &comm->commId, sizeof(comm->commId)) == 0) {
       // Fill in the missingRanks array that follows the comm->ranks.
       struct rasCollCommsMissingRank* missingRanks = (struct rasCollCommsMissingRank*)(comm->ranks+comm->nRanks);

src/ras/rasnet.cc

@@ -365,15 +365,16 @@ ncclResult_t rasNetAcceptNewSocket() {
   NCCLCHECKGOTO(ncclSocketAccept(&sock->sock, &rasNetListeningSocket), ret, fail);
   NCCLCHECKGOTO(ncclSocketReady(&sock->sock, &ready), ret, fail);
 
-  if (sock->sock.fd != -1) {
-    NCCLCHECKGOTO(rasGetNewPollEntry(&sock->pfd), ret, fail);
-    rasPfds[sock->pfd].fd = sock->sock.fd;
-    rasPfds[sock->pfd].events = POLLIN; // Initially we'll just wait for a handshake from the other side.  This also
-                                        // helps the code tell the sides apart.
-    sock->status = RAS_SOCK_CONNECTING;
+  if (sock->sock.fd == -1)
+    goto fail; // We'll return ncclSuccess, but we need to clean up the incomplete socket first.
 
-    INFO(NCCL_RAS, "RAS new incoming socket connection from %s", ncclSocketToString(&sock->sock.addr, rasLine));
-  }
+  NCCLCHECKGOTO(rasGetNewPollEntry(&sock->pfd), ret, fail);
+  rasPfds[sock->pfd].fd = sock->sock.fd;
+  rasPfds[sock->pfd].events = POLLIN; // Initially we'll just wait for a handshake from the other side.  This also
+                                      // helps the code tell the sides apart.
+  sock->status = RAS_SOCK_CONNECTING;
+
+  INFO(NCCL_RAS, "RAS new incoming socket connection from %s", ncclSocketToString(&sock->sock.addr, rasLine));
 
 exit:
   return ret;
@@ -480,7 +481,10 @@ void rasSocksHandleTimeouts(int64_t now, int64_t* nextWakeup) {
 // Once we get an EOF when receiving data, we finalize the termination.
 // For not fully established sockets, we can terminate immediately as there's no useful data to extract.
 void rasSocketTerminate(struct rasSocket* sock, bool finalize, uint64_t startRetryOffset, bool retry) {
-  assert(sock->status != RAS_SOCK_CLOSED);
+  if (sock->status == RAS_SOCK_CLOSED) {
+    INFO(NCCL_RAS, "RAS socket in closed state passed for termination -- internal error?");
+    // The code below can actually handle such a case gracefully.
+  }
   if (sock->conn) {
     struct rasConnection* conn = sock->conn;
     // If the sock of the connection points back to us, it means that we are the current socket of this
@@ -542,8 +546,10 @@ void rasSocketTerminate(struct rasSocket* sock, bool finalize, uint64_t startRet
   } else {
     // Either the caller requested finalization or we cannot receive on it.
     (void)ncclSocketClose(&sock->sock);
-    rasPfds[sock->pfd].fd = -1;
-    rasPfds[sock->pfd].events = rasPfds[sock->pfd].revents = 0;
+    if (sock->pfd != -1) {
+      rasPfds[sock->pfd].fd = -1;
+      rasPfds[sock->pfd].events = rasPfds[sock->pfd].revents = 0;
+    }
     free(sock->recvMsg);
     freeSockEntry(sock);
   }