feat: Move the logic for cleaning up the ncclUniqueId file into the ProcessGroup destructor.

Author: kilinchange

infini_train/include/nn/parallel/process_group.h

@@ -29,6 +29,7 @@ namespace infini_train::nn::parallel {
 class ProcessGroup {
 public:
     explicit ProcessGroup(const std::string &process_group_name, const std::vector<int> &device_indices);
+    ~ProcessGroup();
 
     int GetGroupRank(int thread_rank) const;
 
@@ -69,6 +70,8 @@ class ProcessGroup {
     int world_size_ = 0;
 
     const std::string name_ = "";
+
+    bool is_main_process_ = false;
 };
 #endif
 

infini_train/src/nn/parallel/process_group.cc

@@ -68,6 +68,15 @@ void ReadNcclUniqueId(ncclUniqueId &nccl_id, const std::string &pg_name) {
     ifs.read(reinterpret_cast<char *>(&nccl_id), sizeof(nccl_id));
     ifs.close();
 }
+
+void CleanupNcclIdFile(const std::string &pg_name) {
+    const std::filesystem::path cwd = std::filesystem::current_path();
+    std::string file_path = NcclFileName(pg_name);
+
+    if (std::filesystem::exists(file_path)) {
+        std::filesystem::remove(file_path);
+    }
+}
 #endif
 
 } // namespace
@@ -86,6 +95,12 @@ ProcessGroup::ProcessGroup(const std::string &process_group_name, const std::vec
     }
 }
 
+ProcessGroup::~ProcessGroup() {
+    if (is_main_process_) {
+        CleanupNcclIdFile(name_);
+    }
+}
+
 void ProcessGroup::InitSingleProcess(const std::vector<int> &ranks) {
     comms_.resize(world_size_);
     NCCL_CHECK(ncclCommInitAll(comms_.data(), world_size_, ranks.data()));
@@ -105,8 +120,9 @@ void ProcessGroup::InitMultiProcess(const std::vector<int> &ranks) {
 
     if (std::ranges::min(ranks) < (global::GetGlobalProcRank() + 1) * global::GetNthreadPerProc()
         && std::ranges::min(ranks) >= global::GetGlobalProcRank() * global::GetNthreadPerProc()) {
-        ncclGetUniqueId(&nccl_id);
+        is_main_process_ = true;
 
+        ncclGetUniqueId(&nccl_id);
         WriteNcclUniqueId(nccl_id, name_);
     } else {
         ReadNcclUniqueId(nccl_id, name_);

tools/infini_run/infini_run.cc

@@ -15,20 +15,6 @@ 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)");
 
-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());
-            }
-        }
-    }
-}
-
 int main(int argc, char **argv) {
     gflags::ParseCommandLineFlags(&argc, &argv, true);
     google::InitGoogleLogging(argv[0]);
@@ -70,9 +56,5 @@ int main(int argc, char **argv) {
         wait(&status);
     }
 
-    if (FLAGS_node_rank == 0) {
-        CleanupNcclIdFiles();
-    }
-
     return 0;
 }