diff --git a/inversionson/autoinverter.py b/inversionson/autoinverter.py index 37b3bad..af37e83 100644 --- a/inversionson/autoinverter.py +++ b/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")