Reconnect and send the observation again when the inference socket drops

positronic/offboard/client.py

@@ -2,6 +2,7 @@
 import ssl
 import time
 import urllib.parse
+from collections.abc import Callable
 from enum import Enum
 from http import HTTPStatus
 from typing import Any
@@ -23,41 +24,62 @@
 # per use.
 DEFAULT_INFER_TIMEOUT = 180.0
 
+# How long ``infer`` may spend rebuilding a dropped connection, across every attempt of one reconnect.
+# FOOTGUN: the arm holds its last setpoint throughout, and an attended trial (``timeout_sec=None``) has no
+# episode deadline behind this one to end that stall.
+DEFAULT_RECONNECT_DEADLINE = 45.0
 
-class InferenceSession:
-    def __init__(self, websocket: Connection, infer_timeout: float = DEFAULT_INFER_TIMEOUT):
-        self._websocket = websocket
-        self._infer_timeout = infer_timeout
-        self._metadata = self._handshake()
 
-    def _handshake(self, timeout_per_message: float = 30.0) -> dict[str, Any]:
-        """Receive status updates until server is ready.
+def _handshake(websocket: Connection, timeout_per_message: float = 30.0) -> dict[str, Any]:
+    """Receive status updates until server is ready.
 
-        The server must send an update at least every ``timeout_per_message`` seconds.
-        """
-        try:
-            while True:
-                response = deserialise(self._websocket.recv(timeout=timeout_per_message))
-                if protocol.ERROR in response:
-                    raise RuntimeError(f'Server error: {response[protocol.ERROR]}')
-                try:
-                    status = protocol.ServerStatus(response.get(protocol.STATUS))
-                except ValueError:
-                    raise RuntimeError(f'Unexpected server response: {response}') from None
-
-                if status is protocol.ServerStatus.READY:
-                    return response[protocol.META]
-                if status is protocol.ServerStatus.ERROR:
-                    raise RuntimeError('Server error: Unknown error')
-
-                message = response.get(protocol.MESSAGE, status)
-                logger.info(f'Server status: [{status}] {message}')
+    The server must send an update at least every ``timeout_per_message`` seconds.
+    """
+    try:
+        while True:
+            response = deserialise(websocket.recv(timeout=timeout_per_message))
+            if protocol.ERROR in response:
+                raise RuntimeError(f'Server error: {response[protocol.ERROR]}')
+            try:
+                status = protocol.ServerStatus(response.get(protocol.STATUS))
+            except ValueError:
+                raise RuntimeError(f'Unexpected server response: {response}') from None
 
-        except TimeoutError:
-            raise TimeoutError(
-                f'Server did not send status update within {timeout_per_message}s. '
-                f'Server may have crashed or model loading is taking too long without progress updates.'
-            ) from None
+            if status is protocol.ServerStatus.READY:
+                return response[protocol.META]
+            if status is protocol.ServerStatus.ERROR:
+                raise RuntimeError('Server error: Unknown error')
+
+            message = response.get(protocol.MESSAGE, status)
+            logger.info(f'Server status: [{status}] {message}')
+
+    except TimeoutError:
+        raise TimeoutError(
+            f'Server did not send status update within {timeout_per_message}s. '
+            f'Server may have crashed or model loading is taking too long without progress updates.'
+        ) from None
+
+
+class InferenceSession:
+    """One websocket to one served session, carrying observations out and trajectories back.
+
+    ``reopen`` returns a fresh socket whose handshake has reached ready, and makes a dropped connection
+    recoverable: an ``infer`` that loses the socket reconnects through it and sends the same observation
+    again. A session without one raises on a drop, which is all a caller owning no way to reconnect can do.
+    """
+
+    def __init__(
+        self,
+        websocket: Connection,
+        infer_timeout: float = DEFAULT_INFER_TIMEOUT,
+        *,
+        metadata: dict[str, Any] | None = None,
+        reopen: Callable[[], Connection] | None = None,
+    ):
+        self._websocket = websocket
+        self._infer_timeout = infer_timeout
+        self._reopen = reopen
+        self._metadata = _handshake(websocket) if metadata is None else metadata
 
     @property
     def metadata(self) -> dict[str, Any]:
@@ -69,15 +91,39 @@ def infer(self, obs: dict[str, Any]) -> Any:
         ``obs`` must be wire-serializable: plain-data containers and scalars, plus numeric numpy
         arrays/scalars, and no arbitrary Python objects. The result is whatever the server's session
         returned — canonically a list of action dicts, but a bare dict or ``None`` too.
+
+        A dropped socket reconnects and sends the observation once more, within
+        ``DEFAULT_RECONNECT_DEADLINE``. Every other failure reaches the caller as it is.
         """
         serialised = serialise(obs)
         logger.debug('Size of serialised obs: %1.f KiB', len(serialised) / 1024)
+        try:
+            return self._round_trip(serialised)
+        # The server may still be computing this observation, so a second send doubles the work on a backend
+        # already too slow to answer the first.
+        except TimeoutError:
+            raise
+        # A container recycling under a backend that scales to zero drops the socket, on the send or the recv,
+        # as ``ConnectionClosed`` or a bare ``OSError`` (TLS included). Neither says anything about the
+        # observation, so it goes again on a new socket.
+        except (ConnectionClosed, OSError) as e:
+            if self._reopen is None:
+                raise
+            logger.warning('Inference connection dropped (%s); reconnecting and sending the observation again', e)
+            self._websocket.close()
+            # The URL pins the model, so the session keeps the metadata it opened under: the episode records
+            # that meta once, and it describes the whole episode.
+            self._websocket = self._reopen()
+            # The server built this session a moment ago, so it holds none of the dropped one's state. A second
+            # drop is a backend that cannot serve this observation at all, and reaches the caller.
+            return self._round_trip(serialised)
+
+    def _round_trip(self, serialised: bytes) -> Any:
         # The pair reads as the uplink and then the wait the server's own time sits inside: each span
         # holds the socket alone. A send outlasting its own bytes is an uplink too slow for the payload.
         wire_bytes = {telemetry_keys.ATTR_WIRE_BYTES: len(serialised)}
         with telemetry.span(telemetry_keys.SPAN_WIRE_SEND, **wire_bytes):
             self._websocket.send(serialised)
-
         try:
             with telemetry.span(telemetry_keys.SPAN_WIRE_RECV):
                 received = self._websocket.recv(timeout=self._infer_timeout)
@@ -168,8 +214,10 @@ class InferenceClient:
     out of the URL, which is meant to be safe to hand around.
 
     The timeouts describe this connection, not any one session: ``open_timeout`` bounds the TCP/TLS
-    handshake alone, ``connect_deadline`` how long a cold backend may take to answer across retries, and
-    ``infer_timeout`` one inference round trip.
+    handshake alone, ``connect_deadline`` how long a cold backend may take to answer across retries,
+    ``infer_timeout`` one inference round trip, and ``reconnect_deadline`` how long a session may spend
+    rebuilding a socket dropped mid-inference. The last is much the shortest, because it is spent with the
+    caller's robot live — see ``DEFAULT_RECONNECT_DEADLINE``.
     """
 
     def __init__(
@@ -180,6 +228,7 @@ def __init__(
         open_timeout: float = 10.0,
         connect_deadline: float = 900.0,
         infer_timeout: float = DEFAULT_INFER_TIMEOUT,
+        reconnect_deadline: float = DEFAULT_RECONNECT_DEADLINE,
     ):
         split = urllib.parse.urlsplit(url if '://' in url else f'//{url}')
         if split.scheme not in ('', 'http', 'ws', 'https', 'wss'):
@@ -203,10 +252,26 @@ def __init__(
         self.open_timeout = open_timeout
         self.connect_deadline = connect_deadline
         self.infer_timeout = infer_timeout
+        self.reconnect_deadline = reconnect_deadline
 
     def new_session(self) -> InferenceSession:
         """Creates a new inference session on the model the URL names."""
-        deadline = time.monotonic() + self.connect_deadline
+        ws, metadata = self._open(self.connect_deadline)
+        return InferenceSession(
+            ws,
+            infer_timeout=self.infer_timeout,
+            metadata=metadata,
+            # A drop mid-inference reconnects through the same cold-start retries, on the shorter deadline
+            # a live robot can afford.
+            reopen=lambda: self._open(self.reconnect_deadline)[0],
+        )
+
+    def _open(self, deadline_sec: float) -> tuple[Connection, dict[str, Any]]:
+        """A connected socket whose status handshake has reached ready, and the metadata that handshake read.
+
+        Retries a backend that is not up yet, until ``deadline_sec`` of wall clock has passed.
+        """
+        deadline = time.monotonic() + deadline_sec
         backoff = 1.0
         retries = _ConnectRetries()
         while True:
@@ -221,25 +286,25 @@ def new_session(self) -> InferenceSession:
                     additional_headers=self.headers,
                     ping_interval=20.0,
                 )
-                return InferenceSession(ws, infer_timeout=self.infer_timeout)
+                return ws, _handshake(ws)
             # ``SSLCertVerificationError`` is an ``ssl.SSLError``, but a bad certificate is permanent
             # misconfiguration, not a cold start — surface it immediately instead of retrying to the deadline.
             except ssl.SSLCertVerificationError as e:
                 raise type(e)(f'{e} (connecting to {self.session_url})') from e
-            # A cold backend fails before the session is ready in several ways: the connect times out, the edge
-            # resets TLS (``SSLError``), it rejects or drops the HTTP upgrade (``InvalidHandshake`` — e.g. a
-            # 502/503 while the backend boots), or it accepts the socket and then drops or stalls the status
-            # handshake inside ``InferenceSession`` (``ConnectionClosed``/``TimeoutError``). All mean "not ready
-            # yet", so retry within the deadline instead of letting one kill the run.
+            # Each of these is a backend not up yet — a timed-out connect, a TLS reset at the edge, a 502/503
+            # on the upgrade, a status handshake dropped or stalled — so retry within the deadline.
             except (TimeoutError, ssl.SSLError, ConnectionClosed, InvalidHandshake) as e:
                 if ws is not None:
                     ws.close()
                 if retries.take(e) is _ConnectOutcome.SURFACE:
                     raise
-                if time.monotonic() >= deadline:
+                # The sleep is clipped to what is left, so the deadline bounds the wall clock spent here and
+                # not merely the instant the last attempt starts at.
+                pause = min(backoff, deadline - time.monotonic())
+                if pause <= 0:
                     raise TimeoutError(f'{e} (connecting to {self.session_url})') from e
-                logger.info('Server not ready (cold start?): %s; retrying in %.0fs', e, backoff)
-                time.sleep(backoff)
+                logger.info('Server not ready (cold start?): %s; retrying in %.0fs', e, pause)
+                time.sleep(pause)
                 backoff = min(backoff * 2, 30.0)
             except OSError as e:
                 raise type(e)(f'{e} (connecting to {self.session_url})') from e

positronic/offboard/tests/test_offboard.py

@@ -1,8 +1,11 @@
 from types import MappingProxyType
+from typing import Any, cast
 from unittest.mock import ANY
 
 import numpy as np
 import pytest
+from websockets.exceptions import ConnectionClosedError
+from websockets.sync.connection import Connection
 
 from positronic import keys
 from positronic.drivers.roboarm.command import (
@@ -14,7 +17,8 @@
     to_wire,
 )
 from positronic.geom import Rotation, Transform3D
-from positronic.offboard.client import InferenceClient
+from positronic.offboard import protocol
+from positronic.offboard.client import InferenceClient, InferenceSession
 from positronic.offboard.protocol import deserialise, serialise, typed_commands
 from positronic.utils.serialization import encode_jpeg
 
@@ -269,3 +273,79 @@ def test_a_command_crossing_the_wire_arrives_typed_from_either_shape(self):
 
         assert isinstance(from_envelope, CartesianPosition) and isinstance(from_bare, CartesianPosition)
         np.testing.assert_allclose(from_bare.pose.translation, from_envelope.pose.translation, atol=1e-6)
+
+
+class _FakeConnection:
+    """A websocket answering one canned result, which drops the connection on its first ``send`` if told to."""
+
+    def __init__(self, result: Any, *, drops_first_send: bool = False):
+        self._result = result
+        self._drops = drops_first_send
+        self.sent: list[bytes] = []
+        self.closed = False
+
+    def send(self, data: bytes) -> None:
+        if self._drops:
+            raise ConnectionClosedError(None, None)
+        self.sent.append(data)
+
+    def recv(self, timeout: float | None = None) -> bytes:
+        return serialise({protocol.RESULT: self._result})
+
+    def close(self) -> None:
+        self.closed = True
+
+
+def _session(*sockets: _FakeConnection) -> InferenceSession:
+    """A session over ``sockets[0]``, reconnecting onto each of the rest in turn."""
+    spares = iter(sockets[1:])
+    return InferenceSession(
+        cast(Connection, sockets[0]), metadata={'model_name': 'test'}, reopen=lambda: cast(Connection, next(spares))
+    )
+
+
+class TestInferDropsTheConnection:
+    """A backend that scales to zero drops the socket when its container recycles. The session that can
+    reconnect sends the observation again; the one that cannot says so."""
+
+    def test_a_dropped_send_reconnects_and_serves_the_observation(self):
+        dropped = _FakeConnection(None, drops_first_send=True)
+        fresh = _FakeConnection({'action_data': [1, 2, 3]})
+
+        result = _session(dropped, fresh).infer({'image': 'test'})
+
+        assert result == {'action_data': [1, 2, 3]}
+        assert dropped.closed, 'the dropped socket is closed before the reconnect'
+        assert [deserialise(sent) for sent in fresh.sent] == [{'image': 'test'}]
+
+    def test_a_session_that_cannot_reconnect_raises(self):
+        """``reopen`` is what makes a drop recoverable; without one the caller hears the drop."""
+        dropped = _FakeConnection(None, drops_first_send=True)
+        session = InferenceSession(cast(Connection, dropped), metadata={'model_name': 'test'})
+
+        with pytest.raises(ConnectionClosedError):
+            session.infer({'image': 'test'})
+
+    def test_a_second_drop_reaches_the_caller(self):
+        """One retry, on a session the server has just built. A backend dropping that one too cannot serve
+        this observation at all."""
+        session = _session(_FakeConnection(None, drops_first_send=True), _FakeConnection(None, drops_first_send=True))
+
+        with pytest.raises(ConnectionClosedError):
+            session.infer({'image': 'test'})
+
+    def test_a_stalled_server_is_not_reconnected(self):
+        """A recv timeout leaves the observation with a server that may still be computing it, so sending it
+        again would double the work on a backend already too slow. It surfaces, and the socket stays closed."""
+
+        def stall(timeout: float | None = None) -> bytes:
+            raise TimeoutError
+
+        stalled = _FakeConnection(None)
+        stalled.recv = stall
+        spare = _FakeConnection({'action_data': ['unused']})
+
+        with pytest.raises(TimeoutError):
+            _session(stalled, spare).infer({'image': 'test'})
+
+        assert stalled.closed and spare.sent == [], 'a stall must not reconnect'