Improve HWP agent robustness (#541)

* minor bugfixes to the hwp-gripper agent

* Modified hwp-pid and hwp-pmx agents to handle connection interruptions

* [pre-commit.ci] pre-commit autoupdate (#540)

updates:
- [github.com/pre-commit/pre-commit-hooks: v4.4.0 → v4.5.0](pre-commit/pre-commit-hooks@v4.4.0...v4.5.0)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Changed connection protocal from an infinite loop into a set number of iterations. Added reconnection functionality into the agent acq process

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fixes to pre-commit failures

* Added read=False to set_current_limit meathod

---------

Co-authored-by: Bryce Bixler <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

socs/agents/hwp_gripper/agent.py

@@ -381,8 +381,8 @@ def shutdown(self, session, params=None):
         self.shutdown_mode = True
         return True, 'Shutdown completed'
 
-    def grip_hwp(self, session, params=None):
-        """grip_hwp()
+    def grip(self, session, params=None):
+        """grip()
 
         **Task** - Series of commands to automatically grip the HWP.
         This will return grippers to their home position, then move them each
@@ -442,7 +442,8 @@ def run_and_append(func, *args, **kwargs):
 
                 # Reset alarms. If the warm-limit is hit, the alarm will be triggered
                 # and return_dict['result'] will be True
-                return_dict = run_and_append(self.client.reset, job='grip', check_shutdown=check_shutdown)
+                return_dict = run_and_append(self.client.reset, job='grip',
+                                             check_shutdown=check_shutdown)
 
                 if return_dict['result']:
                     # If the warm-limit is hit, move the actuator outwards bit
@@ -676,9 +677,7 @@ def main(args=None):
                                   args=args)
 
     agent, runner = ocs_agent.init_site_agent(args)
-    gripper_agent = HWPGripperAgent(agent, mcu_ip=args.mcu_ip,
-                                    control_port=args.control_port,
-                                    supervisor_id=args.supervisor_id)
+    gripper_agent = HWPGripperAgent(agent, args)
     agent.register_task('init_connection', gripper_agent.init_connection,
                         startup=True)
     agent.register_process('monitor_state', gripper_agent.monitor_state,

socs/agents/hwp_pid/agent.py

@@ -295,13 +295,17 @@ def acq(self, session, params):
                 data = {'timestamp': time.time(),
                         'block_name': 'HWPPID', 'data': {}}
 
-                current_freq = self.pid.get_freq()
-                target_freq = self.pid.get_target()
-                direction = self.pid.get_direction()
-
-                data['data']['current_freq'] = current_freq
-                data['data']['target_freq'] = target_freq
-                data['data']['direction'] = direction
+                try:
+                    current_freq = self.pid.get_freq()
+                    target_freq = self.pid.get_target()
+                    direction = self.pid.get_direction()
+
+                    data['data']['current_freq'] = current_freq
+                    data['data']['target_freq'] = target_freq
+                    data['data']['direction'] = direction
+                except BaseException:
+                    time.sleep(1)
+                    continue
 
                 self.agent.publish_to_feed('hwppid', data)
 

socs/agents/hwp_pid/drivers/pid_controller.py

@@ -22,15 +22,17 @@ class PID:
 
     def __init__(self, ip, port, verb=False):
         self.verb = verb
+        self.ip = ip
+        self.port = port
         self.hex_freq = '00000'
         self.direction = None
         self.target = 0
         # Need to setup connection before setting direction
-        self.conn = self._establish_connection(ip, int(port))
+        self.conn = self._establish_connection(self.ip, int(self.port))
         self.set_direction('0')
 
     @staticmethod
-    def _establish_connection(ip, port, timeout=5):
+    def _establish_connection(ip, port, timeout=2):
         """Connect to PID controller.
 
         Args:
@@ -45,18 +47,17 @@ def _establish_connection(ip, port, timeout=5):
 
         """
         conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        conn.settimeout(timeout)
         # unit tests might fail on first connection attempt
         attempts = 3
         for attempt in range(attempts):
             try:
                 conn.connect((ip, port))
                 break
-            except ConnectionRefusedError:
+            except (ConnectionRefusedError, OSError):
                 print(f"Failed to connect to device at {ip}:{port}")
-                print(f"Connection attempts remaining: {attempts-attempt-1}")
-            time.sleep(1)
-        conn.settimeout(timeout)
-
+        else:
+            raise RuntimeError('Could not connect to PID controller')
         return conn
 
     @staticmethod
@@ -259,20 +260,21 @@ def send_message(self, msg):
             str: Respnose from the controller.
 
         """
-        self.conn.sendall((msg + '\r\n').encode())
-        time.sleep(0.5)  # Don't send messages too quickly
         for attempt in range(2):
             try:
+                self.conn.sendall((msg + '\r\n').encode())
+                time.sleep(0.5)  # Don't send messages too quickly
                 data = self.conn.recv(4096).decode().strip()
-                break
-            except socket.timeout:
+                return data
+            except (socket.timeout, OSError):
                 print("Caught timeout waiting for response from PID controller. "
                       + "Trying again...")
                 time.sleep(1)
                 if attempt == 1:
-                    raise RuntimeError(
-                        'Response from PID controller timed out.')
-        return data
+                    print("Resetting connection")
+                    self.conn.close()
+                    self.conn = self._establish_connection(self.ip, int(self.port))
+                    return self.send_message(msg)
 
     def return_messages(self, msg):
         """Decode list of responses from PID controller and return useful

socs/agents/hwp_pmx/agent.py

@@ -281,26 +281,31 @@ def acq(self, session, params):
                     'block_name': 'hwppmx',
                     'data': {}
                 }
-                msg, curr = self.dev.meas_current()
-                data['data']['current'] = curr
 
-                msg, volt = self.dev.meas_voltage()
-                data['data']['voltage'] = volt
+                try:
+                    msg, curr = self.dev.meas_current()
+                    data['data']['current'] = curr
 
-                msg, code = self.dev.check_error()
-                data['data']['err_code'] = code
-                data['data']['err_msg'] = msg
+                    msg, volt = self.dev.meas_voltage()
+                    data['data']['voltage'] = volt
 
-                prot_code = self.dev.check_prot()
-                if prot_code != 0:
-                    self.prot = prot_code
+                    msg, code = self.dev.check_error()
+                    data['data']['err_code'] = code
+                    data['data']['err_msg'] = msg
 
-                prot_msg = self.dev.get_prot_msg(self.prot)
-                data['data']['prot_code'] = self.prot
-                data['data']['prot_msg'] = prot_msg
+                    prot_code = self.dev.check_prot()
+                    if prot_code != 0:
+                        self.prot = prot_code
 
-                msg, src = self.dev.check_source()
-                data['data']['source'] = src
+                    prot_msg = self.dev.get_prot_msg(self.prot)
+                    data['data']['prot_code'] = self.prot
+                    data['data']['prot_msg'] = prot_msg
+
+                    msg, src = self.dev.check_source()
+                    data['data']['source'] = src
+                except BaseException:
+                    time.sleep(sleep_time)
+                    continue
 
                 self.agent.publish_to_feed('hwppmx', data)
                 session.data = {'curr': curr,

socs/agents/hwp_pmx/drivers/PMX_ethernet.py

@@ -1,5 +1,5 @@
+import socket
 import time
-from socket import AF_INET, SOCK_STREAM, socket
 
 protection_status_key = [
     'Over voltage',
@@ -21,26 +21,53 @@ class PMX:
     """
 
     def __init__(self, ip, port):
-        self.sock = socket(AF_INET, SOCK_STREAM)
-        self.sock.connect((ip, port))
-        self.sock.settimeout(5)
-
+        self.ip = ip
+        self.port = port
         self.wait_time = 0.01
         self.buffer_size = 128
+        self.conn = self._establish_connection(self.ip, int(self.port))
+
+    def _establish_connection(self, ip, port, timeout=2):
+        conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        conn.settimeout(timeout)
+        attempts = 3
+        for attempt in range(attempts):
+            try:
+                conn.connect((ip, port))
+                break
+            except (ConnectionRefusedError, OSError):
+                print(f"Failed to connect to device at {ip}:{port}")
+        else:
+            raise RuntimeError('Could not connect to PID controller')
+        return conn
 
     def close(self):
-        self.sock.close()
-
-    def read(self):
-        return self.sock.recv(self.buffer_size).decode('utf-8')
+        self.conn.close()
+
+    def send_message(self, msg, read=True):
+        for attempt in range(2):
+            try:
+                self.conn.sendall(msg)
+                time.sleep(0.5)
+                if read:
+                    data = self.conn.recv(self.buffer_size).decode('utf-8')
+                    return data
+                return
+            except (socket.timeout, OSError):
+                print("Caught timeout waiting for responce from PMX. Trying again...")
+                time.sleep(1)
+                if attempt == 1:
+                    print("Resetting connection")
+                    self.conn.close()
+                    self.conn = self._establish_connection(self.ip, int(self.port))
+                    return self.send_message(msg, read=read)
 
     def wait(self):
         time.sleep(self.wait_time)
 
     def check_output(self):
         """ Return the output status """
-        self.sock.sendall(b'output?\n')
-        val = int(self.read())
+        val = int(self.send_message(b'output?\n'))
         msg = "Measured output state = "
         states = {0: 'OFF', 1: 'ON'}
         if val in states:
@@ -51,103 +78,95 @@ def check_output(self):
 
     def check_error(self):
         """ Check oldest error from error queues. Error queues store up to 255 errors """
-        self.sock.sendall(b':system:error?\n')
-        val = self.read()
+        val = self.send_message(b':system:error?\n')
         code, msg = val.split(',')
         code = int(code)
         msg = msg[1:-2]
         return msg, code
 
     def clear_alarm(self):
         """ Clear alarm """
-        self.sock.sendall(b'output:protection:clear\n')
+        self.send_message(b'output:protection:clear\n', read=False)
 
     def turn_on(self):
         """ Turn the PMX on """
-        self.sock.sendall(b'output 1\n')
+        self.send_message(b'output 1\n', read=False)
         self.wait()
         return self.check_output()
 
     def turn_off(self):
         """ Turn the PMX off """
-        self.sock.sendall(b'output 0\n')
+        self.send_message(b'output 0\n', read=False)
         self.wait()
         return self.check_output()
 
     def check_current(self):
         """ Check the current setting """
-        self.sock.sendall(b'curr?\n')
-        val = float(self.read())
+        val = float(self.send_message(b'curr?\n'))
         msg = "Current setting = {:.3f} A".format(val)
         return msg, val
 
     def check_voltage(self):
         """ Check the voltage setting """
-        self.sock.sendall(b'volt?\n')
-        val = float(self.read())
+        val = float(self.send_message(b'volt?\n'))
         msg = "Voltage setting = {:.3f} V".format(val)
         return msg, val
 
     def meas_current(self):
         """ Measure the current """
-        self.sock.sendall(b'meas:curr?\n')
-        val = float(self.read())
+        val = float(self.send_message(b'meas:curr?\n'))
         msg = "Measured current = {:.3f} A".format(val)
         return msg, val
 
     def meas_voltage(self):
         """ Measure the voltage """
-        self.sock.sendall(b'meas:volt?\n')
-        val = float(self.read())
+        val = float(self.send_message(b'meas:volt?\n'))
         msg = "Measured voltage = {:.3f} V".format(val)
         return msg, val
 
     def set_current(self, curr):
         """ Set the current """
-        self.sock.sendall(b'curr %a\n' % curr)
+        self.send_message(b'curr %a\n' % curr, read=False)
         self.wait()
         return self.check_current()
 
     def set_voltage(self, vol):
         """ Set the voltage """
-        self.sock.sendall(b'volt %a\n' % vol)
+        self.send_message(b'volt %a\n' % vol, read=False)
         self.wait()
         return self.check_voltage()
 
     def check_source(self):
         """ Check the source of PMX """
-        self.sock.sendall(b'volt:ext:sour?\n')
-        val = self.read()
+        val = self.send_message(b'volt:ext:sour?\n')
         msg = "Source: " + val
         return msg, val
 
     def use_external_voltage(self):
         """ Set PMX to use external voltage """
-        self.sock.sendall(b'volt:ext:sour volt\n')
+        self.send_message(b'volt:ext:sour volt\n', read=False)
         self.wait()
         return self.check_source()
 
     def ign_external_voltage(self):
         """ Set PMX to ignore external voltage """
-        self.sock.sendall(b'volt:ext:sour none\n')
+        self.send_message(b'volt:ext:sour none\n', read=False)
         self.wait()
         return self.check_source()
 
     def set_current_limit(self, curr_lim):
         """ Set the PMX current limit """
-        self.sock.sendall(b'curr:prot %a\n' % curr_lim)
+        self.send_message(b'curr:prot %a\n' % curr_lim, read=False)
         self.wait()
-        self.sock.sendall(b'curr:prot?\n')
-        val = float(self.read())
+        val = float(self.send_message(b'curr:prot?\n'))
         msg = "Current Limit: {:.3f} A".format(val)
         return msg
 
     def set_voltage_limit(self, vol_lim):
         """ Set the PMX voltage limit """
-        self.sock.sendall(b'volt:prot %a\n' % vol_lim)
+        self.send_message(b'volt:prot %a\n' % vol_lim, read=False)
         self.wait()
-        self.sock.sendall(b'volt:prot?\n')
-        val = float(self.read())
+        val = float(self.send_message(b'volt:prot?\n'))
         msg = "Voltage Limit: {:.3f} V".format(val)
         return msg
 
@@ -156,8 +175,7 @@ def check_prot(self):
         Return:
             val (int): protection status code
         """
-        self.sock.sendall(b'stat:ques?\n')
-        val = int(self.read())
+        val = int(self.send_message(b'stat:ques?\n'))
         return val
 
     def get_prot_msg(self, val):