refactor: refactor the naming of ncclUniqueId files and add a unified cleanup logic

Author: kilinchange

infini_train/src/nn/parallel/process_group.cc

@@ -3,6 +3,7 @@
 #include <algorithm>
 #include <chrono>
 #include <filesystem>
+#include <format>
 #include <fstream>
 #include <iterator>
 #include <memory>
@@ -42,18 +43,28 @@ const std::unordered_map<ReduceOpType, ncclRedOp_t> kNcclReduceOpMap = {
     {ReduceOpType::kAvg, ncclAvg},
 };
 
-void WriteNcclUniqueId(const ncclUniqueId &nccl_id, const std::string &filename) {
-    std::string tmp_path = filename + ".tmp";
+inline std::string NcclFileName(const std::string &name, bool tmp = false) {
+    return std::format("ncclUniqueId_{}.{}", name, tmp ? "tmp" : "bin");
+}
+
+void WriteNcclUniqueId(const ncclUniqueId &nccl_id, const std::string &pg_name) {
+    std::string tmp_path = NcclFileName(pg_name, true);
 
     std::ofstream ofs(tmp_path, std::ios::binary);
     ofs.write(reinterpret_cast<const char *>(&nccl_id), sizeof(nccl_id));
     ofs.close();
 
-    std::rename(tmp_path.c_str(), filename.c_str());
+    std::rename(tmp_path.c_str(), NcclFileName(pg_name).c_str());
 }
 
-void ReadNcclUniqueId(ncclUniqueId &nccl_id, const std::string &filename) {
-    std::ifstream ifs(filename, std::ios::binary);
+void ReadNcclUniqueId(ncclUniqueId &nccl_id, const std::string &pg_name) {
+    std::string file_path = NcclFileName(pg_name);
+
+    while (std::filesystem::exists(file_path) == false) {
+        std::this_thread::sleep_for(std::chrono::microseconds(1000));
+    }
+
+    std::ifstream ifs(file_path, std::ios::binary);
     ifs.read(reinterpret_cast<char *>(&nccl_id), sizeof(nccl_id));
     ifs.close();
 }
@@ -78,9 +89,6 @@ ProcessGroup::ProcessGroup(const std::string &process_group_name, const std::vec
 
         WriteNcclUniqueId(nccl_id, name_);
     } else {
-        while (std::filesystem::exists(name_) == false) {
-            std::this_thread::sleep_for(std::chrono::microseconds(1000));
-        }
         ReadNcclUniqueId(nccl_id, name_);
     }
 

tools/infini_run/infini_run.cc

@@ -1,13 +1,11 @@
+#include <cstdio>
 #include <cstdlib>
-#include <fstream>
-#include <ios>
-#include <sstream>
+#include <filesystem>
+#include <regex>
 #include <string>
 #include <sys/wait.h>
-#include <sys/stat.h>
 #include <unistd.h>
 #include <vector>
-#include <cstdio>
 
 #ifdef USE_NCCL
 #include <nccl.h>
@@ -21,15 +19,19 @@ DEFINE_int32(nproc_per_node, 1, "Number of processes per node");
 DEFINE_int32(node_rank, 0, "Rank of this node");
 DEFINE_string(rdzv_endpoint, "127.0.0.1:29500", "Rendezvous endpoint (host:port)");
 
-#ifdef USE_NCCL
-std::string NcclIdToString(const ncclUniqueId& id) {
-    std::ostringstream oss;
-    for (int i = 0; i < NCCL_UNIQUE_ID_BYTES; ++i) {
-        oss << std::hex << std::uppercase << std::setw(2) << std::setfill('0') << (int)(unsigned char)id.internal[i];
+void CleanupNcclIdFiles() {
+    const std::filesystem::path cwd = std::filesystem::current_path();
+    std::regex pattern(R"(ncclUniqueId_.*\.bin)");
+
+    for (const auto &entry : std::filesystem::directory_iterator(cwd)) {
+        if (entry.is_regular_file()) {
+            const std::string filename = entry.path().filename().string();
+            if (std::regex_match(filename, pattern)) {
+                std::filesystem::remove(entry.path());
+            }
+        }
     }
-    return oss.str();
 }
-#endif
 
 int main(int argc, char **argv) {
     gflags::ParseCommandLineFlags(&argc, &argv, true);
@@ -46,23 +48,6 @@ int main(int argc, char **argv) {
     std::string master_addr = FLAGS_rdzv_endpoint.substr(0, FLAGS_rdzv_endpoint.find(':'));
     std::string master_port = FLAGS_rdzv_endpoint.substr(FLAGS_rdzv_endpoint.find(':') + 1);
 
-    const char* nccl_id_path = "/data/shared/InfiniTrain-dev/data/nccl_id.bin";
-    const char* nccl_id_tmp_path = "/data/shared/InfiniTrain-dev/data/nccl_id.tmp";
-
-#ifdef USE_NCCL
-    if (FLAGS_node_rank == 0) {
-        ncclUniqueId id;
-        ncclGetUniqueId(&id);
-
-        std::ofstream ofs(nccl_id_tmp_path, std::ios::binary);
-        ofs.write((char *)&id, sizeof(id));
-        ofs.close();
-
-        // atomic operation
-        rename(nccl_id_tmp_path, nccl_id_path);
-    }
-#endif
-
     for (int local_proc_rank = 0; local_proc_rank < FLAGS_nproc_per_node; ++local_proc_rank) {
         pid_t pid = fork();
         if (pid == 0) {
@@ -74,20 +59,6 @@ int main(int argc, char **argv) {
             setenv("MASTER_ADDR", master_addr.c_str(), 1);
             setenv("MASTER_PORT", master_port.c_str(), 1);
 
-#ifdef USE_NCCL
-            struct stat st;
-            while (stat(nccl_id_path, &st) != 0) {
-                usleep(1000);
-            }
-
-            ncclUniqueId id;
-            std::ifstream ifs(nccl_id_path, std::ios::binary);
-            ifs.read((char*)&id, sizeof(id));
-
-            std::string id_str = NcclIdToString(id);
-            setenv("NCCL_UNIQUE_ID", id_str.c_str(), 1);
-#endif
-
             execvp(train_program.c_str(), train_argv.data());
             perror("exec failed");
             exit(1);
@@ -101,7 +72,7 @@ int main(int argc, char **argv) {
 
 #ifdef USE_NCCL
     if (FLAGS_node_rank == 0) {
-        std::remove(nccl_id_path);
+        CleanupNcclIdFiles();
     }
 #endif