Murko: publish and look for stop messages (#1686)

* Filter out small x Murko results

* Add test

* Rename confusing variable

* More fixes

* Publish complete message to redis

* Wait for results complete message to stop results loop

* Move forwarding complete message to MurkoCallback

* Fix MurkoResults device re stop messages

* Remove unused signals

* Relax typing

* Fix tests

* Add test

* Remove unused code

* Add test and tidy tests

* Small changes

* Improve test typing

* PR comments

Author: jacob720

src/dodal/devices/i04/murko_results.py

@@ -1,5 +1,6 @@
 import json
 import pickle
+import time
 from dataclasses import dataclass
 from enum import Enum
 from typing import TypedDict
@@ -21,6 +22,7 @@
 from dodal.log import LOGGER
 
 NO_MURKO_RESULT = (-1, -1)
+RESULTS_COMPLETE_MESSAGE = "murko_results_complete"
 
 
 class MurkoMetadata(TypedDict):
@@ -71,7 +73,8 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
     solutions for y and z can be calculated using numpy's linear algebra library.
     """
 
-    TIMEOUT_S = 2
+    GET_MESSAGE_TIMEOUT_S = 2
+    RESULTS_COMPLETE_TIMEOUT_S = 5
     PERCENTAGE_TO_USE = 25
     LEFTMOST_PIXEL_TO_USE = 10
     NUMBER_OF_WRONG_RESULTS_TO_LOG = 5
@@ -82,7 +85,6 @@ def __init__(
         redis_password=RedisConstants.REDIS_PASSWORD,
         redis_db=RedisConstants.MURKO_REDIS_DB,
         name="",
-        stop_angle=350,
     ):
         self.redis_client = StrictRedis(
             host=redis_host,
@@ -91,8 +93,6 @@ def __init__(
         )
         self.pubsub = self.redis_client.pubsub()
         self.sample_id = soft_signal_rw(str)  # Should get from redis
-        self.stop_angle = soft_signal_rw(int, initial_value=stop_angle)
-        self.invert_stop_angle = soft_signal_rw(bool, initial_value=False)
 
         self._reset()
 
@@ -121,15 +121,27 @@ async def unstage(self):
 
     @AsyncStatus.wrap
     async def trigger(self):
-        # Wait for results
         sample_id = await self.sample_id.get_value()
-
-        while not await self.check_if_reached_stop_angle():
+        t_last_result = time.time()
+        while True:
+            if time.time() - t_last_result > self.RESULTS_COMPLETE_TIMEOUT_S:
+                LOGGER.warning(
+                    f"Time since last result > {self.RESULTS_COMPLETE_TIMEOUT_S}, expected to receive {RESULTS_COMPLETE_MESSAGE}"
+                )
+                break
             # waits here for next batch to be received
-            message = await self.pubsub.get_message(timeout=self.TIMEOUT_S)
-            if message is None:
-                continue
-            await self.process_batch(message, sample_id)
+            message = await self.pubsub.get_message(timeout=self.GET_MESSAGE_TIMEOUT_S)
+            if message and message["type"] == "message":
+                t_last_result = time.time()
+                data = pickle.loads(message["data"])
+
+                if data == RESULTS_COMPLETE_MESSAGE:
+                    LOGGER.info(
+                        f"Received results complete message: {RESULTS_COMPLETE_MESSAGE}"
+                    )
+                    break
+
+                await self.process_batch(data, sample_id)
 
         if not self._results:
             raise NoResultsFoundError("No results retrieved from Murko")
@@ -157,22 +169,18 @@ async def trigger(self):
                 f"murko:{sample_id}:metadata", result.uuid, json.dumps(result.metadata)
             )
 
-    async def process_batch(self, message: dict | None, sample_id: str):
-        if message and message["type"] == "message":
-            batch_results: list[dict] = pickle.loads(message["data"])
-            for results in batch_results:
-                for uuid, result in results.items():
-                    if metadata_str := await self.redis_client.hget(  # type: ignore
-                        f"murko:{sample_id}:metadata", uuid
-                    ):
-                        LOGGER.info(
-                            f"Found metadata for uuid {uuid}, processing result"
-                        )
-                        self.process_result(
-                            result, MurkoMetadata(json.loads(metadata_str))
-                        )
-                    else:
-                        LOGGER.info(f"Found no metadata for uuid {uuid}")
+    async def process_batch(
+        self, batch_results: list[tuple[str, dict]], sample_id: str
+    ):
+        for result_with_uuid in batch_results:
+            uuid, result = result_with_uuid
+            if metadata_str := await self.redis_client.hget(  # type: ignore
+                f"murko:{sample_id}:metadata", uuid
+            ):
+                LOGGER.info(f"Found metadata for uuid {uuid}, processing result")
+                self.process_result(result, MurkoMetadata(json.loads(metadata_str)))
+            else:
+                LOGGER.info(f"Found no metadata for uuid {uuid}")
 
     def process_result(self, result: dict, metadata: MurkoMetadata):
         """Uses the 'most_likely_click' coordinates from Murko to calculate the
@@ -253,16 +261,6 @@ def filter_outliers(self):
         LOGGER.info(f"Number of results after filtering: {len(best_x)}")
         return best_x
 
-    async def check_if_reached_stop_angle(self):
-        inverted = await self.invert_stop_angle.get_value()
-        stop_angle = await self.stop_angle.get_value()
-        if self._last_omega is None:
-            return False
-        if inverted:
-            return self._last_omega <= stop_angle
-        else:
-            return self._last_omega >= stop_angle
-
 
 def get_yz_least_squares(vertical_dists: list, omegas: list) -> tuple[float, float]:
     """Get the least squares solution for y and z from the vertical distances and omega angles.