From cd12b67480e98d4aa8214fcacfa7876f050b87c2 Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Mon, 5 Feb 2024 14:54:09 +0000 Subject: [PATCH 1/2] (#1117) Retry writing to the zebra if it fails --- .../device_setup_plans/setup_zebra.py | 26 +++++++++++++ .../device_setup_plans/test_zebra_setup.py | 38 ++++++++++++++++++- 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/src/hyperion/device_setup_plans/setup_zebra.py b/src/hyperion/device_setup_plans/setup_zebra.py index fc9e4e7bb..31974efb2 100644 --- a/src/hyperion/device_setup_plans/setup_zebra.py +++ b/src/hyperion/device_setup_plans/setup_zebra.py @@ -1,4 +1,7 @@ +from typing import Callable + import bluesky.plan_stubs as bps +import bluesky.preprocessors as bpp from dodal.devices.zebra import ( DISCONNECT, IN1_TTL, @@ -20,6 +23,24 @@ from hyperion.log import LOGGER +def bluesky_retry(): + def decorator(func: Callable): + def newfn(*args, **kwargs): + def log_and_retry(exception): + LOGGER.error( + f"Function {func.__name__} failed with {exception}, retrying" + ) + yield from func(*args, **kwargs) + + yield from bpp.contingency_wrapper( + func(*args, **kwargs), except_plan=log_and_retry, auto_raise=False + ) + + return newfn + + return decorator + + def arm_zebra(zebra: Zebra): yield from bps.abs_set(zebra.pc.arm, ArmDemand.ARM, wait=True) @@ -28,6 +49,7 @@ def disarm_zebra(zebra: Zebra): yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, wait=True) +@bluesky_retry() def setup_zebra_for_rotation( zebra: Zebra, axis: I03Axes = I03Axes.OMEGA, @@ -93,6 +115,7 @@ def setup_zebra_for_rotation( yield from bps.wait(group) +@bluesky_retry() def setup_zebra_for_gridscan( zebra: Zebra, group="setup_zebra_for_gridscan", wait=False ): @@ -105,6 +128,7 @@ def setup_zebra_for_gridscan( yield from bps.wait(group) +@bluesky_retry() def set_zebra_shutter_to_manual( zebra: Zebra, group="set_zebra_shutter_to_manual", wait=False ): @@ -115,10 +139,12 @@ def set_zebra_shutter_to_manual( yield from bps.wait(group) +@bluesky_retry() def make_trigger_safe(zebra: Zebra, group="make_zebra_safe", wait=False): yield from bps.abs_set(zebra.inputs.soft_in_1, 0, wait=wait, group=group) +@bluesky_retry() def setup_zebra_for_panda_flyscan( zebra: Zebra, group="setup_zebra_for_panda_flyscan", wait=False ): diff --git a/tests/unit_tests/device_setup_plans/test_zebra_setup.py b/tests/unit_tests/device_setup_plans/test_zebra_setup.py index b69bea212..e8f407b57 100644 --- a/tests/unit_tests/device_setup_plans/test_zebra_setup.py +++ b/tests/unit_tests/device_setup_plans/test_zebra_setup.py @@ -1,7 +1,8 @@ from functools import partial -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest +from bluesky import plan_stubs as bps from dodal.beamlines import i03 from dodal.devices.zebra import ( IN3_TTL, @@ -17,6 +18,7 @@ from hyperion.device_setup_plans.setup_zebra import ( arm_zebra, + bluesky_retry, disarm_zebra, set_zebra_shutter_to_manual, setup_zebra_for_gridscan, @@ -81,3 +83,37 @@ def side_effect(set_armed_to: int, _): with pytest.raises(Exception): zebra.pc.arm.armed.set(1) RE(disarm_zebra(zebra, 0.2)) + + +class MyException(Exception): + pass + + +def test_when_first_try_fails_then_bluesky_retry_tries_again(RE, done_status): + mock_device = MagicMock() + + @bluesky_retry() + def my_plan(value): + yield from bps.abs_set(mock_device, value) + + mock_device.set.side_effect = [MyException(), done_status] + + RE(my_plan(10)) + + assert mock_device.set.mock_calls == [call(10), call(10)] + + +def test_when_all_tries_fail_then_bluesky_retry_throws_error(RE, done_status): + mock_device = MagicMock() + + @bluesky_retry() + def my_plan(value): + yield from bps.abs_set(mock_device, value) + + exception_2 = MyException() + mock_device.set.side_effect = [MyException(), exception_2] + + with pytest.raises(MyException) as e: + RE(my_plan(10)) + + assert e.value == exception_2 From e0a1776ee882618bf8dea1ced1f2c67752d15978 Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Tue, 6 Feb 2024 11:14:56 +0000 Subject: [PATCH 2/2] (#1117) Tidy up bluesky retry decorator --- .../device_setup_plans/setup_zebra.py | 44 +++++++++++-------- .../device_setup_plans/test_zebra_setup.py | 4 +- 2 files changed, 28 insertions(+), 20 deletions(-) diff --git a/src/hyperion/device_setup_plans/setup_zebra.py b/src/hyperion/device_setup_plans/setup_zebra.py index 31974efb2..add7a5328 100644 --- a/src/hyperion/device_setup_plans/setup_zebra.py +++ b/src/hyperion/device_setup_plans/setup_zebra.py @@ -1,3 +1,4 @@ +from functools import wraps from typing import Callable import bluesky.plan_stubs as bps @@ -23,22 +24,29 @@ from hyperion.log import LOGGER -def bluesky_retry(): - def decorator(func: Callable): - def newfn(*args, **kwargs): - def log_and_retry(exception): - LOGGER.error( - f"Function {func.__name__} failed with {exception}, retrying" - ) - yield from func(*args, **kwargs) +def bluesky_retry(func: Callable): + """Decorator that will retry the decorated plan if it fails. - yield from bpp.contingency_wrapper( - func(*args, **kwargs), except_plan=log_and_retry, auto_raise=False - ) + Use this with care as it knows nothing about the state of the world when things fail. + If it is possible that your plan fails when the beamline is in a transient state that + the plan could not act on do not use this decorator without doing some more intelligent + clean up. - return newfn + You should avoid using this decorator often in general production as it hides errors, + instead it should be used only for debugging these underlying errors. + """ + + @wraps(func) + def newfunc(*args, **kwargs): + def log_and_retry(exception): + LOGGER.error(f"Function {func.__name__} failed with {exception}, retrying") + yield from func(*args, **kwargs) + + yield from bpp.contingency_wrapper( + func(*args, **kwargs), except_plan=log_and_retry, auto_raise=False + ) - return decorator + return newfunc def arm_zebra(zebra: Zebra): @@ -49,7 +57,7 @@ def disarm_zebra(zebra: Zebra): yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, wait=True) -@bluesky_retry() +@bluesky_retry def setup_zebra_for_rotation( zebra: Zebra, axis: I03Axes = I03Axes.OMEGA, @@ -115,7 +123,7 @@ def setup_zebra_for_rotation( yield from bps.wait(group) -@bluesky_retry() +@bluesky_retry def setup_zebra_for_gridscan( zebra: Zebra, group="setup_zebra_for_gridscan", wait=False ): @@ -128,7 +136,7 @@ def setup_zebra_for_gridscan( yield from bps.wait(group) -@bluesky_retry() +@bluesky_retry def set_zebra_shutter_to_manual( zebra: Zebra, group="set_zebra_shutter_to_manual", wait=False ): @@ -139,12 +147,12 @@ def set_zebra_shutter_to_manual( yield from bps.wait(group) -@bluesky_retry() +@bluesky_retry def make_trigger_safe(zebra: Zebra, group="make_zebra_safe", wait=False): yield from bps.abs_set(zebra.inputs.soft_in_1, 0, wait=wait, group=group) -@bluesky_retry() +@bluesky_retry def setup_zebra_for_panda_flyscan( zebra: Zebra, group="setup_zebra_for_panda_flyscan", wait=False ): diff --git a/tests/unit_tests/device_setup_plans/test_zebra_setup.py b/tests/unit_tests/device_setup_plans/test_zebra_setup.py index e8f407b57..beffb083e 100644 --- a/tests/unit_tests/device_setup_plans/test_zebra_setup.py +++ b/tests/unit_tests/device_setup_plans/test_zebra_setup.py @@ -92,7 +92,7 @@ class MyException(Exception): def test_when_first_try_fails_then_bluesky_retry_tries_again(RE, done_status): mock_device = MagicMock() - @bluesky_retry() + @bluesky_retry def my_plan(value): yield from bps.abs_set(mock_device, value) @@ -106,7 +106,7 @@ def my_plan(value): def test_when_all_tries_fail_then_bluesky_retry_throws_error(RE, done_status): mock_device = MagicMock() - @bluesky_retry() + @bluesky_retry def my_plan(value): yield from bps.abs_set(mock_device, value)