From 388db6ab2c9271b62637ff4b5e746ff35adb657f Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Tue, 28 Jan 2025 16:39:48 +0000 Subject: [PATCH] Use gpu results if specified (#1025) * Use gpu results if specified * Zocalo device should not let use gpu and compare to both be enabled --- src/dodal/devices/zocalo/zocalo_results.py | 25 ++++- .../devices/unit_tests/test_zocalo_results.py | 103 ++++++++++++------ 2 files changed, 91 insertions(+), 37 deletions(-) diff --git a/src/dodal/devices/zocalo/zocalo_results.py b/src/dodal/devices/zocalo/zocalo_results.py index 1cdea67e17..c0706d7637 100644 --- a/src/dodal/devices/zocalo/zocalo_results.py +++ b/src/dodal/devices/zocalo/zocalo_results.py @@ -129,7 +129,12 @@ class ZocaloResults(StandardReadable, Triggerable): prefix (str): EPICS PV prefix for the device - use_cpu_and_gpu (bool): When True, ZocaloResults will wait for results from the CPU and the GPU, compare them, and provide a warning if the results differ. When False, ZocaloResults will only use results from the CPU + use_cpu_and_gpu (bool): When True, ZocaloResults will wait for results from the + CPU and the GPU, compare them, and provide a warning if the results differ. When + False, ZocaloResults will only use results from the CPU + + use_gpu (bool): When True, ZocaloResults will take the first set of + results that it receives (which are likely the GPU results) """ @@ -142,6 +147,7 @@ def __init__( timeout_s: float = DEFAULT_TIMEOUT, prefix: str = "", use_cpu_and_gpu: bool = False, + use_gpu: bool = False, ) -> None: self.zocalo_environment = zocalo_environment self.sort_key = SortKeys[sort_key] @@ -151,6 +157,7 @@ def __init__( self._raw_results_received: Queue = Queue() self.transport: CommonTransport | None = None self.use_cpu_and_gpu = use_cpu_and_gpu + self.use_gpu = use_gpu self.centre_of_mass, self._com_setter = soft_signal_r_and_setter( Array1D[np.float64], name="centre_of_mass" @@ -213,6 +220,11 @@ async def stage(self): clearing the queue. Plans using this device should wait on ZOCALO_STAGE_GROUP before triggering processing for the experiment""" + if self.use_cpu_and_gpu and self.use_gpu: + raise ValueError( + "Cannot compare GPU and CPU results and use GPU results at the same time." + ) + LOGGER.info("Subscribing to results queue") try: self._subscribe_to_results() @@ -253,8 +265,13 @@ async def trigger(self): raw_results = self._raw_results_received.get(timeout=self.timeout_s) source_of_first_results = source_from_results(raw_results) - # Wait for results from CPU and GPU, warn and continue if only GPU times out. Error if CPU times out + if self.use_gpu and source_of_first_results == ZocaloSource.CPU: + LOGGER.warning( + "Configured to use GPU results but CPU came first, using CPU results." + ) + if self.use_cpu_and_gpu: + # Wait for results from CPU and GPU, warn and continue if only GPU times out. Error if CPU times out if source_of_first_results == ZocaloSource.CPU: LOGGER.warning("Received zocalo results from CPU before GPU") raw_results_two_sources = [raw_results] @@ -303,7 +320,7 @@ async def trigger(self): raise err LOGGER.info( - f"Zocalo results from {ZocaloSource.CPU.value} processing: found {len(raw_results['results'])} crystals." + f"Zocalo results from {source_from_results(raw_results)} processing: found {len(raw_results['results'])} crystals." ) # Sort from strongest to weakest in case of multiple crystals await self._put_results( @@ -335,7 +352,7 @@ def _receive_result( results = message.get("results", []) - if self.use_cpu_and_gpu: + if self.use_cpu_and_gpu or self.use_gpu: self._raw_results_received.put( {"results": results, "recipe_parameters": recipe_parameters} ) diff --git a/tests/devices/unit_tests/test_zocalo_results.py b/tests/devices/unit_tests/test_zocalo_results.py index fd7150f844..fa6efaf810 100644 --- a/tests/devices/unit_tests/test_zocalo_results.py +++ b/tests/devices/unit_tests/test_zocalo_results.py @@ -257,20 +257,7 @@ def zocalo_plan(): recipe_wrapper, {}, { - "results": [ - { - "centre_of_mass": [ - 2.207133058984911, - 1.4175240054869684, - 13.317215363511659, - ], - "max_voxel": [2, 1, 13], - "max_count": 702.0, - "n_voxels": 12, - "total_count": 5832.0, - "bounding_box": [[1, 0, 12], [4, 3, 15]], - } - ], + "results": [TEST_RESULTS[0]], "status": "success", "type": "3d", }, @@ -313,6 +300,18 @@ async def test_if_use_cpu_and_gpu_zocalos_then_wait_twice_for_results( assert zocalo_results._raw_results_received.get.call_count == 2 +async def test_if_use_gpu_then_only_use_first_result( + zocalo_results: ZocaloResults, RE: RunEngine +): + zocalo_results.use_gpu = True + await zocalo_results.stage() + zocalo_results._raw_results_received.put([]) + zocalo_results._raw_results_received.put([]) + zocalo_results._raw_results_received.get = MagicMock() + RE(bps.trigger(zocalo_results, wait=False)) + assert zocalo_results._raw_results_received.get.call_count == 1 + + @patch("dodal.devices.zocalo.zocalo_results.LOGGER") async def test_source_of_zocalo_results_correctly_identified( mock_logger, zocalo_results: ZocaloResults, RE: RunEngine @@ -351,14 +350,17 @@ async def test_if_zocalo_results_timeout_from_gpu_then_warn( ) -async def test_if_zocalo_results_from_gpu_but_not_cpu_then_error( +async def test_given_comparing_results_if_zocalo_results_from_gpu_but_not_cpu_then_error( zocalo_results: ZocaloResults, RE: RunEngine ): zocalo_results.use_cpu_and_gpu = True await zocalo_results.stage() zocalo_results._raw_results_received.get = MagicMock( side_effect=[ - {"recipe_parameters": {"test": 0, "gpu": True}, "results": []}, + { + "recipe_parameters": {"test": 0, "gpu": True}, + "results": [TEST_RESULTS[0]], + }, Empty, ] ) @@ -366,8 +368,26 @@ async def test_if_zocalo_results_from_gpu_but_not_cpu_then_error( await zocalo_results.trigger() +async def test_given_using_gpu_results_if_zocalo_results_from_gpu_but_not_cpu_then_uses_gpu( + zocalo_results: ZocaloResults, RE: RunEngine +): + zocalo_results.use_gpu = True + await zocalo_results.stage() + zocalo_results._raw_results_received.get = MagicMock( + side_effect=[ + { + "recipe_parameters": {"dcgid": 0, "dcid": 0, "gpu": True}, + "results": [TEST_RESULTS[0]], + }, + Empty, + ] + ) + await zocalo_results.trigger() + assert len(await zocalo_results.centre_of_mass.get_value()) + + @patch("dodal.devices.zocalo.zocalo_results.LOGGER") -async def test_if_cpu_results_arrive_before_gpu_then_warn( +async def test_given_comparing_results_if_cpu_results_arrive_before_gpu_then_warn( mock_logger, zocalo_results: ZocaloResults, RE: RunEngine ): zocalo_results.use_cpu_and_gpu = True @@ -418,7 +438,7 @@ async def test_if_cpu_results_arrive_before_gpu_then_warn( ], ) @patch("dodal.devices.zocalo.zocalo_results.LOGGER") -async def test_warning_if_results_are_different( +async def test_given_comparing_results_then_warning_if_results_are_different( mock_logger, zocalo_results: ZocaloResults, RE: RunEngine, dict1, dict2, output ): zocalo_results.use_cpu_and_gpu = True @@ -476,20 +496,7 @@ def zocalo_plan(): recipe_wrapper, {}, { - "results": [ - { - "centre_of_mass": [ - 2.207133058984911, - 1.4175240054869684, - 13.317215363511659, - ], - "max_voxel": [2, 1, 13], - "max_count": 702.0, - "n_voxels": 12, - "total_count": 5832.0, - "bounding_box": [[1, 0, 12], [4, 3, 15]], - } - ], + "results": [TEST_RESULTS[0]], "status": "success", "type": "3d", }, @@ -507,7 +514,7 @@ def zocalo_plan(): RE(zocalo_plan()) -async def test_given_gpu_enabled_when_no_results_found_then_returns_no_results( +async def test_given_comparing_results_when_no_results_found_then_returns_no_results( zocalo_results: ZocaloResults, ): zocalo_results.use_cpu_and_gpu = True @@ -520,3 +527,33 @@ async def test_given_gpu_enabled_when_no_results_found_then_returns_no_results( ) await zocalo_results.trigger() assert len(await zocalo_results.centre_of_mass.get_value()) == 0 + + +@patch("dodal.devices.zocalo.zocalo_results.LOGGER") +async def test_given_using_gpu_results_if_results_from_cpu_first_then_warn_and_use( + mock_logger: MagicMock, zocalo_results: ZocaloResults +): + zocalo_results.use_gpu = True + await zocalo_results.stage() + zocalo_results._raw_results_received.get = MagicMock( + side_effect=[ + { + "recipe_parameters": {"dcgid": 0, "dcid": 0}, + "results": [TEST_RESULTS[0]], + }, + ] + ) + await zocalo_results.trigger() + assert len(await zocalo_results.centre_of_mass.get_value()) + mock_logger.warning.assert_called_with( + "Configured to use GPU results but CPU came first, using CPU results." + ) + + +async def test_given_using_gpu_results_and_comparing_results_both_on_then_error_when_staged( + zocalo_results: ZocaloResults, +): + zocalo_results.use_gpu = True + zocalo_results.use_cpu_and_gpu = True + with pytest.raises(ValueError): + await zocalo_results.stage()