mx-bluesky 1.5.5__py3-none-any.whl → 1.5.7__py3-none-any.whl

mx_bluesky/common/external_interaction/callbacks/xray_centre/ispyb_callback.py

@@ -1,23 +1,28 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Sequence
+from enum import StrEnum
+from math import isclose
 from time import time
 from typing import TYPE_CHECKING, Any, TypeVar
 
 from bluesky import preprocessors as bpp
 from bluesky.utils import MsgGenerator, make_decorator
+from dodal.devices.zocalo import ZocaloStartInfo
 
 from mx_bluesky.common.external_interaction.callbacks.common.ispyb_callback_base import (
     BaseISPyBCallback,
+    D,
 )
 from mx_bluesky.common.external_interaction.callbacks.common.ispyb_mapping import (
     populate_data_collection_group,
     populate_remaining_data_collection_info,
 )
+from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import (
+    ZocaloInfoGenerator,
+)
 from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_mapping import (
     construct_comment_for_gridscan,
-    populate_xy_data_collection_info,
-    populate_xz_data_collection_info,
 )
 from mx_bluesky.common.external_interaction.ispyb.data_model import (
     DataCollectionGridInfo,
@@ -33,14 +38,21 @@ from mx_bluesky.common.external_interaction.ispyb.ispyb_store import (
 )
 from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample
 from mx_bluesky.common.parameters.constants import DocDescriptorNames, PlanNameConstants
-from mx_bluesky.common.parameters.gridscan import (
-    GridCommon,
-)
+from mx_bluesky.common.parameters.gridscan import GridCommon
 from mx_bluesky.common.utils.exceptions import (
     ISPyBDepositionNotMade,
     SampleException,
 )
 from mx_bluesky.common.utils.log import ISPYB_ZOCALO_CALLBACK_LOGGER, set_dcgid_tag
+from mx_bluesky.common.utils.utils import number_of_frames_from_scan_spec
+
+OMEGA_TOLERANCE = 1
+
+
+class GridscanPlane(StrEnum):
+    OMEGA_XY = "0"
+    OMEGA_XZ = "90"
+
 
 if TYPE_CHECKING:
     from event_model import Event, RunStart, RunStop
@@ -89,10 +101,10 @@ class GridscanISPyBCallback(BaseISPyBCallback):
     ) -> None:
         super().__init__(emit=emit)
         self.ispyb: StoreInIspyb
-        self.ispyb_ids: IspybIds = IspybIds()
         self.param_type = param_type
         self._start_of_fgs_uid: str | None = None
         self._processing_start_time: float | None = None
+        self._grid_plane_to_id_map: dict[GridscanPlane, int] = {}
         self.data_collection_group_info: DataCollectionGroupInfo | None
 
     def activity_gated_start(self, doc: RunStart):
@@ -108,6 +120,7 @@ class GridscanISPyBCallback(BaseISPyBCallback):
             mx_bluesky_parameters = doc.get("mx_bluesky_parameters")
             assert isinstance(mx_bluesky_parameters, str)
             self.params = self.param_type.model_validate_json(mx_bluesky_parameters)
+            assert isinstance(self.params, DiffractionExperimentWithSample)
             self.ispyb = StoreInIspyb(self.ispyb_config)
             self.data_collection_group_info = populate_data_collection_group(
                 self.params
@@ -118,8 +131,8 @@ class GridscanISPyBCallback(BaseISPyBCallback):
                     data_collection_info=populate_remaining_data_collection_info(
                         "MX-Bluesky: Xray centring 1 -",
                         None,
-                        populate_xy_data_collection_info(
-                            self.params.detector_params,
+                        DataCollectionInfo(
+                            data_collection_number=self.params.detector_params.run_number,
                         ),
                         self.params,
                     ),
@@ -128,7 +141,11 @@ class GridscanISPyBCallback(BaseISPyBCallback):
                     data_collection_info=populate_remaining_data_collection_info(
                         "MX-Bluesky: Xray centring 2 -",
                         None,
-                        populate_xz_data_collection_info(self.params.detector_params),
+                        DataCollectionInfo(
+                            data_collection_number=(
+                                self.params.detector_params.run_number + 1
+                            ),
+                        ),
                         self.params,
                     )
                 ),
@@ -175,7 +192,6 @@ class GridscanISPyBCallback(BaseISPyBCallback):
         assert self.params, "ISPyB handler didn't receive parameters!"
         assert self.data_collection_group_info, "No data collection group"
         data = doc["data"]
-        data_collection_id = None
         data_collection_info = DataCollectionInfo(
             xtal_snapshot1=data.get("oav-grid_snapshot-last_path_full_overlay"),
             xtal_snapshot2=data.get("oav-grid_snapshot-last_path_outer"),
@@ -214,10 +230,9 @@ class GridscanISPyBCallback(BaseISPyBCallback):
                 f"by {data_collection_grid_info.steps_y} "
             )
 
-        if len(self.ispyb_ids.data_collection_ids) > self._oav_snapshot_event_idx:
-            data_collection_id = self.ispyb_ids.data_collection_ids[
-                self._oav_snapshot_event_idx
-            ]
+        data_collection_id = self.ispyb_ids.data_collection_ids[
+            self._oav_snapshot_event_idx
+        ]
         self._populate_axis_info(data_collection_info, doc["data"]["smargon-omega"])
 
         scan_data_info = ScanDataInfo(
@@ -228,6 +243,11 @@ class GridscanISPyBCallback(BaseISPyBCallback):
         ISPYB_ZOCALO_CALLBACK_LOGGER.info(
             "Updating ispyb data collection after oav snapshot."
         )
+        grid_plane = _smargon_omega_to_xyxz_plane(doc["data"]["smargon-omega"])
+        # Snapshots may be triggered in a different order to gridscans, so save
+        # the mapping to the data collection id in order to trigger Zocalo correctly.
+        self._grid_plane_to_id_map[grid_plane] = data_collection_id
+
         self._oav_snapshot_event_idx += 1
         return [scan_data_info]
 
@@ -294,5 +314,43 @@ class GridscanISPyBCallback(BaseISPyBCallback):
                 self.ispyb_ids.data_collection_group_id,
             )
             self.data_collection_group_info = None
+            self._grid_plane_to_id_map.clear()
             return super().activity_gated_stop(doc)
-        return self._tag_doc(doc)
+        return self.tag_doc(doc)
+
+    def tag_doc(self, doc: D) -> D:
+        doc = super().tag_doc(doc)
+        assert isinstance(doc, dict)
+        if self._grid_plane_to_id_map:
+            doc["grid_plane_to_id_map"] = self._grid_plane_to_id_map
+        return doc  # type: ignore
+
+
+def generate_start_info_from_omega_map() -> ZocaloInfoGenerator:
+    """
+    Generate the zocalo trigger info from bluesky runs where the frame number is
+    computed using metadata added to the document by the ISPyB callback and the
+    run start which together can be used to determine the correct frame numbering.
+    """
+    doc = yield []
+    omega_to_scan_spec = doc["omega_to_scan_spec"]
+    start_frame = 0
+    infos = []
+    for i, omega in enumerate([GridscanPlane.OMEGA_XY, GridscanPlane.OMEGA_XZ]):
+        frames = number_of_frames_from_scan_spec(omega_to_scan_spec[omega])
+        infos.append(
+            ZocaloStartInfo(
+                doc["grid_plane_to_id_map"][omega], None, start_frame, frames, i
+            )
+        )
+        start_frame += frames
+    yield infos
+
+
+def _smargon_omega_to_xyxz_plane(smargon_omega: float) -> GridscanPlane:
+    modulo_180 = abs(smargon_omega) % 180
+    is_xy = isclose(modulo_180, 0, abs_tol=OMEGA_TOLERANCE)
+    assert is_xy or isclose(modulo_180, 90, abs_tol=OMEGA_TOLERANCE), (
+        f"Smargon snapshot omega not in tolerance of compass point {smargon_omega}"
+    )
+    return GridscanPlane.OMEGA_XY if is_xy else GridscanPlane.OMEGA_XZ

mx_bluesky/common/external_interaction/callbacks/xray_centre/ispyb_mapping.py

@@ -1,33 +1,13 @@
 from __future__ import annotations
 
 import numpy
-from dodal.devices.detector import DetectorParams
 from dodal.devices.oav import utils as oav_utils
 
 from mx_bluesky.common.external_interaction.ispyb.data_model import (
     DataCollectionGridInfo,
-    DataCollectionInfo,
 )
 
 
-def populate_xz_data_collection_info(detector_params: DetectorParams):
-    assert (
-        detector_params.omega_start is not None
-        and detector_params.run_number is not None
-    ), "StoreGridscanInIspyb failed to get parameters"
-    run_number = detector_params.run_number + 1
-    info = DataCollectionInfo(
-        data_collection_number=run_number,
-    )
-    return info
-
-
-def populate_xy_data_collection_info(detector_params: DetectorParams):
-    return DataCollectionInfo(
-        data_collection_number=detector_params.run_number,
-    )
-
-
 def construct_comment_for_gridscan(grid_info: DataCollectionGridInfo) -> str:
     assert grid_info is not None, "StoreGridScanInIspyb failed to get parameters"
 

mx_bluesky/common/parameters/components.py

@@ -231,6 +231,7 @@ class DiffractionExperimentWithSample(DiffractionExperiment, WithSample): ...
 
 class MultiXtalSelection(BaseModel):
     name: str
+    ignore_xtal_not_found: bool = False
 
 
 class TopNByMaxCountSelection(MultiXtalSelection):

mx_bluesky/common/parameters/device_composites.py

@@ -8,7 +8,7 @@ from dodal.devices.common_dcm import BaseDCM
 from dodal.devices.detector.detector_motion import DetectorMotion
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.fast_grid_scan import (
-    ZebraFastGridScan,
+    ZebraFastGridScanThreeD,
 )
 from dodal.devices.flux import Flux
 from dodal.devices.mx_phase1.beamstop import Beamstop
@@ -53,7 +53,7 @@ class GridDetectThenXRayCentreComposite(FlyScanEssentialDevices):
     beamstop: Beamstop
     dcm: BaseDCM
     detector_motion: DetectorMotion
-    zebra_fast_grid_scan: ZebraFastGridScan
+    zebra_fast_grid_scan: ZebraFastGridScanThreeD
     flux: Flux
     oav: OAV
     pin_tip_detection: PinTipDetection

mx_bluesky/common/parameters/gridscan.py

@@ -1,15 +1,19 @@
 from __future__ import annotations
 
+from abc import abstractmethod
+from typing import Generic, TypeVar
+
 from dodal.devices.aperturescatterguard import ApertureValue
 from dodal.devices.detector.det_dim_constants import EIGER2_X_9M_SIZE, EIGER2_X_16M_SIZE
 from dodal.devices.detector.detector import DetectorParams
 from dodal.devices.fast_grid_scan import (
-    ZebraGridScanParams,
+    GridScanParamsCommon,
+    ZebraGridScanParamsThreeD,
 )
 from dodal.utils import get_beamline_name
 from pydantic import Field, PrivateAttr
 from scanspec.core import Path as ScanPath
-from scanspec.specs import Line, Static
+from scanspec.specs import Concat, Line, Product, Static
 
 from mx_bluesky.common.parameters.components import (
     DiffractionExperimentWithSample,
@@ -33,6 +37,10 @@ DETECTOR_SIZE_PER_BEAMLINE = {
     "i04": EIGER2_X_16M_SIZE,
 }
 
+GridScanParamType = TypeVar(
+    "GridScanParamType", bound=GridScanParamsCommon, covariant=True
+)
+
 
 class GridCommon(
     DiffractionExperimentWithSample,
@@ -87,36 +95,80 @@ class GridCommon(
         )
 
 
-class SpecifiedGrid(XyzStarts, WithScan):
+class SpecifiedGrid(GridCommon, XyzStarts, WithScan, Generic[GridScanParamType]):
     """A specified grid is one which has defined values for the start position,
     grid and box sizes, etc., as opposed to parameters for a plan which will create
     those parameters at some point (e.g. through optical pin detection)."""
 
+    grid1_omega_deg: float = Field(default=GridscanParamConstants.OMEGA_1)
+    x_step_size_um: float = Field(default=GridscanParamConstants.BOX_WIDTH_UM)
+    y_step_size_um: float = Field(default=GridscanParamConstants.BOX_WIDTH_UM)
+    x_steps: int = Field(gt=0)
+    y_steps: int = Field(gt=0)
+    _set_stub_offsets: bool = PrivateAttr(default_factory=lambda: False)
+
+    @property
+    @abstractmethod
+    def FGS_params(self) -> GridScanParamType: ...
+
+    def do_set_stub_offsets(self, value: bool):
+        self._set_stub_offsets = value
+
+    @property
+    def grid_1_spec(self):
+        x_end = self.x_start_um + self.x_step_size_um * (self.x_steps - 1)
+        y1_end = self.y_start_um + self.y_step_size_um * (self.y_steps - 1)
+        grid_1_x = Line("sam_x", self.x_start_um, x_end, self.x_steps)
+        grid_1_y = Line("sam_y", self.y_start_um, y1_end, self.y_steps)
+        grid_1_z = Static("sam_z", self.z_start_um)
+        return grid_1_y.zip(grid_1_z) * ~grid_1_x
+
+    @property
+    def scan_indices(self) -> list[int]:
+        """The first index of each gridscan, useful for writing nexus files/VDS"""
+        return [
+            0,
+            len(ScanPath(self.grid_1_spec.calculate()).consume().midpoints["sam_x"]),
+        ]
+
+    @property
+    @abstractmethod
+    def scan_spec(self) -> Product[str] | Concat[str]:
+        """A fully specified ScanSpec object representing all grids, with x, y, z and
+        omega positions."""
+
+    @property
+    def scan_points(self):
+        """A list of all the points in the scan_spec."""
+        return ScanPath(self.scan_spec.calculate()).consume().midpoints
+
+    @property
+    def scan_points_first_grid(self):
+        """A list of all the points in the first grid scan."""
+        return ScanPath(self.grid_1_spec.calculate()).consume().midpoints
+
+    @property
+    def num_images(self) -> int:
+        return len(self.scan_points["sam_x"])
+
 
 class SpecifiedThreeDGridScan(
-    GridCommon,
-    SpecifiedGrid,
+    SpecifiedGrid[ZebraGridScanParamsThreeD],
     SplitScan,
     WithOptionalEnergyChange,
 ):
     """Parameters representing a so-called 3D grid scan, which consists of doing a
     gridscan in X and Y, followed by one in X and Z."""
 
-    grid1_omega_deg: float = Field(default=GridscanParamConstants.OMEGA_1)
-    grid2_omega_deg: float = Field(default=GridscanParamConstants.OMEGA_2)
-    x_step_size_um: float = Field(default=GridscanParamConstants.BOX_WIDTH_UM)
-    y_step_size_um: float = Field(default=GridscanParamConstants.BOX_WIDTH_UM)
+    z_steps: int = Field(gt=0)
     z_step_size_um: float = Field(default=GridscanParamConstants.BOX_WIDTH_UM)
     y2_start_um: float
     z2_start_um: float
-    x_steps: int = Field(gt=0)
-    y_steps: int = Field(gt=0)
-    z_steps: int = Field(gt=0)
-    _set_stub_offsets: bool = PrivateAttr(default_factory=lambda: False)
+    grid2_omega_deg: float = Field(default=GridscanParamConstants.OMEGA_2)
 
     @property
-    def FGS_params(self) -> ZebraGridScanParams:
-        return ZebraGridScanParams(
+    def FGS_params(self) -> ZebraGridScanParamsThreeD:
+        return ZebraGridScanParamsThreeD(
             x_steps=self.x_steps,
             y_steps=self.y_steps,
             z_steps=self.z_steps,
@@ -133,18 +185,6 @@ class SpecifiedThreeDGridScan(
             transmission_fraction=self.transmission_frac,
         )
 
-    def do_set_stub_offsets(self, value: bool):
-        self._set_stub_offsets = value
-
-    @property
-    def grid_1_spec(self):
-        x_end = self.x_start_um + self.x_step_size_um * (self.x_steps - 1)
-        y1_end = self.y_start_um + self.y_step_size_um * (self.y_steps - 1)
-        grid_1_x = Line("sam_x", self.x_start_um, x_end, self.x_steps)
-        grid_1_y = Line("sam_y", self.y_start_um, y1_end, self.y_steps)
-        grid_1_z = Static("sam_z", self.z_start_um)
-        return grid_1_y.zip(grid_1_z) * ~grid_1_x
-
     @property
     def grid_2_spec(self):
         x_end = self.x_start_um + self.x_step_size_um * (self.x_steps - 1)
@@ -154,35 +194,13 @@ class SpecifiedThreeDGridScan(
         grid_2_y = Static("sam_y", self.y2_start_um)
         return grid_2_z.zip(grid_2_y) * ~grid_2_x
 
-    @property
-    def scan_indices(self):
-        """The first index of each gridscan, useful for writing nexus files/VDS"""
-        return [
-            0,
-            len(ScanPath(self.grid_1_spec.calculate()).consume().midpoints["sam_x"]),
-        ]
-
     @property
     def scan_spec(self):
         """A fully specified ScanSpec object representing both grids, with x, y, z and
         omega positions."""
         return self.grid_1_spec.concat(self.grid_2_spec)
 
-    @property
-    def scan_points(self):
-        """A list of all the points in the scan_spec."""
-        return ScanPath(self.scan_spec.calculate()).consume().midpoints
-
-    @property
-    def scan_points_first_grid(self):
-        """A list of all the points in the first grid scan."""
-        return ScanPath(self.grid_1_spec.calculate()).consume().midpoints
-
     @property
     def scan_points_second_grid(self):
         """A list of all the points in the second grid scan."""
         return ScanPath(self.grid_2_spec.calculate()).consume().midpoints
-
-    @property
-    def num_images(self) -> int:
-        return len(self.scan_points["sam_x"])

mx_bluesky/hyperion/__main__.py

@@ -1,4 +1,5 @@
 import json
+import signal
 import threading
 from dataclasses import asdict
 from sys import argv
@@ -32,9 +33,10 @@ from mx_bluesky.hyperion.parameters.cli import (
     HyperionMode,
     parse_cli_args,
 )
-from mx_bluesky.hyperion.parameters.constants import CONST
+from mx_bluesky.hyperion.parameters.constants import CONST, HyperionConstants
 from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
 from mx_bluesky.hyperion.plan_runner import PlanRunner
+from mx_bluesky.hyperion.plan_runner_api import create_server_for_udc
 from mx_bluesky.hyperion.runner import (
     GDARunner,
     StatusAndMessage,
@@ -170,7 +172,7 @@ def main():
     """Main application entry point."""
     args = parse_cli_args()
     initialise_globals(args)
-    hyperion_port = 5005
+    hyperion_port = HyperionConstants.HYPERION_PORT
     context = setup_context(dev_mode=args.dev_mode)
 
     if args.mode == HyperionMode.GDA:
@@ -188,7 +190,18 @@ def main():
         )
         runner.wait_on_queue()
     else:
-        run_forever(PlanRunner(context))
+        plan_runner = PlanRunner(context, args.dev_mode)
+        create_server_for_udc(plan_runner)
+        _register_sigterm_handler(plan_runner)
+        run_forever(plan_runner)
+
+
+def _register_sigterm_handler(runner: PlanRunner):
+    def shutdown_on_sigterm(sig_num, frame):
+        LOGGER.info("Received SIGTERM, shutting down...")
+        runner.shutdown()
+
+    signal.signal(signal.SIGTERM, shutdown_on_sigterm)
 
 
 if __name__ == "__main__":

mx_bluesky/hyperion/baton_handler.py

@@ -6,8 +6,14 @@ from blueapi.core.context import BlueskyContext
 from bluesky import plan_stubs as bps
 from bluesky import preprocessors as bpp
 from bluesky.utils import MsgGenerator, RunEngineInterrupted
+from dodal.common.beamlines.commissioning_mode import set_commissioning_signal
+from dodal.devices.aperturescatterguard import ApertureScatterguard
 from dodal.devices.baton import Baton
+from dodal.devices.motors import XYZStage
+from dodal.devices.robot import BartRobot
+from dodal.devices.smargon import Smargon
 
+from mx_bluesky.common.device_setup_plans.robot_load_unload import robot_unload
 from mx_bluesky.common.experiment_plans.inner_plans.udc_default_state import (
     UDCDefaultDevices,
     move_to_udc_default_state,
@@ -77,6 +83,7 @@ def run_udc_when_requested(context: BlueskyContext, runner: PlanRunner):
 
     def acquire_baton() -> MsgGenerator:
         yield from _wait_for_hyperion_requested(baton)
+        LOGGER.debug("Hyperion is now current baton holder.")
         yield from bps.abs_set(baton.current_user, HYPERION_USER)
 
     def collect() -> MsgGenerator:
@@ -96,14 +103,20 @@ def run_udc_when_requested(context: BlueskyContext, runner: PlanRunner):
 
         # re-fetch the baton because the device has been reinstantiated
         baton = _get_baton(context)
+        current_visit: str | None = None
         while (yield from _is_requesting_baton(baton)):
-            yield from _fetch_and_process_agamemnon_instruction(baton, runner)
+            current_visit = yield from _fetch_and_process_agamemnon_instruction(
+                baton, runner, current_visit
+            )
+        if current_visit:
+            yield from _perform_robot_unload(runner.context, current_visit)
 
     def release_baton() -> MsgGenerator:
         # If hyperion has given up the baton itself we need to also release requested
         # user so that hyperion doesn't think we're requested again
         baton = _get_baton(context)
-        previous_requested_user = yield from _safely_release_baton(baton)
+        previous_requested_user = yield from _unrequest_baton(baton)
+        LOGGER.debug("Hyperion no longer current baton holder.")
         yield from bps.abs_set(baton.current_user, NO_USER, wait=True)
         _raise_baton_released_alert(get_alerting_service(), previous_requested_user)
 
@@ -111,11 +124,11 @@ def run_udc_when_requested(context: BlueskyContext, runner: PlanRunner):
         yield from bpp.contingency_wrapper(collect(), final_plan=release_baton)
 
     context.run_engine(acquire_baton())
-    _initialise_udc(context)
+    _initialise_udc(context, runner.is_dev_mode)
     context.run_engine(collect_then_release())
 
 
-def _initialise_udc(context: BlueskyContext):
+def _initialise_udc(context: BlueskyContext, dev_mode: bool):
     """
     Perform all initialisation that happens at the start of UDC just after the
     baton is acquired, but before we execute any plans or move hardware.
@@ -125,21 +138,25 @@ def _initialise_udc(context: BlueskyContext):
     """
     LOGGER.info("Initialising mx-bluesky for UDC start...")
     clear_all_device_caches(context)
-    setup_devices(context, False)
+    LOGGER.debug("Reinitialising beamline devices")
+    setup_devices(context, dev_mode)
+    set_commissioning_signal(_get_baton(context).commissioning)
 
 
 def _wait_for_hyperion_requested(baton: Baton):
+    LOGGER.debug("Hyperion waiting for baton...")
     SLEEP_PER_CHECK = 0.1
     while True:
         requested_user = yield from bps.rd(baton.requested_user)
         if requested_user == HYPERION_USER:
+            LOGGER.debug("Baton requested for Hyperion")
             break
         yield from bps.sleep(SLEEP_PER_CHECK)
 
 
 def _fetch_and_process_agamemnon_instruction(
-    baton: Baton, runner: PlanRunner
-) -> MsgGenerator:
+    baton: Baton, runner: PlanRunner, current_visit: str | None
+) -> MsgGenerator[str | None]:
     parameter_list: Sequence[MxBlueskyParameters] = create_parameters_from_agamemnon()
     if parameter_list:
         for parameters in parameter_list:
@@ -148,6 +165,7 @@ def _fetch_and_process_agamemnon_instruction(
             )
             match parameters:
                 case LoadCentreCollect():
+                    current_visit = parameters.visit
                     devices: Any = create_devices(runner.context)
                     yield from runner.execute_plan(
                         partial(load_centre_collect_full, devices, parameters)
@@ -161,7 +179,8 @@ def _fetch_and_process_agamemnon_instruction(
     else:
         _raise_udc_completed_alert(get_alerting_service())
         # Release the baton for orderly exit from the instruction loop
-        yield from _safely_release_baton(baton)
+        yield from _unrequest_baton(baton)
+    return current_visit
 
 
 def _raise_udc_start_alert(alert_service: AlertService):
@@ -205,7 +224,7 @@ def _get_baton(context: BlueskyContext) -> Baton:
     return find_device_in_context(context, "baton", Baton)
 
 
-def _safely_release_baton(baton: Baton) -> MsgGenerator[str]:
+def _unrequest_baton(baton: Baton) -> MsgGenerator[str]:
     """Relinquish the requested user of the baton if it is not already requested
     by another user.
 
@@ -214,6 +233,17 @@ def _safely_release_baton(baton: Baton) -> MsgGenerator[str]:
     """
     requested_user = yield from bps.rd(baton.requested_user)
     if requested_user == HYPERION_USER:
+        LOGGER.debug("Hyperion no longer requesting baton")
         yield from bps.abs_set(baton.requested_user, NO_USER)
         return NO_USER
     return requested_user
+
+
+def _perform_robot_unload(context: BlueskyContext, visit: str) -> MsgGenerator:
+    robot = find_device_in_context(context, "robot", BartRobot)
+    smargon = find_device_in_context(context, "smargon", Smargon)
+    aperture_scatterguard = find_device_in_context(
+        context, "aperture_scatterguard", ApertureScatterguard
+    )
+    lower_gonio = find_device_in_context(context, "lower_gonio", XYZStage)
+    yield from robot_unload(robot, smargon, aperture_scatterguard, lower_gonio, visit)

mx_bluesky/hyperion/device_setup_plans/smargon.py

@@ -1,6 +1,8 @@
 import numpy as np
 from bluesky import plan_stubs as bps
+from bluesky.utils import FailedStatus
 from dodal.devices.smargon import CombinedMove, Smargon
+from ophyd_async.epics.motor import MotorLimitsException
 
 from mx_bluesky.common.utils.exceptions import SampleException
 
@@ -9,12 +11,15 @@ def move_smargon_warn_on_out_of_range(
     smargon: Smargon, position: np.ndarray | list[float] | tuple[float, float, float]
 ):
     """Throws a SampleException if the specified position is out of range for the
-    smargon. Otherwise moves to that position."""
-    limits = yield from smargon.get_xyz_limits()
-    if not limits.position_valid(position):
-        raise SampleException(
-            "Pin tip centring failed - pin too long/short/bent and out of range"
+    smargon. Otherwise moves to that position. The check is from ophyd-async"""
+    try:
+        yield from bps.mv(
+            smargon, CombinedMove(x=position[0], y=position[1], z=position[2])
         )
-    yield from bps.mv(
-        smargon, CombinedMove(x=position[0], y=position[1], z=position[2])
-    )
+    except FailedStatus as fs:
+        if isinstance(fs.__cause__, MotorLimitsException):
+            raise SampleException(
+                "Pin tip centring failed - pin too long/short/bent and out of range"
+            ) from fs.__cause__
+        else:
+            raise fs