mx-bluesky 1.5.14__py3-none-any.whl → 1.5.16__py3-none-any.whl

mx_bluesky/hyperion/experiment_plans/load_centre_collect_full_plan.py

@@ -13,6 +13,9 @@ from dodal.devices.oav.oav_parameters import OAVParameters
 
 import mx_bluesky.common.xrc_result as flyscan_result
 from mx_bluesky.common.parameters.components import WithSnapshot
+from mx_bluesky.common.parameters.rotation import (
+    RotationScanPerSweep,
+)
 from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.exceptions import CrystalNotFoundError
 from mx_bluesky.common.utils.log import LOGGER
@@ -31,7 +34,6 @@ from mx_bluesky.hyperion.external_interaction.config_server import (
 )
 from mx_bluesky.hyperion.parameters.constants import CONST, I03Constants
 from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
-from mx_bluesky.hyperion.parameters.rotation import RotationScanPerSweep
 
 
 @pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})

mx_bluesky/hyperion/experiment_plans/pin_centre_then_xray_centre_plan.py

@@ -28,7 +28,10 @@ from mx_bluesky.common.experiment_plans.pin_tip_centring_plan import (
 from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
     ispyb_activation_wrapper,
 )
-from mx_bluesky.common.parameters.constants import OavConstants
+from mx_bluesky.common.parameters.constants import OavConstants, PlanNameConstants
+from mx_bluesky.common.preprocessors.preprocessors import (
+    pause_xbpm_feedback_during_collection_at_desired_transmission_decorator,
+)
 from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.common.xrc_result import XRayCentreEventHandler
@@ -103,13 +106,21 @@ def pin_centre_then_flyscan_plan(
         grid_detect_params = create_parameters_for_grid_detection(parameters)
         oav_params = OAVParameters("xrayCentring", oav_config_file)
 
-        yield from detect_grid_and_do_gridscan(
+        @pause_xbpm_feedback_during_collection_at_desired_transmission_decorator(
             composite,
-            grid_detect_params,
-            oav_params,
-            HyperionSpecifiedThreeDGridScan,
-            construct_hyperion_specific_features,
+            parameters.transmission_frac,
+            PlanNameConstants.GRIDSCAN_OUTER,
         )
+        def _grid_detect_plan():
+            yield from detect_grid_and_do_gridscan(
+                composite,
+                grid_detect_params,
+                oav_params,
+                HyperionSpecifiedThreeDGridScan,
+                construct_hyperion_specific_features,
+            )
+
+        yield from _grid_detect_plan()
 
     yield from ispyb_activation_wrapper(_pin_centre_then_flyscan_plan(), parameters)
 

mx_bluesky/hyperion/experiment_plans/robot_load_and_change_energy.py

@@ -27,7 +27,6 @@ from dodal.devices.xbpm_feedback import XBPMFeedback
 from mx_bluesky.common.device_setup_plans.robot_load_unload import (
     do_plan_while_lower_gonio_at_home,
     prepare_for_robot_load,
-    wait_for_smargon_not_disabled,
 )
 from mx_bluesky.hyperion.experiment_plans.set_energy_plan import (
     SetEnergyComposite,
@@ -95,8 +94,6 @@ def do_robot_load(
 
     yield from bps.wait("robot_load")
 
-    yield from wait_for_smargon_not_disabled(composite.smargon)
-
     yield from bps.mv(composite.thawer, OnOff.ON)
 
 

mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py

@@ -1,7 +1,3 @@
-from __future__ import annotations
-
-import dataclasses
-
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
 import pydantic
@@ -25,7 +21,7 @@ from dodal.devices.synchrotron import Synchrotron
 from dodal.devices.thawer import Thawer
 from dodal.devices.undulator import UndulatorInKeV
 from dodal.devices.xbpm_feedback import XBPMFeedback
-from dodal.devices.zebra.zebra import RotationDirection, Zebra
+from dodal.devices.zebra.zebra import Zebra
 from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutter
 from dodal.plan_stubs.check_topup import check_topup_and_wait_if_necessary
 from dodal.plans.preprocessors.verify_undulator_gap import (
@@ -53,20 +49,24 @@ from mx_bluesky.common.experiment_plans.oav_snapshot_plan import (
     oav_snapshot_plan,
     setup_beamline_for_oav,
 )
+from mx_bluesky.common.experiment_plans.rotation.rotation_utils import (
+    RotationMotionProfile,
+    calculate_motion_profile,
+)
 from mx_bluesky.common.parameters.components import WithSnapshot
+from mx_bluesky.common.parameters.rotation import (
+    RotationScan,
+    SingleRotationScan,
+)
 from mx_bluesky.common.preprocessors.preprocessors import (
-    transmission_and_xbpm_feedback_for_collection_decorator,
+    pause_xbpm_feedback_during_collection_at_desired_transmission_decorator,
 )
 from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.hyperion.device_setup_plans.setup_zebra import (
     arm_zebra,
 )
-from mx_bluesky.hyperion.parameters.constants import CONST, I03Constants
-from mx_bluesky.hyperion.parameters.rotation import (
-    RotationScan,
-    SingleRotationScan,
-)
+from mx_bluesky.hyperion.parameters.constants import CONST
 
 
 @pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
@@ -100,120 +100,6 @@ def create_devices(context: BlueskyContext) -> RotationScanComposite:
     return device_composite_from_context(context, RotationScanComposite)
 
 
-DEFAULT_DIRECTION = RotationDirection.NEGATIVE
-DEFAULT_MAX_VELOCITY = 120
-# Use a slightly larger time to acceleration than EPICS as it's better to be cautious
-ACCELERATION_MARGIN = 1.5
-
-
-@dataclasses.dataclass
-class RotationMotionProfile:
-    start_scan_deg: float
-    start_motion_deg: float
-    scan_width_deg: float
-    shutter_time_s: float
-    direction: RotationDirection
-    speed_for_rotation_deg_s: float
-    acceleration_offset_deg: float
-    shutter_opening_deg: float
-    total_exposure_s: float
-    distance_to_move_deg: float
-    max_velocity_deg_s: float
-
-
-def calculate_motion_profile(
-    params: SingleRotationScan,
-    motor_time_to_speed_s: float,
-    max_velocity_deg_s: float,
-) -> RotationMotionProfile:
-    """Calculates the various numbers needed for motions in the rotation scan.
-    Rotates through "scan width" plus twice an "offset" to take into account
-    acceleration at the start and deceleration at the end, plus the number of extra
-    degrees of rotation needed to make sure the fast shutter has fully opened before the
-    detector trigger is sent.
-    See https://github.com/DiamondLightSource/hyperion/wiki/rotation-scan-geometry
-    for a simple pictorial explanation."""
-
-    assert params.rotation_increment_deg > 0
-
-    direction = params.rotation_direction
-    start_scan_deg = params.omega_start_deg
-
-    if I03Constants.OMEGA_FLIP:
-        # If omega_flip is True then the motor omega axis is inverted with respect to the
-        # hyperion coordinate system.
-        start_scan_deg = -start_scan_deg
-        direction = (
-            direction.POSITIVE
-            if direction == direction.NEGATIVE
-            else direction.NEGATIVE
-        )
-
-    num_images = params.num_images
-    shutter_time_s = params.shutter_opening_time_s
-    image_width_deg = params.rotation_increment_deg
-    exposure_time_s = params.exposure_time_s
-    motor_time_to_speed_s *= ACCELERATION_MARGIN
-
-    LOGGER.info("Calculating rotation scan motion profile:")
-    LOGGER.info(
-        f"{num_images=}, {shutter_time_s=}, {image_width_deg=}, {exposure_time_s=}, {direction=}"
-    )
-
-    scan_width_deg = num_images * params.rotation_increment_deg
-    LOGGER.info(f"{scan_width_deg=} = {num_images=} * {params.rotation_increment_deg=}")
-
-    speed_for_rotation_deg_s = image_width_deg / exposure_time_s
-    LOGGER.info("speed_for_rotation_deg_s = image_width_deg / exposure_time_s")
-    LOGGER.info(
-        f"{speed_for_rotation_deg_s=} = {image_width_deg=} / {exposure_time_s=}"
-    )
-
-    acceleration_offset_deg = motor_time_to_speed_s * speed_for_rotation_deg_s
-    LOGGER.info(
-        f"{acceleration_offset_deg=} = {motor_time_to_speed_s=} * {speed_for_rotation_deg_s=}"
-    )
-
-    start_motion_deg = start_scan_deg - (acceleration_offset_deg * direction.multiplier)
-    LOGGER.info(
-        f"{start_motion_deg=} = {start_scan_deg=} - ({acceleration_offset_deg=} * {direction.multiplier=})"
-    )
-
-    shutter_opening_deg = speed_for_rotation_deg_s * shutter_time_s
-    LOGGER.info(
-        f"{shutter_opening_deg=} = {speed_for_rotation_deg_s=} * {shutter_time_s=}"
-    )
-
-    shutter_opening_deg = speed_for_rotation_deg_s * shutter_time_s
-    LOGGER.info(
-        f"{shutter_opening_deg=} = {speed_for_rotation_deg_s=} * {shutter_time_s=}"
-    )
-
-    total_exposure_s = num_images * exposure_time_s
-    LOGGER.info(f"{total_exposure_s=} = {num_images=} * {exposure_time_s=}")
-
-    distance_to_move_deg = (
-        scan_width_deg + shutter_opening_deg + acceleration_offset_deg * 2
-    ) * direction.multiplier
-    LOGGER.info(
-        f"{distance_to_move_deg=} = ({scan_width_deg=} + {shutter_opening_deg=} + {acceleration_offset_deg=} * 2) * {direction=})"
-    )
-
-    return RotationMotionProfile(
-        start_scan_deg=start_scan_deg,
-        start_motion_deg=start_motion_deg,
-        scan_width_deg=scan_width_deg,
-        shutter_time_s=shutter_time_s,
-        direction=direction,
-        speed_for_rotation_deg_s=speed_for_rotation_deg_s,
-        acceleration_offset_deg=acceleration_offset_deg,
-        shutter_opening_deg=shutter_opening_deg,
-        total_exposure_s=total_exposure_s,
-        distance_to_move_deg=distance_to_move_deg,
-        max_velocity_deg_s=max_velocity_deg_s,
-    )
-
-
 def rotation_scan_plan(
     composite: RotationScanComposite,
     params: SingleRotationScan,
@@ -399,7 +285,7 @@ def rotation_scan_internal(
     eiger: EigerDetector = composite.eiger
     eiger.set_detector_parameters(parameters.detector_params)
 
-    @transmission_and_xbpm_feedback_for_collection_decorator(
+    @pause_xbpm_feedback_during_collection_at_desired_transmission_decorator(
         composite,
         parameters.transmission_frac,
     )

mx_bluesky/hyperion/experiment_plans/set_energy_plan.py

@@ -17,7 +17,7 @@ from dodal.devices.xbpm_feedback import XBPMFeedback
 
 from mx_bluesky.common.parameters.constants import PlanNameConstants
 from mx_bluesky.common.preprocessors.preprocessors import (
-    transmission_and_xbpm_feedback_for_collection_wrapper,
+    pause_xbpm_feedback_during_collection_at_desired_transmission_wrapper,
 )
 from mx_bluesky.hyperion.device_setup_plans import dcm_pitch_roll_mirror_adjuster
 
@@ -74,7 +74,7 @@ def set_energy_plan(
     )
 
     if energy_ev:
-        yield from transmission_and_xbpm_feedback_for_collection_wrapper(
+        yield from pause_xbpm_feedback_during_collection_at_desired_transmission_wrapper(
             _set_energy_plan(energy_ev / 1000, composite),
             composite_for_wrapper,
             DESIRED_TRANSMISSION_FRACTION,

mx_bluesky/hyperion/experiment_plans/udc_default_state.py

@@ -17,6 +17,7 @@ from dodal.devices.fluorescence_detector_motion import FluorescenceDetector
 from dodal.devices.fluorescence_detector_motion import InOut as FlouInOut
 from dodal.devices.hutch_shutter import HutchShutter, ShutterDemand
 from dodal.devices.mx_phase1.beamstop import BeamstopPositions
+from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.robot import BartRobot, PinMounted
 from dodal.devices.scintillator import InOut as ScinInOut
 from dodal.devices.scintillator import Scintillator
@@ -49,6 +50,7 @@ class UDCDefaultDevices(BeamstopCheckDevices):
     robot: BartRobot
     scintillator: Scintillator
     smargon: Smargon
+    oav: OAV
 
 
 class UnexpectedSampleError(BeamlineCheckFailureError): ...
@@ -109,7 +111,7 @@ def move_to_udc_default_state(devices: UDCDefaultDevices):
     )
 
     # Wait for all of the above to complete
-    yield from bps.wait(group=_GROUP_PRE_BEAMSTOP_CHECK, timeout=0.1)
+    yield from bps.wait(group=_GROUP_PRE_BEAMSTOP_CHECK, timeout=10)
 
     feature_flags: HyperionFeatureSettings = (
         get_hyperion_config_client().get_feature_flags()
@@ -141,7 +143,11 @@ def move_to_udc_default_state(devices: UDCDefaultDevices):
         devices.cryojet.fine, CryoInOut.IN, group=_GROUP_POST_BEAMSTOP_CHECK
     )
 
-    yield from bps.wait(_GROUP_POST_BEAMSTOP_CHECK)
+    yield from bps.abs_set(
+        devices.oav.zoom_controller, "1.0x", group=_GROUP_POST_BEAMSTOP_CHECK
+    )
+
+    yield from bps.wait(_GROUP_POST_BEAMSTOP_CHECK, timeout=10)
 
 
 def _verify_correct_cryostream_selected(

mx_bluesky/hyperion/external_interaction/agamemnon.py

@@ -12,12 +12,11 @@ import requests
 from deepdiff.diff import DeepDiff
 from dodal.utils import get_beamline_name
 from jsonschema import ValidationError
-from pydantic_extra_types.semantic_version import SemanticVersion
 
 from mx_bluesky.common.parameters.components import (
-    PARAMETER_VERSION,
     MxBlueskyParameters,
     WithVisit,
+    get_param_version,
 )
 from mx_bluesky.common.parameters.constants import (
     GridscanParamConstants,
@@ -88,7 +87,7 @@ def create_parameters_from_agamemnon() -> Sequence[MxBlueskyParameters]:
                     Wait.model_validate(
                         {
                             "duration_s": data,
-                            "parameter_model_version": _get_param_version(),
+                            "parameter_model_version": get_param_version(),
                         }
                     )
                 ]
@@ -228,10 +227,6 @@ def _get_withenergy_parameters_from_agamemnon(parameters: dict) -> dict[str, Any
         return {"demand_energy_ev": None}
 
 
-def _get_param_version() -> SemanticVersion:
-    return SemanticVersion.validate_from_str(str(PARAMETER_VERSION))
-
-
 def _populate_parameters_from_agamemnon(
     agamemnon_params,
 ) -> Sequence[LoadCentreCollect]:
@@ -250,7 +245,7 @@ def _populate_parameters_from_agamemnon(
     return [
         LoadCentreCollect.model_validate(
             {
-                "parameter_model_version": _get_param_version(),
+                "parameter_model_version": get_param_version(),
                 "visit": visit,
                 "detector_distance_mm": detector_distance,
                 "sample_id": agamemnon_params["sample"]["id"],

mx_bluesky/hyperion/external_interaction/callbacks/__main__.py

@@ -1,12 +1,17 @@
 import logging
-from collections.abc import Callable, Sequence
+from abc import abstractmethod
+from collections.abc import Callable
+from contextlib import AbstractContextManager
 from threading import Thread
 from time import sleep  # noqa
 from urllib import request
 from urllib.error import URLError
 
+from blueapi.config import ApplicationConfig, ConfigLoader
 from bluesky.callbacks import CallbackBase
 from bluesky.callbacks.zmq import Proxy, RemoteDispatcher
+from bluesky_stomp.messaging import StompClient
+from bluesky_stomp.models import Broker
 from dodal.log import LOGGER as DODAL_LOGGER
 from dodal.log import set_up_all_logging_handlers
 
@@ -52,8 +57,11 @@ from mx_bluesky.hyperion.external_interaction.callbacks.rotation.nexus_callback
 from mx_bluesky.hyperion.external_interaction.callbacks.snapshot_callback import (
     BeamDrawingCallback,
 )
-from mx_bluesky.hyperion.parameters.cli import parse_callback_dev_mode_arg
-from mx_bluesky.hyperion.parameters.constants import CONST, HyperionConstants
+from mx_bluesky.hyperion.external_interaction.callbacks.stomp.dispatcher import (
+    StompDispatcher,
+)
+from mx_bluesky.hyperion.parameters.cli import CallbackArgs, parse_callback_args
+from mx_bluesky.hyperion.parameters.constants import CONST
 from mx_bluesky.hyperion.parameters.gridscan import (
     GridCommonWithHyperionDetectorParams,
     HyperionSpecifiedThreeDGridScan,
@@ -143,29 +151,72 @@ def log_debug(msg, *args, **kwargs):
     NEXUS_LOGGER.debug(msg, *args, **kwargs)
 
 
-def wait_for_threads_forever(threads: Sequence[Thread]):
-    alive = [t.is_alive() for t in threads]
-    try:
-        log_debug("Trying to wait forever on callback and dispatcher threads")
-        while all(alive):
-            sleep(LIVENESS_POLL_SECONDS)
-            alive = [t.is_alive() for t in threads]
-    except KeyboardInterrupt:
-        log_info("Main thread received interrupt - exiting.")
-    else:
-        log_info("Proxy or dispatcher thread ended - exiting.")
-
-
 class HyperionCallbackRunner:
     """Runs Nexus, ISPyB and Zocalo callbacks in their own process."""
 
-    def __init__(self, dev_mode) -> None:
-        setup_logging(dev_mode)
+    def __init__(self, callback_args: CallbackArgs) -> None:
+        setup_logging(callback_args.dev_mode)
         log_info("Hyperion callback process started.")
         set_alerting_service(LoggingAlertService(CONST.GRAYLOG_STREAM_ID))
 
         self.callbacks = setup_callbacks()
 
+        self.watchdog_thread = Thread(
+            target=run_watchdog,
+            daemon=True,
+            name="Watchdog",
+            args=[callback_args.watchdog_port],
+        )
+
+        self._dispatcher_cm: DispatcherContextMgr
+        if callback_args.stomp_config:
+            self._dispatcher_cm = StompDispatcherContextMgr(
+                callback_args, self.callbacks
+            )
+        else:
+            self._dispatcher_cm = RemoteDispatcherContextMgr(self.callbacks)
+
+    def start(self):
+        log_info(f"Launching threads, with callbacks: {self.callbacks}")
+        self.watchdog_thread.start()
+        with self._dispatcher_cm:
+            ping_watchdog_while_alive(self._dispatcher_cm, self.watchdog_thread)
+
+
+def run_watchdog(watchdog_port: int):
+    log_info("Hyperion watchdog keepalive running")
+    while True:
+        try:
+            with request.urlopen(
+                f"http://localhost:{watchdog_port}/callbackPing",
+                timeout=PING_TIMEOUT_S,
+            ) as response:
+                if response.status != 200:
+                    log_debug(
+                        f"Unable to ping Hyperion liveness endpoint, status {response.status}"
+                    )
+        except URLError as e:
+            log_debug("Unable to ping Hyperion liveness endpoint", exc_info=e)
+        sleep(HYPERION_PING_INTERVAL_S)
+
+
+def main(dev_mode=False) -> None:
+    callback_args = parse_callback_args()
+    callback_args.dev_mode = dev_mode or callback_args.dev_mode
+    print(f"In dev mode: {dev_mode}")
+    runner = HyperionCallbackRunner(callback_args)
+    runner.start()
+
+
+class DispatcherContextMgr(AbstractContextManager):
+    @abstractmethod
+    def is_alive(self) -> bool: ...
+
+
+class RemoteDispatcherContextMgr(DispatcherContextMgr):
+    def __init__(self, callbacks: list[CallbackBase]):
+        super().__init__()
+
         self.proxy = Proxy(*CONST.CALLBACK_0MQ_PROXY_PORTS)
         self.proxy_thread = Thread(
             target=self.proxy.start, daemon=True, name="0MQ Proxy"
@@ -182,47 +233,70 @@ class HyperionCallbackRunner:
 
         self.dispatcher_thread = Thread(
             target=start_dispatcher,
-            args=[self.callbacks],
+            args=[callbacks],
             daemon=True,
             name="0MQ Dispatcher",
         )
-
-        self.watchdog_thread = Thread(target=run_watchdog, daemon=True, name="Watchdog")
         log_info("Created 0MQ proxy and local RemoteDispatcher.")
 
-    def start(self):
-        log_info(f"Launching threads, with callbacks: {self.callbacks}")
+    def __enter__(self):
+        log_info("Proxy and dispatcher thread launched.")
         self.proxy_thread.start()
         self.dispatcher_thread.start()
-        self.watchdog_thread.start()
-        log_info("Proxy and dispatcher thread launched.")
-        wait_for_threads_forever(
-            [self.proxy_thread, self.dispatcher_thread, self.watchdog_thread]
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback, /):
+        self.dispatcher.stop()
+        # proxy has no way to stop
+
+    def is_alive(self):
+        return self.proxy_thread.is_alive() and self.dispatcher_thread.is_alive()
+
+
+class StompDispatcherContextMgr(DispatcherContextMgr):
+    def __init__(self, args: CallbackArgs, callbacks: list[CallbackBase]):
+        super().__init__()
+        loader = ConfigLoader(ApplicationConfig)
+        loader.use_values_from_yaml(args.stomp_config)
+        config = loader.load()
+        log_info(
+            f"Stomp client configured on {config.stomp.url.host}:{config.stomp.url.port}"
+        )
+        self._stomp_client = StompClient.for_broker(
+            broker=Broker(
+                host=config.stomp.url.host,
+                port=config.stomp.url.port,
+                auth=config.stomp.auth,
+            )
         )
+        self._dispatcher = StompDispatcher(self._stomp_client)
+        for cb in callbacks:
+            self._dispatcher.subscribe(cb)
 
+    def is_alive(self) -> bool:
+        return self._stomp_client.is_connected()
 
-def run_watchdog():
-    log_info("Hyperion watchdog keepalive running")
-    while True:
-        try:
-            with request.urlopen(
-                f"http://localhost:{HyperionConstants.HYPERION_PORT}/callbackPing",
-                timeout=PING_TIMEOUT_S,
-            ) as response:
-                if response.status != 200:
-                    log_debug(
-                        f"Unable to ping Hyperion liveness endpoint, status {response.status}"
-                    )
-        except URLError as e:
-            log_debug("Unable to ping Hyperion liveness endpoint", exc_info=e)
-        sleep(HYPERION_PING_INTERVAL_S)
+    def __enter__(self):
+        self._dispatcher.__enter__()
+        return self
 
+    def __exit__(self, exc_type, exc_value, traceback, /):
+        self._dispatcher.__exit__(exc_type, exc_value, traceback)
 
-def main(dev_mode=False) -> None:
-    dev_mode = dev_mode or parse_callback_dev_mode_arg()
-    print(f"In dev mode: {dev_mode}")
-    runner = HyperionCallbackRunner(dev_mode)
-    runner.start()
+
+def ping_watchdog_while_alive(
+    dispatcher_cm: DispatcherContextMgr, watchdog_thread: Thread
+):
+    alive = watchdog_thread.is_alive() and dispatcher_cm.is_alive()
+    try:
+        log_debug("Trying to wait forever on callback and dispatcher threads")
+        while alive:
+            sleep(LIVENESS_POLL_SECONDS)
+            alive = watchdog_thread.is_alive() and dispatcher_cm.is_alive()
+    except KeyboardInterrupt:
+        log_info("Main thread received interrupt - exiting.")
+    else:
+        log_info("Proxy or dispatcher thread ended - exiting.")
 
 
 if __name__ == "__main__":

mx_bluesky/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py

@@ -25,13 +25,15 @@ from mx_bluesky.common.external_interaction.ispyb.ispyb_store import (
     StoreInIspyb,
 )
 from mx_bluesky.common.parameters.components import IspybExperimentType
+from mx_bluesky.common.parameters.rotation import (
+    SingleRotationScan,
+)
 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
 from mx_bluesky.hyperion.external_interaction.callbacks.rotation.ispyb_mapping import (
     populate_data_collection_info_for_rotation,
 )
 from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.rotation import SingleRotationScan
 
 if TYPE_CHECKING:
     from event_model.documents import Event, RunStart, RunStop

mx_bluesky/hyperion/external_interaction/callbacks/rotation/ispyb_mapping.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 from mx_bluesky.common.external_interaction.ispyb.data_model import DataCollectionInfo
-from mx_bluesky.hyperion.parameters.rotation import SingleRotationScan
+from mx_bluesky.common.parameters.rotation import (
+    SingleRotationScan,
+)
 
 
 def populate_data_collection_info_for_rotation(params: SingleRotationScan):

mx_bluesky/hyperion/external_interaction/callbacks/rotation/nexus_callback.py

@@ -11,9 +11,12 @@ from mx_bluesky.common.external_interaction.nexus.nexus_utils import (
     vds_type_based_on_bit_depth,
 )
 from mx_bluesky.common.external_interaction.nexus.write_nexus import NexusWriter
+from mx_bluesky.common.parameters.constants import RotationParamConstants
+from mx_bluesky.common.parameters.rotation import (
+    SingleRotationScan,
+)
 from mx_bluesky.common.utils.log import NEXUS_LOGGER, format_doc_for_log
-from mx_bluesky.hyperion.parameters.constants import CONST, I03Constants
-from mx_bluesky.hyperion.parameters.rotation import SingleRotationScan
+from mx_bluesky.hyperion.parameters.constants import CONST
 
 if TYPE_CHECKING:
     from event_model.documents import Event, EventDescriptor, RunStart
@@ -100,6 +103,6 @@ class RotationNexusFileCallback(PlanReactiveCallback):
                 full_num_of_images=self.full_num_of_images,
                 meta_data_run_number=self.meta_data_run_number,
                 axis_direction=AxisDirection.NEGATIVE
-                if I03Constants.OMEGA_FLIP
+                if RotationParamConstants.OMEGA_FLIP
                 else AxisDirection.POSITIVE,
             )

mx_bluesky/hyperion/external_interaction/callbacks/stomp/dispatcher.py

@@ -0,0 +1,33 @@
+from blueapi.client.event_bus import AnyEvent
+from blueapi.core import DataEvent
+from bluesky.run_engine import Dispatcher
+from bluesky_stomp.messaging import MessageContext, StompClient
+from bluesky_stomp.models import MessageTopic
+from event_model import DocumentNames
+
+from mx_bluesky.common.utils.log import ISPYB_ZOCALO_CALLBACK_LOGGER as LOGGER
+
+BLUEAPI_EVENT_TOPIC = "public.worker.event"
+
+
+class StompDispatcher(Dispatcher):
+    def __init__(self, stomp_client: StompClient):
+        super().__init__()
+        self._client = stomp_client
+
+    def __enter__(self):
+        self._subscription_id = self._client.subscribe(
+            MessageTopic(name=BLUEAPI_EVENT_TOPIC), self._on_event
+        )
+        LOGGER.info("Connecting to stomp broker...")
+        self._client.connect()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        LOGGER.info("Disconnecting from stomp and unsubscribing...")
+        self._client.disconnect()
+        self._client.unsubscribe(self._subscription_id)
+
+    def _on_event(self, event: AnyEvent, context: MessageContext):
+        match event:
+            case DataEvent(name=name, doc=doc):  # type: ignore
+                self.process(DocumentNames[name], doc)