ENH: add helpers to manage "batches" of runs

[tacaswell]

This is the result of a day if discussion on NSLS-II slack.

[tacaswell]

from bluesky import RunEngine
from bluesky.callbacks.core import LiveTable
from bluesky.plan_stubs import open_run, close_run, null
from event_model import RunRouter

from nslsii.batches import setup_batch


def inner_plan(M):
    for j in range(M):
        yield from open_run()
        yield from close_run()


def batch(batch_md, *, N=5, M=1, comment_function=None):
    add_to_batch, close_batch = yield from setup_batch(
        batch_md, comment_function=comment_function
    )
    for j in range(N):
        yield from add_to_batch(inner_plan(M=M))
    yield from close_batch()


RE = RunEngine()


def only_outer(name, doc):
    if doc.get("purpose") == "batch header":
        return [
            LiveTable(["step_uid", "step_comment", "step_index"], default_prec=25)
        ], []
    return [], []


rr = RunRouter([only_outer])

RE(batch({"bob": "ardvark"}, N=5, M=2), rr)

gives

+-----------+------------+----------------------------+----------------------------+----------------------------+
|   seq_num |       time |                   step_uid |               step_comment |                 step_index |
+-----------+------------+----------------------------+----------------------------+----------------------------+
|         1 | 22:55:41.6 |  f27e3969-0ecc-44ce-85d4-f |                     step 0 |                          0 |
|         2 | 22:55:41.6 |  e2e7ee96-7fa8-44aa-a9dc-e |                     step 1 |                          1 |
|         3 | 22:55:41.6 |  48d93ad2-f4d4-4624-aa76-0 |                     step 2 |                          2 |
|         4 | 22:55:41.6 |  2817bb29-e338-4f11-8d24-3 |                     step 3 |                          3 |
|         5 | 22:55:41.6 |  cd04e05e-d6ea-4b60-8849-0 |                     step 4 |                          4 |
|         6 | 22:55:41.6 |  e2ab4ac3-d021-4610-9fe2-e |                     step 5 |                          5 |
|         7 | 22:55:41.6 |  3fe9be55-7b99-4a04-a671-f |                     step 6 |                          6 |
|         8 | 22:55:41.6 |  c00f51a0-ed42-4a29-a187-5 |                     step 7 |                          7 |
|         9 | 22:55:41.6 |  53f0c06f-3eb2-49ea-916f-2 |                     step 8 |                          8 |
|        10 | 22:55:41.6 |  b32c6d79-183e-48eb-8a95-1 |                     step 9 |                          9 |
+-----------+------------+----------------------------+----------------------------+----------------------------+
generator batch ['8ca0d13c'] (scan num: 1)

[cjtitus]

Running with a supplemental data preprocessor (just motors in baseline) gives the following error

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 0x7f9c5363fe50>
   1808 except KeyError as ke:

KeyError: None

Which is, I assume, because the run key needs to be passed to baseline. Not sure how to do this.

[cjtitus]

This is really a problem with set_run_key_wrapper itself, which I have created an issue for bluesky/bluesky#1529

[cjtitus]

Actually, there will be a problem with Monitors and Flyers as well, which are going to be kicked off twice when you nest runs like this.

For this batch scheme to work, I think the easiest thing to do would be to somehow suppress the Baseline/Monitor/Flyer preprocessors, if that is possible.

[tacaswell]

For this batch scheme to work, I think the easiest thing to do would be to somehow suppress the Baseline/Monitor/Flyer preprocessors, if that is possible.

The logic to turn things off like that is hard to get right. If you are going down this route I would remove the pre-processor on the RE and directly decorate / wrap your inner plans.

[tacaswell]

The pre-processors only get to see a stream of Msg and act on them the ordering here is

baseline(set_run_key(plan()))

The baseline wrapper does not know (naively) what set_run_key is doing, it just sees a stream of Msg it is reacting to. All of those wrappers pre-date (by a lot) the concept of run_key . I think that it should be possible for them to get a bit smarter and look at the run-key on the start message and correctly propogate it to the extra messages the wrappers generate.

This would fix baseline, but would still leave you with the double monitor_during / fly_during issue.

[cjtitus]

I'm not actually using any monitors or flyers, so if the wrappers could be updated to propagate the run-key, that would be a work-able solution for me. I opened an issue on the Bluesky github, so anyone is welcome to fix it. Maybe I'll do it and submit a PR. It should be fixed regardless of the Batch stuff.

For the double monitor/fly issue, this would mean foregoing the SupplementalData preprocessor altogether, and directly wrapping all plans in monitor/fly. This is probably not too much of an issue for me, but might be confusing.

I think the only way to do it more cleanly is to find some way for the preprocessors to ignore the batch open_run/close_run. I thought of two ways.

1. Add some kind of special "ignore me" flag to the batch open_run/close_run Msg that the Baseline/Flyer/Monitor wrappers look for (and ignore)
2. Create a new open_batch/close_batch Msg that does basically the same thing as open_run/close_run, but will of course be ignored by all current wrappers.

This second way is something I could actually test relatively rapidly. To first order, I could just register an open_batch command that literally just calls open_run, but wouldn't trigger the preprocessors.

[tacaswell]

For the double monitor/fly issue, this would mean foregoing the SupplementalData preprocessor altogether ...

Yes, however if you are going down the route of these batches things are already pretty custom. You can still close over the SD object in your plans, you just can not hang it off of the RE.

Add some kind of special "ignore me" flag to the batch open_run/close_run Msg that the Baseline/Flyer/Monitor wrappers look for (and ignore)

Baking things into the contents of the message to opt-opt of an optional pre-processor is a code-smell. That said, a local thing you could try is to wrap the SD run-router in another layer that does the filtering. It is a bit ugly, but could work, and it being in application code rather than the library code makes it seem a bit better.

Create a new open_batch/close_batch Msg that does basically the same thing as open_run/close_run, but will of course be ignored by all current wrappers.

You really want another generation of documents in the event model ;) I do not like this option, but it will work and is consistent with the rest of the system.

This second way is something I could actually test relatively rapidly. To first order, I could just register an open_batch command that literally just calls open_run, but wouldn't trigger the preprocessors.

Yes, but you would need to do both open and close I think or the wrappers might become unhappy about seeing a close with no open.