diff --git a/src/mx_bluesky/hyperion/baton_handler.py b/src/mx_bluesky/hyperion/baton_handler.py index 2014929df..4d393461c 100644 --- a/src/mx_bluesky/hyperion/baton_handler.py +++ b/src/mx_bluesky/hyperion/baton_handler.py @@ -6,6 +6,7 @@ from bluesky.utils import MsgGenerator, RunEngineInterrupted from dodal.common.beamlines.commissioning_mode import set_commissioning_signal from dodal.devices.baton import Baton +from dodal.devices.synchrotron import Synchrotron from mx_bluesky.common.external_interaction.alerting import ( AlertService, @@ -33,6 +34,7 @@ HYPERION_USER = "Hyperion" NO_USER = "None" +COUNTDOWN_THRESHOLD_SECONDS = 600 def run_forever(runner: PlanRunner): @@ -85,6 +87,7 @@ def collect() -> MsgGenerator: baton: The baton device runner: The runner """ + _raise_udc_start_alert(get_alerting_service()) yield from bpp.contingency_wrapper( runner.decode_and_execute( @@ -103,6 +106,9 @@ def collect() -> MsgGenerator: baton = _get_baton(context) current_visit: str | None = None while (yield from _is_requesting_baton(baton)): + abort = yield from _abort_if_countdown_too_low(context, baton) + if abort: + return current_visit = yield from _fetch_and_process_agamemnon_instruction( baton, runner, current_visit ) @@ -180,9 +186,7 @@ def _fetch_and_process_agamemnon_instruction( current_visit, parameter_list ) else: - _raise_udc_completed_alert(get_alerting_service()) - # Release the baton for orderly exit from the instruction loop - yield from _unrequest_baton(baton) + yield from _release_baton_on_completed_alert(baton) return current_visit @@ -218,6 +222,10 @@ def _get_baton(context: BlueskyContext) -> Baton: return find_device_in_context(context, "baton", Baton) +def _get_synchrotron(context: BlueskyContext) -> Synchrotron: + return find_device_in_context(context, "synchrotron", Synchrotron) + + def _unrequest_baton(baton: Baton) -> MsgGenerator[str]: """Relinquish the requested user of the baton if it is not already requested by another user. @@ -231,3 +239,24 @@ def _unrequest_baton(baton: Baton) -> MsgGenerator[str]: yield from bps.abs_set(baton.requested_user, NO_USER) return NO_USER return requested_user + + +def _abort_if_countdown_too_low( + context: BlueskyContext, baton: Baton +) -> MsgGenerator[bool]: + synchrotron = _get_synchrotron(context) + countdown = yield from bps.rd(synchrotron.machine_user_countdown) + + LOGGER.info(f"Synchrotron beam countdown is {countdown} seconds") + + if countdown < COUNTDOWN_THRESHOLD_SECONDS: + LOGGER.info("Synchrotron machine countdown too low") + yield from _release_baton_on_completed_alert(baton) + return True + + return False + + +def _release_baton_on_completed_alert(baton) -> MsgGenerator: + _raise_udc_completed_alert(get_alerting_service()) + yield from _unrequest_baton(baton) diff --git a/tests/unit_tests/hyperion/test_baton_handler.py b/tests/unit_tests/hyperion/test_baton_handler.py index 31520ea77..5d2d3d557 100644 --- a/tests/unit_tests/hyperion/test_baton_handler.py +++ b/tests/unit_tests/hyperion/test_baton_handler.py @@ -15,6 +15,7 @@ from bluesky.simulators import RunEngineSimulator, assert_message_and_return_remaining from dodal.devices.baton import Baton from dodal.devices.detector.detector_motion import DetectorMotion +from dodal.devices.synchrotron import Synchrotron from dodal.utils import get_beamline_based_on_environment_variable from ophyd_async.core import get_mock_put, set_mock_value @@ -110,6 +111,7 @@ def bluesky_context( lower_gonio, baton, detector_motion, + synchrotron, use_beamline_t01, ): # Baton for real run engine @@ -125,24 +127,26 @@ def mock_load_module(module, **kwargs): lower_gonio, baton, detector_motion, + synchrotron, ] for device in devices: context.register_device(device) return {d.name: d for d in devices}, {} - context.with_device_manager( - get_beamline_based_on_environment_variable().devices, - mock=True, - ) - - baton_with_requested_user(context, HYPERION_USER) with patch.object(context, "with_device_manager", mock_load_module): + context.with_device_manager( + get_beamline_based_on_environment_variable().devices, + mock=True, + ) + synchrotron_with_countdown(context) + baton_with_requested_user(context, HYPERION_USER) yield context @pytest.fixture def bluesky_context_with_sim_run_engine(sim_run_engine: RunEngineSimulator): baton_requested_user = HYPERION_USER + countdown = 1200 # Baton for sim run engine def get_requested_user(msg): @@ -153,12 +157,20 @@ def set_requested_user(msg): nonlocal baton_requested_user baton_requested_user = msg.args[0] + def machine_user_countdown_read(msg): + return {msg.obj.name: {"value": countdown}} + sim_run_engine.add_handler("locate", get_requested_user, "baton-requested_user") sim_run_engine.add_handler( "set", set_requested_user, # type: ignore "baton-requested_user", ) + sim_run_engine.add_handler( + "read", + machine_user_countdown_read, + "synchrotron-machine_user_countdown", + ) msgs = [] @@ -224,6 +236,14 @@ def baton_with_requested_user( return baton +def synchrotron_with_countdown( + bluesky_context: BlueskyContext, seconds: int = 1200 +) -> Synchrotron: + synchrotron = find_device_in_context(bluesky_context, "synchrotron", Synchrotron) + set_mock_value(synchrotron.machine_user_countdown, seconds) + return synchrotron + + @pytest.fixture() def udc_runner(bluesky_context: BlueskyContext) -> PlanRunner: runner = InProcessRunner(bluesky_context, True) @@ -1032,3 +1052,18 @@ def test_hyperion_doesnt_exit_if_udc_default_state_fails_a_check( mock_move_to_udc_default_state.assert_called_once() assert get_mock_put(baton.requested_user).mock_calls[-1] == call(NO_USER, wait=True) assert get_mock_put(baton.current_user).mock_calls[-1] == call(NO_USER, wait=True) + + +def test_baton_handler_fails_if_synchrotron_machine_countdown_below_threshold( + bluesky_context: BlueskyContext, + udc_runner: PlanRunner, + dont_patch_clear_devices, + caplog, +): + synchrotron = find_device_in_context(bluesky_context, "synchrotron", Synchrotron) + set_mock_value(synchrotron.machine_user_countdown, 5) + + with caplog.at_level("INFO"): + run_udc_when_requested(bluesky_context, udc_runner) + + assert "Synchrotron machine countdown too low" in caplog.text diff --git a/tests/unit_tests/hyperion/test_main_system.py b/tests/unit_tests/hyperion/test_main_system.py index 1a9988622..82fed037a 100644 --- a/tests/unit_tests/hyperion/test_main_system.py +++ b/tests/unit_tests/hyperion/test_main_system.py @@ -18,6 +18,7 @@ from blueapi.core import BlueskyContext from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator from dodal.devices.baton import Baton +from dodal.devices.synchrotron import Synchrotron from dodal.devices.zebra.zebra import Zebra from flask.testing import FlaskClient from ophyd_async.core import set_mock_value @@ -766,7 +767,9 @@ def wait_for_udc_to_start_then_send_sigterm(): plan_runner = mock_create_udc_server.mock_calls[0].args[0] context = plan_runner.context baton = find_device_in_context(context, "baton", Baton) + synchrotron = find_device_in_context(context, "synchrotron", Synchrotron) set_mock_value(baton.requested_user, HYPERION_USER) + set_mock_value(synchrotron.machine_user_countdown, 1200) while len(mock_create_parameters_from_agamemnon.mock_calls) == 0: sleep(0.2) os.kill(os.getpid(), signal.SIGTERM) diff --git a/tests/unit_tests/t01.py b/tests/unit_tests/t01.py index 875dae44f..45e2e122a 100644 --- a/tests/unit_tests/t01.py +++ b/tests/unit_tests/t01.py @@ -8,6 +8,7 @@ from dodal.device_manager import DeviceManager from dodal.devices.baton import Baton +from dodal.devices.synchrotron import Synchrotron from dodal.devices.xbpm_feedback import XBPMFeedback from dodal.utils import BeamlinePrefix, get_beamline_name @@ -25,6 +26,14 @@ def baton() -> Baton: return Baton(f"{PREFIX.beamline_prefix}-CS-BATON-01:") +@devices.factory() +def synchrotron() -> Synchrotron: + """Get the i03 Synchrotron device, instantiate it if it hasn't already been. + If this is called when already instantiated in i03, it will return the existing object. + """ + return Synchrotron() + + @devices.factory() def xbpm_feedback() -> XBPMFeedback: """Get the i03 XBPM feeback device, instantiate it if it hasn't already been.