mx-bluesky 1.5.0__py3-none-any.whl → 1.5.2__py3-none-any.whl

mx_bluesky/common/device_setup_plans/robot_load_unload.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+from bluesky.utils import MsgGenerator
+from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
+from dodal.devices.motors import XYZStage
+from dodal.devices.robot import BartRobot
+from dodal.devices.smargon import CombinedMove, Smargon, StubPosition
+from dodal.plan_stubs.motor_utils import MoveTooLarge, home_and_reset_wrapper
+
+from mx_bluesky.common.utils.log import LOGGER
+from mx_bluesky.hyperion.parameters.constants import CONST
+
+
+def wait_for_smargon_not_disabled(smargon: Smargon, timeout=60):
+    """Waits for the smargon disabled flag to go low. The robot hardware is responsible
+    for setting this to low when it is safe to move. It does this through a physical
+    connection between the robot and the smargon.
+    """
+    LOGGER.info("Waiting for smargon enabled")
+    SLEEP_PER_CHECK = 0.1
+    times_to_check = int(timeout / SLEEP_PER_CHECK)
+    for _ in range(times_to_check):
+        smargon_disabled = yield from bps.rd(smargon.disabled)
+        if not smargon_disabled:
+            LOGGER.info("Smargon now enabled")
+            return
+        yield from bps.sleep(SLEEP_PER_CHECK)
+    raise TimeoutError(
+        "Timed out waiting for smargon to become enabled after robot load"
+    )
+
+
+def _raise_exception_if_moved_out_of_cryojet(exception):
+    yield from bps.null()
+    if isinstance(exception, MoveTooLarge):
+        raise Exception(
+            f"Moving {exception.axis} back to {exception.position} after \
+                        robot load would move it out of the cryojet. The max safe \
+                        distance is {exception.maximum_move}"
+        )
+
+
+def do_plan_while_lower_gonio_at_home(plan: MsgGenerator, lower_gonio: XYZStage):
+    """Moves the lower gonio to home then performs the provided plan and moves it back.
+
+    The lower gonio must be in the correct position for the robot load and we
+    want to put it back afterwards. Note we don't need to wait for the move as the robot
+    is interlocked to the lower gonio and the move is quicker than the robot takes to
+    get to the load position.
+
+    Args:
+        plan (MsgGenerator): The plan to run while the lower gonio is at home.
+        lower_gonio (XYZStage): The lower gonio to home.
+    """
+    yield from bpp.contingency_wrapper(
+        home_and_reset_wrapper(
+            plan,
+            lower_gonio,
+            BartRobot.LOAD_TOLERANCE_MM,
+            CONST.HARDWARE.CRYOJET_MARGIN_MM,
+            "lower_gonio",
+            wait_for_all=False,
+        ),
+        except_plan=_raise_exception_if_moved_out_of_cryojet,
+    )
+    return "reset-lower_gonio"
+
+
+def prepare_for_robot_load(
+    aperture_scatterguard: ApertureScatterguard, smargon: Smargon
+):
+    yield from bps.abs_set(
+        aperture_scatterguard.selected_aperture,
+        ApertureValue.OUT_OF_BEAM,
+        group="prepare_robot_load",
+    )
+
+    yield from bps.mv(smargon.stub_offsets, StubPosition.RESET_TO_ROBOT_LOAD)
+
+    yield from bps.mv(smargon, CombinedMove(x=0, y=0, z=0, chi=0, phi=0, omega=0))
+
+    yield from bps.wait("prepare_robot_load")
+
+
+def robot_unload(
+    robot: BartRobot,
+    smargon: Smargon,
+    aperture_scatterguard: ApertureScatterguard,
+    lower_gonio: XYZStage,
+    visit: str,
+):
+    """Unloads the currently mounted pin into the location that it was loaded from. The
+    loaded location is stored on the robot and so need not be provided.
+    """
+    yield from prepare_for_robot_load(aperture_scatterguard, smargon)
+    sample_id = yield from bps.rd(robot.sample_id)
+
+    @bpp.run_decorator(
+        md={
+            "subplan_name": CONST.PLAN.ROBOT_UNLOAD,
+            "metadata": {"visit": visit, "sample_id": sample_id},
+            "activate_callbacks": [
+                "RobotLoadISPyBCallback",
+            ],
+        },
+    )
+    def do_robot_unload_and_send_to_ispyb():
+        yield from bps.create(name=CONST.DESCRIPTORS.ROBOT_UPDATE)
+        yield from bps.read(robot)
+        yield from bps.save()
+
+        def _unload():
+            yield from bps.trigger(robot.unload, wait=True)
+            yield from wait_for_smargon_not_disabled(smargon)
+
+        gonio_finished = yield from do_plan_while_lower_gonio_at_home(
+            _unload(), lower_gonio
+        )
+        yield from bps.wait(gonio_finished)
+
+    yield from do_robot_unload_and_send_to_ispyb()

mx_bluesky/common/experiment_plans/change_aperture_then_move_plan.py

@@ -4,6 +4,7 @@ from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureVal
 from dodal.devices.smargon import Smargon, StubPosition
 
 from mx_bluesky.common.device_setup_plans.manipulate_sample import move_x_y_z
+from mx_bluesky.common.parameters.constants import PlanGroupCheckpointConstants
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.common.utils.tracing import TRACER
 from mx_bluesky.common.xrc_result import XRayCentreResult
@@ -43,6 +44,7 @@ def change_aperture_then_move_to_xtal(
 def set_aperture_for_bbox_mm(
     aperture_device: ApertureScatterguard,
     bbox_size_mm: list[float] | numpy.ndarray,
+    group=PlanGroupCheckpointConstants.GRID_READY_FOR_DC,
 ):
     """Sets aperture size based on bbox_size.
 
@@ -71,4 +73,6 @@ def set_aperture_for_bbox_mm(
         f"Setting aperture to {new_selected_aperture} based on bounding box size {bbox_size_mm}."
     )
 
-    yield from bps.abs_set(aperture_device.selected_aperture, new_selected_aperture)
+    yield from bps.abs_set(
+        aperture_device.selected_aperture, new_selected_aperture, group=group
+    )

mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py

@@ -22,7 +22,7 @@ from mx_bluesky.common.experiment_plans.inner_plans.do_fgs import (
     ZOCALO_STAGE_GROUP,
     kickoff_and_complete_gridscan,
 )
-from mx_bluesky.common.experiment_plans.read_hardware import (
+from mx_bluesky.common.experiment_plans.inner_plans.read_hardware import (
     read_hardware_plan,
 )
 from mx_bluesky.common.parameters.constants import (
@@ -55,6 +55,25 @@ class BeamlineSpecificFGSFeatures:
     get_xrc_results_from_zocalo: bool
 
 
+def generic_tidy(xrc_composite: FlyScanEssentialDevices, wait=True) -> MsgGenerator:
+    """Tidy Zocalo and turn off Eiger dev/shm. Ran after the beamline-specific tidy plan"""
+
+    LOGGER.info("Tidying up Zocalo")
+    group = "generic_tidy"
+    # make sure we don't consume any other results
+    yield from bps.unstage(xrc_composite.zocalo, group=group)
+
+    # Turn off dev/shm streaming to avoid filling disk, see https://github.com/DiamondLightSource/hyperion/issues/1395
+    LOGGER.info("Turning off Eiger dev/shm streaming")
+    # Fix types in ophyd-async (https://github.com/DiamondLightSource/mx-bluesky/issues/855)
+    yield from bps.abs_set(
+        xrc_composite.eiger.odin.fan.dev_shm_enable,  # type: ignore
+        0,
+        group=group,
+    )
+    yield from bps.wait(group)
+
+
 def construct_beamline_specific_FGS_features(
     setup_trigger_plan: Callable[..., MsgGenerator],
     tidy_plan: Callable[..., MsgGenerator],
@@ -71,7 +90,7 @@ def construct_beamline_specific_FGS_features(
         Ran directly before kicking off the gridscan.
 
         tidy_plan (Callable): Tidy up states of devices. Ran at the end of the flyscan, regardless of
-        whether or not it finished successfully.
+        whether or not it finished successfully. Zocalo and Eiger are cleaned up separately
 
         set_flyscan_params_plan (Callable): Set PV's for the relevant Fast Grid Scan dodal device
 
@@ -135,6 +154,10 @@ def common_flyscan_xray_centre(
     There are a few other useful decorators to use with this plan, see: verify_undulator_gap_before_run_decorator, transmission_and_xbpm_feedback_for_collection_decorator
     """
 
+    def _overall_tidy():
+        yield from beamline_specific.tidy_plan()
+        yield from generic_tidy(composite)
+
     def _decorated_flyscan():
         @bpp.set_run_key_decorator(PlanNameConstants.GRIDSCAN_OUTER)
         @bpp.run_decorator(  # attach experiment metadata to the start document
@@ -146,7 +169,7 @@ def common_flyscan_xray_centre(
                 ],
             }
         )
-        @bpp.finalize_decorator(lambda: beamline_specific.tidy_plan(composite))
+        @bpp.finalize_decorator(lambda: _overall_tidy())
         def run_gridscan_and_tidy(
             fgs_composite: FlyScanEssentialDevices,
             params: SpecifiedThreeDGridScan,
@@ -288,6 +311,7 @@ def _xrc_result_in_boxes_to_result_in_mm(
         ),
         max_count=xrc_result["max_count"],
         total_count=xrc_result["total_count"],
+        sample_id=xrc_result["sample_id"],
     )
 
 

mx_bluesky/common/experiment_plans/common_grid_detect_then_xray_centre_plan.py

@@ -111,6 +111,7 @@ def grid_detect_then_xray_centre(
     )
 
 
+# This function should be private but is currently called by Hyperion, see https://github.com/DiamondLightSource/mx-bluesky/issues/1148
 def detect_grid_and_do_gridscan(
     composite: GridDetectThenXRayCentreComposite,
     parameters: GridCommon,

mx_bluesky/common/experiment_plans/inner_plans/do_fgs.py

@@ -14,7 +14,9 @@ from dodal.log import LOGGER
 from dodal.plan_stubs.check_topup import check_topup_and_wait_if_necessary
 from scanspec.core import AxesPoints, Axis
 
-from mx_bluesky.common.experiment_plans.read_hardware import read_hardware_for_zocalo
+from mx_bluesky.common.experiment_plans.inner_plans.read_hardware import (
+    read_hardware_for_zocalo,
+)
 from mx_bluesky.common.parameters.constants import (
     PlanNameConstants,
 )

mx_bluesky/common/external_interaction/ispyb/exp_eye_store.py

@@ -1,7 +1,9 @@
 import configparser
 from dataclasses import dataclass
 from enum import StrEnum
+from typing import Any, Literal
 
+from event_model.documents import Event
 from requests import JSONDecodeError, patch, post
 from requests.auth import AuthBase
 
@@ -63,6 +65,19 @@ assert all(len(value) <= 20 for value in BLSampleStatus), (
 )
 
 
+def create_update_data_from_event_doc(
+    mapping: dict[str, str], event: Event
+) -> dict[str, Any]:
+    """Given a mapping between bluesky event data and an event itself this function will
+    create a dict that can be used to update exp-eye."""
+    event_data = event["data"]
+    return {
+        target_key: event_data[source_key]
+        for source_key, target_key in mapping.items()
+        if source_key in event_data
+    }
+
+
 class ExpeyeInteraction:
     """Exposes functionality from the Expeye core API"""
 
@@ -74,24 +89,22 @@ class ExpeyeInteraction:
         self._base_url = url
         self._auth = BearerAuth(token)
 
-    def start_load(
+    def start_robot_action(
         self,
+        action_type: Literal["LOAD", "UNLOAD"],
         proposal_reference: str,
         visit_number: int,
         sample_id: int,
-        dewar_location: int,
-        container_location: int,
     ) -> RobotActionID:
-        """Create a robot load entry in ispyb.
+        """Create a robot action entry in ispyb.
 
         Args:
+            action_type ("LOAD" | "UNLOAD"): The robot action being performed
             proposal_reference (str): The proposal of the experiment e.g. cm37235
             visit_number (int): The visit number for the proposal, usually this can be
                                 found added to the end of the proposal e.g. the data for
                                 visit number 2 of proposal cm37235 is in cm37235-2
             sample_id (int): The id of the sample in the database
-            dewar_location (int): Which puck in the dewar the sample is in
-            container_location (int): Which pin in that puck has the sample
 
         Returns:
             RobotActionID: The id of the robot load action that is created
@@ -102,39 +115,28 @@ class ExpeyeInteraction:
 
         data = {
             "startTimestamp": get_current_time_string(),
+            "actionType": action_type,
             "sampleId": sample_id,
-            "actionType": "LOAD",
-            "containerLocation": container_location,
-            "dewarLocation": dewar_location,
         }
         response = _send_and_get_response(self._auth, url, data, post)
         return response["robotActionId"]
 
-    def update_barcode_and_snapshots(
+    def update_robot_action(
         self,
         action_id: RobotActionID,
-        barcode: str,
-        snapshot_before_path: str,
-        snapshot_after_path: str,
+        data: dict[str, Any],
     ):
-        """Update the barcode and snapshots of an existing robot action.
+        """Update an existing robot action to contain additional info.
 
         Args:
             action_id (RobotActionID): The id of the action to update
-            barcode (str): The barcode to give the action
-            snapshot_before_path (str): Path to the snapshot before robot load
-            snapshot_after_path (str): Path to the snapshot after robot load
+            data (dict): The data to update with, where the keys match those expected
+                         by exp-eye.
         """
         url = self._base_url + self.UPDATE_ROBOT_ACTION.format(action_id=action_id)
-
-        data = {
-            "sampleBarcode": barcode,
-            "xtalSnapshotBefore": snapshot_before_path,
-            "xtalSnapshotAfter": snapshot_after_path,
-        }
         _send_and_get_response(self._auth, url, data, patch)
 
-    def end_load(self, action_id: RobotActionID, status: str, reason: str):
+    def end_robot_action(self, action_id: RobotActionID, status: str, reason: str):
         """Finish an existing robot action, providing final information about how it went
 
         Args:

mx_bluesky/common/external_interaction/ispyb/ispyb_store.py

@@ -41,7 +41,6 @@ class IspybIds(BaseModel):
 class StoreInIspyb:
     def __init__(self, ispyb_config: str) -> None:
         self.ISPYB_CONFIG_PATH: str = ispyb_config
-        self._data_collection_group_id: int | None
 
     def begin_deposition(
         self,
@@ -146,11 +145,17 @@ class StoreInIspyb:
     def append_to_comment(
         self, data_collection_id: int, comment: str, delimiter: str = " "
     ) -> None:
-        with ispyb.open(self.ISPYB_CONFIG_PATH) as conn:
-            assert conn is not None, "Failed to connect to ISPyB!"
-            mx_acquisition: MXAcquisition = conn.mx_acquisition
-            mx_acquisition.update_data_collection_append_comments(
-                data_collection_id, comment, delimiter
+        try:
+            with ispyb.open(self.ISPYB_CONFIG_PATH) as conn:
+                assert conn is not None, "Failed to connect to ISPyB!"
+                mx_acquisition: MXAcquisition = conn.mx_acquisition
+                mx_acquisition.update_data_collection_append_comments(
+                    data_collection_id, comment, delimiter
+                )
+        except ispyb.ReadWriteError as e:
+            ISPYB_ZOCALO_CALLBACK_LOGGER.warning(
+                f"Unable to log comment, comment probably exceeded column length: {comment}",
+                exc_info=e,
             )
 
     def update_data_collection_group_table(
@@ -187,7 +192,6 @@ class StoreInIspyb:
             params["parentid"] = data_collection_group_id
             params["endtime"] = end_time
             params["run_status"] = run_status
-
             mx_acquisition.upsert_data_collection(list(params.values()))
 
     def _store_position_table(

mx_bluesky/common/external_interaction/nexus/write_nexus.py

@@ -20,13 +20,13 @@ from mx_bluesky.common.external_interaction.nexus.nexus_utils import (
     create_goniometer_axes,
     get_start_and_predicted_end_time,
 )
-from mx_bluesky.common.parameters.components import DiffractionExperimentWithSample
+from mx_bluesky.common.parameters.components import DiffractionExperiment
 
 
 class NexusWriter:
     def __init__(
         self,
-        parameters: DiffractionExperimentWithSample,
+        parameters: DiffractionExperiment,
         data_shape: tuple[int, int, int],
         scan_points: AxesPoints,
         *,

mx_bluesky/common/parameters/components.py

@@ -238,9 +238,14 @@ class TopNByMaxCountSelection(MultiXtalSelection):
     n: int
 
 
+class TopNByMaxCountForEachSampleSelection(MultiXtalSelection):
+    name: Literal["TopNByMaxCountForEachSample"] = "TopNByMaxCountForEachSample"  #  pyright: ignore [reportIncompatibleVariableOverride]
+    n: int
+
+
 class WithCentreSelection(BaseModel):
-    select_centres: TopNByMaxCountSelection = Field(
-        discriminator="name", default=TopNByMaxCountSelection(n=1)
+    select_centres: TopNByMaxCountSelection | TopNByMaxCountForEachSampleSelection = (
+        Field(discriminator="name", default=TopNByMaxCountSelection(n=1))
     )
 
     @property

mx_bluesky/common/parameters/constants.py

@@ -11,12 +11,13 @@ from mx_bluesky.definitions import ROOT_DIR
 
 BEAMLINE = get_beamline_name("test")
 TEST_MODE = BEAMLINE == "test"
+ZEBRA_STATUS_TIMEOUT = 30
 
 
 @dataclass(frozen=True)
 class DocDescriptorNames:
-    # Robot load event descriptor
-    ROBOT_LOAD = "robot_load"
+    # Robot load/unload event descriptor
+    ROBOT_UPDATE = "robot_update"
     # For callbacks to use
     OAV_ROTATION_SNAPSHOT_TRIGGERED = "rotation_snapshot_triggered"
     OAV_GRID_SNAPSHOT_TRIGGERED = "snapshot_to_ispyb"
@@ -41,8 +42,9 @@ class OavConstants:
 @dataclass(frozen=True)
 class PlanNameConstants:
     LOAD_CENTRE_COLLECT = "load_centre_collect"
-    # Robot load subplan
+    # Robot subplans
     ROBOT_LOAD = "robot_load"
+    ROBOT_UNLOAD = "robot_unload"
     # Gridscan
     GRID_DETECT_AND_DO_GRIDSCAN = "grid_detect_and_do_gridscan"
     GRID_DETECT_INNER = "grid_detect"

mx_bluesky/common/parameters/device_composites.py

@@ -11,7 +11,7 @@ from dodal.devices.fast_grid_scan import (
     ZebraFastGridScan,
 )
 from dodal.devices.flux import Flux
-from dodal.devices.i03 import Beamstop
+from dodal.devices.mx_phase1.beamstop import Beamstop
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
 from dodal.devices.robot import BartRobot

mx_bluesky/common/parameters/gridscan.py

@@ -30,6 +30,7 @@ DETECTOR_SIZE_PER_BEAMLINE = {
     "i02-1": EIGER2_X_9M_SIZE,
     "dev": EIGER2_X_16M_SIZE,
     "i03": EIGER2_X_16M_SIZE,
+    "i04": EIGER2_X_16M_SIZE,
 }
 
 

mx_bluesky/common/xrc_result.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
+from collections import defaultdict
 from collections.abc import Callable, Sequence
 from functools import partial
 
@@ -10,6 +11,7 @@ from event_model import RunStart
 
 from mx_bluesky.common.parameters.components import (
     MultiXtalSelection,
+    TopNByMaxCountForEachSampleSelection,
     TopNByMaxCountSelection,
 )
 
@@ -39,18 +41,21 @@ class XRayCentreResult:
             containing the crystal
         max_count: The maximum spot count encountered in any one grid box in the crystal
         total_count: The total count across all boxes in the crystal.
+        sample_id: The sample id associated with the centre.
     """
 
     centre_of_mass_mm: np.ndarray
     bounding_box_mm: tuple[np.ndarray, np.ndarray]
     max_count: int
     total_count: int
+    sample_id: int | None
 
     def __eq__(self, o):
         return (
             isinstance(o, XRayCentreResult)
             and o.max_count == self.max_count
             and o.total_count == self.total_count
+            and o.sample_id == self.sample_id
             and all(o.centre_of_mass_mm == self.centre_of_mass_mm)
             and all(o.bounding_box_mm[0] == self.bounding_box_mm[0])
             and all(o.bounding_box_mm[1] == self.bounding_box_mm[1])
@@ -64,9 +69,27 @@ def top_n_by_max_count(
     return sorted_hits[:n]
 
 
+def top_n_by_max_count_for_each_sample(
+    unfiltered: Sequence[XRayCentreResult], n: int
+) -> Sequence[XRayCentreResult]:
+    xrc_results_by_sample_id: dict[int | None, list[XRayCentreResult]] = defaultdict(
+        list[XRayCentreResult]
+    )
+    for result in unfiltered:
+        xrc_results_by_sample_id[result.sample_id].append(result)
+    return [
+        result
+        for results in xrc_results_by_sample_id.values()
+        for result in sorted(results, key=lambda x: x.max_count, reverse=True)[:n]
+    ]
+
+
 def resolve_selection_fn(
     params: MultiXtalSelection,
 ) -> Callable[[Sequence[XRayCentreResult]], Sequence[XRayCentreResult]]:
-    if isinstance(params, TopNByMaxCountSelection):
-        return partial(top_n_by_max_count, n=params.n)
+    match params:
+        case TopNByMaxCountSelection():
+            return partial(top_n_by_max_count, n=params.n)
+        case TopNByMaxCountForEachSampleSelection():
+            return partial(top_n_by_max_count_for_each_sample, n=params.n)
     raise ValueError(f"Invalid selection function {params.name}")

mx_bluesky/hyperion/__main__.py

@@ -218,7 +218,7 @@ class RunExperiment(Resource):
                 status_and_message = self.runner.start(plan, params, plan_name)
             except Exception as e:
                 status_and_message = make_error_status_and_message(e)
-                LOGGER.error(format_exception(e))
+                LOGGER.error("".join(format_exception(e)))
 
         elif action == Actions.STOP.value:
             status_and_message = self.runner.stop()

mx_bluesky/hyperion/baton_handler.py

@@ -1,18 +1,26 @@
 from collections.abc import Sequence
 
+from blueapi.core import BlueskyContext
 from bluesky import plan_stubs as bps
 from bluesky import preprocessors as bpp
 from dodal.devices.baton import Baton
 
+from mx_bluesky.common.utils.context import find_device_in_context
 from mx_bluesky.common.utils.exceptions import WarningException
+from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.hyperion.experiment_plans.load_centre_collect_full_plan import (
     LoadCentreCollectComposite,
+    create_devices,
     load_centre_collect_full,
 )
 from mx_bluesky.hyperion.external_interaction.agamemnon import (
     create_parameters_from_agamemnon,
 )
 from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
+from mx_bluesky.hyperion.utils.context import (
+    clear_all_device_caches,
+    setup_devices,
+)
 
 HYPERION_USER = "Hyperion"
 NO_USER = "None"
@@ -59,7 +67,24 @@ def move_to_default_state():
     yield from bps.null()
 
 
-def run_udc_when_requested(baton: Baton, composite: LoadCentreCollectComposite):
+def initialise_udc(context: BlueskyContext, dev_mode: bool = False):
+    """
+    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.
+
+    Beamline devices are unloaded and reloaded in order to pick up any new configuration,
+    bluesky context gets new set of devices.
+    """
+    LOGGER.info("Initialising mx-bluesky for UDC start...")
+    clear_all_device_caches(context)
+    setup_devices(context, dev_mode)
+
+
+def _get_baton(context: BlueskyContext) -> Baton:
+    return find_device_in_context(context, "baton", Baton)
+
+
+def run_udc_when_requested(context: BlueskyContext, dev_mode: bool = False):
     """This will wait for the baton to be handed to hyperion and then run through the
     UDC queue from agamemnon until:
       1. There are no more instructions from agamemnon
@@ -69,21 +94,28 @@ def run_udc_when_requested(baton: Baton, composite: LoadCentreCollectComposite):
     In the case of 1. or 2. hyperion will immediately release the baton. In the case of
     3. the baton will be released after the next collection has finished."""
 
+    baton = _get_baton(context)
     yield from wait_for_hyperion_requested(baton)
     yield from bps.abs_set(baton.current_user, HYPERION_USER)
 
-    def default_state_then_collect():
+    def initialise_then_collect():
+        initialise_udc(context, dev_mode)
         yield from move_to_default_state()
-        yield from main_hyperion_loop(baton, composite)
+
+        # re-fetch the baton because the device has been reinstantiated
+        new_baton = _get_baton(context)
+        composite = create_devices(context)
+        yield from main_hyperion_loop(new_baton, composite)
 
     def release_baton():
         # 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)
         requested_user = yield from bps.rd(baton.requested_user)
         if requested_user == HYPERION_USER:
             yield from bps.abs_set(baton.requested_user, NO_USER)
         yield from bps.abs_set(baton.current_user, NO_USER)
 
     yield from bpp.contingency_wrapper(
-        default_state_then_collect(), final_plan=release_baton
+        initialise_then_collect(), final_plan=release_baton
     )