ensure that all mpi processes quit safely

MaterSim · Sep 4, 2024 · a3cae4e · a3cae4e
1 parent f03d218
commit a3cae4e
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 17 deletions.
diff --git a/pyxtal/optimize/DFS.py b/pyxtal/optimize/DFS.py
@@ -241,16 +241,24 @@ def _run(self, pool=None):
             if self.ff_opt:
                 N_added = self.update_ff_paramters(cur_xtals, engs, N_added)
             else:
+                quit = False
                 if self.rank == 0:
                     if self.ref_pmg is not None:
-                        success_rate = self.success_count(cur_xtals,
-                                                          matches)
+                        success_rate = self.success_count(cur_xtals, matches)
                         if self.early_termination(success_rate):
-                            return success_rate
+                            quit = True
 
                     elif self.ref_pxrd is not None:
                         self.count_pxrd_match(cur_xtals, matches)
 
+                # quit the loop
+                if self.use_mpi:
+                    quit = self.comm.bcast(quit, root=0)
+
+                # Ensure that all ranks exit
+                if quit:
+                    return success_rate
+
         return success_rate
 
 

diff --git a/pyxtal/optimize/QRS.py b/pyxtal/optimize/QRS.py
@@ -270,6 +270,7 @@ def _run(self, pool=None):
             gen_results = self.local_optimization(cur_xtals, qrs=True, pool=pool)
 
             # Summary and Ranking
+            quit = False
             if self.rank == 0:
                 cur_xtals, matches, engs = self.gen_summary(t0, 
                                         gen_results, cur_xtals)
@@ -295,15 +296,19 @@ def _run(self, pool=None):
                     self.sampler = qmc.Sobol(d=len_reps, scramble=False)
 
                 if self.ref_pmg is not None:
-                    success_rate = self.success_count(cur_xtals,
-                                                      matches)
+                    success_rate = self.success_count(cur_xtals, matches)
 
                     if self.early_termination(success_rate):
-                        return success_rate
+                        quit = True
 
                 elif ref_pxrd is not None:
-                    self.count_pxrd_match(cur_xtals,
-                                          matches)
+                    self.count_pxrd_match(cur_xtals, matches)
+
+            if self.use_mpi:
+                quit = self.comm.bcast(quit, root=0)
+
+            if quit:
+                return success_rate
 
         return success_rate
 

diff --git a/pyxtal/optimize/WFS.py b/pyxtal/optimize/WFS.py
@@ -216,16 +216,22 @@ def _run(self, pool=None):
             if self.ff_opt:
                 N_added = self.update_ff_paramters(cur_xtals, engs, N_added)
             else:
+                quit = False
                 if self.rank == 0:
                     if self.ref_pmg is not None:
-                        success_rate = self.success_count(cur_xtals,
-                                                          matches)
+                        success_rate = self.success_count(cur_xtals, matches)
                         if self.early_termination(success_rate):
-                            return success_rate
+                            quit = True
 
                     elif self.ref_pxrd is not None:
-                        self.count_pxrd_match(cur_xtals,
-                                              matches)
+                        self.count_pxrd_match(cur_xtals, matches)
+                # quit the loop
+                if self.use_mpi:
+                    quit = self.comm.bcast(quit, root=0)
+
+                # Ensure that all ranks exit
+                if quit:
+                    return success_rate
 
         return success_rate
 

diff --git a/pyxtal/optimize/base.py b/pyxtal/optimize/base.py
@@ -356,8 +356,9 @@ def run(self, ref_pmg=None, ref_pxrd=None):
 
         if self.rank == 0:
             t = (time() - t0)/60
-            print(f"{self.name:s} COMPLETED in {t:.1f} mins {self.N_struc:d} strucs.")
-
+            strs = f"{self.name:s} {self.workdir} COMPLETED "
+            strs += f"in {t:.1f} mins {self.N_struc:d} strucs."
+            print(strs)
         return results
 
     def select_xtals(self, ref_xtals, ids, N_max):
@@ -383,18 +384,19 @@ def count_pxrd_match(self, xtals, matches):
             matches (list): list of XRD matches
 
         """
+        gen = self.generation
         for i, match in enumerate(matches):
             if match > 0.85:
                 (xtal, tag) = xtals[i]
                 with open(self.matched_cif, "a+") as f:
                     e = xtal.energy / sum(xtal.numMols)
                     try:
-                        label = self.tag + "-g" + str(self.generation) + "-p" + str(i)
+                        label = self.tag + "-g" + str(gen) + "-p" + str(i)
                         label += f"-e{e:.3f}-{tag:s}-{match:4.2f}"
                     except:
                         print("Error in e, tag, match", e, tag, match)
                     f.writelines(xtal.to_file(header=label))
-                    self.matches.append((self.generation, i, xtal, e, match, tag))
+                    self.matches.append((gen, i, xtal, e, match, tag))
 
     def success_count(self, xtals, matches):
         """