From cc3e89cf2cde279a64ee3a0df02a205cf50ddc8f Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Thu, 29 Feb 2024 20:59:20 +0000 Subject: [PATCH 1/4] (DiamondLightSource/hyperion#1091) Allow sending start index and frame number to zocalo --- .../devices/zocalo/zocalo_interaction.py | 22 ++++++++-- .../system_tests/test_zocalo_results.py | 4 +- .../unit_tests/test_zocalo_interaction.py | 41 +++++++++++++------ 3 files changed, 49 insertions(+), 18 deletions(-) diff --git a/src/dodal/devices/zocalo/zocalo_interaction.py b/src/dodal/devices/zocalo/zocalo_interaction.py index 217da22d41..0f9211621f 100644 --- a/src/dodal/devices/zocalo/zocalo_interaction.py +++ b/src/dodal/devices/zocalo/zocalo_interaction.py @@ -42,16 +42,30 @@ def _send_to_zocalo(self, parameters: dict): finally: transport.disconnect() - def run_start(self, data_collection_id: int): + def run_start( + self, + data_collection_id: int, + start_index: int, + number_of_frames: int, + ): """Tells the data analysis pipeline we have started a run. Assumes that appropriate data has already been put into ISPyB Args: - data_collection_id (int): The ID of the data collection representing the - gridscan in ISPyB + data_collection_id (int): The ID of the data collection in ISPyB + start_index (int): The index of the first image of this collection within + the file written by the detector. + number_of_frames (int): The number of frames in this collection. """ LOGGER.info(f"Starting Zocalo job with ispyb id {data_collection_id}") - self._send_to_zocalo({"event": "start", "ispyb_dcid": data_collection_id}) + self._send_to_zocalo( + { + "event": "start", + "ispyb_dcid": data_collection_id, + "start_index": start_index, + "number_of_frames": number_of_frames, + } + ) def run_end(self, data_collection_id: int): """Tells the data analysis pipeline we have finished a run. diff --git a/tests/devices/system_tests/test_zocalo_results.py b/tests/devices/system_tests/test_zocalo_results.py index 18849f0c46..3b40d20980 100644 --- a/tests/devices/system_tests/test_zocalo_results.py +++ b/tests/devices/system_tests/test_zocalo_results.py @@ -39,7 +39,7 @@ async def zocalo_device(): async def test_read_results_from_fake_zocalo(zocalo_device: ZocaloResults): zocalo_device._subscribe_to_results() zc = ZocaloTrigger("dev_artemis") - zc.run_start(0) + zc.run_start(0, 0, 100) zc.run_end(0) zocalo_device.timeout_s = 5 @@ -66,7 +66,7 @@ async def test_stage_unstage_controls_read_results_from_fake_zocalo( def plan(): yield from bps.open_run() - zc.run_start(0) + zc.run_start(0, 0, 100) zc.run_end(0) yield from bps.sleep(0.15) yield from bps.trigger_and_read([zocalo_device]) diff --git a/tests/devices/unit_tests/test_zocalo_interaction.py b/tests/devices/unit_tests/test_zocalo_interaction.py index ea43abfafe..978bc197b7 100644 --- a/tests/devices/unit_tests/test_zocalo_interaction.py +++ b/tests/devices/unit_tests/test_zocalo_interaction.py @@ -13,7 +13,12 @@ SIM_ZOCALO_ENV = "dev_artemis" EXPECTED_DCID = 100 -EXPECTED_RUN_START_MESSAGE = {"event": "start", "ispyb_dcid": EXPECTED_DCID} +EXPECTED_RUN_START_MESSAGE = { + "event": "start", + "ispyb_dcid": EXPECTED_DCID, + "start_index": 0, + "number_of_frames": 100, +} EXPECTED_RUN_END_MESSAGE = { "event": "end", "ispyb_dcid": EXPECTED_DCID, @@ -65,27 +70,39 @@ def with_exception(function_to_run, mock_transport): @mark.parametrize( - "function_to_test,function_wrapper,expected_message", + "function_wrapper,expected_message", [ - (zc.run_start, normally, EXPECTED_RUN_START_MESSAGE), + (normally, EXPECTED_RUN_START_MESSAGE), ( - zc.run_start, with_exception, EXPECTED_RUN_START_MESSAGE, ), - (zc.run_end, normally, EXPECTED_RUN_END_MESSAGE), - (zc.run_end, with_exception, EXPECTED_RUN_END_MESSAGE), ], ) -def test__run_start_and_end( - function_to_test: Callable, function_wrapper: Callable, expected_message: Dict -): +def test_run_start(function_wrapper: Callable, expected_message: Dict): + """ + Args: + function_wrapper (Callable): A wrapper used to test for expected exceptions + expected_message (Dict): The expected dictionary sent to zocalo + """ + function_to_run = partial(zc.run_start, EXPECTED_DCID, 0, 100) + function_to_run = partial(function_wrapper, function_to_run) + _test_zocalo(function_to_run, expected_message) + + +@mark.parametrize( + "function_wrapper,expected_message", + [ + (normally, EXPECTED_RUN_END_MESSAGE), + (with_exception, EXPECTED_RUN_END_MESSAGE), + ], +) +def test__run_start_and_end(function_wrapper: Callable, expected_message: Dict): """ Args: - function_to_test (Callable): The function to test e.g. start/stop zocalo - function_wrapper (Callable): A wrapper around the function, used to test for expected exceptions + function_wrapper (Callable): A wrapper used to test for expected exceptions expected_message (Dict): The expected dictionary sent to zocalo """ - function_to_run = partial(function_to_test, EXPECTED_DCID) + function_to_run = partial(zc.run_end, EXPECTED_DCID) function_to_run = partial(function_wrapper, function_to_run) _test_zocalo(function_to_run, expected_message) From 65609c8b8b2230ca37e1737995e19a84c9716acb Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Tue, 5 Mar 2024 13:42:55 +0000 Subject: [PATCH 2/4] (DiamondLightSource/hyperion#1091) Add filename to zocalo callback --- src/dodal/devices/zocalo/zocalo_interaction.py | 3 +++ tests/devices/unit_tests/test_zocalo_interaction.py | 4 +++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/dodal/devices/zocalo/zocalo_interaction.py b/src/dodal/devices/zocalo/zocalo_interaction.py index 0f9211621f..9229fbb995 100644 --- a/src/dodal/devices/zocalo/zocalo_interaction.py +++ b/src/dodal/devices/zocalo/zocalo_interaction.py @@ -45,6 +45,7 @@ def _send_to_zocalo(self, parameters: dict): def run_start( self, data_collection_id: int, + filename: str, start_index: int, number_of_frames: int, ): @@ -53,6 +54,7 @@ def run_start( Args: data_collection_id (int): The ID of the data collection in ISPyB + filename (str): The name of the file that the detector will store into dev/shm start_index (int): The index of the first image of this collection within the file written by the detector. number_of_frames (int): The number of frames in this collection. @@ -62,6 +64,7 @@ def run_start( { "event": "start", "ispyb_dcid": data_collection_id, + "filename": filename, "start_index": start_index, "number_of_frames": number_of_frames, } diff --git a/tests/devices/unit_tests/test_zocalo_interaction.py b/tests/devices/unit_tests/test_zocalo_interaction.py index 978bc197b7..d1b91172bb 100644 --- a/tests/devices/unit_tests/test_zocalo_interaction.py +++ b/tests/devices/unit_tests/test_zocalo_interaction.py @@ -13,11 +13,13 @@ SIM_ZOCALO_ENV = "dev_artemis" EXPECTED_DCID = 100 +EXPECTED_FILENAME = "test/file" EXPECTED_RUN_START_MESSAGE = { "event": "start", "ispyb_dcid": EXPECTED_DCID, "start_index": 0, "number_of_frames": 100, + "filename": EXPECTED_FILENAME, } EXPECTED_RUN_END_MESSAGE = { "event": "end", @@ -85,7 +87,7 @@ def test_run_start(function_wrapper: Callable, expected_message: Dict): function_wrapper (Callable): A wrapper used to test for expected exceptions expected_message (Dict): The expected dictionary sent to zocalo """ - function_to_run = partial(zc.run_start, EXPECTED_DCID, 0, 100) + function_to_run = partial(zc.run_start, EXPECTED_DCID, EXPECTED_FILENAME, 0, 100) function_to_run = partial(function_wrapper, function_to_run) _test_zocalo(function_to_run, expected_message) From e3ace0eead4ef7e43a02b25f072395d0d42c060c Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Tue, 5 Mar 2024 16:28:44 +0000 Subject: [PATCH 3/4] (DiamondLightSource/hyperion#1091) Created dataclass to hold zocalo info --- src/dodal/devices/zocalo/__init__.py | 5 +-- .../devices/zocalo/zocalo_interaction.py | 44 +++++++++++-------- .../unit_tests/test_zocalo_interaction.py | 10 +++-- 3 files changed, 33 insertions(+), 26 deletions(-) diff --git a/src/dodal/devices/zocalo/__init__.py b/src/dodal/devices/zocalo/__init__.py index edd2579d2e..3644917422 100644 --- a/src/dodal/devices/zocalo/__init__.py +++ b/src/dodal/devices/zocalo/__init__.py @@ -1,6 +1,4 @@ -from dodal.devices.zocalo.zocalo_interaction import ( - ZocaloTrigger, -) +from dodal.devices.zocalo.zocalo_interaction import ZocaloStartInfo, ZocaloTrigger from dodal.devices.zocalo.zocalo_results import ( NoResultsFromZocalo, NoZocaloSubscription, @@ -17,6 +15,7 @@ "ZOCALO_READING_PLAN_NAME", "NoResultsFromZocalo", "NoZocaloSubscription", + "ZocaloStartInfo", ] ZOCALO_READING_PLAN_NAME = "zocalo reading" diff --git a/src/dodal/devices/zocalo/zocalo_interaction.py b/src/dodal/devices/zocalo/zocalo_interaction.py index 9229fbb995..bc4e8dd47e 100644 --- a/src/dodal/devices/zocalo/zocalo_interaction.py +++ b/src/dodal/devices/zocalo/zocalo_interaction.py @@ -1,5 +1,8 @@ +import dataclasses import getpass import socket +from dataclasses import dataclass +from typing import Optional import zocalo.configuration from workflows.transport import lookup @@ -16,6 +19,22 @@ def _get_zocalo_connection(environment): return transport +@dataclass +class ZocaloStartInfo: + """ + data_collection_id (int): The ID of the data collection in ISPyB + filename (str): The name of the file that the detector will store into dev/shm + number_of_frames (int): The number of frames in this collection. + start_index (int): The index of the first image of this collection within the file + written by the detector + """ + + ispyb_dcid: int + filename: Optional[str] + start_index: int + number_of_frames: int + + class ZocaloTrigger: """This class just sends 'run_start' and 'run_end' messages to zocalo, it is intended to be used in bluesky callback classes. To get results from zocalo back @@ -44,31 +63,18 @@ def _send_to_zocalo(self, parameters: dict): def run_start( self, - data_collection_id: int, - filename: str, - start_index: int, - number_of_frames: int, + start_data: ZocaloStartInfo, ): """Tells the data analysis pipeline we have started a run. Assumes that appropriate data has already been put into ISPyB Args: - data_collection_id (int): The ID of the data collection in ISPyB - filename (str): The name of the file that the detector will store into dev/shm - start_index (int): The index of the first image of this collection within - the file written by the detector. - number_of_frames (int): The number of frames in this collection. + start_data (ZocaloStartInfo): Data about the collection to send to zocalo """ - LOGGER.info(f"Starting Zocalo job with ispyb id {data_collection_id}") - self._send_to_zocalo( - { - "event": "start", - "ispyb_dcid": data_collection_id, - "filename": filename, - "start_index": start_index, - "number_of_frames": number_of_frames, - } - ) + LOGGER.info(f"Starting Zocalo job {start_data}") + data = dataclasses.asdict(start_data) + data["event"] = "start" + self._send_to_zocalo(data) def run_end(self, data_collection_id: int): """Tells the data analysis pipeline we have finished a run. diff --git a/tests/devices/unit_tests/test_zocalo_interaction.py b/tests/devices/unit_tests/test_zocalo_interaction.py index d1b91172bb..5767fcd1bf 100644 --- a/tests/devices/unit_tests/test_zocalo_interaction.py +++ b/tests/devices/unit_tests/test_zocalo_interaction.py @@ -9,17 +9,18 @@ from dodal.devices.zocalo import ( ZocaloTrigger, ) +from dodal.devices.zocalo.zocalo_interaction import ZocaloStartInfo SIM_ZOCALO_ENV = "dev_artemis" EXPECTED_DCID = 100 EXPECTED_FILENAME = "test/file" EXPECTED_RUN_START_MESSAGE = { - "event": "start", "ispyb_dcid": EXPECTED_DCID, - "start_index": 0, - "number_of_frames": 100, "filename": EXPECTED_FILENAME, + "number_of_frames": 100, + "start_index": 0, + "event": "start", } EXPECTED_RUN_END_MESSAGE = { "event": "end", @@ -87,7 +88,8 @@ def test_run_start(function_wrapper: Callable, expected_message: Dict): function_wrapper (Callable): A wrapper used to test for expected exceptions expected_message (Dict): The expected dictionary sent to zocalo """ - function_to_run = partial(zc.run_start, EXPECTED_DCID, EXPECTED_FILENAME, 0, 100) + data = ZocaloStartInfo(EXPECTED_DCID, EXPECTED_FILENAME, 0, 100) + function_to_run = partial(zc.run_start, data) function_to_run = partial(function_wrapper, function_to_run) _test_zocalo(function_to_run, expected_message) From 61e841ecbc6b28e1f240c958294c9e1b658d4530 Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Tue, 5 Mar 2024 16:43:00 +0000 Subject: [PATCH 4/4] Fix typo --- src/dodal/devices/zocalo/zocalo_interaction.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/dodal/devices/zocalo/zocalo_interaction.py b/src/dodal/devices/zocalo/zocalo_interaction.py index bc4e8dd47e..48303e5f7d 100644 --- a/src/dodal/devices/zocalo/zocalo_interaction.py +++ b/src/dodal/devices/zocalo/zocalo_interaction.py @@ -22,7 +22,7 @@ def _get_zocalo_connection(environment): @dataclass class ZocaloStartInfo: """ - data_collection_id (int): The ID of the data collection in ISPyB + ispyb_dcid (int): The ID of the data collection in ISPyB filename (str): The name of the file that the detector will store into dev/shm number_of_frames (int): The number of frames in this collection. start_index (int): The index of the first image of this collection within the file