remove parts moved to bluesky_adaptive; cleanup

Author: ronpandolfi

tsuchinoko/execution/bluesky_adaptive.py

@@ -1,11 +1,9 @@
 import pickle
 import time
-from abc import ABC, abstractmethod
-from typing import Tuple, List, Sequence, Dict
+from typing import Tuple, List
 
 import zmq
 from loguru import logger
-from numpy._typing import ArrayLike
 
 from . import Engine
 
@@ -76,7 +74,6 @@ def get_measurements(self) -> List[Tuple]:
             else:
                 assert 'target_measured' in payload
                 x, (y, v) = payload['target_measured']
-                # TODO: Any additional quantities to be interrogated in Tsuchinoko can be included in the trailing dict
                 new_measurements.append((x, y, v, {}))
                 # stash the last position measured as the 'current' position of the instrument
                 self.position = x
@@ -105,158 +102,3 @@ def recv_payload(self, flags=0) -> dict:
             self.update_targets(self._last_targets_sent)
             payload_response = self.recv_payload(flags)
         return payload_response
-
-
-# ----------------------------------------------------------------------------------------------------------------------
-# This is a prototype Agent to be used with bluesky-adaptive. This should be extracted before merge.
-
-from bluesky_adaptive.agents.base import Agent
-
-
-class TsuchinokoBase(ABC):
-    def __init__(self, *args, host: str = '127.0.0.1', port: int = 5557, **kwargs):
-        """
-
-        Parameters
-        ----------
-        args
-            args passed through to `bluesky_adaptive.agents.base.Agent.__init__()`
-        host
-            A host address target for the zmq socket.
-        port
-            The port used for the zmq socket.
-        kwargs
-            kwargs passed through to `bluesky_adaptive.agents.base.Agent.__init__()`
-        """
-
-        super().__init__(*args, **kwargs)
-        self.host = host
-        self.port = port
-        self.outbound_measurements = []
-        self.context = None
-        self.socket = None
-        self.setup_socket()
-        self.last_targets_received = time.time()
-        self.kickstart()
-
-    def kickstart(self):
-        self.send_payload({'send_targets': True})  # kickstart to recover from shutdowns
-        self.last_targets_received = time.time()  # forgive lack of response until now
-
-    def setup_socket(self):
-        self.context = zmq.Context()
-        self.socket = self.context.socket(zmq.PAIR)
-
-        # Attempt to connect, retry every second if fails
-        while True:
-            try:
-                self.socket.connect(f"tcp://{self.host}:{self.port}")
-            except zmq.ZMQError:
-                logger.info(f'Unable to connect to tcp://{self.host}:{self.port}. Retrying in 1 second...')
-                time.sleep(1)
-            else:
-                logger.info(f'Connected to tcp://{self.host}:{self.port}.')
-                break
-
-    def tell(self, x, y, v):
-        """
-        Send measurement to BlueskyAdaptiveEngine
-        """
-        yv = (y, v)
-        payload = {'target_measured': (x, yv)}
-        self.send_payload(payload)
-
-    def ask(self, batch_size: int) -> Sequence[ArrayLike]:
-        """
-        Wait until at least one target is received, also exhaust the queue of received targets, overwriting old ones
-        """
-        payload = None
-        while True:
-            try:
-                payload = self.recv_payload(flags=zmq.NOBLOCK)
-            except zmq.ZMQError:
-                if payload is not None:
-                    break
-                else:
-                    time.sleep(SLEEP_FOR_TSUCHINOKO_TIME)
-                    if time.time() > self.last_targets_received + FORCE_KICKSTART_TIME:
-                        self.kickstart()
-        assert 'targets' in payload
-        self.last_targets_received = time.time()
-        return payload['targets']
-
-    def send_payload(self, payload: dict):
-        logger.info(f'message: {payload}')
-        self.socket.send(pickle.dumps(payload))
-
-    def recv_payload(self, flags=0) -> dict:
-        payload_response = pickle.loads(self.socket.recv(flags=flags))
-        logger.info(f'response: {payload_response}')
-        return payload_response
-
-
-class TsuchinokoAgent(TsuchinokoBase, Agent):
-    """
-    A Bluesky-Adaptive 'Agent'. This Agent communicates with Tsuchinoko over zmq to request new targets and report back
-    measurements. This is an abstract class that must be subclassed.
-
-    A `tsuchinoko.execution.bluesky_adaptive.BlueskyAdaptiveEngine` is required for the Tsuchinoko server to complement
-    one of these `TsuchinokoAgent`.
-    """
-
-    def tell(self, x, y, v) -> Dict[str, ArrayLike]:
-        super().tell(x, y, v)
-        return self.get_tell_document(x, y, v)
-
-    def ask(self, batch_size: int) -> Tuple[Sequence[Dict[str, ArrayLike]], Sequence[ArrayLike]]:
-        targets = super().ask(batch_size)
-        return self.get_ask_documents(targets), targets
-
-    @abstractmethod
-    def get_tell_document(self, x, y, v) -> Dict[str, ArrayLike]:
-        """
-        Return any single document corresponding to 'tell'-ing Tsuchinoko about the newly measured `x`, `y` data
-
-        Parameters
-        ----------
-        x :
-            Independent variable for data observed
-        y :
-            Dependent variable for data observed
-        v :
-            Variance for measurement of y
-
-        Returns
-        -------
-        dict
-            Dictionary to be unpacked or added to a document
-
-        """
-        ...
-
-    @abstractmethod
-    def get_ask_documents(self, targets: Sequence[ArrayLike]) -> Sequence[Dict[str, ArrayLike]]:
-        """
-        Ask the agent for a new batch of points to measure.
-
-        Parameters
-        ----------
-        targets : List[Tuple]
-            The new target positions to be measured received during this `ask`.
-
-        Returns
-        -------
-        docs : Sequence[dict]
-            Documents of key metadata from the ask approach for each point in next_points.
-            Must be length of batch size.
-
-        """
-        ...
-
-
-if __name__ == '__main__':
-    # NOTE: This usage is a primitive mocking of Bluesky-Adaptive's processes
-    agent = TsuchinokoBase()
-    while True:
-        targets = agent.ask(0)
-        agent.tell(targets[0], 1)