Add retry/abort measurement dialog for environment evaluation failures

src/badger/core_subprocess.py

@@ -3,6 +3,7 @@
 import time
 import traceback
 from typing import Any
+from queue import Empty
 from pandas import DataFrame
 import multiprocessing as mp
 import os
@@ -22,6 +23,53 @@
 
 logger = logging.getLogger(__name__)
 
+MEASUREMENT_ERROR_TYPE = "measurement_error"
+MEASUREMENT_ACTION_TYPE = "measurement_action"
+MEASUREMENT_ACTION_RETRY = "retry"
+MEASUREMENT_ACTION_ABORT = "abort"
+
+
+def evaluate_measurement_with_retry(
+    routine: Routine,
+    point: Any,
+    queue: mp.Queue,
+    stop_process: mp.Event,
+) -> DataFrame:
+    while True:
+        try:
+            return routine.evaluate_data(point)
+        except Exception as e:
+            error_title = f"{type(e).__name__}: {e}"
+            error_traceback = traceback.format_exc()
+            logger.error(f"Measurement failed: {error_title}\n{error_traceback}")
+            queue.put(
+                {
+                    "type": MEASUREMENT_ERROR_TYPE,
+                    "title": error_title,
+                    "traceback": error_traceback,
+                }
+            )
+
+            while True:
+                if stop_process.is_set():
+                    raise BadgerRunTerminated
+                try:
+                    msg = queue.get(timeout=0.1)
+                except Empty:
+                    continue
+
+                if (
+                    isinstance(msg, dict)
+                    and msg.get("type") == MEASUREMENT_ACTION_TYPE
+                    and msg.get("action")
+                    in [MEASUREMENT_ACTION_RETRY, MEASUREMENT_ACTION_ABORT]
+                ):
+                    if msg["action"] == MEASUREMENT_ACTION_RETRY:
+                        break
+                    raise BadgerRunTerminated(
+                        f"Run terminated after measurement error: {error_title}"
+                    )
+
 
 def convert_to_solution(result: DataFrame, routine: Routine):
     """
@@ -217,7 +265,9 @@ def run_routine_subprocess(
             logger.info("Evaluating initial points...")
             for _, ele in initial_points.iterrows():
                 logger.debug(f"Evaluating initial point: {ele.to_dict()}")
-                result = routine.evaluate_data(ele.to_dict())
+                result = evaluate_measurement_with_retry(
+                    routine, ele.to_dict(), queue, stop_process
+                )
                 solution = convert_to_solution(result, routine)
                 opt_logger.update(Events.OPTIMIZATION_STEP, solution)
                 if evaluate:
@@ -280,7 +330,9 @@ def run_routine_subprocess(
                 )
                 pause_process.wait()
 
-            result = routine.evaluate_data(candidates)
+            result = evaluate_measurement_with_retry(
+                routine, candidates, queue, stop_process
+            )
             solution = convert_to_solution(result, routine)
             opt_logger.update(Events.OPTIMIZATION_STEP, solution)
 

src/badger/gui/components/routine_runner.py

@@ -4,17 +4,24 @@
 
 import pandas as pd
 from PyQt5.QtCore import pyqtSignal, QObject, QTimer
+from PyQt5.QtWidgets import QDialog
 
 from badger.errors import BadgerRunTerminated
 from badger.tests.utils import get_current_vars
 from badger.routine import calculate_variable_bounds, calculate_initial_points
 from badger.settings import init_settings
 from badger.gui.components.process_manager import ProcessManager
+from badger.gui.windows.measurement_retry_dialog import BadgerMeasurementRetryDialog
 from badger.routine import Routine
 from badger.errors import BadgerError
 
 logger = logging.getLogger(__name__)
 
+MEASUREMENT_ERROR_TYPE = "measurement_error"
+MEASUREMENT_ACTION_TYPE = "measurement_action"
+MEASUREMENT_ACTION_RETRY = "retry"
+MEASUREMENT_ACTION_ABORT = "abort"
+
 
 class BadgerRoutineSignals(QObject):
     env_ready = pyqtSignal(list)
@@ -234,15 +241,35 @@ def check_queue(self) -> None:
 
         if not self.data_and_error_queue.empty():
             try:
-                error_title, error_traceback = self.data_and_error_queue.get()
-                BadgerError(error_title, error_traceback)
+                msg = self.data_and_error_queue.get()
+                if isinstance(msg, dict) and msg.get("type") == MEASUREMENT_ERROR_TYPE:
+                    action = self.handle_measurement_error(msg)
+                    self.data_and_error_queue.put(
+                        {
+                            "type": MEASUREMENT_ACTION_TYPE,
+                            "action": action,
+                        }
+                    )
+                else:
+                    error_title, error_traceback = msg
+                    BadgerError(error_title, error_traceback)
             except ValueError:  # seems to only occur in tests
                 pass
 
         if not self.routine_process.is_alive():
             self.close()
             self.evaluate_queue[1].close()
 
+    def handle_measurement_error(self, msg: dict) -> str:
+        dialog = BadgerMeasurementRetryDialog(
+            text=msg.get("title", "Measurement failed."),
+            detailedText=msg.get("traceback", ""),
+        )
+        result = dialog.exec_()
+        if result == QDialog.Accepted:
+            return MEASUREMENT_ACTION_RETRY
+        return MEASUREMENT_ACTION_ABORT
+
     def after_evaluate(self, results: pd.DataFrame) -> None:
         logger.debug("Received evaluation results from subprocess.")
         """

src/badger/gui/windows/measurement_retry_dialog.py

@@ -0,0 +1,54 @@
+from PyQt5.QtCore import Qt
+from PyQt5.QtGui import QFont, QFontDatabase, QTextOption
+from PyQt5.QtWidgets import (
+    QDialog,
+    QHBoxLayout,
+    QLabel,
+    QPushButton,
+    QScrollArea,
+    QTextEdit,
+    QVBoxLayout,
+)
+
+
+class BadgerMeasurementRetryDialog(QDialog):
+    def __init__(self, text="", detailedText="", parent=None):
+        super().__init__(parent)
+        self.setWindowFlags(self.windowFlags() | Qt.WindowMaximizeButtonHint)
+
+        mainLayout = QVBoxLayout(self)
+
+        self.textLabel = QLabel(
+            text
+            + "\n\nThere was an error setting variables or getting observables."
+            + "\nRetry this measurement?"
+        )
+        self.textLabel.setMinimumWidth(320)
+        self.textLabel.setWordWrap(True)
+        font = QFont()
+        font.setBold(True)
+        self.textLabel.setFont(font)
+        mainLayout.addWidget(self.textLabel)
+
+        scrollArea = QScrollArea()
+        scrollArea.setWidgetResizable(True)
+        self.detailedTextWidget = QTextEdit(detailedText)
+        self.detailedTextWidget.setReadOnly(True)
+        self.detailedTextWidget.setWordWrapMode(QTextOption.NoWrap)
+        monoFont = QFontDatabase.systemFont(QFontDatabase.FixedFont)
+        monoFont.setPointSize(12)
+        self.detailedTextWidget.setFont(monoFont)
+        scrollArea.setWidget(self.detailedTextWidget)
+        mainLayout.addWidget(scrollArea)
+
+        buttonBox = QHBoxLayout()
+        self.retryButton = QPushButton("Retry Measurement")
+        self.stopButton = QPushButton("Stop Run")
+        self.retryButton.clicked.connect(self.accept)
+        self.stopButton.clicked.connect(self.reject)
+        buttonBox.addWidget(self.retryButton)
+        buttonBox.addWidget(self.stopButton)
+        mainLayout.addLayout(buttonBox)
+
+        self.setWindowTitle("Measurement Error")
+        self.resize(560, 360)

src/badger/tests/test_gui_basic.py

@@ -4,6 +4,7 @@
 
 import pytest
 from PyQt5.QtCore import QEventLoop, Qt, QTimer
+from PyQt5.QtWidgets import QDialog
 
 
 @pytest.fixture(scope="session")
@@ -130,6 +131,54 @@ def inner(ins, event):
         window.process_manager.close_proccesses()
 
 
+def test_measurement_retry_dialog_in_app_flow(qtbot, init_multiprocessing):
+    from badger.gui.windows.main_window import BadgerMainWindow
+    from badger.tests.utils import create_routine, fix_path_issues
+
+    class _FakeProcess:
+        def is_alive(self):
+            return True
+
+    fix_path_issues()
+    window = BadgerMainWindow()
+    qtbot.addWidget(window)
+
+    loop = QEventLoop()
+    QTimer.singleShot(1000, loop.quit)
+    loop.exec_()
+
+    routine = create_routine()
+    home_page = window.home_page
+    home_page.current_routine = routine
+    home_page.go_run(-1)
+    home_page.run_monitor.init_routine_runner()
+    runner = home_page.run_monitor.routine_runner
+    runner.data_and_error_queue = multiprocessing.Queue()
+    runner.evaluate_queue = multiprocessing.Pipe()
+    runner.routine_process = _FakeProcess()
+
+    with patch(
+        "badger.gui.windows.measurement_retry_dialog.BadgerMeasurementRetryDialog.exec_",
+        return_value=QDialog.Accepted,
+    ) as exec_mock:
+        runner.data_and_error_queue.put(
+            {
+                "type": "measurement_error",
+                "title": "Injected measurement error",
+                "traceback": "traceback details",
+            }
+        )
+        runner.check_queue()
+        exec_mock.assert_called_once()
+        response = runner.data_and_error_queue.get(timeout=1)
+        assert response == {
+            "type": "measurement_action",
+            "action": "retry",
+        }
+
+    window.process_manager.close_proccesses()
+
+
 # TODO: Check the use_low_noise_prior parameter in the routine
 # once it's running -- currently use_low_noise_prior is not exposed in the GUI
 # so need to check the routine object held by the monitor/runner

src/badger/tests/test_measurement_retry_dialog.py

@@ -0,0 +1,28 @@
+from PyQt5.QtCore import QTimer
+from PyQt5.QtWidgets import QDialog
+
+
+def test_measurement_retry_dialog_retry(qtbot):
+    from badger.gui.windows.measurement_retry_dialog import (
+        BadgerMeasurementRetryDialog,
+    )
+
+    dialog = BadgerMeasurementRetryDialog(text="Test error", detailedText="traceback")
+    qtbot.addWidget(dialog)
+    assert "Test error" in dialog.textLabel.text()
+    assert "traceback" in dialog.detailedTextWidget.toPlainText()
+    QTimer.singleShot(0, dialog.retryButton.click)
+    assert dialog.exec_() == QDialog.Accepted
+
+
+def test_measurement_retry_dialog_stop(qtbot):
+    from badger.gui.windows.measurement_retry_dialog import (
+        BadgerMeasurementRetryDialog,
+    )
+
+    dialog = BadgerMeasurementRetryDialog(text="Test error", detailedText="traceback")
+    qtbot.addWidget(dialog)
+    assert "Test error" in dialog.textLabel.text()
+    assert "traceback" in dialog.detailedTextWidget.toPlainText()
+    QTimer.singleShot(0, dialog.stopButton.click)
+    assert dialog.exec_() == QDialog.Rejected