Speed up RAT calculations in python

rascal2/core/runner.py

@@ -3,57 +3,80 @@
 import os
 from dataclasses import dataclass
 from logging import INFO
-from multiprocessing import Process, Queue
+from multiprocessing import Event, Process, Queue
 
 import ratapi as rat
 from PyQt6 import QtCore
-from ratapi.utils.enums import Procedures
 
 from rascal2.config import MatlabHelper, get_matlab_engine
 
+NUMBER_PROCESSES = 5
+
 
 class RATRunner(QtCore.QObject):
     """Class for running rat."""
 
     event_received = QtCore.pyqtSignal()
     finished = QtCore.pyqtSignal()
     stopped = QtCore.pyqtSignal()
+    go_event = Event()
+    processes_list_go_exit_events = []
 
-    def __init__(self, rat_inputs, procedure: Procedures, display_on: bool):
+    def __init__(self, parent=None, start_runners_early: bool = True, num_processes: int = NUMBER_PROCESSES):
         super().__init__()
+        self.parent = parent
         self.timer = QtCore.QTimer()
         self.timer.setInterval(1)
         self.timer.timeout.connect(self.check_queue)
+        self.num_processes = num_processes
+        self.start_runners_early = start_runners_early
 
         # this queue handles both event data and results
         self.queue = Queue()
-        matlab_helper = MatlabHelper()
-        self.process = Process(
-            target=run,
-            args=(
-                self.queue,
-                rat_inputs,
-                procedure,
-                display_on,
-                matlab_helper.ready_event,
-                matlab_helper.engine_output,
-            ),
-        )
+        self.arg_queue = Queue()
+        self.go_event = Event()
+        self.exit_event = Event()
+        self.rat_inputs = None
+        self.procedure = None
+        self.display_on = None
+        self.processes_list = []
+        self.refresh_process_list()
+        self.process = None
         self.updated_problem = None
         self.results = None
         self.error = None
         self.events = []
 
+    def set_runner_args(self, rat_inputs, procedure, display_on: bool, working_dir: str):
+        self.arg_queue.put((rat_inputs, procedure, display_on, working_dir))
+        self.rat_inputs = rat_inputs
+        self.display_on = display_on
+
     def start(self):
         """Start the calculation."""
-        self.process.start()
+        if not self.process:
+            self.process, (self.go_event, self.exit_event) = self.get_new_process()
+        self.go_event.set()
+        if not self.process.is_alive():
+            self.process.start()
         self.timer.start()
 
+    def get_new_process(self):
+        if not self.processes_list:
+            self.refresh_process_list()
+        return self.processes_list.pop(0), self.processes_list_go_exit_events.pop(0)
+
     def interrupt(self):
         """Interrupt the running process."""
         self.timer.stop()
         self.process.kill()
         self.stopped.emit()
+        self.go_event.clear()
+        self.clear_process()
+
+    def clear_process(self):
+        """Clear the current process."""
+        self.process = None
 
     def check_queue(self):
         """Check for new data in the queue."""
@@ -63,65 +86,151 @@ def check_queue(self):
         for item in iter(self.queue.get, None):
             if isinstance(item, tuple):
                 self.updated_problem, self.results = item
+                self.go_event.clear()
                 self.finished.emit()
             elif isinstance(item, Exception):
                 self.error = item
+                self.go_event.clear()
                 self.stopped.emit()
+            elif isinstance(item, list):
+                return item[0]
             else:  # else, assume item is an event
                 self.events.append(item)
                 self.event_received.emit()
 
-
-def run(queue, rat_inputs: tuple, procedure: str, display: bool, engine_ready, engine_output):
+    def refresh_process_list(self):
+        self.processes_list_go_exit_events = [(Event(), Event()) for _ in range(self.num_processes)]
+        matlab_helper = MatlabHelper()
+        self.processes_list = [
+            Process(
+                target=run,
+                args=(
+                    self.queue,
+                    self.arg_queue,
+                    self.processes_list_go_exit_events[ind][0],
+                    self.processes_list_go_exit_events[ind][1],
+                    matlab_helper.ready_event,
+                    matlab_helper.engine_output,
+                ),
+            )
+            for ind in range(self.num_processes)
+        ]
+
+    def clear_queues(self):
+        self.queue.empty()
+        self.arg_queue.empty()
+        self.events.clear()
+        self.go_event.clear()
+        self.exit_event.clear()
+
+    def start_processes(self):
+        if self.start_runners_early:
+            for process in self.processes_list:
+                process.start()
+
+    def stop_processes(self):
+        self.exit_event.set()
+        self.go_event.set()
+        for go_event, exit_event in self.processes_list_go_exit_events:
+            exit_event.set()
+            go_event.set()
+        for process in self.processes_list:
+            if process.is_alive():
+                process.kill()
+        self.process = None
+        self.processes_list.clear()
+        self.clear_queues()
+        self.processes_list_go_exit_events.clear()
+        self.queue.close()
+        self.arg_queue.close()
+
+
+def init_matlab_engine(problem_definition, engine_ready, engine_output, queue):
+    """Initialise the Matlab engine if using a Matlab custom file and returns the engine future if available."""
+    engine_future = rat.wrappers.MatlabWrapper.loader
+    if engine_future is None and any([file["language"] == "matlab" for file in problem_definition.customFiles.files]):
+        if not engine_output:
+            queue.put(LogData(INFO, "Attempting to start Matlab..."))
+
+        result = get_matlab_engine(engine_ready, engine_output)
+        if isinstance(result, Exception):
+            raise result
+        else:
+            engine_future = result
+            engine_future.result().cd(os.getcwd())
+    return engine_future
+
+
+def stop_matlab_engine(engine_future):
+    """Exit the Matlab engine future if present."""
+    if engine_future is not None:
+        engine_future.result().exit()
+
+
+def run(queue: Queue, arg_queue: Queue, go_event, exit_event, engine_ready, engine_output):
     """Run RAT and put the result into the queue.
 
     Parameters
     ----------
     queue : Queue
         The interprocess queue for the RATRunner.
-    rat_inputs : tuple
-        The C++ inputs for rat.
-    procedure : str
-        The optimisation procedure.
-    display : bool
-        Whether to display events.
+    arg_queue :
+        A queue of arguments used to initialize the RAT process, passed from the Main Presenter
 
     """
-    problem_definition, cpp_controls = rat_inputs
-
-    if display:
-        rat.events.register(rat.events.EventTypes.Message, queue.put)
-        rat.events.register(rat.events.EventTypes.Progress, queue.put)
-        rat.events.register(rat.events.EventTypes.Plot, queue.put)
-        queue.put(LogData(INFO, "Starting RAT"))
-
+    engine_future = None
+    while True:
+        go_event.wait()
+        if exit_event.is_set():
+            stop_matlab_engine(engine_future)
+            return
+        rat_inputs, procedure, display, working_dir = arg_queue.get()
+        os.chdir(working_dir)
+        problem_definition, cpp_controls = rat_inputs
+
+        if display:
+            rat.events.register(rat.events.EventTypes.Message, queue.put)
+            rat.events.register(rat.events.EventTypes.Progress, queue.put)
+            rat.events.register(rat.events.EventTypes.Plot, queue.put)
+            queue.put(LogData(INFO, "Starting RAT"))
+
+        try:
+            engine_future = init_matlab_engine(problem_definition, engine_ready, engine_output, queue)
+            problem_definition, output_results, bayes_results = rat.rat_core.RATMain(problem_definition, cpp_controls)
+            results = rat.outputs.make_results(procedure, output_results, bayes_results)
+        except Exception as err:
+            queue.put(err)
+            go_event.clear()
+            continue
+
+        if display:
+            queue.put(LogData(INFO, "Finished RAT"))
+            rat.events.clear()
+
+        queue.put((problem_definition, results))
+        go_event.clear()
+
+
+def run_matlab_init_engine(queue, engine_output, engine_ready, display_on):
+    """Get the engine future from the matlab engine and put in queue if successfully."""
     try:
-        engine_future = None
-        if any([file["language"] == "matlab" for file in problem_definition.customFiles.files]):
-            if not engine_output:
-                queue.put(LogData(INFO, "Attempting to start Matlab..."))
-
-            result = get_matlab_engine(engine_ready, engine_output)
-            if isinstance(result, Exception):
-                raise result
-            else:
-                engine_future = result
-                engine_future.result().cd(os.getcwd())
-
-        problem_definition, output_results, bayes_results = rat.rat_core.RATMain(problem_definition, cpp_controls)
-        results = rat.outputs.make_results(procedure, output_results, bayes_results)
-        if engine_future is not None:
-            engine_future.result().exit()
+        if not engine_output and display_on:
+            queue.put(LogData(INFO, "Attempting to start Matlab..."))
+
+        result = get_matlab_engine(engine_ready, engine_output)
+        if display_on:
+            queue.put(LogData(INFO, "Got Matlab engine"))
+        if isinstance(result, Exception):
+            raise result
+        else:
+            engine_future = result
+            engine_future.result().cd(os.getcwd())
+            queue.put([engine_future])
+
     except Exception as err:
         queue.put(err)
         return
 
-    if display:
-        queue.put(LogData(INFO, "Finished RAT"))
-        rat.events.clear()
-
-    queue.put((problem_definition, results))
-
 
 @dataclass
 class LogData:

rascal2/ui/presenter.py

@@ -5,6 +5,7 @@
 
 import ratapi as rat
 import ratapi.wrappers
+from PyQt6.QtCore import QCoreApplication
 
 from rascal2.config import LOGGER, SETTINGS, MatlabHelper
 from rascal2.core import commands
@@ -15,6 +16,8 @@
 
 from .model import MainWindowModel
 
+START_PROCESSES = bool(os.getenv("START_PROCESSES", "True"))
+
 
 class MainWindowPresenter:
     """Facilitates interaction between View and Model.
@@ -29,6 +32,11 @@ def __init__(self, view):
         self.view = view
         self.model = MainWindowModel()
         self.worker = None
+        self.runner = RATRunner(self, start_runners_early=START_PROCESSES)
+        self.runner.finished.connect(self.handle_results)
+        self.runner.stopped.connect(self.handle_interrupt)
+        self.runner.event_received.connect(self.handle_event)
+        self.runner.start_processes()
 
     def create_project(self, name: str, save_path: str):
         """Create a new RAT project and controls object then initialise UI.
@@ -227,13 +235,10 @@ def run(self):
         self.view.plot_widget.bayes_plots_button.setVisible(False)
 
         self.model.controls.initialise_IPC()
+        working_dir = os.getcwd()
         rat_inputs = rat.inputs.make_input(self.model.project, self.model.controls)
         display_on = self.model.controls.display != rat.utils.enums.Display.Off
-
-        self.runner = RATRunner(rat_inputs, self.model.controls.procedure, display_on)
-        self.runner.finished.connect(self.handle_results)
-        self.runner.stopped.connect(self.handle_interrupt)
-        self.runner.event_received.connect(self.handle_event)
+        self.runner.set_runner_args(rat_inputs, self.model.controls.procedure, display_on, working_dir)
         self.view.terminal_widget.write("Initializing RAT Process...")
         self.runner.start()
 
@@ -249,6 +254,7 @@ def handle_results(self):
         )
         self.view.handle_results(self.runner.results)
         self.model.controls.delete_IPC()
+        self.runner.clear_queues()
 
     def handle_interrupt(self):
         """Handle a RAT run being interrupted."""
@@ -274,6 +280,7 @@ def handle_event(self):
                 self.view.plot_widget.plot_with_blit(event)
             case LogData():
                 LOGGER.log(event.level, event.msg)
+        QCoreApplication.processEvents()
 
     def edit_project(self, updated_project: dict, preview: bool = True) -> None:
         """Edit the Project with a dictionary of attributes.

rascal2/ui/view.py

@@ -81,6 +81,8 @@ def closeEvent(self, event):
             event.accept()
         else:
             event.ignore()
+        self.presenter.runner.stop_processes()
+        event.accept()
 
     def show_project_dialog(self, dialog: StartupDialog):
         """Show a startup dialog of a given type.

tests/conftest.py

@@ -19,6 +19,18 @@ def global_setting():
     return GLOBAL_SETTING
 
 
+@pytest.fixture(autouse=True)
+@patch("rascal2.core.runner.NUMBER_PROCESSES")
+def fix_num_processes(num_processes):
+    num_processes.return_value = 1
+    yield
+
+
+@pytest.fixture(scope="function", autouse=True)
+def mock_start_processes_setting(monkeypatch):
+    monkeypatch.setenv("START_PROCESSES", "False")
+
+
 @pytest.fixture(scope="session", autouse=True)
 def mock_setting(request):
     global GLOBAL_SETTING