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..a39a0928 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: @@ -345,48 +345,48 @@ class BeamlineConfig: 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) + 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", expire=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", expire=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", expire=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", expire=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", expire=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", 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()) + 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