Working RMG Adapter

Author: calvinp0

t3/main.py

@@ -631,45 +631,17 @@ def run_rmg(self, restart_rmg: bool = False):
                                 walltime=self.t3['options']['max_RMG_walltime'],
                                 max_iterations=self.t3['options']['max_rmg_iterations'],
                                 verbose=self.verbose,
-                                t3_project_name=self.project
+                                t3_project_name=self.project,
+                                restart_rmg=restart_rmg,
                                 )
         rmg_adapter.run_rmg()
         if not os.path.isfile(self.paths['RMG input']):
             raise ValueError(f"The RMG input file {self.paths['RMG input']} could not be written.")
         tic = datetime.datetime.now()
 
         max_rmg_exceptions_allowed = self.t3['options']['max_RMG_exceptions_allowed']
-    
 
-
-
-        
-        
-    #
-
-        write_rmg_input_file(
-            rmg=self.rmg,
-            t3=self.t3,
-            iteration=self.iteration,
-            path=self.paths['RMG input'],
-            walltime=self.t3['options']['max_RMG_walltime'],
-        )
-        if not os.path.isfile(self.paths['RMG input']):
-            raise ValueError(f"The RMG input file {self.paths['RMG input']} could not be written.")
-        tic = datetime.datetime.now()
-
-        max_rmg_exceptions_allowed = self.t3['options']['max_RMG_exceptions_allowed']
-        rmg_exception_encountered = rmg_runner(rmg_input_file_path=self.paths['RMG input'],
-                                               job_log_path=self.paths['RMG job log'],
-                                               logger=self.logger,
-                                               memory=self.rmg['memory'] * 1000 if self.rmg['memory'] is not None else None,
-                                               max_iterations=self.t3['options']['max_rmg_iterations'],
-                                               verbose=self.verbose,
-                                               t3_project_name=self.project,
-                                               rmg_execution_type=self.rmg['rmg_execution_type'],
-                                               restart_rmg=restart_rmg,
-                                               )
-        if rmg_exception_encountered:
+        if rmg_adapter.rmg_exception_encountered:
             self.rmg_exceptions_counter += 1
             if self.rmg_exceptions_counter > max_rmg_exceptions_allowed:
                 self.logger.error(f'This is the {get_number_with_ordinal_indicator(self.rmg_exceptions_counter)} '
@@ -681,10 +653,14 @@ def run_rmg(self, restart_rmg: bool = False):
                                     f'This is the {get_number_with_ordinal_indicator(self.rmg_exceptions_counter)} '
                                     f'exception raised by RMG.\n'
                                     f'The maximum number of exceptions allowed is {max_rmg_exceptions_allowed}.')
-
         elapsed_time = time_lapse(tic)
         self.logger.info(f'RMG terminated, execution time: {elapsed_time}')
 
+
+
+        
+        
+    #
     def determine_species_and_reactions_to_calculate(self) -> bool:
         """
         Determine which species and reactions in the executed RMG job should be calculated.

t3/runners/rmg_adapter.py

@@ -3,6 +3,7 @@
 import math
 import shutil
 import datetime
+import time
 
 from mako.template import Template
 from t3.utils.writer import to_camel_case
@@ -18,6 +19,7 @@
               }
 
 RMG_EXECUTION_TYPE = settings['execution_type']['rmg']
+MAX_RMG_RUNS_PER_ITERATION = 5 # TODO: Why is this hard-coded?
 submit_filenames = settings['submit_filenames']
 rmg_memory = settings['rmg_initial_memory']
 if RMG_EXECUTION_TYPE == 'queue':
@@ -70,37 +72,68 @@ def __init__(self,
         self.server = server or SERVER
         self.testing = testing
         self.logger = logger
-        
+        self.previous_job_status = None
+        self.time_running = 0
+        self.restart_rmg = restart_rmg
 
         if not os.path.isdir(local_t3_path):
             os.makedirs(local_t3_path)
 
         self.files_to_upload = list()
         self.folder_to_download = None
-        
+        self.rmg_errors = list()
+        self.rmg_run_count = 0
+
     def run_rmg(self):
         """
         Run RMG
         """
-        self.set_cpu_and_mem()
-        self.set_file_paths()
-        self.set_files()
-
         if self.rmg_execution_type == 'incore':
             self.execute_incore()
-        elif self.rmg_execution_type == 'queue':
-            self.execute_queue()
-            # While the job is running, periodically check the status of the job
-            while self.job_status == 'running':
-                self.determine_rmg_job_status()
-            # Once the job is done, download the results
-            if self.job_status == 'done':
-                self.download_files()
+        elif self.rmg_execution_type == 'local' or self.rmg_execution_type == 'queue':
+            while self.cont_run_rmg:
+                self.set_cpu_and_mem()
+                self.set_file_paths()
+                self.set_files()
+
+
+                if self.rmg_execution_type == 'queue':
+                    self.rmg_run_count += 1
+                    self.execute_queue()
+                    # While the job is running, periodically check the status of the job
+                    while self.job_status == 'running':
+                        # Wait for 5 minutes before checking again
+                        time.sleep(300)
+                        self.time_running += 300
+                        self.determine_rmg_job_status()
+                        # Log the status of the job
+                        # If we do it every 5 mins, the log file will be flooded with the same message
+                        # So only log if the status has changed or if 30 mins have passed
 
+                        if self.job_status != self.previous_job_status or self.time_running % 1800 == 0:
+                            self.logger.info(f'RMG-{self.iteration}_iteration job status: {self.job_status}')
+                            self.previous_job_status = self.job_status
 
+                    # Once the job is done, download the results
+                    if self.job_status == 'done':
+                        # Log that the job is done and will download the results
+                        self.logger.info(f'RMG-{self.iteration}_iteration job status: {self.job_status}, downloading results...')
+                        self.download_files()
+
+                # Need to check for convergence or errors
+                self.check_convergance()
+                # Get local err file path
+                err_path = os.path.join(self.local_rmg_path, 'err.txt')
+                if os.path.isfile(err_path):
+                    os.rename(err_path, os.path.join(self.local_rmg_path, f'err_{datetime.datetime.now().strftime("%b%d_%Y_%H:%M:%S")}.txt'))
+                self.rmg_errors.append(self.error)
         else:
             raise ValueError(f'RMG execution type {self.rmg_execution_type} is not supported.')
 
+        # set self.rmg_exceptions the opposite of self.rmg_converged
+        self.rmg_exception_encountered = not self.rmg_converged
+        
+        
     def write_rmg_input_file(self):
         """
         Write an RMG input file to the given file path.
@@ -424,7 +457,6 @@ def set_cpu_and_mem(self):
             self.submit_script_memory = math.ceil(total_submit_script_memory)  # in MB
         self.set_input_file_memory()
 
-
     def set_files(self) -> None:
         """
         Set files to be uploaded and downloaded. Writes the files if needed.
@@ -444,17 +476,35 @@ def set_files(self) -> None:
         if self.rmg_execution_type != 'incore':
             # We need a submit script (submitted local or SSH)
             self.write_submit_script()
+
             self.files_to_upload.append(self.get_file_property_dictionary(
                 file_name=submit_filenames[CLUSTER_SOFT]))
         # 1.2. RMG input file
         self.write_rmg_input_file()
+        # If this a restart, we need to upload the restart file
+        if self.restart_rmg:
+            restart_string = "restartFromSeed(path='seed')"
+            with open(self.rmg_input_file_path, 'r') as f:
+                content = f.read()
+            seed_path = os.path.join(self.local_rmg_path, 'seed')
+            if restart_string not in content and os.path.isdir(seed_path) and os.listdir(seed_path):
+                if os.path.isfile(os.path.join(self.local_rmg_path, 'restart_from_seed.py')):
+                    os.rename(src=os.path.join(self.local_rmg_path, 'input.py'),
+                                dst=os.path.join(self.local_rmg_path, 'input.py.old'))
+                    os.rename(src=os.path.join(self.local_rmg_path, 'restart_from_seed.py'),
+                            dst=os.path.join(self.local_rmg_path, 'input.py'))
+                elif os.path.isfile(os.path.join(self.local_rmg_path, 'input.py')):
+                    with open(os.path.join(self.local_rmg_path, 'input.py'), 'r') as f:
+                        content = f.read()
+                    with open(os.path.join(self.local_rmg_path, 'input.py'), 'w') as f:
+                        f.write(restart_string + '\n\n' + content)
+        
         self.files_to_upload.append(self.get_file_property_dictionary(file_name='input.py'))
 
         # 2. ** Download **
         # 2.1. RMG output folder
         self.folder_to_download = self.remote_path
 
-    
     def set_additional_file_paths(self) -> None:
         """
         Set additional file paths to be uploaded and downloaded.
@@ -589,6 +639,7 @@ def upload_files(self) -> None:
                         except shutil.SameFileError:
                             pass
             self.initial_time = datetime.datetime.now()
+    
     def determine_rmg_job_status(self) -> None:
         """
         Determine the RMG job status.
@@ -608,3 +659,45 @@ def download_files(self) -> None:
            # Also, even if the job is submitted to the queue, no need to download files if the server is local.
             with SSHClient(self.server) as ssh:
                 ssh.download_folder(remote_folder_path=self.remote_path, local_folder_path=self.local_rmg_path)
+
+    def check_convergance(self) -> None:
+        self.rmg_converged, self.error = False, None
+        rmg_log_path = os.path.join(self.local_rmg_path, 'RMG.log')
+        rmg_err_path = os.path.join(self.local_rmg_path, 'err.txt')
+        if os.path.isfile(rmg_log_path):
+            with open(rmg_log_path, 'r') as f:
+                # Read the lines of the log file in reverse order
+                lines = f.readlines()[::-1]
+                for line in lines:
+                    if 'MODEL GENERATION COMPLETED' in line:
+                        self.rmg_converged = True
+                        break
+        if not self.rmg_converged and os.path.isfile(rmg_err_path):
+            with open(rmg_err_path, 'r') as f:
+                lines = f.readlines()[::-1]
+            for line in lines:
+                if 'Error' in line:
+                    self.error = line.strip()
+                    break
+        return self.rmg_converged, self.error
+
+    def convergance_failure(self) -> None:
+
+        if not self.rmg_converged:
+            if self.error is not None:
+                self.logger.error(f'RMG job {self.job_id} failed with error: {self.error}')
+            # Check if memory was the failure - TODO: this is not working
+            #new_memory = self.get_new_memory_rmg_run()
+
+        self.cont_run_rmg = not self.rmg_converged \
+                          and self.rmg_run_count < MAX_RMG_RUNS_PER_ITERATION \
+                          and not(len(self.rmg_errors)) >=2 and self.error is not None \
+                          and self.error == self.rmg_errors[-2]
+        self.restart_rmg = False if self.error is not None and 'Could not find one or more of the required files/directories ' \
+                                                         'for restarting from a seed mechanism' in self.error else True
+    
+    def get_new_memory_rmg_run(self) -> None:
+        """
+        Get a new memory for the RMG job.
+        """
+        pass

t3/utils/ssh.py

@@ -11,6 +11,7 @@
 import os
 import time
 from typing import Any, Callable, List, Optional, Tuple, Union
+import stat
 
 import paramiko
 
@@ -226,12 +227,18 @@ def download_folder(self,
                 remote_filepath = remote_folder_path + '/' + filename
                 local_filepath = os.path.join(local_folder_path, filename)
 
+
+
                  # Download files
-                if item.isfile():
+                 # stat.S_ISREG(item.st_mode) is True if item is a file (not a folder)
+                if stat.S_ISREG(item.st_mode):
                     self._sftp.get(remote_filepath, local_filepath)
 
                 # Recursively download folders
-                elif item.isdir():
+                # stat.S_ISDIR(item.st_mode) is True if item is a folder
+                elif stat.S_ISDIR(item.st_mode):
+                    # create the folder in the local path
+                    os.makedirs(local_filepath, exist_ok=True)
                     self.download_folder(remote_filepath, local_filepath)
         except IOError:
             logger.warning(f'Got an IOError when trying to download file '