From 93d98357b4358d2152ad25f192c2201dc2c2bb95 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 8 Sep 2026 10:54:49 +0200 Subject: [PATCH 1/4] feat: remove all busy-state handling --- pyproject.toml | 1 - scripts/free_busy_state.py | 5 - src/aare/daq/config.py | 64 +----- src/aare/daq/daq.py | 310 ++++++-------------------- src/aare/daq/server.py | 20 -- tests/unit/daq/test_face_detection.py | 3 - 6 files changed, 79 insertions(+), 324 deletions(-) delete mode 100644 scripts/free_busy_state.py diff --git a/pyproject.toml b/pyproject.toml index 124fa5cb..dbf57401 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,6 @@ dependencies = [ "requests", "pyepics~=3.5", "redis", - "python-redis-lock", "fastapi", "uvicorn", "aaredb>=0.83.1", diff --git a/scripts/free_busy_state.py b/scripts/free_busy_state.py deleted file mode 100644 index 27eb9e31..00000000 --- a/scripts/free_busy_state.py +++ /dev/null @@ -1,5 +0,0 @@ -from aaredaq.config import BeamlineConfig -from aaredaqlib.beamline import MXBeamline - -c = BeamlineConfig(MXBeamline.X06DA) -c.state_busy = False diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index 1d0bfdda..6c0bb1ca 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -2,16 +2,13 @@ import base64 import io import json import time -import traceback from dataclasses import asdict, is_dataclass from datetime import datetime import numpy as np import redis -import redis_lock from aarecommon.config.beamline import cfg_get from aarecommon.config.logger import setup_logger -from aarecommon.errors.exception_handler import BeamlineBusyException from aarecommon.math.coordinate import AerotechCoordinate, Coordinate from aarecommon.models.auth import ( BatonHolderInfo, @@ -40,6 +37,7 @@ from aarecommon.models.models import ( ZoomModel, zoom_manager, ) +from redis.lock import Lock as RedisLock from aare.daq.config_model import LocalContactConfigModel @@ -344,7 +342,7 @@ class BeamlineConfig: return SessionsStateEnum.OwnedByElse def try_set_active_session(self, session: int, expiry_sec: int) -> None: - with redis_lock.Lock(self._client, f"{self._bl}:active_session_lock", expire=10): + with RedisLock(self.redis, f"{self._bl}:active_session_lock", expire=10): active = self.active_session if active is None: self._client.set(f"{self._bl}:active_session", session) @@ -356,7 +354,7 @@ class BeamlineConfig: # TODO finish setting this up! def try_extend_active_session(self, session: int, expiry_sec: int) -> None: - with redis_lock.Lock(self._client, f"{self._bl}:active_session_lock", expire=10): + with RedisLock.Lock(self.redis, f"{self._bl}:active_session_lock", expire=10): active = self.active_session if active is None: raise RuntimeError("There is no active session with given id. Try again later.") @@ -378,7 +376,7 @@ class BeamlineConfig: ) def end_active_session(self, session: int) -> None: - with redis_lock.Lock(self._client, f"{self._bl}:active_session_lock", expire=10): + with RedisLock.Lock(self.redis, f"{self._bl}:active_session_lock", expire=10): active = self.active_session if active is None: return @@ -388,9 +386,9 @@ class BeamlineConfig: def force_set_active_session(self, session: int, expiry_sec: int) -> None: # Ensure that there is no active try-set for active session - with redis_lock.Lock(self._client, f"{self._bl}:active_session_lock", expire=10): - self._client.set(f"{self._bl}:active_session", session) - self._client.expire(f"{self._bl}:active_session", expiry_sec) + with RedisLock.Lock(self.redis, f"{self._bl}:active_session_lock", expire=10): + self.redis.set(f"{self._bl}:active_session", session) + self.redis.expire(f"{self._bl}:active_session", expiry_sec) # ========== BATON SYSTEM ========== @@ -463,7 +461,7 @@ class BeamlineConfig: # Can't transfer while beamline is busy # Add automation queue check here when you implement it # return not (self.state_busy or self.automation_queue_running) - return not self.state_busy + raise NotImplementedError("Checking for busy state is not implemented") def execute_baton_transfer( self, @@ -477,7 +475,7 @@ class BeamlineConfig: Atomically transfer the baton to a new holder. Use existing active_session_lock for consistency. """ - with redis_lock.Lock(self._client, f"{self._bl}:active_session_lock", expire=10): + with RedisLock(self._client, f"{self._bl}:active_session_lock", expire=10): self._client.set(f"{self._bl}:active_session", to_session) self._client.expire(f"{self._bl}:active_session", expiry_sec) self.baton_holder = BatonHolderInfo( @@ -532,31 +530,11 @@ class BeamlineConfig: else: self._client.delete(f"{self._bl}:commissioning_mode") - # Beamline state management - # Atomic check if beamline is busy and if not set state to busy - def try_set_busy(self, timeout: int | None = None): - with redis_lock.Lock(self._client, f"{self._bl}:move_state_lock", expire=10): - if self.state_busy: - raise BeamlineBusyException("Beamline is busy") - self.state_busy = True - if timeout is not None: - self._client.expire(f"{self._bl}:busy", timeout) - - def set_busy(self, target: BeamlineStateEnum, timeout: int | None = None): - self.try_set_busy(timeout=timeout) - curr_state = self.state - if curr_state != target: - self.state_busy = False - raise RuntimeError("Beamline is not in a proper state") - def start_moving( self, target: BeamlineStateEnum, timeout: int | None = None ) -> BeamlineStateEnum: - self.try_set_busy(timeout=timeout) curr_state = self.state - if target == curr_state: - self.state_busy = False - else: + if target != curr_state: self.state = BeamlineStateEnum.Moving return curr_state @@ -578,23 +556,6 @@ class BeamlineConfig: def state(self, state: BeamlineStateEnum): self._client.set(f"{self._bl}:state", state.value) - @property - def state_busy(self) -> bool: - return self._client.get(f"{self._bl}:busy") is not None - - @state_busy.setter - def state_busy(self, i: bool): - - logger.debug( - f"Busy flag switched to: {i}, at:\n{''.join(traceback.format_stack(limit=5)[:-2])}" - ) - if i: - self._client.set(f"{self._bl}:busy", "1") - else: - self._client.delete(f"{self._bl}:busy") - - # Other beamline settings - @property def tell_mount_count(self) -> int: return int(self._client.incr(f"{self._bl}:tell_mount_count")) @@ -667,12 +628,12 @@ class BeamlineConfig: @property def settings(self) -> BeamlineSettingsModel: - with redis_lock.Lock(self._client, f"{self._bl}:settings_lock", expire=10): + with RedisLock.Lock(self._client, f"{self._bl}:settings_lock", expire=10): return self._get_settings() @settings.setter def settings(self, data: BeamlineSettingsModel): - with redis_lock.Lock(self._client, f"{self._bl}:settings_lock", expire=10): + with RedisLock.Lock(self._client, f"{self._bl}:settings_lock", expire=10): current = self._get_settings() updated_data = current.model_copy(update=data.model_dump(exclude_unset=True)) self._client.set(f"{self._bl}:settings", updated_data.model_dump_json()) @@ -1270,6 +1231,5 @@ if __name__ == "__main__": cfg = BeamlineConfig(bl=mx_beamline()) # cfg.allow_non_staff_request_from_staff = True - # cfg.state_busy = False # fg.abr_meas_pos = AerotechCoordinate(at_mm=Coordinate(x=0.0,y=0.0,z=0.0)) cfg.dtz_safe_position = 300.0 diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index a9adf94e..f7ff08e1 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -368,50 +368,26 @@ class AareDAQ: return self._devs.read_current_state_from_bec() def restart_bec_worker(self) -> dict[str, object]: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.restart_bec_worker(simulated=self._cfg.simulate_bec) - return {"ok": True, "device": "bec", "simulated": self._cfg.simulate_bec} - finally: - self._cfg.state_busy = False + self._devs.restart_bec_worker(simulated=self._cfg.simulate_bec) + return {"ok": True, "device": "bec", "simulated": self._cfg.simulate_bec} def restart_detector(self) -> dict[str, object]: - self._cfg.try_set_busy(timeout=360) - try: - beamline = MXBeamline.SIMULATED if self._cfg.simulated_detector else self._beamline - logger.info(f"Restarting JFJoch wrapper with simulated={self._cfg.simulated_detector}") - self._jfjoch = JFJochWrapper(beamline) - return { - "ok": True, - "device": "detector", - "simulated": bool(self._cfg.simulated_detector), - } - finally: - self._cfg.state_busy = False + beamline = MXBeamline.SIMULATED if self._cfg.simulated_detector else self._beamline + logger.info(f"Restarting JFJoch wrapper with simulated={self._cfg.simulated_detector}") + self._jfjoch = JFJochWrapper(beamline) + return {"ok": True, "device": "detector", "simulated": bool(self._cfg.simulated_detector)} def restart_tell(self) -> dict[str, object]: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.restart_tell(simulated=self._cfg.simulate_tell) - return {"ok": True, "device": "tell", "simulated": self._cfg.simulate_tell} - finally: - self._cfg.state_busy = False + self._devs.restart_tell(simulated=self._cfg.simulate_tell) + return {"ok": True, "device": "tell", "simulated": self._cfg.simulate_tell} def restart_aerotech(self) -> dict[str, object]: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.restart_aerotech(simulated=self._cfg.simulate_aerotech) - return {"ok": True, "device": "aerotech", "simulated": self._cfg.simulate_aerotech} - finally: - self._cfg.state_busy = False + self._devs.restart_aerotech(simulated=self._cfg.simulate_aerotech) + return {"ok": True, "device": "aerotech", "simulated": self._cfg.simulate_aerotech} def restart_smargon(self) -> dict[str, object]: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.restart_smargon(simulated=self._cfg.simulate_smargon) - return {"ok": True, "device": "smargon", "simulated": self._cfg.simulate_smargon} - finally: - self._cfg.state_busy = False + self._devs.restart_smargon(simulated=self._cfg.simulate_smargon) + return {"ok": True, "device": "smargon", "simulated": self._cfg.simulate_smargon} def set_runtime_simulation(self, device: str, enabled: bool) -> dict[str, object]: device = str(device).strip().lower() @@ -1561,10 +1537,6 @@ class AareDAQ: def state(self) -> BeamlineStateEnum: return self._cfg.state - @property - def busy(self) -> bool: - return self._cfg.state_busy - @state.setter def state(self, target: BeamlineStateEnum): if target == BeamlineStateEnum.Moving: @@ -1575,11 +1547,7 @@ class AareDAQ: raise RuntimeError("Cannot explicitly move to busy state") start = time.perf_counter() - self._cfg.try_set_busy(timeout=300) - try: - self._set_state(target) - finally: - self._cfg.state_busy = False + self._set_state(target) end = time.perf_counter() @@ -1675,22 +1643,17 @@ class AareDAQ: if -2000 < val < 2000: try: self._devs.aerotech_omega = val - self._cfg.state_busy = False except Exception: logger.exception("Omega move timed out") - self._cfg.state_busy = False else: - self._cfg.state_busy = False logger.error("Omega has to be between -2000 and 2000 degrees") raise ValueError("Omega has to be between -2000 and 2000 degrees (for now)") @omega.setter def omega(self, val: float): - self._cfg.set_busy(BeamlineStateEnum.SampleAlignment) self._omega(val) def omega_rel(self, val: float): - self._cfg.set_busy(BeamlineStateEnum.SampleAlignment) curr_omega = self._devs.aerotech_omega self._omega(curr_omega + val) @@ -1783,52 +1746,26 @@ class AareDAQ: self._devs.samcam_settings = s def tweak_abr_meas_pos(self, c: AerotechCoordinate): - self._cfg.set_busy(BeamlineStateEnum.SampleAlignment) - try: - new_meas_pos = AerotechCoordinate(at_mm=self._cfg.abr_meas_pos.at_mm + c.at_mm) - self._cfg.abr_meas_pos = new_meas_pos - self._devs.aerotech_pos = new_meas_pos - self._saved_box = None - self._cfg.state_busy = False - except Exception: - self._cfg.state_busy = False - raise + new_meas_pos = AerotechCoordinate(at_mm=self._cfg.abr_meas_pos.at_mm + c.at_mm) + self._cfg.abr_meas_pos = new_meas_pos + self._devs.aerotech_pos = new_meas_pos + self._saved_box = None def save_abr_meas_pos(self): - self._cfg.set_busy(BeamlineStateEnum.SampleAlignment) - try: - self._cfg.abr_meas_pos = AerotechCoordinate(at_mm=self._devs.aerotech_pos.at_mm) - self._devs.bec_worker.save_current_aerotech_position() - self._cfg.state_busy = False - except Exception: - self._cfg.state_busy = False - raise + self._cfg.abr_meas_pos = AerotechCoordinate(at_mm=self._devs.aerotech_pos.at_mm) + self._devs.bec_worker.save_current_aerotech_position() def goto_abr_meas_pos(self): - self._cfg.set_busy(BeamlineStateEnum.SampleAlignment) - try: - self._devs.aerotech_pos = self._cfg.abr_meas_pos - self._cfg.state_busy = False - except Exception: - self._cfg.state_busy = False - raise + self._devs.aerotech_pos = self._cfg.abr_meas_pos def create_sample(self, target: SampleShortInfo): + curr_sample = self._cfg.current_sample - self._cfg.try_set_busy(timeout=360) + if curr_sample is not None and curr_sample.location is not None: + raise RuntimeError("Sample from TELL is loaded") - try: - curr_sample = self._cfg.current_sample - - if curr_sample is not None and curr_sample.location is not None: - raise RuntimeError("Sample from TELL is loaded") - - self._aare.create_manual_sample(target) - self._cfg.current_sample = target - self._cfg.state_busy = False - except Exception: - self._cfg.state_busy = False - raise + self._aare.create_manual_sample(target) + self._cfg.current_sample = target def check_tell_mount_start_conditions(self) -> None: self._devs.tell.validate_mount_start_conditions() @@ -1838,12 +1775,9 @@ class AareDAQ: def park_and_dry(self, park=True): """External API for dry and park only""" - self._cfg.try_set_busy(timeout=360) try: self._execute_dry(park=park, unmount=False) - self._cfg.state_busy = False except Exception as e: - self._cfg.state_busy = False logger.error(f"Failed to park and dry: {e}") raise @@ -1854,27 +1788,20 @@ class AareDAQ: logger.exception("Failed to turn off blower") def initialise_smargon(self): - self._cfg.try_set_busy(timeout=360) try: self._devs.smargon_initialize() - self._cfg.state_busy = False except Exception as e: - self._cfg.state_busy = False logger.error(f"Failed to initialise Smargon: {e}") raise def initialise_detector(self): - self._cfg.try_set_busy(timeout=360) try: self._jfjoch.initialize() - self._cfg.state_busy = False except Exception as e: - self._cfg.state_busy = False logger.error(f"Failed to initialise detector: {e}") raise def recovery_unmount_sample(self) -> None: - self._cfg.try_set_busy(timeout=360) try: self._set_state(BeamlineStateEnum.RobotSampleExchange) self._devs.tell.check_enable_motion() @@ -1883,14 +1810,11 @@ class AareDAQ: self._devs.tell.unmount(wait=True, timeout=360.0) self._cfg.current_sample = None self._set_state(BeamlineStateEnum.SampleAlignment) - self._cfg.state_busy = False except Exception: - self._cfg.state_busy = False raise @sample.setter def sample(self, target: SampleShortInfo | None): - self._cfg.try_set_busy(timeout=360) try: logger.debug(f"Mount target {target}") @@ -1909,10 +1833,8 @@ class AareDAQ: raise MountingFailed(f"Failed to {operation_name.lower()} {sample_name}") logger.info(f"Sample operation completed: {target}") - self._cfg.state_busy = False except Exception as e: - self._cfg.state_busy = False logger.debug(f"Failed to change mounted sample: {e}") raise @@ -1946,11 +1868,9 @@ class AareDAQ: self._cfg.set_busy(BeamlineStateEnum.SampleAlignment) try: best_z = self._auto_focus(settings) - self._cfg.state_busy = False return best_z except Exception as e: logger.error(f"Autofocus failed: {e}") - self._cfg.state_busy = False raise def auto_exposure(self): @@ -2057,18 +1977,15 @@ class AareDAQ: sample_log_context(self.sample), raster_request_log_context(request) ), ) - self._cfg.try_set_busy(timeout=ceil(360)) try: result = self._execute_raster_sequence(request, auto_center=auto_center) if result is None: raise RasterScanException("Raster scan failed") self._set_state(BeamlineStateEnum.SampleAlignment) - self._cfg.state_busy = False return result except Exception: self._set_state(BeamlineStateEnum.SampleAlignment) - self._cfg.state_busy = False raise def _rotation(self, request: RotationScanRequest) -> CompletedRotationScan: @@ -2157,7 +2074,6 @@ class AareDAQ: rotation_request_log_context(request, total_time_s=total_time), ), ) - self._cfg.try_set_busy(timeout=ceil(total_time + 360)) try: result = self._execute_rotation_sequence(request) @@ -2166,16 +2082,13 @@ class AareDAQ: logger.error("Rotation scan failed, no result returned") raise DataCollectionException("Rotation scan failed, no result returned") self._set_state(BeamlineStateEnum.SampleAlignment) - self._cfg.state_busy = False return result finally: try: - if self._cfg.state_busy and self._cfg.state != BeamlineStateEnum.Maintenance: - self._set_state(BeamlineStateEnum.SampleAlignment) + self._set_state(BeamlineStateEnum.SampleAlignment) except Exception: logger.exception("Failed to restore SampleAlignment after rotation") - finally: - self._cfg.state_busy = False + raise @property def dtz(self) -> float: @@ -2187,7 +2100,6 @@ class AareDAQ: @dtz.setter def dtz(self, val: float): - self._cfg.try_set_busy(timeout=360) state = self._cfg.state dtz_low = self._cfg.cached_dtz_low @@ -2198,27 +2110,22 @@ class AareDAQ: try: self.refresh_detector_metadata_cache() except Exception as e: - self._cfg.state_busy = False raise RuntimeError(f"DTZ limits unavailable and refresh failed: {e}") from e dtz_low = self._cfg.cached_dtz_low dtz_high = self._cfg.cached_dtz_high if dtz_low is None or dtz_high is None: - self._cfg.state_busy = False raise RuntimeError("DTZ limits are unavailable") if val < dtz_low or val > dtz_high: - self._cfg.state_busy = False raise RuntimeError(f"dtz={val} outside limits {dtz_low} to {dtz_high}") if state == BeamlineStateEnum.DataCollection: - self._cfg.state_busy = False raise RuntimeError("Cannot set dtz during data collection") elif state == BeamlineStateEnum.SampleAlignment: self._devs.set_dtz(val, wait=False) self._cfg.dtz = val - self._cfg.state_busy = False @property def smargon(self) -> SmargonCoordinate: @@ -2226,29 +2133,15 @@ class AareDAQ: @smargon.setter def smargon(self, sc: SmargonCoordinate): - self._cfg.set_busy(BeamlineStateEnum.SampleAlignment) - try: - self._saved_box = None - self._devs.smargon_pos = sc - self._devs.smargon_wait() - self._cfg.state_busy = False - except Exception: - self._cfg.state_busy = False - raise + self._saved_box = None + self._devs.smargon_pos = sc + self._devs.smargon_wait() def mark_beam(self, x_pxl: float, y_pxl: float): - self._cfg.set_busy(BeamlineStateEnum.BeamLocation) - try: - self._cfg.mark_beam(x_pxl, y_pxl, self._devs.zoom) - self._cfg.state_busy = False - except Exception: - self._cfg.state_busy = False - raise + self._cfg.mark_beam(x_pxl, y_pxl, self._devs.zoom) def clear_mark_beam(self): - self._cfg.set_busy(BeamlineStateEnum.BeamLocation) self._cfg.clear_mark_beam() - self._cfg.state_busy = False @property def sample_geometry(self) -> SampleGeometryModel: @@ -2290,24 +2183,17 @@ class AareDAQ: Returns: RasterGridRequest object representing the found bounding box, or None if failed. """ - try: - self._cfg.try_set_busy(timeout=360) - r = get_ml_bounding_box( - mlbox=self._mlbox, - sample=self.sample, - sample_geometry=self.sample_geometry, - filename=filename, - upload_image=self._aare.upload_image, - logger=logger, - max_images=self.AUTO_RASTER_MAX_IMAGES, - min_cell_size_mm=self.AUTO_RASTER_MAX_IMAGES, - skip_if_exceed_max_image_threshold=self.AUTO_RASTER_SKIP_IF_EXCEED_MAX_IMAGE_THRESHOLD, - ) - self._cfg.state_busy = False - return r - except Exception: - self._cfg.state_busy = False - raise + return get_ml_bounding_box( + mlbox=self._mlbox, + sample=self.sample, + sample_geometry=self.sample_geometry, + filename=filename, + upload_image=self._aare.upload_image, + logger=logger, + max_images=self.AUTO_RASTER_MAX_IMAGES, + min_cell_size_mm=self.AUTO_RASTER_MAX_IMAGES, + skip_if_exceed_max_image_threshold=self.AUTO_RASTER_SKIP_IF_EXCEED_MAX_IMAGE_THRESHOLD, + ) def face_detection( self, steps: int = 14, step_size: int = 15, face_min_ratio: float = 0.3 @@ -2323,14 +2209,9 @@ class AareDAQ: Returns: Dictionary containing face detection results, including found samples and fits. """ - self._cfg.try_set_busy(timeout=360) - try: - result = self._execute_face_detection( - steps=steps, step_size=step_size, face_min_ratio=face_min_ratio, report_error=True - ) - return result.payload - finally: - self._cfg.state_busy = False + return self._execute_face_detection( + steps=steps, step_size=step_size, face_min_ratio=face_min_ratio, report_error=True + ).payload @log_timing(logger, "Auto loop center") def auto_loop_center(self, sample: SampleShortInfo | None = None) -> float: @@ -2349,7 +2230,6 @@ class AareDAQ: """ start = time.perf_counter() try: - self._cfg.try_set_busy(timeout=360) if sample is None: if self.sample is None: raise LoopCenteringFailed( @@ -2361,12 +2241,8 @@ class AareDAQ: if not self._execute_loop_centering(sample): raise LoopCenteringFailed - - self._cfg.state_busy = False - except Exception: self._cfg.zoom_mode = ZoomModeEnum.User - self._cfg.state_busy = False raise finally: @@ -2564,23 +2440,16 @@ class AareDAQ: msg += "with an error" logger.error(f"{msg}, time taken {time.perf_counter() - start} seconds.") try: - if self._cfg.state_busy: - self._set_state(BeamlineStateEnum.RobotSampleExchange) - else: - logger.warning( - "Skipping recovery transition to RobotSampleExchange: " - "beamline is no longer busy. The busy key may have expired" - "befor recovery could run" - ) + self._set_state(BeamlineStateEnum.RobotSampleExchange) except Exception: logger.exception( "Failed to transition to RobotSampleExchange during error recovery" ) + raise else: msg += " successfully" logger.info(f"{msg}, time taken {time.perf_counter() - start} seconds.") - self._cfg.state_busy = False end = time.perf_counter() return end - start @@ -2610,7 +2479,6 @@ class AareDAQ: try: logger.info("Automation-measure starting sequence") self._validate_automation_state(context="automation start") - self._cfg.try_set_busy(timeout=self.AUTOMATION_BUSY_TIMEOUT_S) self._validate_automation_state(context="after acquiring automation busy state") logger.info("Cancelling any pending jfjoch operations") @@ -2676,7 +2544,6 @@ class AareDAQ: logger.info( f"Automation-measure - loop Centering done at {time.perf_counter() - start}" ) - logger.debug(f"Automation-measure - current busy-state: {self._cfg.state_busy}") face_detection_result = self._execute_face_detection( steps=7, step_size=30, face_min_ratio=0.3, report_error=True @@ -2937,8 +2804,7 @@ class AareDAQ: extra={"from_state": curr_state, "to_state": target}, ) - if not self._cfg.state_busy: - raise RuntimeError("Beamline should be busy") + # TODO: CHECK BUSY LOCK if target == BeamlineStateEnum.Maintenance: self._cfg.state = BeamlineStateEnum.Maintenance @@ -3137,7 +3003,6 @@ class AareDAQ: extra={"from_state": curr_state, "to_state": target}, ) self._cfg.state = BeamlineStateEnum.Maintenance - self._cfg.state_busy = False raise StateTransitionFailed( f"State transition failed: {curr_state} -> {target}. " f"Beamline moved to Maintenance. Original error: {e}" @@ -3408,32 +3273,16 @@ class AareDAQ: self._jfjoch.cancel() def anneal(self, time_s: float): - self._cfg.set_busy(BeamlineStateEnum.SampleAlignment) - try: - self._devs.anneal(time_s) - finally: - self._cfg.state_busy = False + self._devs.anneal(time_s) def mono_pitch_scan(self, plot: bool = False) -> None: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.bec_worker.mono_pitch_scan_runner(plot=plot) - finally: - self._cfg.state_busy = False + self._devs.bec_worker.mono_pitch_scan_runner(plot=plot) def change_energy(self, value: float, plot: bool = False) -> None: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.bec_worker.change_energy(value=value, plot=plot) - finally: - self._cfg.state_busy = False + self._devs.bec_worker.change_energy(value=value, plot=plot) def bec_load_user_macros(self) -> None: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.bec_worker.load_user_macros() - finally: - self._cfg.state_busy = False + self._devs.bec_worker.load_user_macros() def bec_list_all_user_macros(self) -> list[str]: macros = self._devs.bec_worker.list_all_user_macros() @@ -3448,33 +3297,17 @@ class AareDAQ: return [str(item) for item in devices] def bec_reinitialise_planner_and_position_devices(self, method: str = "auto") -> list[str]: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.bec_worker.load_user_macros() - return self._devs.bec_worker.reinitialise_planner_and_position_devices(method=method) - finally: - self._cfg.state_busy = False + self._devs.bec_worker.load_user_macros() + return self._devs.bec_worker.reinitialise_planner_and_position_devices(method=method) def bec_save_current_bs_pos(self) -> None: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.bec_worker.save_current_bs_pos() - finally: - self._cfg.state_busy = False + self._devs.bec_worker.save_current_bs_pos() def bec_save_current_collimator_pos(self) -> None: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.bec_worker.save_current_collimator_pos() - finally: - self._cfg.state_busy = False + self._devs.bec_worker.save_current_collimator_pos() def bec_save_current_aerotech_position(self) -> None: - self._cfg.try_set_busy(timeout=360) - try: - self._devs.bec_worker.save_current_aerotech_position() - finally: - self._cfg.state_busy = False + self._devs.bec_worker.save_current_aerotech_position() def steer_beam_available(self) -> bool: return "beam_steering" in self._devs.bec_worker.dev @@ -3483,27 +3316,20 @@ class AareDAQ: """Run the routine to move the beam to the sample location. Update the location if provided.""" if "beam_steering" not in self._devs.bec_worker.dev: raise BECCommunicationError("Beam steering device does not exist in the BEC config.") - self._cfg.try_set_busy(timeout=360) - - try: - if not self.shutter: - raise AareException( - "Shutter not open! Please open the shutter before running beam steering." - ) - self._dispatch.bec_macros.auto_exposure() - if x is not None: - self._devs.bec_worker.dev.beam_steering.sample_loc_x_px.set(x).wait() - if y is not None: - self._devs.bec_worker.dev.beam_steering.sample_loc_y_px.set(y).wait() - self._devs.bec_worker.dev.beam_steering.trigger().wait() - finally: - self._cfg.state_busy = False + if not self.shutter: + raise AareException( + "Shutter not open! Please open the shutter before running beam steering." + ) + self._dispatch.bec_macros.auto_exposure() + if x is not None: + self._devs.bec_worker.dev.beam_steering.sample_loc_x_px.set(x).wait() + if y is not None: + self._devs.bec_worker.dev.beam_steering.sample_loc_y_px.set(y).wait() + self._devs.bec_worker.dev.beam_steering.trigger().wait() def fluorimeter_take_spectrum( self, fm: FluorescenceSpectrumParameterModel ) -> FluorescenceSpectrumOutputModel: - self._cfg.try_set_busy(timeout=360) - try: self._set_state(BeamlineStateEnum.XrayFluorescence) @@ -3513,11 +3339,9 @@ class AareDAQ: # TODO: Fill self._set_state(BeamlineStateEnum.SampleAlignment) - self._cfg.state_busy = False return None except Exception: self._set_state(BeamlineStateEnum.SampleAlignment) - self._cfg.state_busy = False raise def get_local_contact_config(self) -> LocalContactConfigModel: diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index ce4a1c6d..b1a8b19d 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -731,34 +731,25 @@ async def bec_save_current_aerotech_position(token: str = Depends(oauth2_scheme) def initialise_aerotech(self): - self._cfg.try_set_busy(timeout=360) try: self._devs.aerotech.home_aerotech() - self._cfg.state_busy = False except Exception as e: - self._cfg.state_busy = False logger.error(f"Failed to initialise Aerotech: {e}") raise def detector_take_pedestal(self): - self._cfg.try_set_busy(timeout=360) try: self._jfjoch.take_pedestal() - self._cfg.state_busy = False except Exception as e: - self._cfg.state_busy = False logger.error(f"Failed to take detector pedestal: {e}") raise def initialise_detector(self): - self._cfg.try_set_busy(timeout=360) try: self._jfjoch.initialize() - self._cfg.state_busy = False except Exception as e: - self._cfg.state_busy = False logger.error(f"Failed to initialise detector: {e}") raise @@ -1500,7 +1491,6 @@ async def force_clear_busy( data = auth.parse_token(token) auth.check_jwt_staff_only(data) _validate_recovery_code(payload.confirmation_code) - cfg.state_busy = False logger.warning( "Beamline busy flag cleared via protected endpoint.", extra={"session": getattr(data, "session", None)}, @@ -1528,10 +1518,8 @@ async def force_maintenance_state( sample_mounted = _sample_is_mounted() prev_state = cfg.state - prev_busy = cfg.state_busy auth.force_current_sesion(cfg, data) - cfg.state_busy = False cfg.state = BeamlineStateEnum.Maintenance logger.warning( @@ -1539,7 +1527,6 @@ async def force_maintenance_state( extra={ "session": getattr(data, "session", None), "previous_state": getattr(prev_state, "name", str(prev_state)), - "previous_busy": prev_busy, "sample_mounted": sample_mounted, }, ) @@ -1548,7 +1535,6 @@ async def force_maintenance_state( "ok": True, "sample_mounted": sample_mounted, "previous_state": getattr(prev_state, "name", str(prev_state)), - "previous_busy": prev_busy, "new_state": BeamlineStateEnum.Maintenance.name, } @@ -1573,12 +1559,6 @@ async def recovery_unmount_sample( auth.force_current_sesion(cfg, data) - if cfg.state_busy: - raise HTTPException( - status_code=api_status.HTTP_409_CONFLICT, - detail="Beamline is busy. Clear or recover the beamline before attempting recovery unmount.", - ) - status = daq.status if not getattr(status, "tell_connected", False): raise HTTPException( diff --git a/tests/unit/daq/test_face_detection.py b/tests/unit/daq/test_face_detection.py index 72a55736..8ccdc090 100644 --- a/tests/unit/daq/test_face_detection.py +++ b/tests/unit/daq/test_face_detection.py @@ -76,8 +76,6 @@ def test_public_face_detection_uses_execute_face_detection(monkeypatch): from aare.daq.daq import AareDAQ daq = object.__new__(AareDAQ) - cfg = types.SimpleNamespace(try_set_busy=lambda timeout=360: None, state_busy=False) - daq._cfg = cfg daq._execute_face_detection = lambda **kwargs: FaceDetectionResult( success=True, @@ -88,4 +86,3 @@ def test_public_face_detection_uses_execute_face_detection(monkeypatch): assert result["running"] is False assert result["samples"] == [{"angle": 45}] - assert cfg.state_busy is False -- 2.54.0 From a79207f2c21003b0d7e605450451730a276ae90e Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 8 Sep 2026 11:55:48 +0200 Subject: [PATCH 2/4] feat: add hardware lock to server --- pyproject.toml | 4 +- src/aare/daq/config.py | 260 +++++++++--------- src/aare/daq/daq.py | 30 +- src/aare/daq/server.py | 157 +++++++---- src/aare/devices/jfjoch.py | 5 - src/aare/gui/panels/local_contact_panel.py | 15 +- src/aare/gui/threads/daq_worker.py | 8 +- .../test_automation_progress_state_manager.py | 2 +- 8 files changed, 261 insertions(+), 220 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index dbf57401..dcc0fd5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,9 +6,6 @@ readme = "README.md" requires-python = ">=3.11" dependencies = [ "uv", - "gunicorn", - # >=0.7: DataCollectionParameters.transmission is a 0-1 fraction, which - # the scan panels rely on (older releases held an int percentage). "aarecommon>=0.7.3", "pydantic>=2.11", "numpy", @@ -112,6 +109,7 @@ ignore = [ "DTZ005", "DTZ006", ] +isort.split-on-trailing-comma=false [tool.ruff.format] skip-magic-trailing-comma = true diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index 6c0bb1ca..60902e67 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -86,8 +86,8 @@ class BeamlineConfig: operations are performed with atomic safety using Redis locking mechanisms. Attributes: - __bl (str): The beamline's unique identifier or name. - __client (redis.Redis): Redis client instance used for interacting with the datastore. + _bl (str): The beamline's unique identifier or name. + redis (redis.Redis): Redis client instance used for interacting with the datastore. """ GUI_SESSION_EXPIRE_SECONDS = 60 * 10 @@ -99,7 +99,7 @@ class BeamlineConfig: host = "localhost" else: host = cfg_get("daq.hardware.redis_url", f"{self._bl}-redis.psi.ch") - self._client = redis.Redis(host=host, port=6379, db=0, decode_responses=True) + self.redis = redis.Redis(host=host, port=6379, db=0, decode_responses=True) self.simulated_detector = bl is MXBeamline.SIMULATED self._initialize_optional_yaml_defaults() @@ -136,7 +136,7 @@ class BeamlineConfig: ) -> None: expiry = int(expiry_sec or self.GUI_SESSION_EXPIRE_SECONDS) - pipe = self._client.pipeline() + pipe = self.redis.pipeline() pipe.set(self._gui_session_key(payload.session), payload.model_dump_json()) pipe.expire(self._gui_session_key(payload.session), expiry) pipe.sadd(self._gui_sessions_index_key(), payload.session) @@ -144,7 +144,7 @@ class BeamlineConfig: def _current_gui_session_ttl(self, session: int) -> int | None: try: - ttl = int(self._client.ttl(self._gui_session_key(session))) + ttl = int(self.redis.ttl(self._gui_session_key(session))) except Exception: logger.debug("Could not read the GUI session TTL", exc_info=True) return None @@ -154,7 +154,7 @@ class BeamlineConfig: return None def _read_gui_session(self, session: int) -> OpenGuiSessionInfo | None: - raw = self._client.get(self._gui_session_key(session)) + raw = self.redis.get(self._gui_session_key(session)) if raw is None: return None try: @@ -192,7 +192,7 @@ class BeamlineConfig: ) -> OpenGuiSessionInfo | None: payload = self._read_gui_session(session) if payload is None: - self._client.srem(self._gui_sessions_index_key(), session) + self.redis.srem(self._gui_sessions_index_key(), session) return None payload.last_interaction_ts = last_interaction_ts @@ -205,7 +205,7 @@ class BeamlineConfig: ) -> OpenGuiSessionInfo | None: payload = self._read_gui_session(session) if payload is None: - self._client.srem(self._gui_sessions_index_key(), session) + self.redis.srem(self._gui_sessions_index_key(), session) return None payload.close_requested = True @@ -220,7 +220,7 @@ class BeamlineConfig: def clear_gui_close_request(self, session: int) -> None: payload = self._read_gui_session(session) if payload is None: - self._client.srem(self._gui_sessions_index_key(), session) + self.redis.srem(self._gui_sessions_index_key(), session) return payload.close_requested = False @@ -232,29 +232,29 @@ class BeamlineConfig: self._write_gui_session(payload, expiry_sec=ttl) def remove_gui_session(self, session: int) -> None: - pipe = self._client.pipeline() + pipe = self.redis.pipeline() pipe.delete(self._gui_session_key(session)) pipe.srem(self._gui_sessions_index_key(), session) pipe.execute() self.purge_expired_gui_sessions() def purge_expired_gui_sessions(self) -> None: - session_ids = self._client.smembers(self._gui_sessions_index_key()) + session_ids = self.redis.smembers(self._gui_sessions_index_key()) if not session_ids: return expired_ids: list[str] = [] for session_id in session_ids: - if not self._client.exists(self._gui_session_key(int(session_id))): + if not self.redis.exists(self._gui_session_key(int(session_id))): expired_ids.append(session_id) if expired_ids: - self._client.srem(self._gui_sessions_index_key(), *expired_ids) + self.redis.srem(self._gui_sessions_index_key(), *expired_ids) def get_open_gui_sessions(self) -> list[OpenGuiSessionInfo]: self.purge_expired_gui_sessions() - session_ids = self._client.smembers(self._gui_sessions_index_key()) + session_ids = self.redis.smembers(self._gui_sessions_index_key()) if not session_ids: return [] @@ -273,7 +273,7 @@ class BeamlineConfig: def get_gui_session(self, session: int) -> OpenGuiSessionInfo | None: payload = self._read_gui_session(session) if payload is None: - self._client.srem(self._gui_sessions_index_key(), session) + self.redis.srem(self._gui_sessions_index_key(), session) return None holder = self.baton_holder @@ -284,7 +284,7 @@ class BeamlineConfig: @property def allow_non_staff_request_from_staff(self) -> bool: - raw = self._client.get(f"{self._bl}:allow_non_staff_request_from_staff") + raw = self.redis.get(f"{self._bl}:allow_non_staff_request_from_staff") if raw is None: return False return str(raw).strip().lower() in {"1", "true", "yes", "on"} @@ -292,12 +292,12 @@ class BeamlineConfig: @allow_non_staff_request_from_staff.setter def allow_non_staff_request_from_staff(self, enabled: bool) -> None: if enabled: - self._client.set(f"{self._bl}:allow_non_staff_request_from_staff", "1") + self.redis.set(f"{self._bl}:allow_non_staff_request_from_staff", "1") else: - self._client.delete(f"{self._bl}:allow_non_staff_request_from_staff") + self.redis.delete(f"{self._bl}:allow_non_staff_request_from_staff") def generate_session(self) -> int: - return int(self._client.incr(f"{self._bl}:session")) + return int(self.redis.incr(f"{self._bl}:session")) @property def active_session(self) -> int | None: @@ -342,51 +342,51 @@ class BeamlineConfig: return SessionsStateEnum.OwnedByElse def try_set_active_session(self, session: int, expiry_sec: int) -> None: - with RedisLock(self.redis, f"{self._bl}:active_session_lock", expire=10): + with RedisLock(self.redis, f"{self._bl}:active_session_lock", timeout=10): active = self.active_session if active is None: - self._client.set(f"{self._bl}:active_session", session) + self.redis.set(f"{self._bl}:active_session", session) elif active != session: raise RuntimeError( "There is already active session with different id. Try again later." ) - self._client.expire(f"{self._bl}:active_session", expiry_sec) + self.redis.expire(f"{self._bl}:active_session", expiry_sec) # TODO finish setting this up! def try_extend_active_session(self, session: int, expiry_sec: int) -> None: - with RedisLock.Lock(self.redis, f"{self._bl}:active_session_lock", expire=10): + with RedisLock(self.redis, f"{self._bl}:active_session_lock", timeout=10): active = self.active_session if active is None: raise RuntimeError("There is no active session with given id. Try again later.") # if active == session: - # self._client.expire(f"{self._bl}:active_session", expiry_sec, gt=True) + # self.redis.expire(f"{self._bl}:active_session", expiry_sec, gt=True) if active == session: key = f"{self._bl}:active_session" - current_ttl = self._client.ttl(key) + current_ttl = self.redis.ttl(key) # Redis compatibility: # some Redis servers do not support EXPIRE with GT option. # Emulate "extend only if greater" manually while holding the lock. if current_ttl is None or current_ttl < 0 or current_ttl < expiry_sec: - self._client.expire(key, expiry_sec) + self.redis.expire(key, expiry_sec) else: raise RuntimeError( "There is already active session with different id. Try again later." ) def end_active_session(self, session: int) -> None: - with RedisLock.Lock(self.redis, f"{self._bl}:active_session_lock", expire=10): + with RedisLock(self.redis, f"{self._bl}:active_session_lock", timeout=10): active = self.active_session if active is None: return if active == session: - self._client.delete(f"{self._bl}:active_session") - self._client.delete(f"{self._bl}:baton_holder") + self.redis.delete(f"{self._bl}:active_session") + self.redis.delete(f"{self._bl}:baton_holder") def force_set_active_session(self, session: int, expiry_sec: int) -> None: # Ensure that there is no active try-set for active session - with RedisLock.Lock(self.redis, f"{self._bl}:active_session_lock", expire=10): + with RedisLock(self.redis, f"{self._bl}:active_session_lock", timeout=10): self.redis.set(f"{self._bl}:active_session", session) self.redis.expire(f"{self._bl}:active_session", expiry_sec) @@ -395,7 +395,7 @@ class BeamlineConfig: @property def baton_holder(self) -> BatonHolderInfo | None: """Get information about the current baton holder.""" - tmp = self._client.get(f"{self._bl}:baton_holder") + tmp = self.redis.get(f"{self._bl}:baton_holder") if tmp is None: return None try: @@ -407,14 +407,14 @@ class BeamlineConfig: @baton_holder.setter def baton_holder(self, info: BatonHolderInfo | None) -> None: if info is None: - self._client.delete(f"{self._bl}:baton_holder") + self.redis.delete(f"{self._bl}:baton_holder") else: - self._client.set(f"{self._bl}:baton_holder", info.model_dump_json()) + self.redis.set(f"{self._bl}:baton_holder", info.model_dump_json()) @property def pending_baton_request(self) -> BatonRequest | None: """Get the current pending baton request, if any.""" - tmp = self._client.get(f"{self._bl}:baton_request") + tmp = self.redis.get(f"{self._bl}:baton_request") if tmp is None: return None try: @@ -428,19 +428,19 @@ class BeamlineConfig: ) -> None: """Set a pending baton request with auto-expiry for timeout.""" if request is None: - self._client.delete(f"{self._bl}:baton_request") + self.redis.delete(f"{self._bl}:baton_request") else: - self._client.set(f"{self._bl}:baton_request", request.model_dump_json()) + self.redis.set(f"{self._bl}:baton_request", request.model_dump_json()) # Add a few seconds buffer so we can detect timeout vs expiry - self._client.expire(f"{self._bl}:baton_request", timeout_sec + 5) + self.redis.expire(f"{self._bl}:baton_request", timeout_sec + 5) def clear_pending_baton_request(self) -> None: - self._client.delete(f"{self._bl}:baton_request") + self.redis.delete(f"{self._bl}:baton_request") @property def queued_baton_transfer(self) -> BatonTransferQueue | None: """Get queued transfer waiting for beamline to be available.""" - tmp = self._client.get(f"{self._bl}:baton_transfer_queue") + tmp = self.redis.get(f"{self._bl}:baton_transfer_queue") if tmp is None: return None try: @@ -452,9 +452,9 @@ class BeamlineConfig: @queued_baton_transfer.setter def queued_baton_transfer(self, transfer: BatonTransferQueue | None) -> None: if transfer is None: - self._client.delete(f"{self._bl}:baton_transfer_queue") + self.redis.delete(f"{self._bl}:baton_transfer_queue") else: - self._client.set(f"{self._bl}:baton_transfer_queue", transfer.model_dump_json()) + self.redis.set(f"{self._bl}:baton_transfer_queue", transfer.model_dump_json()) def can_transfer_baton_now(self) -> bool: """Check if baton can be transferred (beamline not mid-operation).""" @@ -475,9 +475,9 @@ class BeamlineConfig: Atomically transfer the baton to a new holder. Use existing active_session_lock for consistency. """ - with RedisLock(self._client, f"{self._bl}:active_session_lock", expire=10): - self._client.set(f"{self._bl}:active_session", to_session) - self._client.expire(f"{self._bl}:active_session", expiry_sec) + with RedisLock(self.redis, f"{self._bl}:active_session_lock", timeout=10): + self.redis.set(f"{self._bl}:active_session", to_session) + self.redis.expire(f"{self._bl}:active_session", expiry_sec) self.baton_holder = BatonHolderInfo( username=to_username, session=to_session, is_staff=to_is_staff, pgroup=to_pgroup ) @@ -508,27 +508,27 @@ class BeamlineConfig: @property def pgroup(self) -> str | None: - tmp = self._client.get(f"{self._bl}:pgroup") + tmp = self.redis.get(f"{self._bl}:pgroup") return tmp @pgroup.setter def pgroup(self, pgroup: str | None) -> None: if pgroup is None: - self._client.delete(f"{self._bl}:pgroup") + self.redis.delete(f"{self._bl}:pgroup") else: - self._client.set(f"{self._bl}:pgroup", pgroup) + self.redis.set(f"{self._bl}:pgroup", pgroup) @property def commissioning_mode(self) -> bool: - tmp = self._client.get(f"{self._bl}:commissioning_mode") + tmp = self.redis.get(f"{self._bl}:commissioning_mode") return tmp is not None @commissioning_mode.setter def commissioning_mode(self, commisioning_mode: bool) -> None: if commisioning_mode: - self._client.set(f"{self._bl}:commissioning_mode", "1") + self.redis.set(f"{self._bl}:commissioning_mode", "1") else: - self._client.delete(f"{self._bl}:commissioning_mode") + self.redis.delete(f"{self._bl}:commissioning_mode") def start_moving( self, target: BeamlineStateEnum, timeout: int | None = None @@ -540,7 +540,7 @@ class BeamlineConfig: @property def state(self) -> BeamlineStateEnum: - raw_value = self._client.get(f"{self._bl}:state") + raw_value = self.redis.get(f"{self._bl}:state") if raw_value is None: return BeamlineStateEnum.Maintenance @@ -554,11 +554,11 @@ class BeamlineConfig: @state.setter def state(self, state: BeamlineStateEnum): - self._client.set(f"{self._bl}:state", state.value) + self.redis.set(f"{self._bl}:state", state.value) @property def tell_mount_count(self) -> int: - return int(self._client.incr(f"{self._bl}:tell_mount_count")) + return int(self.redis.incr(f"{self._bl}:tell_mount_count")) def pixel_to_mm(self, zoom: float) -> float: cfg = self.settings @@ -600,8 +600,8 @@ class BeamlineConfig: @property def beam_size_mm(self) -> Coordinate: - tmp_x = self._client.get(f"{self._bl}:beam_size_x") - tmp_y = self._client.get(f"{self._bl}:beam_size_y") + tmp_x = self.redis.get(f"{self._bl}:beam_size_x") + tmp_y = self.redis.get(f"{self._bl}:beam_size_y") if tmp_x: val_x = float(tmp_x) else: @@ -614,11 +614,11 @@ class BeamlineConfig: @beam_size_mm.setter def beam_size_mm(self, data: Coordinate): - self._client.set(f"{self._bl}:beam_size_x", data.x) - self._client.set(f"{self._bl}:beam_size_y", data.y) + self.redis.set(f"{self._bl}:beam_size_x", data.x) + self.redis.set(f"{self._bl}:beam_size_y", data.y) def _get_settings(self) -> BeamlineSettingsModel: - tmp = self._client.get(f"{self._bl}:settings") + tmp = self.redis.get(f"{self._bl}:settings") if tmp is None: return BeamlineSettingsModel() @@ -628,19 +628,19 @@ class BeamlineConfig: @property def settings(self) -> BeamlineSettingsModel: - with RedisLock.Lock(self._client, f"{self._bl}:settings_lock", expire=10): + with RedisLock(self.redis, f"{self._bl}:settings_lock", timeout=10): return self._get_settings() @settings.setter def settings(self, data: BeamlineSettingsModel): - with RedisLock.Lock(self._client, f"{self._bl}:settings_lock", expire=10): + with RedisLock(self.redis, f"{self._bl}:settings_lock", timeout=10): current = self._get_settings() updated_data = current.model_copy(update=data.model_dump(exclude_unset=True)) - self._client.set(f"{self._bl}:settings", updated_data.model_dump_json()) + self.redis.set(f"{self._bl}:settings", updated_data.model_dump_json()) @property def cryojet_settings(self) -> CryojetSettingsModel: - tmp = self._client.get(f"{self._bl}:cryojet_settings") + tmp = self.redis.get(f"{self._bl}:cryojet_settings") if tmp is None: return CryojetSettingsModel() @@ -649,17 +649,17 @@ class BeamlineConfig: @cryojet_settings.setter def cryojet_settings(self, data: CryojetSettingsModel): - self._client.set(f"{self._bl}:cryojet_settings", data.model_dump_json()) + self.redis.set(f"{self._bl}:cryojet_settings", data.model_dump_json()) def get_alc_bkg(self, zoom: float, exp: float, gain: float) -> np.ndarray | None: - return base64_to_numpy(self._client.get(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}")) + return base64_to_numpy(self.redis.get(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}")) def put_alc_bkg(self, zoom: float, exp: float, gain: float, data: np.ndarray): - self._client.set(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}", numpy_to_base64(data)) + self.redis.set(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}", numpy_to_base64(data)) @property def spreadsheet(self) -> SampleShortInfoList: - tmp = self._client.get(f"{self._bl}:sample_spreadsheet") + tmp = self.redis.get(f"{self._bl}:sample_spreadsheet") if tmp is None: return SampleShortInfoList(s=[]) @@ -673,17 +673,17 @@ class BeamlineConfig: @spreadsheet.setter def spreadsheet(self, data: SampleShortInfoList): - self._client.set(f"{self._bl}:sample_spreadsheet", data.model_dump_json()) + self.redis.set(f"{self._bl}:sample_spreadsheet", data.model_dump_json()) def listen_changes_spreadsheet(self) -> redis.client.PubSub: - self._client.config_set("notify-keyspace-events", "KEA") - pubsub = self._client.pubsub() + self.redis.config_set("notify-keyspace-events", "KEA") + pubsub = self.redis.pubsub() pubsub.psubscribe(f"__keyspace@0__:{self._bl}:sample_spreadsheet") return pubsub @property def reference_tools(self) -> SampleShortInfoList: - tmp = self._client.get(f"{self._bl}:reference-tools") + tmp = self.redis.get(f"{self._bl}:reference-tools") if tmp is None: return SampleShortInfoList(s=[]) @@ -692,17 +692,17 @@ class BeamlineConfig: @reference_tools.setter def reference_tools(self, data: SampleShortInfoList): - self._client.set(f"{self._bl}:reference-tools", data.model_dump_json()) + self.redis.set(f"{self._bl}:reference-tools", data.model_dump_json()) def listen_changes_reference_tools(self) -> redis.client.PubSub: - self._client.config_set("notify-keyspace-events", "KEA") - pubsub = self._client.pubsub() + self.redis.config_set("notify-keyspace-events", "KEA") + pubsub = self.redis.pubsub() pubsub.psubscribe(f"__keyspace@0__:{self._bl}:reference-tools") return pubsub @property def current_sample(self) -> SampleShortInfo | None: - tmp = self._client.get(f"{self._bl}:current_sample") + tmp = self.redis.get(f"{self._bl}:current_sample") if tmp is None: return None @@ -712,13 +712,13 @@ class BeamlineConfig: @current_sample.setter def current_sample(self, sample: SampleShortInfo | None): if sample is None: - self._client.delete(f"{self._bl}:current_sample") + self.redis.delete(f"{self._bl}:current_sample") else: - self._client.set(f"{self._bl}:current_sample", sample.model_dump_json()) + self.redis.set(f"{self._bl}:current_sample", sample.model_dump_json()) @property def beam_mark_coeff(self) -> BeamMarkCoeffModel: - tmp = self._client.get(f"{self._bl}:beam_center_camera") + tmp = self.redis.get(f"{self._bl}:beam_center_camera") if tmp is None: return BeamMarkCoeffModel() @@ -727,12 +727,12 @@ class BeamlineConfig: @beam_mark_coeff.setter def beam_mark_coeff(self, data: BeamMarkCoeffModel): - self._client.set(f"{self._bl}:beam_center_camera", data.model_dump_json()) + self.redis.set(f"{self._bl}:beam_center_camera", data.model_dump_json()) # TODO tidy up zoom functions @property def zoom_mode(self) -> ZoomModeEnum: - raw_value = self._client.get(f"{self._bl}:zoom_mode") + raw_value = self.redis.get(f"{self._bl}:zoom_mode") if raw_value is None: print("no zoom mode given, defaulting to user mode") return ZoomModeEnum.User @@ -746,7 +746,7 @@ class BeamlineConfig: @zoom_mode.setter def zoom_mode(self, mode: ZoomModeEnum): - self._client.set(f"{self._bl}:zoom_mode", mode.value) + self.redis.set(f"{self._bl}:zoom_mode", mode.value) @staticmethod def zoom_setting_string(mode: ZoomModeEnum = ZoomModeEnum.User) -> str: @@ -766,7 +766,7 @@ class BeamlineConfig: mode = self.zoom_mode if not mode or not isinstance(mode, ZoomModeEnum): raise ValueError("incorrect zoom settings mode used") - tmp = self._client.get(f"{self._bl}:{self.zoom_setting_string(mode)}") + tmp = self.redis.get(f"{self._bl}:{self.zoom_setting_string(mode)}") if tmp is None: return zoom_manager(mode, self._mxb) data_dict = json.loads(tmp) @@ -777,7 +777,7 @@ class BeamlineConfig: mode = self.zoom_mode if not mode or not isinstance(mode, ZoomModeEnum): raise ValueError("incorrect zoom settings mode used") - self._client.set(f"{self._bl}:{self.zoom_setting_string(mode)}", data.model_dump_json()) + self.redis.set(f"{self._bl}:{self.zoom_setting_string(mode)}", data.model_dump_json()) def save_zoom_camera_setting( self, zoom_value: float, settings: SampleCameraSettings, mode: ZoomModeEnum | None = None @@ -787,14 +787,14 @@ class BeamlineConfig: active zoom mode; pass ``mode`` to target a specific one explicitly.""" mode = mode or self.zoom_mode key = f"{self._bl}:{self.zoom_setting_string(mode)}" - tmp = self._client.get(key) + tmp = self.redis.get(key) model = ZoomModel(**json.loads(tmp)) if tmp is not None else zoom_manager(mode, self._mxb) model.z[zoom_value] = settings - self._client.set(key, model.model_dump_json()) + self.redis.set(key, model.model_dump_json()) @property def abr_meas_pos(self) -> AerotechCoordinate: - tmp = self._client.get(f"{self._bl}:abr_meas_pos") + tmp = self.redis.get(f"{self._bl}:abr_meas_pos") if tmp is None: return ABR_POS_MOUNT @@ -803,33 +803,33 @@ class BeamlineConfig: @abr_meas_pos.setter def abr_meas_pos(self, data: AerotechCoordinate): - self._client.set(f"{self._bl}:abr_meas_pos", data.model_dump_json()) + self.redis.set(f"{self._bl}:abr_meas_pos", data.model_dump_json()) @property def dtz(self) -> float | None: - tmp = self._client.get(f"{self._bl}:dtz") + tmp = self.redis.get(f"{self._bl}:dtz") if tmp is None: return None return float(tmp) @dtz.setter def dtz(self, dtz: float): - self._client.set(f"{self._bl}:dtz", dtz) + self.redis.set(f"{self._bl}:dtz", dtz) @property def dtz_safe_position(self) -> float | None: - tmp = self._client.get(f"{self._bl}:dtz_safe_position") + tmp = self.redis.get(f"{self._bl}:dtz_safe_position") if tmp is None: return None return float(tmp) @dtz_safe_position.setter def dtz_safe_position(self, dtz: float): - self._client.set(f"{self._bl}:dtz_safe_position", dtz) + self.redis.set(f"{self._bl}:dtz_safe_position", dtz) @property def xrf(self) -> FluorescenceSpectrumOutputModel | None: - tmp = self._client.get(f"{self._bl}:xrf") + tmp = self.redis.get(f"{self._bl}:xrf") if tmp is None: return None data_dict = json.loads(tmp) @@ -838,18 +838,18 @@ class BeamlineConfig: @xrf.setter def xrf(self, data: FluorescenceSpectrumOutputModel | None): if data is None: - self._client.delete(f"{self._bl}:xrf") + self.redis.delete(f"{self._bl}:xrf") else: - self._client.set(f"{self._bl}:xrf", data.model_dump_json()) + self.redis.set(f"{self._bl}:xrf", data.model_dump_json()) def clear_mark_beam(self): - self._client.delete(f"{self._bl}:beam_mark") + self.redis.delete(f"{self._bl}:beam_mark") def mark_beam(self, x_pxl: float, y_pxl: float, zoom: float): - self._client.hset( + self.redis.hset( f"{self._bl}:beam_mark", mapping={f"{zoom}": json.dumps({"x": x_pxl, "y": y_pxl})} ) - vals = self._client.hgetall(f"{self._bl}:beam_mark") + vals = self.redis.hgetall(f"{self._bl}:beam_mark") if len(vals) >= 3: zooms = [] @@ -874,7 +874,7 @@ class BeamlineConfig: @property def crystal_size(self) -> CrystalSize: - tmp = self._client.get(f"{self._bl}:crystal_size") + tmp = self.redis.get(f"{self._bl}:crystal_size") if tmp is None: return CrystalSize(x=0, y=0, z=0) data_dict = json.loads(tmp) @@ -882,11 +882,11 @@ class BeamlineConfig: @crystal_size.setter def crystal_size(self, xtal_size: CrystalSize): - self._client.set(f"{self._bl}:crystal_size", xtal_size.model_dump_json()) + self.redis.set(f"{self._bl}:crystal_size", xtal_size.model_dump_json()) @property def last_best_res(self) -> float | None: - tmp = self._client.get(f"{self._bl}:last_best_res") + tmp = self.redis.get(f"{self._bl}:last_best_res") if tmp is None: return None return float(tmp) @@ -894,13 +894,13 @@ class BeamlineConfig: @last_best_res.setter def last_best_res(self, best_res: float | None): if best_res is None: - self._client.delete(f"{self._bl}:last_best_res") + self.redis.delete(f"{self._bl}:last_best_res") else: - self._client.set(f"{self._bl}:last_best_res", best_res) + self.redis.set(f"{self._bl}:last_best_res", best_res) @property def last_best_b_factor(self) -> float | None: - tmp = self._client.get(f"{self._bl}:last_best_b_factor") + tmp = self.redis.get(f"{self._bl}:last_best_b_factor") if tmp is None: return None return float(tmp) @@ -908,22 +908,22 @@ class BeamlineConfig: @last_best_b_factor.setter def last_best_b_factor(self, last_best_b_factor: float | None): if last_best_b_factor is None: - self._client.delete(f"{self._bl}:last_best_b_factor") + self.redis.delete(f"{self._bl}:last_best_b_factor") else: - self._client.set(f"{self._bl}:last_best_b_factor", last_best_b_factor) + self.redis.set(f"{self._bl}:last_best_b_factor", last_best_b_factor) def _mount_failure_streak_key(self) -> str: return f"{self._bl}:mount_fail_count" def get_mount_failure_streak(self) -> int: - value = self._client.get(self._mount_failure_streak_key()) + value = self.redis.get(self._mount_failure_streak_key()) return int(value) if value else 0 def increment_mount_failure_streak(self) -> int: - return int(self._client.incr(self._mount_failure_streak_key())) + return int(self.redis.incr(self._mount_failure_streak_key())) def reset_mount_failure_streak(self) -> None: - self._client.delete(self._mount_failure_streak_key()) + self.redis.delete(self._mount_failure_streak_key()) def get_mount_fail_count(self) -> int: return self.get_mount_failure_streak() @@ -936,7 +936,7 @@ class BeamlineConfig: @property def simple_input_parameters(self) -> SimpleStrategyInputModel | None: - tmp = self._client.get(f"{self._bl}:simple_input_params") + tmp = self.redis.get(f"{self._bl}:simple_input_params") if tmp is None: return None data_dict = json.loads(tmp) @@ -945,13 +945,13 @@ class BeamlineConfig: @simple_input_parameters.setter def simple_input_parameters(self, input_params: SimpleStrategyInputModel | None): if input_params is None: - self._client.delete(f"{self._bl}:simple_input_params") + self.redis.delete(f"{self._bl}:simple_input_params") else: - self._client.set(f"{self._bl}:simple_input_params", input_params.model_dump_json()) + self.redis.set(f"{self._bl}:simple_input_params", input_params.model_dump_json()) @property def auto_params(self) -> SimpleScanParameters | None: - tmp = self._client.get(f"{self._bl}:auto_params") + tmp = self.redis.get(f"{self._bl}:auto_params") if tmp is None: logger.debug(f"auto_params missing in redis key {self._bl}:auto_params") return None @@ -965,9 +965,9 @@ class BeamlineConfig: @auto_params.setter def auto_params(self, params: SimpleScanParameters | None): if params is None: - self._client.delete(f"{self._bl}:auto_params") + self.redis.delete(f"{self._bl}:auto_params") else: - self._client.set(f"{self._bl}:auto_params", params.model_dump_json()) + self.redis.set(f"{self._bl}:auto_params", params.model_dump_json()) def _automation_progress_key(self) -> str: return f"{self._bl}:automation_progress" @@ -976,12 +976,12 @@ class BeamlineConfig: return f"{self._bl}:automation_progress_seq" def reset_automation_progress(self) -> None: - self._client.set(self._automation_progress_seq_key(), 0) - self._client.delete(self._automation_progress_key()) + self.redis.set(self._automation_progress_seq_key(), 0) + self.redis.delete(self._automation_progress_key()) def get_automation_progress_state(self) -> dict: - seq_raw = self._client.get(self._automation_progress_seq_key()) - payload_raw = self._client.get(self._automation_progress_key()) + seq_raw = self.redis.get(self._automation_progress_seq_key()) + payload_raw = self.redis.get(self._automation_progress_key()) seq = int(seq_raw) if seq_raw is not None else 0 progress = json.loads(payload_raw) if payload_raw else None @@ -1003,8 +1003,8 @@ class BeamlineConfig: else: raise TypeError(f"Unsupported automation progress type: {type(progress).__name__}") - next_seq = int(self._client.incr(self._automation_progress_seq_key())) - self._client.set( + next_seq = int(self.redis.incr(self._automation_progress_seq_key())) + self.redis.set( self._automation_progress_key(), json.dumps(payload, separators=(",", ":"), default=_json_default), ) @@ -1013,7 +1013,7 @@ class BeamlineConfig: @property def failed_mount_count(self) -> int: - tmp = self._client.get(f"{self._bl}:failed_mount_count") + tmp = self.redis.get(f"{self._bl}:failed_mount_count") if tmp is None: return 0 try: @@ -1028,18 +1028,18 @@ class BeamlineConfig: @failed_mount_count.setter def failed_mount_count(self, count: int): if count == 0: - self._client.delete(f"{self._bl}:failed_mount_count") + self.redis.delete(f"{self._bl}:failed_mount_count") else: - self._client.set(f"{self._bl}:failed_mount_count", count) + self.redis.set(f"{self._bl}:failed_mount_count", count) def increment_failed_mount_count(self) -> int: - return int(self._client.incr(f"{self._bl}:failed_mount_count")) + return int(self.redis.incr(f"{self._bl}:failed_mount_count")) def _runtime_sim_key(self, name: str) -> str: return f"{self._bl}:runtime:simulate:{name}" def get_runtime_simulated(self, name: str, default: bool = False) -> bool: - raw = self._client.get(self._runtime_sim_key(name)) + raw = self.redis.get(self._runtime_sim_key(name)) if raw is None: return default return str(raw).strip().lower() in {"1", "true", "yes", "on"} @@ -1047,9 +1047,9 @@ class BeamlineConfig: def set_runtime_simulated(self, name: str, enabled: bool) -> None: key = self._runtime_sim_key(name) if enabled: - self._client.set(key, "1") + self.redis.set(key, "1") else: - self._client.delete(key) + self.redis.delete(key) @property def simulate_bec(self) -> bool: @@ -1119,7 +1119,7 @@ class BeamlineConfig: def get_detector_metadata(self) -> dict: try: - raw = self._client.get(self._detector_metadata_key()) + raw = self.redis.get(self._detector_metadata_key()) if raw in (None, "", b""): return {} @@ -1140,7 +1140,7 @@ class BeamlineConfig: safe_payload.get("pixel_size_mm") ) safe_payload["updated_at"] = datetime.now().isoformat(timespec="seconds") - self._client.set(self._detector_metadata_key(), json.dumps(safe_payload)) + self.redis.set(self._detector_metadata_key(), json.dumps(safe_payload)) return safe_payload @property @@ -1190,7 +1190,7 @@ class BeamlineConfig: try: redis_key = f"{self._bl}:local_contact_config" - raw_value = self._client.get(redis_key) + raw_value = self.redis.get(redis_key) if raw_value in (None, "", b""): return default @@ -1209,7 +1209,7 @@ class BeamlineConfig: validated = LocalContactConfigModel.model_validate(config) try: redis_key = f"{self._bl}:local_contact_config" - self._client.set(redis_key, validated.model_dump_json()) + self.redis.set(redis_key, validated.model_dump_json()) logger.info(f"Saved Local Contact config to Redis: {redis_key}") except Exception as e: logger.error(f"Failed to write Local Contact config to Redis: {e}") diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index f7ff08e1..47570a32 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -4,7 +4,6 @@ import secrets import time from collections.abc import Callable from datetime import UTC, datetime -from math import ceil from pathlib import Path from typing import Any @@ -68,6 +67,7 @@ from aarecommon.models.raster_grid import CompletedRasterGrid, RasterGridRequest from aarecommon.models.rotation_scan import CompletedRotationScan, RotationScanRequest from aarecommon.models.tell import TellPhaseEnum, TellStateModel from aareDB import SampleEventType +from redis.lock import Lock as RedisLock from aare.beamline_dispatch.protocols import BeamlineDispatch from aare.daq import workflows @@ -256,8 +256,6 @@ class _FaceDetectionProgressReporter: self._daq._emit_face_detection_progress(payload) -# TODO tidy up DAQ - migrate functions into different scripts, to reduce size? -# TODO investigate using a state machine within each operation to reduce callbacks? class AareDAQ: """ Main Data Acquisition class for the Aare system. @@ -273,7 +271,10 @@ class AareDAQ: AUTO_RASTER_MIN_CELL_SIZE_MM = 0.005 AUTO_RASTER_SKIP_IF_EXCEED_MAX_IMAGE_THRESHOLD = True - def __init__(self, cfg: BeamlineConfig, bl: MXBeamline, dispatch: BeamlineDispatch): + def __init__( + self, cfg: BeamlineConfig, bl: MXBeamline, dispatch: BeamlineDispatch, hw_lock: RedisLock + ): + self._hw_lock = hw_lock self.last_time = 0.0 self._dispatch = dispatch self._cfg = cfg @@ -304,6 +305,10 @@ class AareDAQ: pgroup_provider=_DAQPGroupProvider(self), ) + @property + def busy(self) -> bool: + return self._hw_lock.locked() + def shutdown(self): self._devs.bec_worker.shutdown() @@ -1802,16 +1807,13 @@ class AareDAQ: raise def recovery_unmount_sample(self) -> None: - try: - self._set_state(BeamlineStateEnum.RobotSampleExchange) - self._devs.tell.check_enable_motion() - self._devs.tell.wait_not_busy() - self._devs.tell.set_in_mount_position(True) - self._devs.tell.unmount(wait=True, timeout=360.0) - self._cfg.current_sample = None - self._set_state(BeamlineStateEnum.SampleAlignment) - except Exception: - raise + self._set_state(BeamlineStateEnum.RobotSampleExchange) + self._devs.tell.check_enable_motion() + self._devs.tell.wait_not_busy() + self._devs.tell.set_in_mount_position(True) + self._devs.tell.unmount(wait=True, timeout=360.0) + self._cfg.current_sample = None + self._set_state(BeamlineStateEnum.SampleAlignment) @sample.setter def sample(self, target: SampleShortInfo | None): diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index b1a8b19d..7f6c192f 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -1,4 +1,5 @@ import asyncio +import functools import hmac import importlib import json @@ -6,18 +7,18 @@ import os import time from collections.abc import AsyncGenerator from contextlib import asynccontextmanager -from typing import Any, ClassVar import uvicorn from aarecommon.config.beamline import mx_beamline from aarecommon.config.logger import get_uvicorn_logging_config, setup_logger from aarecommon.errors.codes import AareErrorCode, export_error_codes_grouped from aarecommon.errors.exception_handler import ( + BeamlineBusyException, MaintenanceStateException, SampleException, UserRightsException, ) -from aarecommon.math.coordinate import AerotechCoordinate, Coordinate, SmargonCoordinate +from aarecommon.math.coordinate import AerotechCoordinate, SmargonCoordinate from aarecommon.math.sample_geometry import SampleGeometryModel from aarecommon.models.auth import BatonRequestStatus, BatonStatus from aarecommon.models.automation import AutomationProgress @@ -45,8 +46,8 @@ from fastapi import Depends, FastAPI, HTTPException, Request from fastapi import status as api_status from fastapi.concurrency import run_in_threadpool from fastapi.security import OAuth2PasswordBearer +from redis.lock import Lock as RedisLock from starlette.responses import StreamingResponse -from uvicorn.workers import UvicornWorker # deprecated shim, present in pinned 0.34.2 from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch from aare.beamline_dispatch.protocols import BeamlineDispatch @@ -65,6 +66,7 @@ oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") cfg: BeamlineConfig daq: AareDAQ bl_dispatch: BeamlineDispatch +hardware_busy_lock: RedisLock _all_pgroups_cache: dict[str, tuple[list[str], float]] = {} _ALL_PGROUPS_TTL_S = 60.0 # adjust TTL as needed @@ -82,13 +84,26 @@ _automation_progress_state: dict = {"seq": 0, "progress": None} _automation_progress_state_lock = asyncio.Lock() -class AareUvicornWorker(UvicornWorker): - # CONFIG_KWARGS merged last into uvicorn Config (uvicorn/workers.py:69) → - # keeps our access-log filter + proxy_headers under gunicorn. - CONFIG_KWARGS: ClassVar[dict[str, Any]] = { - "log_config": get_uvicorn_logging_config(), - "proxy_headers": False, - } +@asynccontextmanager +async def _lock_hw(): + if hardware_busy_lock.owned(): + yield + else: + if not hardware_busy_lock.acquire(blocking=False): + raise BeamlineBusyException("Beamline hardware lock is held by another worker") + logger.debug("Hardware lock acquired by process") + yield + hardware_busy_lock.release() + logger.debug("Hardware lock released by process") + + +def needs_hw_lock(func): + @functools.wraps(func) + async def wrapper(*args, **kwargs): + async with _lock_hw(): + return await func(*args, **kwargs) + + return wrapper @asynccontextmanager @@ -98,7 +113,7 @@ async def lifespan(application: FastAPI): All stateful / connection-opening initialisation belongs here so that each worker gets its own fresh Redis, BEC, EPICS, and TELL connections. """ - global cfg, daq, bl_dispatch + global cfg, daq, bl_dispatch, hardware_busy_lock logger.info(f"Worker {os.getpid()} setting up JWT authentication...") bl_dispatch = get_beamline_dispatch() @@ -108,7 +123,9 @@ async def lifespan(application: FastAPI): # ── Core objects (Redis, EPICS PVs, BEC, TELL, JFJoch, etc.) ── bl = mx_beamline() cfg = BeamlineConfig(bl) - daq = AareDAQ(cfg, bl, bl_dispatch) + hardware_busy_lock = RedisLock(cfg.redis, name=f"{bl}:hardware_busy_lock") + daq = AareDAQ(cfg, bl, bl_dispatch, hardware_busy_lock) + cfg.state = daq.read_current_state_from_bec() try: @@ -377,6 +394,7 @@ async def sample_geometry(token: str = Depends(oauth2_scheme)) -> SampleGeometry @app.put("/beamline/omega") +@needs_hw_lock async def set_omega_abs(val: float, token: str = Depends(oauth2_scheme)): """ Set the omega angle of the goniometer. @@ -395,6 +413,7 @@ async def set_omega_abs(val: float, token: str = Depends(oauth2_scheme)): @app.put("/beamline/omega_rel") +@needs_hw_lock async def set_omega_rel(val: float, token: str = Depends(oauth2_scheme)): """ Move the omega angle relatively. @@ -413,6 +432,7 @@ async def set_omega_rel(val: float, token: str = Depends(oauth2_scheme)): @app.put("/beamline/front_light") +@needs_hw_lock async def front_light(val: float, token: str = Depends(oauth2_scheme)): """ Set the front light intensity. @@ -431,6 +451,7 @@ async def front_light(val: float, token: str = Depends(oauth2_scheme)): @app.put("/beamline/back_light") +@needs_hw_lock async def back_light(val: float, token: str = Depends(oauth2_scheme)): """ Set the back light intensity. @@ -449,6 +470,7 @@ async def back_light(val: float, token: str = Depends(oauth2_scheme)): @app.put("/beamline/zoom") +@needs_hw_lock async def zoom(val: float, token: str = Depends(oauth2_scheme)): """ Set the camera zoom level. @@ -467,6 +489,7 @@ async def zoom(val: float, token: str = Depends(oauth2_scheme)): @app.post("/beamline/mono_pitch_scan") +@needs_hw_lock async def mono_pitch_scan(plot: bool = False, token: str = Depends(oauth2_scheme)): """ Run a monochromator pitch scan. Staff only. @@ -500,6 +523,7 @@ async def steer_beam_available(token: str = Depends(oauth2_scheme)): @app.post("/beamline/steer_beam") +@needs_hw_lock async def steer_beam( x: int | None = None, y: int | None = None, token: str = Depends(oauth2_scheme) ): @@ -540,6 +564,7 @@ async def change_energy(value: float, plot: bool = False, token: str = Depends(o @app.put("/beamline/smargon") +@needs_hw_lock async def smargon(val: SmargonCoordinate, token: str = Depends(oauth2_scheme)): """ Move the Smargon goniometer to specified coordinates. @@ -559,6 +584,7 @@ async def smargon(val: SmargonCoordinate, token: str = Depends(oauth2_scheme)): @app.post("/beamline/tweak_abr_meas_pos") +@needs_hw_lock async def tweak_abr_meas_pos(val: AerotechCoordinate, token: str = Depends(oauth2_scheme)): """ Tweak the Aerotech measurement position. Staff only. @@ -576,6 +602,7 @@ async def tweak_abr_meas_pos(val: AerotechCoordinate, token: str = Depends(oauth @app.post("/beamline/save_abr_meas_pos") +@needs_hw_lock async def save_abr_meas_pos(token: str = Depends(oauth2_scheme)): """ Save the current Aerotech measurement position. Staff only. @@ -593,6 +620,7 @@ async def save_abr_meas_pos(token: str = Depends(oauth2_scheme)): @app.post("/beamline/save_beam_location_camera_setting") +@needs_hw_lock async def save_beam_location_camera_setting(token: str = Depends(oauth2_scheme)): """ Persist the camera's current gain/exposure as the beam-location preset for @@ -607,6 +635,7 @@ async def save_beam_location_camera_setting(token: str = Depends(oauth2_scheme)) @app.post("/beamline/anneal") +@needs_hw_lock async def anneal(time_s: float, token: str = Depends(oauth2_scheme)): """ Perform sample annealing for a specified duration. @@ -625,6 +654,7 @@ async def anneal(time_s: float, token: str = Depends(oauth2_scheme)): @app.post("/smargon/initialize") +@needs_hw_lock async def initialise_smargon(token: str = Depends(oauth2_scheme)) -> dict: """ Initialise Smargon. Staff only. @@ -642,6 +672,7 @@ async def initialise_smargon(token: str = Depends(oauth2_scheme)) -> dict: @app.post("/bec/load_user_macros") +@needs_hw_lock async def bec_load_user_macros(token: str = Depends(oauth2_scheme)) -> dict: """ Load BEC user macros. Staff only. @@ -673,6 +704,7 @@ async def bec_list_all_devices(token: str = Depends(oauth2_scheme)) -> list: @app.post("/bec/reinitialise_planner_and_position_devices") +@needs_hw_lock async def bec_reinitialise_planner_and_position_devices( method: str = "auto", token: str = Depends(oauth2_scheme) ) -> dict: @@ -698,6 +730,7 @@ async def bec_reinitialise_planner_and_position_devices( @app.post("/bec/save_current_bs_pos") +@needs_hw_lock async def bec_save_current_bs_pos(token: str = Depends(oauth2_scheme)) -> dict: """ Save the current BEC beamstop work position. Staff only. @@ -709,6 +742,7 @@ async def bec_save_current_bs_pos(token: str = Depends(oauth2_scheme)) -> dict: @app.post("/bec/save_current_collimator_pos") +@needs_hw_lock async def bec_save_current_collimator_pos(token: str = Depends(oauth2_scheme)) -> dict: """ Save the current BEC collimator work position. Staff only. @@ -720,6 +754,7 @@ async def bec_save_current_collimator_pos(token: str = Depends(oauth2_scheme)) - @app.post("/bec/save_current_aerotech_position") +@needs_hw_lock async def bec_save_current_aerotech_position(token: str = Depends(oauth2_scheme)) -> dict: """ Save the current BEC aerotech work position and reload device config. Staff only. @@ -730,27 +765,23 @@ async def bec_save_current_aerotech_position(token: str = Depends(oauth2_scheme) return {"ok": True, "message": "Saved current BEC aerotech work position and reloaded devices."} -def initialise_aerotech(self): +@app.post("/aerotech/initialize") +@needs_hw_lock +def initialize_aerotech(self): try: self._devs.aerotech.home_aerotech() except Exception as e: - logger.error(f"Failed to initialise Aerotech: {e}") + logger.error(f"Failed to initialize Aerotech: {e}") raise -def detector_take_pedestal(self): - try: - self._jfjoch.take_pedestal() - except Exception as e: - logger.error(f"Failed to take detector pedestal: {e}") - raise - - -def initialise_detector(self): +@app.post("/detector/initialize") +@needs_hw_lock +def initialize_detector(self): try: self._jfjoch.initialize() except Exception as e: - logger.error(f"Failed to initialise detector: {e}") + logger.error(f"Failed to initialize detector: {e}") raise @@ -838,6 +869,7 @@ async def local_contact_restart_device(device: str, token: str = Depends(oauth2_ @app.post("/local_contact/resync/detector_metadata") +@needs_hw_lock async def local_contact_resync_detector_metadata(token: str = Depends(oauth2_scheme)) -> dict: """ Refresh cached detector metadata and DTZ limits. Staff only. @@ -872,6 +904,7 @@ async def local_contact_set_config( @app.post("/beamline/goto_abr_meas_pos") +@needs_hw_lock async def goto_abr_meas_pos(token: str = Depends(oauth2_scheme)): """ Move the Aerotech to the saved measurement position. @@ -889,6 +922,7 @@ async def goto_abr_meas_pos(token: str = Depends(oauth2_scheme)): @app.post("/beam_mark/add") +@needs_hw_lock async def mark_beam(x: float, y: float, token: str = Depends(oauth2_scheme)): """ Mark the beam position on the camera image. Staff only. @@ -908,6 +942,7 @@ async def mark_beam(x: float, y: float, token: str = Depends(oauth2_scheme)): @app.post("/beam_mark/clear") +@needs_hw_lock async def clear_beam_mark(token: str = Depends(oauth2_scheme)): """ Clear the beam mark from the camera image. Staff only. @@ -924,26 +959,8 @@ async def clear_beam_mark(token: str = Depends(oauth2_scheme)): return "OK" -@app.post("/beamline/beam_size_mm") -async def beam_size_mm(x: float, y: float, token: str = Depends(oauth2_scheme)): - """ - Set the beam size in millimeters. Staff only. - - Args: - x: Beam width in mm. - y: Beam height in mm. - token: OAuth2 access token. - - Returns: - "OK" on success. - """ - logger.debug(f"Beam Size {x}, {y}") - auth.check_jwt_staff(cfg, auth.parse_token(token)) - daq.beam_size_mm = Coordinate(x=x, y=y) - return "OK" - - @app.put("/beamline/samcam") +@needs_hw_lock async def samcam_settings(s: SampleCameraSettings, token: str = Depends(oauth2_scheme)): """ Update the sample camera settings (exposure, gain, etc.). @@ -962,6 +979,7 @@ async def samcam_settings(s: SampleCameraSettings, token: str = Depends(oauth2_s @app.put("/beamline/autoexposure") +@needs_hw_lock async def samcam_autoexposure(token: str = Depends(oauth2_scheme)): """ Update the sample camera settings (exposure, gain, etc.). @@ -979,6 +997,7 @@ async def samcam_autoexposure(token: str = Depends(oauth2_scheme)): @app.post("/samcam/autofocus") +@needs_hw_lock async def samcam_autofocus(s: AutofocusSettings, token: str = Depends(oauth2_scheme)): """ Trigger the sample camera autofocus procedure. @@ -997,6 +1016,7 @@ async def samcam_autofocus(s: AutofocusSettings, token: str = Depends(oauth2_sch @app.post("/beamline/shutter") +@needs_hw_lock async def shutter(val: bool, token: str = Depends(oauth2_scheme)): """ Open or close the beamline shutter. @@ -1044,6 +1064,7 @@ async def sample(token: str = Depends(oauth2_scheme)) -> SampleShortInfo: @app.post("/tell/park_and_dry") +@needs_hw_lock async def park_and_dry(token: str = Depends(oauth2_scheme)): """ Execute the 'park and dry' procedure for the sample changer (TELL). @@ -1061,6 +1082,7 @@ async def park_and_dry(token: str = Depends(oauth2_scheme)): @app.post("/tell/toggle_blower") +@needs_hw_lock async def tell_toggle_blower(token: str = Depends(oauth2_scheme)) -> dict: """ Toggle the TELL blower. Staff only. @@ -1073,11 +1095,12 @@ async def tell_toggle_blower(token: str = Depends(oauth2_scheme)) -> dict: """ data = auth.parse_token(token) auth.check_jwt_staff_only(data) - daq.blower_control() + daq.tell_toggle_blower() return {"ok": True, "message": "TELL blower toggled."} @app.post("/sample/mount") +@needs_hw_lock async def mount(dbid: int, token: str = Depends(oauth2_scheme), reference: bool = False): """ Mount a sample from the spreadsheet onto the goniometer. @@ -1118,6 +1141,7 @@ async def mount(dbid: int, token: str = Depends(oauth2_scheme), reference: bool @app.post("/sample/unmount") +@needs_hw_lock async def unmount(token: str = Depends(oauth2_scheme)): """ Unmount the current sample from the goniometer. @@ -1135,6 +1159,7 @@ async def unmount(token: str = Depends(oauth2_scheme)): @app.post("/sample/manual") +@needs_hw_lock async def manual(s: SampleShortInfo, token: str = Depends(oauth2_scheme)): """ Manually create or update a sample. @@ -1151,6 +1176,7 @@ async def manual(s: SampleShortInfo, token: str = Depends(oauth2_scheme)): @app.post("/sample/resync") +@needs_hw_lock async def sample_resync(token: str = Depends(oauth2_scheme)) -> dict: """ Manually trigger a resynchronization of the sample information from the changer (TELL). @@ -1167,7 +1193,7 @@ async def sample_resync(token: str = Depends(oauth2_scheme)) -> dict: return {"ok": True, "message": "TELL sample cache resynced."} -def get_spreadsheet(data: TokenData) -> SampleShortInfoList: +def _get_spreadsheet(data: TokenData) -> SampleShortInfoList: """ Get the sample spreadsheet for the given user/pgroup. @@ -1183,7 +1209,7 @@ def get_spreadsheet(data: TokenData) -> SampleShortInfoList: return cfg.spreadsheet_pgroup(data.pgroups) -def get_reference_tools() -> SampleShortInfoList: +def _get_reference_tools() -> SampleShortInfoList: """ Get the list of reference tools. @@ -1202,7 +1228,7 @@ async def reference_tools_event_stream() -> AsyncGenerator[str, None]: """ try: while True: - yield get_reference_tools().model_dump_json() + yield _get_reference_tools().model_dump_json() await asyncio.sleep(10) except asyncio.CancelledError: return @@ -1220,7 +1246,7 @@ async def spreadsheet_event_stream(data: TokenData) -> AsyncGenerator[str, None] """ try: while True: - yield get_spreadsheet(data).model_dump_json() + yield _get_spreadsheet(data).model_dump_json() await asyncio.sleep(10) except asyncio.CancelledError: return @@ -1281,7 +1307,7 @@ async def spreadsheet(token: str = Depends(oauth2_scheme)) -> SampleShortInfoLis Returns: SampleShortInfoList. """ - return get_spreadsheet(auth.parse_token(token)) + return _get_spreadsheet(auth.parse_token(token)) @app.get("/sample/reference_tools") @@ -1296,11 +1322,12 @@ async def reference_tools(token: str = Depends(oauth2_scheme)) -> SampleShortInf SampleShortInfoList. """ auth.check_jwt_ro(cfg, auth.parse_token(token)) - return get_reference_tools() + return _get_reference_tools() # State transitions @app.post("/state/dewar_exchange") +@needs_hw_lock async def dewar_exchange(token: str = Depends(oauth2_scheme)): """ Transition beamline state to DewarTransfer. @@ -1317,6 +1344,7 @@ async def dewar_exchange(token: str = Depends(oauth2_scheme)): @app.post("/state/sample_exchange") +@needs_hw_lock async def sample_exchange(token: str = Depends(oauth2_scheme)): """ Transition beamline state to SampleExchange. @@ -1333,6 +1361,7 @@ async def sample_exchange(token: str = Depends(oauth2_scheme)): @app.post("/state/sample_alignment") +@needs_hw_lock async def sample_alignment(token: str = Depends(oauth2_scheme)): """ Transition beamline state to SampleAlignment. @@ -1348,6 +1377,7 @@ async def sample_alignment(token: str = Depends(oauth2_scheme)): @app.post("/state/beam_location") +@needs_hw_lock async def beam_location(token: str = Depends(oauth2_scheme)): """ Transition beamline state to BeamLocation. Staff only. @@ -1360,6 +1390,7 @@ async def beam_location(token: str = Depends(oauth2_scheme)): @app.post("/state/beamstop_alignment") +@needs_hw_lock async def beamstop_alignment(token: str = Depends(oauth2_scheme)): """ Transition beamline state to BeamstopAlignment. @@ -1376,6 +1407,7 @@ async def beamstop_alignment(token: str = Depends(oauth2_scheme)): @app.post("/state/flux_measurement") +@needs_hw_lock async def flux_measurement(token: str = Depends(oauth2_scheme)): """ Transition beamline state to FluxMeasurement. @@ -1392,6 +1424,7 @@ async def flux_measurement(token: str = Depends(oauth2_scheme)): @app.post("/state/data_collection") +@needs_hw_lock async def data_collection(token: str = Depends(oauth2_scheme)): """ Transition beamline state to DewarTransfer. @@ -1408,6 +1441,7 @@ async def data_collection(token: str = Depends(oauth2_scheme)): @app.post("/state/robot_sample_exchange") +@needs_hw_lock async def robot_sample_exchange(token: str = Depends(oauth2_scheme)): """ Transition beamline state to SampleExchange. @@ -1424,6 +1458,7 @@ async def robot_sample_exchange(token: str = Depends(oauth2_scheme)): @app.post("/state/xray_fluorescence") +@needs_hw_lock async def xray_fluorescence(token: str = Depends(oauth2_scheme)): """ Transition beamline state to SampleAlignment. @@ -1439,6 +1474,7 @@ async def xray_fluorescence(token: str = Depends(oauth2_scheme)): @app.post("/state/xtal_snapshot") +@needs_hw_lock async def xtal_snapshot(token: str = Depends(oauth2_scheme)): """ Transition beamline state to BeamLocation. Staff only. @@ -1451,6 +1487,7 @@ async def xtal_snapshot(token: str = Depends(oauth2_scheme)): @app.post("/access/take_over_beamline") +@needs_hw_lock async def take_over_beamline( payload: RecoveryActionRequest, token: str = Depends(oauth2_scheme) ) -> str: @@ -1491,6 +1528,7 @@ async def force_clear_busy( data = auth.parse_token(token) auth.check_jwt_staff_only(data) _validate_recovery_code(payload.confirmation_code) + cfg.redis.delete(f"{cfg._mxb}:hardware_busy_lock") logger.warning( "Beamline busy flag cleared via protected endpoint.", extra={"session": getattr(data, "session", None)}, @@ -1540,6 +1578,7 @@ async def force_maintenance_state( @app.post("/recovery/unmount_sample") +@needs_hw_lock async def recovery_unmount_sample( payload: RecoveryActionRequest, token: str = Depends(oauth2_scheme) ) -> dict: @@ -1592,6 +1631,7 @@ async def recovery_unmount_sample( # Scans @app.post("/scan/raster") +@needs_hw_lock async def raster( val: RasterGridRequest, auto_center: bool = False, token: str = Depends(oauth2_scheme) ) -> CompletedRasterGrid: @@ -1611,6 +1651,7 @@ async def raster( @app.post("/scan/rotation") +@needs_hw_lock async def rotation( val: RotationScanRequest, token: str = Depends(oauth2_scheme) ) -> CompletedRotationScan: @@ -1629,6 +1670,7 @@ async def rotation( @app.post("/scan/auto") +@needs_hw_lock async def auto(s: SampleShortInfo, token: str = Depends(oauth2_scheme)): """ Execute a fully automated measurement sequence for a sample. @@ -1674,6 +1716,7 @@ async def set_smart_params(p: SimpleScanParameters, token: str = Depends(oauth2_ @app.post("/scan/cancel") +@needs_hw_lock async def cancel(token: str = Depends(oauth2_scheme)): """ Cancel the currently running scan or automation. @@ -1688,6 +1731,7 @@ async def cancel(token: str = Depends(oauth2_scheme)): # ALC routines @app.post("/alc/center_loop") +@needs_hw_lock async def alc_center_loop(token: str = Depends(oauth2_scheme)) -> str: """ Trigger the automated loop centering procedure (ALC). @@ -1705,6 +1749,7 @@ async def alc_center_loop(token: str = Depends(oauth2_scheme)) -> str: @app.post("/alc/ml_bounding_box") +@needs_hw_lock async def alc_ml_bounding_box(token: str = Depends(oauth2_scheme)) -> RasterGridRequest | None: """ Request an ML-based bounding box for the sample. @@ -1720,6 +1765,7 @@ async def alc_ml_bounding_box(token: str = Depends(oauth2_scheme)) -> RasterGrid @app.post("/face_detection/run") +@needs_hw_lock async def face_detection_run( steps: int, step_size: int, token: str = Depends(oauth2_scheme) ) -> dict: @@ -1808,6 +1854,7 @@ async def pgroup(token: str = Depends(oauth2_scheme)) -> str: @app.put("/access/pgroup") +@needs_hw_lock async def set_pgroup(val: str, token: str = Depends(oauth2_scheme)) -> str: """ Set the active pgroup. @@ -1850,6 +1897,7 @@ async def del_pgroup(token: str = Depends(oauth2_scheme)) -> str: @app.put("/beamline/commissioning_mode") +@needs_hw_lock async def set_commissioning_mode(val: bool, token: str = Depends(oauth2_scheme)) -> str: """ Set the commissioning mode. Staff only. @@ -2225,6 +2273,7 @@ async def get_cryo_settings(token: str = Depends(oauth2_scheme)) -> CryojetSetti @app.put("/beamline/cryo_settings") +@needs_hw_lock async def put_cryo_settings(s: CryojetSettingsModel, token: str = Depends(oauth2_scheme)): """ Update the cryojet settings. Staff only. @@ -2284,6 +2333,7 @@ async def get_all_pgroups(token: str = Depends(oauth2_scheme)): @app.post("/fluorimeter/spectrum") +@needs_hw_lock async def fluorimeter_spectrum( input: FluorescenceSpectrumParameterModel, token: str = Depends(oauth2_scheme) ) -> FluorescenceSpectrumOutputModel: @@ -2302,6 +2352,7 @@ async def fluorimeter_spectrum( @app.post("/fluorimeter/start") +@needs_hw_lock async def fluorimeter_start(erase: bool = False, token: str = Depends(oauth2_scheme)) -> str: """ Start the fluorimeter measurement. @@ -2319,6 +2370,7 @@ async def fluorimeter_start(erase: bool = False, token: str = Depends(oauth2_sch @app.post("/fluorimeter/stop") +@needs_hw_lock async def fluorimeter_stop(token: str = Depends(oauth2_scheme)) -> str: """ Stop the fluorimeter measurement. @@ -2335,6 +2387,7 @@ async def fluorimeter_stop(token: str = Depends(oauth2_scheme)) -> str: @app.get("/fluorimeter/status") +@needs_hw_lock async def fluorimeter_status(token: str = Depends(oauth2_scheme)) -> int | None: """ Get the current fluorimeter status. @@ -2350,6 +2403,7 @@ async def fluorimeter_status(token: str = Depends(oauth2_scheme)) -> int | None: @app.get("/fluorimeter/data") +@needs_hw_lock async def fluorimeter_data(token: str = Depends(oauth2_scheme)) -> list[int] | None: """ Get the latest fluorimeter data. @@ -2365,6 +2419,7 @@ async def fluorimeter_data(token: str = Depends(oauth2_scheme)) -> list[int] | N @app.get("/fluorimeter/background") +@needs_hw_lock async def fluorimeter_background(token: str = Depends(oauth2_scheme)) -> list[int] | None: """ Get the fluorimeter background data. @@ -2453,6 +2508,7 @@ async def sse_fluorimeter(token: str = Depends(oauth2_scheme)): @app.post("/samcam/send_screenshot_db") +@needs_hw_lock async def send_screenshot_db( filename: str | None = None, message: str | None = None, token: str = Depends(oauth2_scheme) ) -> str: @@ -2486,6 +2542,7 @@ async def send_message_db( @app.post("/state/maintenance") +@needs_hw_lock async def maintenance(token: str = Depends(oauth2_scheme)) -> str: """ Transition beamline state to Maintenance. Staff only. diff --git a/src/aare/devices/jfjoch.py b/src/aare/devices/jfjoch.py index d5183100..163a04c8 100644 --- a/src/aare/devices/jfjoch.py +++ b/src/aare/devices/jfjoch.py @@ -360,11 +360,6 @@ class JFJochWrapper: endpoint="config_select_detector_get", ) from e - def take_pedestal(self): - raise NotImplementedError( - "take_pedestal is not implemented in DAQ through the JFJoch API yet" - ) - @needs_init def get_diffraction_image( self, diff --git a/src/aare/gui/panels/local_contact_panel.py b/src/aare/gui/panels/local_contact_panel.py index 3361d978..b8534858 100644 --- a/src/aare/gui/panels/local_contact_panel.py +++ b/src/aare/gui/panels/local_contact_panel.py @@ -466,16 +466,16 @@ class LocalContactPanel(QFrame): "Initialise", [ self._make_button( - "Initialise detector", - self._daq.initialise_detector, - "Initialising detector.", + "Initialize detector", + self._daq.initialize_detector, + "Initializing detector.", ), self._make_button( "Initialise Smargon", self._daq.initialise_smargon, "Initialising Smargon." ), self._make_button( "Initialise Aerotech", - self._daq.initialise_aerotech, + self._daq.initialize_aerotech, "Initialising Aerotech.", ), ], @@ -521,13 +521,6 @@ class LocalContactPanel(QFrame): row, 0, ) - grid.addWidget( - self._make_button( - "Take pedestal", self._daq.detector_take_pedestal, "Requesting detector pedestal." - ), - row, - 1, - ) grid.addWidget( self._make_button( "Resync detector/DTZ hardware cache", diff --git a/src/aare/gui/threads/daq_worker.py b/src/aare/gui/threads/daq_worker.py index 81579a2b..a5989aa7 100644 --- a/src/aare/gui/threads/daq_worker.py +++ b/src/aare/gui/threads/daq_worker.py @@ -1777,16 +1777,12 @@ class DAQWorker(QObject): self.generic_post("beamline/save_beam_location_camera_setting") @Slot() - def initialise_aerotech(self): + def initialize_aerotech(self): logger.info("initisalisation does not initisalise aareSCAN but runs homing script") self.generic_post("aerotech/initialize") @Slot() - def detector_take_pedestal(self): - self.generic_post("detector/take_pedestal") - - @Slot() - def initialise_detector(self): + def initialize_detector(self): self.generic_post("detector/initialize") @Slot() diff --git a/tests/unit/daq/test_automation_progress_state_manager.py b/tests/unit/daq/test_automation_progress_state_manager.py index 8d6f2b3c..c03c3df2 100644 --- a/tests/unit/daq/test_automation_progress_state_manager.py +++ b/tests/unit/daq/test_automation_progress_state_manager.py @@ -32,7 +32,7 @@ class _FakeRedis: def _make_config_with_fake_redis() -> BeamlineConfig: cfg = BeamlineConfig.__new__(BeamlineConfig) cfg._bl = "testbeamline" - cfg._client = _FakeRedis() + cfg.redis = _FakeRedis() return cfg -- 2.54.0 From 6312804a65137032a61208ab7e6b5239122cfa1f Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 8 Sep 2026 15:44:47 +0200 Subject: [PATCH 3/4] style: ignore pyright issues in redis return values --- src/aare/daq/auth.py | 4 +- src/aare/daq/config.py | 107 +++++++++++++++++++---------------------- src/aare/daq/daq.py | 2 +- src/aare/daq/server.py | 9 ++-- 4 files changed, 58 insertions(+), 64 deletions(-) diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index aeadb149..68fd2ddc 100644 --- a/src/aare/daq/auth.py +++ b/src/aare/daq/auth.py @@ -99,8 +99,8 @@ def authenticate_user(cfg: BeamlineConfig, username: str) -> str: def parse_token(token: str = Depends(oauth2_scheme)) -> TokenData: try: payload = jwt.decode(token, jwt_key(), algorithms=[ALGORITHM]) - token = TokenData(**payload) - return token + _token = TokenData(**payload) + return _token except jwt.PyJWTError as e: raise AuthenticationException( message="Invalid token", diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index 60902e67..14db0716 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -4,9 +4,9 @@ import json import time from dataclasses import asdict, is_dataclass from datetime import datetime +from typing import Any import numpy as np -import redis from aarecommon.config.beamline import cfg_get from aarecommon.config.logger import setup_logger from aarecommon.math.coordinate import AerotechCoordinate, Coordinate @@ -37,6 +37,7 @@ from aarecommon.models.models import ( ZoomModel, zoom_manager, ) +from redis.client import PubSub, Redis from redis.lock import Lock as RedisLock from aare.daq.config_model import LocalContactConfigModel @@ -99,7 +100,8 @@ class BeamlineConfig: host = "localhost" else: host = cfg_get("daq.hardware.redis_url", f"{self._bl}-redis.psi.ch") - self.redis = redis.Redis(host=host, port=6379, db=0, decode_responses=True) + self.hw_lock = RedisLock(cfg.redis, name=f"{bl}:hardware_busy_lock") + self.redis = Redis(host=host, port=6379, db=0, decode_responses=True) self.simulated_detector = bl is MXBeamline.SIMULATED self._initialize_optional_yaml_defaults() @@ -144,7 +146,7 @@ class BeamlineConfig: def _current_gui_session_ttl(self, session: int) -> int | None: try: - ttl = int(self.redis.ttl(self._gui_session_key(session))) + ttl = int(self.redis.ttl(self._gui_session_key(session))) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 except Exception: logger.debug("Could not read the GUI session TTL", exc_info=True) return None @@ -158,7 +160,7 @@ class BeamlineConfig: if raw is None: return None try: - return OpenGuiSessionInfo(**json.loads(raw)) + return OpenGuiSessionInfo(**json.loads(raw)) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 except Exception: logger.warning(f"Failed to parse GUI session info for session {session}", exc_info=True) return None @@ -239,7 +241,7 @@ class BeamlineConfig: self.purge_expired_gui_sessions() def purge_expired_gui_sessions(self) -> None: - session_ids = self.redis.smembers(self._gui_sessions_index_key()) + session_ids: set[str] = self.redis.smembers(self._gui_sessions_index_key()) # pyright: ignore[reportArgumentType, reportAssignmentType] # using sync client - remove ignore on upgrade to redis v8 if not session_ids: return @@ -297,7 +299,7 @@ class BeamlineConfig: self.redis.delete(f"{self._bl}:allow_non_staff_request_from_staff") def generate_session(self) -> int: - return int(self.redis.incr(f"{self._bl}:session")) + return int(self.redis.incr(f"{self._bl}:session")) # pyright: ignore[reportArgumentType, reportAssignmentType] # using sync client - remove ignore on upgrade to redis v8 @property def active_session(self) -> int | None: @@ -368,7 +370,7 @@ class BeamlineConfig: # Redis compatibility: # some Redis servers do not support EXPIRE with GT option. # Emulate "extend only if greater" manually while holding the lock. - if current_ttl is None or current_ttl < 0 or current_ttl < expiry_sec: + if current_ttl is None or float(current_ttl) < 0 or float(current_ttl) < expiry_sec: # pyright: ignore[reportArgumentType, reportAssignmentType] # using sync client - remove ignore on upgrade to redis v8 self.redis.expire(key, expiry_sec) else: raise RuntimeError( @@ -458,10 +460,7 @@ class BeamlineConfig: def can_transfer_baton_now(self) -> bool: """Check if baton can be transferred (beamline not mid-operation).""" - # Can't transfer while beamline is busy - # Add automation queue check here when you implement it - # return not (self.state_busy or self.automation_queue_running) - raise NotImplementedError("Checking for busy state is not implemented") + return not self.hw_lock.locked() def execute_baton_transfer( self, @@ -508,8 +507,7 @@ class BeamlineConfig: @property def pgroup(self) -> str | None: - tmp = self.redis.get(f"{self._bl}:pgroup") - return tmp + return self.redis.get(f"{self._bl}:pgroup") # pyright: ignore[reportReturnType] # using sync client - remove ignore on upgrade to redis v8 @pgroup.setter def pgroup(self, pgroup: str | None) -> None: @@ -520,8 +518,7 @@ class BeamlineConfig: @property def commissioning_mode(self) -> bool: - tmp = self.redis.get(f"{self._bl}:commissioning_mode") - return tmp is not None + return self.redis.get(f"{self._bl}:commissioning_mode") is not None @commissioning_mode.setter def commissioning_mode(self, commisioning_mode: bool) -> None: @@ -545,8 +542,8 @@ class BeamlineConfig: if raw_value is None: return BeamlineStateEnum.Maintenance try: - int_value = int(raw_value) # Ensure it's an integer - return BeamlineStateEnum(int_value) # Convert to BeamlineStateEnum + int_value = int(raw_value) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + return BeamlineStateEnum(int_value) except (ValueError, KeyError): raise ValueError( f"Invalid 'mx_state' value: {raw_value}. Expected integer corresponding to a BeamlineStateEnum." @@ -558,12 +555,12 @@ class BeamlineConfig: @property def tell_mount_count(self) -> int: - return int(self.redis.incr(f"{self._bl}:tell_mount_count")) + return int(self.redis.incr(f"{self._bl}:tell_mount_count")) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def pixel_to_mm(self, zoom: float) -> float: cfg = self.settings base_pixel_in_mm = 1.0 / ( - cfg.camera_translation_factor_b * np.exp(cfg.camera_translation_factor_a * zoom) + cfg.camera_translation_factor_b * np.exp(cfg.camera_translation_factor_a * zoom) # pyright: ignore[reportOptionalOperand] ) # Apply lens magnification correction relative to the default 10x lens. # A lower magnification lens (e.g. 5x) makes each pixel cover more physical space. @@ -600,17 +597,9 @@ class BeamlineConfig: @property def beam_size_mm(self) -> Coordinate: - tmp_x = self.redis.get(f"{self._bl}:beam_size_x") - tmp_y = self.redis.get(f"{self._bl}:beam_size_y") - if tmp_x: - val_x = float(tmp_x) - else: - val_x = 0.04 - if tmp_y: - val_y = float(tmp_y) - else: - val_y = 0.04 - return Coordinate(x=val_x, y=val_y) + x = float(self.redis.get(f"{self._bl}:beam_size_x") or 0.4) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + y = float(self.redis.get(f"{self._bl}:beam_size_y") or 0.4) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + return Coordinate(x=x, y=y) @beam_size_mm.setter def beam_size_mm(self, data: Coordinate): @@ -622,7 +611,7 @@ class BeamlineConfig: if tmp is None: return BeamlineSettingsModel() - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 model = BeamlineSettingsModel(**data_dict) return model @@ -644,7 +633,7 @@ class BeamlineConfig: if tmp is None: return CryojetSettingsModel() - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return CryojetSettingsModel(**data_dict) @cryojet_settings.setter @@ -652,7 +641,7 @@ class BeamlineConfig: self.redis.set(f"{self._bl}:cryojet_settings", data.model_dump_json()) def get_alc_bkg(self, zoom: float, exp: float, gain: float) -> np.ndarray | None: - return base64_to_numpy(self.redis.get(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}")) + return base64_to_numpy(self.redis.get(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}")) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def put_alc_bkg(self, zoom: float, exp: float, gain: float, data: np.ndarray): self.redis.set(f"{self._bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}", numpy_to_base64(data)) @@ -663,7 +652,7 @@ class BeamlineConfig: if tmp is None: return SampleShortInfoList(s=[]) - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return SampleShortInfoList(**data_dict) def spreadsheet_pgroup(self, pgroups: list[str]) -> SampleShortInfoList: @@ -675,7 +664,7 @@ class BeamlineConfig: def spreadsheet(self, data: SampleShortInfoList): self.redis.set(f"{self._bl}:sample_spreadsheet", data.model_dump_json()) - def listen_changes_spreadsheet(self) -> redis.client.PubSub: + def listen_changes_spreadsheet(self) -> PubSub: self.redis.config_set("notify-keyspace-events", "KEA") pubsub = self.redis.pubsub() pubsub.psubscribe(f"__keyspace@0__:{self._bl}:sample_spreadsheet") @@ -687,7 +676,7 @@ class BeamlineConfig: if tmp is None: return SampleShortInfoList(s=[]) - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return SampleShortInfoList(**data_dict) @reference_tools.setter @@ -706,7 +695,7 @@ class BeamlineConfig: if tmp is None: return None - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return SampleShortInfo(**data_dict) @current_sample.setter @@ -737,8 +726,8 @@ class BeamlineConfig: print("no zoom mode given, defaulting to user mode") return ZoomModeEnum.User try: - int_value = int(raw_value) # Ensure it's an integer - return ZoomModeEnum(int_value) # Convert to ZoomModeEnum + int_value = int(raw_value) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + return ZoomModeEnum(int_value) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 except (ValueError, KeyError): raise ValueError( f"Invalid 'mx_zoom_mode' value: {raw_value}. Expected integer corresponding to a ZoomModeEnum." @@ -788,9 +777,9 @@ class BeamlineConfig: mode = mode or self.zoom_mode key = f"{self._bl}:{self.zoom_setting_string(mode)}" tmp = self.redis.get(key) - model = ZoomModel(**json.loads(tmp)) if tmp is not None else zoom_manager(mode, self._mxb) - model.z[zoom_value] = settings - self.redis.set(key, model.model_dump_json()) + model = ZoomModel(**json.loads(tmp)) if tmp is not None else zoom_manager(mode, self._mxb) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + model.z[zoom_value] = settings # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + self.redis.set(key, model.model_dump_json()) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 @property def abr_meas_pos(self) -> AerotechCoordinate: @@ -798,7 +787,7 @@ class BeamlineConfig: if tmp is None: return ABR_POS_MOUNT - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return AerotechCoordinate(**data_dict) @abr_meas_pos.setter @@ -810,7 +799,7 @@ class BeamlineConfig: tmp = self.redis.get(f"{self._bl}:dtz") if tmp is None: return None - return float(tmp) + return float(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 @dtz.setter def dtz(self, dtz: float): @@ -821,7 +810,7 @@ class BeamlineConfig: tmp = self.redis.get(f"{self._bl}:dtz_safe_position") if tmp is None: return None - return float(tmp) + return float(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 @dtz_safe_position.setter def dtz_safe_position(self, dtz: float): @@ -851,11 +840,11 @@ class BeamlineConfig: ) vals = self.redis.hgetall(f"{self._bl}:beam_mark") - if len(vals) >= 3: + if len(vals) >= 3: # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 zooms = [] x_pxls = [] y_pxls = [] - for k, v in vals.items(): + for k, v in vals.items(): # pyright: ignore[reportArgumentType,reportAttributeAccessIssue] # using sync client - remove ignore on upgrade to redis v8 zooms.append(float(k)) x_pxls.append(float(json.loads(v)["x"])) y_pxls.append(float(json.loads(v)["y"])) @@ -917,10 +906,10 @@ class BeamlineConfig: def get_mount_failure_streak(self) -> int: value = self.redis.get(self._mount_failure_streak_key()) - return int(value) if value else 0 + return int(value) if value else 0 # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def increment_mount_failure_streak(self) -> int: - return int(self.redis.incr(self._mount_failure_streak_key())) + return int(self.redis.incr(self._mount_failure_streak_key())) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def reset_mount_failure_streak(self) -> None: self.redis.delete(self._mount_failure_streak_key()) @@ -979,16 +968,18 @@ class BeamlineConfig: self.redis.set(self._automation_progress_seq_key(), 0) self.redis.delete(self._automation_progress_key()) - def get_automation_progress_state(self) -> dict: + def get_automation_progress_state(self) -> dict[str, Any]: seq_raw = self.redis.get(self._automation_progress_seq_key()) payload_raw = self.redis.get(self._automation_progress_key()) - seq = int(seq_raw) if seq_raw is not None else 0 - progress = json.loads(payload_raw) if payload_raw else None + seq = int(seq_raw) if seq_raw is not None else 0 # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 + progress = json.loads(payload_raw) if payload_raw else None # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return {"seq": seq, "progress": progress} - def set_automation_progress_state(self, progress: AutomationProgress | dict) -> dict: + def set_automation_progress_state( + self, progress: AutomationProgress | dict[str, Any] + ) -> dict[str, Any]: def _json_default(value): if isinstance(value, datetime): return value.isoformat() @@ -1003,7 +994,7 @@ class BeamlineConfig: else: raise TypeError(f"Unsupported automation progress type: {type(progress).__name__}") - next_seq = int(self.redis.incr(self._automation_progress_seq_key())) + next_seq = int(self.redis.incr(self._automation_progress_seq_key())) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 self.redis.set( self._automation_progress_key(), json.dumps(payload, separators=(",", ":"), default=_json_default), @@ -1017,7 +1008,7 @@ class BeamlineConfig: if tmp is None: return 0 try: - return int(tmp) + return int(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 except (TypeError, ValueError): logger.warning( "Failed Mount Count is not an integer, resetting to 0.", @@ -1033,7 +1024,7 @@ class BeamlineConfig: self.redis.set(f"{self._bl}:failed_mount_count", count) def increment_failed_mount_count(self) -> int: - return int(self.redis.incr(f"{self._bl}:failed_mount_count")) + return int(self.redis.incr(f"{self._bl}:failed_mount_count")) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 def _runtime_sim_key(self, name: str) -> str: return f"{self._bl}:runtime:simulate:{name}" @@ -1132,7 +1123,7 @@ class BeamlineConfig: logger.warning(f"Failed to read detector metadata from Redis: {e}", exc_info=True) return {} - def set_detector_metadata(self, payload: dict) -> dict: + def set_detector_metadata(self, payload: dict[str, Any]) -> dict[str, Any]: safe_payload = dict(payload or {}) safe_payload["dtz_low"] = self._coerce_optional_float(safe_payload.get("dtz_low")) safe_payload["dtz_high"] = self._coerce_optional_float(safe_payload.get("dtz_high")) @@ -1144,7 +1135,7 @@ class BeamlineConfig: return safe_payload @property - def cached_detector_metadata(self) -> dict: + def cached_detector_metadata(self) -> dict[str, Any]: return self.get_detector_metadata() @property diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 47570a32..4f301e4e 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -2190,7 +2190,7 @@ class AareDAQ: sample=self.sample, sample_geometry=self.sample_geometry, filename=filename, - upload_image=self._aare.upload_image, + upload_image=self._aare.upload_image, # pyright: ignore[reportArgumentType] logger=logger, max_images=self.AUTO_RASTER_MAX_IMAGES, min_cell_size_mm=self.AUTO_RASTER_MAX_IMAGES, diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index 7f6c192f..b7c39029 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import asyncio import functools import hmac @@ -7,6 +9,7 @@ import os import time from collections.abc import AsyncGenerator from contextlib import asynccontextmanager +from typing import Any import uvicorn from aarecommon.config.beamline import mx_beamline @@ -123,7 +126,7 @@ async def lifespan(application: FastAPI): # ── Core objects (Redis, EPICS PVs, BEC, TELL, JFJoch, etc.) ── bl = mx_beamline() cfg = BeamlineConfig(bl) - hardware_busy_lock = RedisLock(cfg.redis, name=f"{bl}:hardware_busy_lock") + hardware_busy_lock = cfg.hw_lock daq = AareDAQ(cfg, bl, bl_dispatch, hardware_busy_lock) cfg.state = daq.read_current_state_from_bec() @@ -655,7 +658,7 @@ async def anneal(time_s: float, token: str = Depends(oauth2_scheme)): @app.post("/smargon/initialize") @needs_hw_lock -async def initialise_smargon(token: str = Depends(oauth2_scheme)) -> dict: +async def initialise_smargon(token: str = Depends(oauth2_scheme)) -> dict[str, Any]: """ Initialise Smargon. Staff only. @@ -673,7 +676,7 @@ async def initialise_smargon(token: str = Depends(oauth2_scheme)) -> dict: @app.post("/bec/load_user_macros") @needs_hw_lock -async def bec_load_user_macros(token: str = Depends(oauth2_scheme)) -> dict: +async def bec_load_user_macros(token: str = Depends(oauth2_scheme)) -> dict[str, Any]: """ Load BEC user macros. Staff only. """ -- 2.54.0 From f0d3b4c9dd701a777b0b73a1a59c3dd6e42de49c Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 8 Sep 2026 16:42:27 +0200 Subject: [PATCH 4/4] fix: add tests and cleanup after exceptions --- pyproject.toml | 1 + src/aare/daq/auth.py | 2 + src/aare/daq/config.py | 6 +-- src/aare/daq/server.py | 20 +++++--- tests/unit/daq/test_server_hw_lock.py | 73 +++++++++++++++++++++++++++ 5 files changed, 91 insertions(+), 11 deletions(-) create mode 100644 tests/unit/daq/test_server_hw_lock.py diff --git a/pyproject.toml b/pyproject.toml index dcc0fd5e..3ad6b2d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ test = [ "pytest-qt", "pytest-asyncio", "pytest-timeout", + "fakeredis[lua]", "diff-cover", "ruff>=0.15" ] diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index 68fd2ddc..2269fb8f 100644 --- a/src/aare/daq/auth.py +++ b/src/aare/daq/auth.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import grp import logging import os diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index 14db0716..060fd155 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -100,8 +100,8 @@ class BeamlineConfig: host = "localhost" else: host = cfg_get("daq.hardware.redis_url", f"{self._bl}-redis.psi.ch") - self.hw_lock = RedisLock(cfg.redis, name=f"{bl}:hardware_busy_lock") self.redis = Redis(host=host, port=6379, db=0, decode_responses=True) + self.hw_lock = RedisLock(self.redis, name=f"{bl}:hardware_busy_lock") self.simulated_detector = bl is MXBeamline.SIMULATED self._initialize_optional_yaml_defaults() @@ -683,7 +683,7 @@ class BeamlineConfig: def reference_tools(self, data: SampleShortInfoList): self.redis.set(f"{self._bl}:reference-tools", data.model_dump_json()) - def listen_changes_reference_tools(self) -> redis.client.PubSub: + def listen_changes_reference_tools(self) -> PubSub: self.redis.config_set("notify-keyspace-events", "KEA") pubsub = self.redis.pubsub() pubsub.psubscribe(f"__keyspace@0__:{self._bl}:reference-tools") @@ -711,7 +711,7 @@ class BeamlineConfig: if tmp is None: return BeamMarkCoeffModel() - data_dict = json.loads(tmp) + data_dict = json.loads(tmp) # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8 return BeamMarkCoeffModel(**data_dict) @beam_mark_coeff.setter diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index b7c39029..60b293af 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -95,8 +95,10 @@ async def _lock_hw(): if not hardware_busy_lock.acquire(blocking=False): raise BeamlineBusyException("Beamline hardware lock is held by another worker") logger.debug("Hardware lock acquired by process") - yield - hardware_busy_lock.release() + try: + yield + finally: + hardware_busy_lock.release() logger.debug("Hardware lock released by process") @@ -734,7 +736,7 @@ async def bec_reinitialise_planner_and_position_devices( @app.post("/bec/save_current_bs_pos") @needs_hw_lock -async def bec_save_current_bs_pos(token: str = Depends(oauth2_scheme)) -> dict: +async def bec_save_current_bs_pos(token: str = Depends(oauth2_scheme)) -> dict[str, Any]: """ Save the current BEC beamstop work position. Staff only. """ @@ -746,7 +748,7 @@ async def bec_save_current_bs_pos(token: str = Depends(oauth2_scheme)) -> dict: @app.post("/bec/save_current_collimator_pos") @needs_hw_lock -async def bec_save_current_collimator_pos(token: str = Depends(oauth2_scheme)) -> dict: +async def bec_save_current_collimator_pos(token: str = Depends(oauth2_scheme)) -> dict[str, Any]: """ Save the current BEC collimator work position. Staff only. """ @@ -758,7 +760,7 @@ async def bec_save_current_collimator_pos(token: str = Depends(oauth2_scheme)) - @app.post("/bec/save_current_aerotech_position") @needs_hw_lock -async def bec_save_current_aerotech_position(token: str = Depends(oauth2_scheme)) -> dict: +async def bec_save_current_aerotech_position(token: str = Depends(oauth2_scheme)) -> dict[str, Any]: """ Save the current BEC aerotech work position and reload device config. Staff only. """ @@ -873,7 +875,9 @@ async def local_contact_restart_device(device: str, token: str = Depends(oauth2_ @app.post("/local_contact/resync/detector_metadata") @needs_hw_lock -async def local_contact_resync_detector_metadata(token: str = Depends(oauth2_scheme)) -> dict: +async def local_contact_resync_detector_metadata( + token: str = Depends(oauth2_scheme), +) -> dict[str, Any]: """ Refresh cached detector metadata and DTZ limits. Staff only. """ @@ -1086,7 +1090,7 @@ async def park_and_dry(token: str = Depends(oauth2_scheme)): @app.post("/tell/toggle_blower") @needs_hw_lock -async def tell_toggle_blower(token: str = Depends(oauth2_scheme)) -> dict: +async def tell_toggle_blower(token: str = Depends(oauth2_scheme)) -> dict[str, Any]: """ Toggle the TELL blower. Staff only. @@ -1180,7 +1184,7 @@ async def manual(s: SampleShortInfo, token: str = Depends(oauth2_scheme)): @app.post("/sample/resync") @needs_hw_lock -async def sample_resync(token: str = Depends(oauth2_scheme)) -> dict: +async def sample_resync(token: str = Depends(oauth2_scheme)) -> dict[str, Any]: """ Manually trigger a resynchronization of the sample information from the changer (TELL). diff --git a/tests/unit/daq/test_server_hw_lock.py b/tests/unit/daq/test_server_hw_lock.py new file mode 100644 index 00000000..9113f162 --- /dev/null +++ b/tests/unit/daq/test_server_hw_lock.py @@ -0,0 +1,73 @@ +from unittest.mock import patch + +import fakeredis +import pytest +from redis.lock import Lock as RedisLock + +LOCK_NAME = "SIMULATED:hardware_busy_lock" +AUTH = {"Authorization": "Bearer fake-token"} + +THREAD_LOCAL = False + + +@pytest.fixture +def fake_redis(): + return fakeredis.FakeRedis(decode_responses=True) + + +@pytest.fixture +def hw_lock(server_module, client, fake_redis, monkeypatch): + """This worker's lock. + + Depends on ``client`` so that TestClient's lifespan (which overwrites + ``hardware_busy_lock``) has already run by the time we patch. + """ + lock = RedisLock(fake_redis, name=LOCK_NAME, thread_local=THREAD_LOCAL) + monkeypatch.setattr(server_module, "hardware_busy_lock", lock) + return lock + + +@pytest.fixture +def other_worker(fake_redis): + """A second uvicorn worker's view of the same lock in Redis.""" + return RedisLock(fake_redis, name=LOCK_NAME, thread_local=THREAD_LOCAL) + + +def test_endpoint_acquires_and_releases_the_lock(client, hw_lock, other_worker): + with patch("aare.daq.auth.check_jwt_rw"), patch("aare.daq.server.daq"): + response = client.put("/beamline/omega?val=10.5", headers=AUTH) + + assert response.status_code == 200 + # Released on the way out, so a second worker can take it. + assert not hw_lock.owned() + assert other_worker.acquire(blocking=False) + + +def test_endpoint_is_503_while_another_worker_holds_the_lock(client, hw_lock, other_worker): + assert other_worker.acquire(blocking=False) + + with patch("aare.daq.auth.check_jwt_rw"), patch("aare.daq.server.daq"): + response = client.put("/beamline/omega?val=10.5", headers=AUTH) + + assert response.status_code == 503 + assert response.json()["exception_class"] == "BeamlineBusyException" + # The rejected request must not have stolen or released the other worker's lock. + assert other_worker.owned() + assert not hw_lock.owned() + + +async def test_lock_hw_is_reentrant_for_the_owning_worker(server_module, hw_lock): + assert hw_lock.acquire(blocking=False) + + async with server_module._lock_hw(): + assert hw_lock.owned() + + assert hw_lock.owned(), "re-entrant exit released a lock it did not acquire" + + +async def test_lock_is_released_when_the_handler_raises(server_module, hw_lock): + with pytest.raises(RuntimeError): + async with server_module._lock_hw(): + raise RuntimeError("handler blew up") + + assert not hw_lock.owned(), "hardware lock leaked after an exception" -- 2.54.0