Trigger Zocalo and turn on feedback after all collections in MSP

docs/developer/general/explanations/callback_and_run_logic.rst

@@ -0,0 +1,23 @@
+Callbacks are used to trigger external services:
+* Ispyb deposition
+* Nexus writing
+* Zocalo triggering
+
+These are linked in that to trigger zocalo you need to have made an ispyb deposition, written a nexus file and have finished writing raw data to disk. Nexus files and ispyb depositions can be made at anytime, we do not need to have necessarily finished writing raw data. 
+
+Currently, the requirement of needing to have written to ispyb is explicit as the ispyb callback will emit to the zocalo callback. The nexus file is written when the hardware is read during a collection and so its ordering is implied. When instantiated the zocalo callback is told on which plan to trigger and it is up to the plan developer to make sure this plan finishes after data is written to the detector.
+
+In general, the ordering flow of when callbacks are triggered is controlled by emitting documents with the expected plan name and data.
+
+Rotation Scans
+==============
+
+There are currently two ways of doing rotation scans. A single scan creates one hdf file, one ispyb deposition and then triggers zocalo once. Multi rotation scans create one hdf file for all rotations but then N nexus files, N ispyb depositions and triggers zocalo N times.
+
+Single scans will be removed in https://github.com/DiamondLightSource/mx-bluesky/issues/847
+
+For multi rotations this is does by starting 1+2*N different runs:
+
+1. ``CONST.PLAN.ROTATION_MULTI``: This is emitted once for the whole multiple rotation. It is used by the nexus callback to get the full number of images and meta_data_run_number so that it knows which hdf file to use. When this is finished zocalo end is triggered.
+2. ``CONST.PLAN.ROTATION_OUTER``: Emitted N times, inside a ``CONST.PLAN.ROTATION_MULTI`` run. This is used to create the initial ispyb deposition and create the nexus writer (but not actually write the file)
+3. ``CONST.PLAN.ROTATION_MAIN``: Emitted N times, inside ``CONST.PLAN.ROTATION_OUTER`` run. Used to finish writing to ispyb (i.e. write success/failure) and to send collection information to zocalo.

docs/developer/general/index.rst

@@ -41,6 +41,7 @@ Documentation is split into four categories, and each is also accessible from li
             :maxdepth: 1
 
             explanations/decisions
+            explanations/callback_and_run_logic
 
         +++
 

src/mx_bluesky/common/external_interaction/callbacks/common/zocalo_callback.py

@@ -18,10 +18,11 @@
 
 class ZocaloCallback(CallbackBase):
     """Callback class to handle the triggering of Zocalo processing.
-    Sends zocalo a run_start signal on receiving a start document for the specified
-    sub-plan, and sends a run_end signal on receiving a stop document for the same plan.
+    Will start listening for collections when {triggering_plan} has been started.
 
-    The metadata of the sub-plan this starts on must include a zocalo_environment.
+    For every ispyb deposition that occurs inside this run the callback will send zocalo
+    a run_start signal. Once the {triggering_plan} has ended the callback will send a
+    run_end signal for all collections.
 
     Shouldn't be subscribed directly to the RunEngine, instead should be passed to the
     `emit` argument of an ISPyB callback which appends DCIDs to the relevant start doc.
@@ -30,7 +31,9 @@ class ZocaloCallback(CallbackBase):
     def _reset_state(self):
         self.run_uid: str | None = None
         self.zocalo_info: list[ZocaloStartInfo] = []
+        self._started_zocalo_collections: list[ZocaloStartInfo] = []
         self.descriptors: dict[str, EventDescriptor] = {}
+        self.start_frame = 0
 
     def __init__(self, triggering_plan: str, zocalo_environment: str):
         super().__init__()
@@ -42,26 +45,21 @@ def start(self, doc: RunStart):
         ISPYB_ZOCALO_CALLBACK_LOGGER.info("Zocalo handler received start document.")
         if self.triggering_plan and doc.get("subplan_name") == self.triggering_plan:
             self.run_uid = doc.get("uid")
-            assert isinstance(scan_points := doc.get("scan_points"), list)
+        if self.run_uid:
             if (
-                isinstance(ispyb_ids := doc.get("ispyb_dcids"), tuple)
+                isinstance(scan_points := doc.get("scan_points"), list)
+                and isinstance(ispyb_ids := doc.get("ispyb_dcids"), tuple)
                 and len(ispyb_ids) > 0
             ):
                 ISPYB_ZOCALO_CALLBACK_LOGGER.info(f"Zocalo triggering for {ispyb_ids}")
                 ids_and_shape = list(zip(ispyb_ids, scan_points, strict=False))
-                start_frame = 0
-                self.zocalo_info = []
                 for idx, id_and_shape in enumerate(ids_and_shape):
                     id, shape = id_and_shape
                     num_frames = number_of_frames_from_scan_spec(shape)
                     self.zocalo_info.append(
-                        ZocaloStartInfo(id, None, start_frame, num_frames, idx)
+                        ZocaloStartInfo(id, None, self.start_frame, num_frames, idx)
                     )
-                    start_frame += num_frames
-            else:
-                raise ISPyBDepositionNotMade(
-                    f"No ISPyB IDs received by the start of {self.triggering_plan=}"
-                )
+                    self.start_frame += num_frames
 
     def descriptor(self, doc: EventDescriptor):
         self.descriptors[doc["uid"]] = doc
@@ -73,14 +71,19 @@ def event(self, doc: Event) -> Event:
             for start_info in self.zocalo_info:
                 start_info.filename = filename
                 self.zocalo_interactor.run_start(start_info)
+                self._started_zocalo_collections.append(start_info)
+            self.zocalo_info = []
         return doc
 
     def stop(self, doc: RunStop):
         if doc.get("run_start") == self.run_uid:
             ISPYB_ZOCALO_CALLBACK_LOGGER.info(
                 f"Zocalo handler received stop document, for run {doc.get('run_start')}."
             )
-            assert self.zocalo_interactor is not None
-            for info in self.zocalo_info:
+            if not self._started_zocalo_collections:
+                raise ISPyBDepositionNotMade(
+                    f"No ISPyB IDs received by the end of {self.triggering_plan=}"
+                )
+            for info in self._started_zocalo_collections:
                 self.zocalo_interactor.run_end(info.ispyb_dcid)
             self._reset_state()

src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py

@@ -461,6 +461,8 @@ def rotation_scan_core(
 
             yield from rotation_scan_core(single_scan)
 
+        yield from bps.unstage(eiger)
+
     LOGGER.info("setting up and staging eiger...")
     yield from start_preparing_data_collection_then_do_plan(
         composite.beamstop,
@@ -470,4 +472,3 @@ def rotation_scan_core(
         _multi_rotation_scan(),
         group=CONST.WAIT.ROTATION_READY_FOR_DC,
     )
-    yield from bps.unstage(eiger)

src/mx_bluesky/hyperion/external_interaction/callbacks/__main__.py

@@ -67,7 +67,7 @@ def create_rotation_callbacks() -> tuple[
     return (
         RotationNexusFileCallback(),
         RotationISPyBCallback(
-            emit=ZocaloCallback(CONST.PLAN.ROTATION_MAIN, CONST.ZOCALO_ENV)
+            emit=ZocaloCallback(CONST.PLAN.ROTATION_MULTI, CONST.ZOCALO_ENV)
         ),
     )
 

tests/unit_tests/common/external_interaction/callbacks/test_zocalo_handler.py

@@ -5,6 +5,7 @@
 
 from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import (
     ZocaloCallback,
+    ZocaloTrigger,
 )
 from mx_bluesky.common.external_interaction.ispyb.ispyb_store import (
     IspybIds,
@@ -14,65 +15,69 @@
 from mx_bluesky.hyperion.external_interaction.callbacks.__main__ import (
     create_gridscan_callbacks,
 )
-from mx_bluesky.hyperion.parameters.constants import CONST
 
 from .....conftest import TestData
 
-EXPECTED_DCID = 100
-EXPECTED_RUN_START_MESSAGE = {"event": "start", "ispyb_dcid": EXPECTED_DCID}
-EXPECTED_RUN_END_MESSAGE = {
-    "event": "end",
-    "ispyb_dcid": EXPECTED_DCID,
-    "ispyb_wait_for_runstatus": "1",
-}
+EXPECTED_RUN_START_MESSAGE = {"subplan_name": "test_plan_name", "uid": "my_uuid"}
+EXPECTED_RUN_END_MESSAGE = {"event": "end", "run_start": "my_uuid"}
 
 td = TestData()
 
 
-def start_dict(plan_name: str = "test_plan_name", env: str = "test_env"):
-    return {CONST.TRIGGER.ZOCALO: plan_name, "zocalo_environment": env}
-
-
 class TestZocaloHandler:
     def _setup_handler(self):
         zocalo_handler = ZocaloCallback("test_plan_name", "test_env")
         return zocalo_handler
 
     def test_handler_doesnt_trigger_on_wrong_plan(self):
         zocalo_handler = self._setup_handler()
-        zocalo_handler.start(start_dict("_not_test_plan_name"))  # type: ignore
-
-    def test_handler_raises_on_right_plan_with_wrong_metadata(self):
-        zocalo_handler = self._setup_handler()
-        with pytest.raises(AssertionError):
-            zocalo_handler.start({"subplan_name": "test_plan_name"})  # type: ignore
+        zocalo_handler.start({"sybplan_name": "_not_test_plan_name"})  # type: ignore
 
-    def test_handler_raises_on_right_plan_with_no_ispyb_ids(self):
+    @patch(
+        "mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback.ZocaloTrigger",
+        autospec=True,
+    )
+    def test_handler_stores_collection_if_ispyb_ids_come_in_with_triggering_plan(
+        self, zocalo_trigger: ZocaloTrigger
+    ):
         zocalo_handler = self._setup_handler()
-        with pytest.raises(ISPyBDepositionNotMade):
-            zocalo_handler.start(
-                {
-                    "subplan_name": "test_plan_name",
-                    "zocalo_environment": "test_env",
-                    "scan_points": [{"test": [1, 2, 3]}],
-                }  # type: ignore
-            )
+        assert not zocalo_handler.zocalo_info
+        zocalo_handler.start(
+            {
+                **EXPECTED_RUN_START_MESSAGE,
+                "ispyb_dcids": (135, 139),
+                "scan_points": [{"test": [1, 2, 3]}, {"test": [2, 3, 4]}],
+            }  # type: ignore
+        )
+        assert len(zocalo_handler.zocalo_info) == 2
 
     @patch(
         "mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback.ZocaloTrigger",
         autospec=True,
     )
-    def test_handler_inits_zocalo_trigger_on_right_plan(self, zocalo_trigger):
+    def test_handler_stores_collection_ispyb_ids_come_in_as_subplan(
+        self, zocalo_trigger: ZocaloTrigger
+    ):
         zocalo_handler = self._setup_handler()
+        assert not zocalo_handler.zocalo_info
+        zocalo_handler.start(EXPECTED_RUN_START_MESSAGE)  # type: ignore
+        assert not zocalo_handler.zocalo_info
         zocalo_handler.start(
             {
-                "subplan_name": "test_plan_name",
-                "zocalo_environment": "test_env",
-                "ispyb_dcids": (135, 139),
+                "subplan_name": "other_plan",
+                "ispyb_dcids": (135,),
                 "scan_points": [{"test": [1, 2, 3]}],
             }  # type: ignore
         )
-        assert zocalo_handler.zocalo_interactor is not None
+        zocalo_handler.start(
+            {
+                "subplan_name": "other_plan",
+                "ispyb_dcids": (563,),
+                "scan_points": [{"test": [2, 3, 4]}],
+            }  # type: ignore
+        )
+
+        assert len(zocalo_handler.zocalo_info) == 2
 
     @patch(
         "mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback.ZocaloTrigger",
@@ -133,3 +138,15 @@ def test_execution_of_do_fgs_triggers_zocalo_calls(
         assert zocalo_handler.zocalo_interactor.run_end.call_count == len(dc_ids)  # type: ignore
 
         zocalo_handler._reset_state.assert_called()
+
+    @patch(
+        "mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback.ZocaloTrigger",
+        autospec=True,
+    )
+    def test_handler_raises_on_the_end_of_a_plan_with_no_depositions(
+        self, zocalo_trigger: ZocaloTrigger
+    ):
+        zocalo_handler = self._setup_handler()
+        zocalo_handler.start(EXPECTED_RUN_START_MESSAGE)  # type: ignore
+        with pytest.raises(ISPyBDepositionNotMade):
+            zocalo_handler.stop(EXPECTED_RUN_END_MESSAGE)  # type: ignore