Optimize multiprocessing worker performance and improve state management

Author: PierreRaybaut

datalab/gui/processor/base.py

@@ -13,6 +13,7 @@
 import time
 import warnings
 from dataclasses import dataclass
+from enum import Enum, auto
 from multiprocessing.pool import Pool
 from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Optional
 
@@ -94,12 +95,81 @@ def run_with_env(func: Callable, args: tuple, env_json: str) -> CompOut:
     return wng_err_func(func, args)
 
 
+class WorkerState(Enum):
+    """Worker states for computation lifecycle."""
+
+    IDLE = auto()  # Ready to start new computation
+    STARTING = auto()  # Computation starting (prevents race conditions)
+    RUNNING = auto()  # Computation in progress
+    FINISHED = auto()  # Computation completed, result available
+
+
+class WorkerStateMachine:
+    """State machine for managing worker computation lifecycle.
+
+    This class handles state transitions for worker computations,
+    ensuring valid state flow and preventing invalid operations.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the state machine in IDLE state."""
+        self._current_state = WorkerState.IDLE
+
+    @property
+    def current_state(self) -> WorkerState:
+        """Get the current state.
+
+        Returns:
+            Current WorkerState.
+        """
+        return self._current_state
+
+    def transition_to(self, target_state: WorkerState) -> None:
+        """Transition to the specified target state.
+
+        Args:
+            target_state: The state to transition to.
+
+        Raises:
+            ValueError: If the transition is not valid from the current state.
+        """
+        # Define valid state transitions
+        valid_transitions = {
+            WorkerState.IDLE: {WorkerState.STARTING},
+            WorkerState.STARTING: {WorkerState.RUNNING},
+            WorkerState.RUNNING: {WorkerState.FINISHED},
+            WorkerState.FINISHED: {WorkerState.IDLE},
+        }
+
+        # Allow transitions to the same state (no-op)
+        if target_state == self._current_state:
+            return
+
+        # Check if transition is valid
+        allowed_targets = valid_transitions.get(self._current_state, set())
+        if target_state not in allowed_targets:
+            raise ValueError(
+                f"Invalid transition from {self._current_state} to {target_state}. "
+                f"Valid transitions: {allowed_targets}"
+            )
+
+        self._current_state = target_state
+
+    def reset_to_idle(self) -> None:
+        """Reset state to IDLE unconditionally.
+
+        This is used for restart/cancel operations where we need
+        to force the state back to IDLE regardless of current state.
+        """
+        self._current_state = WorkerState.IDLE
+
+
 class Worker:
     """Multiprocessing worker, to run long-running tasks in a separate process"""
 
     def __init__(self) -> None:
         self.asyncresult: AsyncResult = None
-        self.result: Any = None
+        self.state_machine = WorkerStateMachine()
 
     @staticmethod
     def create_pool() -> None:
@@ -129,47 +199,128 @@ def terminate_pool(wait: bool = False) -> None:
     def restart_pool(self) -> None:
         """Terminate and recreate the pool"""
         # Terminate the process and stop the timer
-        self.terminate_pool(wait=False)
+        Worker.terminate_pool(wait=False)
         # Recreate the pool for the next computation
-        self.create_pool()
+        Worker.create_pool()
+        # Reset worker state after pool restart
+        self.asyncresult = None
+        self.state_machine.reset_to_idle()
 
     def run(self, func: Callable, args: tuple[Any]) -> None:
         """Run computation.
 
         Args:
             func: function to run
             args: arguments
+
+        Raises:
+            ValueError: If not in IDLE state or pool is not available.
         """
+        # Check if we can start computation
+        if self.state_machine.current_state != WorkerState.IDLE:
+            current_state = self.state_machine.current_state
+            raise ValueError(f"Cannot start computation from {current_state} state")
+
+        # Transition to starting state
+        self.state_machine.transition_to(WorkerState.STARTING)
+
         global POOL  # pylint: disable=global-statement,global-variable-not-assigned
-        assert POOL is not None
+        if POOL is None:
+            raise ValueError("Multiprocessing pool is not available")
+
+        # Start the computation
         env_json = sigima_options.get_env()
         self.asyncresult = POOL.apply_async(run_with_env, (func, args, env_json))
 
+        # Transition to running state
+        self.state_machine.transition_to(WorkerState.RUNNING)
+
+    def restart(self) -> None:
+        """Restart/cancel current computation"""
+        current_state = self.state_machine.current_state
+
+        if current_state == WorkerState.IDLE:
+            return  # Already idle, nothing to restart
+        elif current_state == WorkerState.STARTING:
+            # If we're still starting, just go back to idle
+            self.asyncresult = None
+        elif current_state == WorkerState.RUNNING:
+            # Cancel the running computation - use restart_pool for consistency
+            self.restart_pool()
+            return  # restart_pool already handles state reset
+        elif current_state == WorkerState.FINISHED:
+            # Clean up and go to idle
+            self.asyncresult = None
+
+        # Let state machine handle the transition to idle
+        self.state_machine.reset_to_idle()
+
     def close(self) -> None:
         """Close worker: close pool properly and wait for all tasks to finish"""
         # Close multiprocessing Pool properly, but only if no computation is running,
         # to avoid blocking the GUI at exit (so, when wait=True, we wait for the
         # task to finish before closing the pool but there is actually no task running,
         # so the pool is closed immediately but *properly*)
-        self.terminate_pool(wait=self.asyncresult is None)
+        Worker.terminate_pool(wait=self.asyncresult is None)
 
     def is_computation_finished(self) -> bool:
         """Return True if computation is finished.
 
         Returns:
             bool: True if computation is finished
         """
-        return self.asyncresult.ready()
+        current_state = self.state_machine.current_state
+
+        if current_state == WorkerState.IDLE:
+            return True  # No computation has been started
+        elif current_state == WorkerState.STARTING:
+            return False  # Computation is starting, not finished yet
+        elif current_state == WorkerState.FINISHED:
+            return True  # Already finished
+        elif current_state == WorkerState.RUNNING:
+            if self.asyncresult is None:
+                return False  # Should not happen, but defensive
+            finished = self.asyncresult.ready()
+            if finished:
+                # Transition to finished state
+                self.state_machine.transition_to(WorkerState.FINISHED)
+            return finished
+        else:
+            raise ValueError(f"Invalid worker state: {current_state}")
 
     def get_result(self) -> CompOut:
         """Return computation result.
 
         Returns:
             CompOut: computation result
+
+        Raises:
+            ValueError: If not in FINISHED state or no result available.
         """
-        self.result = self.asyncresult.get()
-        self.asyncresult = None
-        return self.result
+        # Check if we can get result
+        if self.state_machine.current_state != WorkerState.FINISHED:
+            current_state = self.state_machine.current_state
+            raise ValueError(f"Cannot get result from {current_state} state")
+
+        if self.asyncresult is None:
+            raise ValueError("No result available")
+
+        # Get result and clean up (ensure cleanup happens even if exception occurs)
+        try:
+            result = self.asyncresult.get()
+            return result
+        finally:
+            # Always clean up, even if get() raises an exception
+            self.asyncresult = None
+            self.state_machine.transition_to(WorkerState.IDLE)
+
+    def has_result_available(self) -> bool:
+        """Check if computation finished successfully and result is available.
+
+        Returns:
+            True if computation completed successfully and result can be retrieved.
+        """
+        return self.state_machine.current_state == WorkerState.FINISHED
 
 
 def is_pairwise_mode() -> bool:
@@ -268,7 +419,7 @@ def set_process_isolation_enabled(self, enabled: bool) -> None:
         if enabled:
             if self.worker is None:
                 self.worker = Worker()
-                self.worker.create_pool()
+                Worker.create_pool()
         else:
             if self.worker is not None:
                 self.worker.terminate_pool()
@@ -469,19 +620,14 @@ def __exec_func(
                 return wng_err_func(func, args)
             # Process isolation: run function in a separate process
             self.worker.run(func, args)
-            # Adaptive sleep to balance responsiveness and performance
-            sleep_duration = 0.001  # Start with 1ms
-            max_sleep = 0.01  # Cap at 10ms
-            sleep_increment = 0.001  # Increase by 1ms each iteration
             while not self.worker.is_computation_finished():
-                QW.QApplication.processEvents()  # Keep UI responsive
-                time.sleep(sleep_duration)  # Adaptive sleep to avoid busy waiting
-                # Gradually increase sleep time for longer computations
-                sleep_duration = min(sleep_duration + sleep_increment, max_sleep)
+                QW.QApplication.processEvents()
+                time.sleep(0)  # Just yields to other threads - no forced delay
                 if progress.wasCanceled():  # User canceled the operation
-                    self.worker.restart_pool()  # Terminate and recreate the pool
+                    self.worker.restart()  # Cancel computation and reset to idle
                     break
-            if self.worker.is_computation_finished():
+            # Only get result if computation actually finished (not canceled)
+            if self.worker.has_result_available():
                 return self.worker.get_result()
         return None