Add robust error handling to plotter and heatmap threads

Introduces try/except blocks and disables plotting on exceptions in Plotter and Heatmap threads, logging errors and preventing further UI updates after failures. Refactors runner thread to encapsulate step logic and improve exception handling, ensuring proper cleanup and error reporting. Updates tests to skip parameter validation for mock parameters.

Author: caidish

src/measureit/_internal/plotter_thread.py

@@ -9,6 +9,8 @@
 from PyQt5.QtGui import QFont
 from PyQt5.QtWidgets import QHBoxLayout, QLabel, QProgressBar, QVBoxLayout, QWidget
 
+from ..logging_utils import get_sweep_logger
+
 # Configure PyQtGraph for better performance
 pg.setConfigOptions(
     antialias=False,  # Disable antialiasing for better performance
@@ -82,6 +84,7 @@ def __init__(self, sweep, plot_bin=1):
         self.last_pass = False
         self.figs_set = False
         self.kill_flag = False
+        self.logger = getattr(sweep, "logger", None) or get_sweep_logger("plotter")
 
         # PyQtGraph-specific attributes
         self.widget = None
@@ -284,7 +287,11 @@ def create_figs(self):
         self.widget.closeEvent = self.handle_close
 
         # Show the widget
-        self.widget.show()
+        try:
+            self.widget.show()
+        except Exception as e:
+            self._disable_plotting(f"Plotter widget failed to show: {e}")
+            return
 
     @pyqtSlot(object, int)
     def add_data(self, data, direction):
@@ -297,11 +304,16 @@ def add_data(self, data, direction):
         direction:
             The direction of the sweep (0 or 1).
         """
-        self.data_queue.append((data, direction))
+        if self.kill_flag:
+            return
+        try:
+            self.data_queue.append((data, direction))
 
-        # Only update plots when we have enough data points or it's been a while
-        # This prevents excessive update calls that slow down the plotting
-        self.update_plots(force=data is None)
+            # Only update plots when we have enough data points or it's been a while
+            # This prevents excessive update calls that slow down the plotting
+            self.update_plots(force=data is None)
+        except Exception as e:
+            self._disable_plotting(f"Plotter add_data failed: {e}")
 
     @pyqtSlot(bool)
     def update_plots(self, force=False):
@@ -313,79 +325,83 @@ def update_plots(self, force=False):
         force:
             If True, process all queued data immediately.
         """
-        if not self.figs_set:
-            return
-
-        # Check if we should update - either we have enough data or force is True
-        if len(self.data_queue) < self.plot_bin and not force:
+        if not self.figs_set or self.kill_flag:
             return
-
-        # Process all queued data at once for better performance
-        while len(self.data_queue) > 0:
-            temp = self.data_queue.popleft()
-            if temp[0] is None:
-                break
-            data = deque(temp[0])
-            direction = temp[1]
-
-            # Get time data
-            time_data = data.popleft()
-
-            # Handle set parameter plot
-            set_param_data = None
-            if self.sweep.set_param is not None:
-                set_param_data = data.popleft()
-
-                # Add to set parameter data arrays
-                # Ensure scalars by flattening any arrays
-                time_val = time_data[1]
-                if hasattr(time_val, "flatten"):
-                    time_val = float(np.array(time_val).flatten()[0])
-
-                set_val = set_param_data[1]
-                if hasattr(set_val, "flatten"):
-                    set_val = float(np.array(set_val).flatten()[0])
-
-                self.set_data_arrays["x"].append(time_val)
-                self.set_data_arrays["y"].append(set_val)
-
-            # Determine x-axis data for followed parameters
-            x_data_value = (
-                time_data[1]
-                if self.sweep.x_axis == 1
-                else (
-                    set_param_data[1]
-                    if self.sweep.set_param is not None
-                    else time_data[1]
+        try:
+            # Check if we should update - either we have enough data or force is True
+            if len(self.data_queue) < self.plot_bin and not force:
+                return
+
+            # Process all queued data at once for better performance
+            while len(self.data_queue) > 0:
+                temp = self.data_queue.popleft()
+                if temp[0] is None:
+                    break
+                data = deque(temp[0])
+                direction = temp[1]
+
+                # Get time data
+                time_data = data.popleft()
+
+                # Handle set parameter plot
+                set_param_data = None
+                if self.sweep.set_param is not None:
+                    set_param_data = data.popleft()
+
+                    # Add to set parameter data arrays
+                    # Ensure scalars by flattening any arrays
+                    time_val = time_data[1]
+                    if hasattr(time_val, "flatten"):
+                        time_val = float(np.array(time_val).flatten()[0])
+
+                    set_val = set_param_data[1]
+                    if hasattr(set_val, "flatten"):
+                        set_val = float(np.array(set_val).flatten()[0])
+
+                    self.set_data_arrays["x"].append(time_val)
+                    self.set_data_arrays["y"].append(set_val)
+
+                # Determine x-axis data for followed parameters
+                x_data_value = (
+                    time_data[1]
+                    if self.sweep.x_axis == 1
+                    else (
+                        set_param_data[1]
+                        if self.sweep.set_param is not None
+                        else time_data[1]
+                    )
                 )
-            )
-            # Ensure x_data_value is scalar
-            if hasattr(x_data_value, "flatten"):
-                x_data_value = float(np.array(x_data_value).flatten()[0])
+                # Ensure x_data_value is scalar
+                if hasattr(x_data_value, "flatten"):
+                    x_data_value = float(np.array(x_data_value).flatten()[0])
 
-            # Add data to arrays for followed parameters
-            for i, data_pair in enumerate(data):
-                param = self.sweep._params[i]
+                # Add data to arrays for followed parameters
+                for i, data_pair in enumerate(data):
+                    param = self.sweep._params[i]
 
-                if param in self.data_arrays:
-                    direction_key = "forward" if direction == 0 else "backward"
+                    if param in self.data_arrays:
+                        direction_key = "forward" if direction == 0 else "backward"
 
-                    # Ensure y_data is scalar
-                    y_value = data_pair[1]
-                    if hasattr(y_value, "flatten"):
-                        y_value = float(np.array(y_value).flatten()[0])
+                        # Ensure y_data is scalar
+                        y_value = data_pair[1]
+                        if hasattr(y_value, "flatten"):
+                            y_value = float(np.array(y_value).flatten()[0])
 
-                    self.data_arrays[param][direction_key]["x"].append(x_data_value)
-                    self.data_arrays[param][direction_key]["y"].append(y_value)
+                        self.data_arrays[param][direction_key]["x"].append(x_data_value)
+                        self.data_arrays[param][direction_key]["y"].append(y_value)
 
-        # Now update all plots at once (much more efficient)
-        self._update_plot_displays()
-        self.update_progress_widgets()
+            # Now update all plots at once (much more efficient)
+            self._update_plot_displays()
+            self.update_progress_widgets()
+        except Exception as e:
+            self._disable_plotting(f"Plotter update_plots failed: {e}")
 
     def _update_plot_displays(self):
         """Efficiently updates all plot displays using stored data arrays.
         Optimized for large datasets - converts to numpy arrays efficiently.
         """
+        if self.kill_flag:
+            return
         # Update set parameter plot
         if self.sweep.set_param is not None and self.set_plot_item is not None:
             if self.set_data_arrays["x"] and self.set_data_arrays["y"]:
@@ -497,6 +513,23 @@ def get_plot_data(self, param_index):
                 }
         return None
 
+    def _disable_plotting(self, reason: str):
+        """Disable plotting after an exception and log the error."""
+        try:
+            self.logger.error(reason, exc_info=True)
+        except Exception:
+            pass
+        self.kill_flag = True
+        try:
+            self.sweep.plot_data = False
+        except Exception:
+            pass
+        try:
+            if self.figs_set:
+                self.clear()
+        except Exception:
+            pass
+
     @pyqtSlot()
     def run(self):
         """Creates figures if they have not already been set."""

src/measureit/_internal/runner_thread.py

@@ -120,6 +120,32 @@ def add_plotter(self, plotter):
         self.plotter = plotter
         self.send_data.connect(self.plotter.add_data)
 
+    def _run_step(self) -> bool:
+        """Run a single sweep step; return False when the loop should exit."""
+        t = time.monotonic()
+        state = getattr(self.sweep.progressState, "state", None)
+
+        if state == SweepState.RUNNING:
+            data = self.sweep.update_values()
+            self.sweep.update_progress()
+
+            if self.plotter is not None and self.sweep.plot_data is True:
+                self.send_data.emit(data, self.sweep.direction)
+
+        # Smart sleep: compensate for time spent executing update
+        sleep_time = self.sweep.inter_delay - (time.monotonic() - t)
+        if sleep_time > 0:
+            time.sleep(sleep_time)
+
+        # Refresh state after possible updates
+        state = getattr(self.sweep.progressState, "state", None)
+        if state in (SweepState.DONE, SweepState.KILLED, SweepState.ERROR):
+            if self.sweep.save_data is True and self.datasaver is not None:
+                self.datasaver.flush_data_to_database()
+            return False
+
+        return True
+
     def _set_parent(self, sweep):
         """Sets a parent sweep if the Runner Thread is created independently.
 
@@ -197,49 +223,42 @@ def run(self):
                 return  # Exit run() early - no point entering the main loop
 
         # print(f"called runner from thread: {QThread.currentThreadId()}")
-        while True:
-            t = time.monotonic()
-            state = getattr(self.sweep.progressState, "state", None)
-
-            data = None
-            if state == SweepState.RUNNING:
+        try:
+            while True:
                 try:
-                    data = self.sweep.update_values()
+                    should_continue = self._run_step()
                 except ParameterException as e:
                     # safe_set already retried once and gave up - immediately transition to ERROR
                     # This prevents the sweep from continuing to the next setpoint with bad data
                     # Don't emit completed signal here - defer it until after loop exits
                     # to avoid blocking the main event loop
                     self.sweep.mark_error(
                         f"Parameter operation failed: {e}",
-                        _from_runner=True
+                        _from_runner=True,
                     )
                     break  # Exit loop immediately to avoid race conditions
-                self.sweep.update_progress()
-
-                if (
-                    self.plotter is not None
-                    and self.sweep.plot_data is True
-                ):
-                    self.send_data.emit(data, self.sweep.direction)
-
-            # Smart sleep: compensate for time spent executing update
-            sleep_time = self.sweep.inter_delay - (time.monotonic() - t)
-
-            if sleep_time > 0:
-                time.sleep(sleep_time)
-
-            # Refresh state after possible updates
-            state = getattr(self.sweep.progressState, "state", None)
-            if state in (SweepState.DONE, SweepState.KILLED, SweepState.ERROR):
-                if self.sweep.save_data is True and self.datasaver is not None:
-                    self.datasaver.flush_data_to_database()
-                if state in (SweepState.KILLED, SweepState.ERROR):
+                except (KeyboardInterrupt, SystemExit):
+                    raise
+                except Exception as e:
+                    # Catch-all: ensure unexpected exceptions mark the sweep as ERROR
+                    logger.error(
+                        "Unhandled exception in runner thread: %s\n%s",
+                        e,
+                        traceback.format_exc(),
+                    )
+                    try:
+                        self.sweep.mark_error(
+                            f"Unhandled runner error: {e}",
+                            _from_runner=True,
+                        )
+                    except Exception:
+                        pass
                     break
-                # Allow DONE sweeps to exit after flushing
-                break
 
-        self.exit_datasaver()
+                if not should_continue:
+                    break
+        finally:
+            self.exit_datasaver()
 
         # Emit completed signal for ERROR state after loop exits
         # This is deferred to avoid blocking the main event loop during exception handling
@@ -249,8 +268,14 @@ def run(self):
 
     def exit_datasaver(self):
         if self.datasaver is not None:
-            if self.sweep.save_data is True:
-                self.datasaver.flush_data_to_database()
-
-            self.runner.__exit__(None, None, None)
+            try:
+                if self.sweep.save_data is True:
+                    self.datasaver.flush_data_to_database()
+            except Exception:
+                logger.error("Failed to flush datasaver:\n%s", traceback.format_exc())
+            try:
+                if self.runner is not None:
+                    self.runner.__exit__(None, None, None)
+            except Exception:
+                logger.error("Failed to close datasaver:\n%s", traceback.format_exc())
             self.datasaver = None