ZocaloResults: add parameter to use results from GPU

src/dodal/devices/zocalo/zocalo_results.py

@@ -61,7 +61,7 @@
     be triggered from a plan-subscribed callback using the run_start() and run_end()
     methods on dodal.devices.zocalo.ZocaloTrigger.
 
-    See https://github.com/DiamondLightSource/dodal/wiki/How-to-Interact-with-Zocalo"""
+    See https://diamondlightsource.github.io/dodal/main/how-to/zocalo.html"""
 
     def __init__(
         self,
@@ -71,6 +71,7 @@
         sort_key: str = DEFAULT_SORT_KEY.value,
         timeout_s: float = DEFAULT_TIMEOUT,
         prefix: str = "",
+        use_fastest_zocalo_result: bool = False
     ) -> None:
         self.zocalo_environment = zocalo_environment
         self.sort_key = SortKeys[sort_key]
@@ -79,6 +80,7 @@
         self._prefix = prefix
         self._raw_results_received: Queue = Queue()
         self.transport: CommonTransport | None = None
+        self.use_fastest_zocalo_result = use_fastest_zocalo_result
 
         self.results, self._results_setter = soft_signal_r_and_setter(
             list[XrcResult], name="results"
@@ -165,7 +167,10 @@
             )
 
             raw_results = self._raw_results_received.get(timeout=self.timeout_s)
-            LOGGER.info(f"Zocalo: found {len(raw_results['results'])} crystals.")
+            source_of_results = "CPU" if not raw_results["ispyb_ids"].get("gpu") else "GPU"
+            if source_of_results == "CPU" and self.use_fastest_zocalo_result:
+                LOGGER.warn("Recieved zocalo results from CPU before GPU")
+            LOGGER.info(f"Zocalo results from {source_of_results} processing: found {len(raw_results['results'])} crystals.")
             # Sort from strongest to weakest in case of multiple crystals
             await self._put_results(
                 sorted(
@@ -237,9 +242,18 @@
             self.transport.ack(header)  # type: ignore # we create transport here
 
             results = message.get("results", [])
-            self._raw_results_received.put(
-                {"results": results, "ispyb_ids": recipe_parameters}
-            )
+
+            if self.use_fastest_zocalo_result:
+                self._raw_results_received.put(
+                    {"results": results, "ispyb_ids": recipe_parameters}
+                )
+            else:
+                #Only add to queue if results are from CPU
+                if recipe_parameters.get('gpu') == None:
+                    self._raw_results_received.put(
+                    {"results": results, "ispyb_ids": recipe_parameters}
+                )
+
 
         subscription = workflows.recipe.wrap_subscribe(
             self.transport,

tests/devices/unit_tests/test_zocalo_results.py

@@ -207,11 +207,14 @@ async def test_zocalo_results_trigger_log_message(
         name="zocalo", zocalo_environment="dev_artemis", timeout_s=2
     )
 
+    recipe_wrapper = MagicMock()
+    recipe_wrapper.recipe_step = {"parameters": {}}
+
     def zocalo_plan():
         yield from bps.stage(zocalo_results)
         receive_result = mock_wrap_subscribe.mock_calls[0].args[2]
         receive_result(
-            MagicMock(autospec=RecipeWrapper),
+            recipe_wrapper,
             {},
             {
                 "results": [
@@ -235,7 +238,7 @@ def zocalo_plan():
         yield from bps.trigger(zocalo_results)
 
     RE(zocalo_plan())
-    mock_logger.info.assert_has_calls([call("Zocalo: found 1 crystals.")])
+    mock_logger.info.assert_has_calls([call("Zocalo results from CPU processing: found 1 crystals.")])
 
 
 @patch("dodal.devices.zocalo.zocalo_results._get_zocalo_connection", autospec=True)