mx-bluesky 1.5.1__py3-none-any.whl → 1.5.3__py3-none-any.whl

mx_bluesky/hyperion/external_interaction/config_server.py

@@ -1,35 +1,16 @@
 from functools import cache
 
-from daq_config_server.client import ConfigServer
+from mx_bluesky.common.external_interaction.config_server import MXConfigClient
+from mx_bluesky.hyperion.parameters.constants import (
+    HyperionFeatureSetting,
+    HyperionFeatureSettingources,
+)
 
-from mx_bluesky.common.external_interaction.config_server import FeatureFlags
-from mx_bluesky.common.utils.log import LOGGER
-from mx_bluesky.hyperion.parameters.constants import CONST
 
-
-class HyperionFeatureFlags(FeatureFlags):
-    """
-    Feature flags specific to Hyperion.
-
-    Attributes:
-        use_panda_for_gridscan:         If True then the PandA is used for gridscans, otherwise the zebra is used
-        use_gpu_results:                If True then GPU result processing is enabled
-            and the GPU result is taken.
-        set_stub_offsets:               If True then set the stub offsets after moving to the crystal (ignored for
-            multi-centre)
-        omega_flip:                     If True then invert the smargon omega motor rotation commands with respect to
-         the hyperion request. See "Hyperion Coordinate Systems" in the documentation.
-        alternate_rotation_direction:   If True then the for multi-sample pins the rotation direction of
-         successive rotation scans is alternated between positive and negative.
-    """
-
-    @staticmethod
-    @cache
-    def get_config_server() -> ConfigServer:
-        return ConfigServer(CONST.CONFIG_SERVER_URL, LOGGER)
-
-    use_panda_for_gridscan: bool = CONST.I03.USE_PANDA_FOR_GRIDSCAN
-    use_gpu_results: bool = CONST.I03.USE_GPU_RESULTS
-    set_stub_offsets: bool = CONST.I03.SET_STUB_OFFSETS
-    omega_flip: bool = CONST.I03.OMEGA_FLIP
-    alternate_rotation_direction: bool = CONST.I03.ALTERNATE_ROTATION_DIRECTION
+@cache
+def get_hyperion_config_client() -> MXConfigClient[HyperionFeatureSetting]:
+    return MXConfigClient(
+        feature_sources=HyperionFeatureSettingources,
+        feature_dc=HyperionFeatureSetting,
+        url="https://daq-config.diamond.ac.uk",
+    )

mx_bluesky/hyperion/parameters/cli.py

@@ -1,12 +1,19 @@
 import argparse
+from enum import StrEnum
 
 from pydantic.dataclasses import dataclass
 
 from mx_bluesky._version import version
 
 
+class HyperionMode(StrEnum):
+    GDA = "gda"
+    UDC = "udc"
+
+
 @dataclass
 class HyperionArgs:
+    mode: HyperionMode
     dev_mode: bool = False
 
 
@@ -39,7 +46,12 @@ def parse_cli_args() -> HyperionArgs:
         action="version",
         version=version,
     )
-    args = parser.parse_args()
-    return HyperionArgs(
-        dev_mode=args.dev or False,
+    parser.add_argument(
+        "--mode",
+        help="Launch in the specified mode (default is 'gda')",
+        default=HyperionMode.GDA,
+        type=HyperionMode,
+        choices=HyperionMode.__members__.values(),
     )
+    args = parser.parse_args()
+    return HyperionArgs(dev_mode=args.dev or False, mode=args.mode)

mx_bluesky/hyperion/parameters/components.py

@@ -1,8 +1,10 @@
-from pydantic import Field
+from mx_bluesky.common.parameters.components import MxBlueskyParameters
 
-from mx_bluesky.common.parameters.components import WithPandaGridScan
-from mx_bluesky.hyperion.external_interaction.config_server import HyperionFeatureFlags
 
+class Wait(MxBlueskyParameters):
+    """Represents an instruction from Agamemnon for Hyperion to wait for a specified time
+    Attributes:
+        duration_s: duration to wait in seconds
+    """
 
-class WithHyperionUDCFeatures(WithPandaGridScan):
-    features: HyperionFeatureFlags = Field(default=HyperionFeatureFlags())
+    duration_s: float

mx_bluesky/hyperion/parameters/constants.py

@@ -8,6 +8,8 @@ from mx_bluesky.common.parameters.constants import (
     DocDescriptorNames,
     EnvironmentConstants,
     ExperimentParamConstants,
+    FeatureSetting,
+    FeatureSettingources,
     HardwareConstants,
     OavConstants,
     PlanGroupCheckpointConstants,
@@ -24,18 +26,30 @@ class I03Constants:
     INSERTION_PREFIX = "SR03S" if TEST_MODE else "SR03I"
     OAV_CENTRING_FILE = OavConstants.OAV_CONFIG_JSON
     SHUTTER_TIME_S = 0.06
-    USE_PANDA_FOR_GRIDSCAN = False
-    SET_STUB_OFFSETS = False
+    USE_GPU_RESULTS = True
     OMEGA_FLIP = True
     ALTERNATE_ROTATION_DIRECTION = True
 
-    # Turns on GPU processing for zocalo and uses the results that come back
-    USE_GPU_RESULTS = True
+
+# These currently exist in GDA domain.properties
+class HyperionFeatureSettingources(FeatureSettingources):
+    USE_GPU_RESULTS = "gda.mx.hyperion.xrc.use_gpu_results"
+    USE_PANDA_FOR_GRIDSCAN = "gda.mx.hyperion.use_panda_for_gridscans"
+    SET_STUB_OFFSETS = "gda.mx.hyperion.do_stub_offsets"
+    PANDA_RUNUP_DISTANCE_MM = "gda.mx.hyperion.panda_runup_distance_mm"
+
+
+# Use these defaults if we can't read from the config server
+@dataclass
+class HyperionFeatureSetting(FeatureSetting):
+    USE_GPU_RESULTS: bool = True
+    USE_PANDA_FOR_GRIDSCAN: bool = False
+    SET_STUB_OFFSETS: bool = False
+    PANDA_RUNUP_DISTANCE_MM: float = 0.16
 
 
 @dataclass(frozen=True)
 class HyperionConstants:
-    DESCRIPTORS = DocDescriptorNames()
     ZOCALO_ENV = EnvironmentConstants.ZOCALO_ENV
     HARDWARE = HardwareConstants()
     I03 = I03Constants()
@@ -49,7 +63,8 @@ class HyperionConstants:
         if TEST_MODE
         else "https://daq-config.diamond.ac.uk/api"
     )
-    GRAYLOG_PORT = 12232
+    GRAYLOG_PORT = 12232  # Hyperion stream
+    GRAYLOG_STREAM_ID = "66264f5519ccca6d1c9e4e03"
     PARAMETER_SCHEMA_DIRECTORY = "src/hyperion/parameters/schemas/"
     LOG_FILE_NAME = "hyperion.log"
     DEVICE_SETTINGS_CONSTANTS = DeviceSettingsConstants()

mx_bluesky/hyperion/parameters/gridscan.py

@@ -9,10 +9,12 @@ from mx_bluesky.common.parameters.gridscan import (
     GridCommon,
     SpecifiedThreeDGridScan,
 )
-from mx_bluesky.hyperion.parameters.components import WithHyperionUDCFeatures
+from mx_bluesky.hyperion.external_interaction.config_server import (
+    get_hyperion_config_client,
+)
 
 
-class GridCommonWithHyperionDetectorParams(GridCommon, WithHyperionUDCFeatures):
+class GridCommonWithHyperionDetectorParams(GridCommon):
     """Used by models which require detector parameters but have no specifications of the grid"""
 
     # These detector params only exist so that we can properly select enable_dev_shm. Remove in
@@ -20,11 +22,13 @@ class GridCommonWithHyperionDetectorParams(GridCommon, WithHyperionUDCFeatures):
     @property
     def detector_params(self):
         params = super().detector_params
-        params.enable_dev_shm = self.features.use_gpu_results
+        params.enable_dev_shm = (
+            get_hyperion_config_client().get_feature_flags().USE_GPU_RESULTS
+        )
         return params
 
 
-class HyperionSpecifiedThreeDGridScan(WithHyperionUDCFeatures, SpecifiedThreeDGridScan):
+class HyperionSpecifiedThreeDGridScan(SpecifiedThreeDGridScan):
     """Hyperion's 3D grid scan deviates from the common class due to: optionally using a PandA, optionally using dev_shm for GPU analysis, and using a config server for features"""
 
     # These detector params only exist so that we can properly select enable_dev_shm. Remove in
@@ -33,7 +37,9 @@ class HyperionSpecifiedThreeDGridScan(WithHyperionUDCFeatures, SpecifiedThreeDGr
     @property
     def detector_params(self):
         params = super().detector_params
-        params.enable_dev_shm = self.features.use_gpu_results
+        params.enable_dev_shm = (
+            get_hyperion_config_client().get_feature_flags().USE_GPU_RESULTS
+        )
         return params
 
     # Relative to common grid scan, stub offsets are defined by config server
@@ -51,7 +57,9 @@ class HyperionSpecifiedThreeDGridScan(WithHyperionUDCFeatures, SpecifiedThreeDGr
             z1_start_mm=self.z_start_um / 1000,
             y2_start_mm=self.y2_start_um / 1000,
             z2_start_mm=self.z2_start_um / 1000,
-            set_stub_offsets=self.features.set_stub_offsets,
+            set_stub_offsets=get_hyperion_config_client()
+            .get_feature_flags()
+            .SET_STUB_OFFSETS,
             dwell_time_ms=self.exposure_time_s * 1000,
             transmission_fraction=self.transmission_frac,
         )
@@ -75,8 +83,12 @@ class HyperionSpecifiedThreeDGridScan(WithHyperionUDCFeatures, SpecifiedThreeDGr
             z1_start_mm=self.z_start_um / 1000,
             y2_start_mm=self.y2_start_um / 1000,
             z2_start_mm=self.z2_start_um / 1000,
-            set_stub_offsets=self.features.set_stub_offsets,
-            run_up_distance_mm=self.panda_runup_distance_mm,
+            set_stub_offsets=get_hyperion_config_client()
+            .get_feature_flags()
+            .SET_STUB_OFFSETS,
+            run_up_distance_mm=get_hyperion_config_client()
+            .get_feature_flags()
+            .PANDA_RUNUP_DISTANCE_MM,
             transmission_fraction=self.transmission_frac,
         )
 
@@ -84,13 +96,9 @@ class HyperionSpecifiedThreeDGridScan(WithHyperionUDCFeatures, SpecifiedThreeDGr
 class OddYStepsException(Exception): ...
 
 
-class PinTipCentreThenXrayCentre(
-    GridCommonWithHyperionDetectorParams, WithHyperionUDCFeatures
-):
+class PinTipCentreThenXrayCentre(GridCommonWithHyperionDetectorParams):
     tip_offset_um: float = 0
 
 
-class GridScanWithEdgeDetect(
-    GridCommonWithHyperionDetectorParams, WithHyperionUDCFeatures
-):
+class GridScanWithEdgeDetect(GridCommonWithHyperionDetectorParams):
     pass

mx_bluesky/hyperion/parameters/load_centre_collect.py

@@ -8,7 +8,6 @@ from mx_bluesky.common.parameters.components import (
     WithSample,
     WithVisit,
 )
-from mx_bluesky.hyperion.parameters.components import WithHyperionUDCFeatures
 from mx_bluesky.hyperion.parameters.robot_load import (
     RobotLoadThenCentre,
 )
@@ -28,7 +27,6 @@ class LoadCentreCollect(
     WithVisit,
     WithSample,
     WithCentreSelection,
-    WithHyperionUDCFeatures,
 ):
     """Experiment parameters to perform the combined robot load,
     pin-tip centre and rotation scan operations."""
@@ -39,6 +37,7 @@ class LoadCentreCollect(
     @model_validator(mode="before")
     @classmethod
     def validate_model(cls, values):
+        values = values.copy()
         allowed_keys = (
             LoadCentreCollect.model_fields.keys()
             | RobotLoadThenCentre.model_fields.keys()
@@ -50,12 +49,6 @@ class LoadCentreCollect(
             f"Unexpected fields found in LoadCentreCollect {disallowed_keys}"
         )
 
-        assert "features" not in values["robot_load_then_centre"], (
-            "Features flags must be specified at top-level in LoadCentreCollect"
-        )
-        assert "features" not in values["multi_rotation_scan"], (
-            "Features flags must be specified at top-level in LoadCentreCollect"
-        )
         keys_from_outer_load_centre_collect = (
             MxBlueskyParameters.model_fields.keys()
             | WithSample.model_fields.keys()
@@ -89,12 +82,6 @@ class LoadCentreCollect(
         values["robot_load_then_centre"] = new_robot_load_then_centre_params
         return values
 
-    @model_validator(mode="after")
-    def _ensure_features_are_internally_consistent(self) -> Self:
-        self.robot_load_then_centre.features = self.features
-        self.multi_rotation_scan.features = self.features
-        return self
-
     @model_validator(mode="after")
     def _check_rotation_start_xyz_is_not_specified(self) -> Self:
         for scan in self.multi_rotation_scan.single_rotation_scans:

mx_bluesky/hyperion/parameters/robot_load.py

@@ -10,7 +10,6 @@ from mx_bluesky.common.parameters.components import (
 from mx_bluesky.common.parameters.constants import (
     HardwareConstants,
 )
-from mx_bluesky.hyperion.parameters.components import WithHyperionUDCFeatures
 from mx_bluesky.hyperion.parameters.gridscan import (
     GridCommonWithHyperionDetectorParams,
     PinTipCentreThenXrayCentre,
@@ -23,9 +22,7 @@ class RobotLoadAndEnergyChange(
     thawing_time: float = Field(default=HardwareConstants.THAWING_TIME)
 
 
-class RobotLoadThenCentre(
-    GridCommonWithHyperionDetectorParams, WithHyperionUDCFeatures
-):
+class RobotLoadThenCentre(GridCommonWithHyperionDetectorParams):
     thawing_time: float = Field(default=HardwareConstants.THAWING_TIME)
 
     @property

mx_bluesky/hyperion/parameters/rotation.py

@@ -28,7 +28,6 @@ from mx_bluesky.common.parameters.components import (
     WithSample,
     WithScan,
 )
-from mx_bluesky.hyperion.parameters.components import WithHyperionUDCFeatures
 from mx_bluesky.hyperion.parameters.constants import (
     CONST,
     I03Constants,
@@ -56,7 +55,7 @@ class RotationScanPerSweep(OptionalGonioAngleStarts, OptionalXyzStarts, WithSamp
     nexus_vds_start_img: int = Field(default=0, ge=0)
 
 
-class RotationExperiment(DiffractionExperiment, WithHyperionUDCFeatures):
+class RotationExperiment(DiffractionExperiment):
     shutter_opening_time_s: float = Field(default=CONST.I03.SHUTTER_TIME_S)
     rotation_increment_deg: float = Field(default=0.1, gt=0)
     ispyb_experiment_type: IspybExperimentType = Field(

mx_bluesky/hyperion/plan_runner.py

@@ -0,0 +1,78 @@
+import threading
+from collections.abc import Callable
+
+from blueapi.core import BlueskyContext
+from bluesky.utils import MsgGenerator, RequestAbort
+
+from mx_bluesky.common.parameters.constants import Status
+from mx_bluesky.common.utils.exceptions import WarningException
+from mx_bluesky.common.utils.log import LOGGER
+from mx_bluesky.hyperion.runner import BaseRunner
+
+
+class PlanException(Exception):
+    """Identifies an exception that was encountered during plan execution."""
+
+    pass
+
+
+class PlanRunner(BaseRunner):
+    """Runner that executes experiments from inside a running Bluesky plan"""
+
+    def __init__(
+        self,
+        context: BlueskyContext,
+    ) -> None:
+        super().__init__(context)
+        self.current_status: Status = Status.IDLE
+
+    def execute_plan(
+        self,
+        experiment: Callable[[], MsgGenerator],
+    ) -> MsgGenerator:
+        """Execute the specified experiment plan.
+        Args:
+            experiment: The experiment to run
+        Raises:
+            PlanException: If the plan raised an exception
+            RequestAbort: If the RunEngine aborted during execution"""
+
+        self.current_status = Status.BUSY
+
+        try:
+            yield from experiment()
+            self.current_status = Status.IDLE
+        except WarningException as e:
+            LOGGER.warning("Plan failed with warning", exc_info=e)
+            self.current_status = Status.FAILED
+        except RequestAbort:
+            # This will occur when the run engine processes an abort when we shut down
+            LOGGER.info("UDC Runner aborting")
+            raise
+        except Exception as e:
+            LOGGER.error("Plan failed with exception", exc_info=e)
+            self.current_status = Status.FAILED
+            raise PlanException("Exception thrown in plan execution") from e
+
+    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.RE.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()
+            return

mx_bluesky/hyperion/runner.py

@@ -0,0 +1,189 @@
+import threading
+from abc import abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass
+from queue import Empty, Queue
+from typing import Any
+
+from blueapi.core import BlueskyContext
+from bluesky.callbacks.zmq import Publisher
+from bluesky.utils import MsgGenerator
+
+from mx_bluesky.common.external_interaction.callbacks.common.log_uid_tag_callback import (
+    LogUidTaggingCallback,
+)
+from mx_bluesky.common.parameters.components import MxBlueskyParameters
+from mx_bluesky.common.parameters.constants import Actions, Status
+from mx_bluesky.common.utils.exceptions import WarningException
+from mx_bluesky.common.utils.log import LOGGER
+from mx_bluesky.common.utils.tracing import TRACER
+from mx_bluesky.hyperion.experiment_plans.experiment_registry import PLAN_REGISTRY
+from mx_bluesky.hyperion.parameters.constants import CONST
+
+
+@dataclass
+class Command:
+    action: Actions
+    devices: Any | None = None
+    experiment: Callable[[Any, Any], MsgGenerator] | None = None
+
+    def __str__(self):
+        return f"Command({self.action}, {self.parameters}"
+
+    parameters: MxBlueskyParameters | None = None
+
+
+@dataclass
+class StatusAndMessage:
+    status: str
+    message: str = ""
+
+    def __init__(self, status: Status, message: str = "") -> None:
+        self.status = status.value
+        self.message = message
+
+
+@dataclass
+class ErrorStatusAndMessage(StatusAndMessage):
+    exception_type: str = ""
+
+
+def make_error_status_and_message(exception: Exception):
+    return ErrorStatusAndMessage(
+        status=Status.FAILED.value,
+        message=repr(exception),
+        exception_type=type(exception).__name__,
+    )
+
+
+class BaseRunner:
+    @abstractmethod
+    def shutdown(self):
+        """Performs orderly prompt shutdown.
+        Aborts the run engine and terminates the loop waiting for messages."""
+        pass
+
+    def __init__(self, context: BlueskyContext):
+        self.context: BlueskyContext = context
+        self.RE = context.run_engine
+        # These references are necessary to maintain liveness of callbacks because RE
+        # only keeps a weakref
+        self._logging_uid_tag_callback = LogUidTaggingCallback()
+        self._publisher = Publisher(f"localhost:{CONST.CALLBACK_0MQ_PROXY_PORTS[0]}")
+
+        self.RE.subscribe(self._logging_uid_tag_callback)
+        LOGGER.info("Connecting to external callback ZMQ proxy...")
+        self.RE.subscribe(self._publisher)
+
+
+class GDARunner(BaseRunner):
+    """Runner that executes plans submitted by Flask requests from GDA."""
+
+    def __init__(
+        self,
+        context: BlueskyContext,
+    ) -> None:
+        super().__init__(context)
+        self.current_status: StatusAndMessage = StatusAndMessage(Status.IDLE)
+        self._last_run_aborted: bool = False
+        self._command_queue: Queue[Command] = Queue()
+
+    def start(
+        self,
+        experiment: Callable,
+        parameters: MxBlueskyParameters,
+        plan_name: str | None = None,
+    ) -> StatusAndMessage:
+        """Start a new bluesky plan
+        Args:
+            experiment: A bluesky plan
+            parameters: The parameters to be submitted
+            plan_name: Name of the plan that will be used to resolve the composite factory
+                to supply devices for the plan, if any are needed"""
+        LOGGER.info(f"Started with parameters: {parameters.model_dump_json(indent=2)}")
+
+        devices: Any = (
+            PLAN_REGISTRY[plan_name]["setup"](self.context) if plan_name else None
+        )
+
+        if (
+            self.current_status.status == Status.BUSY.value
+            or self.current_status.status == Status.ABORTING.value
+        ):
+            return StatusAndMessage(Status.FAILED, "Bluesky already running")
+        else:
+            self.current_status = StatusAndMessage(Status.BUSY)
+            self._command_queue.put(
+                Command(
+                    action=Actions.START,
+                    devices=devices,
+                    experiment=experiment,
+                    parameters=parameters,
+                )
+            )
+            return StatusAndMessage(Status.SUCCESS)
+
+    def stop(self) -> StatusAndMessage:
+        """Stop the currently executing plan."""
+        if self.current_status.status == Status.ABORTING.value:
+            return StatusAndMessage(Status.FAILED, "Bluesky already stopping")
+        else:
+            self.current_status = StatusAndMessage(Status.ABORTING)
+            stopping_thread = threading.Thread(target=self._stopping_thread)
+            stopping_thread.start()
+            self._last_run_aborted = True
+            return StatusAndMessage(Status.ABORTING)
+
+    def shutdown(self):
+        """Stops the run engine and the loop waiting for messages."""
+        print("Shutting down: Stopping the run engine gracefully")
+        self.stop()
+        self._command_queue.put(Command(action=Actions.SHUTDOWN))
+
+    def _stopping_thread(self):
+        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.RE.abort()
+            self.current_status = StatusAndMessage(Status.IDLE)
+        except Exception as e:
+            self.current_status = make_error_status_and_message(e)
+
+    def fetch_next_command(self) -> Command:
+        """Fetch the next command from the queue, blocks if queue is empty."""
+        return self._command_queue.get()
+
+    def try_fetch_next_command(self) -> Command | None:
+        """Fetch the next command from the queue or return None if no command available."""
+        try:
+            return self._command_queue.get(block=False)
+        except Empty:
+            return None
+
+    def wait_on_queue(self):
+        while True:
+            command = self.fetch_next_command()
+            if command.action == Actions.SHUTDOWN:
+                return
+            elif command.action == Actions.START:
+                if command.experiment is None:
+                    raise ValueError("No experiment provided for START")
+                try:
+                    with TRACER.start_span("do_run"):
+                        self.RE(command.experiment(command.devices, command.parameters))
+
+                    self.current_status = StatusAndMessage(Status.IDLE)
+
+                    self._last_run_aborted = False
+                except WarningException as exception:
+                    LOGGER.warning("Warning Exception", exc_info=True)
+                    self.current_status = make_error_status_and_message(exception)
+                except Exception as exception:
+                    LOGGER.error("Exception on running plan", exc_info=True)
+
+                    if self._last_run_aborted:
+                        # Aborting will cause an exception here that we want to swallow
+                        self._last_run_aborted = False
+                    else:
+                        self.current_status = make_error_status_and_message(exception)

mx_bluesky/hyperion/utils/context.py

@@ -1,5 +1,6 @@
 from blueapi.core import BlueskyContext
-from dodal.utils import get_beamline_based_on_environment_variable
+from dodal.common.beamlines.beamline_utils import clear_devices
+from dodal.utils import collect_factories, get_beamline_based_on_environment_variable
 
 import mx_bluesky.hyperion.experiment_plans as hyperion_plans
 from mx_bluesky.common.utils.log import LOGGER
@@ -9,11 +10,24 @@ def setup_context(dev_mode: bool = False) -> BlueskyContext:
     context = BlueskyContext()
     context.with_plan_module(hyperion_plans)
 
-    context.with_dodal_module(
-        get_beamline_based_on_environment_variable(),
-        mock=dev_mode,
-    )
+    setup_devices(context, dev_mode)
 
     LOGGER.info(f"Plans found in context: {context.plan_functions.keys()}")
 
     return context
+
+
+def clear_all_device_caches(context: BlueskyContext):
+    context.unregister_all_devices()
+    clear_devices()
+
+    for f in collect_factories(get_beamline_based_on_environment_variable()).values():
+        if hasattr(f, "cache_clear"):
+            f.cache_clear()  # type: ignore
+
+
+def setup_devices(context: BlueskyContext, dev_mode: bool):
+    context.with_dodal_module(
+        get_beamline_based_on_environment_variable(),
+        mock=dev_mode,
+    )

mx_bluesky/phase1_zebra/__init__.py

@@ -0,0 +1 @@
+"""The zebra on i03 and i04 (phase1 beamlines) are configured in a specific way which are compatible to the plans in this module"""