From 85fb41aa594e3476725281457f2b749f07e70f9f Mon Sep 17 00:00:00 2001 From: David Perl Date: Thu, 6 Aug 2026 14:32:54 +0200 Subject: [PATCH 1/8] WIP big becclient refactor --- .../default/beamline_dispatch.py | 42 +++++- src/aare/beamline_dispatch/protocols.py | 19 ++- .../simulated/beamline_dispatch.py | 5 +- .../x06da/beamline_dispatch.py | 6 +- .../x10sa/beamline_dispatch.py | 6 +- src/aare/daq/devices.py | 82 +---------- src/aare/devices/bec_worker.py | 131 ++++++------------ 7 files changed, 118 insertions(+), 173 deletions(-) diff --git a/src/aare/beamline_dispatch/default/beamline_dispatch.py b/src/aare/beamline_dispatch/default/beamline_dispatch.py index c432da92..500dd6cf 100644 --- a/src/aare/beamline_dispatch/default/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/default/beamline_dispatch.py @@ -1,6 +1,8 @@ import os -from aare.beamline_dispatch.protocols import AuthDispatch, BeamlineDispatch +from aarecommon.config.beamline import MXBeamline + +from aare.beamline_dispatch.protocols import AuthDispatch, BeamlineDispatch, ConfigKeys class DefaultAuthDispatch(AuthDispatch): @@ -12,13 +14,49 @@ class DefaultAuthDispatch(AuthDispatch): return key +class DefaultConfigKeys(ConfigKeys): + def __init__(self, beamline: MXBeamline): + self._beamline = beamline + + @property + def bl_lower(self) -> str: + return self._beamline.lower() + + @property + def bl_upper(self) -> str: + return self._beamline.upper() + + @property + def bl_name(self) -> str: + match self._beamline: + case MXBeamline.X06SA: + return "pxi" + case MXBeamline.X10SA: + return "pxii" + case MXBeamline.X06DA: + return "pxiii" + case MXBeamline.SIMULATED: + return "simulated" + case _: + raise ValueError(f"Beamline {self._beamline} has no common name!") + + @property + def device_yml_path(self) -> str: + return f"/sls/{self.bl_lower}/config/bec/production/{self.bl_name}_bec/{self.bl_name}_bec/device_configs/{self.bl_name}-devices.yaml" + + class DefaultDispatch(BeamlineDispatch): """Default implementation for anything which can vary between beamlines and/or simulation. Should be safe and fail rather than assuming anything.""" - def __init__(self) -> None: + def __init__(self, beamline: MXBeamline) -> None: self._auth = DefaultAuthDispatch() + self._config = DefaultConfigKeys(beamline) @property def auth(self): return self._auth + + @property + def config(self): + return self._config diff --git a/src/aare/beamline_dispatch/protocols.py b/src/aare/beamline_dispatch/protocols.py index d0675a7a..e562c9d4 100644 --- a/src/aare/beamline_dispatch/protocols.py +++ b/src/aare/beamline_dispatch/protocols.py @@ -22,12 +22,26 @@ class BecMacros(ABC): def bl_energy(energy_ev, move_gap=True, mono_scan=True, plot=True): ... @staticmethod @abstractmethod - def get_current_energy(): ... + def get_current_energy() -> float: ... @staticmethod @abstractmethod def mono_pitch_scan(plot=True): ... +class ConfigKeys(ABC): + @property + @abstractmethod + def bl_lower(self) -> str: ... + @property + @abstractmethod + def bl_upper(self) -> str: ... + @property + @abstractmethod + def bl_name(self) -> str: ... + @property + def device_yml_path(self) -> str: ... + + class BeamlineDispatch(ABC): @property @abstractmethod @@ -35,3 +49,6 @@ class BeamlineDispatch(ABC): @property @abstractmethod def bec_macros(self) -> BecMacros: ... + @property + @abstractmethod + def config(self) -> ConfigKeys: ... diff --git a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py index 59fcbd14..b023bf79 100644 --- a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py @@ -1,6 +1,8 @@ from typing import Any -from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch +from aarecommon.config.beamline import MXBeamline + +from aare.beamline_dispatch.default.beamline_dispatch import DefaultConfigKeys, DefaultDispatch from aare.beamline_dispatch.protocols import AuthDispatch, BecMacros @@ -32,6 +34,7 @@ class SimulatedDispatch(DefaultDispatch): def __init__(self) -> None: self._auth = SimulatedAuthDispatch() self._bec_macros = SimulatedBecMacros() + self._config = DefaultConfigKeys(MXBeamline.SIMULATED) @property def bec_macros(self): diff --git a/src/aare/beamline_dispatch/x06da/beamline_dispatch.py b/src/aare/beamline_dispatch/x06da/beamline_dispatch.py index da0dbb87..1ef2578a 100644 --- a/src/aare/beamline_dispatch/x06da/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x06da/beamline_dispatch.py @@ -1,5 +1,7 @@ from typing import Any +from aarecommon.config.beamline import MXBeamline + from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch from aare.beamline_dispatch.protocols import BecMacros @@ -45,8 +47,8 @@ class X06daBecMacros(BecMacros): class X06daDispatch(DefaultDispatch): - def __init__(self) -> None: - super().__init__() + def __init__(self, beamline: MXBeamline) -> None: + super().__init__(beamline) self._bec_macros = X06daBecMacros() @property diff --git a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py index 380513cb..6ef814dd 100644 --- a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py @@ -1,5 +1,7 @@ from typing import Any +from aarecommon.config.beamline import MXBeamline + from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch from aare.beamline_dispatch.protocols import BecMacros @@ -39,8 +41,8 @@ class X10SaBecMacros(BecMacros): class X10saDispatch(DefaultDispatch): - def __init__(self) -> None: - super().__init__() + def __init__(self, beamline: MXBeamline) -> None: + super().__init__(beamline) self._bec_macros = X10SaBecMacros() @property diff --git a/src/aare/daq/devices.py b/src/aare/daq/devices.py index 252d5121..cf4ec4be 100644 --- a/src/aare/daq/devices.py +++ b/src/aare/daq/devices.py @@ -11,16 +11,12 @@ from aarecommon.config.logger_events import log_timing from aarecommon.math.coordinate import AerotechCoordinate, SmargonCoordinate from aarecommon.models.beamline import MXBeamline from aarecommon.models.models import SampleCameraSettings, StagePositionEnum -from epics import PV from aare.devices import aerotech, smargon from aare.devices.area_detector import AutoEnum, epicsAD from aare.devices.bec_worker import BECClientWorker -from aare.devices.enum_pv import EnumPV from aare.devices.experimental_hutch_shutter import ExperimentalHutchShutter -from aare.devices.my_motor import MyMotor from aare.devices.pss_state import PssState -from aare.devices.set_get_pv import PredefinedPV, SetGetPV from aare.devices.tell_client import make_tell_client logger = setup_logger("aareDAQ") @@ -37,78 +33,14 @@ class BeamlineDevices: logger.debug("initialising BEC worker done") self._smargon = smargon.Smargon(beamline) self.exp_shutter = ExperimentalHutchShutter(beamline=self._beamline) + self.dtz_mod = cfg_get("daq.detector_distance_limit_modifier", 1.0) + + # TODO convert epics pvs to BEC -> need to get all samcam signals into a device + # BEC device for pss state also good # Personnel Safety System: gates whether the robot is allowed to move. self.pss = PssState(beamline=self._beamline) - - # faster to define the dtz object here than in functions and then use - self._dtz = self.bec_worker.dev.det_z - self.dtz_mod = cfg_get("daq.detector_distance_limit_modifier", 1.0) - # TODO convert epics pvs to BEC self._sample_cam = epicsAD(f"{BEAMLINE}-ES-MS:") - self._front_light = PredefinedPV( - name="front_light", - setpv=f"{BEAMLINE}-ES-FL:SET", - getpv=f"{BEAMLINE}-ES-FL:SET", - predefs={"off": 1.49, "half": 2.0, "max": 3.0}, - timeout=10.0, - ) - self._back_light = PredefinedPV( - name="back_light", - setpv=f"{BEAMLINE}-ES-BL:SET", - getpv=f"{BEAMLINE}-ES-BL:SET", - predefs={"off": 0, "half": 0.98, "max": 1.2}, - timeout=10.0, - ) - - # self._front_light = self.bec_worker.dev.fl_bright # need wrapper on bec_worker layer - # self._back_light = self.bec_worker.dev.bl_bright #need wrapper on bec_worker layer - - self._back_light_pos = EnumPV( - name="back_light_pos", - setpv=f"{BEAMLINE}-ES-BL:POS-SET", - getpv=f"{BEAMLINE}-ES-BL:POS-GET", - timeout=10.0, - ) - - self._ringcurrent = self.bec_worker.ring_current - - self._zoom = SetGetPV( - name="zoom", setpv=f"{BEAMLINE}-ES-MS:ZOOM.VAL", getpv=f"{BEAMLINE}-ES-MS:ZOOM.RBV" - ) - - self._cryojet_pos = EnumPV( - name="cryojet_pos", - setpv=f"{BEAMLINE}-ES-CS:POS-SET", - getpv=f"{BEAMLINE}-ES-CS:POS-GET", - timeout=10.0, - ) - - self._cryojet_x = MyMotor(f"{BEAMLINE}-ES-CS:TRX") # currently in is 5 out is 15? - - self._cryojet_temperature_get = PV(f"{BEAMLINE}-ES-CS:TEMP_RBV") - self._cryojet_temperature_set = PV(f"{BEAMLINE}-ES-CS:TEMP.VAL") - - if self._beamline == MXBeamline.X06DA: - self._transmission = SetGetPV( - name="transmission", - setpv=f"{BEAMLINE}-ES-BCFI:TRANSM-SET", - getpv=f"{BEAMLINE}-ES-BCFI:TRANSM-GET", - ) - else: - self._transmission = SetGetPV( - name="transmission", - setpv=f"{BEAMLINE}-ES-SSFI:TRANSM-SET", - getpv=f"{BEAMLINE}-ES-SSFI:TRANSM-GET", - ) - - self._fast_shutter = PV(f"{BEAMLINE}-ES-SHUTTER:SET") - - self.magnet_position_sensor = PV(f"{BEAMLINE}-ES-DFS:CBOX-CMP1") - self.magnet_position_sensor_readout = PV(f"{BEAMLINE}-ES-DFS:CBOX-USER1") - # self.magnet_position_sensor_readout = PV(f"{BEAMLINE}-ES-DFS:CBOX-REFVAL1") - self.magnet_position_sensor_state = PV(f"{BEAMLINE}-ES-DFS:CBOX-STATE") - def restart_bec_worker(self, simulated: bool = False) -> None: try: if getattr(self, "bec_worker", None) is not None and not simulated: @@ -141,7 +73,7 @@ class BeamlineDevices: # Transmission @property def transmission(self) -> float: - return self._transmission.value + return self.bec_worker.tra @transmission.setter def transmission(self, value: float): @@ -178,14 +110,14 @@ class BeamlineDevices: # Zoom @property def zoom(self) -> float: - return self._zoom.value + return self._samcam_zoom.get() @zoom.setter def zoom(self, value: float): self.set_zoom(value, wait=True) def set_zoom(self, value: float, /, wait: bool = True): - self._zoom.move(value, wait=wait) + self._samcam_zoom.put(value, wait=wait) # Optics @property diff --git a/src/aare/devices/bec_worker.py b/src/aare/devices/bec_worker.py index 28e908b1..2b3d3384 100644 --- a/src/aare/devices/bec_worker.py +++ b/src/aare/devices/bec_worker.py @@ -1,15 +1,16 @@ import sys import time -from enum import Enum +from enum import Enum, StrEnum from aarecommon.config.beamline import cfg_get, mx_beamline from aarecommon.config.logger import setup_logger from aarecommon.config.logger_events import log_timing -from aarecommon.errors.exception_handler import BECCommunicationError +from aarecommon.errors.exception_handler import BECCommunicationError, BECException from aarecommon.models.beamline import MXBeamline from bec_ipython_client import BECIPythonClient from bec_ipython_client.signals import OperationMode from bec_lib.device import RPCError, ScanRequestError +from bec_lib.devicemanager import DeviceContainer from bec_lib.procedures.helper import FrontendProcedureHelper from bec_lib.service_config import ServiceConfig @@ -49,39 +50,49 @@ class BeamlineState(str, Enum): XTAL_SNAPSHOT = "xtal_snapshot" +class BecDevices(StrEnum): + DET_Z = "det_z" + FL_BRIGHT = "fl_bright" + BL_BRIGHT = "bl_bright" + BL_POS = "bl_pos" + SCAM_ZOOM = "scam_zoom" + CRYO_POS = "cryo_pos" + CRYO_TEMP = "cryo_temp" + CRYO_X = "cryo_x" + TRANSM = "transm" + BCU_SHUTTER = "bcu_shutter" + MAG_OK = "mag_ok" + MAG_RBV = "mag_rbv" + + +class MissingDevicesError(BECException): ... + + class BECClientWorker: def __init__(self, beamline: MXBeamline, name: str = "default"): - BEAMLINE = beamline.value.lower() - self.beamline = beamline - if self.beamline is MXBeamline.X06DA: - self._beamline_name = "pxiii" - elif self.beamline is MXBeamline.X10SA: - self._beamline_name = "pxii" - elif self.beamline is MXBeamline.X06DA: - self._beamline_name = "pxi" - elif self.beamline is MXBeamline.SIMULATED: - self._beamline_name = "SIMULATED" - else: - raise ValueError(f"Unknown beamline: {beamline}") - self.dispatch = get_beamline_dispatch() - if self.beamline is MXBeamline.SIMULATED: + if beamline is MXBeamline.SIMULATED: self.simulated = True else: self.simulated = False - logger.debug(f"Initializing BECClientWorker for {BEAMLINE} beamline") - host = cfg_get("daq.hardware.bec_url", f"{BEAMLINE}-bec-001.psi.ch") + logger.debug( + f"Initializing BECClientWorker for {self.dispatch.config.bl_lower} beamline" + ) + host = cfg_get( + "daq.hardware.bec_url", f"{self.dispatch.config.bl_lower}-bec-001.psi.ch" + ) service_config = ServiceConfig(redis={"host": host, "port": 6379}) service_config.config["log_writer"]["base_path"] = "/tmp/logs" - # service_config.config["user_macros"]["base_path"]=f'/sls/{BEAMLINE}/config/bec/production/pxiii_bec/pxiii_bec' - # print(service_config.config) self.client = BECIPythonClient(config=service_config, mode=OperationMode.Procedure) self.client.start() - # self.client.config.update_session_with_file("/sls/x10sa/config/bec/production/bec/bec_lib/bec_lib/config_helper.py") - self.dev = self.client.device_manager.devices - print(self.dev.keys()) + self.dev: DeviceContainer = self.client.device_manager.devices + missing_devices = {str(d) for d in list(BecDevices) if d not in self.dev} + if missing_devices: + raise MissingDevicesError( + f"{missing_devices} were not found in the BEC session. Available: {list(self.dev.keys())}. If no devices are available, load a BEC device config." + ) self.scans = self.client.scans self.macros = self.dispatch.bec_macros self.helper = FrontendProcedureHelper(self.client.connector) @@ -94,26 +105,11 @@ class BECClientWorker: logger.debug(f"simulated is {self.simulated}") def _init_beamline_environment(self): - try: - self.position_devices, self.planner = self.macros.init_beamline_environment() - self._backlight_brightness = self.position_devices["bl_bright"] - self._frontlight_brightness = self.position_devices["fl_bright"] - self._zoom = self.dev.scam_zoom - self._ring_current = self.dev.sls_current - except Exception as e: - logger.error(f"Error initialising planar and position devices: {e}") - self.position_devices = None - self.planner = None - self._backlight_brightness = None - self._frontlight_brightness = None - try: - self._zoom = self.dev.scam_zoom - self._ring_current = self.dev.sls_current - except Exception as e: - logger.exception("Error initialising zoom and ring_current") - self._zoom = None - self.ring_current = None - raise RuntimeError(f"Error initialising BEC devices: {e}") from e + self.position_devices, self.planner = self.macros.init_beamline_environment() + self._backlight_brightness = self.position_devices["bl_bright"] + self._frontlight_brightness = self.position_devices["fl_bright"] + self._zoom = self.dev[BecDevices.SCAM_ZOOM] + self._ring_current = self.dev.sls_current def _bec_error( self, exc: Exception, *, operation: str, tags: list[str] | None = None @@ -219,28 +215,6 @@ class BECClientWorker: except Exception: logger.exception("Error sending scilog message") - def run_macro(self, macro_name: str, *args, queue: str = "default", **kwargs): - if self.simulated: - logger.debug(f"Simulating macro {macro_name}") - return None - try: - return self.client.proc.run_macro(macro_name, *args, queue=queue) - except Exception as e: - raise self._bec_error(e, operation=f"run_macro:{macro_name}") from e - - def run_macro_blocked(self, macro_name: str, *args, queue: str = "default", **kwargs): - if self.simulated: - logger.debug(f"Simulating macro {macro_name}") - return None - try: - status = self.run_macro(macro_name, *args, queue=queue) - print(status) - status.wait() - print(status) - return status - except Exception as e: - raise self._bec_error(e, operation=f"run_macro_blocked:{macro_name}") from e - @log_timing(logger, "BEC move_to") def move_to(self, state: BeamlineState): start = time.perf_counter() @@ -299,27 +273,6 @@ class BECClientWorker: except Exception as e: raise self._bec_error(e, operation="list_all_user_macros") from e - def _list_all_macros(self): - result = self.client.macros.list_user_macros() - if result is None: - return [] - return result - - def load_user_macros(self): - if self.simulated: - logger.debug("Simulating load_user_macros") - return None - try: - return self._load_user_macros() - except Exception as e: - raise self._bec_error(e, operation="load_user_macros") from e - - def _load_user_macros(self): - result = self.macros.load_all_user_macros() - if result is None: - logger.warning("BEC load_all_user_macros returned None") - return result - def reinitialise_planner_and_position_devices(self, method: str = "auto"): """ Reinitialise BEC position devices and planner. @@ -337,9 +290,7 @@ class BECClientWorker: logger.debug(f"Simulating reinitialise_planner_and_position_devices(method={method})") return [] try: - self.client.config.update_session_with_file( - f"/sls/{self.beamline}/config/bec/production/{self._beamline_name}_bec/{self._beamline_name}_bec/device_configs/{self._beamline_name}-devices.yaml" - ) + self.client.config.update_session_with_file(self.dispatch.config.device_yml_path) self._init_beamline_environment() logger.info(f"Reinitialised BEC planner and position devices using method={method}") return self.list_position_devices() @@ -356,7 +307,7 @@ class BECClientWorker: self.macros.mono_pitch_scan(plot) except Exception as e: raise self._bec_error(e, operation="mono_pitch_scan", tags=["mono_pitch_scan"]) from e - if self.beamline is MXBeamline.X06DA: + if self.dispatch.config.bl_upper == MXBeamline.X06DA: addtional_text = [f"New dcm_pitch position: {self.dev.dcm_pitch.position:5f}"] else: addtional_text = [f"New dcm_theta2 position: {self.dev.dccm_theta2.position:5f}"] -- 2.54.0 From 56546617f8d6ae731723216725f2465da0acb4d4 Mon Sep 17 00:00:00 2001 From: David Perl Date: Thu, 6 Aug 2026 14:34:14 +0200 Subject: [PATCH 2/8] WIP remove pv modules --- src/aare/devices/enum_pv.py | 36 ------- src/aare/devices/my_motor.py | 104 ------------------- src/aare/devices/set_get_pv.py | 85 --------------- tests/unit/devices/test_enum_pv.py | 96 ----------------- tests/unit/devices/test_my_motor.py | 156 ---------------------------- 5 files changed, 477 deletions(-) delete mode 100644 src/aare/devices/enum_pv.py delete mode 100644 src/aare/devices/my_motor.py delete mode 100644 src/aare/devices/set_get_pv.py delete mode 100644 tests/unit/devices/test_enum_pv.py delete mode 100644 tests/unit/devices/test_my_motor.py diff --git a/src/aare/devices/enum_pv.py b/src/aare/devices/enum_pv.py deleted file mode 100644 index fd725e17..00000000 --- a/src/aare/devices/enum_pv.py +++ /dev/null @@ -1,36 +0,0 @@ -from enum import Enum -from typing import Any - -from aare.devices.set_get_pv import MoveResult, SetGetPV - - -class EnumPV(SetGetPV): - def __init__(self, name: str, setpv: str, getpv: str, **kwargs): - super().__init__(name, setpv, getpv, **kwargs) - if not self.setpoint_pv.enum_strs: - raise RuntimeError(f"{setpv} is not an ENUM PV") - - @property - def position(self) -> str: - return self.readback_pv.get(as_string=True) - - def _resolve(self, x: Any) -> MoveResult: - # accept Enum member - if isinstance(x, Enum): - x = x.name - - # accept index - if isinstance(x, int): - try: - return MoveResult(target=self.setpoint_pv.enum_strs[x], name=None) - except Exception as e: - raise ValueError(f"Bad enum index {x}") from e - - # accept name -> match against enum strings (case-insensitive) - if isinstance(x, str): - for s in self.setpoint_pv.enum_strs: - if s.strip().lower() == x.strip().lower(): - return MoveResult(target=s, name=s) - raise ValueError(f"'{x}' not in {list(self.setpoint_pv.enum_strs)}") - - raise TypeError(f"Unsupported enum command type: {type(x).__name__}") diff --git a/src/aare/devices/my_motor.py b/src/aare/devices/my_motor.py deleted file mode 100644 index 7878f38c..00000000 --- a/src/aare/devices/my_motor.py +++ /dev/null @@ -1,104 +0,0 @@ -import time - -from epics import Motor - - -class MyMotor(Motor): - """Wrapper for the EPICS motor PV.""" - - def __init__(self, name, timeout=5.0): - super().__init__(name.upper(), timeout=timeout) - - @property - def speed(self): - """Gets the current motor.slew_speed value""" - return self.get("VELO") - - @speed.setter - def speed(self, v): - """Sets the motor slew speed""" - self.put("VELO", v) - - @property - def position(self): - """Gets the current motor readback value""" - return self.readback - - @property - def value(self): - """Gets the current motor.drive value""" - return self.drive - - @value.setter - def value(self, v): - """Sets motor.drive to a set value""" - self.drive = v - - def stop(self): - """Stops the motor""" - self.stop_motor() - - @property - def moving(self): - """Returns True if the motor is moving""" - return bool(self.get("MOVN")) - - @property - def units(self): - """Returns the units of the motor as a string""" - return self.get("EGU", as_string=True) - - @property - def limits(self): - """Returns (low_limit, high_limit)""" - return self.get("HLM"), self.get("LLM") - - @limits.setter - def limits(self, limits): - """Sets (low_limit, high_limit)""" - low, high = limits - self.put("LLM", low) - self.put("HLM", high) - - def move_motor(self, val, relative=False, wait=False, timeout=300.0): - """ - Moves the motor to an absolute or relative position. - :param val: Position to move to - :param relative: If True, moves relative to current position - :param wait: If True, waits for completion (synchronous) - :param timeout: Maximum time to wait for completion - """ - - return self.move(val, relative=relative, wait=wait, timeout=timeout) - - def home(self, direction="forward", wait=False): - """ - Homes the motor. - :param direction: 'forward' or 'reverse' - """ - field = "HOMF" if direction == "forward" else "HOMR" - self.put(field, 1) - if wait: - self.wait_for_stop() - - def wait_for_stop(self, timeout=300.0, poll_rate=0.01): - """ - Synchronous wait until the motor stops moving. - """ - start_time = time.time() - while self.moving: - time.sleep(poll_rate) - if time.time() - start_time > timeout: - raise RuntimeError(f"Timeout waiting for motor {self.name} to stop") - - async def wait_for_stop_async(self, timeout=300.0, poll_rate=0.01): - """ - Asynchronous wait until the motor stops moving. - """ - import asyncio - - start_time = time.time() - while self.moving: - await asyncio.sleep(poll_rate) - if time.time() - start_time > timeout: - raise RuntimeError(f"Timeout waiting for motor {self.name} to stop") diff --git a/src/aare/devices/set_get_pv.py b/src/aare/devices/set_get_pv.py deleted file mode 100644 index fa849073..00000000 --- a/src/aare/devices/set_get_pv.py +++ /dev/null @@ -1,85 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable, Mapping -from dataclasses import dataclass -from typing import Any - -from epics import PV - -from aare.devices.mx_lib import pv_wait - -RawValue = str | float | int -ResolverValue = ( - RawValue - | tuple[Callable[..., RawValue], tuple[Any, ...]] # (func, args) pattern you already use -) - - -@dataclass -class MoveResult: - target: RawValue - name: str | None = None - - -class SetGetPV: - def __init__( - self, - name: str, - setpv: str, - getpv: str, - *, - timeout: float = 60.0, - tolerance: float | None = None, - ): - self.name = name - self.setpoint_pv = PV(setpv) - self.readback_pv = PV(getpv) - self.default_timeout = timeout - self.tolerance = tolerance - self._last_target: RawValue | None = None - - @property - def value(self) -> Any: - return self.readback_pv.get() - - def _resolve(self, x: Any) -> MoveResult: - # Default: treat input as raw value - return MoveResult(target=x, name=None) - - def move(self, x: Any, *, wait: bool = False, timeout: float | None = None) -> MoveResult: - res = self._resolve(x) - self._last_target = res.target - self.setpoint_pv.put(res.target) - if wait: - self.wait(timeout=timeout) - return res - - def wait(self, *, timeout: float | None = None): - if self._last_target is None: - return - pv_wait( - self.readback_pv, - self._last_target, - timeout=timeout or self.default_timeout, - tolerance=self.tolerance, - ) - - -class PredefinedPV(SetGetPV): - def __init__( - self, name: str, setpv: str, getpv: str, predefs: Mapping[str, ResolverValue], **kwargs - ): - super().__init__(name, setpv, getpv, **kwargs) - self._predefs = dict(predefs) - - @property - def positions(self) -> list[str]: - return list(self._predefs.keys()) - - def _resolve(self, x: Any) -> MoveResult: - if isinstance(x, str) and x in self._predefs: - v = self._predefs[x] - if isinstance(v, tuple) and callable(v[0]): - v = v[0](*v[1]) - return MoveResult(target=v, name=x) - return MoveResult(target=x, name=None) diff --git a/tests/unit/devices/test_enum_pv.py b/tests/unit/devices/test_enum_pv.py deleted file mode 100644 index 89658a5e..00000000 --- a/tests/unit/devices/test_enum_pv.py +++ /dev/null @@ -1,96 +0,0 @@ -from enum import Enum -from unittest.mock import patch - -import pytest - -from aare.devices.enum_pv import EnumPV - - -class MockPV: - def __init__(self, name): - self.pvname = name - self.enum_strs = None - self.value = None - - def get(self, as_string=False): - if as_string: - return str(self.value) - return self.value - - def put(self, value): - self.value = value - - -@pytest.fixture -def mock_pvs(): - with patch("aare.devices.set_get_pv.PV") as mock_pv: - set_pv = MockPV("SET") - get_pv = MockPV("GET") - mock_pv.side_effect = lambda name: set_pv if name == "SET" else get_pv - yield set_pv, get_pv - - -def test_enum_pv_init_fail(mock_pvs): - set_pv, _get_pv = mock_pvs - set_pv.enum_strs = None - with pytest.raises(RuntimeError): - EnumPV("test", "SET", "GET") - - -def test_enum_pv_init_success(mock_pvs): - set_pv, _get_pv = mock_pvs - set_pv.enum_strs = ("State1", "State2") - epv = EnumPV("test", "SET", "GET") - assert epv.name == "test" - - -def test_enum_pv_position(mock_pvs): - set_pv, get_pv = mock_pvs - set_pv.enum_strs = ("State1", "State2") - get_pv.value = "State1" - epv = EnumPV("test", "SET", "GET") - assert epv.position == "State1" - - -def test_enum_pv_resolve_enum(mock_pvs): - set_pv, _get_pv = mock_pvs - set_pv.enum_strs = ("State1", "State2") - epv = EnumPV("test", "SET", "GET") - - class MyEnum(Enum): - State1 = 1 - - res = epv._resolve(MyEnum.State1) - assert res.target == "State1" - - -def test_enum_pv_resolve_int(mock_pvs): - set_pv, _get_pv = mock_pvs - set_pv.enum_strs = ("State1", "State2") - epv = EnumPV("test", "SET", "GET") - - res = epv._resolve(0) - assert res.target == "State1" - - with pytest.raises(ValueError): - epv._resolve(5) - - -def test_enum_pv_resolve_str(mock_pvs): - set_pv, _get_pv = mock_pvs - set_pv.enum_strs = (" State1 ", "State2") - epv = EnumPV("test", "SET", "GET") - - res = epv._resolve("state1") - assert res.target == " State1 " - - with pytest.raises(ValueError): - epv._resolve("Missing") - - -def test_enum_pv_resolve_invalid_type(mock_pvs): - set_pv, _get_pv = mock_pvs - set_pv.enum_strs = ("State1", "State2") - epv = EnumPV("test", "SET", "GET") - with pytest.raises(TypeError): - epv._resolve(1.5) diff --git a/tests/unit/devices/test_my_motor.py b/tests/unit/devices/test_my_motor.py deleted file mode 100644 index c6cb1258..00000000 --- a/tests/unit/devices/test_my_motor.py +++ /dev/null @@ -1,156 +0,0 @@ -from unittest.mock import MagicMock, PropertyMock, patch - -import pytest - -from aare.devices.my_motor import MyMotor - - -@pytest.fixture -def mock_motor_base(): - with patch("epics.motor.Motor.__init__", return_value=None) as mock_init: - # We need to yield Mocks that will be used by the instance methods - mock_get = MagicMock() - mock_put = MagicMock() - mock_move = MagicMock() - - # Patch them on the base class epics.motor.Motor - with ( - patch("epics.motor.Motor.get", mock_get), - patch("epics.motor.Motor.put", mock_put), - patch("epics.motor.Motor.move", mock_move), - patch("epics.motor.Motor.PV"), - patch("epics.motor.Motor.readback", create=True), - patch("epics.motor.Motor.drive", create=True), - patch("epics.motor.Motor.stop_motor", create=True), - patch("epics.motor.Motor.moving_flag", create=True), - ): - yield mock_init, mock_get, mock_put, mock_move - - -def test_my_motor_init(mock_motor_base): - mock_init, _, _, _ = mock_motor_base - _ = MyMotor("X10SA-DI-MTR-01") - mock_init.assert_called_with("X10SA-DI-MTR-01", timeout=5.0) - - -def test_my_motor_speed(mock_motor_base): - _, mock_get, _mock_put, _ = mock_motor_base - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - mock_get.return_value = 1.0 - assert m.speed == 1.0 - mock_get.assert_called_with("VELO") - - # When m.speed = 2.0 is called, it might use epics.motor.Motor.put - # Let's see if we can just test that it doesn't crash for now if we can't capture the call - m.speed = 2.0 - - -def test_my_motor_properties(mock_motor_base): - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - with patch("epics.motor.Motor.readback", new_callable=PropertyMock, create=True) as mock_rb: - mock_rb.return_value = 10.0 - assert m.position == 10.0 - - with patch("epics.motor.Motor.drive", new_callable=PropertyMock, create=True) as mock_drive: - mock_drive.return_value = 5.0 - assert m.value == 5.0 - # m.value = 6.0 calls self.drive = 6.0. - # In epics.Motor, drive is a PVProperty. - m.value = 6.0 - - -def test_my_motor_stop(mock_motor_base): - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - with patch("epics.motor.Motor.stop_motor", create=True) as mock_stop: - m.stop() - mock_stop.assert_called_once() - - -def test_my_motor_moving(mock_motor_base): - _, mock_get, _, _ = mock_motor_base - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - mock_get.return_value = 1 - assert m.moving is True - mock_get.return_value = 0 - assert m.moving is False - - -def test_my_motor_units(mock_motor_base): - _, mock_get, _, _ = mock_motor_base - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - mock_get.return_value = "mm" - assert m.units == "mm" - mock_get.assert_called_with("EGU", as_string=True) - - -def test_my_motor_limits(mock_motor_base): - _, mock_get, _mock_put, _ = mock_motor_base - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - mock_get.side_effect = [100, 0] - assert m.limits == (100, 0) - - m.limits = (0, 100) - - -def test_my_motor_move_motor(mock_motor_base): - _, _, _, mock_move = mock_motor_base - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - m.move_motor(10, relative=True, wait=True) - mock_move.assert_called_with(10, relative=True, wait=True, timeout=300.0) - - -def test_my_motor_home(mock_motor_base): - _, _, mock_put, _ = mock_motor_base - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - m.home(direction="forward") - mock_put.assert_called_with("HOMF", 1) - - m.home(direction="reverse") - mock_put.assert_called_with("HOMR", 1) - - -def test_my_motor_wait_for_stop(mock_motor_base): - _, mock_get, _, _ = mock_motor_base - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - # Simulate moving then stopping - # First 2 calls return 1 (moving), 3rd call returns 0 (stopped) - mock_get.side_effect = [1, 1, 0] - m.wait_for_stop(timeout=1.0, poll_rate=0.001) - - -def test_my_motor_wait_for_stop_timeout(mock_motor_base): - _, mock_get, _, _ = mock_motor_base - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - mock_get.return_value = 1 - with pytest.raises(RuntimeError): - m.wait_for_stop(timeout=0.01, poll_rate=0.001) - - -@pytest.mark.asyncio -async def test_my_motor_wait_for_stop_async(mock_motor_base): - _, mock_get, _, _ = mock_motor_base - m = MyMotor("MTR") - with patch.object(MyMotor, "name", create=True, new_callable=PropertyMock) as mock_name: - mock_name.return_value = "MTR" - mock_get.side_effect = [1, 0] - await m.wait_for_stop_async(timeout=1.0, poll_rate=0.001) -- 2.54.0 From 295db27f80c8e5c34a9426e0227e3f1844387b7d Mon Sep 17 00:00:00 2001 From: David Perl Date: Thu, 6 Aug 2026 17:14:12 +0200 Subject: [PATCH 3/8] wip fix bp, remove more cruft --- pyproject.toml | 1 + .../beamline_dispatch/beamline_dispatch.py | 6 +- .../default/beamline_dispatch.py | 30 +- src/aare/beamline_dispatch/protocols.py | 11 + .../simulated/beamline_dispatch.py | 17 +- .../x06da/beamline_dispatch.py | 60 ++-- .../x06sa/beamline_dispatch.py | 4 +- .../x10sa/beamline_dispatch.py | 54 ++-- src/aare/devices/bec_worker.py | 227 ++------------ src/aare/devices/mx_lib.py | 289 ------------------ tests/unit/devices/test_mx_lib.py | 92 ------ 11 files changed, 139 insertions(+), 652 deletions(-) delete mode 100644 src/aare/devices/mx_lib.py delete mode 100644 tests/unit/devices/test_mx_lib.py diff --git a/pyproject.toml b/pyproject.toml index 5ae6974f..22d48a80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -101,6 +101,7 @@ multi_line_output = 3 line-length = 100 [tool.ruff.lint] +isort.split-on-trailing-comma=false ignore = [ # Timestamps here are local beamline wall-clock and are stored/displayed as # naive ISO strings; adding a tzinfo would change what aaredb and the GUI see. diff --git a/src/aare/beamline_dispatch/beamline_dispatch.py b/src/aare/beamline_dispatch/beamline_dispatch.py index 80e1bb47..879ff005 100644 --- a/src/aare/beamline_dispatch/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/beamline_dispatch.py @@ -13,12 +13,12 @@ def get_beamline_dispatch() -> BeamlineDispatch: case MXBeamline.X06DA: from .x06da import X06daDispatch - return X06daDispatch() + return X06daDispatch(MXBeamline.X06DA) case MXBeamline.X06SA: from .x06sa import X06saDispatch - return X06saDispatch() + return X06saDispatch(MXBeamline.X06SA) # pyright: ignore case MXBeamline.X10SA: from .x10sa import X10saDispatch - return X10saDispatch() + return X10saDispatch(MXBeamline.X10SA) diff --git a/src/aare/beamline_dispatch/default/beamline_dispatch.py b/src/aare/beamline_dispatch/default/beamline_dispatch.py index 500dd6cf..a2830190 100644 --- a/src/aare/beamline_dispatch/default/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/default/beamline_dispatch.py @@ -1,8 +1,15 @@ import os +from abc import ABC from aarecommon.config.beamline import MXBeamline +from aarecommon.config.logger import setup_logger +from bec_ipython_client import BECIPythonClient +from bec_ipython_client.signals import OperationMode +from bec_lib.client import ServiceConfig -from aare.beamline_dispatch.protocols import AuthDispatch, BeamlineDispatch, ConfigKeys +from aare.beamline_dispatch.protocols import AuthDispatch, BeamlineDispatch, BecSetup, ConfigKeys + +logger = setup_logger("aareDAQ") class DefaultAuthDispatch(AuthDispatch): @@ -45,13 +52,28 @@ class DefaultConfigKeys(ConfigKeys): return f"/sls/{self.bl_lower}/config/bec/production/{self.bl_name}_bec/{self.bl_name}_bec/device_configs/{self.bl_name}-devices.yaml" -class DefaultDispatch(BeamlineDispatch): +class DefaultBecSetup(BecSetup): + def __init__(self, beamline: MXBeamline) -> None: + super().__init__(beamline) + self._beamline = beamline + + def create_client(self) -> BECIPythonClient: + bl = self._beamline.value.lower() + logger.debug(f"Initializing BECClientWorker for {bl} beamline") + host = f"{bl}-bec-001.psi.ch" + service_config = ServiceConfig(redis={"host": host, "port": 6379}) + service_config.config["log_writer"]["base_path"] = "/tmp/logs" + return BECIPythonClient(config=service_config, mode=OperationMode.Procedure) + + +class DefaultDispatch(BeamlineDispatch, ABC): """Default implementation for anything which can vary between beamlines and/or simulation. Should be safe and fail rather than assuming anything.""" def __init__(self, beamline: MXBeamline) -> None: self._auth = DefaultAuthDispatch() self._config = DefaultConfigKeys(beamline) + self._bec_setup = DefaultBecSetup(beamline) @property def auth(self): @@ -60,3 +82,7 @@ class DefaultDispatch(BeamlineDispatch): @property def config(self): return self._config + + @property + def bec_setup(self): + return self._bec_setup diff --git a/src/aare/beamline_dispatch/protocols.py b/src/aare/beamline_dispatch/protocols.py index e562c9d4..61f2ebcc 100644 --- a/src/aare/beamline_dispatch/protocols.py +++ b/src/aare/beamline_dispatch/protocols.py @@ -1,6 +1,9 @@ from abc import ABC, abstractmethod from typing import Any +from aarecommon.models.models import MXBeamline +from bec_ipython_client import BECIPythonClient + class AuthDispatch(ABC): @abstractmethod @@ -28,6 +31,11 @@ class BecMacros(ABC): def mono_pitch_scan(plot=True): ... +class BecSetup(ABC): + def __init__(self, beamline: MXBeamline) -> None: ... + def create_client(self) -> BECIPythonClient: ... + + class ConfigKeys(ABC): @property @abstractmethod @@ -51,4 +59,7 @@ class BeamlineDispatch(ABC): def bec_macros(self) -> BecMacros: ... @property @abstractmethod + def bec_setup(self) -> BecSetup: ... + @property + @abstractmethod def config(self) -> ConfigKeys: ... diff --git a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py index b023bf79..b2521b1e 100644 --- a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py @@ -1,9 +1,11 @@ from typing import Any +from unittest.mock import MagicMock from aarecommon.config.beamline import MXBeamline +from bec_ipython_client import BECIPythonClient from aare.beamline_dispatch.default.beamline_dispatch import DefaultConfigKeys, DefaultDispatch -from aare.beamline_dispatch.protocols import AuthDispatch, BecMacros +from aare.beamline_dispatch.protocols import AuthDispatch, BecMacros, BecSetup class SimulatedAuthDispatch(AuthDispatch): @@ -25,13 +27,22 @@ class SimulatedBecMacros(BecMacros): @staticmethod def bl_energy(energy_ev, move_gap=True, mono_scan=True, plot=True): ... @staticmethod - def get_current_energy(): ... + def get_current_energy() -> float: ... @staticmethod def mono_pitch_scan(plot=True): ... +class SimulatedBecSetup(BecSetup): + def __init__(self, beamline: MXBeamline) -> None: + super().__init__(beamline) + self._beamline = beamline + + def create_client(self) -> BECIPythonClient: + return MagicMock(spec=BECIPythonClient) + + class SimulatedDispatch(DefaultDispatch): - def __init__(self) -> None: + def __init__(self) -> None: # pyright: ignore[reportMissingSuperCall] self._auth = SimulatedAuthDispatch() self._bec_macros = SimulatedBecMacros() self._config = DefaultConfigKeys(MXBeamline.SIMULATED) diff --git a/src/aare/beamline_dispatch/x06da/beamline_dispatch.py b/src/aare/beamline_dispatch/x06da/beamline_dispatch.py index 1ef2578a..f92d4466 100644 --- a/src/aare/beamline_dispatch/x06da/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x06da/beamline_dispatch.py @@ -1,49 +1,47 @@ from typing import Any from aarecommon.config.beamline import MXBeamline +from pxiii_bec.macros.build_devices import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] + save_and_reload, + save_current_position, +) +from pxiii_bec.macros.init_beamline import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] + init_beamline_environment, +) +from pxiii_bec.macros.pxiii_energy import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] + bl_energy, + get_current_energy, + mono_pitch_scan, +) from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch from aare.beamline_dispatch.protocols import BecMacros class X06daBecMacros(BecMacros): - def __init__(self) -> None: - from pxiii_bec.macros.build_devices import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] - save_and_reload, - save_current_position, - ) - from pxiii_bec.macros.init_beamline import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] - init_beamline_environment, - ) - from pxiii_bec.macros.pxiii_energy import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] - bl_energy, - get_current_energy, - mono_pitch_scan, - ) - - self.save_and_reload = save_and_reload - self.save_current_position = save_current_position - self.init_beamline_environment = init_beamline_environment - - def _bl_energy(energy_ev, move_gap=True, mono_scan=True, plot=True): - return bl_energy(energy_ev=energy_ev, plot=plot) - - self.bl_energy = _bl_energy - self.get_current_energy = get_current_energy - self.mono_pitch_scan = mono_pitch_scan + @staticmethod + def save_and_reload() -> tuple[Any, Any]: + return save_and_reload() @staticmethod - def save_and_reload() -> tuple[Any, Any]: ... + def save_current_position(device, position, axis=None, force=False, max_delta=0.5) -> None: + return save_current_position + @staticmethod - def save_current_position(device, position, axis=None, force=False, max_delta=0.5) -> None: ... + def init_beamline_environment() -> tuple[Any, Any]: + return init_beamline_environment() + @staticmethod - def init_beamline_environment() -> tuple[Any, Any]: ... + def bl_energy(energy_ev, move_gap=True, mono_scan=True, plot=True): + return bl_energy(energy_ev=energy_ev, plot=plot) + @staticmethod - def bl_energy(energy_ev, move_gap=True, mono_scan=True, plot=True): ... + def get_current_energy() -> float: + return get_current_energy() + @staticmethod - def get_current_energy(): ... - @staticmethod - def mono_pitch_scan(plot=True): ... + def mono_pitch_scan(plot=True): + return mono_pitch_scan(plot) class X06daDispatch(DefaultDispatch): diff --git a/src/aare/beamline_dispatch/x06sa/beamline_dispatch.py b/src/aare/beamline_dispatch/x06sa/beamline_dispatch.py index fcae2cec..9224cbd0 100644 --- a/src/aare/beamline_dispatch/x06sa/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x06sa/beamline_dispatch.py @@ -1,4 +1,6 @@ +from abc import ABC + from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch -class X06saDispatch(DefaultDispatch): ... +class X06saDispatch(DefaultDispatch, ABC): ... diff --git a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py index 6ef814dd..b12a9344 100644 --- a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py @@ -1,43 +1,47 @@ from typing import Any from aarecommon.config.beamline import MXBeamline +from pxii_bec.macros.build_devices import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] + save_and_reload, + save_current_position, +) +from pxii_bec.macros.init_beamline import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] + init_beamline_environment, +) +from pxii_bec.macros.pxii_energy import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] + bl_energy, + get_current_energy, + mono_pitch_scan, +) from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch from aare.beamline_dispatch.protocols import BecMacros class X10SaBecMacros(BecMacros): - def __init__(self) -> None: - from pxii_bec.macros.build_devices import ( # type: ignore - save_and_reload, - save_current_position, - ) - from pxii_bec.macros.init_beamline import init_beamline_environment # type: ignore - from pxii_bec.macros.pxii_energy import ( # type: ignore - bl_energy, - get_current_energy, - mono_pitch_scan, - ) - - self.save_and_reload = save_and_reload - self.save_current_position = save_current_position - self.init_beamline_environment = init_beamline_environment - self.bl_energy = bl_energy - self.get_current_energy = get_current_energy - self.mono_pitch_scan = mono_pitch_scan + @staticmethod + def save_and_reload() -> tuple[Any, Any]: + return save_and_reload() @staticmethod - def save_and_reload() -> tuple[Any, Any]: ... + def save_current_position(device, position, axis=None, force=False, max_delta=0.5) -> None: + return save_current_position + @staticmethod - def save_current_position(device, position, axis=None, force=False, max_delta=0.5) -> None: ... + def init_beamline_environment() -> tuple[Any, Any]: + return init_beamline_environment() + @staticmethod - def init_beamline_environment() -> tuple[Any, Any]: ... + def bl_energy(energy_ev, move_gap=True, mono_scan=True, plot=True): + return bl_energy(energy_ev=energy_ev, move_gap=move_gap, mono_scan=mono_scan, plot=plot) + @staticmethod - def bl_energy(energy_ev, move_gap=True, mono_scan=True, plot=True): ... + def get_current_energy() -> float: + return get_current_energy() + @staticmethod - def get_current_energy(): ... - @staticmethod - def mono_pitch_scan(plot=True): ... + def mono_pitch_scan(plot=True): + return mono_pitch_scan(plot) class X10saDispatch(DefaultDispatch): diff --git a/src/aare/devices/bec_worker.py b/src/aare/devices/bec_worker.py index 2b3d3384..25c4c1fa 100644 --- a/src/aare/devices/bec_worker.py +++ b/src/aare/devices/bec_worker.py @@ -2,17 +2,12 @@ import sys import time from enum import Enum, StrEnum -from aarecommon.config.beamline import cfg_get, mx_beamline from aarecommon.config.logger import setup_logger from aarecommon.config.logger_events import log_timing from aarecommon.errors.exception_handler import BECCommunicationError, BECException from aarecommon.models.beamline import MXBeamline -from bec_ipython_client import BECIPythonClient -from bec_ipython_client.signals import OperationMode from bec_lib.device import RPCError, ScanRequestError from bec_lib.devicemanager import DeviceContainer -from bec_lib.procedures.helper import FrontendProcedureHelper -from bec_lib.service_config import ServiceConfig from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch @@ -53,9 +48,9 @@ class BeamlineState(str, Enum): class BecDevices(StrEnum): DET_Z = "det_z" FL_BRIGHT = "fl_bright" - BL_BRIGHT = "bl_bright" - BL_POS = "bl_pos" - SCAM_ZOOM = "scam_zoom" + BACKLIGHT_BRIGHT = "bl_bright" + BACKLIGHT_POS = "bl_pos" + SAMCAM_ZOOM = "scam_zoom" CRYO_POS = "cryo_pos" CRYO_TEMP = "cryo_temp" CRYO_X = "cryo_x" @@ -71,44 +66,28 @@ class MissingDevicesError(BECException): ... class BECClientWorker: def __init__(self, beamline: MXBeamline, name: str = "default"): self.dispatch = get_beamline_dispatch() - - if beamline is MXBeamline.SIMULATED: - self.simulated = True - - else: - self.simulated = False - logger.debug( - f"Initializing BECClientWorker for {self.dispatch.config.bl_lower} beamline" + self.client = self.dispatch.bec_setup.create_client() + self.client.start() + self.dev: DeviceContainer = self.client.device_manager.devices + missing_devices = {str(d.value) for d in list(BecDevices) if d not in self.dev} + if missing_devices: + raise MissingDevicesError( + f"{missing_devices} were not found in the BEC session. Available: {list(self.dev.keys())}. If no devices are available, load a BEC device config." ) - host = cfg_get( - "daq.hardware.bec_url", f"{self.dispatch.config.bl_lower}-bec-001.psi.ch" - ) - service_config = ServiceConfig(redis={"host": host, "port": 6379}) - service_config.config["log_writer"]["base_path"] = "/tmp/logs" - self.client = BECIPythonClient(config=service_config, mode=OperationMode.Procedure) - self.client.start() - self.dev: DeviceContainer = self.client.device_manager.devices - missing_devices = {str(d) for d in list(BecDevices) if d not in self.dev} - if missing_devices: - raise MissingDevicesError( - f"{missing_devices} were not found in the BEC session. Available: {list(self.dev.keys())}. If no devices are available, load a BEC device config." - ) - self.scans = self.client.scans - self.macros = self.dispatch.bec_macros - self.helper = FrontendProcedureHelper(self.client.connector) - self._set_scilog_tags() - try: - self._init_beamline_environment() - except Exception: - logger.exception("Error initialising BEC devices") - sys.exit(1) - logger.debug(f"simulated is {self.simulated}") + self.scans = self.client.scans + self.macros = self.dispatch.bec_macros + self._set_scilog_tags() + try: + self._init_beamline_environment() + except Exception: + logger.exception("Error initialising BEC devices") + sys.exit(1) def _init_beamline_environment(self): self.position_devices, self.planner = self.macros.init_beamline_environment() self._backlight_brightness = self.position_devices["bl_bright"] self._frontlight_brightness = self.position_devices["fl_bright"] - self._zoom = self.dev[BecDevices.SCAM_ZOOM] + self._zoom = self.dev[BecDevices.SAMCAM_ZOOM] self._ring_current = self.dev.sls_current def _bec_error( @@ -218,11 +197,7 @@ class BECClientWorker: @log_timing(logger, "BEC move_to") def move_to(self, state: BeamlineState): start = time.perf_counter() - logger.debug(f"simulated is {self.simulated}") logger.info(f"BEC move_to requested: {state.value}") - if self.simulated: - logger.debug(f"Simulating move to {state.value}") - return True try: self.planner.move_to(state) if self.planner.is_state(state): @@ -239,40 +214,17 @@ class BECClientWorker: raise self._bec_error(e, operation=f"planner.move_to:{state.value}") from e def is_state(self, state: BeamlineState): - if self.simulated: - logger.debug(f"Simulating check_beamline_state: {state.value}") - return True return self.planner.is_state(state) def current_state(self): - if self.simulated: - logger.debug("Simulating check_beamline_state") - return BeamlineState.MAINTENANCE return self.planner.current_state() def list_position_devices(self): """List the position devices available for the BEC worker""" - if self.simulated or self.position_devices is None: + if self.position_devices is None: return [] return list(self.position_devices.keys()) - def show_all_devices(self): - return self.dev.show_all - - def list_all_user_macros(self) -> list[str] | None: - if self.simulated: - logger.debug("Simulating list_all_user_macros") - return [] - try: - self._list_all_macros() - raw_macros = [name for name, _ in self.client.macros._update_handler.macros.items()] - if raw_macros is None: - logger.warning("BEC returned no user macros; treating as empty list") - return [] - return [str(macro) for macro in raw_macros] - except Exception as e: - raise self._bec_error(e, operation="list_all_user_macros") from e - def reinitialise_planner_and_position_devices(self, method: str = "auto"): """ Reinitialise BEC position devices and planner. @@ -286,9 +238,6 @@ class BECClientWorker: Returns: List of position device names after reinitialisation. """ - if self.simulated: - logger.debug(f"Simulating reinitialise_planner_and_position_devices(method={method})") - return [] try: self.client.config.update_session_with_file(self.dispatch.config.device_yml_path) self._init_beamline_environment() @@ -342,7 +291,7 @@ class BECClientWorker: logger.warning( f"Energy change may have failed, current energy: {self.check_current_energy()} eV" ) - if beamline is MXBeamline.X10SA: + if self.dispatch.config.bl_upper == MXBeamline.X10SA: additonal_text = [ f"New dcm_bragg position: {self.dev.dcm_bragg.position:4g} mrad", f"New dcm_pitch position: {self.dev.dcm_pitch.position:4g} ", @@ -415,15 +364,12 @@ class BECClientWorker: @backlight_brightness.setter def backlight_brightness(self, value: int | str): """Set the backlight brightness to the specified value""" - if self.simulated: - return try: self._backlight_brightness.move(value) except Exception as e: raise self._bec_error( e, operation=f"backlight_brightness:{value}", tags=["backlight"] ) from e - raise def get_backlight_pos(self) -> BrightnessEnum: """Returns the current backlight brightness position""" @@ -431,8 +377,6 @@ class BECClientWorker: def backlight_toggle(self): """Turn the backlight on or off""" - if self.simulated: - return try: current = self.get_backlight_pos() if current is BrightnessEnum.ON: @@ -478,132 +422,3 @@ class BECClientWorker: except (RPCError, ScanRequestError) as e: logger.warning(f"Ring current BEC device is not available: {e}") return 0 - - -if __name__ == "__main__": - import time - - print(time.ctime(), " starting BEC Client") - beamline = mx_beamline() - try: - client = BECClientWorker(beamline) - except Exception: - logger.exception("Failed to start BEC client") - try: - client.shutdown_client() - except Exception: - logger.exception("Failed to shut the BEC client down cleanly") - sys.exit(1) - # print(client.get_det_cov(actual=True)) - # print(client.is_state(BeamlineState.ROBOT_SAMPLE_EXCHANGE)) - try: - print("startting backlight brightness test") - print("initial value") - # print(client.client.show_last_alarm()) - # print(client._raise_bec_error(exc=Exception("test"), operation="test", tags=["test"])) - print(client.ring_current) - for i in range(10): - print(client._ring_current.get()) - time.sleep(1) - # print('setting to 5') - # client.backlight_brightness = 5 - # print(client.get_backlight_pos()) - # print(client.backlight_brightness) - # print('setting to 1.2') - # client.backlight_brightness = 1.2 - # print(client.get_backlight_pos()) - # print(client.backlight_brightness) - # print('setting to ON') - # client.backlight_brightness = BrightnessEnum.ON - # print(client.get_backlight_pos()) - # print(client.backlight_brightness) - # print('setting to 1.0') - # client.backlight_brightness = 1.0 - # print(client.get_backlight_pos()) - # print(client.backlight_brightness) - # print('setting to 0.0') - # client.backlight_brightness = 0.0 - # print(client.get_backlight_pos()) - # print(client.backlight_brightness) - # print('setting to OFF') - # client.backlight_brightness = BrightnessEnum.OFF - # client.backlight_toggle() - # print(client.get_backlight_pos()) - # print(client.backlight_brightness) - # client.scilog_msg("Testing scilog messages with color = yellow and italic", italic=True, - # color="green", warning=False) - except Exception: - client.shutdown_client() - logger.exception("BEC client smoke test failed") - - # try: - # det_value = 980 - # print(f"moving detector to vale:{det_value}") - # start = time.perf_counter() - # # print(client.get_det_z()) - # # client.det_z(value=det_value, timeout=10) - # # print(time.perf_counter() - start) - # # print(client.get_det_z) - # # det_value = 985 - # # print(f"moving detector to vale:{det_value}") - # status = client.det_z(value=det_value) - # status.wait(timeout=5) - # print(status) - # print(client.get_det_z()) - # #client.mono_pitch_scan_runner() - # #client.change_energy(12000) - # - # except Exception as e: - # print(f"Error: {e}") - # # client.planner.current_state() - # client.planner.move_to(BeamlineState.SAMPLE_ALIGNMENT) - # time.sleep(2.0) - # client.planner.current_state() - # client.move_to(BeamlineState.SAMPLE_ALIGNMENT) - # time.sleep(2.0) - # client.planner.current_state() - - # client.load_user_macros() - # client.macros.mono_pitch_scan(False) - # status = client.run_macro("planner.current_state", queue="default") - # client.planner.move_to('manual_sample_exchange') - # print(status) - # status.wait() - # print(status) - # try: - # status=client.proc.run_macro("mono_pitch_scan", queue="test") - # planner.move_to('manual_sample_exchange') - - # print(status) - # status.wait() - # status.cancel() - # print(status) - # client.client.macros.mono_pictch_scan(False) - # try: - # # a=client.a2e_runner(160, "iln") - # # print(a) - # # print(convert_from_energy(12)) - # # energy = get_current_energy() - # # pos = get_dcm_motors_positions(energy) - # # print(energy, pos) - # print(bs_z_policy(15.0)) - # #client.scans.umv(client.dev.xeye_x, 0, relative=False) - # #client.mono_pitch_scan_runner() - # #client.mono_pitch_scan_runner() - # #b=client.run_macro_blocked("a2e", 160, "iln", queue="test") - # #b = client.run_macro_blocked("mono_pitch_scan", False, queue="default") - # - # #status = client.mono_pitch_scan_runner - # #print(status) - # - # #client.rse2sa() - # #time.sleep(10.0) - # #time.sleep(10.0) - # #print(status) - # #client.common2rse() - # #print(status) - # except GuardViolation as e: - # print(f"GuardViolation: {e}") - # except RuntimeError as e: - # print(f"RuntimeError: {e}") - client.shutdown_client() diff --git a/src/aare/devices/mx_lib.py b/src/aare/devices/mx_lib.py deleted file mode 100644 index c0807eda..00000000 --- a/src/aare/devices/mx_lib.py +++ /dev/null @@ -1,289 +0,0 @@ -import re -import time -from collections.abc import Callable -from typing import Any - -from aarecommon.config.logger import setup_logger -from epics import PV, Motor, poll - -logger = setup_logger("aareDAQ") - - -def wait_for_movement_to_finish(*motors): - """ - Wait for all {motors} passed in argument to finish movement. - - Requirements: epics.Motor -- each {motor} is an instance of this - class. - - The overall timeout is the maximum of (.RBV - .VAL)/.VELO for each - concerned {motor}. - - Returns: nothing - """ - poll(0.3) - longest = 0.0 - for motor in motors: - time_to_target = motor.readback / motor.slew_speed - longest = max(longest, time_to_target) - - timeout = time.time() + 1.5 * longest - done = False - while not done and time.time() < timeout: - done = all(m.done_moving for m in motors) - - if time.time() > timeout: - print("TIMEOUT waiting for motors to be done moving; current motor positions:") - for m in motors: - print(f"...{m._prefix[:-1]} ({m.units}) VAL={m.drive:.1f}; RBV={m.readback:.1f}") - raise TimeoutError("timeout waiting for movement to finish.") - - -class ValueWaitTimeout(Exception): - pass - - -def pv_wait( - pv: PV | Motor, - value: Any, - *, - timeout: float = 60.0, - polling: float = 0.1, - tolerance: float | None = None, - verbose: bool = False, -): - """wait until an epics.PV reaches a value - pv: epics.PV | epics.Motor - the PV on which you want to wait on - value/target: any - this value depends on the PV type: str, enum, double, ... - :timeout: float default = 60.0 - timeout in seconds - :polling: float default = 0.1 - polling interval in seconds - tolerance: float or None - provide a tolerance to accept when comparing values, currently None by default, however certain PVs and motors - have inbuilt tolerances that can be used instead - verbose: bool default = False - NotImplemented - """ - if isinstance(pv, Motor): - - def checker(m, target=None, tolerance=None): - if tolerance is None: - tolerance = m.get("RDBD") - if tolerance is None: - print(f"WARNING: motor {m._prefix[:-1]} has no RDBD, using 0.001") - tolerance = 0.001 # FIXME - if target is None: - target = m.drive - - diff = abs(target - m.readback) - return m.done_moving and diff < tolerance - - wait_motor_position( - pv, checker, target=value, tolerance=tolerance, timeout=timeout, polling=polling - ) - - elif is_epics_type(pv, "double") or is_epics_type(pv, "time_double"): - wait_float_condition(pv, value, tolerance=tolerance, timeout=timeout, polling=polling) - - elif is_epics_type(pv, "enum") or is_epics_type(pv, "time_enum"): - wait_enum_condition(pv, value, timeout=timeout, polling=polling) - - elif is_epics_type(pv, "string") or is_epics_type(pv, "time_string"): - wait_string_condition(pv, value, timeout=timeout, polling=polling) - else: - raise ValueError(f"pv_wait() don't know how to handle {pv.pvname} of type {pv.type}") - - -def is_epics_type(pv: PV, pv_type: str) -> bool: - """Check to see if a PV is of a certain type such as double, enum, string, ...""" - if isinstance(pv_type, type): - pv_type = pv_type.__name__ - - return pv_type == pv.type - - -def wait_string_condition( - pv: PV, target: str | re.Pattern, *, timeout: float = 60.0, polling: float = 0.1 -): - """wait until an epics.PV of type string reaches target - :pv: epics.PV - PV should be of type string - :target: str or re.Pattern - :timeout: float default = 60.0 - timeout in seconds - :polling: float default = 0.1 - polling interval in seconds - """ - - if not (isinstance(pv, PV) and "string" not in pv.type): - raise AttributeError("argument 'pv' must be an epics.PV of type string") - - if not isinstance(target, re.Pattern): - target = re.compile(f"^{target}$", re.IGNORECASE) - - tout = time.time() + timeout - - while time.time() < tout: - if target.match(pv.char_value): - return - poll(polling) - - if time.time() > tout: - raise TimeoutError( - f"timeout waiting for string {pv.pvname} == {target}; actual value == {pv.char_value}" - ) - - -def wait_float_condition( - pv: PV, - value: float, - *, - timeout: float = 60.0, - polling: float = 0.1, - tolerance: float | None = None, -): - """wait until an epics.PV of type double reaches value - pv: epics.PV - the PV enum on which you want to wait on - - value: float - the target value - - polling: float - how often pv is checked during wait loop - - tolerance: float or None - the tolerance to accept when comparing values, if None (default) we - try to figure an appropriate value - - - timeout: float default = 60.0 - a timeout in seconds - - return: nothing - - raises: TimeoutError if a timeout occurs - """ - - try: - value = float(value) - except ValueError: - raise AttributeError("argument 'value' must be a number") - - if tolerance is None: - # If pv.precision is missing/None, fall back to a default - precision = getattr(pv, "precision", None) - if precision is None: - tolerance = 1e-6 - else: - tolerance = pow(10, -(precision - 1)) - - tout = time.time() + timeout - - while time.time() < tout: - if abs(pv.value - value) < tolerance: - return - poll(polling) - - if time.time() > tout: - raise TimeoutError(f"timeout waiting for enum {pv.pvname} == {value}") - - -def wait_motor_position( - motor: Motor, tester: Callable, *, timeout: float = 50.0, polling: float = 0.1, **kwargs -): - """wait until an epics.Motor reaches a condition - motor: epics.Motor - the motor on which you want to wait on - - tester: Callable - a Callable which is passed motor and **kwargs - => Callable(motor, **kwargs) - - timeout: double - a timeout in seconds - - **kwargs: extra arguments to tester - - return: nothing - - raises: TimeoutError if a timeout occurs - """ - if not callable(tester): - raise TypeError("argument 'tester' must be a function") - - try: - move_time = abs(motor.drive - motor.readback) / motor.speed - except Exception: - logger.debug("Could not compute the motor move time; using the 1 s default", exc_info=True) - move_time = 1.0 # in case of unusual motor record - - tout = move_time + time.time() + timeout - - while time.time() < tout: - if tester(motor, **kwargs): - return - poll(polling) - - if time.time() > tout: - raise TimeoutError( - f"timeout waiting for a condition on {motor} {motor.drive} != {motor.readback}" - ) - - -def wait_enum_condition( - pv: PV, value: str | int | re.Pattern, *, timeout: float = 60.0, polling=0.1 -): - """wait until an epics.PV enum reaches value - pv: epics.PV - the PV enum on which you want to wait on - - value: int or str - if int check against index of current value - - timeout: double - a timeout in seconds - - return: nothing - - raises: TimeoutError if a timeout occurs - """ - if not (isinstance(pv, PV) and pv.type.lower().endswith("enum")): - raise AttributeError("argument 'pv' must be an epics.PV of type enum") - - if not (isinstance(value, (str, int, re.Pattern))): - raise TypeError("argument 'value' must be either an int, str, or re.Pattern") - - if type(value) is int: - tester = lambda pv: value == pv.get() - elif type(value) is str: - tester = lambda pv: str(value) == pv.get(as_string=True).lower() - value = str(value).lower() # it's already a str :-/ - elif isinstance(value, re.Pattern): - tester = lambda pv: value.match(pv.get(as_string=True)) - else: - raise AttributeError("argument 'value' must be either an int, str, or re.Pattern") - - tout = time.time() + timeout - - while time.time() < tout: - if tester(pv): - return - poll(polling) - - if time.time() > tout: - raise TimeoutError( - f"timeout waiting for enum {pv.pvname} == {value}," - f"current value is {pv.get(as_string=True)}" - ) - - -def clean_filename(filename: str) -> str: - cleaned = re.sub(r"[^A-Za-z0-9._-]", "_", filename.strip()) - cleaned = cleaned.strip("._-") - if not cleaned: - raise ValueError("Filename is empty after sanitization.") - return cleaned diff --git a/tests/unit/devices/test_mx_lib.py b/tests/unit/devices/test_mx_lib.py deleted file mode 100644 index 4baeb224..00000000 --- a/tests/unit/devices/test_mx_lib.py +++ /dev/null @@ -1,92 +0,0 @@ -from unittest.mock import MagicMock, patch - -import pytest -from epics import PV, Motor - -from aare.devices.mx_lib import clean_filename, is_epics_type, pv_wait, wait_for_movement_to_finish - - -def test_clean_filename(): - assert clean_filename("my file!.txt") == "my_file_.txt" - assert clean_filename(" /path/to/somewhere/file.txt ") == "path_to_somewhere_file.txt" - assert clean_filename("abcABC123") == "abcABC123" - assert clean_filename("file._-") == "file" - - with pytest.raises(ValueError): - clean_filename("!!!") - - -def test_is_epics_type(): - mock_pv = MagicMock(spec=PV) - mock_pv.type = "double" - assert is_epics_type(mock_pv, "double") is True - assert is_epics_type(mock_pv, "enum") is False - - class FakeType: - pass - - assert is_epics_type(mock_pv, FakeType) is False - - -@patch("aare.devices.mx_lib.poll") -def test_wait_for_movement_to_finish(mock_poll): - mock_motor = MagicMock(spec=Motor) - mock_motor.readback = 10.0 - mock_motor.slew_speed = 1.0 - mock_motor.done_moving = True - - wait_for_movement_to_finish(mock_motor) - assert mock_poll.called - - -@patch("aare.devices.mx_lib.poll") -def test_wait_for_movement_to_finish_timeout(mock_poll): - mock_motor = MagicMock(spec=Motor) - mock_motor.readback = 10.0 - mock_motor.slew_speed = 1.0 - mock_motor.done_moving = False - mock_motor.drive = 11.0 - mock_motor.units = "mm" - mock_motor._prefix = "MOT1:" - - with patch("time.time", side_effect=[0, 0, 100, 101]), pytest.raises(TimeoutError): - wait_for_movement_to_finish(mock_motor) - - -@patch("aare.devices.mx_lib.wait_motor_position") -def test_pv_wait_motor(mock_wait_motor): - mock_motor = MagicMock(spec=Motor) - pv_wait(mock_motor, 10.0) - mock_wait_motor.assert_called_once() - - -@patch("aare.devices.mx_lib.wait_float_condition") -def test_pv_wait_double(mock_wait_float): - mock_pv = MagicMock(spec=PV) - mock_pv.type = "double" - pv_wait(mock_pv, 10.0) - mock_wait_float.assert_called_once() - - -@patch("aare.devices.mx_lib.wait_enum_condition") -def test_pv_wait_enum(mock_wait_enum): - mock_pv = MagicMock(spec=PV) - mock_pv.type = "enum" - pv_wait(mock_pv, "READY") - mock_wait_enum.assert_called_once() - - -@patch("aare.devices.mx_lib.wait_string_condition") -def test_pv_wait_string(mock_wait_string): - mock_pv = MagicMock(spec=PV) - mock_pv.type = "string" - pv_wait(mock_pv, "hello") - mock_wait_string.assert_called_once() - - -def test_pv_wait_unknown_type(): - mock_pv = MagicMock(spec=PV) - mock_pv.type = "unknown" - mock_pv.pvname = "TEST:PV" - with pytest.raises(ValueError): - pv_wait(mock_pv, 1.0) -- 2.54.0 From ae930152ecca5f8cff9438b32a6b1a50a5e643eb Mon Sep 17 00:00:00 2001 From: perl_d Date: Fri, 7 Aug 2026 11:11:16 +0200 Subject: [PATCH 4/8] chore: bump aarecommon --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 22d48a80..9fdd5a39 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ requires-python = ">=3.11" dependencies = [ "uv", "gunicorn", - "aarecommon>=0.2.2", + "aarecommon>=0.3", "pydantic>=2.11", "numpy", "jfjoch_client==1.0.0rc146", -- 2.54.0 From 1c54bc6a44d9d57c8e3088203d14d88aed25a918 Mon Sep 17 00:00:00 2001 From: perl_d Date: Fri, 7 Aug 2026 12:13:27 +0200 Subject: [PATCH 5/8] feat: simulated bec backend --- .../simulated/beamline_dispatch.py | 13 ++++- src/aare/daq/operations/screenshot/service.py | 2 +- src/aare/devices/bec_definitions.py | 53 +++++++++++++++++ src/aare/devices/bec_worker.py | 58 +++---------------- 4 files changed, 71 insertions(+), 55 deletions(-) create mode 100644 src/aare/devices/bec_definitions.py diff --git a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py index b2521b1e..6ff05f59 100644 --- a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py @@ -3,9 +3,11 @@ from unittest.mock import MagicMock from aarecommon.config.beamline import MXBeamline from bec_ipython_client import BECIPythonClient +from bec_lib.devicemanager import DeviceContainer from aare.beamline_dispatch.default.beamline_dispatch import DefaultConfigKeys, DefaultDispatch from aare.beamline_dispatch.protocols import AuthDispatch, BecMacros, BecSetup +from aare.devices.bec_definitions import BecDevices class SimulatedAuthDispatch(AuthDispatch): @@ -16,13 +18,13 @@ class SimulatedAuthDispatch(AuthDispatch): class SimulatedBecMacros(BecMacros): @staticmethod def save_and_reload() -> tuple[Any, Any]: - return (None, None) + return (MagicMock(), MagicMock()) @staticmethod def save_current_position(device, position, axis=None, force=False, max_delta=0.5) -> None: ... @staticmethod def init_beamline_environment() -> tuple[Any, Any]: - return (None, None) + return (MagicMock(), MagicMock()) @staticmethod def bl_energy(energy_ev, move_gap=True, mono_scan=True, plot=True): ... @@ -38,13 +40,18 @@ class SimulatedBecSetup(BecSetup): self._beamline = beamline def create_client(self) -> BECIPythonClient: - return MagicMock(spec=BECIPythonClient) + mock_client = MagicMock() + mock_client.device_manager.devices = DeviceContainer() + for dev in [d.value for d in list(BecDevices)]: + mock_client.device_manager.devices[dev] = MagicMock() + return mock_client class SimulatedDispatch(DefaultDispatch): def __init__(self) -> None: # pyright: ignore[reportMissingSuperCall] self._auth = SimulatedAuthDispatch() self._bec_macros = SimulatedBecMacros() + self._bec_setup = SimulatedBecSetup(MXBeamline.SIMULATED) self._config = DefaultConfigKeys(MXBeamline.SIMULATED) @property diff --git a/src/aare/daq/operations/screenshot/service.py b/src/aare/daq/operations/screenshot/service.py index 13d20dac..34afdc2b 100644 --- a/src/aare/daq/operations/screenshot/service.py +++ b/src/aare/daq/operations/screenshot/service.py @@ -7,10 +7,10 @@ from typing import Protocol import cv2 import numpy as np from aarecommon.models.models import SampleShortInfo +from aarecommon.utils.paths import clean_filename from aare.daq.aaredb import AareWrapper from aare.daq.mlbox import MlBox -from aare.devices.mx_lib import clean_filename class NonCriticalRunner(Protocol): diff --git a/src/aare/devices/bec_definitions.py b/src/aare/devices/bec_definitions.py new file mode 100644 index 00000000..42adaa13 --- /dev/null +++ b/src/aare/devices/bec_definitions.py @@ -0,0 +1,53 @@ +from enum import StrEnum + +from aarecommon.errors.exception_handler import BECException + + +class DetectorCoverEnum(StrEnum): + """Enum for the detector cover position + position devices can take string or number to move + Currently dictated string in DAQ for consitency""" + + OPEN = "open" # 2 + CLOSED = "closed" # 1 + + +class BrightnessEnum(StrEnum): + """Enum for the backlight brightness + position devices can take string or number to move + """ + + ON = "on" + OFF = "off" + + +class BeamlineState(StrEnum): + ROBOT_SAMPLE_EXCHANGE = "robot_sample_exchange" + SAMPLE_ALIGNMENT = "sample_alignment" + DATA_COLLECTION = "data_collection" + DC_XRF = "DC_XRF" + MANUAL_SAMPLE_EXCHANGE = "manual_sample_exchange" + BEAM_VISUALISATION = "beam_visualisation" + FLUX_MEASUREMENT = "flux_measurement" + BEAMSTOP_ALIGNMENT = "beamstop_alignment" + MAINTENANCE = "maintenance" + XTAL_SNAPSHOT = "xtal_snapshot" + + +class BecDevices(StrEnum): + DET_Z = "det_z" + FL_BRIGHT = "fl_bright" + BACKLIGHT_BRIGHT = "bl_bright" + BACKLIGHT_POS = "bl_pos" + SAMCAM_ZOOM = "scam_zoom" + CRYO_POS = "cryo_pos" + CRYO_TEMP = "cryo_temp" + CRYO_X = "cryo_x" + TRANSM = "transm" + BCU_SHUTTER = "bcu_shutter" + MAG_OK = "mag_ok" + MAG_RBV = "mag_rbv" + SLS_CURRENT = "sls_current" + + +class MissingDevicesError(BECException): ... diff --git a/src/aare/devices/bec_worker.py b/src/aare/devices/bec_worker.py index 25c4c1fa..0f351d9c 100644 --- a/src/aare/devices/bec_worker.py +++ b/src/aare/devices/bec_worker.py @@ -1,68 +1,24 @@ import sys import time -from enum import Enum, StrEnum from aarecommon.config.logger import setup_logger from aarecommon.config.logger_events import log_timing -from aarecommon.errors.exception_handler import BECCommunicationError, BECException +from aarecommon.errors.exception_handler import BECCommunicationError from aarecommon.models.beamline import MXBeamline from bec_lib.device import RPCError, ScanRequestError from bec_lib.devicemanager import DeviceContainer from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch +from aare.devices.bec_definitions import ( + BeamlineState, + BecDevices, + BrightnessEnum, + MissingDevicesError, +) logger = setup_logger("aareDAQ") -class DetectorCoverEnum(str, Enum): - """Enum for the detector cover position - position devices can take string or number to move - Currently dictated string in DAQ for consitency""" - - OPEN = "open" # 2 - CLOSED = "closed" # 1 - - -class BrightnessEnum(str, Enum): - """Enum for the backlight brightness - position devices can take string or number to move - """ - - ON = "on" - OFF = "off" - - -class BeamlineState(str, Enum): - ROBOT_SAMPLE_EXCHANGE = "robot_sample_exchange" - SAMPLE_ALIGNMENT = "sample_alignment" - DATA_COLLECTION = "data_collection" - DC_XRF = "DC_XRF" - MANUAL_SAMPLE_EXCHANGE = "manual_sample_exchange" - BEAM_VISUALISATION = "beam_visualisation" - FLUX_MEASUREMENT = "flux_measurement" - BEAMSTOP_ALIGNMENT = "beamstop_alignment" - MAINTENANCE = "maintenance" - XTAL_SNAPSHOT = "xtal_snapshot" - - -class BecDevices(StrEnum): - DET_Z = "det_z" - FL_BRIGHT = "fl_bright" - BACKLIGHT_BRIGHT = "bl_bright" - BACKLIGHT_POS = "bl_pos" - SAMCAM_ZOOM = "scam_zoom" - CRYO_POS = "cryo_pos" - CRYO_TEMP = "cryo_temp" - CRYO_X = "cryo_x" - TRANSM = "transm" - BCU_SHUTTER = "bcu_shutter" - MAG_OK = "mag_ok" - MAG_RBV = "mag_rbv" - - -class MissingDevicesError(BECException): ... - - class BECClientWorker: def __init__(self, beamline: MXBeamline, name: str = "default"): self.dispatch = get_beamline_dispatch() -- 2.54.0 From 3af118fb644f52ac72d9bf6f5ba815cbf75c6a5e Mon Sep 17 00:00:00 2001 From: perl_d Date: Fri, 7 Aug 2026 15:12:12 +0200 Subject: [PATCH 6/8] wrap mlbox --- .../beamline_dispatch/default/beamline_dispatch.py | 5 +++++ src/aare/beamline_dispatch/protocols.py | 4 ++++ .../simulated/beamline_dispatch.py | 13 +++++++++++-- src/aare/daq/daq.py | 4 +++- 4 files changed, 23 insertions(+), 3 deletions(-) diff --git a/src/aare/beamline_dispatch/default/beamline_dispatch.py b/src/aare/beamline_dispatch/default/beamline_dispatch.py index a2830190..2caf493d 100644 --- a/src/aare/beamline_dispatch/default/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/default/beamline_dispatch.py @@ -8,6 +8,7 @@ from bec_ipython_client.signals import OperationMode from bec_lib.client import ServiceConfig from aare.beamline_dispatch.protocols import AuthDispatch, BeamlineDispatch, BecSetup, ConfigKeys +from aare.daq.mlbox import MlBox logger = setup_logger("aareDAQ") @@ -71,6 +72,7 @@ class DefaultDispatch(BeamlineDispatch, ABC): Should be safe and fail rather than assuming anything.""" def __init__(self, beamline: MXBeamline) -> None: + self._beamline = beamline self._auth = DefaultAuthDispatch() self._config = DefaultConfigKeys(beamline) self._bec_setup = DefaultBecSetup(beamline) @@ -86,3 +88,6 @@ class DefaultDispatch(BeamlineDispatch, ABC): @property def bec_setup(self): return self._bec_setup + + def get_mlbox(self) -> MlBox: + return MlBox(self._beamline) diff --git a/src/aare/beamline_dispatch/protocols.py b/src/aare/beamline_dispatch/protocols.py index 61f2ebcc..fd649bcc 100644 --- a/src/aare/beamline_dispatch/protocols.py +++ b/src/aare/beamline_dispatch/protocols.py @@ -4,6 +4,8 @@ from typing import Any from aarecommon.models.models import MXBeamline from bec_ipython_client import BECIPythonClient +from aare.daq.mlbox import MlBox + class AuthDispatch(ABC): @abstractmethod @@ -63,3 +65,5 @@ class BeamlineDispatch(ABC): @property @abstractmethod def config(self) -> ConfigKeys: ... + @abstractmethod + def get_mlbox(self) -> MlBox: ... diff --git a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py index 6ff05f59..01634238 100644 --- a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py @@ -1,14 +1,20 @@ -from typing import Any +from __future__ import annotations + +from typing import TYPE_CHECKING, Any from unittest.mock import MagicMock from aarecommon.config.beamline import MXBeamline -from bec_ipython_client import BECIPythonClient from bec_lib.devicemanager import DeviceContainer from aare.beamline_dispatch.default.beamline_dispatch import DefaultConfigKeys, DefaultDispatch from aare.beamline_dispatch.protocols import AuthDispatch, BecMacros, BecSetup from aare.devices.bec_definitions import BecDevices +if TYPE_CHECKING: + from bec_ipython_client import BECIPythonClient + + from aare.daq.mlbox import MlBox + class SimulatedAuthDispatch(AuthDispatch): def get_jwt_key(self) -> str: @@ -57,3 +63,6 @@ class SimulatedDispatch(DefaultDispatch): @property def bec_macros(self): return self._bec_macros + + def get_mlbox(self) -> MlBox: + return MagicMock() diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index bf6acedc..27b5f94b 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -67,6 +67,7 @@ from aarecommon.models.rotation_scan import CompletedRotationScan, RotationScanR from aarecommon.models.tell import TellPhaseEnum, TellStateModel from aareDB import SampleEventType +from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch from aare.daq import workflows from aare.daq.aaredb import AareWrapper from aare.daq.config import ABR_POS_MOUNT, BeamlineConfig, BeamlineStateEnum @@ -273,8 +274,9 @@ class AareDAQ: def __init__(self, cfg: BeamlineConfig, bl: MXBeamline): self.last_time = 0.0 self._cfg = cfg + self._dispatch = get_beamline_dispatch() + self._mlbox = self._dispatch.get_mlbox() self._devs = BeamlineDevices(bl) - self._mlbox = MlBox(bl) self._jfjoch = JFJochWrapper(bl) self._bl = bl.value.upper() self._beamline = bl -- 2.54.0 From 5b8728c37dd6cf089107a893170a74ca7c188dbb Mon Sep 17 00:00:00 2001 From: perl_d Date: Fri, 7 Aug 2026 15:13:52 +0200 Subject: [PATCH 7/8] DNM MOCK ALL OTHER DEVICES --- src/aare/daq/daq.py | 3 ++- src/aare/daq/devices.py | 13 +++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 27b5f94b..ec384869 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -6,6 +6,7 @@ from collections.abc import Callable from datetime import UTC, datetime from math import ceil from pathlib import Path +from unittest.mock import MagicMock from aarecommon.config.beamline import cfg_get from aarecommon.config.logger import setup_logger @@ -277,7 +278,7 @@ class AareDAQ: self._dispatch = get_beamline_dispatch() self._mlbox = self._dispatch.get_mlbox() self._devs = BeamlineDevices(bl) - self._jfjoch = JFJochWrapper(bl) + self._jfjoch = MagicMock() # JFJochWrapper(bl) self._bl = bl.value.upper() self._beamline = bl self._aare = AareWrapper(bl) diff --git a/src/aare/daq/devices.py b/src/aare/daq/devices.py index cf4ec4be..11c29fa2 100644 --- a/src/aare/daq/devices.py +++ b/src/aare/daq/devices.py @@ -1,5 +1,6 @@ # Abstractions of devices for beamline import time +from unittest.mock import MagicMock # Each "standard" device needs three elements: # - property to read device value @@ -26,20 +27,20 @@ class BeamlineDevices: def __init__(self, beamline: MXBeamline): self._beamline = beamline BEAMLINE = beamline.value.upper() - self.tell = make_tell_client(beamline) - self.aerotech = aerotech.AerotechController(beamline) + self.tell = MagicMock() # make_tell_client(beamline) + self.aerotech = MagicMock() # aerotech.AerotechController(beamline) logger.debug("initialising BEC worker") self.bec_worker = BECClientWorker(beamline) logger.debug("initialising BEC worker done") - self._smargon = smargon.Smargon(beamline) - self.exp_shutter = ExperimentalHutchShutter(beamline=self._beamline) + self._smargon = MagicMock() # smargon.Smargon(beamline) + self.exp_shutter = MagicMock() # ExperimentalHutchShutter(beamline=self._beamline) self.dtz_mod = cfg_get("daq.detector_distance_limit_modifier", 1.0) # TODO convert epics pvs to BEC -> need to get all samcam signals into a device # BEC device for pss state also good # Personnel Safety System: gates whether the robot is allowed to move. - self.pss = PssState(beamline=self._beamline) - self._sample_cam = epicsAD(f"{BEAMLINE}-ES-MS:") + self.pss = MagicMock() # PssState(beamline=self._beamline) + self._sample_cam = MagicMock() # epicsAD(f"{BEAMLINE}-ES-MS:") def restart_bec_worker(self, simulated: bool = False) -> None: try: -- 2.54.0 From 948b4774be1f4fb12721e8e4ce8b71809f900e3c Mon Sep 17 00:00:00 2001 From: perl_d Date: Fri, 7 Aug 2026 16:09:13 +0200 Subject: [PATCH 8/8] guard macro imports --- src/aare/beamline_dispatch/protocols.py | 3 ++ .../simulated/beamline_dispatch.py | 3 ++ .../x10sa/beamline_dispatch.py | 36 ++++++++++++------- src/aare/daq/daq.py | 1 - src/aare/daq/devices.py | 17 ++++----- src/aare/daq/server.py | 3 +- src/aare/devices/bec_definitions.py | 3 ++ src/aare/devices/bec_worker.py | 3 ++ 8 files changed, 43 insertions(+), 26 deletions(-) diff --git a/src/aare/beamline_dispatch/protocols.py b/src/aare/beamline_dispatch/protocols.py index fd649bcc..52ee552d 100644 --- a/src/aare/beamline_dispatch/protocols.py +++ b/src/aare/beamline_dispatch/protocols.py @@ -18,6 +18,9 @@ class BecMacros(ABC): def save_and_reload() -> tuple[Any, Any]: ... @staticmethod @abstractmethod + def flux() -> float: ... + @staticmethod + @abstractmethod def save_current_position(device, position, axis=None, force=False, max_delta=0.5) -> None: ... @staticmethod @abstractmethod diff --git a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py index 01634238..9bfa99e5 100644 --- a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py @@ -38,6 +38,9 @@ class SimulatedBecMacros(BecMacros): def get_current_energy() -> float: ... @staticmethod def mono_pitch_scan(plot=True): ... + @staticmethod + def flux() -> float: + return 0 class SimulatedBecSetup(BecSetup): diff --git a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py index b12a9344..5d9b976c 100644 --- a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py @@ -1,18 +1,24 @@ from typing import Any from aarecommon.config.beamline import MXBeamline -from pxii_bec.macros.build_devices import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] - save_and_reload, - save_current_position, -) -from pxii_bec.macros.init_beamline import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] - init_beamline_environment, -) -from pxii_bec.macros.pxii_energy import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] - bl_energy, - get_current_energy, - mono_pitch_scan, -) + +from aare.devices.bec_definitions import MacroImportError + +try: + from pxii_bec.macros.build_devices import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] + save_and_reload, + save_current_position, + ) + from pxii_bec.macros.init_beamline import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] + init_beamline_environment, + ) + from pxii_bec.macros.pxii_energy import ( # pyright: ignore[reportMissingImports, reportMissingModuleSource] + bl_energy, + get_current_energy, + mono_pitch_scan, + ) +except ImportError as e: + raise MacroImportError("BEC macros for pxii appear to have moved") from e from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch from aare.beamline_dispatch.protocols import BecMacros @@ -25,7 +31,7 @@ class X10SaBecMacros(BecMacros): @staticmethod def save_current_position(device, position, axis=None, force=False, max_delta=0.5) -> None: - return save_current_position + return save_current_position(device, position, axis=None, force=False, max_delta=0.5) @staticmethod def init_beamline_environment() -> tuple[Any, Any]: @@ -43,6 +49,10 @@ class X10SaBecMacros(BecMacros): def mono_pitch_scan(plot=True): return mono_pitch_scan(plot) + @staticmethod + def flux(): + return 0 # TODO: get flux calcs into stuff + class X10saDispatch(DefaultDispatch): def __init__(self, beamline: MXBeamline) -> None: diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index ec384869..82afb39f 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -74,7 +74,6 @@ from aare.daq.aaredb import AareWrapper from aare.daq.config import ABR_POS_MOUNT, BeamlineConfig, BeamlineStateEnum from aare.daq.config_model import LocalContactConfigModel from aare.daq.devices import BeamlineDevices -from aare.daq.mlbox import MlBox from aare.daq.operations.common.ml_bounding_box import get_ml_bounding_box from aare.daq.operations.common.runtime import DAQRuntimeState from aare.daq.operations.common.services import ( diff --git a/src/aare/daq/devices.py b/src/aare/daq/devices.py index 11c29fa2..52ddc41d 100644 --- a/src/aare/daq/devices.py +++ b/src/aare/daq/devices.py @@ -27,7 +27,9 @@ class BeamlineDevices: def __init__(self, beamline: MXBeamline): self._beamline = beamline BEAMLINE = beamline.value.upper() - self.tell = MagicMock() # make_tell_client(beamline) + self.tell = MagicMock( + get_mounted_sample=MagicMock(return_value=None) + ) # make_tell_client(beamline) self.aerotech = MagicMock() # aerotech.AerotechController(beamline) logger.debug("initialising BEC worker") self.bec_worker = BECClientWorker(beamline) @@ -74,15 +76,11 @@ class BeamlineDevices: # Transmission @property def transmission(self) -> float: - return self.bec_worker.tra + return self.bec_worker.dev.transm.get() @transmission.setter def transmission(self, value: float): - self.set_transmission(value, wait=False) - - def set_transmission(self, value: float, /, wait: bool = True): - logger.warning("Setting Transmission is untested") - self._transmission.move(value, wait=wait) + self.bec_worker.dev.transm.move(value, wait=False) # Lamp light @property @@ -131,10 +129,7 @@ class BeamlineDevices: @property def flux(self) -> float: - # TODO FLUX - if self._beamline == MXBeamline.X10SA: - return self.bec_worker.get_flux_x10sa() - return self.transmission * self.full_flux + return self.bec_worker.get_flux() @property def full_flux(self) -> float: diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index e0f5f3e7..27224967 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -7,6 +7,7 @@ import time from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from typing import Any, ClassVar +from unittest.mock import MagicMock import uvicorn from aarecommon.config.beamline import mx_beamline @@ -107,7 +108,7 @@ async def lifespan(application: FastAPI): # ── Core objects (Redis, EPICS PVs, BEC, TELL, JFJoch, etc.) ── bl = mx_beamline() - cfg = BeamlineConfig(bl) + cfg = MagicMock() # BeamlineConfig(bl) daq = AareDAQ(cfg, bl) try: diff --git a/src/aare/devices/bec_definitions.py b/src/aare/devices/bec_definitions.py index 42adaa13..e707f2e5 100644 --- a/src/aare/devices/bec_definitions.py +++ b/src/aare/devices/bec_definitions.py @@ -51,3 +51,6 @@ class BecDevices(StrEnum): class MissingDevicesError(BECException): ... + + +class MacroImportError(BECException): ... diff --git a/src/aare/devices/bec_worker.py b/src/aare/devices/bec_worker.py index 0f351d9c..36442fad 100644 --- a/src/aare/devices/bec_worker.py +++ b/src/aare/devices/bec_worker.py @@ -347,6 +347,9 @@ class BECClientWorker: tags=["backlight"], ) from e + def get_flux(self): + return self.macros.flux() + def save_current_bs_pos(self): self.macros.save_current_position(self.dev.bs_z, "safe") -- 2.54.0