Enhance mini GUI run controls

src/badger/core_subprocess.py

@@ -31,13 +31,18 @@
     MEASUREMENT_ACTION_TYPE,
     MEASUREMENT_ACTION_RETRY,
     MEASUREMENT_ACTION_ABORT,
+    TERMINATION_REACHED_TYPE,
+    TERMINATION_ACTION_TYPE,
+    TERMINATION_ACTION_CONTINUE,
+    TERMINATION_ACTION_END,
 )
 from badger.logger import _get_default_logger
 from badger.logger.event import Events
 from badger.routine import Routine
 from badger.log import configure_process_logging
 from xopt.errors import FeasibilityError, XoptError
 from xopt.vocs import select_best
+from xopt.generators.sequential import SequentialGenerator
 
 
 logger = logging.getLogger(__name__)
@@ -89,6 +94,119 @@ def evaluate_measurement_with_retry(
                     )
 
 
+def pause_for_termination_dialog_action(
+    queue: mp.Queue,
+    stop_process: mp.Event,
+    pause_process: mp.Event,
+    dialog_action_queue: mp.Queue,
+    tc_condition: dict,
+) -> None:
+    """Pause the run and wait for user action when run-until condition is reached."""
+    queue.put(
+        {
+            "type": TERMINATION_REACHED_TYPE,
+            "tc_condition": tc_condition,
+        }
+    )
+
+    while True:
+        if stop_process.is_set():
+            raise BadgerRunTerminated
+
+        try:
+            msg = dialog_action_queue.get(
+                timeout=0.1
+            )  # short timeout here, so we can make checks for stop_process
+        except Empty:
+            continue
+
+        if (
+            isinstance(msg, dict)
+            and msg.get("type") == TERMINATION_ACTION_TYPE
+            and msg.get("action")
+            in [TERMINATION_ACTION_CONTINUE, TERMINATION_ACTION_END]
+        ):
+            if msg["action"] == TERMINATION_ACTION_CONTINUE:
+                pause_process.set()
+                return
+
+            raise BadgerRunTerminated(
+                "Run terminated after termination condition reached"
+            )
+
+
+def check_termination_condition(
+    termination_condition: dict,
+    start_time: float,
+    routine: Routine,
+    queue: mp.Queue,
+    stop_process: mp.Event,
+    pause_process: mp.Event,
+    dialog_action_queue: mp.Queue,
+) -> bool:
+    """
+    Check whether termination conditon has been reached.
+    Pause for user action when a configured termination condition is reached.
+    """
+    if not termination_condition or not start_time:
+        return False
+
+    tc_config = termination_condition
+    idx = tc_config["tc_idx"]
+    if idx == 0:
+        max_eval = tc_config["max_eval"]
+        if routine.data is not None:
+            if "live" in routine.data.columns:
+                # Only count number of live data points
+                count = sum(1 for live_val in routine.data["live"] if live_val == 1)
+            else:
+                count = len(routine.data)
+            logger.debug(f"Checking max_eval termination: {count} >= {max_eval}")
+        else:
+            count = 0
+
+        if count >= max_eval:
+            logger.info(
+                "Max evaluations reached. Pausing optimization and waiting for user action."
+            )
+            pause_process.clear()
+            pause_for_termination_dialog_action(
+                queue=queue,
+                stop_process=stop_process,
+                pause_process=pause_process,
+                dialog_action_queue=dialog_action_queue,
+                tc_condition={
+                    "type": "max_eval",
+                    "config": max_eval,
+                    "state": count,
+                },
+            )
+            return True
+    elif idx == 1:
+        max_time = tc_config["max_time"]
+        dt = time.time() - start_time
+        logger.debug(f"Checking max_time termination: {dt} >= {max_time}")
+        if dt >= max_time:
+            logger.info(
+                "Max time reached. Pausing optimization and waiting for user action."
+            )
+            pause_process.clear()
+            pause_for_termination_dialog_action(
+                queue=queue,
+                stop_process=stop_process,
+                pause_process=pause_process,
+                dialog_action_queue=dialog_action_queue,
+                tc_condition={
+                    "type": "max_time",
+                    "config": max_time,
+                    "state": dt,
+                },
+            )
+            return True
+
+    return False
+
+
 def convert_to_solution(result: DataFrame, routine: Routine):
     """
     This method is passed the latest evaluated solution and converts that to a printable format for the terminal.
@@ -141,6 +259,7 @@ def convert_to_solution(result: DataFrame, routine: Routine):
 
 
 def run_routine_subprocess(
+    args_queue: mp.Queue,
     queue: mp.Queue,
     evaluate_queue: mp.Pipe,
     stop_process: mp.Event,
@@ -192,7 +311,7 @@ def run_routine_subprocess(
 
     args: dict[str, Any] = {}
     try:
-        args = queue.get(timeout=1)
+        args = args_queue.get(timeout=1)
         logger.debug(f"Received args from queue: {args}")
     except Exception as e:
         logger.error(f"Error in subprocess queue.get: {type(e).__name__}, {str(e)}")
@@ -214,6 +333,9 @@ def run_routine_subprocess(
             if routine.data is not None:
                 logger.info("Resetting routine data")
                 routine.data = routine.data.iloc[0:0]  # reset the data
+        else:
+            if isinstance(routine.generator, SequentialGenerator):
+                routine.generator.add_data(routine.data.copy())
 
     except Exception as e:
         error_title = f"{type(e).__name__}: {e}"
@@ -302,37 +424,26 @@ def run_routine_subprocess(
                 logger.info("Pause process not set. Waiting...")
                 pause_process.wait()
 
-            if termination_condition and start_time:
-                tc_config = termination_condition
-                idx = tc_config["tc_idx"]
-                if idx == 0:
-                    max_eval = tc_config["max_eval"]
-                    if routine.data is not None:
-                        if "live" in routine.data.columns:
-                            # Only count number of live data points
-                            count = sum(
-                                1 for live_val in routine.data["live"] if live_val == 1
-                            )
-                        else:
-                            count = len(routine.data)
-                        logger.debug(
-                            f"Checking max_eval termination: {count} >= {max_eval}"
-                        )
-                    else:
-                        count = 0
-
-                    if count >= max_eval:
-                        logger.info(
-                            "Max evaluations reached. Terminating optimization."
-                        )
-                        raise BadgerRunTerminated
-                elif idx == 1:
-                    max_time = tc_config["max_time"]
-                    dt = time.time() - start_time
-                    logger.debug(f"Checking max_time termination: {dt} >= {max_time}")
-                    if dt >= max_time:
-                        logger.info("Max time reached. Terminating optimization.")
-                        raise BadgerRunTerminated
+            if check_termination_condition(
+                termination_condition,
+                start_time,
+                routine,
+                queue,
+                stop_process,
+                pause_process,
+                dialog_action_queue,
+            ):
+                # termination_condition = extend_termination_condition(
+                #    termination_condition, original_termination_condition
+                # )
+                termination_condition = None
+                queue.put(
+                    {
+                        "type": "termination_extended",
+                        "termination_condition": termination_condition,
+                    }
+                )
+                continue
 
             candidates = routine.generator.generate(1)[0]
             logger.debug(f"Generated candidates: {candidates}")

src/badger/errors.py

@@ -121,3 +121,10 @@ def __init__(self, message="Optimization run has been terminated!"):
 MEASUREMENT_ACTION_TYPE = "measurement_action"
 MEASUREMENT_ACTION_RETRY = "retry"
 MEASUREMENT_ACTION_ABORT = "abort"
+
+# Constants for run-until termination dialog feature.
+# Used in communication between routine runner and subprocess.
+TERMINATION_REACHED_TYPE = "termination_reached"
+TERMINATION_ACTION_TYPE = "termination_action"
+TERMINATION_ACTION_CONTINUE = "continue"
+TERMINATION_ACTION_END = "end"