Repository: bluesky/bluesky
Pull request #1623 (2023-10-20 09:30:59)

Remove group from groups after Msg('wait_for'), not before, in RunEngine._wait
RunEngine._wait now waits until wait_for message has been produced before removing futures from groups/status objects


tacaswell: I think that if the wait for raises we will not clear the state?
rosesyrett: I would like to add a test to this; however, that is currently non-trivial, because it would require adding a dependency on ophyd-async, because the best way of testing this is doing something along the lines of:

```python
def test_wait_propagates_error():
    from ophyd_async.core import Device

    class DummyDevice(Device):
        async def set(self, value:int):
            await asyncio.sleep(0.01)
            raise RuntimeError()
    dummy = DummyDevice("somename")
    RE(abs_set(dummy, 1))
```

This way, when abs set immediately sets it, there is no exception, so that line passes. But when we get into the run_engine._wait, the error is caught and reported on (instead of silently failing, which is what we found was happening...)

So the question is, do we want to rely on ophyd-async?
coretl: How about we make it so that we pop unless we get a `TimeoutError`?
tacaswell: There is a difference if the wait times out or if the Status object times out and fail itself.
coretl: Ok, how about this:

- We make a new `class WaitForTimeoutError(TimeoutError):`
- In: https://github.com/bluesky/bluesky/blob/44579d5201280bcfee20574a72fa063bf7d0160a/bluesky/run_engine.py#L1741C25-L1741C25 we raise this
- In https://github.com/bluesky/bluesky/blob/44579d5201280bcfee20574a72fa063bf7d0160a/bluesky/run_engine.py#L2205 the only time we don't pop is on `WaitForTimeoutError`

tacaswell: I miss-remembered how the raising for failed status worked https://github.com/bluesky/bluesky/blob/44579d5201280bcfee20574a72fa063bf7d0160a/bluesky/run_engine.py#L2214-L2235

so the sad-path on the status object failing itself is the happy-path here so that is not a problem.

However, a second thing I thought of here is `CancelError` s being through in by the event loop (for example if the user pauses while waiting).  

Can we go with the `TimeoutError` subclass and then if we catch that put the futures back?  You could argue that on pause is some other part of the code's job to clean up all of the stale status objects, but I think it is clearer if `_wait` cleans up after itself and "re-arms" the group on a plan-side timeout.
coretl: > Can we go with the TimeoutError subclass and then if we catch that put the futures back? 

Sounds good, I'll update the PR

Patch:
diff --git a/bluesky/bundlers.py b/bluesky/bundlers.py
index f7ae5db9..dd6100a4 100644
--- a/bluesky/bundlers.py
+++ b/bluesky/bundlers.py
@@ -1005,6 +1005,10 @@ class RunBundler:
                     "A `collect` message on a device that isn't EventCollectable or EventPageCollectable "
                     "requires a `name=stream_name` argument"
                 )
+            if not indices_difference:
+                raise RuntimeError(
+                    "A `collect` message returned no events/event_pages or stream_datums"
+                )
 
             # Since there are no events or event_pages incrementing the sequence counter, we do it ourselves.
             self._sequence_counters[message_stream_name] += indices_difference
diff --git a/bluesky/preprocessors.py b/bluesky/preprocessors.py
index 859bdcfd..50a6c272 100644
--- a/bluesky/preprocessors.py
+++ b/bluesky/preprocessors.py
@@ -11,8 +11,7 @@ from .utils import (get_hinted_fields, normalize_subs_input, root_ancestor,
                     short_uid as _short_uid, make_decorator,
                     RunEngineControlException, merge_axis)
 from functools import wraps
-from .plan_stubs import (open_run, close_run, mv, pause, trigger_and_read,
-                         declare_stream, stage_all, unstage_all)
+from .plan_stubs import (open_run, close_run, mv, pause, trigger_and_read, declare_stream)
 
 
 def plan_mutator(plan, msg_proc):
@@ -918,11 +917,12 @@ def lazily_stage_wrapper(plan):
         else:
             return None, None
 
-    def inner_unstage_all():
-        yield from unstage_all(*reversed(devices_staged))
+    def unstage_all():
+        for device in reversed(devices_staged):
+            yield Msg('unstage', device)
 
     return (yield from finalize_wrapper(plan_mutator(plan, inner),
-                                        inner_unstage_all()))
+                                        unstage_all()))
 
 
 def stage_wrapper(plan, devices):
@@ -950,10 +950,12 @@ def stage_wrapper(plan, devices):
     devices = separate_devices(root_ancestor(device) for device in devices)
 
     def stage_devices():
-        yield from stage_all(*devices)
+        for d in devices:
+            yield Msg('stage', d)
 
     def unstage_devices():
-        yield from unstage_all(*reversed(devices))
+        for d in reversed(devices):
+            yield Msg('unstage', d)
 
     def inner():
         yield from stage_devices()
diff --git a/bluesky/run_engine.py b/bluesky/run_engine.py
index 0a1e89a8..1ef5f568 100644
--- a/bluesky/run_engine.py
+++ b/bluesky/run_engine.py
@@ -46,6 +46,10 @@ class _RunEnginePanic(Exception):
     ...
 
 
+class WaitForTimeoutError(TimeoutError):
+    ...
+
+
 @dataclass
 class RunEngineResult():
     """
@@ -1738,7 +1742,7 @@ class RunEngine:
         futs = [asyncio.ensure_future(f()) for f in futs]
         _, pending = await asyncio.wait(futs, **self._loop_for_kwargs, **msg.kwargs)
         if pending:
-            raise TimeoutError("Plan failed to complete in the specified time")
+            raise WaitForTimeoutError("Plan failed to complete in the specified time")
 
     async def _open_run(self, msg):
         """Instruct the RunEngine to start a new "run"
@@ -2203,6 +2207,11 @@ class RunEngine:
                     # bar.
                     self.waiting_hook(status_objs)
                 await self._wait_for(Msg('wait_for', None, futs, timeout=msg.kwargs.get("timeout", None)))
+            except WaitForTimeoutError:
+                # We might wait to call wait again, so put the futures and status objects back in
+                self._groups[group] = futs
+                self._status_objs[group] = status_objs
+                raise
             finally:
                 if self.waiting_hook is not None:
                     # Notify the waiting_hook function that we have moved on by


Test patch:
diff --git a/bluesky/tests/test_new_examples.py b/bluesky/tests/test_new_examples.py
index e98232ce..d2ec5322 100644
--- a/bluesky/tests/test_new_examples.py
+++ b/bluesky/tests/test_new_examples.py
@@ -1,7 +1,7 @@
 import asyncio
 from collections import defaultdict
 import time as ttime
-from typing import Dict, List
+from typing import Dict
 import pytest
 from bluesky import Msg, RunEngineInterrupted
 from bluesky.plan_stubs import (
@@ -375,20 +375,11 @@ def test_lazily_stage(hw):
 
     processed_plan = list(lazily_stage_wrapper(plan()))
 
-    expected_plan: List[Msg] = [
-        Msg('stage', det1), Msg('read', det1), Msg('read', det1),
-        Msg('stage', det2), Msg('read', det2), Msg('unstage', det2),
-        Msg('unstage', det1)
-    ]
+    expected = [Msg('stage', det1), Msg('read', det1), Msg('read', det1),
+                Msg('stage', det2), Msg('read', det2), Msg('unstage', det2),
+                Msg('unstage', det1)]
 
-    # Prevent issue with unstage_all creating a randomly assigned group
-    assert len(processed_plan) == len(expected_plan)
-    for i in range(len(expected_plan)):
-        expected, actual = expected_plan[i], processed_plan[i]
-        assert actual.command == expected.command
-        assert actual.obj == expected.obj
-    groups = {msg.kwargs["group"] for msg in processed_plan if msg.command == "unstage"}
-    assert len(groups) == 1  # All unstage messages are in the same group
+    assert processed_plan == expected
 
 
 def test_subs():
diff --git a/bluesky/tests/test_run_engine.py b/bluesky/tests/test_run_engine.py
index 6c6da572..b88357a2 100644
--- a/bluesky/tests/test_run_engine.py
+++ b/bluesky/tests/test_run_engine.py
@@ -1399,6 +1399,30 @@ def test_exceptions_kill_run_and_failed_status_holds_detail(RE):
         assert isinstance(exc.args[0], UnknownStatusFailure)
 
 
+@requires_ophyd
+def test_wait_with_timeout_can_be_repeated(RE):
+    from ophyd import StatusBase
+
+    class Dummy:
+        name = 'dummy'
+
+        def set(self, val):
+            self.st = StatusBase()
+            return self.st
+
+    def plan():
+        dummy = Dummy()
+        yield Msg('set', dummy, 1, group='test')
+        with pytest.raises(TimeoutError):
+            yield Msg('wait', group='test', timeout=0.1)
+        with pytest.raises(TimeoutError):
+            yield Msg('wait', group='test', timeout=0.1)
+        dummy.st.set_finished()
+        yield Msg('wait', group='test')
+
+    RE(plan())
+
+
 @requires_ophyd
 def test_status_propagates_exception_through_run_engine(RE):
     # just to make sure that 'pardon_failures' does not block *real* failures