Update autoinverter.py

inversionson/autoinverter.py

@@ -11,6 +11,7 @@
 from colorama import Fore, Style
 from typing import Union, List
 from inversionson.remote_helpers.helpers import preprocess_remote_gradient
+from salvus.flow import api
 
 init()
 
@@ -64,6 +65,64 @@ def _send_whatsapp_announcement(self):
         client.messages.create(
             body=string, from_=from_whatsapp, to=to_whatsapp
         )
+
+    def jobs_on_remote_site(self, site_name: str, job_number: int) -> bool:
+        """
+        Check how many jobs are running or pending on the remote site
+
+        :param site_name: Name of remote site
+        :type site_name: str
+        :param job_numer: Number of submitting jobs
+        :type job_number: int
+        """
+        max_submitted_jobs_per_user = 30
+
+        # Check job_status
+        job_list = api.get_jobs(
+            limit=100,
+            site_name = site_name, 
+            job_status = ["running", "pending"],
+            update_jobs = True,
+        )
+
+        finished_job = 0
+        for job in job_list:
+            status = job.update_status(force_update=False)
+            if status.name == "finished":
+                finished_job += 1
+
+        remaining_job = len(job_list) - finished_job
+
+        # check job_array_status
+        job_array_list = api.get_job_arrays(
+            limit=100,
+            site_name = site_name, 
+            job_array_status = ["running", "pending"], 
+            update_job_arrays = True,
+        )
+
+        remaining_job_array = 0
+        for job_array in job_array_list:
+            status = job_array.update_status()
+            for job_status in status:
+                if job_status.name == "running" or job_status.name == "pending":
+                    remaining_job_array += 1
+
+        remaining_job += remaining_job_array
+        if max_submitted_jobs_per_user - remaining_job >= job_number:
+            boolean = True
+        else:
+            boolean = False
+
+        if not boolean:
+            print(
+                f"{remaining_job} jobs are on {site_name}: " 
+                "Sleeping for a minute before checking again"
+            )
+            time.sleep(60)
+            return self.jobs_on_remote_site(site_name, job_number)
+
+        return boolean
 
     def prepare_iteration(self, first=False, validation=False):
         """
@@ -257,7 +316,12 @@ def run_forward_simulation(self, event: str):
             and self.comm.project.meshes == "mono-mesh"
         ):
             w.set_mesh(self.comm.project.remote_mesh)
-
+
+        boolean = self.jobs_on_remote_site(self.comm.project.site_name, 1)
+        if boolean:
+            print(f"Submit forward simulation to {self.comm.project.site_name}")
+            pass
+
         self.comm.salvus_flow.submit_job(
             event=event,
             simulation=w,
@@ -346,7 +410,12 @@ def run_adjoint_simulation(self, event: str):
             and self.comm.project.meshes == "mono-mesh"
         ):
             w_adjoint.set_mesh(self.comm.project.remote_mesh)
-
+
+        boolean = self.jobs_on_remote_site(self.comm.project.site_name, 1)
+        if boolean:
+            print(f"Submit adjoint simulation to {self.comm.project.site_name}")
+            pass 
+
         self.comm.salvus_flow.submit_job(
             event=event,
             simulation=w_adjoint,
@@ -719,6 +788,14 @@ def smooth_gradient(self, event: str):
         if self.comm.project.remote_gradient_processing:
             self.comm.smoother.run_remote_smoother(event=event)
         else:
+            boolean = self.jobs_on_remote_site(
+                self.comm.project.smoothing_site_name, 
+                len(self.comm.project.inversion_params),
+            )
+            if boolean:
+                print(f"Submit smooth simulation to {self.comm.project.site_name}")
+                pass
+
             smoothing_config = self.comm.smoother.\
                 generate_smoothing_config(event=event
             )
@@ -1425,6 +1502,7 @@ def compute_misfit_and_gradient(self, task: str, verbose: str):
                 )
             )
 
+            print(f"Event: {event} \n")         
             self.run_forward_simulation(event)
 
             print(Fore.RED + "\n =========================== \n")
@@ -1492,6 +1570,8 @@ def compute_misfit_and_gradient(self, task: str, verbose: str):
                         ":rocket: | Run adjoint simulation", use_aliases=True
                     )
                 )
+
+                print(f"Event: {event} \n")
                 self.run_adjoint_simulation(event)
 
         self.compute_misfit_on_validation_data()
@@ -1711,7 +1791,8 @@ def compute_misfit(self, task: str, verbose: str):
                     ":rocket: | Run forward simulation", use_aliases=True
                 )
             )
-
+
+            print(f"Event: {event} \n")
             self.run_forward_simulation(event)
 
             print(Fore.RED + "\n =========================== \n")