Repository: bluesky/bluesky
Pull request #1620 (2023-10-10 13:31:23)

Wait for async staging in stage_wrapper
## Description

Use stage_all and unstage_all to wait for staging to complete when stage returns a Status rather than blocking

## Motivation and Context

When we added async stage we didn't update all the plan preprocessors to `Msg('wait', group=g)` after a series of`Msg(`stage`, group=g)`. This is the cheapest change, although we will have to also fix `lazy_stage_wrapper` in another PR, and there may be other places emitting `Msg('stage')` directly rather than using `bps.stage()`. 

Is this correct? Or should we make the run engine wait for stage/unstage always unless it gets a group in the stage message?

## How Has This Been Tested?
Locally, no additional tests written yet


Patch:
diff --git a/bluesky/preprocessors.py b/bluesky/preprocessors.py
index 50a6c272..859bdcfd 100644
--- a/bluesky/preprocessors.py
+++ b/bluesky/preprocessors.py
@@ -11,7 +11,8 @@ 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)
+from .plan_stubs import (open_run, close_run, mv, pause, trigger_and_read,
+                         declare_stream, stage_all, unstage_all)
 
 
 def plan_mutator(plan, msg_proc):
@@ -917,12 +918,11 @@ def lazily_stage_wrapper(plan):
         else:
             return None, None
 
-    def unstage_all():
-        for device in reversed(devices_staged):
-            yield Msg('unstage', device)
+    def inner_unstage_all():
+        yield from unstage_all(*reversed(devices_staged))
 
     return (yield from finalize_wrapper(plan_mutator(plan, inner),
-                                        unstage_all()))
+                                        inner_unstage_all()))
 
 
 def stage_wrapper(plan, devices):
@@ -950,12 +950,10 @@ def stage_wrapper(plan, devices):
     devices = separate_devices(root_ancestor(device) for device in devices)
 
     def stage_devices():
-        for d in devices:
-            yield Msg('stage', d)
+        yield from stage_all(*devices)
 
     def unstage_devices():
-        for d in reversed(devices):
-            yield Msg('unstage', d)
+        yield from unstage_all(*reversed(devices))
 
     def inner():
         yield from stage_devices()


Test patch:
diff --git a/bluesky/tests/test_new_examples.py b/bluesky/tests/test_new_examples.py
index d2ec5322..e98232ce 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
+from typing import Dict, List
 import pytest
 from bluesky import Msg, RunEngineInterrupted
 from bluesky.plan_stubs import (
@@ -375,11 +375,20 @@ def test_lazily_stage(hw):
 
     processed_plan = list(lazily_stage_wrapper(plan()))
 
-    expected = [Msg('stage', det1), Msg('read', det1), Msg('read', det1),
-                Msg('stage', det2), Msg('read', det2), Msg('unstage', det2),
-                Msg('unstage', det1)]
+    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)
+    ]
 
-    assert processed_plan == expected
+    # 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
 
 
 def test_subs():