From af4b8720e57605c2004066b8aad947069024aa56 Mon Sep 17 00:00:00 2001 From: Jack Lashner Date: Tue, 20 Feb 2024 14:39:06 -0500 Subject: [PATCH] Supervisor actions (#628) * Add device_emulator reconnect on lost connection * HWP Supervisor Action system * docs, misc changes * Separate "gripper" and "driver" ibootbars * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Makes tcp reconnection optional * Docs updates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Addresses Brian's feedback --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/agents/hwp_supervisor_agent.rst | 110 ++++- socs/agents/hwp_supervisor/agent.py | 488 +++++++++++++------- socs/testing/device_emulator.py | 24 +- socs/testing/hwp_emulator.py | 9 +- tests/agents/hwp_pid/test_pid_controller.py | 2 +- 5 files changed, 463 insertions(+), 170 deletions(-) diff --git a/docs/agents/hwp_supervisor_agent.rst b/docs/agents/hwp_supervisor_agent.rst index bdd9cb119..519a03e70 100644 --- a/docs/agents/hwp_supervisor_agent.rst +++ b/docs/agents/hwp_supervisor_agent.rst @@ -67,17 +67,113 @@ This agent has two main purposes: - Serve as a host for high-level HWP operations that require coordinated control of various HWP subsystems, such as "Begin rotating at 2 Hz" -Right now only the first point is implemented, but operations can be added here -as we need them. - -HWP subsystems should implement a ``monitor_shutdown`` process that uses the -``get_op_data`` function to get the hwp-supervisor's session data, and check -the ``action`` field to determine if shutdown should be initiated. - +For the first point, the supervisor agent implements the ``monitor`` process, +which monitors HWP-related processes to compile full state info for the HWP, and uses +that to make a determination if the HWP should be shutdown. + +For high-level control, the HWP supervisor implements a state machine that +is used to perform complex operations with HWP agents that depend on the global +state of the HWP and related hardware. + +Control States and Actions +````````````````````````````` +In the context of the HWP supervisor state machine, a *Control State* is a python +dataclass that contains data to dictate what the supervisor will do on each +update call. For example, while in the ``WaitForTargetFreq`` state, the +supervisor will do nothing until the HWP frequency is within tolerance of a +specified freq for a specified period of time, at which point it will transition +into the Done state. + +A *Control Action* is a user-requested operation, in which a starting control +state is requested, that then transitions through any number of subsequent +states before completing. The action object contains its current state, state +history, completion status, and whether it considers itself successful. The +action is considered "complete" when it transitions into a "completion state", +which can be ``Done``, ``Error``, ``Abort``, or ``Idle``, at which point no more +state transitions will occur. In between update calls, a control action may be +aborted by the state-machine, where the action will transition into the +completed "Abort" state, and no further action will be taken. + +OCS agent operations are generally one-to-one with control actions, where each +operation begins a new action, and sleeps until that action is complete. +If an operation is started while there is already an action is in progress, the +current action will be aborted at the next opportunity and replaced with the new +requested action. The ``abort_action`` task can be used to abort the current +action without beginning a new one. + +Examples +``````````` +Below is an example client script that runs the PID to freq operation, and waits +until the target freq has been reached. + +.. code-block:: python + + supervisor = OCSClient("hwp-supervisor") + + result = supervisor.pid_to_freq(target_freq=2.0) + print(result.session['data']['action']) + + >> {'action_id': 6, + 'completed': True, + 'cur_state': {'class': 'Done', 'msg': None, 'success': True}, + 'state_history': [{'class': 'PIDToFreq', + 'direction': '0', + 'freq_tol': 0.05, + 'freq_tol_duration': 10.0, + 'target_freq': 2.0}, + {'class': 'WaitForTargetFreq', + 'freq_tol': 0.05, + 'freq_tol_duration': 10.0, + 'freq_within_tol_start': 1706829677.7613404, + 'target_freq': 2.0}, + {'class': 'Done', 'msg': None, 'success': True}], + 'success': True} + +Below is an example of a client script that starts to PID the HWP to 2 Hz, then +aborts the PID action, and shuts off the PMX power supply. Note that the +``abort_action`` here is technically redundant, since starting the new action +would abort the active action in the same manner. + +.. code-block:: python + + supervisor = OCSClient("hwp-supervisor") + + supervisor.pid_to_freq.start(target_freq=2.0) + supervisor.abort_action() + res1 = supervisor.pid_to_freq.wait() + res2 = supervisor.pmx_off() + + print("Result 1:") + print(res1.session['data']['action']) + print("Result 2: ") + print(res2.session['data']['action']) + + >> + Result 1: + {'action_id': 1, + 'completed': True, + 'cur_state': {'class': 'Abort'}, + 'state_history': [{'class': 'PIDToFreq', + 'direction': '0', + 'freq_tol': 0.05, + 'freq_tol_duration': 10.0, + 'target_freq': 2.0}, + {'class': 'Abort'}], + 'success': False} + Result 2: + {'action_id': 3, + 'completed': True, + 'cur_state': {'class': 'Done', 'msg': None, 'success': True}, + 'state_history': [{'class': 'PmxOff', 'success': True}, + {'class': 'Done', 'msg': None, 'success': True}], + 'success': True} Agent API ----------- +.. autoclass:: socs.agents.hwp_supervisor.agent.ControlAction + :members: + .. autoclass:: socs.agents.hwp_supervisor.agent.HWPSupervisor :members: diff --git a/socs/agents/hwp_supervisor/agent.py b/socs/agents/hwp_supervisor/agent.py index e845cab91..06e2bb98c 100644 --- a/socs/agents/hwp_supervisor/agent.py +++ b/socs/agents/hwp_supervisor/agent.py @@ -4,7 +4,7 @@ import time import traceback from dataclasses import asdict, dataclass, field -from typing import Optional +from typing import Dict, List, Optional import numpy as np import ocs @@ -14,6 +14,8 @@ from ocs.ocs_client import OCSClient, OCSReply from ocs.ocs_twisted import Pacemaker +client_cache = {} + def get_op_data(agent_id, op_name, log=None, test_mode=False): """ @@ -78,7 +80,12 @@ def get_op_data(agent_id, op_name, log=None, test_mode=False): if 'SITE_HTTP' in os.environ: args += [f"--site-http={os.environ['SITE_HTTP']}"] - client = site_config.get_control_client(agent_id, args=args) + if agent_id in client_cache: + client = client_cache[agent_id] + else: + client = site_config.get_control_client(agent_id, args=args) + client_cache[agent_id] = client + try: _, _, session = OCSReply(*client.request('status', op_name)) except client_http.ControlClientError as e: @@ -102,7 +109,32 @@ class HWPClients: pid: Optional[OCSClient] = None ups: Optional[OCSClient] = None lakeshore: Optional[OCSClient] = None - iboot: Optional[OCSClient] = None + gripper_iboot: Optional[OCSClient] = None + driver_iboot: Optional[OCSClient] = None + + +@dataclass +class IBootState: + instance_id: str + outlets: List[int] + outlet_state: Dict[int, Optional[int]] = None + op_data: Optional[Dict] = None + + def __post_init__(self): + self.outlet_state = {o: None for o in self.outlets} + self.outlet_labels = {o: f'outletStatus_{o}' for o in self.outlets} + + def update(self): + op = get_op_data(self.instance_id, 'acq', test_mode=False) + self.op_data = op + if op['status'] != 'ok': + self.outlet_state = {o: None for o in self.outlets} + return + + self.outlet_state = { + outlet: op['data'][label]['status'] + for outlet, label in self.outlet_labels.items() + } @dataclass @@ -135,10 +167,21 @@ class HWPState: last_quad: Optional[float] = None last_quad_time: Optional[float] = None - iboot_outlet1: Optional[int] = None - iboot_outlet2: Optional[int] = None - iboot_outlet1_state: Optional[int] = None - iboot_outlet2_state: Optional[int] = None + gripper_iboot: Optional[IBootState] = None + driver_iboot: Optional[IBootState] = None + + @classmethod + def from_args(cls, args: argparse.Namespace): + self = cls( + temp_field=args.ybco_temp_field, + temp_thresh=args.ybco_temp_thresh, + ups_minutes_remaining_thresh=args.ups_minutes_remaining_thresh, + ) + if args.gripper_iboot_id is not None: + self.gripper_iboot = IBootState(args.gripper_iboot_id, args.gripper_iboot_outlets) + if args.driver_iboot_id is not None: + self.driver_iboot = IBootState(args.driver_iboot_id, args.driver_iboot_outlets) + return self def _update_from_keymap(self, op, keymap): if op['status'] != 'ok': @@ -265,29 +308,6 @@ def update_ups_state(self, op): self.ups_last_connection_attempt = data['ups_connection']['last_attempt'] self.ups_connected = data['ups_connection']['connected'] - def update_iboot_state(self, op): - """ - Updates state values from the IBoot acq operation results. - - Args - ----- - op : dict - Dict containing the operations (from get_op_data) from the IBoot - ``acq`` process - """ - iboot_keymap = { - 'iboot_outlet1_state': (f'outletStatus_{self.iboot_outlet1}', 'status'), - 'iboot_outlet2_state': (f'outletStatus_{self.iboot_outlet2}', 'status'), - } - - if op['status'] != 'ok': - for k in iboot_keymap: - setattr(self, k, None) - return - - for k, f in iboot_keymap.items(): - setattr(self, k, op['data'][f[0]][f[1]]) - @property def pmx_action(self): """ @@ -337,16 +357,24 @@ def gripper_action(self): class ControlState: - """ - Namespace for HWP control state definitions - """ + """Namespace for HWP control state definitions""" @dataclass - class Idle: + class Base: + def __init__(self): + super().__init__() + self.start_time = time.time() + + def encode(self): + d = {'class': self.__class__.__name__} + d.update(asdict(self)) + return d + + @dataclass + class Idle(Base): """Does nothing""" - start_time: float = field(default_factory=time.time) @dataclass - class PIDToFreq: + class PIDToFreq(Base): """ Configures PID and PMX agents to PID to a target frequency. @@ -361,17 +389,14 @@ class PIDToFreq: to consider the target frequency reached. freq_tol_duration : float Duration in seconds that the frequency must be within the tolerance - start_time : float - Time that the state was entered """ target_freq: float direction: str freq_tol: float freq_tol_duration: float - start_time: float = field(default_factory=time.time) @dataclass - class WaitForTargetFreq: + class WaitForTargetFreq(Base): """ Wait until HWP reaches its target frequency before transitioning to the Done state. @@ -394,10 +419,9 @@ class WaitForTargetFreq: freq_tol: float freq_tol_duration: float freq_within_tol_start: Optional[float] = None - start_time: float = field(default_factory=time.time) @dataclass - class ConstVolt: + class ConstVolt(Base): """ Configure PMX agent to output a constant voltage. @@ -410,10 +434,9 @@ class ConstVolt: """ voltage: float direction: str - start_time: float = field(default_factory=time.time) @dataclass - class Done: + class Done(Base): """ Signals the last state has completed @@ -428,10 +451,9 @@ class Done: """ success: bool msg: str = None - start_time: float = field(default_factory=time.time) @dataclass - class Error: + class Error(Base): """ Signals the last state update threw an error @@ -446,16 +468,15 @@ class Error: start_time: float = field(default_factory=time.time) @dataclass - class Brake: + class Brake(Base): """ Configure the PID and PMX agents to actively brake the HWP """ freq_tol: float freq_tol_duration: float - start_time: float = field(default_factory=time.time) @dataclass - class WaitForBrake: + class WaitForBrake(Base): """ Waits until the HWP has slowed before shutting off PMX @@ -468,10 +489,9 @@ class WaitForBrake: min_freq: float init_quad: float prev_freq: float = None - start_time: float = field(default_factory=time.time) @dataclass - class PmxOff: + class PmxOff(Base): """ Turns off the PMX @@ -480,14 +500,82 @@ class PmxOff: start_time : float Time that the state was entered """ - freq_tol: float = 0.05 - freq_tol_duration: float = 20 - start_time: float = field(default_factory=time.time) + success: bool = True + + @dataclass + class Abort(Base): + """Abort current action""" + pass + + completed_states = (Done, Error, Abort, Idle) + + +class ControlAction: + """ + This is a class to contain data regarding a single HWP control action. + This groups together states that are part of a single action, and whether + the action is completed and successful. + """ + _cur_action_id: int = 0 + _id_lock = threading.Lock() + + def __init__(self, state: ControlState.Base): + with ControlAction._id_lock: + self.action_id = ControlAction._cur_action_id + ControlAction._cur_action_id += 1 + + self.completed = False + self.success = False + self.state_history = [] + self.set_state(state) + + def set_state(self, state: ControlState.Base): + """ + Sets state for the current action. If this is a `completed_state`, + will mark as complete. + """ + self.state_history.append(state) + self.cur_state = state + if isinstance(state, ControlState.completed_states): + self.completed = True + if isinstance(state, ControlState.Done): + self.success = state.success + + def encode(self): + """Encodes this as a dict""" + return dict( + action_id=self.action_id, + completed=self.completed, + success=self.success, + cur_state=self.cur_state.encode(), + state_history=[s.encode() for s in self.state_history], + ) + + def sleep_until_complete(self, session=None, dt=1): + """ + Sleeps until the action is complete. + + Args + ----- + session: OpSession, optional + If specified, this will set `session.data['action']` to the encoded + action on each iteration to keep session data up to date + dt: float, optional + Time to sleep between iterations. + """ + while True: + if session is not None: + session.data.update({'action': self.encode()}) + if self.completed: + return + time.sleep(dt) class ControlStateMachine: def __init__(self): - self.state = ControlState.Idle() + self.action: ControlAction = ControlAction(ControlState.Idle()) + self.action_history: List[ControlAction] = [] + self.max_action_history_count = 100 self.log = txaio.make_logger() # pylint: disable=E1101 self.lock = threading.Lock() @@ -532,87 +620,80 @@ def run_and_validate(self, op, kwargs=None, timeout=10, log=None): return session - def _set_state(self, state): - self.log.info("Changing from {s1} to {s2}", s1=self.state, s2=state) - self.state = state - def update(self, clients, hwp_state): """Run the next series of actions for the current state""" try: self.lock.acquire() + state = self.action.cur_state - if isinstance(self.state, ControlState.PIDToFreq): + if isinstance(state, ControlState.PIDToFreq): self.run_and_validate(clients.pid.set_direction, - kwargs={'direction': self.state.direction}) + kwargs={'direction': state.direction}) self.run_and_validate(clients.pid.declare_freq, - kwargs={'freq': self.state.target_freq}) + kwargs={'freq': state.target_freq}) self.run_and_validate(clients.pmx.use_ext) self.run_and_validate(clients.pmx.set_on) self.run_and_validate(clients.pid.tune_freq) - # kwargs={'freq': self.state.target_freq}) + # kwargs={'freq': state.target_freq}) - self._set_state(ControlState.WaitForTargetFreq( - target_freq=self.state.target_freq, - freq_tol=self.state.freq_tol, - freq_tol_duration=self.state.freq_tol_duration + self.action.set_state(ControlState.WaitForTargetFreq( + target_freq=state.target_freq, + freq_tol=state.freq_tol, + freq_tol_duration=state.freq_tol_duration )) - elif isinstance(self.state, ControlState.WaitForTargetFreq): + elif isinstance(state, ControlState.WaitForTargetFreq): # Check if we are close enough to the target frequency. # This will make sure we remain within the frequency threshold for # ``self.freq_tol_duration`` seconds before switching to DONE f = hwp_state.pid_current_freq if f is None: - self.state.freq_within_tol_start = None + state.freq_within_tol_start = None return - if np.abs(f - self.state.target_freq) > self.state.freq_tol: - self.state.freq_within_tol_start = None + if np.abs(f - state.target_freq) > state.freq_tol: + state.freq_within_tol_start = None return # If within tolerance for freq_tol_duration, switch to Done - if self.state.freq_within_tol_start is None: - self.state.freq_within_tol_start = time.time() + if state.freq_within_tol_start is None: + state.freq_within_tol_start = time.time() - time_within_tol = time.time() - self.state.freq_within_tol_start - if time_within_tol > self.state.freq_tol_duration: - self._set_state(ControlState.Done(success=True)) + time_within_tol = time.time() - state.freq_within_tol_start + if time_within_tol > state.freq_tol_duration: + self.action.set_state(ControlState.Done(success=True)) - elif isinstance(self.state, ControlState.ConstVolt): + elif isinstance(state, ControlState.ConstVolt): self.run_and_validate(clients.pmx.set_on) self.run_and_validate(clients.pid.set_direction, - kwargs={'direction': self.state.direction}) + kwargs={'direction': state.direction}) self.run_and_validate(clients.pmx.ign_ext) self.run_and_validate(clients.pmx.set_v, - kwargs={'volt': self.state.voltage}) - self._set_state(ControlState.Done(success=True)) + kwargs={'volt': state.voltage}) + self.action.set_state(ControlState.Done(success=True)) - elif isinstance(self.state, ControlState.PmxOff): + elif isinstance(state, ControlState.PmxOff): self.run_and_validate(clients.pmx.set_off) self.run_and_validate(clients.pid.declare_freq, kwargs={'freq': 0}) self.run_and_validate(clients.pid.tune_freq) - self._set_state(ControlState.WaitForTargetFreq( - target_freq=0, - freq_tol=self.state.freq_tol, - freq_tol_duration=self.state.freq_tol_duration, - )) + self.action.set_state(ControlState.Done(success=state.success)) - elif isinstance(self.state, ControlState.Brake): + elif isinstance(state, ControlState.Brake): init_quad = hwp_state.last_quad init_quad_time = hwp_state.last_quad_time if init_quad is None or init_quad_time is None: self.log.warn("Could not determine direction from Encoder agent") self.log.warn("Setting PMX Off") - self._set_state(ControlState.PmxOff()) + self.action.set_state(ControlState.PmxOff(success=False)) return quad_last_updated = time.time() - init_quad_time if quad_last_updated > 10.0: self.log.warn(f"Quad has not been updated in last {quad_last_updated} sec") self.log.warn("Setting PMX Off, since can't confirm direction") - self._set_state(ControlState.PmxOff()) + self.action.set_state(ControlState.PmxOff(success=False)) return self.run_and_validate(clients.pmx.ign_ext) @@ -629,13 +710,13 @@ def update(self, clients, hwp_state): self.run_and_validate(clients.pid.set_direction, kwargs=dict(direction=new_d)) - self._set_state(ControlState.WaitForBrake( + self.action.set_state(ControlState.WaitForBrake( init_quad=init_quad, min_freq=0.5, prev_freq=hwp_state.enc_freq )) - elif isinstance(self.state, ControlState.WaitForBrake): + elif isinstance(state, ControlState.WaitForBrake): quad = hwp_state.last_quad quad_time = hwp_state.last_quad_time freq = hwp_state.enc_freq @@ -643,44 +724,53 @@ def update(self, clients, hwp_state): if quad is None or quad_time is None: self.log.warn("Could not determine direction from Encoder agent") self.log.warn("Setting PMX Off") - self._set_state(ControlState.PmxOff()) + self.action.set_state(ControlState.PmxOff()) return quad_last_updated = time.time() - quad_time if quad_last_updated > 10.0: self.log.warn(f"Quad has not been updated in last {quad_last_updated} sec") self.log.warn("Setting PMX Off, since can't confirm direction") - self._set_state(ControlState.PmxOff()) + self.action.set_state(ControlState.PmxOff()) return - if freq - self.state.prev_freq > 0: + if freq - state.prev_freq > 0: self.log.warn("HWP Freq is increasing! Setting PMX Off") - self._set_state(ControlState.PmxOff()) + self.action.set_state(ControlState.PmxOff()) return - quad_diff = np.abs(quad - self.state.init_quad) - if freq < self.state.min_freq or quad_diff > 0.1: + quad_diff = np.abs(quad - state.init_quad) + if freq < state.min_freq or quad_diff > 0.1: self.run_and_validate(clients.pmx.set_off) - self._set_state(ControlState.WaitForTargetFreq( + self.action.set_state(ControlState.WaitForTargetFreq( target_freq=0, freq_tol=0.1, freq_tol_duration=10, )) - self.prev_freq = freq + state.prev_freq = freq return except Exception: tb = traceback.format_exc() self.log.error("Error updating state:\n{tb}", tb=tb) - self._set_state(ControlState.Error(traceback=tb)) + self.action.set_state(ControlState.Error(traceback=tb)) finally: self.lock.release() - def request_state(self, state): + def request_new_action(self, state: ControlState.Base): + """ + Requests that a new action is started with a given state. + If an action is already in progress, it will be aborted. + """ with self.lock: - self._set_state(state) - return True + if not self.action.completed: + self.action.set_state(ControlState.Abort()) + if len(self.action_history) > self.max_action_history_count: + self.action_history.pop(0) + self.action = ControlAction(state) + self.action_history.append(self.action) + return self.action class HWPSupervisor: @@ -722,17 +812,9 @@ def __init__(self, agent, args): self.hwp_pid_id = args.hwp_pid_id self.ups_id = args.ups_id - self.iboot_id = args.iboot_id - self.iboot_outlet1 = args.iboot_outlet1 - 1 - self.iboot_outlet2 = args.iboot_outlet2 - 1 - - self.hwp_state = HWPState( - temp_field=self.ybco_temp_field, - temp_thresh=args.ybco_temp_thresh, - ups_minutes_remaining_thresh=args.ups_minutes_remaining_thresh, - iboot_outlet1=self.iboot_outlet1, - iboot_outlet2=self.iboot_outlet2, - ) + self.hwp_state = HWPState.from_args(args) + self.gripper_iboot_id = args.gripper_iboot_id + self.driver_iboot_id = args.driver_iboot_id self.control_state_machine = ControlStateMachine() self.forward_is_cw = args.forward_dir == 'cw' @@ -755,7 +837,8 @@ def get_client(id): pid=get_client(self.hwp_pid_id), ups=get_client(self.ups_id), lakeshore=get_client(self.ybco_lakeshore_id), - iboot=get_client(self.iboot_id), + gripper_iboot=get_client(self.gripper_iboot_id), + driver_iboot=get_client(self.driver_iboot_id), ) @ocs_agent.param('test_mode', type=bool, default=False) @@ -832,9 +915,8 @@ def monitor(self, session, params): temp_op = get_op_data(self.ybco_lakeshore_id, 'acq', **kw) enc_op = get_op_data(self.hwp_encoder_id, 'acq', **kw) pmx_op = get_op_data(self.hwp_pmx_id, 'acq', **kw) - pid_op = get_op_data(self.hwp_pid_id, 'acq', **kw) + pid_op = get_op_data(self.hwp_pid_id, 'main', **kw) ups_op = get_op_data(self.ups_id, 'acq', **kw) - iboot_op = get_op_data(self.iboot_id, 'acq', **kw) session.data['monitored_sessions'] = { 'temperature': temp_op, @@ -842,7 +924,6 @@ def monitor(self, session, params): 'pmx': pmx_op, 'pid': pid_op, 'ups': ups_op, - 'iboot': iboot_op } # gather state info @@ -851,7 +932,12 @@ def monitor(self, session, params): self.hwp_state.update_temp_state(temp_op) self.hwp_state.update_ups_state(ups_op) self.hwp_state.update_enc_state(enc_op) - self.hwp_state.update_iboot_state(iboot_op) + + if self.hwp_state.driver_iboot is not None: + self.hwp_state.driver_iboot.update() + if self.hwp_state.gripper_iboot is not None: + self.hwp_state.gripper_iboot.update() + session.data['hwp_state'] = asdict(self.hwp_state) # Get actions for each hwp subsystem @@ -877,22 +963,31 @@ def spin_control(self, session, params): **Process** - Process to manage the spin-state for HWP agents. This will issue commands to various HWP agents depending on the current control state. + + Notes + -------- + + Example of ``session.data``:: + + >>> session['data'] + { + 'current_action': , + 'action_history': List[Encoded action] + 'timestammp': + } """ clients = self._get_hwp_clients() session.set_status('running') while session.status in ['starting', 'running']: self.control_state_machine.update(clients, self.hwp_state) - - s = self.control_state_machine.state - state_dict = asdict(s) - state_dict['state_name'] = s.__class__.__name__ - session.data = { - 'state': state_dict, - 'timestamp': time.time(), + 'current_action': self.control_state_machine.action.encode(), + 'action_history': [a.encode() for a in self.control_state_machine.action_history], + 'timestamp': time.time() } time.sleep(1) + return True, "Finished spin control process" def _stop_spin_control(self, session, params): session.status = 'stopping' @@ -917,6 +1012,20 @@ def pid_to_freq(self, session, params): freq_thresh_duration : float Duration (seconds) for which the HWP must be within ``freq_thresh`` of the ``target_freq`` to be considered successful. + + Notes + -------- + + Example of ``session.data``:: + + >>> session['data'] + {'action': + {'action_id': 3, + 'completed': True, + 'cur_state': {'class': 'Done', 'msg': None, 'success': True}, + 'state_history': List[ConrolState], + 'success': True} + } """ if params['target_freq'] >= 0: d = '0' if self.forward_is_cw else '1' @@ -929,11 +1038,9 @@ def pid_to_freq(self, session, params): freq_tol_duration=params['freq_tol_duration'], direction=d ) - success = self.control_state_machine.request_state(state) - if success: - return True, f"Set state to {state}" - else: - return False, "Failed to update state" + action = self.control_state_machine.request_new_action(state) + action.sleep_until_complete(session=session) + return action.success, f"Completed with state: {action.cur_state}" @ocs_agent.param('voltage', type=float) @ocs_agent.param('direction', type=str, choices=['cw', 'ccw'], default='cw') @@ -950,6 +1057,20 @@ def set_const_voltage(self, session, params): Direction of the HWP. Must be one of ``cw`` or ``ccw``, corresponding to the clockwise and counter-clockwise directions of the HWP, as seen when looking at the cryostat from the sky. + + Notes + -------- + + Example of ``session.data``:: + + >>> session['data'] + {'action': + {'action_id': 3, + 'completed': True, + 'cur_state': {'class': 'Done', 'msg': None, 'success': True}, + 'state_history': List[ConrolState], + 'success': True} + } """ if params['direction'] == 'cw': d = '0' if self.forward_is_cw else '1' @@ -959,11 +1080,9 @@ def set_const_voltage(self, session, params): voltage=params['voltage'], direction=d ) - success = self.control_state_machine.request_state(state) - if success: - return True, f"Set state to {state}" - else: - return False, "Failed to update state" + action = self.control_state_machine.request_new_action(state) + action.sleep_until_complete(session=session) + return action.success, f"Completed with state: {action.cur_state}" @ocs_agent.param('freq_tol', type=float, default=0.05) @ocs_agent.param('freq_tol_duration', type=float, default=10) @@ -979,34 +1098,76 @@ def brake(self, session, params): freq_thresh_duration : float Duration (seconds) for which the HWP must be within ``freq_thresh`` of the ``target_freq`` to be considered successful. + + Notes + -------- + + Example of ``session.data``:: + + >>> session['data'] + {'action': + {'action_id': 3, + 'completed': True, + 'cur_state': {'class': 'Done', 'msg': None, 'success': True}, + 'state_history': List[ConrolState], + 'success': True} + } """ state = ControlState.Brake( freq_tol=params['freq_tol'], freq_tol_duration=params['freq_tol_duration'] ) - success = self.control_state_machine.request_state(state) - if success: - return True, f"Set state to {state}" - else: - return False, "Failed to update state" + action = self.control_state_machine.request_new_action(state) + action.sleep_until_complete(session=session) + return action.success, f"Completed with state: {action.cur_state}" - @ocs_agent.param('freq_tol', type=float, default=None) - @ocs_agent.param('freq_tol_duration', type=float, default=None) def pmx_off(self, session, params): """pmx_off() **Task** - Sets the control state to turn off the PMX. + + Notes + -------- + + Example of ``session.data``:: + + >>> session['data'] + {'action': + {'action_id': 3, + 'completed': True, + 'cur_state': {'class': 'Done', 'msg': None, 'success': True}, + 'state_history': List[ConrolState], + 'success': True} + } """ - kw = {} - for p in ['freq_tol', 'freq_tol_duration']: - if params[p] is not None: - kw[p] = params[p] - state = ControlState.PmxOff(**kw) - success = self.control_state_machine.request_state(state) - if success: - return True, f"Set state to {state}" - else: - return False, "Failed to update state" + state = ControlState.PmxOff() + action = self.control_state_machine.request_new_action(state) + action.sleep_until_complete(session=session) + return action.success, f"Completed with state: {action.cur_state}" + + def abort_action(self, session, params): + """abort_action() + + **Task** - Aborts the current action, setting the control state to Idle + + Notes + -------- + + Example of ``session.data``:: + + >>> session['data'] + {'action': + {'action_id': 3, + 'completed': True, + 'cur_state': {'class': 'Idle'}, + 'state_history': List[ConrolState], + 'success': False} + } + """ + state = ControlState.Idle() + action = self.control_state_machine.request_new_action(state) + session.data['action'] = action.encode() + return True, "Set state to idle" def make_parser(parser=None): @@ -1031,11 +1192,21 @@ def make_parser(parser=None): pgroup.add_argument('--ups-minutes-remaining-thresh', type=float, help="Threshold for UPS minutes remaining before a " "shutdown is triggered") - pgroup.add_argument('--iboot-id', help="Instance ID for IBoot-PDU agent") - pgroup.add_argument('--iboot-outlet1', type=int, default=1, - help="IBoot-PDU outlet connected to gripper drive") - pgroup.add_argument('--iboot-outlet2', type=int, default=2, - help="IBoot-PDU outlet connected to gripper control") + + pgroup.add_argument( + '--driver-iboot-id', + help="Instance ID for IBoot-PDU agent that powers the HWP Driver board") + pgroup.add_argument( + '--driver-iboot-outlets', nargs='+', type=int, + help="Outlets for driver iboot power") + + pgroup.add_argument( + '--gripper-iboot-id', + help="Instance ID for IBoot-PDU agent that powers the gripper controller") + pgroup.add_argument( + '--gripper-iboot-outlets', nargs='+', type=int, + help="Outlets for gripper iboot power") + pgroup.add_argument('--forward-dir', choices=['cw', 'ccw'], default="cw", help="Whether the PID 'forward' direction is cw or ccw") return parser @@ -1058,6 +1229,7 @@ def main(args=None): agent.register_task('set_const_voltage', hwp.set_const_voltage) agent.register_task('brake', hwp.brake) agent.register_task('pmx_off', hwp.pmx_off) + agent.register_task('abort_action', hwp.abort_action) runner.run(agent, auto_reconnect=True) diff --git a/socs/testing/device_emulator.py b/socs/testing/device_emulator.py index 0999d3672..88e933db6 100644 --- a/socs/testing/device_emulator.py +++ b/socs/testing/device_emulator.py @@ -12,7 +12,8 @@ import serial -def create_device_emulator(responses, relay_type, port=9001, encoding='utf-8'): +def create_device_emulator(responses, relay_type, port=9001, encoding='utf-8', + reconnect=False): """Create a device emulator fixture. This provides a device emulator that can be used to mock a device during @@ -27,6 +28,8 @@ def create_device_emulator(responses, relay_type, port=9001, encoding='utf-8'): encoding (str): Encoding for the messages and responses. See :func:`socs.testing.device_emulator.DeviceEmulator` for more details. + reconnect (bool): If True, on TCP client disconnect, the emulator will + listen for new incoming connections instead of quitting Returns: function: @@ -40,7 +43,7 @@ def create_device_emulator(responses, relay_type, port=9001, encoding='utf-8'): @pytest.fixture() def create_device(): - device = DeviceEmulator(responses, encoding) + device = DeviceEmulator(responses, encoding, reconnect=reconnect) if relay_type == 'serial': device.create_serial_relay() @@ -66,6 +69,8 @@ class DeviceEmulator: given encoding. No encoding is used if set to None. That can be useful if you need to use raw data from your hardware. Defaults to 'utf-8'. + reconnect (bool): If True, on TCP client disconnect, the emulator will + listen for new incoming connections instead of quitting Attributes: responses (dict): Current set of responses the DeviceEmulator would @@ -83,10 +88,11 @@ class DeviceEmulator: """ - def __init__(self, responses, encoding='utf-8'): + def __init__(self, responses, encoding='utf-8', reconnect=False): self.responses = deepcopy(responses) self.default_response = None self.encoding = encoding + self.reconnect = reconnect self._type = None self._read = True self._conn = None @@ -254,6 +260,18 @@ def _read_socket(self, port): while self._read: try: msg = self._conn.recv(4096) + if not msg: + self.logger.info("Client disconnected") + if self.reconnect: + self.logger.info("Waiting for new connection") + # attempt to reconnect + self._conn, client_address = self._sock.accept() + self.logger.info(f"Client connection made from {client_address}") + continue + else: + self.logger.info("Shutting down") + break + # Was seeing this on tests in the cryomech agent except ConnectionResetError: self.logger.info('Caught connection reset on Agent clean up') diff --git a/socs/testing/hwp_emulator.py b/socs/testing/hwp_emulator.py index 9a942ec96..af975e6dc 100644 --- a/socs/testing/hwp_emulator.py +++ b/socs/testing/hwp_emulator.py @@ -1,6 +1,7 @@ """ HWP Emulation module """ +import argparse import logging import threading import time @@ -105,6 +106,7 @@ def update_loop(self): with s.lock: if s.pmx.source == "volt": s.cur_freq = lerp(s.cur_freq, s.pid.freq_setpoint, 0.3) + time.sleep(1) def process_pmx_msg(self, data): """Process messages for PMX emulator""" @@ -214,7 +216,12 @@ def create_emulator(): if __name__ == "__main__": - hwp_em = HWPEmulator() + parser = argparse.ArgumentParser() + parser.add_argument("--pid-port", type=int, default=5025) + parser.add_argument("--pmx-port", type=int, default=5026) + args = parser.parse_args() + + hwp_em = HWPEmulator(pid_port=args.pid_port, pmx_port=args.pmx_port) try: hwp_em.start() while True: diff --git a/tests/agents/hwp_pid/test_pid_controller.py b/tests/agents/hwp_pid/test_pid_controller.py index 5dd5e36ab..3e2f6816c 100644 --- a/tests/agents/hwp_pid/test_pid_controller.py +++ b/tests/agents/hwp_pid/test_pid_controller.py @@ -2,7 +2,7 @@ from socs.testing.device_emulator import create_device_emulator pid_emu = create_device_emulator( - {'*W02400000': 'W02\r'}, relay_type='tcp', port=3003) + {'*W02400000': 'W02\r'}, relay_type='tcp', port=3003, reconnect=False) def test_send_message(pid_emu):