Additional system tests for supervisor

rtuck99 · rtuck99 · commit dc0b3e4b36bb · 2025-12-18T15:54:52.000Z
diff --git a/src/mx_bluesky/hyperion/plan_runner.py b/src/mx_bluesky/hyperion/plan_runner.py
@@ -40,8 +40,13 @@ class PlanError(Exception):
 
 
 class PlanRunner(BaseRunner):
+    EXTERNAL_CALLBACK_POLL_INTERVAL_S = 1
+    EXTERNAL_CALLBACK_WATCHDOG_TIMER_S = 60
+
     def __init__(self, context: BlueskyContext, dev_mode: bool):
         super().__init__(context)
+        self._callbacks_started = False
+        self._callback_watchdog_expiry = time.monotonic()
         self.is_dev_mode = dev_mode
 
     @abstractmethod
@@ -50,27 +55,63 @@ def decode_and_execute(
     ) -> MsgGenerator:
         pass
 
-    @abstractmethod
     def reset_callback_watchdog_timer(self):
-        pass
+        """Called periodically to reset the watchdog timer when the external callbacks ping us."""
+        self._callbacks_started = True
+        self._callback_watchdog_expiry = (
+            time.monotonic() + self.EXTERNAL_CALLBACK_WATCHDOG_TIMER_S
+        )
 
     @property
     @abstractmethod
     def current_status(self) -> Status:
         pass
 
+    def check_external_callbacks_are_alive(self):
+        callback_expiry = time.monotonic() + self.EXTERNAL_CALLBACK_WATCHDOG_TIMER_S
+        while time.monotonic() < callback_expiry:
+            if self._callbacks_started:
+                break
+            # If on first launch the external callbacks aren't started yet, wait until they are
+            LOGGER.info("Waiting for external callbacks to start")
+            yield from bps.sleep(self.EXTERNAL_CALLBACK_POLL_INTERVAL_S)
+        else:
+            raise RuntimeError("External callbacks not running - try restarting")
+
+        if not self._external_callbacks_are_alive():
+            raise RuntimeError(
+                "External callback watchdog timer expired, check external callbacks are running."
+            )
+
+    def request_run_engine_abort(self):
+        """Asynchronously request an abort from the run engine. This cannot be done from
+        inside the main thread."""
+
+        def issue_abort():
+            try:
+                # abort() causes the run engine to throw a RequestAbort exception
+                # inside the plan, which will propagate through the contingency wrappers.
+                # When the plan returns, the run engine will raise RunEngineInterrupted
+                self.run_engine.abort()
+            except Exception as e:
+                LOGGER.warning(
+                    "Exception encountered when issuing abort() to RunEngine:",
+                    exc_info=e,
+                )
+
+        stopping_thread = threading.Thread(target=issue_abort)
+        stopping_thread.start()
+
+    def _external_callbacks_are_alive(self) -> bool:
+        return time.monotonic() < self._callback_watchdog_expiry
+
 
 class InProcessRunner(PlanRunner):
     """Runner that executes experiments from inside a running Bluesky plan"""
 
-    EXTERNAL_CALLBACK_WATCHDOG_TIMER_S = 60
-    EXTERNAL_CALLBACK_POLL_INTERVAL_S = 1
-
     def __init__(self, context: BlueskyContext, dev_mode: bool) -> None:
         super().__init__(context, dev_mode)
         self._current_status: Status = Status.IDLE
-        self._callbacks_started = False
-        self._callback_watchdog_expiry = time.monotonic()
 
     def decode_and_execute(
         self, current_visit: str | None, parameter_list: Sequence[MxBlueskyParameters]
@@ -115,20 +156,7 @@ def execute_plan(
         self._current_status = Status.BUSY
 
         try:
-            callback_expiry = time.monotonic() + self.EXTERNAL_CALLBACK_WATCHDOG_TIMER_S
-            while time.monotonic() < callback_expiry:
-                if self._callbacks_started:
-                    break
-                # If on first launch the external callbacks aren't started yet, wait until they are
-                LOGGER.info("Waiting for external callbacks to start")
-                yield from bps.sleep(self.EXTERNAL_CALLBACK_POLL_INTERVAL_S)
-            else:
-                raise RuntimeError("External callbacks not running - try restarting")
-
-            if not self._external_callbacks_are_alive():
-                raise RuntimeError(
-                    "External callback watchdog timer expired, check external callbacks are running."
-                )
+            yield from self.check_external_callbacks_are_alive()
             yield from experiment()
             self._current_status = Status.IDLE
         except WarningError as e:
@@ -147,35 +175,12 @@ def shutdown(self):
         """Performs a prompt shutdown. Aborts the run engine and terminates the loop
         waiting for messages."""
 
-        def issue_abort():
-            try:
-                # abort() causes the run engine to throw a RequestAbort exception
-                # inside the plan, which will propagate through the contingency wrappers.
-                # When the plan returns, the run engine will raise RunEngineInterrupted
-                self.run_engine.abort()
-            except Exception as e:
-                LOGGER.warning(
-                    "Exception encountered when issuing abort() to RunEngine:",
-                    exc_info=e,
-                )
-
         LOGGER.info("Shutting down: Stopping the run engine gracefully")
         if self.current_status != Status.ABORTING:
             self._current_status = Status.ABORTING
-            stopping_thread = threading.Thread(target=issue_abort)
-            stopping_thread.start()
+            self.request_run_engine_abort()
             return
 
-    def reset_callback_watchdog_timer(self):
-        """Called periodically to reset the watchdog timer when the external callbacks ping us."""
-        self._callbacks_started = True
-        self._callback_watchdog_expiry = (
-            time.monotonic() + self.EXTERNAL_CALLBACK_WATCHDOG_TIMER_S
-        )
-
-    def _external_callbacks_are_alive(self) -> bool:
-        return time.monotonic() < self._callback_watchdog_expiry
-
     @property
     def current_status(self) -> Status:
         return self._current_status
diff --git a/src/mx_bluesky/hyperion/supervisor/_supervisor.py b/src/mx_bluesky/hyperion/supervisor/_supervisor.py
@@ -2,6 +2,7 @@
 from pathlib import Path
 
 from blueapi.client.client import BlueapiClient
+from blueapi.client.event_bus import BlueskyStreamingError
 from blueapi.config import ApplicationConfig, ConfigLoader
 from blueapi.core import BlueskyContext
 from blueapi.service.model import TaskRequest
@@ -13,7 +14,7 @@
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.hyperion.parameters.components import UDCCleanup, UDCDefaultState, Wait
 from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
-from mx_bluesky.hyperion.plan_runner import PlanRunner
+from mx_bluesky.hyperion.plan_runner import PlanError, PlanRunner
 
 
 def create_context() -> BlueskyContext:
@@ -36,49 +37,72 @@ def __init__(
     ):
         super().__init__(bluesky_context, dev_mode)
         self.blueapi_client = BlueapiClient.from_config(client_config)
+        self._current_status = Status.IDLE
 
     def decode_and_execute(
         self, current_visit: str | None, parameter_list: Sequence[MxBlueskyParameters]
     ) -> MsgGenerator:
+        try:
+            yield from self.check_external_callbacks_are_alive()
+        except Exception as e:
+            raise PlanError(f"Exception raised during plan execution: {e}") from e
         # TODO determine what is the instrument session for udc_default_state etc.
         instrument_session = current_visit or "NO_VISIT"
-        for parameters in parameter_list:
-            LOGGER.info(
-                f"Executing plan with parameters: {parameters.model_dump_json(indent=2)}"
-            )
-            match parameters:
-                case LoadCentreCollect():  # TODO
-                    pass
-                case Wait():
-                    yield from bps.sleep(parameters.duration_s)
-                case UDCDefaultState():
-                    task_request = TaskRequest(
-                        name="move_to_udc_default_state",
-                        params={},
-                        instrument_session=instrument_session,
-                    )
-                    self.blueapi_client.run_task(task_request)
-                    yield from bps.null()
-                case UDCCleanup():
-                    task_request = TaskRequest(
-                        name="clean_up_udc",
-                        params={"visit": current_visit},
-                        instrument_session=instrument_session,
-                    )
-                    self.blueapi_client.run_task(task_request)
-                    yield from bps.null()
-                case _:
-                    raise AssertionError(
-                        f"Unsupported instruction decoded from agamemnon {type(parameters)}"
-                    )
+        try:
+            if self._current_status == Status.ABORTING:
+                raise PlanError("Plan execution cancelled, supervisor is shutting down")
+            self._current_status = Status.BUSY
+            for parameters in parameter_list:
+                LOGGER.info(
+                    f"Executing plan with parameters: {parameters.model_dump_json(indent=2)}"
+                )
+                match parameters:
+                    case LoadCentreCollect():  # TODO
+                        pass
+                    case Wait():
+                        yield from bps.sleep(parameters.duration_s)
+                    case UDCDefaultState():
+                        task_request = TaskRequest(
+                            name="move_to_udc_default_state",
+                            params={},
+                            instrument_session=instrument_session,
+                        )
+                        self.blueapi_client.run_task(task_request)
+                        yield from bps.null()
+                    case UDCCleanup():
+                        task_request = TaskRequest(
+                            name="clean_up_udc",
+                            params={"visit": current_visit},
+                            instrument_session=instrument_session,
+                        )
+                        try:
+                            self.blueapi_client.run_task(task_request)
+                        except BlueskyStreamingError as e:
+                            # We will receive a BlueskyStreamingError if the remote server
+                            # processed an abort during plan execution, but this is not
+                            # the only possible cause.
+                            if self.current_status == Status.ABORTING:
+                                LOGGER.info("Aborting local runner...")
+                                self.request_run_engine_abort()
+                            else:
+                                raise PlanError(
+                                    f"Exception raised during plan execution: {e}"
+                                ) from e
+                        yield from bps.null()
+                    case _:
+                        raise AssertionError(
+                            f"Unsupported instruction decoded from agamemnon {type(parameters)}"
+                        )
+        except:
+            self._current_status = Status.FAILED
+            raise
+        else:
+            self._current_status = Status.IDLE
         return current_visit
 
-    def reset_callback_watchdog_timer(self):
-        pass
-
     @property
     def current_status(self) -> Status:
-        return Status.FAILED  # TODO
+        return self._current_status
 
     def is_connected(self) -> bool:
         try:
@@ -89,4 +113,11 @@ def is_connected(self) -> bool:
         return True
 
     def shutdown(self):
-        pass
+        LOGGER.info(
+            "Hyperion supervisor received shutdown request, signalling abort to BlueAPI server..."
+        )
+        if self.current_status != Status.BUSY:
+            self.request_run_engine_abort()
+        else:
+            self._current_status = Status.ABORTING
+            self.blueapi_client.abort()
diff --git a/tests/system_tests/hyperion/supervisor/dummy_plans.py b/tests/system_tests/hyperion/supervisor/dummy_plans.py
@@ -7,6 +7,8 @@
 from dodal.devices.robot import BartRobot
 from dodal.devices.smargon import Smargon
 
+from mx_bluesky.common.utils.exceptions import WarningError
+
 
 def publish_event(plan_name: str):
     yield from bps.open_run(md={"plan_name": plan_name})
@@ -23,3 +25,11 @@ def clean_up_udc(
     detector_motion: DetectorMotion = inject("detector_motion"),
 ) -> MsgGenerator:
     yield from publish_event("clean_up_udc")
+    match visit:
+        case "raise_warning_error":
+            raise WarningError("Test warning error")
+        case "raise_other_error":
+            raise RuntimeError("Test unexpected error")
+        case "wait_for_abort":
+            while True:
+                yield from bps.sleep(1)
diff --git a/tests/system_tests/hyperion/supervisor/test_supervisor.py b/tests/system_tests/hyperion/supervisor/test_supervisor.py
