Add IPC domain rank detection

CMakeLists.txt

@@ -226,6 +226,10 @@ if(MSCCLPP_USE_CUDA)
     else()
         set(GPU_LIBRARIES CUDA::cudart CUDA::cuda_driver)
     endif()
+    if(NOT TARGET CUDA::nvml)
+        message(FATAL_ERROR "CUDA NVML target CUDA::nvml is required for MSCCLPP CUDA builds.")
+    endif()
+    list(APPEND GPU_LIBRARIES CUDA::nvml)
 else()
     set(CMAKE_HIP_STANDARD 20)
     set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -Wall -Wextra")

include/mscclpp/core.hpp

@@ -46,6 +46,10 @@ class Bootstrap {
   /// @return The total number of ranks per node.
   virtual int getNranksPerNode() const = 0;
 
+  /// Return the number of ranks in this rank's GPU IPC domain.
+  /// @return The number of ranks in the GPU IPC domain.
+  virtual int getNranksPerIpcDomain() const;
+
   /// Send arbitrary data to another process.
   ///
   /// Data sent via `send(senderBuff, size, receiverRank, tag)` can be received via `recv(receiverBuff, size,
@@ -144,6 +148,9 @@ class TcpBootstrap : public Bootstrap {
   /// Return the total number of ranks per node.
   int getNranksPerNode() const override;
 
+  /// Return the number of ranks in this rank's GPU IPC domain.
+  int getNranksPerIpcDomain() const override;
+
   /// Send arbitrary data to another process.
   ///
   /// Data sent via `send(senderBuff, size, receiverRank, tag)` can be received via `recv(receiverBuff, size,

python/csrc/core_py.cpp

@@ -56,6 +56,7 @@ void register_core(nb::module_& m) {
       .def("get_rank", &Bootstrap::getRank)
       .def("get_n_ranks", &Bootstrap::getNranks)
       .def("get_n_ranks_per_node", &Bootstrap::getNranksPerNode)
+      .def("get_n_ranks_per_ipc_domain", &Bootstrap::getNranksPerIpcDomain)
       .def(
           "send",
           [](Bootstrap* self, uintptr_t ptr, size_t size, int peer, int tag) {

python/mscclpp/_core/comm.py

@@ -73,6 +73,7 @@ def __init__(
         self.my_rank = self.bootstrap.get_rank()
         self.nranks = self.bootstrap.get_n_ranks()
         self.nranks_per_node = self.bootstrap.get_n_ranks_per_node()
+        self.nranks_per_ipc_domain = self.bootstrap.get_n_ranks_per_ipc_domain()
 
     def barrier(self):
         self.bootstrap.barrier()

src/core/bootstrap/bootstrap.cc

@@ -50,6 +50,8 @@ MSCCLPP_API_CPP void Bootstrap::groupBarrier(const std::vector<int>& ranks) {
   }
 }
 
+MSCCLPP_API_CPP int Bootstrap::getNranksPerIpcDomain() const { return getNranksPerNode(); }
+
 MSCCLPP_API_CPP void Bootstrap::send(const std::vector<char>& data, int peer, int tag) {
   size_t size = data.size();
   send((void*)&size, sizeof(size_t), peer, tag);
@@ -83,6 +85,7 @@ class TcpBootstrap::Impl {
   int getRank();
   int getNranks();
   int getNranksPerNode();
+  int getNranksPerIpcDomain();
   void allGather(void* allData, int size);
   void broadcast(void* data, int size, int root);
   void send(void* data, int size, int peer, int tag);
@@ -95,6 +98,7 @@ class TcpBootstrap::Impl {
   int rank_;
   int nRanks_;
   int nRanksPerNode_;
+  int nRanksPerIpcDomain_;
   bool netInitialized;
   std::unique_ptr<Socket> listenSockRoot_;
   std::unique_ptr<Socket> listenSock_;
@@ -148,6 +152,7 @@ TcpBootstrap::Impl::Impl(int rank, int nRanks)
     : rank_(rank),
       nRanks_(nRanks),
       nRanksPerNode_(0),
+      nRanksPerIpcDomain_(0),
       netInitialized(false),
       peerCommAddresses_(nRanks, SocketAddress()),
       barrierArr_(nRanks, 0),
@@ -451,6 +456,42 @@ int TcpBootstrap::Impl::getNranksPerNode() {
   return nRanksPerNode_;
 }
 
+int TcpBootstrap::Impl::getNranksPerIpcDomain() {
+  if (nRanksPerIpcDomain_ > 0) return nRanksPerIpcDomain_;
+  std::vector<uint64_t> ipcDomainHashes(nRanks_);
+  ipcDomainHashes[rank_] = getIpcDomainHash();
+  allGather(ipcDomainHashes.data(), sizeof(uint64_t));
+
+  std::unordered_map<uint64_t, int> ipcDomainCounts;
+  for (uint64_t ipcDomainHash : ipcDomainHashes) {
+    ipcDomainCounts[ipcDomainHash]++;
+  }
+
+  const int nRanksPerIpcDomain = ipcDomainCounts[ipcDomainHashes[rank_]];
+  const std::string invalidIpcDomainLayout =
+      "IPC domain ranks must have the same size and be arranged in consecutive rank order";
+  if (nRanks_ % nRanksPerIpcDomain != 0) {
+    throw Error(invalidIpcDomainLayout + ": rank count is not divisible by IPC domain size", ErrorCode::InvalidUsage);
+  }
+
+  for (const auto& entry : ipcDomainCounts) {
+    if (entry.second != nRanksPerIpcDomain) {
+      throw Error(invalidIpcDomainLayout + ": IPC domains have different sizes", ErrorCode::InvalidUsage);
+    }
+  }
+  for (int i = 0; i < nRanks_; ++i) {
+    const int ipcDomainFirstRank = i - i % nRanksPerIpcDomain;
+    if (ipcDomainHashes[i] != ipcDomainHashes[ipcDomainFirstRank]) {
+      throw Error(invalidIpcDomainLayout + ": ranks are not grouped by IPC domain", ErrorCode::InvalidUsage);
+    }
+  }
+
+  INFO(MSCCLPP_INIT, "rank %d IPC domain fabric hash 0x%016llx nRanksPerIpcDomain %d", rank_,
+       static_cast<unsigned long long>(ipcDomainHashes[rank_]), nRanksPerIpcDomain);
+  nRanksPerIpcDomain_ = nRanksPerIpcDomain;
+  return nRanksPerIpcDomain_;
+}
+
 void TcpBootstrap::Impl::allGather(void* allData, int size) {
   char* data = static_cast<char*>(allData);
   int rank = rank_;
@@ -592,6 +633,8 @@ MSCCLPP_API_CPP int TcpBootstrap::getNranks() const { return pimpl_->getNranks()
 
 MSCCLPP_API_CPP int TcpBootstrap::getNranksPerNode() const { return pimpl_->getNranksPerNode(); }
 
+MSCCLPP_API_CPP int TcpBootstrap::getNranksPerIpcDomain() const { return pimpl_->getNranksPerIpcDomain(); }
+
 MSCCLPP_API_CPP void TcpBootstrap::send(void* data, int size, int peer, int tag) {
   pimpl_->send(data, size, peer, tag);
 }

src/core/include/utils_internal.hpp

@@ -37,6 +37,7 @@ int64_t busIdToInt64(const std::string busId);
 uint64_t getHash(const char* string, int n);
 uint64_t getHostHash();
 uint64_t getPidHash();
+uint64_t getIpcDomainHash();
 void getRandomData(void* buffer, size_t bytes);
 
 struct netIf {

src/core/utils_internal.cc

@@ -6,6 +6,10 @@
 #include <signal.h>
 #include <unistd.h>
 
+#if defined(MSCCLPP_USE_CUDA)
+#include <nvml.h>
+#endif
+
 #include <cstring>
 #include <fstream>
 #include <iostream>
@@ -175,6 +179,83 @@ uint64_t getPidHash(void) {
   return *pidHash;
 }
 
+#if defined(MSCCLPP_USE_CUDA) && defined(NVML_GPU_FABRIC_UUID_LEN)
+namespace {
+
+class NvmlState {
+ public:
+  NvmlState() : initialized_(nvmlInit_v2() == NVML_SUCCESS) {}
+
+  NvmlState(const NvmlState&) = delete;
+  NvmlState& operator=(const NvmlState&) = delete;
+  NvmlState(NvmlState&&) = delete;
+  NvmlState& operator=(NvmlState&&) = delete;
+
+  ~NvmlState() {
+    if (initialized_) {
+      (void)nvmlShutdown();
+    }
+  }
+
+  bool isInitialized() const { return initialized_; }
+
+ private:
+  bool initialized_ = false;
+};
+
+template <typename FabricInfo>
+uint64_t getFabricHash(const FabricInfo& fabricInfo) {
+  char hashData[NVML_GPU_FABRIC_UUID_LEN + sizeof(fabricInfo.cliqueId)];
+  std::memcpy(hashData, fabricInfo.clusterUuid, NVML_GPU_FABRIC_UUID_LEN);
+  std::memcpy(hashData + NVML_GPU_FABRIC_UUID_LEN, &fabricInfo.cliqueId, sizeof(fabricInfo.cliqueId));
+  return getHash(hashData, sizeof(hashData));
+}
+
+bool tryGetNvmlIpcDomainHash(uint64_t& ipcDomainHash) {
+  // Use the current CUDA device; callers must set the rank's device before querying.
+  int deviceId;
+  char pciBusId[] = "00000000:00:00.0";
+  if (cudaGetDevice(&deviceId) != cudaSuccess ||
+      cudaDeviceGetPCIBusId(pciBusId, sizeof(pciBusId), deviceId) != cudaSuccess) {
+    return false;
+  }
+
+  static NvmlState nvml;
+  nvmlDevice_t nvmlDevice;
+  if (!nvml.isInitialized() || nvmlDeviceGetHandleByPciBusId_v2(pciBusId, &nvmlDevice) != NVML_SUCCESS) {
+    return false;
+  }
+
+#if defined(nvmlGpuFabricInfo_v2)
+  nvmlGpuFabricInfoV_t fabricInfo = {};
+  fabricInfo.version = nvmlGpuFabricInfo_v2;
+  nvmlReturn_t result = nvmlDeviceGetGpuFabricInfoV(nvmlDevice, &fabricInfo);
+#else
+  nvmlGpuFabricInfo_t fabricInfo = {};
+  nvmlReturn_t result = nvmlDeviceGetGpuFabricInfo(nvmlDevice, &fabricInfo);
+#endif
+  if (result != NVML_SUCCESS || fabricInfo.state != NVML_GPU_FABRIC_STATE_COMPLETED ||
+      fabricInfo.status != NVML_SUCCESS) {
+    return false;
+  }
+
+  ipcDomainHash = getFabricHash(fabricInfo);
+  return true;
+}
+
+}  // namespace
+#endif
+
+uint64_t getIpcDomainHash(void) {
+#if defined(MSCCLPP_USE_CUDA) && defined(NVML_GPU_FABRIC_UUID_LEN)
+  uint64_t ipcDomainHash;
+  if (tryGetNvmlIpcDomainHash(ipcDomainHash)) {
+    return ipcDomainHash;
+  }
+#endif
+  return getHostHash();
+}
+
 int parseStringList(const char* string, netIf* ifList, int maxList) {
   if (!string) return 0;
 

test/mp_unit/bootstrap_tests.cc

@@ -42,10 +42,17 @@ void BootstrapTest::bootstrapTestSendRecv(std::shared_ptr<mscclpp::Bootstrap> bo
   }
 }
 
+void BootstrapTest::bootstrapTestIpcDomain(std::shared_ptr<mscclpp::Bootstrap> bootstrap) {
+  int nRanksPerIpcDomain = bootstrap->getNranksPerIpcDomain();
+  EXPECT_GT(nRanksPerIpcDomain, 0);
+  EXPECT_LE(nRanksPerIpcDomain, bootstrap->getNranks());
+}
+
 void BootstrapTest::bootstrapTestAll(std::shared_ptr<mscclpp::Bootstrap> bootstrap) {
   bootstrapTestAllGather(bootstrap);
   bootstrapTestBarrier(bootstrap);
   bootstrapTestSendRecv(bootstrap);
+  bootstrapTestIpcDomain(bootstrap);
 }
 
 TEST(BootstrapTest, WithId) {
@@ -127,6 +134,7 @@ class MPIBootstrap : public mscclpp::Bootstrap {
     MPI_Comm_size(shmcomm, &shmrank);
     return shmrank;
   }
+  int getNranksPerIpcDomain() const override { return getNranksPerNode(); }
   void allGather(void* sendbuf, int size) override {
     MPI_Allgather(MPI_IN_PLACE, 0, MPI_BYTE, sendbuf, size, MPI_BYTE, MPI_COMM_WORLD);
   }

test/mp_unit/mp_unit_tests.hpp

@@ -56,6 +56,8 @@ class BootstrapTest : public MultiProcessTest {
 
   void bootstrapTestSendRecv(std::shared_ptr<mscclpp::Bootstrap> bootstrap);
 
+  void bootstrapTestIpcDomain(std::shared_ptr<mscclpp::Bootstrap> bootstrap);
+
   void bootstrapTestAll(std::shared_ptr<mscclpp::Bootstrap> bootstrap);
 
   // Each test case should finish within 30 seconds.