2.12.7-1

Add network communication through another GPU connected with NVLink
(PXN).
Add aggregation of messages coming from different local GPUs through
PXN and going to the same destination.
Add new v5 plugin API with grouped receives and tags.
Add compat for v4 plugins.
Add naming of NCCL threads to help debugging.
Fix NVLink detection and avoid data corruption when some NVLinks are
down.
Add support for Relaxed Ordering for IB.
Add profiling and timing infrastructure.

src/bootstrap.cc

@@ -1,5 +1,5 @@
 /*************************************************************************
- * Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -9,13 +9,13 @@
 #include "utils.h"
 #include "bootstrap.h"
 #include "net.h"
-#include "socket.h"
 #include <unistd.h>
 #include <sys/types.h>
+#include "proxy.h"
 
 /* Init functions */
 static char bootstrapNetIfName[MAX_IF_NAME_SIZE+1];
-static union socketAddress bootstrapNetIfAddr;
+static union ncclSocketAddress bootstrapNetIfAddr;
 static int bootstrapNetInitDone = 0;
 pthread_mutex_t bootstrapNetLock = PTHREAD_MUTEX_INITIALIZER;
 
@@ -25,17 +25,17 @@ ncclResult_t bootstrapNetInit() {
     if (bootstrapNetInitDone == 0) {
       char* env = getenv("NCCL_COMM_ID");
       if (env) {
-        union socketAddress remoteAddr;
-        if (GetSocketAddrFromString(&remoteAddr, env) != ncclSuccess) {
+        union ncclSocketAddress remoteAddr;
+        if (ncclGetSocketAddrFromString(&remoteAddr, env) != ncclSuccess) {
           WARN("Invalid NCCL_COMM_ID, please use format: <ipv4>:<port> or [<ipv6>]:<port> or <hostname>:<port>");
           return ncclInvalidArgument;
         }
-        if (findInterfaceMatchSubnet(bootstrapNetIfName, &bootstrapNetIfAddr, &remoteAddr, MAX_IF_NAME_SIZE, 1) <= 0) {
+        if (ncclFindInterfaceMatchSubnet(bootstrapNetIfName, &bootstrapNetIfAddr, &remoteAddr, MAX_IF_NAME_SIZE, 1) <= 0) {
           WARN("NET/Socket : No usable listening interface found");
           return ncclSystemError;
         }
       } else {
-        int nIfs = findInterfaces(bootstrapNetIfName, &bootstrapNetIfAddr, MAX_IF_NAME_SIZE, 1);
+        int nIfs = ncclFindInterfaces(bootstrapNetIfName, &bootstrapNetIfAddr, MAX_IF_NAME_SIZE, 1);
         if (nIfs <= 0) {
           WARN("Bootstrap : no socket interface found");
           return ncclInternalError;
@@ -43,7 +43,7 @@ ncclResult_t bootstrapNetInit() {
       }
       char line[SOCKET_NAME_MAXLEN+MAX_IF_NAME_SIZE+2];
       sprintf(line, " %s:", bootstrapNetIfName);
-      socketToString(&bootstrapNetIfAddr, line+strlen(line));
+      ncclSocketToString(&bootstrapNetIfAddr, line+strlen(line));
       INFO(NCCL_INIT, "Bootstrap : Using%s", line);
       bootstrapNetInitDone = 1;
     }
@@ -55,35 +55,28 @@ ncclResult_t bootstrapNetInit() {
 /* Socket Interface Selection type */
 enum bootstrapInterface_t { findSubnetIf = -1, dontCareIf = -2 };
 
-static ncclResult_t bootstrapNetAccept(int listenFd, int* recvFd, union socketAddress *addr) {
-  struct sockaddr *saddr = &addr->sa;
-  socklen_t socklen = sizeof(union socketAddress);
-  SYSCHECKVAL(accept(listenFd, saddr, &socklen), "accept", *recvFd);
-  return ncclSuccess;
-}
-
 // Additional sync functions
-static ncclResult_t bootstrapNetSend(int fd, union socketAddress *addr, void* data, int size) {
-  NCCLCHECK(socketSend(fd, addr, &size, sizeof(int)));
-  NCCLCHECK(socketSend(fd, addr, data, size));
+static ncclResult_t bootstrapNetSend(struct ncclSocket* sock, void* data, int size) {
+  NCCLCHECK(ncclSocketSend(sock, &size, sizeof(int)));
+  NCCLCHECK(ncclSocketSend(sock, data, size));
   return ncclSuccess;
 }
-static ncclResult_t bootstrapNetRecv(int fd, union socketAddress *addr, void* data, int size) {
+static ncclResult_t bootstrapNetRecv(struct ncclSocket* sock, void* data, int size) {
   int recvSize;
-  NCCLCHECK(socketRecv(fd, addr, &recvSize, sizeof(int)));
+  NCCLCHECK(ncclSocketRecv(sock, &recvSize, sizeof(int)));
   if (recvSize > size) {
     WARN("Message truncated : received %d bytes instead of %d", recvSize, size);
     return ncclInternalError;
   }
-  NCCLCHECK(socketRecv(fd, addr, data, std::min(recvSize, size)));
+  NCCLCHECK(ncclSocketRecv(sock, data, std::min(recvSize, size)));
   return ncclSuccess;
 }
 
 struct extInfo {
   int rank;
   int nranks;
-  union socketAddress extAddressListenRoot;
-  union socketAddress extAddressListen;
+  union ncclSocketAddress extAddressListenRoot;
+  union ncclSocketAddress extAddressListen;
 };
 
 #include <sys/resource.h>
@@ -97,24 +90,24 @@ static ncclResult_t setFilesLimit() {
 }
 
 static void *bootstrapRoot(void* args) {
-  int listenFd = (uint64_t)args;
+  struct ncclSocket* listenSock = (struct ncclSocket*)args;
   ncclResult_t res = ncclSuccess;
   int nranks = 0, c = 0;
   struct extInfo info;
-  union socketAddress *rankAddresses = NULL;
-  union socketAddress *rankAddressesRoot = NULL; // for initial rank <-> root information exchange
-  union socketAddress *zero = NULL;
+  union ncclSocketAddress *rankAddresses = NULL;
+  union ncclSocketAddress *rankAddressesRoot = NULL; // for initial rank <-> root information exchange
+  union ncclSocketAddress *zero = NULL;
   NCCLCHECKGOTO(ncclCalloc(&zero, 1), res, out);
   setFilesLimit();
 
   TRACE(NCCL_INIT, "BEGIN");
   /* Receive addresses from all ranks */
   do {
-    int tmpFd;
-    union socketAddress addr;
-    NCCLCHECKGOTO(bootstrapNetAccept(listenFd, &tmpFd, &addr), res, out);
-    NCCLCHECKGOTO(bootstrapNetRecv(tmpFd, &addr, &info, sizeof(info)), res, out);
-    close(tmpFd);
+    struct ncclSocket sock;
+    sock.abortFlag = NULL;
+    NCCLCHECKGOTO(ncclSocketAccept(&sock, listenSock), res, out);
+    NCCLCHECKGOTO(bootstrapNetRecv(&sock, &info, sizeof(info)), res, out);
+    close(sock.fd);
 
     if (c == 0) {
       nranks = info.nranks;
@@ -127,14 +120,14 @@ static void *bootstrapRoot(void* args) {
       goto out;
     }
 
-    if (memcmp(zero, &rankAddressesRoot[info.rank], sizeof(union socketAddress)) != 0) {
+    if (memcmp(zero, &rankAddressesRoot[info.rank], sizeof(union ncclSocketAddress)) != 0) {
       WARN("Bootstrap Root : rank %d of %d ranks has already checked in", info.rank, nranks);
       goto out;
     }
 
     // Save the connection handle for that rank
-    memcpy(rankAddressesRoot+info.rank, &info.extAddressListenRoot, sizeof(union socketAddress));
-    memcpy(rankAddresses+info.rank, &info.extAddressListen, sizeof(union socketAddress));
+    memcpy(rankAddressesRoot+info.rank, &info.extAddressListenRoot, sizeof(union ncclSocketAddress));
+    memcpy(rankAddresses+info.rank, &info.extAddressListen, sizeof(union ncclSocketAddress));
 
     ++c;
     TRACE(NCCL_INIT, "Received connect from rank %d total %d/%d",  info.rank, c, nranks);
@@ -144,15 +137,18 @@ static void *bootstrapRoot(void* args) {
   // Send the connect handle for the next rank in the AllGather ring
   for (int r=0; r<nranks; ++r) {
     int next = (r+1) % nranks;
-    int tmpSendFd;
-    NCCLCHECKGOTO(connectAddress(&tmpSendFd, rankAddressesRoot+r), res, out);
-    NCCLCHECKGOTO(bootstrapNetSend(tmpSendFd, rankAddressesRoot+r, rankAddresses+next, sizeof(union socketAddress)), res, out);
-    close(tmpSendFd);
+    struct ncclSocket sock;
+    sock.abortFlag = NULL;
+    memcpy(&sock.addr, rankAddressesRoot+r, sizeof(union ncclSocketAddress));
+    NCCLCHECKGOTO(ncclSocketConnect(&sock), res, out);
+    NCCLCHECKGOTO(bootstrapNetSend(&sock, rankAddresses+next, sizeof(union ncclSocketAddress)), res, out);
+    close(sock.fd);
   }
   TRACE(NCCL_INIT, "SENT OUT ALL %d HANDLES", nranks);
 
 out:
-  close(listenFd);
+  close(listenSock->fd);
+  free(listenSock);
   if (rankAddresses) free(rankAddresses);
   if (rankAddressesRoot) free(rankAddressesRoot);
   if (zero) free(zero);
@@ -162,28 +158,31 @@ out:
 }
 
 ncclResult_t bootstrapCreateRoot(ncclUniqueId* id, bool idFromEnv) {
-  union socketAddress* connectAddr = (union socketAddress*) id;
-  int listenFd;
-  NCCLCHECK(createListenSocket(&listenFd, connectAddr));
+  struct ncclSocket* listenSock;
+  NCCLCHECK(ncclCalloc(&listenSock, 1));
+  memcpy(&listenSock->addr, id, sizeof(union ncclSocketAddress));
+  NCCLCHECK(ncclSocketListen(listenSock));
+  memcpy(id, &listenSock->addr, sizeof(union ncclSocketAddress));
   pthread_t thread;
-  pthread_create(&thread, NULL, bootstrapRoot, (void*)(uint64_t)listenFd);
+  pthread_create(&thread, NULL, bootstrapRoot, (void*)listenSock);
+  ncclSetThreadName(thread, "NCCL BootstrapR");
   return ncclSuccess;
 }
 
 ncclResult_t bootstrapGetUniqueId(ncclUniqueId* id) {
-  static_assert(sizeof(union socketAddress) < sizeof(ncclUniqueId), "NetId does not fit inside ncclUniqueId");
+  static_assert(sizeof(union ncclSocketAddress) < sizeof(ncclUniqueId), "NetId does not fit inside ncclUniqueId");
   memset(id, 0, sizeof(ncclUniqueId));
-  union socketAddress* connectAddr = (union socketAddress*) id;
+  union ncclSocketAddress* connectAddr = (union ncclSocketAddress*) id;
 
   char* env = getenv("NCCL_COMM_ID");
   if (env) {
     INFO(NCCL_ENV, "NCCL_COMM_ID set by environment to %s", env);
-    if (GetSocketAddrFromString(connectAddr, env) != ncclSuccess) {
+    if (ncclGetSocketAddrFromString(connectAddr, env) != ncclSuccess) {
       WARN("Invalid NCCL_COMM_ID, please use format: <ipv4>:<port> or [<ipv6>]:<port> or <hostname>:<port>");
       return ncclInvalidArgument;
     }
   } else {
-    memcpy(id, &bootstrapNetIfAddr, sizeof(union socketAddress));
+    memcpy(id, &bootstrapNetIfAddr, sizeof(union ncclSocketAddress));
     NCCLCHECK(bootstrapCreateRoot(id, false));
   }
 
@@ -193,157 +192,51 @@ ncclResult_t bootstrapGetUniqueId(ncclUniqueId* id) {
 struct unexConn {
   int peer;
   int tag;
-  int fd;
-  union socketAddress addr;
+  struct ncclSocket sock;
   struct unexConn* next;
 };
 
-// Remote allocator state
-struct remAllocState {
-  int cudaDev;
-  int listenFd;
-  volatile int stop;
-};
-
-struct extState {
-  int extListenFd;
-  int extRingRecvFd;
-  int extRingSendFd;
-  union socketAddress extRingRecvAddr, extRingSendAddr;
-  union socketAddress* peerCommAddresses;
-  union socketAddress* peerAllocAddresses;
+struct bootstrapState {
+  struct ncclSocket listenSock;
+  struct ncclSocket ringRecvSocket;
+  struct ncclSocket ringSendSocket;
+  union ncclSocketAddress* peerCommAddresses;
+  union ncclSocketAddress* peerProxyAddresses;
   struct unexConn* unexpectedConnections;
   int cudaDev;
   int rank;
   int nranks;
-
-  // Intermediate memory allocation service
-  struct remAllocState* allocState;
-  pthread_t allocThread;
+  volatile uint32_t *abortFlag;
 };
 
-#define MAX_SEGMENTS 128
-
-static ncclResult_t remoteAlloc(void** ptr, int fd, union socketAddress *addr) {
-  size_t size;
-  NCCLCHECK(socketRecv(fd, addr, &size, sizeof(size_t)));
-  cudaIpcMemHandle_t devIpc;
-  NCCLCHECK(ncclCudaCalloc((char**)ptr, size));
-  cudaError_t res = cudaIpcGetMemHandle(&devIpc, *ptr);
-  if (res != cudaSuccess) {
-    WARN("[Rem Allocator] cudaIpcGetMemHandle failed : %s", cudaGetErrorString(res));
-    cudaFree(*ptr);
-    CUDACHECK(res);
-  }
-  // The CUDA IPC
-  NCCLCHECK(socketSend(fd, addr, &devIpc, sizeof(cudaIpcMemHandle_t)));
-  // And the direct pointer
-  NCCLCHECK(socketSend(fd, addr, ptr, sizeof(void*)));
-  return ncclSuccess;
-}
-
-#include <poll.h>
-
-// Service thread to allocate memory for other GPUs, used as intermediate step.
-void* ncclRemoteMemAllocationService(void* args) {
-  struct remAllocState* state = (struct remAllocState *) args;
-  if (cudaSetDevice(state->cudaDev) != cudaSuccess) {
-    WARN("[Rem Allocator] Failed to set CUDA device %d", state->cudaDev);
-  }
-
-  // Prepare poll descriptor
-  void* segments[MAX_SEGMENTS];
-  struct pollfd pollfds[MAX_SEGMENTS+1];
-  for (int s=0; s<MAX_SEGMENTS; s++) segments[s] = NULL;
-  for (int s=0; s<MAX_SEGMENTS; s++) {
-    pollfds[s].fd = -1;
-    pollfds[s].events = POLLIN;
-  }
-  pollfds[MAX_SEGMENTS].fd = state->listenFd;
-  pollfds[MAX_SEGMENTS].events = POLLIN;
-
-  int nbuffers = 0;
-  while (state->stop == 0 || (state->stop == 1 && nbuffers > 0)) {
-    if (int error = poll(pollfds, MAX_SEGMENTS+1, 100/*ms*/) < 0) {
-      WARN("[Rem Allocator] Poll failed with error %d", error);
-      return NULL;
-    }
-    if (pollfds[MAX_SEGMENTS].revents) {
-      int s = 0;
-      union socketAddress addr;
-      while (segments[s] != NULL && s < MAX_SEGMENTS) s++;
-      if (bootstrapNetAccept(pollfds[MAX_SEGMENTS].fd, &pollfds[s].fd, &addr) != ncclSuccess) {
-        pollfds[s].fd = -1;
-      } else {
-        if (s == MAX_SEGMENTS || (remoteAlloc(segments+s, pollfds[s].fd, &addr) != ncclSuccess)) {
-          WARN("[Rem Allocator] Allocation failed (segment %d, fd %d)", s, pollfds[s].fd);
-          close(pollfds[s].fd);
-          pollfds[s].fd = -1;
-        } else {
-          nbuffers++;
-        }
-      }
-    }
-    for (int s=0; s<MAX_SEGMENTS; s++) {
-      if (pollfds[s].revents & (POLLIN|POLLHUP)) {
-        if (cudaFree(segments[s]) != cudaSuccess) {
-          WARN("[Rem Allocator] cudaFree %p failed", segments[s]);
-        }
-        segments[s] = NULL;
-        close(pollfds[s].fd);
-        pollfds[s].fd = -1;
-        nbuffers--;
-      }
-    }
-  }
-  for (int s=0; s<MAX_SEGMENTS; s++) {
-    if (segments[s]) cudaFree(segments[s]);
-    close(pollfds[s].fd);
-  }
-  close(state->listenFd);
-  free(state);
-  return NULL;
-}
-
-ncclResult_t bootstrapRemAlloc(size_t size, int rank, void* commState, int* id, cudaIpcMemHandle_t* ipc, void** ptr) {
-  struct extState* state = (struct extState*)commState;
-  int fd;
-  ncclResult_t res;
-  *id = -1;
-  union socketAddress *addr = state->peerAllocAddresses+rank;
-  NCCLCHECK(connectAddress(&fd, addr));
-  NCCLCHECKGOTO(socketSend(fd, addr, &size, sizeof(size_t)), res, end);
-  NCCLCHECKGOTO(socketRecv(fd, addr, ipc, sizeof(cudaIpcMemHandle_t)), res, end);
-  NCCLCHECKGOTO(socketRecv(fd, addr, ptr, sizeof(void*)), res, end);
-  *id = fd;
-end:
-  return res;
-}
-
-ncclResult_t bootstrapRemFree(int id, int rank, void* commState) {
-  SYSCHECK(close(id), "close");
-  return ncclSuccess;
-}
-
-ncclResult_t bootstrapInit(ncclUniqueId * id, int rank, int nranks, void** commState) {
-  struct extState* state;
+ncclResult_t bootstrapInit(ncclUniqueId * id, struct ncclComm* comm) {
+  int rank = comm->rank;
+  int nranks = comm->nRanks;
+  struct bootstrapState* state;
   NCCLCHECK(ncclCalloc(&state, 1));
   state->rank = rank;
   state->nranks = nranks;
-  *commState = state;
+  state->abortFlag = comm->abortFlag;
+  comm->bootstrap = state;
 
   TRACE(NCCL_INIT, "rank %d nranks %d", rank, nranks);
 
   struct extInfo info = { 0 };
   info.rank = rank;
   info.nranks = nranks;
-  int tmpSendFd, tmpRecvFd;
+  struct ncclSocket sock, listenSockRoot;
+  sock.abortFlag = listenSockRoot.abortFlag = comm->abortFlag;
+  sock.asyncFlag = listenSockRoot.asyncFlag = 0;
 
-  int extListenFdRoot;
-  memcpy(&info.extAddressListen,     &bootstrapNetIfAddr, sizeof(union socketAddress));
-  memcpy(&info.extAddressListenRoot, &bootstrapNetIfAddr, sizeof(union socketAddress));
-  NCCLCHECK(createListenSocket(&state->extListenFd, &info.extAddressListen));
-  NCCLCHECK(createListenSocket(&extListenFdRoot, &info.extAddressListenRoot));
+  // Create socket for other ranks to contact me
+  memcpy(&state->listenSock.addr, &bootstrapNetIfAddr, sizeof(union ncclSocketAddress));
+  NCCLCHECK(ncclSocketListen(&state->listenSock));
+  memcpy(&info.extAddressListen, &state->listenSock.addr, sizeof(union ncclSocketAddress));
+
+  // Create socket for root to contact me
+  memcpy(&listenSockRoot.addr, &bootstrapNetIfAddr, sizeof(union ncclSocketAddress));
+  NCCLCHECK(ncclSocketListen(&listenSockRoot));
+  memcpy(&info.extAddressListenRoot, &listenSockRoot.addr, sizeof(union ncclSocketAddress));
 
   // stagger connection times to avoid an overload of the root
   if (nranks > 128) {
@@ -356,35 +249,36 @@ ncclResult_t bootstrapInit(ncclUniqueId * id, int rank, int nranks, void** commS
   }
 
   // send info on my listening socket to root
-  union socketAddress* rootAddr = (union socketAddress*)id;
-  NCCLCHECK(connectAddress(&tmpSendFd, rootAddr));
-  NCCLCHECK(bootstrapNetSend(tmpSendFd, rootAddr,  &info, sizeof(info)));
-  close(tmpSendFd);
+  memcpy(&sock.addr, id, sizeof(union ncclSocketAddress));
+  NCCLCHECK(ncclSocketConnect(&sock));
+  NCCLCHECK(bootstrapNetSend(&sock, &info, sizeof(info)));
+  close(sock.fd);
 
   // get info on my "next" rank in the bootstrap ring from root
-  union socketAddress addr;
-  NCCLCHECK(bootstrapNetAccept(extListenFdRoot, &tmpRecvFd, &addr));
-  NCCLCHECK(bootstrapNetRecv(tmpRecvFd, &addr, &state->extRingSendAddr, sizeof(state->extRingSendAddr)));
-  close(tmpRecvFd);
-  close(extListenFdRoot);
+  NCCLCHECK(ncclSocketAccept(&sock, &listenSockRoot));
+  NCCLCHECK(bootstrapNetRecv(&sock, &state->ringSendSocket.addr, sizeof(union ncclSocketAddress)));
+  close(sock.fd);
+  close(listenSockRoot.fd);
 
-  NCCLCHECK(connectAddress(&state->extRingSendFd, &state->extRingSendAddr));
+  NCCLCHECK(ncclSocketConnect(&state->ringSendSocket));
   // Accept the connect request from the previous rank in the AllGather ring
-  NCCLCHECK(bootstrapNetAccept(state->extListenFd, &state->extRingRecvFd, &state->extRingRecvAddr));
+  NCCLCHECK(ncclSocketAccept(&state->ringRecvSocket, &state->listenSock));
 
   // AllGather all listen handlers
   NCCLCHECK(ncclCalloc(&state->peerCommAddresses, nranks));
-  memcpy(state->peerCommAddresses+rank, &info.extAddressListen, sizeof(union socketAddress));
-  NCCLCHECK(bootstrapAllGather(state, state->peerCommAddresses, sizeof(union socketAddress)));
+  memcpy(state->peerCommAddresses+rank, &state->listenSock.addr, sizeof(union ncclSocketAddress));
+  NCCLCHECK(bootstrapAllGather(state, state->peerCommAddresses, sizeof(union ncclSocketAddress)));
 
-  // Create the memory allocation service
-  NCCLCHECK(ncclCalloc(&state->peerAllocAddresses, nranks));
-  memcpy(state->peerAllocAddresses+rank, &bootstrapNetIfAddr, sizeof(union socketAddress));
-  NCCLCHECK(ncclCalloc(&state->allocState, 1));
-  CUDACHECK(cudaGetDevice(&state->allocState->cudaDev));
-  NCCLCHECK(createListenSocket(&state->allocState->listenFd, state->peerAllocAddresses+rank));
-  pthread_create(&state->allocThread, NULL, ncclRemoteMemAllocationService, state->allocState);
-  NCCLCHECK(bootstrapAllGather(state, state->peerAllocAddresses, sizeof(union socketAddress)));
+  // Create the service proxy
+  NCCLCHECK(ncclCalloc(&state->peerProxyAddresses, nranks));
+  struct ncclSocket* proxySocket;
+  NCCLCHECK(ncclCalloc(&proxySocket, 1));
+  proxySocket->abortFlag = NULL; // proxy is aborted through a message
+  memcpy(&proxySocket->addr, &bootstrapNetIfAddr, sizeof(union ncclSocketAddress));
+  NCCLCHECK(ncclSocketListen(proxySocket));
+  memcpy(state->peerProxyAddresses+rank, &proxySocket->addr, sizeof(union ncclSocketAddress));
+  NCCLCHECK(bootstrapAllGather(state, state->peerProxyAddresses, sizeof(union ncclSocketAddress)));
+  NCCLCHECK(ncclProxyInit(comm, proxySocket, state->peerProxyAddresses));
 
   TRACE(NCCL_INIT, "rank %d nranks %d - DONE", rank, nranks);
 
@@ -392,7 +286,7 @@ ncclResult_t bootstrapInit(ncclUniqueId * id, int rank, int nranks, void** commS
 }
 
 ncclResult_t bootstrapAllGather(void* commState, void* allData, int size) {
-  struct extState* state = (struct extState*)commState;
+  struct bootstrapState* state = (struct bootstrapState*)commState;
   char* data = (char*)allData;
   int rank = state->rank;
   int nranks = state->nranks;
@@ -408,9 +302,9 @@ ncclResult_t bootstrapAllGather(void* commState, void* allData, int size) {
     size_t sslice = (rank - i + nranks) % nranks;
 
     // Send slice to the right
-    NCCLCHECK(bootstrapNetSend(state->extRingSendFd, &state->extRingSendAddr, data+sslice*size, size));
+    NCCLCHECK(bootstrapNetSend(&state->ringSendSocket, data+sslice*size, size));
     // Recv slice from the left
-    NCCLCHECK(bootstrapNetRecv(state->extRingRecvFd, &state->extRingRecvAddr, data+rslice*size, size));
+    NCCLCHECK(bootstrapNetRecv(&state->ringRecvSocket, data+rslice*size, size));
   }
 
   TRACE(NCCL_INIT, "rank %d nranks %d size %d - DONE", rank, nranks, size);
@@ -418,14 +312,15 @@ ncclResult_t bootstrapAllGather(void* commState, void* allData, int size) {
 }
 
 ncclResult_t bootstrapSend(void* commState, int peer, int tag, void* data, int size) {
-  struct extState* state = (struct extState*)commState;
-  int tmpSendFd;
-  union socketAddress *addr = state->peerCommAddresses+peer;
-  NCCLCHECK(connectAddress(&tmpSendFd, addr));
-  NCCLCHECK(bootstrapNetSend(tmpSendFd, addr, &state->rank, sizeof(int)));
-  NCCLCHECK(bootstrapNetSend(tmpSendFd, addr, &tag, sizeof(int)));
-  NCCLCHECK(bootstrapNetSend(tmpSendFd, addr, data, size));
-  close(tmpSendFd);
+  struct bootstrapState* state = (struct bootstrapState*)commState;
+  struct ncclSocket sock;
+  sock.abortFlag = state->abortFlag;
+  memcpy(&sock.addr, state->peerCommAddresses+peer, sizeof(union ncclSocketAddress));
+  NCCLCHECK(ncclSocketConnect(&sock));
+  NCCLCHECK(bootstrapNetSend(&sock, &state->rank, sizeof(int)));
+  NCCLCHECK(bootstrapNetSend(&sock, &tag, sizeof(int)));
+  NCCLCHECK(bootstrapNetSend(&sock, data, size));
+  close(sock.fd);
   return ncclSuccess;
 }
 
@@ -466,14 +361,13 @@ ncclResult_t bootstrapIntraNodeAllGather(void* commState, int *ranks, int rank,
   return ncclSuccess;
 }
 
-ncclResult_t unexpectedEnqueue(struct extState* state, int peer, int tag, int fd, union socketAddress *addr) {
+ncclResult_t unexpectedEnqueue(struct bootstrapState* state, int peer, int tag, struct ncclSocket* sock) {
   // New unex
   struct unexConn* unex;
   NCCLCHECK(ncclCalloc(&unex, 1));
   unex->peer = peer;
   unex->tag = tag;
-  unex->fd = fd;
-  unex->addr = *addr;
+  memcpy(&unex->sock, sock, sizeof(struct ncclSocket));
 
   // Enqueue
   struct unexConn* list = state->unexpectedConnections;
@@ -486,7 +380,7 @@ ncclResult_t unexpectedEnqueue(struct extState* state, int peer, int tag, int fd
   return ncclSuccess;
 }
 
-int unexpectedDequeue(struct extState* state, int peer, int tag, union socketAddress *addr) {
+ncclResult_t unexpectedDequeue(struct bootstrapState* state, int peer, int tag, struct ncclSocket* sock) {
   struct unexConn* elem = state->unexpectedConnections;
   struct unexConn* prev = NULL;
   while (elem) {
@@ -496,79 +390,72 @@ int unexpectedDequeue(struct extState* state, int peer, int tag, union socketAdd
       } else {
         prev->next = elem->next;
       }
-      int fd = elem->fd;
-      *addr = elem->addr;
+      memcpy(sock, &elem->sock, sizeof(struct ncclSocket));
       free(elem);
-      return fd;
+      return ncclSuccess;
     }
     prev = elem;
     elem = elem->next;
   }
-  return -1;
+  sock->fd = -1;
+  return ncclSuccess;
 }
 
 // We can't know who we'll receive from, so we need to receive everything at once
 ncclResult_t bootstrapRecv(void* commState, int peer, int tag, void* data, int size) {
-  struct extState* state = (struct extState*)commState;
+  struct bootstrapState* state = (struct bootstrapState*)commState;
 
-  int tmpRecvFd;
-  union socketAddress addr;
+  struct ncclSocket sock;
+  sock.abortFlag = state->abortFlag;
 
   // Search unexpected connections first
-  if ((tmpRecvFd = unexpectedDequeue(state, peer, tag, &addr)) != -1) {
-    NCCLCHECK(bootstrapNetRecv(tmpRecvFd, &addr, ((char*)data), size));
-    close(tmpRecvFd);
+  NCCLCHECK(unexpectedDequeue(state, peer, tag, &sock));
+  if (sock.fd != -1) {
+    NCCLCHECK(bootstrapNetRecv(&sock, ((char*)data), size));
+    close(sock.fd);
     return ncclSuccess;
   }
 
   // Then look for new connections
   while (1) {
-    union socketAddress addr;
-    NCCLCHECK(bootstrapNetAccept(state->extListenFd, &tmpRecvFd, &addr));
+    NCCLCHECK(ncclSocketAccept(&sock, &state->listenSock));
     int newPeer, newTag;
-    NCCLCHECK(bootstrapNetRecv(tmpRecvFd, &addr, &newPeer, sizeof(int)));
-    NCCLCHECK(bootstrapNetRecv(tmpRecvFd, &addr, &newTag, sizeof(int)));
+    NCCLCHECK(bootstrapNetRecv(&sock, &newPeer, sizeof(int)));
+    NCCLCHECK(bootstrapNetRecv(&sock, &newTag, sizeof(int)));
     if (newPeer == peer && newTag == tag) {
-      NCCLCHECK(bootstrapNetRecv(tmpRecvFd, &addr, ((char*)data), size));
-      close(tmpRecvFd);
+      NCCLCHECK(bootstrapNetRecv(&sock, ((char*)data), size));
+      close(sock.fd);
       return ncclSuccess;
     }
     // Unexpected connection. Save for later.
-    NCCLCHECK(unexpectedEnqueue(state, newPeer, newTag, tmpRecvFd, &addr));
+    NCCLCHECK(unexpectedEnqueue(state, newPeer, newTag, &sock));
   }
 }
 
 ncclResult_t bootstrapClose(void* commState) {
-  struct extState* state = (struct extState*)commState;
+  struct bootstrapState* state = (struct bootstrapState*)commState;
   if (state->unexpectedConnections != NULL) {
     WARN("Unexpected connections are not empty");
     return ncclInternalError;
   }
-  close(state->extListenFd);
-  close(state->extRingSendFd);
-  close(state->extRingRecvFd);
-
-  state->allocState->stop = 1;
-
-  // Join the allocThread so we catch resource leaks as being hung here
-  // pthread_join(state->allocThread, nullptr);
+  close(state->listenSock.fd);
+  close(state->ringSendSocket.fd);
+  close(state->ringRecvSocket.fd);
 
   free(state->peerCommAddresses);
-  free(state->peerAllocAddresses);
   free(state);
 
   return ncclSuccess;
 }
 
 ncclResult_t bootstrapAbort(void* commState) {
-  struct extState* state = (struct extState*)commState;
+  struct bootstrapState* state = (struct bootstrapState*)commState;
   if (commState == NULL) return ncclSuccess;
-  if (state->extListenFd) close(state->extListenFd);
-  if (state->extRingSendFd) close(state->extRingSendFd);
-  if (state->extRingRecvFd) close(state->extRingRecvFd);
-  if (state->allocState) state->allocState->stop = 2;
+  if (state->listenSock.fd) close(state->listenSock.fd);
+  if (state->ringSendSocket.fd) close(state->ringSendSocket.fd);
+  if (state->ringRecvSocket.fd) close(state->ringRecvSocket.fd);
   free(state->peerCommAddresses);
-  free(state->peerAllocAddresses);
+  free(state->peerProxyAddresses);
   free(state);
   return ncclSuccess;
 }