Update simple_cerebros_random_search.py

Add an option to purge model storage after a trial to keep memory from accumulating in a containerized run.

Author: david-thrower

cerebros/simplecerebrosrandomsearch/simple_cerebros_random_search.py

@@ -14,6 +14,7 @@
 from multiprocessing import Process, Lock
 import os
 from gc import collect
+from shutil import rmtree
 
 
 # import optuna
@@ -585,8 +586,14 @@ def run_random_search(self):
         print(f"Best model name: {self.best_model_path}")
         return best
 
-    def get_best_model(self):
+    def purge_model_storage(self):
+        path_0 = f"{self.project_name}/models"
+        rmtree(path_0)
+
+    def get_best_model(self, purge_model_storage_files=False):
         best_model = tf.keras.models.load_model(self.best_model_path)
+        if purge_model_storage:
+            self.purge_model_storage()        
         return best_model
 
 # ->