Allow terminating simulation programmatically (#1828)

Author: tamuri

src/tlo/performance.py

@@ -24,6 +24,7 @@ class MagpieOptions:
     save_sim: bool = False
     save_sim_freq: int = 6
     save_sim_on_end: bool = False
+    wall_time_limit_minutes: float | None = None
 
 
 class Magpie(Module):
@@ -37,6 +38,7 @@ class Magpie(Module):
             save_sim=True,           # save the simulation to a pickle file...
             save_sim_freq=3,         # ...every n months
             save_sim_on_end=True     # save the simulation to a pickle file at the end of the simulation
+            wall_time_limit_minutes=60  # stop the simulation if it runs for more than n minutes (checked every month)
         )
     ```
 
@@ -63,6 +65,9 @@ def initialise_simulation(self, sim: Simulation):
         if self.options.save_sim:
             sim.schedule_event(SaveSimulation(self, self.options.save_sim_freq), sim.start_date)
 
+        if self.options.wall_time_limit_minutes is not None:
+            sim.schedule_event(WallTimeLimit(self, self.options.wall_time_limit_minutes), sim.start_date)
+
     def on_birth(self, mother_id, child_id):
         pass
 
@@ -94,6 +99,23 @@ def apply(self, population):
         logger.info(key="stats", data=data)
 
 
+class WallTimeLimit(RegularEvent, PopulationScopeEventMixin):
+    def __init__(self, module, wall_time_limit_minutes):
+        super().__init__(module, frequency=DateOffset(months=1))
+        self.wall_time_limit_minutes = wall_time_limit_minutes
+        self.start_time = time.time()
+
+    def apply(self, population):
+        now = time.time()
+        elapsed_minutes = (now - self.start_time) / 60
+        if elapsed_minutes >= self.wall_time_limit_minutes:
+            logger.info(
+                key="msg",
+                data=f"Terminating simulation after {self.wall_time_limit_minutes} minutes of runtime"
+            )
+            self.sim.terminate()
+
+
 def make_pickle_filename(sim_date):
     strftime = "%Y%m%dT%H%M%S"
     timestamp = time.strftime(strftime)

src/tlo/simulation.py

@@ -138,6 +138,9 @@ def __init__(
         # Whether simulation has been initialised
         self._initialised = False
 
+        # To allow for graceful termination
+        self._terminate = False
+
     def _configure_logging(
         self,
         filename: Optional[str] = None,
@@ -383,7 +386,7 @@ def run_simulation_to(self, *, to_date: Date) -> None:
         if self.show_progress_bar:
             progress_bar = self._initialise_progress_bar(to_date)
         while (
-            len(self.event_queue) > 0 and self.event_queue.date_of_next_event < to_date
+            len(self.event_queue) > 0 and self.event_queue.date_of_next_event < to_date and not self._terminate
         ):
             event, date = self.event_queue.pop_next_event_and_date()
             if self.show_progress_bar:
@@ -506,6 +509,9 @@ def load_from_pickle(
             simulation._log_filepath = simulation._configure_logging(**log_config)
         return simulation
 
+    def terminate(self):
+        self._terminate = True
+
 
 class EventQueue:
     """A simple priority queue for events.