diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 1859a9222..0af4d6664 100644 --- a/positronic/offboard/client.py +++ b/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 diff --git a/positronic/offboard/tests/test_offboard.py b/positronic/offboard/tests/test_offboard.py index 6967f70a3..73dac5da2 100644 --- a/positronic/offboard/tests/test_offboard.py +++ b/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' diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index add6375c3..97d4f3fe0 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -95,19 +95,21 @@ def test_new_session_passes_additional_headers(self): headers = {'Modal-Key': 'k', 'Modal-Secret': 's'} with ( patch('positronic.offboard.client.connect') as mock_connect, - patch('positronic.offboard.client.InferenceSession') as mock_session_cls, + patch('positronic.offboard.client._handshake') as mock_handshake, ): client = InferenceClient('localhost:8000', headers=headers) - client.new_session() + session = client.new_session() mock_connect.assert_called_once() assert mock_connect.call_args.kwargs['additional_headers'] == headers - mock_session_cls.assert_called_once_with(mock_connect.return_value, infer_timeout=DEFAULT_INFER_TIMEOUT) + assert session._websocket is mock_connect.return_value + assert session.metadata is mock_handshake.return_value + assert session._infer_timeout == DEFAULT_INFER_TIMEOUT def test_new_session_without_headers_passes_none(self): with ( patch('positronic.offboard.client.connect') as mock_connect, - patch('positronic.offboard.client.InferenceSession'), + patch('positronic.offboard.client._handshake'), ): client = InferenceClient('localhost:8000') client.new_session() @@ -197,7 +199,7 @@ def test_unknown_scheme_rejected(self): def test_every_session_dials_the_session_url(self): with ( patch('positronic.offboard.client.connect') as mock_connect, - patch('positronic.offboard.client.InferenceSession'), + patch('positronic.offboard.client._handshake'), ): client = InferenceClient('localhost:8000/api/v1/session/10000?fps=10') client.new_session() @@ -220,13 +222,13 @@ def test_a_403_retries_and_the_session_that_follows_is_returned(self): patch( 'positronic.offboard.client.connect', side_effect=[_refused(HTTPStatus.FORBIDDEN), MagicMock()] ) as mock_connect, - patch('positronic.offboard.client.InferenceSession') as mock_session_cls, + patch('positronic.offboard.client._handshake') as mock_handshake, patch('positronic.offboard.client.time.sleep'), ): session = InferenceClient('localhost:8000').new_session() assert mock_connect.call_count == 2 - assert session is mock_session_cls.return_value + assert session.metadata is mock_handshake.return_value def test_a_403_gives_up_once_its_attempts_are_spent(self): with ( @@ -234,7 +236,7 @@ def test_a_403_gives_up_once_its_attempts_are_spent(self): 'positronic.offboard.client.connect', side_effect=[_refused(HTTPStatus.FORBIDDEN)] * (_ConnectRetries.MAX_FORBIDDEN_ATTEMPTS + 5), ) as mock_connect, - patch('positronic.offboard.client.InferenceSession'), + patch('positronic.offboard.client._handshake'), patch('positronic.offboard.client.time.sleep'), pytest.raises(InvalidStatus), ): @@ -246,7 +248,7 @@ def test_a_403_gives_up_once_its_attempts_are_spent(self): def test_a_refusal_that_no_warm_up_clears_is_raised_at_once(self, status): with ( patch('positronic.offboard.client.connect', side_effect=_refused(status)) as mock_connect, - patch('positronic.offboard.client.InferenceSession'), + patch('positronic.offboard.client._handshake'), patch('positronic.offboard.client.time.sleep'), pytest.raises(InvalidStatus), ): @@ -259,7 +261,7 @@ def test_each_session_opens_on_a_full_budget(self): one_session = [_refused(HTTPStatus.FORBIDDEN)] * (_ConnectRetries.MAX_FORBIDDEN_ATTEMPTS - 1) + [MagicMock()] with ( patch('positronic.offboard.client.connect', side_effect=one_session * 2) as mock_connect, - patch('positronic.offboard.client.InferenceSession'), + patch('positronic.offboard.client._handshake'), patch('positronic.offboard.client.time.sleep'), ): client = InferenceClient('localhost:8000')