Beam suspender: replace the plain SuspendFloor with one that checks the ring's operating mode

src/hextools/profiles/collection.py

@@ -29,12 +29,13 @@
 from ophyd_async.fastcs.panda import HDFPanda
 from tiled.client import from_uri, simple
 from bluesky import plans as bp, plan_stubs as bps, preprocessors as bpp
-from bluesky.suspenders import SuspendFloor
+from bluesky.suspenders import SuspendFloor  # noqa: F401 - operator namespace
 from hextools.utils import show_docs
 
 from hextools.detectors.phantom import PhantomDetector
 from hextools.detectors.kinetix import kinetix_factory
 from hextools.machine import NSLS2OpsMode, NSLS2StorageRing
+from hextools.suspenders import SuspendFloorUnlessOpsMode
 from hextools.motors import (
     CollimatorTable,
     FOV_2_4_mm_Camera,
@@ -236,7 +237,18 @@
     sim_z = SimMotor(name="sim_z", instant=False)
 
 
-RE.install_suspender(SuspendFloor(storage_ring.beam_current, 100, resume_thresh=390))
+# Pause the RunEngine if the beam current drops below 100 mA and resume above
+# 390 mA. Always installed; it ignores low beam while the ring is in maintenance
+# or shutdown, checking the mode each time rather than once at startup.
+RE.install_suspender(
+    SuspendFloorUnlessOpsMode(
+        storage_ring.beam_current,
+        100,
+        resume_thresh=390,
+        ops_mode=storage_ring.operating_mode,
+        skip_modes=(NSLS2OpsMode.MAINTENANCE, NSLS2OpsMode.SHUTDOWN),
+    )
+)
 
 # Configure baseline supplemental data to include in the metadata of every run.
 sd = bpp.SupplementalData(

src/hextools/suspenders.py

@@ -0,0 +1,130 @@
+"""RunEngine suspenders for the HEX beamline."""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable, Iterable
+from typing import Any
+
+from bluesky.protocols import Subscribable
+from bluesky.suspenders import SuspendFloor
+
+
+def _call_on_loop(loop, func: Callable[[], None]) -> None:
+    """Run ``func`` on the RunEngine's event loop and wait for it."""
+    if threading.get_ident() == getattr(loop, "_thread_id", None):
+        func()
+        return
+    done = threading.Event()
+    errors: list[BaseException] = []
+
+    def call():
+        try:
+            func()
+        except BaseException as err:
+            errors.append(err)
+        finally:
+            done.set()
+
+    loop.call_soon_threadsafe(call)
+    done.wait()
+    if errors:
+        raise errors[0]
+
+
+class SuspendFloorUnlessOpsMode(SuspendFloor):
+    """Suspend when a signal falls below a threshold, unless in a skipped ring mode.
+
+    Install it whatever the ring's mode: the mode is checked every time the
+    suspender is evaluated, so a ring entering maintenance releases a tripped
+    suspender, and a ring returning to operations with beam still low trips it.
+    Until the mode has been read, the suspender acts like a plain `SuspendFloor`.
+
+    Parameters
+    ----------
+    signal : Signal
+        The signal to watch, e.g. the storage ring beam current.
+    suspend_thresh : float
+        Suspend if the signal falls below this value.
+    ops_mode : Signal
+        The ring's operating mode. Like ``signal``, either an ophyd-async signal
+        or one following ophyd's subscription pattern.
+    skip_modes : Iterable
+        Modes in which a low signal does not suspend the RunEngine.
+    **kwargs
+        Passed to `SuspendFloor` (``resume_thresh``, ``sleep``, ...).
+    """
+
+    def __init__(
+        self,
+        signal,
+        suspend_thresh: float,
+        *,
+        ops_mode,
+        skip_modes: Iterable[Any],
+        **kwargs,
+    ):
+        super().__init__(signal, suspend_thresh, **kwargs)
+        self._ops_mode = ops_mode
+        self._skip_modes = frozenset(skip_modes)
+        self._mode = None
+        self._mode_cid = None
+        # The base _lock is not reentrant, and _set_mode calls back into __call__.
+        self._eval_lock = threading.RLock()
+
+    @property
+    def _skipping(self) -> bool:
+        return self._mode in self._skip_modes
+
+    def _should_suspend(self, value):
+        return not self._skipping and super()._should_suspend(value)
+
+    def _should_resume(self, value):
+        return self._skipping or super()._should_resume(value)
+
+    def _on_mode_reading(self, reading):
+        self._set_mode(reading[self._ops_mode.name]["value"])
+
+    def _on_mode_value(self, value=None, **kwargs):
+        self._set_mode(value)
+
+    def __call__(self, value, **kwargs):
+        with self._eval_lock:
+            super().__call__(value, **kwargs)
+
+    def _set_mode(self, mode):
+        with self._eval_lock:
+            self._mode = mode
+            if self._last_value is None:
+                return
+            value = self._last_value
+            if self._implements_protocol:
+                value = {self._sig.name: {"value": value}}
+            self(value=value)
+
+    def install(self, RE, *, event_type=None):
+        # The mode first, so the first beam reading is judged against it.
+        if isinstance(self._ops_mode, Subscribable):
+            _call_on_loop(
+                RE.loop, lambda: self._ops_mode.subscribe_reading(self._on_mode_reading)
+            )
+        else:
+            self._mode_cid = self._ops_mode.subscribe(self._on_mode_value, run=True)
+        super().install(RE, event_type=event_type)
+
+    def remove(self):
+        RE = self.RE
+        super().remove()
+        if self._mode_cid is not None:
+            self._ops_mode.unsubscribe(self._mode_cid)
+            self._mode_cid = None
+        elif RE is not None and isinstance(self._ops_mode, Subscribable):
+            _call_on_loop(
+                RE.loop, lambda: self._ops_mode.clear_sub(self._on_mode_reading)
+            )
+
+    def _get_justification(self):
+        just = super()._get_justification()
+        if just and self._mode is not None:
+            just = f"{just} (ring mode: {getattr(self._mode, 'value', self._mode)})"
+        return just

tests/test_suspenders.py

@@ -0,0 +1,146 @@
+from collections.abc import Iterator
+
+import pytest
+from bluesky.run_engine import RunEngine
+from ophyd_async.core import init_devices, set_mock_value
+
+from hextools.machine import NSLS2OpsMode, NSLS2StorageRing
+from hextools.suspenders import SuspendFloorUnlessOpsMode
+
+SKIP = (NSLS2OpsMode.MAINTENANCE, NSLS2OpsMode.SHUTDOWN)
+
+
+@pytest.fixture
+def ring(RE: RunEngine) -> NSLS2StorageRing:
+    with init_devices(mock=True):
+        storage_ring = NSLS2StorageRing()
+    return storage_ring
+
+
+@pytest.fixture
+def suspender(
+    RE: RunEngine, ring: NSLS2StorageRing
+) -> Iterator[SuspendFloorUnlessOpsMode]:
+    # Installed whatever the mode at startup: the mode is checked each time it trips.
+    susp = SuspendFloorUnlessOpsMode(
+        ring.beam_current,
+        100,
+        resume_thresh=390,
+        ops_mode=ring.operating_mode,
+        skip_modes=SKIP,
+    )
+    RE.install_suspender(susp)
+    yield susp
+    RE.remove_suspender(susp)
+
+
+def test_trips_on_low_beam_in_operations(ring, suspender):
+    set_mock_value(ring.beam_current, 50)
+    assert suspender.tripped
+
+
+@pytest.mark.parametrize("mode", SKIP)
+def test_ignores_low_beam_in_maintenance_or_shutdown(ring, suspender, mode):
+    set_mock_value(ring.operating_mode, mode)
+    set_mock_value(ring.beam_current, 50)
+    assert not suspender.tripped
+
+
+def test_releases_when_ring_enters_maintenance_while_tripped(ring, suspender):
+    set_mock_value(ring.beam_current, 50)
+    assert suspender.tripped
+    set_mock_value(ring.operating_mode, NSLS2OpsMode.MAINTENANCE)
+    assert not suspender.tripped
+
+
+def test_trips_when_operations_resume_with_beam_still_low(ring, suspender):
+    set_mock_value(ring.operating_mode, NSLS2OpsMode.SHUTDOWN)
+    set_mock_value(ring.beam_current, 0)
+    assert not suspender.tripped
+    set_mock_value(ring.operating_mode, NSLS2OpsMode.OPERATIONS)
+    assert suspender.tripped
+
+
+def test_resumes_on_beam_recovery_in_operations(ring, suspender):
+    set_mock_value(ring.beam_current, 50)
+    set_mock_value(ring.beam_current, 395)
+    assert not suspender.tripped
+
+
+@pytest.mark.parametrize("legacy", ["beam", "mode", "both"])
+def test_legacy_ophyd_signals(RE: RunEngine, ring: NSLS2StorageRing, legacy):
+    # SuspendFloor also takes ophyd signals, whose callbacks pass value=...
+    # rather than a reading; both sides must work in either style (PR 93 review).
+    from ophyd import Signal
+
+    beam = Signal(name="beam", value=450.0) if legacy != "mode" else ring.beam_current
+    mode = (
+        Signal(name="mode", value=NSLS2OpsMode.OPERATIONS)
+        if legacy != "beam"
+        else ring.operating_mode
+    )
+
+    def set_beam(v):
+        beam.put(v) if isinstance(beam, Signal) else set_mock_value(beam, v)
+
+    def set_mode(m):
+        mode.put(m) if isinstance(mode, Signal) else set_mock_value(mode, m)
+
+    susp = SuspendFloorUnlessOpsMode(
+        beam, 100, resume_thresh=390, ops_mode=mode, skip_modes=SKIP
+    )
+    RE.install_suspender(susp)
+    try:
+        set_beam(50)
+        assert susp.tripped
+        set_mode(NSLS2OpsMode.MAINTENANCE)
+        assert not susp.tripped
+        set_mode(NSLS2OpsMode.OPERATIONS)
+        assert susp.tripped
+    finally:
+        RE.remove_suspender(susp)
+
+
+def test_beam_update_during_mode_change_is_not_lost(RE: RunEngine):
+    # PR 93 review: a beam reading arriving on another thread while a mode
+    # change re-evaluates must not be overwritten by the stale replayed value.
+    import threading
+
+    from ophyd import Signal
+
+    beam = Signal(name="beam", value=0.0)
+    mode = Signal(name="mode", value=NSLS2OpsMode.OPERATIONS)
+
+    class Racing(SuspendFloorUnlessOpsMode):
+        armed = False
+        reads = 0
+
+        @property
+        def _last_value(self):
+            value = self.__dict__.get("_lv")
+            if self.armed and threading.current_thread() is threading.main_thread():
+                self.reads += 1
+                if self.reads == 2:  # after the None check, before the replay
+                    self.armed = False
+                    other = threading.Thread(target=beam.put, args=(450.0,))
+                    other.start()
+                    other.join(timeout=0.5)
+                    self._racer = other
+            return value
+
+        @_last_value.setter
+        def _last_value(self, value):
+            self.__dict__["_lv"] = value
+
+    susp = Racing(beam, 100, resume_thresh=390, ops_mode=mode, skip_modes=SKIP)
+    RE.install_suspender(susp)
+    try:
+        beam.put(50.0)
+        assert susp.tripped
+        susp.armed = True
+        mode.put(NSLS2OpsMode.MAINTENANCE)
+        susp._racer.join()
+        mode.put(NSLS2OpsMode.OPERATIONS)
+        assert not susp.tripped, "judged against the stale beam current"
+    finally:
+        RE.remove_suspender(susp)