From 5eba4b16bfe6a1295b4f89a8bb1d8cb8c323e7b8 Mon Sep 17 00:00:00 2001 From: Tom Willemsen Date: Sat, 30 Nov 2024 17:21:14 +0000 Subject: [PATCH] Simplify/genericise --- doc/plan_stubs/matplotlib_helpers.md | 27 +++++----- manual_system_tests/dae_scan.py | 4 +- src/ibex_bluesky_core/plan_stubs/__init__.py | 47 +++++++++-------- src/ibex_bluesky_core/run_engine/__init__.py | 6 +-- .../run_engine/_msg_handlers.py | 50 ++++++++---------- tests/test_plan_stubs.py | 51 +++++++++---------- 6 files changed, 91 insertions(+), 94 deletions(-) diff --git a/doc/plan_stubs/matplotlib_helpers.md b/doc/plan_stubs/matplotlib_helpers.md index 3c6788a..74d66b1 100644 --- a/doc/plan_stubs/matplotlib_helpers.md +++ b/doc/plan_stubs/matplotlib_helpers.md @@ -1,7 +1,7 @@ -# `matplotlib` helpers +# `call_qt_aware` (matplotlib helpers) When attempting to use `matplotlib` UI functions directly in a plan, and running `matplotlib` using a `Qt` -backend (e.g. in a standalone shell outside IBEX), you may see an error of the form: +backend (e.g. in a standalone shell outside IBEX), you may see a hang or an error of the form: ``` UserWarning: Starting a Matplotlib GUI outside of the main thread will likely fail. @@ -11,34 +11,34 @@ UserWarning: Starting a Matplotlib GUI outside of the main thread will likely fa This is because the `RunEngine` runs plans in a worker thread, not in the main thread, which then requires special handling when calling functions that will update a UI. -The following plan stubs provide Qt-safe wrappers around some matplotlib functions to avoid this error. +The {py:obj}`ibex_bluesky_core.plan_stubs.call_qt_aware` plan stub can call `matplotlib` functions in a +Qt-aware context, which allows them to be run directly from a plan. It allows the same arguments and +keyword-arguments as the underlying matplotlib function it is passed. ```{note} Callbacks such as `LivePlot` and `LiveFitPlot` already route UI calls to the appropriate UI thread by default. The following plan stubs are only necessary if you need to call functions which will create or update a matplotlib -plot from a plan directly. +plot from a plan directly - for example to create or close a set of axes before passing them to callbacks. ``` -## `matplotlib_subplots` - -The {py:obj}`ibex_bluesky_core.plan_stubs.matplotlib_subplots` plan stub is a Qt-safe wrapper -around `matplotlib.pyplot.subplots()`. It allows the same arguments and keyword-arguments as the -underlying matplotlib function. - Usage example: ```python -from ibex_bluesky_core.plan_stubs import matplotlib_subplots +import matplotlib.pyplot as plt +from ibex_bluesky_core.plan_stubs import call_qt_aware from ibex_bluesky_core.callbacks.plotting import LivePlot from bluesky.callbacks import LiveFitPlot from bluesky.preprocessors import subs_decorator + def my_plan(): - # BAD + # BAD - likely to either crash or hang the plan. + # plt.close("all") # fig, ax = plt.subplots() # GOOD - fig, ax = yield from matplotlib_subplots() + yield from call_qt_aware(plt.close, "all") + fig, ax = yield from call_qt_aware(plt.subplots) # Pass the matplotlib ax object to other callbacks @subs_decorator([ @@ -48,4 +48,3 @@ def my_plan(): def inner_plan(): ... ``` - diff --git a/manual_system_tests/dae_scan.py b/manual_system_tests/dae_scan.py index c0675b2..a800a72 100644 --- a/manual_system_tests/dae_scan.py +++ b/manual_system_tests/dae_scan.py @@ -27,7 +27,7 @@ GoodFramesNormalizer, ) from ibex_bluesky_core.devices.simpledae.waiters import GoodFramesWaiter -from ibex_bluesky_core.plan_stubs import matplotlib_subplots +from ibex_bluesky_core.plan_stubs import call_qt_aware from ibex_bluesky_core.run_engine import get_run_engine NUM_POINTS: int = 3 @@ -72,7 +72,7 @@ def dae_scan_plan() -> Generator[Msg, None, None]: controller.run_number.set_name("run number") reducer.intensity.set_name("normalized counts") - _, ax = yield from matplotlib_subplots() + _, ax = yield from call_qt_aware(plt.subplots) lf = LiveFit( Linear.fit(), y=reducer.intensity.name, x=block.name, yerr=reducer.intensity_stddev.name diff --git a/src/ibex_bluesky_core/plan_stubs/__init__.py b/src/ibex_bluesky_core/plan_stubs/__init__.py index 06fefee..3e1b4fb 100644 --- a/src/ibex_bluesky_core/plan_stubs/__init__.py +++ b/src/ibex_bluesky_core/plan_stubs/__init__.py @@ -1,22 +1,20 @@ """Core plan stubs.""" from collections.abc import Generator -from typing import Any, Callable, ParamSpec, TypeVar, cast +from typing import Callable, ParamSpec, TypeVar, cast import bluesky.plan_stubs as bps -import matplotlib.pyplot as plt from bluesky.utils import Msg -from matplotlib.figure import Figure P = ParamSpec("P") T = TypeVar("T") CALL_SYNC_MSG_KEY = "ibex_bluesky_core_call_sync" -CALL_QT_SAFE_MSG_KEY = "ibex_bluesky_core_call_qt_safe" +CALL_QT_AWARE_MSG_KEY = "ibex_bluesky_core_call_qt_aware" -__all__ = ["call_sync", "matplotlib_subplots"] +__all__ = ["call_qt_aware", "call_sync"] def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Generator[Msg, None, T]: @@ -42,32 +40,41 @@ def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Genera Args: func: A callable to run. - args: Arbitrary arguments to be passed to the wrapped function - kwargs: Arbitrary keyword arguments to be passed to the wrapped function + *args: Arbitrary arguments to be passed to the wrapped function + **kwargs: Arbitrary keyword arguments to be passed to the wrapped function + + Returns: + The return value of the wrapped function """ yield from bps.clear_checkpoint() return cast(T, (yield Msg(CALL_SYNC_MSG_KEY, func, *args, **kwargs))) -def matplotlib_subplots( - *args: Any, # noqa: ANN401 - pyright doesn't understand we're wrapping mpl API. - **kwargs: Any, # noqa: ANN401 - pyright doesn't understand we're wrapping mpl API. -) -> Generator[Msg, None, tuple[Figure, Any]]: - """Create a new matplotlib figure and axes, using matplotlib.pyplot.subplots, from a plan. +def call_qt_aware( + func: Callable[P, T], *args: P.args, **kwargs: P.kwargs +) -> Generator[Msg, None, T]: + """Call a matplotlib function in a Qt-aware context, from within a plan. + + If matplotlib is using a Qt backend then UI operations are run on the Qt thread via Qt signals. - This is done in a Qt-safe way, such that if matplotlib is using a Qt backend then - UI operations are run on the Qt thread via Qt signals. + Only matplotlib functions may be run using this plan stub. Args: - args: Arbitrary arguments, passed through to matplotlib.pyplot.subplots - kwargs: Arbitrary keyword arguments, passed through to matplotlib.pyplot.subplots + func: A matplotlib function reference. + *args: Arbitrary arguments, passed through to matplotlib.pyplot.subplots + **kwargs: Arbitrary keyword arguments, passed through to matplotlib.pyplot.subplots + + Raises: + ValueError: if the passed function is not a matplotlib function. Returns: - tuple of (figure, axes) - as per matplotlib.pyplot.subplots() + The return value of the wrapped function """ + # Limit potential for misuse - constrain to just running matplotlib functions. + if not getattr(func, "__module__", "").startswith("matplotlib"): + raise ValueError("Only matplotlib functions should be passed to call_qt_aware") + yield from bps.clear_checkpoint() - return cast( - tuple[Figure, Any], (yield Msg(CALL_QT_SAFE_MSG_KEY, plt.subplots, *args, **kwargs)) - ) + return cast(T, (yield Msg(CALL_QT_AWARE_MSG_KEY, func, *args, **kwargs))) diff --git a/src/ibex_bluesky_core/run_engine/__init__.py b/src/ibex_bluesky_core/run_engine/__init__.py index f9d1a3c..19355bc 100644 --- a/src/ibex_bluesky_core/run_engine/__init__.py +++ b/src/ibex_bluesky_core/run_engine/__init__.py @@ -17,8 +17,8 @@ __all__ = ["get_run_engine"] -from ibex_bluesky_core.plan_stubs import CALL_QT_SAFE_MSG_KEY, CALL_SYNC_MSG_KEY -from ibex_bluesky_core.run_engine._msg_handlers import call_qt_safe_handler, call_sync_handler +from ibex_bluesky_core.plan_stubs import CALL_QT_AWARE_MSG_KEY, CALL_SYNC_MSG_KEY +from ibex_bluesky_core.run_engine._msg_handlers import call_qt_aware_handler, call_sync_handler logger = logging.getLogger(__name__) @@ -97,7 +97,7 @@ def get_run_engine() -> RunEngine: RE.subscribe(log_callback) RE.register_command(CALL_SYNC_MSG_KEY, call_sync_handler) - RE.register_command(CALL_QT_SAFE_MSG_KEY, call_qt_safe_handler) + RE.register_command(CALL_QT_AWARE_MSG_KEY, call_qt_aware_handler) RE.preprocessors.append(functools.partial(bpp.plan_mutator, msg_proc=add_rb_number_processor)) diff --git a/src/ibex_bluesky_core/run_engine/_msg_handlers.py b/src/ibex_bluesky_core/run_engine/_msg_handlers.py index 3f426bb..8801fc9 100644 --- a/src/ibex_bluesky_core/run_engine/_msg_handlers.py +++ b/src/ibex_bluesky_core/run_engine/_msg_handlers.py @@ -3,11 +3,10 @@ Not intended for user use. """ -import asyncio import ctypes import logging import threading -from asyncio import CancelledError, Event, get_running_loop, wait_for +from asyncio import CancelledError, Event, get_running_loop from typing import Any from bluesky.callbacks.mpl_plotting import QtAwareCallback @@ -17,9 +16,6 @@ logger = logging.getLogger(__name__) -CALL_QT_SAFE_TIMEOUT = 5 - - class _ExternalFunctionInterrupted(BaseException): """An external sync function running in a worker thread is being interrupted.""" @@ -34,7 +30,7 @@ async def call_sync_handler(msg: Msg) -> Any: # noqa: ANN401 def _wrapper() -> Any: # noqa: ANN401 nonlocal ret, exc - logger.info("Running '{func.__name__}' with args=({msg.args}), kwargs=({msg.kwargs})") + logger.info("Running '%s' with args=(%s), kwargs=(%s)", func.__name__, msg.args, msg.kwargs) try: ret = func(*msg.args, **msg.kwargs) logger.debug("Running '%s' successful", func.__name__) @@ -102,18 +98,8 @@ def _wrapper() -> Any: # noqa: ANN401 return ret -async def call_qt_safe_handler(msg: Msg) -> Any: # noqa: ANN401 - """Handle ibex_bluesky_core.plan_stubs.call_qt_safe. - - This functionality does not get exposed in a generic way to users, as this is - a tricky area. This *relies* on the passed function being "fast", but - we would have no way to prove that for a generic user-supplied function. So we only - expose known cases, like matplotlib.pyplot.subplots for example. - - Slow functions will cause various undesirable side effects, like stalling the event loop - and interfering with ctrl-c handling. - - """ +async def call_qt_aware_handler(msg: Msg) -> Any: # noqa: ANN401 + """Handle ibex_bluesky_core.plan_stubs.call_sync.""" func = msg.obj done_event = Event() result: Any = None @@ -126,28 +112,34 @@ async def call_qt_safe_handler(msg: Msg) -> Any: # noqa: ANN401 class _Cb(QtAwareCallback): def start(self, doc: RunStart) -> None: nonlocal result, exc - # Note: Qt/UI operations must be "fast", so don't worry too much about timeout or - # interruption cases here. - # Any attempt to forcibly interrupt a function while it's doing UI operations/using - # Qt signals is highly likely to be a bad idea. Don't do that here. try: + logger.info( + "Running '%s' with args=(%s), kwargs=(%s) (Qt)", + func.__name__, + msg.args, + msg.kwargs, + ) result = func(*msg.args, **msg.kwargs) + logger.debug("Running '%s' (Qt) successful", func.__name__) except BaseException as e: + logger.error( + "Running '%s' failed with %s: %s", func.__name__, e.__class__.__name__, e + ) exc = e finally: loop.call_soon_threadsafe(done_event.set) cb = _Cb() # Send fake event to our callback to trigger it (actual contents unimportant) + # If not using Qt, this will run synchronously i.e. block until complete + # If using Qt, this will be sent off to the Qt teleporter which will execute it asynchronously, + # and we have to wait for the event to be set. cb("start", {"time": 0, "uid": ""}) - try: - await wait_for(done_event.wait(), CALL_QT_SAFE_TIMEOUT) - except asyncio.TimeoutError as e: # pragma: no cover (CI on linux doesn't have a UI/Qt backend) - raise TimeoutError( - f"Long-running function '{func.__name__}' passed to call_qt_safe_handler. Functions " - f"passed to call_qt_safe_handler must be faster than {CALL_QT_SAFE_TIMEOUT}s." - ) from e + # Attempting to forcibly interrupt a function while it's doing UI operations/using + # Qt signals is highly likely to be a bad idea. Don't do that here. No special ctrl-c handling. + await done_event.wait() + if exc is not None: raise exc return result diff --git a/tests/test_plan_stubs.py b/tests/test_plan_stubs.py index d633001..c322772 100644 --- a/tests/test_plan_stubs.py +++ b/tests/test_plan_stubs.py @@ -1,14 +1,13 @@ # pyright: reportMissingParameterType=false -import re import time from asyncio import CancelledError -from unittest.mock import patch +from unittest.mock import MagicMock, patch -import matplotlib +import matplotlib.pyplot as plt import pytest from bluesky.utils import Msg -from ibex_bluesky_core.plan_stubs import CALL_QT_SAFE_MSG_KEY, call_sync, matplotlib_subplots +from ibex_bluesky_core.plan_stubs import CALL_QT_AWARE_MSG_KEY, call_qt_aware, call_sync from ibex_bluesky_core.run_engine._msg_handlers import call_sync_handler @@ -69,54 +68,54 @@ def f(): assert end - start == pytest.approx(1, abs=0.2) -def test_call_qt_safe_returns_result(RE): +def test_call_qt_aware_returns_result(RE): def f(arg, keyword_arg): assert arg == "foo" assert keyword_arg == "bar" return 123 def plan(): - return (yield Msg(CALL_QT_SAFE_MSG_KEY, f, "foo", keyword_arg="bar")) + return (yield Msg(CALL_QT_AWARE_MSG_KEY, f, "foo", keyword_arg="bar")) result = RE(plan()) assert result.plan_result == 123 -def test_call_qt_safe_throws_exception(RE): +def test_call_qt_aware_throws_exception(RE): def f(): raise ValueError("broke it") def plan(): - return (yield Msg(CALL_QT_SAFE_MSG_KEY, f)) + return (yield Msg(CALL_QT_AWARE_MSG_KEY, f)) with pytest.raises(ValueError, match="broke it"): RE(plan()) -@pytest.mark.skipif("qt" not in matplotlib.get_backend().lower(), reason="Qt not available") -def test_call_qt_safe_blocking_causes_descriptive_timeout_error(RE): - def f(): - time.sleep(0.2) +def test_call_qt_aware_matplotlib_function(RE): + mock = MagicMock(spec=plt.close) + mock.__module__ = "matplotlib.pyplot" + mock.return_value = 123 def plan(): - return (yield Msg(CALL_QT_SAFE_MSG_KEY, f)) + return (yield from call_qt_aware(mock, "all")) - with patch("ibex_bluesky_core.run_engine._msg_handlers.CALL_QT_SAFE_TIMEOUT", new=0.1): - with pytest.raises( - TimeoutError, - match=re.escape( - "Long-running function 'f' passed to call_qt_safe_handler. " - "Functions passed to call_qt_safe_handler must be faster than 0.1s." - ), - ): - RE(plan()) + result = RE(plan()) + assert result.plan_result == 123 + mock.assert_called_once_with("all") -def test_matplotlib_subplots_calls_pyplot_subplots(RE): +def test_call_qt_aware_non_matplotlib_function(RE): + mock = MagicMock() + mock.__module__ = "some_random_module" + def plan(): - return (yield from matplotlib_subplots("foo", keyword="bar")) + return (yield from call_qt_aware(mock, "arg", keyword_arg="kwarg")) - with patch("ibex_bluesky_core.plan_stubs.plt.subplots") as mock_pyplot_subplots: + with pytest.raises( + ValueError, match="Only matplotlib functions should be passed to call_qt_aware" + ): RE(plan()) - mock_pyplot_subplots.assert_called_once_with("foo", keyword="bar") + + mock.assert_not_called()