Merge pull request #1291 from DiamondLightSource/1017_robot_load_ispy…

…b_exp_eye

Add ispyb entries for robot load

src/hyperion/experiment_plans/experiment_registry.py

@@ -17,6 +17,7 @@
 from hyperion.external_interaction.callbacks.common.callback_util import (
     CallbacksFactory,
     create_gridscan_callbacks,
+    create_robot_load_and_centre_callbacks,
     create_rotation_callbacks,
 )
 from hyperion.parameters.plan_specific.grid_scan_with_edge_detect_params import (
@@ -101,7 +102,7 @@ class ExperimentRegistryEntry(TypedDict):
         "setup": robot_load_then_centre_plan.create_devices,
         "internal_param_type": RobotLoadThenCentreInternalParameters,
         "experiment_param_type": RobotLoadThenCentreParams,
-        "callbacks_factory": create_gridscan_callbacks,
+        "callbacks_factory": create_robot_load_and_centre_callbacks,
     },
 }
 EXPERIMENT_NAMES = list(PLAN_REGISTRY.keys())

src/hyperion/experiment_plans/robot_load_then_centre_plan.py

@@ -5,6 +5,7 @@
 from typing import cast
 
 import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
 from blueapi.core import BlueskyContext, MsgGenerator
 from dodal.devices.aperturescatterguard import AperturePositions, ApertureScatterguard
 from dodal.devices.attenuator import Attenuator
@@ -45,6 +46,7 @@
     set_energy_plan,
 )
 from hyperion.log import LOGGER
+from hyperion.parameters.constants import CONST
 from hyperion.parameters.plan_specific.pin_centre_then_xray_centre_params import (
     PinCentreThenXrayCentreInternalParameters,
 )
@@ -128,7 +130,6 @@ def prepare_for_robot_load(composite: RobotLoadThenCentreComposite):
                       composite.smargon.z, 0,
                       composite.smargon.omega, 0,
                       composite.smargon.chi, 0,
-
                       composite.smargon.phi, 0)
     # fmt: on
 
@@ -141,24 +142,45 @@ def robot_load_then_centre_plan(
 ):
     yield from prepare_for_robot_load(composite)
 
-    yield from bps.abs_set(
-        composite.robot,
-        SampleLocation(
-            parameters.experiment_params.sample_puck,
-            parameters.experiment_params.sample_pin,
-        ),
-        group="robot_load",
+    @bpp.run_decorator(
+        md={
+            "subplan_name": CONST.PLAN.ROBOT_LOAD,
+            "metadata": {
+                "visit_path": parameters.hyperion_params.ispyb_params.visit_path,
+                "sample_id": parameters.hyperion_params.ispyb_params.sample_id,
+                "sample_puck": parameters.experiment_params.sample_puck,
+                "sample_pin": parameters.experiment_params.sample_pin,
+            },
+            "activate_callbacks": [
+                "RobotLoadISPyBCallback",
+            ],
+        }
     )
-
-    if parameters.experiment_params.requested_energy_kev:
-        yield from set_energy_plan(
-            parameters.experiment_params.requested_energy_kev,
-            cast(SetEnergyComposite, composite),
+    def robot_load():
+        yield from bps.abs_set(
+            composite.robot,
+            SampleLocation(
+                parameters.experiment_params.sample_puck,
+                parameters.experiment_params.sample_pin,
+            ),
+            group="robot_load",
         )
 
-    yield from bps.wait("robot_load")
+        if parameters.experiment_params.requested_energy_kev:
+            yield from set_energy_plan(
+                parameters.experiment_params.requested_energy_kev,
+                cast(SetEnergyComposite, composite),
+            )
+
+        yield from bps.wait("robot_load")
+
+        yield from bps.create(name=CONST.PLAN.ROBOT_LOAD)
+        yield from bps.read(composite.robot.barcode)
+        yield from bps.save()
+
+        yield from wait_for_smargon_not_disabled(composite.smargon)
 
-    yield from wait_for_smargon_not_disabled(composite.smargon)
+    yield from robot_load()
 
     params_json = json.loads(parameters.json())
     pin_centre_params = PinCentreThenXrayCentreInternalParameters(**params_json)

src/hyperion/external_interaction/callbacks/common/callback_util.py

@@ -2,6 +2,9 @@
 
 from bluesky.callbacks import CallbackBase
 
+from hyperion.external_interaction.callbacks.robot_load.ispyb_callback import (
+    RobotLoadISPyBCallback,
+)
 from hyperion.external_interaction.callbacks.rotation.ispyb_callback import (
     RotationISPyBCallback,
 )
@@ -16,7 +19,17 @@
 )
 from hyperion.external_interaction.callbacks.zocalo_callback import ZocaloCallback
 
-CallbacksFactory = Callable[[], Tuple[CallbackBase, CallbackBase]]
+CallbacksFactory = Callable[[], Tuple[CallbackBase, ...]]
+
+
+def create_robot_load_and_centre_callbacks() -> (
+    Tuple[GridscanNexusFileCallback, GridscanISPyBCallback, RobotLoadISPyBCallback]
+):
+    return (
+        GridscanNexusFileCallback(),
+        GridscanISPyBCallback(emit=ZocaloCallback()),
+        RobotLoadISPyBCallback(),
+    )
 
 
 def create_gridscan_callbacks() -> (

src/hyperion/external_interaction/callbacks/common/ispyb_mapping.py

@@ -92,6 +92,12 @@ def populate_remaining_data_collection_info(
     return data_collection_info
 
 
+def get_proposal_and_session_from_visit_string(visit_string: str) -> tuple[str, int]:
+    visit_parts = visit_string.split("-")
+    assert len(visit_parts) == 2, f"Unexpected visit string {visit_string}"
+    return visit_parts[0], int(visit_parts[1])
+
+
 def get_visit_string_from_path(path: Optional[str]) -> Optional[str]:
     match = re.search(VISIT_PATH_REGEX, path) if path else None
     return str(match.group(1)) if match else None

src/hyperion/external_interaction/callbacks/robot_load/ispyb_callback.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Optional
+
+from event_model.documents import EventDescriptor
+
+from hyperion.external_interaction.callbacks.common.ispyb_mapping import (
+    get_proposal_and_session_from_visit_string,
+    get_visit_string_from_path,
+)
+from hyperion.external_interaction.callbacks.plan_reactive_callback import (
+    PlanReactiveCallback,
+)
+from hyperion.external_interaction.ispyb.exp_eye_store import (
+    ExpeyeInteraction,
+    RobotActionID,
+)
+from hyperion.log import ISPYB_LOGGER
+from hyperion.parameters.constants import CONST
+
+if TYPE_CHECKING:
+    from event_model.documents import Event, EventDescriptor, RunStart, RunStop
+
+
+class RobotLoadISPyBCallback(PlanReactiveCallback):
+    def __init__(self) -> None:
+        ISPYB_LOGGER.debug("Initialising ISPyB Robot Load Callback")
+        super().__init__(log=ISPYB_LOGGER)
+        self.run_uid: Optional[str] = None
+        self.descriptors: Dict[str, EventDescriptor] = {}
+        self.action_id: RobotActionID | None = None
+        self.expeye = ExpeyeInteraction()
+
+    def activity_gated_start(self, doc: RunStart):
+        ISPYB_LOGGER.debug("ISPyB robot load callback received start document.")
+        if doc.get("subplan_name") == CONST.PLAN.ROBOT_LOAD:
+            self.run_uid = doc.get("uid")
+            assert isinstance(metadata := doc.get("metadata"), Dict)
+            assert isinstance(
+                visit := get_visit_string_from_path(metadata["visit_path"]), str
+            )
+            proposal, session = get_proposal_and_session_from_visit_string(visit)
+            self.action_id = self.expeye.start_load(
+                proposal,
+                session,
+                metadata["sample_id"],
+                metadata["sample_puck"],
+                metadata["sample_pin"],
+            )
+        return super().activity_gated_start(doc)
+
+    def activity_gated_descriptor(self, doc: EventDescriptor) -> EventDescriptor | None:
+        self.descriptors[doc["uid"]] = doc
+
+    def activity_gated_event(self, doc: Event) -> Event | None:
+        event_descriptor = self.descriptors.get(doc["descriptor"])
+        if event_descriptor and event_descriptor.get("name") == CONST.PLAN.ROBOT_LOAD:
+            assert (
+                self.action_id is not None
+            ), "ISPyB Robot load callback event called unexpectedly"
+            barcode = doc["data"]["robot-barcode"]
+            self.expeye.update_barcode(self.action_id, barcode)
+
+        return super().activity_gated_event(doc)
+
+    def activity_gated_stop(self, doc: RunStop) -> RunStop | None:
+        ISPYB_LOGGER.debug("ISPyB robot load callback received stop document.")
+        if doc.get("run_start") == self.run_uid:
+            assert (
+                self.action_id is not None
+            ), "ISPyB Robot load callback stop called unexpectedly"
+            exit_status = (
+                doc.get("exit_status") or "Exit status not available in stop document!"
+            )
+            reason = doc.get("reason") or ""
+            self.expeye.end_load(self.action_id, exit_status, reason)
+            self.action_id = None
+        return super().activity_gated_stop(doc)

src/hyperion/external_interaction/ispyb/exp_eye_store.py

@@ -0,0 +1,114 @@
+import configparser
+from typing import Dict, Tuple
+
+from requests import patch, post
+from requests.auth import AuthBase
+
+from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade
+from hyperion.external_interaction.ispyb.ispyb_utils import (
+    get_current_time_string,
+    get_ispyb_config,
+)
+
+RobotActionID = int
+
+
+class BearerAuth(AuthBase):
+    def __init__(self, token):
+        self.token = token
+
+    def __call__(self, r):
+        r.headers["authorization"] = "Bearer " + self.token
+        return r
+
+
+def _get_base_url_and_token() -> Tuple[str, str]:
+    config = configparser.ConfigParser()
+    conf = get_ispyb_config()
+    config.read(conf)
+    expeye_config = config["expeye"]
+    return expeye_config["url"], expeye_config["token"]
+
+
+class ExpeyeInteraction:
+    CREATE_ROBOT_ACTION = "/proposals/{proposal}/sessions/{visit_number}/robot-actions"
+    UPDATE_ROBOT_ACTION = "/robot-actions/{action_id}"
+
+    def __init__(self) -> None:
+        url, token = _get_base_url_and_token()
+        self.base_url = url + "/core"
+        self.auth = BearerAuth(token)
+
+    def _send_and_get_response(self, url, data, send_func) -> Dict:
+        response = send_func(url, auth=self.auth, json=data)
+        if not response.ok:
+            raise ISPyBDepositionNotMade(f"Could not write {data} to {url}: {response}")
+        return response.json()
+
+    def start_load(
+        self,
+        proposal_reference: str,
+        visit_number: int,
+        sample_id: int,
+        dewar_location: int,
+        container_location: int,
+    ) -> RobotActionID:
+        """Create a robot load entry in ispyb.
+
+        Args:
+            proposal_reference (str): The proposal of the experiment e.g. cm37235
+            visit_number (int): The visit number for the proposal, usually this can be
+                                found added to the end of the proposal e.g. the data for
+                                visit number 2 of proposal cm37235 is in cm37235-2
+            sample_id (int): The id of the sample in the database
+            dewar_location (int): Which puck in the dewar the sample is in
+            container_location (int): Which pin in that puck has the sample
+
+        Returns:
+            RobotActionID: The id of the robot load action that is created
+        """
+        url = self.base_url + self.CREATE_ROBOT_ACTION.format(
+            proposal=proposal_reference, visit_number=visit_number
+        )
+
+        data = {
+            "startTimestamp": get_current_time_string(),
+            "sampleId": sample_id,
+            "actionType": "LOAD",
+            "containerLocation": container_location,
+            "dewarLocation": dewar_location,
+        }
+        response = self._send_and_get_response(url, data, post)
+        return response["robotActionId"]
+
+    def update_barcode(self, action_id: RobotActionID, barcode: str):
+        """Update the barcode of an existing robot action.
+
+        Args:
+            action_id (RobotActionID): The id of the action to update
+            barcode (str): The barcode to give the action
+        """
+        url = self.base_url + self.UPDATE_ROBOT_ACTION.format(action_id=action_id)
+
+        data = {"sampleBarcode": barcode}
+        self._send_and_get_response(url, data, patch)
+
+    def end_load(self, action_id: RobotActionID, status: str, reason: str):
+        """Finish an existing robot action, providing final information about how it went
+
+        Args:
+            action_id (RobotActionID): The action to finish.
+            status (str): The status of the action at the end, "success" for success,
+                          otherwise error
+            reason (str): If the status is in error than the reason for that error
+        """
+        url = self.base_url + self.UPDATE_ROBOT_ACTION.format(action_id=action_id)
+
+        run_status = "SUCCESS" if status == "success" else "ERROR"
+
+        data = {
+            "endTimestamp": get_current_time_string(),
+            "status": run_status,
+            "message": reason,
+        }
+        self._send_and_get_response(url, data, patch)

src/hyperion/parameters/constants.py

@@ -11,7 +11,7 @@ class SimConstants:
     # this one is for unit tests
     ISPYB_CONFIG = "tests/test_data/test_config.cfg"
     # this one is for system tests
-    DEV_ISPYB_DATABASE_CFG = "/dls_sw/dasc/mariadb/credentials/ispyb-dev.cfg"
+    DEV_ISPYB_DATABASE_CFG = "/dls_sw/dasc/mariadb/credentials/ispyb-hyperion-dev.cfg"
 
 
 @dataclass(frozen=True)
@@ -21,6 +21,8 @@ class PlanNameConstants:
     ISPYB_HARDWARE_READ = "ispyb_reading_hardware"
     ISPYB_TRANSMISSION_FLUX_READ = "ispyb_update_transmission_flux"
     ZOCALO_HW_READ = "zocalo_read_hardware_plan"
+    # Robot load
+    ROBOT_LOAD = "robot_load"
     # Gridscan
     GRIDSCAN_OUTER = "run_gridscan_move_and_tidy"
     GRIDSCAN_AND_MOVE = "run_gridscan_and_move"

tests/system_tests/external_interaction/test_exp_eye_dev.py

@@ -0,0 +1,41 @@
+import os
+from time import sleep
+
+import pytest
+from requests import get
+
+from hyperion.external_interaction.ispyb.exp_eye_store import ExpeyeInteraction
+from hyperion.parameters.constants import CONST
+
+
+@pytest.mark.s03
+def test_start_and_end_robot_load():
+    os.environ["ISPYB_CONFIG_PATH"] = CONST.SIM.DEV_ISPYB_DATABASE_CFG
+
+    SAMPLE_ID = 5289780
+    BARCODE = "test_barcode"
+
+    expeye = ExpeyeInteraction()
+
+    robot_action_id = expeye.start_load("cm37235", 2, SAMPLE_ID, 40, 3)
+
+    sleep(0.5)
+
+    print(f"Created {robot_action_id}")
+
+    expeye.update_barcode(robot_action_id, BARCODE)
+
+    sleep(0.5)
+
+    expeye.end_load(robot_action_id, "fail", "Oh no!")
+
+    get_robot_data_url = f"{expeye.base_url}/robot-actions/{robot_action_id}"
+    response = get(get_robot_data_url, auth=expeye.auth)
+
+    assert response.ok
+
+    response = response.json()
+    assert response["robotActionId"] == robot_action_id
+    assert response["status"] == "ERROR"
+    assert response["sampleId"] == SAMPLE_ID
+    assert response["sampleBarcode"] == BARCODE