Repository: bluesky/bluesky (python)

Issue:
set_run_key_wrapper breaks if baseline is present
<!--- Provide a general summary of the issue in the Title above -->
If you have a baseline preprocessor, using `set_run_key_wrapper` on a plan will lead to an error, due to the baseline not being aware of the run key


## Current Behavior
```python
In [41]: RE.preprocessors
Out[41]: 
[SupplementalData(baseline=[Manipulator(prefix='', name='manipulator', settle_time=0.0, timeout=None, egu='', limits=(0, 0), source='computed', read_attrs=['sx', 'sx.readback', 'sx.setpoint', 'sy', 'sy.readback', 'sy.setpoint', 'sz', 'sz.readback', 'sz.setpoint', 'sr', 'sr.readback', 'sr.setpoint', 'x', 'y', 'z', 'r'], configuration_attrs=['sx', 'sy', 'sz', 'sr'], concurrent=True), SoftPositioner(name='eslit', settle_time=0.0, timeout=None, egu='', limits=(0, 0), source='computed'), SoftPositioner(name='i0upAu', settle_time=0.0, timeout=None, egu='', limits=(0, 0), source='computed'), SoftPositioner(name='tesz', settle_time=0.0, timeout=None, egu='', limits=(0, 0), source='computed')], monitors=[], flyers=[]),
 SupplementalData(baseline=[], monitors=[], flyers=[])]

In [42]: RE(set_run_key_wrapper(count([tes, ucal_i400], 5), "test_run_key"))


Transient Scan ID: 7     Time: 2022-06-30 12:04:46
Persistent Unique Scan ID: '4d35d081-e899-46ae-ae4d-520bd3561499'



An exception has occurred, use '%tb verbose' to see the full traceback.
KeyError: None

The above exception was the direct cause of the following exception:
IllegalMessageSequence: Cannot bundle readings without an open run. That is, 'create' must be preceded by 'open_run'.

See /home/ctitus/.cache/bluesky/log/bluesky.log for the full traceback.

In [43]: %tb verbose
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/miniconda3/envs/2022-2.3-py39-tiled/lib/python3.9/site-packages/bluesky/run_engine.py:1807, in RunEngine._create(self=<bluesky.run_engine.RunEngine object>, msg=Msg('create', obj=None, args=(), kwargs={'name': 'baseline'}, run=None))
   1806 try:
-> 1807     current_run = self._run_bundlers[run_key]
        run_key = None
        self._run_bundlers = {}
        self = <bluesky.run_engine.RunEngine object at 0x7f1fdae11eb0>
   1808 except KeyError as ke:

KeyError: None
```
## Possible Solution
<!--- Not obligatory, but suggest a fix/reason for the bug, -->
<!--- or ideas how to implement the addition or change -->
Somehow, preprocessors need to be made aware of the run key, or set_run_key_wrapper needs to have a different mechanism (the new run key gets set as the default?)

## Steps to Reproduce (for bugs)
<!--- Provide a link to a live example, or an unambiguous set of steps to -->
<!--- reproduce this bug. Include code to reproduce, if relevant -->
1. Register a baseline signal
2. Run a plan that has been wrapped with set_run_key_wrapper
3. Profit??

Pull request description:
FIX: support baseline when using run keys
closes #1529

<!--- Provide a general summary of your changes in the Title above -->

## Description
<!--- Describe your changes in detail -->

in baseline wrapper use the `msg.run` from the open/close for the inserted messages.


## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here. -->

The baseline wrapper was written before we adding the "run key" support for muliti run it was naive to it.  This means that the baseline wrapper always tries to put events into the "None" run (default) which is either not what you want or fails with invalid sequence orders.


## How Has This Been Tested?
<!--- Please describe in detail how you tested your changes. -->
<!--- Include details of your testing environment, and the tests you ran to -->
<!--- see how your change affects other areas of the code, etc. -->

Added test based on @lsammut 's reproducer.

<!--
## Screenshots (if appropriate):
-->


Patch:
diff --git a/src/bluesky/consolidators.py b/src/bluesky/consolidators.py
index 35880e7d..c12a8875 100644
--- a/src/bluesky/consolidators.py
+++ b/src/bluesky/consolidators.py
@@ -1,7 +1,6 @@
 import collections
 import dataclasses
 import enum
-import os
 from typing import Any, Optional
 
 import numpy as np
@@ -272,24 +271,21 @@ class HDF5Consolidator(ConsolidatorBase):
         return {"dataset": self._sres_parameters["dataset"].strip("/").split("/"), "swmr": self.swmr}
 
 
-class MultipartRelatedConsolidator(ConsolidatorBase):
-    def __init__(
-        self, permitted_extensions: set[str], stream_resource: StreamResource, descriptor: EventDescriptor
-    ):
+class TIFFConsolidator(ConsolidatorBase):
+    supported_mimetypes = {"multipart/related;type=image/tiff"}
+
+    def __init__(self, stream_resource: StreamResource, descriptor: EventDescriptor):
         super().__init__(stream_resource, descriptor)
-        self.permitted_extensions: set[str] = permitted_extensions
         self.data_uris: list[str] = []
-        self.template = self._sres_parameters["template"]
 
     def get_datum_uri(self, indx: int):
-        """Return a full uri for a datum (an individual image file) based on its index in the sequence.
+        """Return a full uri for a datum (an individual TIFF file) based on its index in the sequence.
 
         This relies on the `template` parameter passed in the StreamResource, which is a string either in the "new"
         Python formatting style that can be evaluated to a file name using the `.format(indx)` method given an
-        integer index, e.g. "{:05d}.ext".
+        integer index, e.g. "{:05d}.tif".
         """
-        assert os.path.splitext(self.template)[1] in self.permitted_extensions
-        return self.uri + self.template.format(indx)
+        return self.uri + self._sres_parameters["template"].format(indx)
 
     def consume_stream_datum(self, doc: StreamDatum):
         # Determine the indices in the names of tiff files from indices of frames and number of frames per file
@@ -309,24 +305,9 @@ class MultipartRelatedConsolidator(ConsolidatorBase):
         super().consume_stream_datum(doc)
 
 
-class TIFFConsolidator(MultipartRelatedConsolidator):
-    supported_mimetypes = {"multipart/related;type=image/tiff"}
-
-    def __init__(self, stream_resource: StreamResource, descriptor: EventDescriptor):
-        super().__init__({".tif", ".tiff"}, stream_resource, descriptor)
-
-
-class JPEGConsolidator(MultipartRelatedConsolidator):
-    supported_mimetypes = {"multipart/related;type=image/jpeg"}
-
-    def __init__(self, stream_resource: StreamResource, descriptor: EventDescriptor):
-        super().__init__({".jpeg", ".jpg"}, stream_resource, descriptor)
-
-
 CONSOLIDATOR_REGISTRY = {
     "application/x-hdf5": HDF5Consolidator,
     "multipart/related;type=image/tiff": TIFFConsolidator,
-    "multipart/related;type=image/jpeg": JPEGConsolidator,
 }
 
 
diff --git a/src/bluesky/preprocessors.py b/src/bluesky/preprocessors.py
index 524c5a54..4f46ef72 100644
--- a/src/bluesky/preprocessors.py
+++ b/src/bluesky/preprocessors.py
@@ -1204,12 +1204,17 @@ def baseline_wrapper(plan, devices, name="baseline"):
 
     def insert_baseline(msg):
         if msg.command == "open_run":
-            return None, head()
-
+            if msg.run is not None:
+                return None, set_run_key_wrapper(head(), msg.run)
+            else:
+                return None, head()
         elif msg.command == "close_run":
 
             def post_baseline():
-                yield from trigger_and_read(devices, name=name)
+                if msg.run is not None:
+                    yield from set_run_key_wrapper(trigger_and_read(devices, name=name), run=msg.run)
+                else:
+                    yield from trigger_and_read(devices, name=name)
                 return (yield msg)
 
             return post_baseline(), None


Test patch:
diff --git a/src/bluesky/tests/test_consolidators.py b/src/bluesky/tests/test_consolidators.py
index b73e6f0c..a78d1fab 100644
--- a/src/bluesky/tests/test_consolidators.py
+++ b/src/bluesky/tests/test_consolidators.py
@@ -2,7 +2,7 @@ from math import ceil
 
 import pytest
 
-from bluesky.consolidators import HDF5Consolidator, consolidator_factory
+from bluesky.consolidators import HDF5Consolidator, TIFFConsolidator
 
 
 @pytest.fixture
@@ -59,12 +59,12 @@ def hdf5_stream_resource_factory():
 
 
 @pytest.fixture
-def image_seq_stream_resource_factory():
-    return lambda image_format, data_key, chunk_shape: {
+def tiff_stream_resource_factory():
+    return lambda data_key, chunk_shape: {
         "data_key": data_key,
-        "mimetype": f"multipart/related;type=image/{image_format}",
+        "mimetype": "multipart/related;type=image/tiff",
         "uri": "file://localhost/test/file/path",
-        "parameters": {"chunk_shape": chunk_shape, "template": "img_{:06d}." + image_format},
+        "parameters": {"chunk_shape": chunk_shape, "template": "img_{:06d}.tiff"},
         "uid": f"stream-resource-uid-{data_key}",
     }
 
@@ -99,25 +99,13 @@ def test_hdf5_shape(descriptor, hdf5_stream_resource_factory, stream_datum_facto
     assert cons.shape == expected
 
 
-supported_image_seq_formats = ["jpeg", "tiff"]
-
-
 @pytest.mark.parametrize("data_key, expected", shape_testdata)
-@pytest.mark.parametrize("image_format", supported_image_seq_formats)
 @pytest.mark.parametrize("files_per_stream_datum", [1, 2, 3, 5])
 def test_tiff_shape(
-    descriptor,
-    image_seq_stream_resource_factory,
-    stream_datum_factory,
-    image_format,
-    data_key,
-    expected,
-    files_per_stream_datum,
+    descriptor, tiff_stream_resource_factory, stream_datum_factory, data_key, expected, files_per_stream_datum
 ):
-    stream_resource = image_seq_stream_resource_factory(
-        image_format=image_format, data_key=data_key, chunk_shape=(1,)
-    )
-    cons = consolidator_factory(stream_resource, descriptor)
+    stream_resource = tiff_stream_resource_factory(data_key=data_key, chunk_shape=(1,))
+    cons = TIFFConsolidator(stream_resource, descriptor)
     assert cons.shape == (0, *expected[1:])
     for i in range(ceil(5 / files_per_stream_datum)):
         doc = stream_datum_factory(
diff --git a/src/bluesky/tests/test_multi_runs.py b/src/bluesky/tests/test_multi_runs.py
index 4054fed2..ccc5d633 100644
--- a/src/bluesky/tests/test_multi_runs.py
+++ b/src/bluesky/tests/test_multi_runs.py
@@ -2,6 +2,7 @@ import pytest
 
 import bluesky.preprocessors as bsp
 from bluesky import plan_stubs as bps
+from bluesky import preprocessors as bpp
 from bluesky.preprocessors import set_run_key_wrapper as srkw
 from bluesky.tests.utils import DocCollector
 
@@ -138,3 +139,35 @@ def test_multirun_smoke_fail(RE, hw):
     for v in dc.stop.values():
         assert v["exit_status"] == "fail"
         assert v["reason"] == "womp womp"
+
+
+def test_multirun_baseline(RE, hw):
+    det1, det2, det3 = hw.det1, hw.det2, hw.det3
+
+    @bpp.set_run_key_decorator("run_2")
+    @bpp.run_decorator(md={})
+    def sim_plan_inner(npts):
+        for j in range(npts):
+            yield from bps.mov(hw.motor1, j * 0.1 + 1, hw.motor2, j * 0.2 - 2)
+            yield from bps.trigger_and_read([hw.motor1, hw.motor2, hw.det2])
+
+    @bpp.set_run_key_decorator("run_1")
+    @bpp.run_decorator(md={})
+    def sim_plan_outer(npts):
+        for j in range(int(npts / 2)):
+            yield from bps.mov(hw.motor, j * 0.2)
+            yield from bps.trigger_and_read([hw.motor, hw.det])
+
+        yield from sim_plan_inner(npts + 1)
+
+        for j in range(int(npts / 2), npts):
+            yield from bps.mov(hw.motor, j * 0.2)
+            yield from bps.trigger_and_read([hw.motor, hw.det])
+
+    # add baseline to RE
+    baseline = [det1, det2, det3]
+    sd = bpp.SupplementalData(baseline=baseline)
+    RE.preprocessors.append(sd)  # comment out this line to avoid error
+
+    # run the plan
+    RE(sim_plan_outer(10))