[WIP] feat: support SSH-based multi-node training launch

CMakeLists.txt

@@ -73,6 +73,11 @@ endif()
 file(GLOB_RECURSE SRC ${PROJECT_SOURCE_DIR}/infini_train/src/*.cc)
 list(FILTER SRC EXCLUDE REGEX ".*kernels/cpu/.*")
 
+set(RENDEZVOUS_SERVER_SOURCE
+  ${PROJECT_SOURCE_DIR}/infini_train/src/core/ccl/rendezvous_server.cc
+)
+list(REMOVE_ITEM SRC ${RENDEZVOUS_SERVER_SOURCE})
+
 # Exclude backend-specific runtime/ccl translation units when the corresponding
 # backend is disabled. This keeps each build self-contained and avoids pulling
 # in headers (e.g. <cuda_runtime.h> / <mcr/mc_runtime.h>) that are not on the
@@ -86,6 +91,18 @@ endif()
 # CPU kernels (*.cc)
 file(GLOB_RECURSE CPU_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/cpu/*.cc)
 
+# ------------------------------------------------------------------------------
+# Rendezvous library
+# ------------------------------------------------------------------------------
+
+add_library(infini_train_rendezvous STATIC ${RENDEZVOUS_SERVER_SOURCE})
+add_library(InfiniTrain::rendezvous ALIAS infini_train_rendezvous)
+target_compile_features(infini_train_rendezvous PUBLIC cxx_std_20)
+target_include_directories(infini_train_rendezvous PUBLIC
+  "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>"
+)
+target_link_libraries(infini_train_rendezvous PUBLIC glog)
+
 # ------------------------------------------------------------------------------
 # CPU kernels library
 # ------------------------------------------------------------------------------

infini_train/include/core/ccl/ccl_utils.h

@@ -5,11 +5,8 @@
 namespace infini_train::core {
 class CclUniqueId;
 
-void WriteUniqueIdFile(const CclUniqueId &unique_id,
-                       const std::string &pg_name);
+void WriteUniqueId(const CclUniqueId &unique_id, const std::string &pg_name);
 
-void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name);
-
-void CleanupUniqueIdFile(const std::string &pg_name);
+void ReadUniqueId(CclUniqueId *unique_id, const std::string &pg_name);
 
 }  // namespace infini_train::core

infini_train/include/core/ccl/rendezvous_protocol.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <stdint.h>
+
+namespace infini_train::core::rendezvous {
+inline constexpr uint32_t kMagic = 0x49545244;
+inline constexpr uint32_t kMaxKeySize = 4096;
+inline constexpr uint32_t kMaxValueSize = 1024 * 1024;
+
+enum class Operation : uint32_t {
+  kPut = 1,
+  kGet = 2,
+};
+
+struct Header {
+  uint32_t magic;
+  uint32_t operation;
+  uint32_t key_size;
+  uint32_t value_size;
+};
+
+struct Response {
+  uint32_t status;
+  uint32_t value_size;
+};
+}  // namespace infini_train::core::rendezvous

infini_train/include/core/ccl/rendezvous_server.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <memory>
+#include <string>
+
+namespace infini_train::core {
+class RendezvousServer {
+ public:
+  RendezvousServer(const std::string &host, int port, int timeout);
+  ~RendezvousServer();
+
+  RendezvousServer(const RendezvousServer &) = delete;
+  RendezvousServer &operator=(const RendezvousServer &) = delete;
+
+  int port() const;
+
+ private:
+  class Impl;
+  std::unique_ptr<Impl> impl_;
+};
+}  // namespace infini_train::core

infini_train/include/nn/parallel/process_group.h

@@ -139,7 +139,6 @@ class ProcessGroup {
 
   const std::string name_ = "";
 
-  bool is_main_process_ = false;
   Device::DeviceType backend_ = Device::DeviceType::kInvalid;
 
   // Save impl for convenience

infini_train/src/core/ccl/ccl_utils.cc

@@ -1,81 +1,135 @@
 #include "infini_train/include/core/ccl/ccl_utils.h"
 
+#include <arpa/inet.h>
+#include <stdint.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+#include <cerrno>
 #include <chrono>
-#include <cstdio>
+#include <compare>
 #include <cstdlib>
-#include <filesystem>
-#include <fstream>
-#include <iterator>
+#include <limits>
+#include <ostream>
+#include <string_view>
 #include <thread>
 
 #include "glog/logging.h"
 
 #include "infini_train/include/core/ccl/ccl_common.h"
+#include "infini_train/include/core/ccl/rendezvous_protocol.h"
+
+#include <netdb.h>
 
 namespace infini_train::core {
 namespace {
-std::string UniqueIdFileName(const std::string &pg_name) {
-  const char *run_id = std::getenv("INFINI_RUN_ID");
-  const std::string prefix = run_id == nullptr ? "" : std::string(run_id) + "_";
-  return "cclUniqueId_" + prefix + pg_name + ".bin";
+int GetEnvInt(const char *name, int fallback) {
+  const char *value = std::getenv(name);
+  return value == nullptr ? fallback : std::stoi(value);
 }
 
-std::string UniqueIdTmpFileName(const std::string &pg_name) {
-  const char *run_id = std::getenv("INFINI_RUN_ID");
-  const std::string prefix = run_id == nullptr ? "" : std::string(run_id) + "_";
-  return "cclUniqueId_" + prefix + pg_name + ".tmp";
+bool ReadAll(int fd, void *buffer, size_t size) {
+  auto *bytes = static_cast<char *>(buffer);
+  while (size > 0) {
+    const ssize_t count = read(fd, bytes, size);
+    if (count < 0 && errno == EINTR) continue;
+    if (count <= 0) return false;
+    bytes += count;
+    size -= static_cast<size_t>(count);
+  }
+  return true;
 }
-}  // namespace
 
-void WriteUniqueIdFile(const CclUniqueId &unique_id,
-                       const std::string &pg_name) {
-  const std::string tmp_path = UniqueIdTmpFileName(pg_name);
-
-  std::ofstream ofs(tmp_path, std::ios::binary);
-  CHECK(ofs.good()) << "Failed to open unique_id tmp file for write: "
-                    << tmp_path;
-  const size_t size = unique_id.Size();
-  ofs.write(reinterpret_cast<const char *>(unique_id.Data()),
-            static_cast<std::streamsize>(size));
-  ofs.close();
-
-  const std::string file_path = UniqueIdFileName(pg_name);
-  CHECK_EQ(std::rename(tmp_path.c_str(), file_path.c_str()), 0)
-      << "Failed to rename unique_id file from " << tmp_path << " to "
-      << file_path;
+bool WriteAll(int fd, const void *buffer, size_t size) {
+  const auto *bytes = static_cast<const char *>(buffer);
+  while (size > 0) {
+    const ssize_t count = write(fd, bytes, size);
+    if (count < 0 && errno == EINTR) continue;
+    if (count <= 0) return false;
+    bytes += count;
+    size -= static_cast<size_t>(count);
+  }
+  return true;
 }
 
-void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name) {
-  CHECK_NOTNULL(unique_id);
-  const std::string file_path = UniqueIdFileName(pg_name);
-
-  while (!std::filesystem::exists(file_path)) {
-    std::this_thread::sleep_for(std::chrono::microseconds(1000));
+int ConnectRendezvous() {
+  const char *host = std::getenv("MASTER_ADDR");
+  const char *port = std::getenv("MASTER_PORT");
+  CHECK(host != nullptr && port != nullptr)
+      << "MASTER_ADDR and MASTER_PORT are required for multi-node rendezvous";
+  addrinfo hints{};
+  hints.ai_family = AF_UNSPEC;
+  hints.ai_socktype = SOCK_STREAM;
+  addrinfo *addresses = nullptr;
+  CHECK_EQ(getaddrinfo(host, port, &hints, &addresses), 0)
+      << "Failed to resolve rendezvous endpoint " << host << ':' << port;
+  const int timeout = GetEnvInt("INFINI_RDZV_TIMEOUT_SEC", 120);
+  const auto deadline =
+      std::chrono::steady_clock::now() + std::chrono::seconds(timeout);
+  int fd = -1;
+  while (std::chrono::steady_clock::now() < deadline && fd < 0) {
+    for (addrinfo *address = addresses; address != nullptr;
+         address = address->ai_next) {
+      const int candidate = socket(address->ai_family, address->ai_socktype,
+                                   address->ai_protocol);
+      if (candidate < 0) continue;
+      if (connect(candidate, address->ai_addr, address->ai_addrlen) == 0) {
+        fd = candidate;
+        break;
+      }
+      close(candidate);
+    }
+    if (fd < 0) std::this_thread::sleep_for(std::chrono::milliseconds(100));
   }
+  freeaddrinfo(addresses);
+  CHECK_GE(fd, 0) << "Timed out connecting to rendezvous " << host << ':'
+                  << port;
+  return fd;
+}
 
-  std::ifstream ifs(file_path, std::ios::binary);
-  CHECK(ifs.good()) << "Failed to open unique_id file for read: " << file_path;
-
-  std::string bytes((std::istreambuf_iterator<char>(ifs)),
-                    std::istreambuf_iterator<char>());
-  ifs.close();
+std::string RendezvousRequest(rendezvous::Operation operation,
+                              const std::string &key, std::string_view value) {
+  CHECK_LE(key.size(), std::numeric_limits<uint32_t>::max());
+  CHECK_LE(value.size(), std::numeric_limits<uint32_t>::max());
+  const int fd = ConnectRendezvous();
+  const rendezvous::Header header{htonl(rendezvous::kMagic),
+                                  htonl(static_cast<uint32_t>(operation)),
+                                  htonl(static_cast<uint32_t>(key.size())),
+                                  htonl(static_cast<uint32_t>(value.size()))};
+  CHECK(WriteAll(fd, &header, sizeof(header)) &&
+        WriteAll(fd, key.data(), key.size()) &&
+        (value.empty() || WriteAll(fd, value.data(), value.size())))
+      << "Failed to send rendezvous request for " << key;
+  rendezvous::Response response{};
+  CHECK(ReadAll(fd, &response, sizeof(response)));
+  CHECK_EQ(ntohl(response.status), 0)
+      << "Rendezvous request failed for " << key;
+  std::string result(ntohl(response.value_size), '\0');
+  CHECK(result.empty() || ReadAll(fd, result.data(), result.size()));
+  close(fd);
+  return result;
+}
 
-  CHECK_EQ(bytes.size(), unique_id->Size())
-      << "Mismatched unique_id size in file. expected=" << unique_id->Size()
-      << ", got=" << bytes.size();
-  unique_id->Load(bytes.data(), bytes.size());
+std::string RendezvousKey(const std::string &pg_name) {
+  const char *run_id = std::getenv("INFINI_RUN_ID");
+  CHECK(run_id != nullptr) << "INFINI_RUN_ID is required";
+  return std::string(run_id) + "/ccl/" + pg_name;
 }
 
-void CleanupUniqueIdFile(const std::string &pg_name) {
-  const std::string file_path = UniqueIdFileName(pg_name);
-  if (std::filesystem::exists(file_path)) {
-    std::filesystem::remove(file_path);
-  }
+}  // namespace
 
-  const std::string tmp_path = UniqueIdTmpFileName(pg_name);
-  if (std::filesystem::exists(tmp_path)) {
-    std::filesystem::remove(tmp_path);
-  }
+void WriteUniqueId(const CclUniqueId &unique_id, const std::string &pg_name) {
+  RendezvousRequest(
+      rendezvous::Operation::kPut, RendezvousKey(pg_name),
+      std::string_view(static_cast<const char *>(unique_id.Data()),
+                       unique_id.Size()));
 }
 
+void ReadUniqueId(CclUniqueId *unique_id, const std::string &pg_name) {
+  CHECK_NOTNULL(unique_id);
+  const std::string bytes = RendezvousRequest(rendezvous::Operation::kGet,
+                                              RendezvousKey(pg_name), {});
+  CHECK_EQ(bytes.size(), unique_id->Size());
+  unique_id->Load(bytes.data(), bytes.size());
+}
 }  // namespace infini_train::core