Added SLURM queue check intervals into SlurmConfig class Fixes #70

Author: Roy-Haolin-Du

a3fe/analyse/mbar.py

@@ -234,7 +234,7 @@ def submit_mbar_slurm(
     if wait:
         for job in jobs:
             while job in virtual_queue.queue:
-                _sleep(30)
+                _sleep(slurm_config.job_submission_wait)
                 virtual_queue.update()
 
     return jobs, mbar_out_files, tmp_files

a3fe/configuration/slurm_config.py

@@ -33,6 +33,8 @@ class SlurmConfig(_BaseModel):
     extra_options: _Dict[str, str] = _Field(
         {}, description="Extra options to pass to SLURM. For example, {'account': 'qt'}"
     )
+    queue_check_interval: int = _Field(30, ge=1, description="Interval in seconds between SLURM queue status checks.")
+    job_submission_wait: int = _Field(300, ge=1, description="Wait time in seconds for job submission to SLURM queue.")
 
     model_config = _ConfigDict(validate_assignment=True)
 

a3fe/run/_simulation_runner.py

@@ -667,9 +667,9 @@ def running(self) -> bool:
     def wait(self) -> None:
         f"""Wait for the {self.__class__.__name__} to finish running."""
         # Give the simulation runner a chance to start
-        _sleep(30)
+        _sleep(self.slurm_config.queue_check_interval)
         while self.running:
-            _sleep(30)  # Check every 30 seconds
+            _sleep(self.slurm_config.queue_check_interval)  # Check every 30 seconds
 
     def get_tot_simtime(self, run_nos: _Optional[_List[int]] = None) -> float:
         f"""

a3fe/run/leg.py

@@ -788,14 +788,14 @@ def _run_slurm(
             self._logger.info(
                 f"Waiting for job {job} to be submitted to the real slurm queue"
             )
-            _sleep(5 * 60)
+            _sleep(self.slurm_config.job_submission_wait)
             self.virtual_queue.update()
 
         # Wait for the job to complete if we've specified wait
         if wait:
             while job in self.virtual_queue.queue:
                 self._logger.info(f"Waiting for job {job} to complete")
-                _sleep(30)
+                _sleep(self.slurm_config.queue_check_interval)
                 self.virtual_queue.update()
 
     def run(

a3fe/run/stage.py

@@ -1158,19 +1158,19 @@ def wait(self) -> None:
         # Override the base class method so that we can update the
         # virtual queue
         # Give the simulations a chance to start
-        _sleep(30)
+        _sleep(self.slurm_config.queue_check_interval)
         self.virtual_queue.update()
         while self.running:
-            _sleep(30)  # Check every 30 seconds
+            _sleep(self.slurm_config.queue_check_interval)  # Check every 30 seconds
             self.virtual_queue.update()
 
     def _wait_ignoring_thread(self) -> None:
         """Wait for the stage to finish running, ignoring the thread."""
-        _sleep(30)
+        _sleep(self.slurm_config.queue_check_interval)
         self.virtual_queue.update()
         # Superclass implementation of running ignores the thread
         while super().running:
-            _sleep(30)
+            _sleep(self.slurm_config.queue_check_interval)
             self.virtual_queue.update()
 
     @property