Trial using threading to run events simultaneously

src/scripts/profiling/scale_run.py

@@ -15,6 +15,7 @@
 from tlo import Date, Simulation, logging
 from tlo.analysis.utils import parse_log_file as parse_log_file_fn
 from tlo.methods.fullmodel import fullmodel
+from tlo.threaded_simulation import ThreadedSimulation
 
 _TLO_ROOT: Path = Path(__file__).parents[3].resolve()
 _TLO_OUTPUT_DIR: Path = (_TLO_ROOT / "outputs").resolve()
@@ -55,6 +56,7 @@ def scale_run(
     ignore_warnings: bool = False,
     log_final_population_checksum: bool = True,
     profiler: Optional["Profiler"] = None,
+    n_threads: Optional[int] = 0,
 ) -> Simulation:
     if ignore_warnings:
         warnings.filterwarnings("ignore")
@@ -74,12 +76,16 @@ def scale_run(
         "suppress_stdout": disable_log_output_to_stdout,
     }
 
-    sim = Simulation(
-        start_date=start_date,
-        seed=seed,
-        log_config=log_config,
-        show_progress_bar=show_progress_bar,
-    )
+    sim_args =  {
+        "start_date": start_date,
+        "seed": seed,
+        "log_config": log_config,
+        "show_progress_bar": show_progress_bar,
+    }
+    if n_threads:
+        sim = ThreadedSimulation(n_threads=n_threads, **sim_args)
+    else:
+        sim = Simulation(**sim_args)
 
     # Register the appropriate modules with the arguments passed through
     sim.register(
@@ -269,6 +275,12 @@ def scale_run(
         ),
         action="store_true",
     )
+    parser.add_argument(
+        "--n-threads",
+        help="Run a threaded simulation using the given number of threaded workers",
+        type=int,
+        default=0,
+    )
     args = parser.parse_args()
     args_dict = vars(args)
 

src/tlo/simulation.py

@@ -1,12 +1,13 @@
 """The main simulation controller."""
+from __future__ import annotations
 
 import datetime
 import heapq
 import itertools
 import time
 from collections import OrderedDict
 from pathlib import Path
-from typing import Dict, Optional, Union
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 
 import numpy as np
 
@@ -15,11 +16,14 @@
 from tlo.events import Event, IndividualScopeEventMixin
 from tlo.progressbar import ProgressBar
 
+if TYPE_CHECKING:
+    from tlo import Module
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
 
-class Simulation:
+class _BaseSimulation:
     """The main control centre for a simulation.
 
     This class contains the core simulation logic and event queue, and holds
@@ -41,8 +45,15 @@ class Simulation:
         The simulation-level random number generator.
         Note that individual modules also have their own random number generator
         with independent state.
+
+    The `step_through_events` method is implemented by the `Simulation` and
+    `ThreadedSimulation` classes, which controls how the simulation events are
+    fired.
     """
 
+    __name__: str = "_BaseSimulation"
+    modules: OrderedDict[str, Module]
+
     def __init__(self, *, start_date: Date, seed: int = None, log_config: dict = None,
                  show_progress_bar=False):
         """Create a new simulation.
@@ -63,6 +74,7 @@ def __init__(self, *, start_date: Date, seed: int = None, log_config: dict = Non
         self.population: Optional[Population] = None
 
         self.show_progress_bar = show_progress_bar
+        self.progress_bar = None
 
         # logging
         if log_config is None:
@@ -205,37 +217,18 @@ def simulate(self, *, end_date):
         for module in self.modules.values():
             module.initialise_simulation(self)
 
-        progress_bar = None
         if self.show_progress_bar:
-            num_simulated_days = (end_date - self.start_date).days
-            progress_bar = ProgressBar(
+            num_simulated_days = (self.end_date - self.start_date).days
+            self.progress_bar = ProgressBar(
                 num_simulated_days, "Simulation progress", unit="day")
-            progress_bar.start()
-
-        while self.event_queue:
-            event, date = self.event_queue.next_event()
+            self.progress_bar.start()
 
-            if self.show_progress_bar:
-                simulation_day = (date - self.start_date).days
-                stats_dict = {
-                    "date": str(date.date()),
-                    "dataframe size": str(len(self.population.props)),
-                    "queued events": str(len(self.event_queue)),
-                }
-                if "HealthSystem" in self.modules:
-                    stats_dict["queued HSI events"] = str(
-                        len(self.modules["HealthSystem"].HSI_EVENT_QUEUE)
-                    )
-                progress_bar.update(simulation_day, stats_dict=stats_dict)
-
-            if date >= end_date:
-                self.date = end_date
-                break
-            self.fire_single_event(event, date)
+        # Run the simulation by firing events in the queue
+        self.step_through_events()
 
         # The simulation has ended.
         if self.show_progress_bar:
-            progress_bar.stop()
+            self.progress_bar.stop()
 
         for module in self.modules.values():
             module.on_simulation_end()
@@ -253,6 +246,17 @@ def simulate(self, *, end_date):
             finally:
                 self.output_file.release()
 
+    def step_through_events(self) -> None:
+        """
+        Method for forward-propagating the simulation, by executing
+        the scheduled events in the queue. This is overwritten by
+        inheriting classes.
+        """
+        raise NotImplementedError(
+            f"{self.__name__} is not intended to be simulated, "
+            "use either Simulation or ThreadedSimulation to run a simulation."
+        )
+
     def schedule_event(self, event, date):
         """Schedule an event to happen on the given future date.
 
@@ -269,15 +273,6 @@ def schedule_event(self, event, date):
 
         self.event_queue.schedule(event=event, date=date)
 
-    def fire_single_event(self, event, date):
-        """Fires the event once for the given date
-
-        :param event: :py:class:`Event` to fire
-        :param date: the date of the event
-        """
-        self.date = date
-        event.run()
-
     def do_birth(self, mother_id):
         """Create a new child person.
 
@@ -308,6 +303,23 @@ def find_events_for_person(self, person_id: int):
 
         return person_events
 
+    def update_progress_bar(self, new_date: Date):
+        """
+        Updates the simulation's progress bar, if this is in use.
+        """
+        if self.show_progress_bar:
+            simulation_day = (new_date - self.start_date).days
+            stats_dict = {
+                "date": str(new_date.date()),
+                "dataframe size": str(len(self.population.props)),
+                "queued events": str(len(self.event_queue)),
+            }
+            if "HealthSystem" in self.modules:
+                stats_dict["queued HSI events"] = str(
+                    len(self.modules["HealthSystem"].HSI_EVENT_QUEUE)
+                )
+            self.progress_bar.update(simulation_day, stats_dict=stats_dict)
+
 
 class EventQueue:
     """A simple priority queue for events.
@@ -329,7 +341,7 @@ def schedule(self, event, date):
         entry = (date, event.priority, next(self.counter), event)
         heapq.heappush(self.queue, entry)
 
-    def next_event(self):
+    def next_event(self) -> Tuple[Event, Date]:
         """Get the earliest event in the queue.
 
         :returns: an (event, date) pair
@@ -340,3 +352,35 @@ def next_event(self):
     def __len__(self):
         """:return: the length of the queue"""
         return len(self.queue)
+
+
+class Simulation(_BaseSimulation):
+    """
+    Default simulation type, which runs a serial simulation.
+    Events in the event_queue are executed in sequence, one
+    after the other, in the order they appear in the queue.
+
+    See `_BaseSimulation` for more details.
+    """
+
+    def step_through_events(self) -> None:
+        """Serial simulation: events are executed in the
+        order they occur in the queue."""
+        while self.event_queue:
+            event, date = self.event_queue.next_event()
+
+            self.update_progress_bar(date)
+
+            if date >= self.end_date:
+                self.date = self.end_date
+                break
+            self.fire_single_event(event, date)
+
+    def fire_single_event(self, event, date):
+        """Fires the event once for the given date
+
+        :param event: :py:class:`Event` to fire
+        :param date: the date of the event
+        """
+        self.date = date
+        event.run()