NCCL源码解析1——生成ncclUniqueId

NCCL源码解析1——生成ncclUniqueId

——lvyilong316

NCCL 通信库通过 ncclGetUniqueId 函数生成并获取 「ncclUniqueId」。本质上，ncclUniqueId 是一个大小为 128 字节数组，如下所示。

#define NCCL_UNIQUE_ID_BYTES 128

typedef struct { char internal[NCCL_UNIQUE_ID_BYTES]; } ncclUniqueId;

ncclGetUniqueId逻辑如下，ncclGetUniqueId 首先调用 ncclInit 函数初始化网络组件，接着调用 bootstrapGetUniqueId 函数生成 ncclUniqueId。

ncclResult_t ncclGetUniqueId(ncclUniqueId* out) {
  NCCLCHECK(ncclInit());
  NCCLCHECK(PtrCheck(out, "GetUniqueId", "out"));
  ncclResult_t res = bootstrapGetUniqueId((struct ncclBootstrapHandle*)out);
  TRACE_CALL("ncclGetUniqueId(0x%llx)", (unsigned long long)hashUniqueId(*out));
  return res;
}

l  ncclInit

我们首先来看 ncclInit 初始化函数。NCCL 使用两套网络，一个是 bootstrap 网络，主要用于初始化时交互一些简单的信息，如机器 IP 地址和端口等，且大部分只在初始化阶段执行一次，因此通常采用 TCP 实现。另一个是通信网络，用于实际通信数据的交互，因此优先使用 GdrCopy 和 RDMA。ncclInit 函数主要用于初始化 bootstrap 网络的 IP 地址信息。该函数调用 initOnceFunc 函数，并通过pthread_once保证 initOnceFunc 在多线程环境中只执行一次。

static ncclResult_t ncclInit() {
  pthread_once(&initOnceControl, initOnceFunc);
  return initResult;
}

在 initOnceFunc 函数中，首先通过 initEnv 函数读取 ~/.nccl.conf 或者 /etc/nccl.conf 文件中的 NCCL 环境变量，这和在用户脚本中通过 export 设置环境变量是等价的。

static void initOnceFunc() {
  initEnv(); //读取 ~/.nccl.conf 或者 /etc/nccl.conf 文件中的 NCCL 环境变量
  initGdrCopy(); //加载libgdrapi.so，初始化GDR相关操作函数
  // Always initialize bootstrap network
  NCCLCHECKGOTO(bootstrapNetInit(), initResult, exit);

  initNvtxRegisteredEnums();
exit:;
}

下面，给出 nccl.conf 文件的一个简单示例：

NCCL_SOCKET_IFNAME=eth0

NCCL_DEBUG=WARN

NCCL_RINGS=0 1 2 3|3 2 1 0

接着，调用 initGdrCopy 函数，而initGdrCopy又调用ncclGdrInit初始化 Gdr相关函数。具体来说，ncclGdrInit函数加载 libgdrapi.so动态库，并初始化相关的操作函数，如 gdr_open、gdr_close等。接着，判断 GdrCopy 库和驱动版本的兼容性，当前 NCCL 要求 GDRAPI 2.1 及其之后版本。

ncclResult_t initGdrCopy() {
  if (ncclParamGdrCopyEnable() == 1) {
    ncclGdrCopy = ncclGdrInit();
  }
  return ncclSuccess;
}

之后，调用 bootstrapNetInit 初始化 bootstrap 网络。bootstrapNetInit 函数用于初始化 bootstrap 网络。如下所示：

ncclResult_t bootstrapNetInit() {
  if (bootstrapNetInitDone == 0) {
    pthread_mutex_lock(&bootstrapNetLock);
    if (bootstrapNetInitDone == 0) {
      /* 如果设置了 NCCL_COMM_ID 环境变量（NCCL_COMM_ID的格式为“nccl_socket:ip:port”，其中ip指定要使用的IP地址，port指定要使用的端口号），则查找一个和该环境变量中指定的 IP 地址处于同一子网的网卡作为 booststrap 网络通信所使用的网卡 bootstrapNetIfAddr。 */
      const char* env = ncclGetEnv("NCCL_COMM_ID");
      if (env) {
        union ncclSocketAddress remoteAddr;
        if (ncclSocketGetAddrFromString(&remoteAddr, env) != ncclSuccess) {
          WARN("Invalid NCCL_COMM_ID, please use format: : or []: or :");
          pthread_mutex_unlock(&bootstrapNetLock);
          return ncclInvalidArgument;
        }
        if (ncclFindInterfaceMatchSubnet(bootstrapNetIfName, &bootstrapNetIfAddr, &remoteAddr, MAX_IF_NAME_SIZE, 1) <= 0) {
          WARN("NET/Socket : No usable listening interface found");
          pthread_mutex_unlock(&bootstrapNetLock);
          return ncclSystemError;
        }
      } else {
        int nIfs = ncclFindInterfaces(bootstrapNetIfName, &bootstrapNetIfAddr, MAX_IF_NAME_SIZE, 1);
        if (nIfs <= 0) {
          WARN("Bootstrap : no socket interface found");
          pthread_mutex_unlock(&bootstrapNetLock);
          return ncclInternalError;
        }
      }
      char line[SOCKET_NAME_MAXLEN+MAX_IF_NAME_SIZE+2];
      sprintf(line, " %s:", bootstrapNetIfName);
      ncclSocketToString(&bootstrapNetIfAddr, line+strlen(line));
      INFO(NCCL_INIT, "Bootstrap : Using%s", line);
      bootstrapNetInitDone = 1;
    }
    pthread_mutex_unlock(&bootstrapNetLock);
  }
  return ncclSuccess;
}

    首先，如果设置了 NCCL_COMM_ID 环境变量，则查找一个和该环境变量中指定的 IP 地址处于同一子网的网卡作为 booststrap 网络通信所使用的网卡 bootstrapNetIfAddr。

否则，使用 ncclFindInterfaces 函数选择一个合适的网卡。这里，如果用户通过 NCCL_SOCKET_IFNAME 环境变量指定了特定的网卡名称，则选择该网卡。

否则网卡选择的具体顺序如下：

l  首先，选择名称以 "ib" 为起始字符串的网卡；如果没找到，

l  接着，选择与 NCCL_COMM_ID 中指定 IP 地址处于同一子网的网卡；否则，

l  其次，选择名称不以 "docker" 和 "lo" 为起始字符串的网卡；否则，

l  再次，选择名称以 "docker" 为起始字符串的网卡；否则，

l  {BANNED}最佳后，选择名称以 "lo" 为起始字符串的网卡。

当不指定环境变量时，就是上述红色字体部分选择出的网卡。这里就确定了bootstrap网络的接口名称bootstrapNetIfName和接口地址bootstrapNetIfAddr。

{BANNED}最佳后，通过 initNvtxRegisteredEnums 函数为基于 nvtx 的工具注册额外的信息。具体来说，注册 reduce 类通信的算子类型。

    回到ncclGetUniqueId函数，调用完ncclInit后，则调用bootstrapGetUniqueId。

l  bootstrapGetUniqueId

我们看bootstrapGetUniqueId的调用方式如下，其返回参数out即ncclUniqueId结构，同时可以转换为ncclBootstrapHandle结构，所以其实ncclUniqueId的本质就是ncclBootstrapHandle。

ncclResult_t res = bootstrapGetUniqueId((struct ncclBootstrapHandle*)out);

所以，我们先来看一下ncclBootstrapHandle这个结构体，如下面的代码片段所示。ncclBootstrapHandle 包含一个魔术值和记录 socket 地址的数据结构。也就是说，ncclUniqueId 是一个包含魔术值和root进程的 socket 地址的数据结构，包括 root 进程的监听地址和端口。

struct ncclBootstrapHandle {
  uint64_t magic;
  union ncclSocketAddress addr;
};
/* Common socket address storage structure for IPv4/IPv6 */
union ncclSocketAddress {
  struct sockaddr sa;
  struct sockaddr_in sin;
  struct sockaddr_in6 sin6;
};

    下面看bootstrapGetUniqueId的代码逻辑

ncclResult_t bootstrapGetUniqueId(struct ncclBootstrapHandle* handle) {
  memset(handle, 0, sizeof(ncclBootstrapHandle));

  const char* env = ncclGetEnv("NCCL_COMM_ID");
  if (env) {
    INFO(NCCL_ENV, "NCCL_COMM_ID set by environment to %s", env);
    /* 首先判断是不是设置了 「NCCL_COMM_ID」 环境变量。NCCL_COMM_ID 中包含 IP 地址和端口号，比如"ip:port"。*/
    if (ncclSocketGetAddrFromString(&handle->addr, env) != ncclSuccess) {
      WARN("Invalid NCCL_COMM_ID, please use format: : or []: or :");
      return ncclInvalidArgument;
    }
    handle->magic = NCCL_MAGIC;
  } else {
    NCCLCHECK(getRandomData(&handle->magic, sizeof(handle->magic)));
    memcpy(&handle->addr, &bootstrapNetIfAddr, sizeof(union ncclSocketAddress));
    NCCLCHECK(bootstrapCreateRoot(handle, false));
  }

  return ncclSuccess;
}

bootstrapGetUniqueId函数首先判断是不是设置了NCCL_COMM_ID环境变量。NCCL_COMM_ID 中包含 IP 地址和端口号，比如"ip:port"。上面我们提到，ncclUniqueId 是一个包含魔术值和 root 的 socket 地址的数据结构。那么，如果设置了 NCCL_COMM_ID 环境变量，那么我直接使用 NCCL_COMM_ID中的 IP 地址和端口初始化 ncclUniqueId 数据结构，并使用默认值 NCCL_MAGIC初始化魔术值。

反之，函数通过读取 /dev/urandom获取随机魔术值。接着，使用前面bootstrapNetInit

函数中获取的bootstrapNetIfAddr继续初始化 ncclUniqueId 结构。这里bootstrapNetIfAddr 是在 ncclInit ->bootstrapNetIfAddr函数中构建的用于 bootstrap 网络的 IP 地址。另外，需要注意的是，bootstrapCreateRoot 函数可能会进一步初始化 ncclUniqueId，以包含监听的端口信息。

接下来，我们来看一看 bootstrapCreateRoot 函数, 该函数在所有进程间建立环形可达的通信链路。为了便于理解后面的代码，我们先这里简单介绍下环形可达通信链路的含义。具体来说，每个进程都建立一个监听端，并且知道其后序进程的监听地址。这样，就构成了一个如下图所示的环形通信链。例如，rank=1 的进程知道 rank=2 的进程的监听地址，这样，rank=1 的进程可以和 rank=2 的后序进程通信。

ncclResult_t bootstrapCreateRoot(struct ncclBootstrapHandle* handle, bool idFromEnv) {
  struct ncclSocket* listenSock;
  struct bootstrapRootArgs* args;
  pthread_t thread;

  NCCLCHECK(ncclCalloc(&listenSock, 1));
  NCCLCHECK(ncclSocketInit(listenSock, &handle->addr, handle->magic, ncclSocketTypeBootstrap, NULL, 0));
  NCCLCHECK(ncclSocketListen(listenSock));
  NCCLCHECK(ncclSocketGetAddr(listenSock, &handle->addr));

  NCCLCHECK(ncclCalloc(&args, 1));
  args->listenSock = listenSock;
  args->magic = handle->magic;
  NEQCHECK(pthread_create(&thread, NULL, bootstrapRoot, (void*)args), 0);
  ncclSetThreadName(thread, "NCCL BootstrapR");
  NEQCHECK(pthread_detach(thread), 0); // will not be pthread_join()'d
  return ncclSuccess;
}

 bootstrapCreateRoot 函数首先调用 ncclSocketInit 函数初始化 listenSock。接着，调用 ncclSocketListen 函数绑定并在 listenSock 监听。这里使用的是handle->addr，即前文所获取的bootstrapNetIfAddr。此外，ncclSocketListen 还通过 bind 函数确定了监听端口。ncclSocketGetAddr 函数将监听的网络地址和端口拷贝到 ncclUniqueId 数据结构中。也就是说，此时 ncclUniqueId 包含了 root 进程监听的 IP 地址和端口。（这里其实主要是端口信息，地址在调用bootstrapCreateRoot之前已经对handle->addr进行了赋值）

在实际使用时，用户需要将 ncclUniqueId 广播到所有进程。因此，其它进程都可以通过 root 进程的监听地址连接 root 进程。这样，其它进程可以将它们自己的监听地址发送到 root 进程，从而构建上面提到的环形通信链。

{BANNED}最佳后，bootstrapCreateRoot 函数还启动 bootstrapRoot 线程，在所有进程间建立环形通信链。

l  bootstrapRoot 线程

bootstrapRoot 在所有进程间建立环形通信链，注意这里是bootstrap网络的环形通信。首先，bootstrapRoot 通过下面的代码片段接收其它所有进程连接，获取与 root 进程连接的 sockaddr 和与其它进程连接的 sockaddr，分别保存在 rankAddressesRoot和 rankAddresses中。

static void *bootstrapRoot(void* rargs) {
......
  /* Receive addresses from all ranks */
  do {
    struct ncclSocket sock;
    NCCLCHECKGOTO(ncclSocketInit(&sock), res, out);
    NCCLCHECKGOTO(ncclSocketAccept(&sock, listenSock), res, out);
    NCCLCHECKGOTO(bootstrapNetRecv(&sock, &info, sizeof(info)), res, out);
    NCCLCHECKGOTO(ncclSocketClose(&sock), res, out);

    if (c == 0) {
      nranks = info.nranks;
      NCCLCHECKGOTO(ncclCalloc(&rankAddresses, nranks), res, out);
      NCCLCHECKGOTO(ncclCalloc(&rankAddressesRoot, nranks), res, out);
    }

    if (nranks != info.nranks) {
      WARN("Bootstrap Root : mismatch in rank count from procs %d : %d", nranks, info.nranks);
      goto out;
    }

    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 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);
  } while (c < nranks);
  TRACE(NCCL_INIT, "COLLECTED ALL %d HANDLES", nranks);

上述过程如下图所示。

这里可能有人会疑惑，怎么上来就接收呢？发送方是谁呢？不要忘了bootstrapRoot 是一个独立的线程，主线程和其他rank还在并行执行，他们得到ncclUniqueId后就会调用MPI_Bcast广播了。

接着，root 进程分别连接每一个 rank=i 进程，并将其后序 rank=i+1 进程的监听地址发送给它。这样，每个进程就可以连接其后序进程，从而构成一个环形通信链。

// Send the connect handle for the next rank in the AllGather ring
  for (int r=0; r<nranks; ++r) {
    int next = (r+1) % nranks;
    struct ncclSocket sock;
    NCCLCHECKGOTO(ncclSocketInit(&sock,rankAddressesRoot+r,magic, ncclSocketTypeBootstrap), res, out);
    NCCLCHECKGOTO(ncclSocketConnect(&sock), res, out);
    NCCLCHECKGOTO(bootstrapNetSend(&sock,rankAddresses+next,sizeof(union ncclSocketAddress)), res, out);
    NCCLCHECKGOTO(ncclSocketClose(&sock), res, out);
  }
  TRACE(NCCL_INIT, "SENT OUT ALL %d HANDLES", nranks);

如下图所示，root 进程连接 rank=2 进程，并将 rank=3 进程的监听地址发送给它。

    {BANNED}最佳后回到nccl-test的run函数中，如上所述，在调用完ncclGetUniqueId后，调用MPI_Bcast将生成的ncclUniqueId广播给所有的rank。这里要注意我们是以root rank为视角讲解的，就如同上面所讲的bootstrapRoot 线程首先通过接收其它所有进程连接。所有一定有其他rank向root发送的过程，那么这一点我们将在后面部分讲解。

小结

在{BANNED}中国第一部分，我们主要分析了 ncclGetUniqueId 函数了实现。整体流程如所示：

该函数在 root 进程中执行，根据用户设置或者系统选择决定所使用的 bootstrap 网络的网卡，并初始化 bootstrap 网络。这里，获取的 ncclUniqueId 实际上是包含魔术值和使用上述网卡 IP 地址和绑定端口的信息。通过广播该 ncclUniqueId，所有进程都可以连接 root 进程，并进一步创建环形通信链路。

 

我们将在下一篇分析 ncclCommInitRank 函数，说明非 root 进程如何和 root 进程协作完成 bootstrap 网络的初始化、communicator 的创建和更为复杂的 NCCL 拓扑探测过程。