From ab69c53b6204183ee85f99bf23f82663d31dc1da Mon Sep 17 00:00:00 2001
From: qzhu2017 <alecfans@gmail.com>
Date: Sun, 8 Sep 2024 16:18:02 -0400
Subject: [PATCH] add the print out of memory

---
 pyxtal/lego/builder.py  | 19 +++++++++++++++----
 pyxtal/optimize/WFS.py  |  2 +-
 pyxtal/optimize/base.py | 17 ++++++++++++-----
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/pyxtal/lego/builder.py b/pyxtal/lego/builder.py
index fb2c4899..f8f10dba 100644
--- a/pyxtal/lego/builder.py
+++ b/pyxtal/lego/builder.py
@@ -335,8 +335,11 @@ def print_fun(x, f, accepted):
 
     # Extract the optimized xtal
     xtal = pyxtal()
-    xtal.from_1d_rep(res.x, sites, dim=dim)
-    return xtal, (x0, res.x)
+    try:
+        xtal.from_1d_rep(res.x, sites, dim=dim)
+        return xtal, (x0, res.x)
+    except:
+        return None
 
 
 def calculate_dSdx(x, xtal, des_ref, f, eps=1e-4, symmetry=True, verbose=False):
@@ -592,6 +595,12 @@ def __str__(self):
     def __repr__(self):
         return str(self)
 
+    def print_memory_usage(self):
+        import psutil
+        process = psutil.Process(os.getpid())
+        mem = process.memory_info().rss / 1024 ** 2
+        self.logging.info(f"Rank {self.rank} memory: {mem:.1f} MB")
+
     def set_descriptor_calculator(self, dtype='SO3', mykwargs={}):
         """
         Set up the calculator for descriptor computation.
@@ -808,6 +817,7 @@ def optimize_xtals_mproc(self, xtals, ncpu, args):
             start, end = i, min([i+N_batches, len(xtals)])
             ids = list(range(start, end))
             print(f"Rank {self.rank} minibatch {start} {end}")
+            self.print_memory_usage()
 
             def generate_args():
                 """
@@ -825,7 +835,8 @@ def generate_args():
                            self.ref_environments, opt_type, T, niter,
                            early_quit, minimizers)
             # Use the generator to pass args to reduce memory usage
-            for result in pool.imap_unordered(minimize_from_x_par, args_list):
+            _xtal, _xs = None, None
+            for result in pool.imap_unordered(minimize_from_x_par, generate_args()):
                 if result is not None:
                     (_xtals, _xs) = result
                     valid_xtals = self.process_xtals(
@@ -837,7 +848,7 @@ def generate_args():
             self.db.clean_structures_spg_topology(dim=self.dim)
 
             # After each minibatch, delete the local variables and run garbage collection
-            del ids, wp_libs, _xtals, _xs
+            del ids, _xtals, _xs
             gc.collect()  # Explicitly call garbage collector to free memory
 
         xtals_opt = list(xtals_opt)
diff --git a/pyxtal/optimize/WFS.py b/pyxtal/optimize/WFS.py
index 283b372c..db8c28f9 100644
--- a/pyxtal/optimize/WFS.py
+++ b/pyxtal/optimize/WFS.py
@@ -199,7 +199,7 @@ def _run(self, pool=None):
 
             # Local optimization
             gen_results = self.local_optimization(cur_xtals, pool=pool)
-            self.logging.info(f"Rank {self.rank} finishes local_opt")
+            self.logging.info(f"Rank {self.rank} finishes local_opt {len(gen_results)}")
 
             prev_xtals = None
             if self.rank == 0:
diff --git a/pyxtal/optimize/base.py b/pyxtal/optimize/base.py
index 5a3388c8..faf46a1b 100644
--- a/pyxtal/optimize/base.py
+++ b/pyxtal/optimize/base.py
@@ -52,7 +52,6 @@ def handler(signum, frame):
         return None  # or some other placeholder for timeout results
 
 def process_task(args):
-    #logger = args[-1]
     #logger.info(f"Rank start processing task")
     result = run_optimizer_with_timeout(args)#[:-1])
     #logger.info(f"Rank finished processing task")
@@ -330,6 +329,13 @@ def __str__(self):
     def __repr__(self):
         return str(self)
 
+    def print_memory_usage(self):
+        import psutil
+        process = psutil.Process(os.getpid())
+        mem = process.memory_info().rss / 1024 ** 2
+        gen = self.generation
+        self.logging.info(f"Rank {self.rank} memory: {mem:.1f} MB in gen {gen}")
+
     def new_struc(self, xtal, xtals):
         return new_struc(xtal, xtals)
 
@@ -967,7 +973,7 @@ def local_optimization(self, xtals, qrs=False, pool=None):
         elif self.ncpu == 1:
             return self.local_optimization_serial(xtals, qrs)
         else:
-            print(f"Local optimization by multi-threads {ncpu}")
+            print(f"Local optimization by multi-threads {self.ncpu}")
             return self.local_optimization_mproc(xtals, self.ncpu, qrs=qrs, pool=pool)
 
     def local_optimization_serial(self, xtals, qrs=False):
@@ -1057,7 +1063,6 @@ def local_optimization_mproc(self, xtals, ncpu, ids=None, qrs=False, pool=None):
             ids = range(len(xtals))
 
         N_cycle = int(np.ceil(len(xtals) / ncpu))
-
         # Generator to create arg_lists for multiprocessing tasks
         def generate_args_lists():
             for i in range(ncpu):
@@ -1071,7 +1076,8 @@ def generate_args_lists():
                 my_args = [_xtals, _ids, mutates, job_tags, *args, self.timeout]
                 yield tuple(my_args)  # Yield args instead of appending to a list
 
-        self.logging.info(f"Rank {self.rank} assign args in local_opt_mproc")
+        self.print_memory_usage()
+        self.logging.info(f"Rank {self.rank} assign local_opt")
 
         gen_results = []
         # Stream the results to avoid holding too much in memory at once
@@ -1081,8 +1087,9 @@ def generate_args_lists():
                 for _res in result:
                     gen_results.append(_res)
             # Explicitly delete the result and call garbage collection
-            def result
+            del result
             gc.collect()
+        self.logging.info(f"Rank {self.rank} finish local_opt {len(gen_results)}")
 
         return gen_results