From fb4ca26927b1f1fd977c1ecb5d5366cc996b163e Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Mon, 31 Aug 2026 10:06:12 +0000 Subject: [PATCH] Reconnect and send the observation again when the inference socket drops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A backend that scales to zero drops the websocket when its container recycles. The rig read that as fatal: one dropped send ended the World mid-episode and the run recorded nothing. `InferenceSession.infer` now reconnects once through the client's own cold-start retries and sends the same observation on the new socket. The handshake moves out of `InferenceSession.__init__` into `_handshake`, so `InferenceClient._open` can connect, handshake and retry as one step and serve both a session open and a reconnect. A stall keeps raising: the server may still be computing the observation, so a second send doubles the work on a backend already too slow. So does a second drop, which is a backend that cannot serve the observation at all. The reconnect gets 45s of wall clock against `connect_deadline`'s 900s, because it is spent with the arm holding its last setpoint and an attended trial carries no episode deadline behind it. The connect loop now clips its backoff sleep to what is left, so that deadline bounds the wall clock spent rather than the instant the last attempt starts at. Ticket: none — unblocks a customer rollout that died on the rig --- positronic/offboard/client.py | 153 +++++++++++++----- positronic/offboard/tests/test_offboard.py | 82 +++++++++- .../offboard/tests/test_remote_policy.py | 22 +-- 3 files changed, 202 insertions(+), 55 deletions(-) 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')