From cd048df9901e767ec675fab187b55f3d83dd9860 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 00:06:51 +0000 Subject: [PATCH 01/46] Add a gRPC wire beside the websocket for inference sessions `positronic/offboard` serves a session over two wires now. `wire.py` holds the seam and the websocket wire; `grpc_wire.py` holds the gRPC one. Every frame stays the same msgpack message, and a wire carries it as opaque bytes, so `protocol.py` does not change. A gRPC session is one bidirectional stream. A generic handler with no serialiser carries the frames, so there is no protobuf schema and no generated code. The session path, the query and the `Authorization` header cross as gRPC metadata, so both wires accept the same model ids and the same session params. `grpc://host:port/api/v1/session/` selects the wire, and `PolicyServer(grpc_port=...)` serves it beside the websocket one. Both wires run on one event loop, so they share the model slot, the inference lock and the idle watchdog. Python's websocket stack costs about 30 ms per 846 KiB observation, which gRPC does in about 1 ms. The websocket wire stays the default, because a managed HTTPS front drops the HTTP/2 frame detail gRPC needs. Ticket: none - a transport measurement with no ticket; the numbers are in the pull request --- positronic/offboard/README.md | 48 ++++- positronic/offboard/client.py | 149 ++++++++----- positronic/offboard/grpc_wire.py | 200 ++++++++++++++++++ positronic/offboard/server.py | 120 +++++++---- positronic/offboard/tests/conftest.py | 22 +- positronic/offboard/tests/test_grpc_wire.py | 174 +++++++++++++++ .../offboard/tests/test_remote_policy.py | 4 +- positronic/offboard/tests/test_server.py | 10 +- positronic/offboard/wire.py | 112 ++++++++++ pyproject.toml | 3 + uv.lock | 43 ++++ 11 files changed, 769 insertions(+), 116 deletions(-) create mode 100644 positronic/offboard/grpc_wire.py create mode 100644 positronic/offboard/tests/test_grpc_wire.py create mode 100644 positronic/offboard/wire.py diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index ddf7140e1..832cd3b4b 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -4,13 +4,37 @@ This package implements the protocol and utilities for offboard policy inference ## Protocol v1 -The unified WebSocket protocol is built to enable ANY hardware to connect to ANY model. All Positronic inference servers (LeRobot, GR00T, OpenPI) implement this protocol, allowing a single `.remote` policy client to work across all vendors. +The unified protocol is built to enable ANY hardware to connect to ANY model. All Positronic inference servers (LeRobot, GR00T, OpenPI) implement this protocol, allowing a single `.remote` policy client to work across all vendors. + +### Wires + +The protocol is a sequence of msgpack frames, and two wires carry them. Both carry the same frames in +the same order, so everything below holds on each. + +| Wire | URL | Port | +|---|---|---| +| WebSocket | `ws://host:8000/api/v1/session[/]` | the server's `port`, beside the HTTP routes | +| gRPC | `grpc://host:9000/api/v1/session[/]` | the server's `grpc_port`, sessions alone | + +The WebSocket wire is the default, and a server serves gRPC only when `grpc_port` names a port. A +gRPC session is one bidirectional stream of the same frames, so no `.proto` file describes them. +The session path and the query cross as the `positronic-session-path` and `positronic-session-query` +metadata, and `Authorization` crosses as the `authorization` metadata. + +Python's WebSocket stack costs about 30 ms per 846 KiB observation in framing and reassembly, which +gRPC does in about 1 ms. Take the gRPC wire on an endpoint a client reaches directly. A managed HTTPS +front usually translates HTTP into its own protocol and drops the HTTP/2 frame detail gRPC needs, so +an endpoint behind one keeps the WebSocket wire. + +`/api/v1/models` is an HTTP route, so it stays on the server's `port`. `InferenceClient.list_models` +over a `grpc://` URL says so. ### Authentication `PolicyServer(auth_token=...)` gates every route below on `Authorization: Bearer `, answering -`401` on the HTTP route and refusing the WebSocket upgrade before the session opens. `serve` — the -entry point every vendor CLI exposes — takes that token from the `AUTH_TOKEN` environment variable, so +`401` on the HTTP route, refusing the WebSocket upgrade before the session opens, and answering +`PERMISSION_DENIED` on the gRPC wire. `serve` — the entry point every vendor CLI exposes — takes that +token from the `AUTH_TOKEN` environment variable, so a secret never lands in the process arguments. No token serves open, which is the usual shape on a trusted LAN; an empty one is a broken secret and refuses to start. `InferenceClient(headers=...)` carries the header, and `positronic.cfg.policy.authed_remote` fills it in from the same variable. @@ -67,14 +91,14 @@ Rules: - **The model source is fixed at launch.** Params that would change it (e.g. `?source.checkpoint=...`) are rejected; the only way to get a different model is the path. - **Only config-launched servers accept params.** All vendor servers qualify; a `PolicyServer` built from an already-instantiated pipeline rejects every param. -Any violation — including an unknown key — fails at connect: the server sends `{"status": "error", "error": ...}` and closes the socket (code 1008) before anything moves, and the Python client raises `RuntimeError`. Overrides apply per session, and the `local_stack` declared in the ready handshake reflects them. +Any violation — including an unknown key — fails at connect: the server sends `{"status": "error", "error": ...}` and ends the session before anything moves, and the Python client raises `RuntimeError`. Overrides apply per session, and the `local_stack` declared in the ready handshake reflects them. Because the whole session configuration fits in the URL, one string is a complete endpoint description: `--policy=.remote --policy.url='gpu-host:8000?codec.fps=10'` accepts `host`, `host:port`, and full -`http(s)`/`ws(s)` URLs — optionally with `/api/v1/session/` — and forwards the query string verbatim. +`http(s)`/`ws(s)`/`grpc` URLs — optionally with `/api/v1/session/` — and forwards the query string verbatim. Credentials are the exception and stay a separate `headers` argument, so the URL itself is safe to hand around. -### WebSocket Flow +### Session Flow #### 1. Handshake Upon connection, the server sends a ready packet with metadata: @@ -202,9 +226,9 @@ uv run positronic eval run --eval=.sim.positronic.stack_cubes \ **Status Streaming:** Long model loads are handled gracefully with progress updates. -**Server-side recording:** Servers accept an optional `recording_dir`. When set, each WebSocket session writes a rerun `.rrd` file that taps both sides of the codec: `raw` captures the obs/action at the wire boundary, and `inference` captures the encoded observation and raw model output. +**Server-side recording:** Servers accept an optional `recording_dir`. When set, each session writes a rerun `.rrd` file that taps both sides of the codec: `raw` captures the obs/action at the wire boundary, and `inference` captures the encoded observation and raw model output. -**Python Client:** We provide a Python client (`positronic.offboard.client.InferenceClient`) that handles the WebSocket protocol automatically. While the API is currently in alpha and may change, we'll do our best to maintain backward compatibility for the inference client. +**Python Client:** We provide a Python client (`positronic.offboard.client.InferenceClient`) that handles the protocol automatically. While the API is currently in alpha and may change, we'll do our best to maintain backward compatibility for the inference client. ## Classes @@ -220,15 +244,15 @@ pipeline = ChunkedSchedule() | remote | PolicySource(my_policy) PolicyServer(pipeline, host='0.0.0.0', port=8000).serve() ``` -`PolicySource` serves one ready in-process policy; vendors instead define a `ModelSource` over a checkpoint directory. Passing a `cfn.Config` that builds the pipeline — as the vendor servers do with their named pipelines — enables [session parameters](#session-parameters); an instantiated pipeline serves exactly as launched. `recording_dir` enables the per-session recording taps described above, and `idle_timeout_min` shuts the server down after that many minutes without activity. +`PolicySource` serves one ready in-process policy; vendors instead define a `ModelSource` over a checkpoint directory. Passing a `cfn.Config` that builds the pipeline — as the vendor servers do with their named pipelines — enables [session parameters](#session-parameters); an instantiated pipeline serves exactly as launched. `recording_dir` enables the per-session recording taps described above, `grpc_port` adds the gRPC wire, and `idle_timeout_min` shuts the server down after that many minutes without activity. ### `server.serve` -The CLI entry point every vendor server exposes. A vendor binds `pipeline` to each of its named pipelines and lists the results as subcommands, so `-server ` launches one. Only `--host`, `--port`, `--recording_dir` and `--idle_timeout_min` are flags of `serve` itself; everything the served model is — codec, source, checkpoint directory — is reached through the pipeline (`--pipeline.source.checkpoints_dir=...`), which is also where a deployment preset binds it. +The CLI entry point every vendor server exposes. A vendor binds `pipeline` to each of its named pipelines and lists the results as subcommands, so `-server ` launches one. Only `--host`, `--port`, `--grpc_port`, `--recording_dir` and `--idle_timeout_min` are flags of `serve` itself; everything the served model is — codec, source, checkpoint directory — is reached through the pipeline (`--pipeline.source.checkpoints_dir=...`), which is also where a deployment preset binds it. ### `client.InferenceClient` A Python client for connecting to an inference server. One URL addresses it, in the same forms `RemotePolicy` accepts: an omitted port is the scheme's own, 443 for `https`/`wss` and 80 otherwise. The URL -fixes the model and the session params, so serving another model means another client. +fixes the wire, the model and the session params, so serving another model means another client. ```python from positronic.offboard.client import InferenceClient @@ -237,6 +261,8 @@ from positronic.offboard.client import InferenceClient client = InferenceClient('localhost:8000') # A named model, tuned for every session this client opens # client = InferenceClient('localhost:8000/api/v1/session/model_a?codec.fps=10') +# The same session on the gRPC wire +# client = InferenceClient('grpc://localhost:9000/api/v1/session/model_a') session = client.new_session() meta = session.metadata diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 98369c4f7..5c2d15a75 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -6,12 +6,12 @@ from http import HTTPStatus from typing import Any +import grpc import httpx from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidStatus from websockets.sync.client import connect -from websockets.sync.connection import Connection -from . import protocol +from . import grpc_wire, protocol, wire from .protocol import deserialise, serialise, typed_commands logger = logging.getLogger(__name__) @@ -23,8 +23,10 @@ class InferenceSession: - def __init__(self, websocket: Connection, infer_timeout: float = DEFAULT_INFER_TIMEOUT): - self._websocket = websocket + """One session over one open connection, whichever wire carries it.""" + + def __init__(self, conn: wire.ClientConnection, infer_timeout: float = DEFAULT_INFER_TIMEOUT): + self._conn = conn self._infer_timeout = infer_timeout self._metadata = self._handshake() @@ -35,7 +37,7 @@ def _handshake(self, timeout_per_message: float = 30.0) -> dict[str, Any]: """ try: while True: - response = deserialise(self._websocket.recv(timeout=timeout_per_message)) + response = deserialise(self._conn.recv(timeout=timeout_per_message)) if protocol.ERROR in response: raise RuntimeError(f'Server error: {response[protocol.ERROR]}') try: @@ -71,14 +73,14 @@ def infer(self, obs: dict[str, Any]) -> Any: serialised = serialise(obs) logger.debug('Size of serialised obs: %1.f KiB', len(serialised) / 1024) - self._websocket.send(serialised) + self._conn.send(serialised) try: - response = deserialise(self._websocket.recv(timeout=self._infer_timeout)) + response = deserialise(self._conn.recv(timeout=self._infer_timeout)) except TimeoutError: # The observation is in flight but unanswered; the server's late response would sit in the socket and # the next ``recv`` would pair it with a future observation. Close so the desynced session can't be # reused — a subsequent ``infer`` fails loudly on the closed socket instead. - self._websocket.close() + self._conn.close() raise TimeoutError( f'No inference response within {self._infer_timeout}s — server stalled or connection half-open' ) from None @@ -90,15 +92,7 @@ def infer(self, obs: dict[str, Any]) -> Any: return typed_commands(response[protocol.RESULT]) def close(self): - state_before_close = self._websocket.state.name - self._websocket.close() - # A close that times out still reaches CLOSED locally; only the close code says the server answered. - logger.info( - 'InferenceSession.close: state %s -> %s, close code %s', - state_before_close, - self._websocket.state.name, - self._websocket.close_code, - ) + logger.info('InferenceSession.close: %s', self._conn.close()) def _session_path(path: str, url: str) -> str: @@ -107,10 +101,10 @@ def _session_path(path: str, url: str) -> str: A URL naming no model — a bare host, or the endpoint with or without a trailing slash — addresses the endpoint itself, which serves whatever the server pinned. """ - if path.rstrip('/') in ('', '/api/v1/session'): - return '/api/v1/session' - if not path.startswith('/api/v1/session/'): - raise ValueError(f'Unexpected path {path!r} in {url!r}; expected /api/v1/session[/]') + if path.rstrip('/') in ('', wire.SESSION_PATH): + return wire.SESSION_PATH + if not path.startswith(f'{wire.SESSION_PATH}/'): + raise ValueError(f'Unexpected path {path!r} in {url!r}; expected {wire.SESSION_PATH}[/]') # Kept as written, percent-encoding included, so the server decodes exactly the id whoever handed out # the URL meant: a trailing slash is part of that id, and an id may itself be a path (a HuggingFace # repo, say), whose own slashes stay separators. @@ -122,11 +116,44 @@ class _ConnectOutcome(Enum): SURFACE = 'surface' +class _Refusal(Enum): + """What a refused connect says about the server.""" + + COLD = 'cold' # still coming up; retry to the deadline + FORBIDDEN = 'forbidden' # a cold backend, or a refused credential; a few attempts, then surface + FINAL = 'final' # the endpoint is saying no; surface at once + + +_COLD_GRPC_CODES = (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.DEADLINE_EXCEEDED) + + +def _refusal(e: Exception) -> _Refusal: + """How to read a refused connect, over either wire. + + Each gRPC code stands for the HTTP status its wire twin answers: ``PERMISSION_DENIED`` for 403, + ``UNAVAILABLE`` for 503, ``RESOURCE_EXHAUSTED`` for 429. + """ + if isinstance(e, InvalidStatus): + status = e.response.status_code + if status == HTTPStatus.FORBIDDEN: + return _Refusal.FORBIDDEN + if status >= HTTPStatus.INTERNAL_SERVER_ERROR or status == HTTPStatus.TOO_MANY_REQUESTS: + return _Refusal.COLD + return _Refusal.FINAL + # A gRPC error carries its code as a `Call`; anything else says nothing about the server. + if isinstance(e, grpc.Call): + code = e.code() + if code is grpc.StatusCode.PERMISSION_DENIED: + return _Refusal.FORBIDDEN + return _Refusal.COLD if code in _COLD_GRPC_CODES else _Refusal.FINAL + return _Refusal.COLD + + class _ConnectRetries: """The retry policy over one ``new_session``'s connect attempts. - 403 is both a cold backend and a refused credential, so it gets a few attempts rather than the whole - ``connect_deadline``. + A refusal both wires answer for a cold backend and for a refused credential — HTTP 403, gRPC + ``PERMISSION_DENIED`` — gets a few attempts rather than the whole ``connect_deadline``. """ MAX_FORBIDDEN_ATTEMPTS = 3 @@ -136,17 +163,19 @@ def __init__(self) -> None: def take(self, e: Exception) -> _ConnectOutcome: """Spend a refused connect against the budget.""" - if not isinstance(e, InvalidStatus): - return _ConnectOutcome.RETRY - status = e.response.status_code - if status == HTTPStatus.FORBIDDEN: + refusal = _refusal(e) + if refusal is _Refusal.FORBIDDEN: self._forbidden_attempts += 1 again = self._forbidden_attempts < self.MAX_FORBIDDEN_ATTEMPTS else: - again = status >= HTTPStatus.INTERNAL_SERVER_ERROR or status == HTTPStatus.TOO_MANY_REQUESTS + again = refusal is _Refusal.COLD return _ConnectOutcome.RETRY if again else _ConnectOutcome.SURFACE +# The URL scheme that puts a session on the gRPC wire. +_GRPC_SCHEME = 'grpc' + + class InferenceClient: """The wire connection to one inference server, addressed by one URL. @@ -156,6 +185,9 @@ class InferenceClient: session — the model id it names and the query it carries as session params — reaches the server exactly as written, so every session opened here serves that model with those params. + ``grpc://`` names the same session on the gRPC wire, which the server offers on a port of its own. That + port carries sessions alone, so ``list_models`` needs the HTTP URL. + ``headers`` carry auth, whether the server checks it or a proxy in front of it does — credentials stay out of the URL, which is meant to be safe to hand around. @@ -174,12 +206,12 @@ def __init__( infer_timeout: float = DEFAULT_INFER_TIMEOUT, ): split = urllib.parse.urlsplit(url if '://' in url else f'//{url}') - if split.scheme not in ('', 'http', 'ws', 'https', 'wss'): + if split.scheme not in ('', 'http', 'ws', 'https', 'wss', _GRPC_SCHEME): raise ValueError(f'Unsupported scheme {split.scheme!r} in {url!r}') if not split.hostname: raise ValueError(f'No host in {url!r}') secure = split.scheme in ('https', 'wss') - ws_scheme = 'wss' if secure else 'ws' + session_scheme = _GRPC_SCHEME if split.scheme == _GRPC_SCHEME else ('wss' if secure else 'ws') http_scheme = 'https' if secure else 'http' default_port = 443 if secure else 80 # urlsplit strips the brackets an IPv6 host needs back in a netloc. @@ -189,43 +221,60 @@ def __init__( # Forwarded verbatim: the server reads each param value as a JSON literal, and only whoever wrote # the URL knows whether `true` means the bool or the string. query = f'?{split.query}' if split.query else '' - self.session_url = f'{ws_scheme}://{netloc}{_session_path(split.path, url)}{query}' - self.api_url = f'{http_scheme}://{netloc}/api/v1' + self._session_path = _session_path(split.path, url) + self._query = split.query + self._grpc_target = f'{host}:{port}' if split.scheme == _GRPC_SCHEME else None + self.session_url = f'{session_scheme}://{netloc}{self._session_path}{query}' + self.api_url = None if self._grpc_target else f'{http_scheme}://{netloc}/api/v1' self.headers = dict(headers) if headers else None self.open_timeout = open_timeout self.connect_deadline = connect_deadline self.infer_timeout = infer_timeout + def _connect(self) -> wire.ClientConnection: + """One session's connection, over the wire the URL names.""" + if self._grpc_target is not None: + return grpc_wire.GrpcClientConnection( + self._grpc_target, self._session_path, self._query, self.headers, self.open_timeout + ) + # A proxy between here and the server closes a connection it has read nothing from, often + # after 60s — well inside one ``infer_timeout`` inference, which sends nothing until it + # answers. The pings keep it open. + websocket = connect( + self.session_url, open_timeout=self.open_timeout, additional_headers=self.headers, ping_interval=20.0 + ) + return wire.WebsocketClientConnection(websocket) + def new_session(self) -> InferenceSession: """Creates a new inference session on the model the URL names.""" deadline = time.monotonic() + self.connect_deadline backoff = 1.0 retries = _ConnectRetries() while True: - ws = None + conn = None try: - # A proxy between here and the server closes a connection it has read nothing from, often - # after 60s — well inside one ``infer_timeout`` inference, which sends nothing until it - # answers. The pings keep it open. - ws = connect( - self.session_url, - open_timeout=self.open_timeout, - additional_headers=self.headers, - ping_interval=20.0, - ) - return InferenceSession(ws, infer_timeout=self.infer_timeout) + conn = self._connect() + return InferenceSession(conn, infer_timeout=self.infer_timeout) # ``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. - except (TimeoutError, ssl.SSLError, ConnectionClosed, InvalidHandshake) as e: - if ws is not None: - ws.close() + # 502/503 while the backend boots), it refuses the gRPC call (``RpcError``), or it accepts the + # connection and then drops or stalls the status handshake inside ``InferenceSession`` + # (``ConnectionClosed``/``PeerDisconnected``/``TimeoutError``). All mean "not ready yet", so retry + # within the deadline instead of letting one kill the run. + except ( + TimeoutError, + ssl.SSLError, + ConnectionClosed, + InvalidHandshake, + grpc.RpcError, + wire.PeerDisconnected, + ) as e: + if conn is not None: + conn.close() if retries.take(e) is _ConnectOutcome.SURFACE: raise if time.monotonic() >= deadline: @@ -238,6 +287,8 @@ def new_session(self) -> InferenceSession: def list_models(self) -> list[str]: """List available models from the server.""" + if self.api_url is None: + raise ValueError(f'{self.session_url} names the gRPC session port; list the models over HTTP') response = httpx.get(f'{self.api_url}/models', headers=self.headers) response.raise_for_status() return response.json()['models'] diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py new file mode 100644 index 000000000..d1b78184a --- /dev/null +++ b/positronic/offboard/grpc_wire.py @@ -0,0 +1,200 @@ +"""The gRPC wire: one bidirectional stream per session, carrying the same ``protocol`` frames. + +The stream is untyped bytes on both sides, so there is no protobuf schema and no generated code: a +generic handler with no serialiser hands each frame over as it arrived. +""" + +import logging +import queue +import threading +import urllib.parse +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping + +import grpc +import grpc.aio +from starlette.datastructures import QueryParams + +from . import wire + +logger = logging.getLogger(__name__) + +# The one method every session runs on. gRPC routes by this path alone. +SERVICE = 'positronic.offboard.v1.Inference' +METHOD = 'Session' +METHOD_PATH = f'/{SERVICE}/{METHOD}' + +# What the websocket wire says in the URL, said here in the session metadata. +SESSION_PATH_HEADER = 'positronic-session-path' +SESSION_QUERY_HEADER = 'positronic-session-query' + +# An observation is a stack of camera frames, and the gRPC default of 4 MiB refuses one. 16 MiB is the +# ceiling uvicorn already gives the websocket wire (``ws_max_size``), so both wires carry the same frame. +_MAX_MESSAGE_BYTES = 16 * 1024 * 1024 + +_MESSAGE_SIZE_OPTIONS = [ + ('grpc.max_receive_message_length', _MAX_MESSAGE_BYTES), + ('grpc.max_send_message_length', _MAX_MESSAGE_BYTES), +] + +# How long ``close`` waits for the server to end the stream, so its own session cleanup runs. +_CLOSE_TIMEOUT_SEC = 5.0 + + +class GrpcClientConnection: + """A client's end of one gRPC session. + + A reader thread drains the response stream into a queue, because the stream itself has no + per-message timeout and ``recv`` needs one. + """ + + def __init__( + self, + target: str, + session_path: str, + query: str, + headers: Mapping[str, str] | None = None, + open_timeout: float = 10.0, + ): + self._target = target + self._channel = grpc.insecure_channel(target, options=_MESSAGE_SIZE_OPTIONS) + try: + grpc.channel_ready_future(self._channel).result(timeout=open_timeout) + except grpc.FutureTimeoutError: + self._channel.close() + raise TimeoutError(f'gRPC channel to {target} is not ready within {open_timeout}s') from None + # gRPC metadata keys are lower case, and they are the same header names the websocket wire sends. + metadata = tuple((key.lower(), value) for key, value in (headers or {}).items()) + ( + (SESSION_PATH_HEADER, session_path), + (SESSION_QUERY_HEADER, query), + ) + self._outbox: queue.SimpleQueue[bytes | None] = queue.SimpleQueue() + self._inbox: queue.SimpleQueue[bytes | BaseException] = queue.SimpleQueue() + self._closed = False + call = self._channel.stream_stream(METHOD_PATH, request_serializer=None, response_deserializer=None) + self._responses = call(self._requests(), metadata=metadata) + self._reader = threading.Thread(target=self._read, name='grpc-session-reader', daemon=True) + self._reader.start() + + def _requests(self): + """The outbound frames. ``None`` ends the stream, which half-closes the session.""" + while (message := self._outbox.get()) is not None: + yield message + + def _read(self) -> None: + """Drain the response stream into the inbox, ending it with what stopped it.""" + try: + for message in self._responses: + self._inbox.put(message) + self._inbox.put(wire.PeerDisconnected(f'{self._target} ended the session')) + except Exception as e: + self._inbox.put(e) + finally: + self._responses.cancel() + + def send(self, message: bytes) -> None: + self._outbox.put(message) + + def recv(self, timeout: float | None = None) -> bytes: + try: + answer = self._inbox.get(timeout=timeout) + except queue.Empty: + raise TimeoutError(f'No message from {self._target} within {timeout}s') from None + if isinstance(answer, BaseException): + raise answer + return answer + + def close(self) -> str: + if self._closed: + return 'already closed' + self._closed = True + self._outbox.put(None) + # The half-close ends the server's session, and the server then ends the stream. Waiting for + # that lets the server release its model slot; closing the channel now would cut it short. + self._reader.join(timeout=_CLOSE_TIMEOUT_SEC) + server_ended_stream = not self._reader.is_alive() + self._channel.close() + # The websocket wire reads the same two facts off a close code. A stream the server never ended + # means it still holds this session, so the next one's handshake waits on a slot nobody released. + return f'peer had ended the stream {self._ended}, server ended it within {_CLOSE_TIMEOUT_SEC}s {server_ended_stream}' + + +def model_id_of(session_path: str) -> str | None: + """The model a session path names, or ``None`` where it names the model the server pinned.""" + prefix = f'{wire.SESSION_PATH}/' + if session_path == wire.SESSION_PATH: + return None + if not session_path.startswith(prefix): + raise ValueError(f'Unexpected session path {session_path!r}; expected {wire.SESSION_PATH}[/]') + return urllib.parse.unquote(session_path[len(prefix) :]) + + +class GrpcServerConnection(wire.ServerConnection): + """A server's end of one gRPC session.""" + + def __init__(self, requests: AsyncIterator[bytes], context: grpc.aio.ServicerContext, headers: Mapping[str, str]): + self._requests = requests + self._context = context + self._headers = headers + + @property + def peer(self) -> str: + return self._context.peer() + + @property + def session_path(self) -> str: + return self._headers.get(SESSION_PATH_HEADER, wire.SESSION_PATH) + + @property + def query_params(self) -> QueryParams: + return QueryParams(self._headers.get(SESSION_QUERY_HEADER, '')) + + async def send(self, message: bytes) -> None: + await self._context.write(message) + + async def receive(self) -> bytes: + try: + return await anext(self._requests) + except StopAsyncIteration: + raise wire.PeerDisconnected(f'{self.peer} ended the session') from None + + async def refuse(self, reason: str) -> None: + self._context.set_code(grpc.StatusCode.ABORTED) + self._context.set_details(reason) + + +def _headers(context: grpc.aio.ServicerContext) -> dict[str, str]: + """The session metadata, as the header names both wires share. A ``-bin`` key carries no header.""" + return {key: value for key, value in (context.invocation_metadata() or ()) if isinstance(value, str)} + + +async def serve( + serve_session: Callable[[GrpcServerConnection], Awaitable[None]], + authorized: Callable[[Mapping[str, str]], bool], + host: str, + port: int, +) -> grpc.aio.Server: + """Start a gRPC server that gives every accepted session to ``serve_session``. + + ``authorized`` reads the session headers and refuses before the session opens, as the websocket + wire refuses the upgrade. + """ + + async def _serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerContext) -> None: + headers = _headers(context) + if not authorized(headers): + await context.abort(grpc.StatusCode.PERMISSION_DENIED, 'Invalid or missing bearer token') + try: + await serve_session(GrpcServerConnection(requests, context, headers)) + except Exception as e: + # The session itself reports what it can over the stream; anything reaching here happened + # before or beyond that, so the client learns of it from the status alone. + logger.error(f'Failed gRPC session: {e}', exc_info=True) + await context.abort(grpc.StatusCode.INTERNAL, str(e)) + + handler = grpc.stream_stream_rpc_method_handler(_serve_one, request_deserializer=None, response_serializer=None) + server = grpc.aio.server(options=_MESSAGE_SIZE_OPTIONS) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler(SERVICE, {METHOD: handler}),)) + bound = server.add_insecure_port(f'{host}:{port}') + await server.start() + logger.info(f'gRPC sessions on {host}:{bound}') + return server diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index fbe7d33b7..84ab72963 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -7,14 +7,15 @@ import os import time from collections import Counter -from collections.abc import Callable +from collections.abc import Callable, Mapping from importlib.metadata import version as _pkg_version from typing import Any import configuronic as cfn +import grpc.aio import pos3 import uvicorn -from fastapi import Depends, FastAPI, Header, HTTPException, WebSocket, WebSocketDisconnect, WebSocketException, status +from fastapi import Depends, FastAPI, Header, HTTPException, WebSocket, WebSocketException, status from starlette.datastructures import QueryParams from positronic.offboard import keys as offboard_keys @@ -23,7 +24,7 @@ from positronic.policy.executor import blocking from positronic.policy.spec import ModelSource, Pipeline, split -from . import protocol +from . import grpc_wire, protocol, wire from .protocol import deserialise, serialise logger = logging.getLogger(__name__) @@ -38,7 +39,7 @@ def bearer(token: str) -> str: return f'Bearer {token}' -async def _acquire_with_keepalives(lock: asyncio.Lock, websocket: WebSocket | None, message: str): +async def _acquire_with_keepalives(lock: asyncio.Lock, conn: wire.ServerConnection | None, message: str): """Acquire ``lock``, emitting ``waiting`` keepalives while queued behind another holder. A peer may hold the lock for a slow load, first-call compile or inference; a silent wait here @@ -49,10 +50,8 @@ async def _acquire_with_keepalives(lock: asyncio.Lock, websocket: WebSocket | No await asyncio.wait_for(lock.acquire(), timeout=10.0) return except TimeoutError: - if websocket is not None: - await websocket.send_bytes( - serialise({protocol.STATUS: protocol.ServerStatus.WAITING, protocol.MESSAGE: message}) - ) + if conn is not None: + await conn.send(serialise({protocol.STATUS: protocol.ServerStatus.WAITING, protocol.MESSAGE: message})) class PolicyManager: @@ -70,8 +69,8 @@ def __init__(self, source: ModelSource): self._lock = asyncio.Lock() self._condition = asyncio.Condition(self._lock) - async def get_policy(self, checkpoint_id: str, websocket: WebSocket | None = None) -> Policy: - await _acquire_with_keepalives(self._lock, websocket, 'Waiting for the model slot') + async def get_policy(self, checkpoint_id: str, conn: wire.ServerConnection | None = None) -> Policy: + await _acquire_with_keepalives(self._lock, conn, 'Waiting for the model slot') try: if self.current_checkpoint_id != checkpoint_id: logger.info(f'Switching policy from {self.current_checkpoint_id} to {checkpoint_id}') @@ -79,8 +78,8 @@ async def get_policy(self, checkpoint_id: str, websocket: WebSocket | None = Non while self.active_sessions > 0: message = f'Waiting for {self.active_sessions} active session(s) to finish...' logger.info(message) - if websocket: - await websocket.send_bytes( + if conn: + await conn.send( serialise({protocol.STATUS: protocol.ServerStatus.WAITING, protocol.MESSAGE: message}) ) @@ -96,8 +95,8 @@ async def get_policy(self, checkpoint_id: str, websocket: WebSocket | None = Non self.current_policy = None self.current_checkpoint_id = None - if websocket: - await websocket.send_bytes( + if conn: + await conn.send( serialise({ protocol.STATUS: protocol.ServerStatus.LOADING, protocol.MESSAGE: f'Loading checkpoint {checkpoint_id}...', @@ -105,34 +104,31 @@ async def get_policy(self, checkpoint_id: str, websocket: WebSocket | None = Non ) logger.info(f'Loading policy {checkpoint_id}') - on_progress = self._progress_callback(websocket) + on_progress = self._progress_callback(conn) self.current_policy = await asyncio.to_thread(self._source.load, checkpoint_id, on_progress) self.current_checkpoint_id = checkpoint_id assert self.current_policy is not None - if websocket: + if conn: self.active_sessions += 1 return self.current_policy finally: self._lock.release() @staticmethod - def _progress_callback(websocket: WebSocket | None) -> Callable[[str], None] | None: + def _progress_callback(conn: wire.ServerConnection | None) -> Callable[[str], None] | None: """Sync callback for the loader thread, marshaling ``loading`` messages onto the event loop. Blocks the loader until each message is on the wire, so one emitted at the very end of a load cannot overtake the ``ready`` that follows it and be read as the first inference result. """ - if websocket is None: + if conn is None: return None loop = asyncio.get_running_loop() def on_progress(msg: str) -> None: asyncio.run_coroutine_threadsafe( - websocket.send_bytes( - serialise({protocol.STATUS: protocol.ServerStatus.LOADING, protocol.MESSAGE: msg}) - ), - loop, + conn.send(serialise({protocol.STATUS: protocol.ServerStatus.LOADING, protocol.MESSAGE: msg})), loop ).result() return on_progress @@ -187,16 +183,20 @@ class PolicyServer: ``local_stack`` spec in the ``ready`` handshake for the rig to build, alongside the marker's own wire settings. The source is the only model loader and is fixed at launch. - When ``pipeline`` is a ``cfn.Config``, query params on the session websocket URL become dotted + When ``pipeline`` is a ``cfn.Config``, query params on the session URL become dotted overrides into the pipeline config (e.g. ``?codec.fps=10``), applied and instantiated per session. Values must be JSON literals (unparseable values pass through as strings) and are applied with ``Config.override_data``, so a param can tune an argument but never name a Python object to import; params that change the model source are rejected too. A server built from an already-instantiated ``Pipeline`` rejects all session params. - The WebSocket session flow is: + The session flow is: accept → session params → resolve → load via manager → remote-half wrap → reset → inference loop + A session runs over one of two wires (see ``positronic.offboard.wire``): the websocket, on ``port``, + and gRPC, on ``grpc_port``. Both carry the same frames, so the flow above is the same on each. The + HTTP routes stay on ``port``; a ``grpc_port`` of ``None`` serves the websocket alone. + On startup (before accepting connections): resolve(None) → load. The default checkpoint is resolved once, at startup, and pinned for every request that names no @@ -212,6 +212,7 @@ def __init__( recording_dir: str | None = None, idle_timeout_min: float | None = None, auth_token: str | None = None, + grpc_port: int | None = None, ): self._pipeline_cfg = pipeline if isinstance(pipeline, cfn.Config) else None self._pipeline = pipeline.instantiate() if isinstance(pipeline, cfn.Config) else pipeline @@ -226,6 +227,7 @@ def __init__( self._manager = PolicyManager(self._source) self.host = host self.port = port + self.grpc_port = grpc_port self.metadata: dict[str, Any] = {offboard_keys.HOST: host, offboard_keys.PORT: port} # Synced once; each session builds its own ``Recorder`` so concurrent streams never mix. self._recording_dir = pos3.sync(recording_dir) if recording_dir else None @@ -249,10 +251,10 @@ def __init__( self.app = FastAPI() http_auth, ws_auth = [Depends(self._require_http_auth)], [Depends(self._require_ws_auth)] self.app.get('/api/v1/models', dependencies=http_auth)(self.get_models) - self.app.websocket('/api/v1/session', dependencies=ws_auth)(self.default_session) + self.app.websocket(wire.SESSION_PATH, dependencies=ws_auth)(self.default_session) # ``:path`` so an id that is itself a path (a HuggingFace repo, say) opens under the name # ``/api/v1/models`` advertises. - self.app.websocket('/api/v1/session/{model_id:path}', dependencies=ws_auth)(self.model_session) + self.app.websocket(f'{wire.SESSION_PATH}/{{model_id:path}}', dependencies=ws_auth)(self.model_session) def _authorized(self, authorization: str | None) -> bool: if self._auth_token is None: @@ -294,27 +296,32 @@ def _session_pipeline(self, params: dict[str, Any]) -> Pipeline: async def default_session(self, websocket: WebSocket): """Serves the model pinned at startup. Naming a model is the path's job, so every query param here is a pipeline override.""" - await self._serve_session(websocket, None) + await websocket.accept() + await self._serve_session(wire.WebsocketServerConnection(websocket), None) async def model_session(self, websocket: WebSocket, model_id: str): - await self._serve_session(websocket, model_id) - - async def _serve_session(self, websocket: WebSocket, model_id: str | None): await websocket.accept() - logger.info(f'Connected to {websocket.client} requesting {model_id or "default"}') + await self._serve_session(wire.WebsocketServerConnection(websocket), model_id) + + async def grpc_session(self, conn: grpc_wire.GrpcServerConnection): + """Serves one gRPC session, on the model the session path names.""" + await self._serve_session(conn, grpc_wire.model_id_of(conn.session_path)) + + async def _serve_session(self, conn: wire.ServerConnection, model_id: str | None): + logger.info(f'Connected to {conn.peer} requesting {model_id or "default"}') self._active_sessions += 1 self._last_activity = time.monotonic() policy: Policy | None = None session = None try: - pipeline = self._session_pipeline(_session_params(websocket.query_params)) + pipeline = self._session_pipeline(_session_params(conn.query_params)) local, border, remote_half = split(pipeline) local_spec = _declared_stack(local) rid = self._source.resolve(model_id) if model_id is not None else self._default_id assert rid is not None - policy = await self._manager.get_policy(rid, websocket) + policy = await self._manager.get_policy(rid, conn) # A request has no control loop to answer ``None`` to. This goes innermost, so every layer # above it sees one call per answer rather than one per call the answer took. answered = blocking(policy) @@ -329,7 +336,7 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): served = remote_half.wrap(answered) if remote_half is not None else answered # ``new_session`` resets the shared backend client, so it must not interleave with an in-flight # inference. Keepalives here: queuing behind a peer would otherwise trip the handshake timeout. - await _acquire_with_keepalives(self._infer_lock, websocket, 'Waiting for inference slot') + await _acquire_with_keepalives(self._infer_lock, conn, 'Waiting for inference slot') try: session = await asyncio.to_thread(served.new_session) finally: @@ -345,11 +352,11 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): offboard_keys.COMPRESS_IMAGES: border.compress_images, offboard_keys.POSITRONIC_VERSION: _pkg_version('positronic'), } - await websocket.send_bytes(serialise({protocol.STATUS: protocol.ServerStatus.READY, protocol.META: meta})) + await conn.send(serialise({protocol.STATUS: protocol.ServerStatus.READY, protocol.META: meta})) try: while True: - message = await websocket.receive_bytes() + message = await conn.receive() self._last_activity = time.monotonic() try: raw_obs = deserialise(message) @@ -358,20 +365,18 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): async with self._infer_lock: # The server's clock is not the rig's. actions = await asyncio.to_thread(session, raw_obs, time.time_ns()) - await websocket.send_bytes(serialise({protocol.RESULT: actions})) + await conn.send(serialise({protocol.RESULT: actions})) except Exception as e: logger.error(f'Error processing message: {e}', exc_info=True) - await websocket.send_bytes(serialise({protocol.ERROR: str(e)})) - except WebSocketDisconnect: + await conn.send(serialise({protocol.ERROR: str(e)})) + except wire.PeerDisconnected: logger.info('Client disconnected') except Exception as e: logger.error(f'Failed session: {e}', exc_info=True) try: - await websocket.send_bytes( - serialise({protocol.STATUS: protocol.ServerStatus.ERROR, protocol.ERROR: str(e)}) - ) - await websocket.close(code=1008, reason=str(e)[:100]) + await conn.send(serialise({protocol.STATUS: protocol.ServerStatus.ERROR, protocol.ERROR: str(e)})) + await conn.refuse(str(e)) except Exception: logger.debug('Failed to send error to client', exc_info=True) finally: @@ -407,6 +412,15 @@ async def _idle_watchdog(self, server: uvicorn.Server): server.should_exit = True return + async def _start_grpc(self) -> grpc.aio.Server: + """Start the gRPC wire on ``grpc_port``, sharing this server's model slot and inference lock.""" + assert self.grpc_port is not None + + def authorized(headers: Mapping[str, str]) -> bool: + return self._authorized(headers.get(AUTH_HEADER.lower())) + + return await grpc_wire.serve(self.grpc_session, authorized, self.host, self.grpc_port) + def serve(self): async def _run(): await self._startup() @@ -414,6 +428,7 @@ async def _run(): server = uvicorn.Server(config) self._last_activity = time.monotonic() watchdog = None + grpc_server = await self._start_grpc() if self.grpc_port is not None else None if self.idle_timeout_min and self.idle_timeout_min > 0: watchdog = asyncio.create_task(self._idle_watchdog(server)) try: @@ -421,6 +436,8 @@ async def _run(): finally: if watchdog is not None: watchdog.cancel() + if grpc_server is not None: + await grpc_server.stop(grace=None) try: asyncio.run(_run()) @@ -430,14 +447,24 @@ async def _run(): self._manager.close() -@cfn.config(host='0.0.0.0', port=8000, recording_dir=None, idle_timeout_min=None) -def serve(pipeline: cfn.Config, host: str, port: int, recording_dir: str | None, idle_timeout_min: float | None): +@cfn.config(host='0.0.0.0', port=8000, recording_dir=None, idle_timeout_min=None, grpc_port=None) +def serve( + pipeline: cfn.Config, + host: str, + port: int, + recording_dir: str | None, + idle_timeout_min: float | None, + grpc_port: int | None, +): """The CLI entry point every vendor server exposes: bind ``pipeline``, and the commands are configs of this. - Only the socket and the recording taps are flags of their own; everything the served model is — + Only the sockets and the recording taps are flags of their own; everything the served model is — codec, source, checkpoint directory — is reached through the pipeline itself (``--pipeline.source.checkpoints_dir=...``), so each of those values has exactly one name. + ``grpc_port`` adds the gRPC wire beside the websocket one. HTTP/2 needs the frames a websocket-only + front drops, so give it a port a client reaches directly. + The bearer token gating the server comes from ``AUTH_TOKEN_ENV`` rather than a flag, which would put a secret in the process arguments; unset serves open. """ @@ -448,4 +475,5 @@ def serve(pipeline: cfn.Config, host: str, port: int, recording_dir: str | None, recording_dir=recording_dir, idle_timeout_min=idle_timeout_min, auth_token=os.environ.get(AUTH_TOKEN_ENV), + grpc_port=grpc_port, ).serve() diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 54dfcbb7a..d1bd5dd29 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -26,16 +26,26 @@ def _find_free_port() -> int: @pytest.fixture def start_server() -> Generator[StartServer, None, None]: - """Factory serving pipelines on daemon threads; every started server is stopped and joined at teardown.""" + """Factory serving pipelines on daemon threads; every started server is stopped and joined at teardown. + + ``grpc=True`` also serves the gRPC wire, on a port of its own that ``PolicyServer.grpc_port`` names. + """ running: list[tuple[uvicorn.Server, threading.Thread]] = [] - def start(pipeline, **server_kwargs) -> tuple[str, int, PolicyServer]: - server = PolicyServer(pipeline, host='localhost', port=_find_free_port(), **server_kwargs) + def start(pipeline, *, grpc: bool = False, **server_kwargs) -> tuple[str, int, PolicyServer]: + grpc_port = _find_free_port() if grpc else None + server = PolicyServer(pipeline, host='localhost', port=_find_free_port(), grpc_port=grpc_port, **server_kwargs) uv_server = uvicorn.Server(uvicorn.Config(server.app, host=server.host, port=server.port, log_level='warning')) async def _run(): await server._startup() - await uv_server.serve() + # Started first, so the websocket port answering means both wires are up. + grpc_server = await server._start_grpc() if grpc_port is not None else None + try: + await uv_server.serve() + finally: + if grpc_server is not None: + await grpc_server.stop(grace=None) thread = threading.Thread(target=asyncio.run, args=(_run(),), daemon=True) thread.start() @@ -104,7 +114,7 @@ def make_mock_policy() -> Callable[..., MagicMock]: return _make_mock_policy -class _DictSource(ModelSource): +class DictSource(ModelSource): """Multi-model source over ready policies; the dict's first key is the default.""" def __init__(self, policies: Mapping[str, Policy]): @@ -153,5 +163,5 @@ def inference_server(start_server: StartServer, mock_policy: MagicMock) -> tuple def multi_policy_server( start_server: StartServer, mock_policy_registry: dict[str, MagicMock] ) -> tuple[str, int, dict[str, MagicMock]]: - host, port, _server = start_server(ChunkedSchedule() | remote | _DictSource(mock_policy_registry)) + host, port, _server = start_server(ChunkedSchedule() | remote | DictSource(mock_policy_registry)) return host, port, mock_policy_registry diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py new file mode 100644 index 000000000..5837abaa6 --- /dev/null +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -0,0 +1,174 @@ +"""The gRPC wire: one session runs over it exactly as it runs over the websocket.""" + +from unittest.mock import ANY, MagicMock + +import configuronic as cfn +import grpc +import pytest + +from positronic.offboard import grpc_wire, wire +from positronic.offboard.client import InferenceClient, _ConnectRetries +from positronic.offboard.server import AUTH_HEADER, PolicyServer, bearer +from positronic.offboard.tests.conftest import DictSource, StartServer +from positronic.policy.layers import ChunkedSchedule, TemporalStack +from positronic.policy.spec import ModelSource, PolicySource, remote + +_TOKEN = 'test-secret-token' + + +def grpc_url(server: PolicyServer, path: str = '') -> str: + return f'grpc://{server.host}:{server.grpc_port}{path}' + + +@pytest.fixture +def both_wires(start_server: StartServer, make_mock_policy) -> tuple[PolicyServer, MagicMock]: + """A server offering both wires over one policy.""" + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + _host, _port, server = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True) + return server, policy + + +def test_a_grpc_session_handshakes_and_infers(both_wires): + server, policy = both_wires + session = InferenceClient(grpc_url(server)).new_session() + try: + assert session.metadata['model_name'] == 'stub' + obs = {'image': 'test'} + assert session.infer(obs) == [{'action': [1, 2, 3]}] + policy._mock_session.assert_called_with(obs, ANY) + finally: + session.close() + + +def test_both_wires_answer_one_observation_alike(both_wires): + server, _policy = both_wires + obs = {'image': 'test'} + over_ws = InferenceClient(f'{server.host}:{server.port}').new_session() + over_grpc = InferenceClient(grpc_url(server)).new_session() + try: + assert over_grpc.metadata == over_ws.metadata + assert over_grpc.infer(obs) == over_ws.infer(obs) + finally: + over_ws.close() + over_grpc.close() + + +def test_closing_a_session_ends_it_on_the_server(both_wires): + """``close`` half-closes the stream and waits, so the server releases the session before it returns.""" + server, _policy = both_wires + session = InferenceClient(grpc_url(server)).new_session() + assert server._active_sessions == 1 + session.close() + assert server._active_sessions == 0 + + +def test_a_failed_inference_reaches_the_client_as_an_exception(both_wires): + server, policy = both_wires + session = InferenceClient(grpc_url(server)).new_session() + try: + policy._mock_session.side_effect = RuntimeError('no such joint') + with pytest.raises(RuntimeError, match='no such joint'): + session.infer({'image': 'test'}) + finally: + session.close() + + +def test_a_session_that_cannot_open_reaches_the_client_as_an_exception(start_server, make_mock_policy): + """A model the source refuses fails in the handshake, before the session serves anything.""" + policies = {'alpha': make_mock_policy([{'action': [1]}], {'model_name': 'alpha'})} + _host, _port, server = start_server(ChunkedSchedule() | remote | DictSource(policies), grpc=True) + with pytest.raises(RuntimeError, match='Unknown model'): + InferenceClient(grpc_url(server, f'{wire.SESSION_PATH}/beta')).new_session() + + +def test_the_session_path_names_the_model(start_server, make_mock_policy): + policies = { + 'alpha': make_mock_policy([{'action': ['alpha']}], {'model_name': 'alpha'}), + 'beta': make_mock_policy([{'action': ['beta']}], {'model_name': 'beta'}), + } + _host, _port, server = start_server(ChunkedSchedule() | remote | DictSource(policies), grpc=True) + session = InferenceClient(grpc_url(server, f'{wire.SESSION_PATH}/beta')).new_session() + try: + assert session.metadata['model_name'] == 'beta' + assert session.infer({'obs': 'beta'}) == [{'action': ['beta']}] + finally: + session.close() + + +def _tunable_pipe(source: ModelSource, offsets: tuple[float, ...] = (-0.1, 0.0)): + return TemporalStack(keys=('x',), offsets_sec=offsets) | ChunkedSchedule() | remote | source + + +def test_the_query_carries_the_session_params(start_server, make_mock_policy): + policies = {'alpha': make_mock_policy([{'action': ['alpha']}], {'model_name': 'alpha'})} + pipe = cfn.Config(_tunable_pipe, source=cfn.Config(DictSource, policies=policies)) + _host, _port, server = start_server(pipe, grpc=True) + session = InferenceClient(grpc_url(server, f'{wire.SESSION_PATH}?offsets=[-0.5, 0.0]')).new_session() + try: + assert session.metadata['local_stack']['seq'][0]['args']['offsets_sec'] == [-0.5, 0.0] + finally: + session.close() + + +@pytest.fixture +def authed_server(start_server: StartServer, make_mock_policy) -> PolicyServer: + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + _host, _port, server = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True, auth_token=_TOKEN) + return server + + +def test_the_grpc_wire_gates_on_the_bearer_token(authed_server): + session = InferenceClient(grpc_url(authed_server), headers={AUTH_HEADER: bearer(_TOKEN)}).new_session() + try: + assert session.metadata['model_name'] == 'stub' + finally: + session.close() + + +@pytest.mark.parametrize('header', [None, bearer('wrong'), _TOKEN]) +def test_the_grpc_wire_refuses_a_session_without_the_token(authed_server, header, monkeypatch): + # A refused credential and a cold backend answer alike, so the client spends attempts on it; one + # is enough to see the refusal. + monkeypatch.setattr(_ConnectRetries, 'MAX_FORBIDDEN_ATTEMPTS', 1) + headers = None if header is None else {AUTH_HEADER: header} + with pytest.raises(grpc.RpcError) as refused: + InferenceClient(grpc_url(authed_server), headers=headers).new_session() + assert refused.value.code() is grpc.StatusCode.PERMISSION_DENIED + + +@pytest.mark.parametrize('url', ['grpcs://gpu-host:9000', 'tcp://gpu-host:9000']) +def test_an_unknown_scheme_is_refused(url): + with pytest.raises(ValueError, match='Unsupported scheme'): + InferenceClient(url) + + +def test_a_grpc_url_names_the_session_port_alone(): + client = InferenceClient('grpc://gpu-host:9000') + assert client.session_url == 'grpc://gpu-host:9000/api/v1/session' + with pytest.raises(ValueError, match='gRPC session port'): + client.list_models() + + +@pytest.mark.parametrize( + ('session_path', 'model_id'), + [ + (wire.SESSION_PATH, None), + (f'{wire.SESSION_PATH}/10000', '10000'), + (f'{wire.SESSION_PATH}/GEAR-Dreams/DreamZero-DROID', 'GEAR-Dreams/DreamZero-DROID'), + (f'{wire.SESSION_PATH}/s3%3A//bucket/ckpt-1', 's3://bucket/ckpt-1'), + ], +) +def test_the_session_path_decodes_as_the_websocket_route_does(session_path, model_id): + assert grpc_wire.model_id_of(session_path) == model_id + + +def test_a_path_outside_the_session_route_is_refused(): + with pytest.raises(ValueError, match='Unexpected session path'): + grpc_wire.model_id_of('/api/v2/session/10000') + + +def test_a_port_that_never_answers_is_named_at_the_deadline(): + """Nothing listens on port 1, so the channel never becomes ready and the connect deadline passes.""" + client = InferenceClient('grpc://localhost:1', open_timeout=0.2, connect_deadline=0.0) + with pytest.raises(TimeoutError, match='grpc://localhost:1'): + client.new_session() diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 9f2f6201e..060644a38 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -102,7 +102,9 @@ def test_new_session_passes_additional_headers(self): 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) + conn = mock_session_cls.call_args.args[0] + assert conn._websocket is mock_connect.return_value + assert mock_session_cls.call_args.kwargs['infer_timeout'] == DEFAULT_INFER_TIMEOUT def test_new_session_without_headers_passes_none(self): with ( diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 2bf3bea0a..6bad1fd82 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -19,6 +19,7 @@ from positronic.offboard.server import AUTH_HEADER, AUTH_TOKEN_ENV, PolicyServer, bearer from positronic.offboard.server_utils import warmup from positronic.offboard.tests.conftest import round_trip +from positronic.offboard.wire import WebsocketClientConnection from positronic.policy import Codec, Policy, RemotePolicy, Session from positronic.policy.base import Runtime from positronic.policy.codec import ActionTimestamp @@ -304,7 +305,7 @@ def _tunable_pipe(source: ModelSource, offsets: tuple[float, ...] = (-0.1, 0.0), def _param_session(host: str, port: int, query: list[tuple[str, str]]) -> InferenceSession: uri = f'ws://{host}:{port}/api/v1/session?' + urllib.parse.urlencode(query) - return InferenceSession(connect(uri)) + return InferenceSession(WebsocketClientConnection(connect(uri))) @pytest.fixture @@ -359,7 +360,8 @@ def test_model_id_is_named_by_path_not_query(param_server): with pytest.raises(RuntimeError, match='model_id'): _param_session(host, port, [('model_id', 'other')]) - session = InferenceSession(connect(f'ws://{host}:{port}/api/v1/session/other?pad_start=false')) + uri = f'ws://{host}:{port}/api/v1/session/other?pad_start=false' + session = InferenceSession(WebsocketClientConnection(connect(uri))) try: assert session.metadata['checkpoint_id'] == 'other' assert session.metadata['local_stack']['seq'][0]['args']['pad_start'] is False @@ -484,7 +486,9 @@ def test_session_outlives_an_idle_ingress_window(authed_endpoint): session = InferenceClient(url, headers={AUTH_HEADER: bearer(token)}).new_session() try: time.sleep(_IDLE_WINDOW_SEC) - assert session._websocket.ping().wait(timeout=30.0) + conn = session._conn + assert isinstance(conn, WebsocketClientConnection), "the idle window is the websocket wire's" + assert conn._websocket.ping().wait(timeout=30.0) finally: session.close() diff --git a/positronic/offboard/wire.py b/positronic/offboard/wire.py new file mode 100644 index 000000000..cdf770cef --- /dev/null +++ b/positronic/offboard/wire.py @@ -0,0 +1,112 @@ +"""The transports one session runs over, and the two ends of an open one. + +A wire carries the ``protocol`` frames as opaque bytes and reads none of them, so the handshake and +the inference loop read the same over every wire. ``grpc_wire`` holds the gRPC one. +""" + +import abc +from typing import Protocol + +from fastapi import WebSocket, WebSocketDisconnect +from starlette.datastructures import QueryParams +from websockets.sync.connection import Connection + +# The route a session opens on. The websocket wire puts it in the URL; the gRPC wire names it in the +# session metadata, so both wires address a model the same way. +SESSION_PATH = '/api/v1/session' + + +class PeerDisconnected(Exception): + """The peer ended the session.""" + + +class ClientConnection(Protocol): + """A client's end of one open session.""" + + def send(self, message: bytes) -> None: ... + + def recv(self, timeout: float | None = None) -> bytes: + """The next message. Raises ``TimeoutError`` when none arrives in time.""" + ... + + def close(self) -> str: + """Close this end, and report what the wire saw, for the log. + + A peer that answered the close leaves a different trace from one that had already gone while the + server still held the session, and the second is what strands the next session's handshake. Only + the wire can tell the two apart, and each says it in its own terms. + """ + ... + + +class WebsocketClientConnection: + """A client's end of one websocket session.""" + + def __init__(self, websocket: Connection): + self._websocket = websocket + + def send(self, message: bytes) -> None: + self._websocket.send(message) + + def recv(self, timeout: float | None = None) -> bytes: + message = self._websocket.recv(timeout=timeout) + assert isinstance(message, bytes), f'A frame is bytes, and this one is {type(message).__name__}' + return message + + def close(self) -> str: + state_before_close = self._websocket.state.name + self._websocket.close() + # A close that times out still reaches CLOSED locally; only the close code says the server answered. + return f'state {state_before_close} -> {self._websocket.state.name}, close code {self._websocket.close_code}' + + +class ServerConnection(abc.ABC): + """A server's end of one open session.""" + + @property + @abc.abstractmethod + def peer(self) -> str: + """Whom this session serves, for the log.""" + + @property + @abc.abstractmethod + def query_params(self) -> QueryParams: + """The session params the client asked for.""" + + @abc.abstractmethod + async def send(self, message: bytes) -> None: ... + + @abc.abstractmethod + async def receive(self) -> bytes: + """The next message. Raises ``PeerDisconnected`` once the client ends the session.""" + + @abc.abstractmethod + async def refuse(self, reason: str) -> None: + """End a session the server cannot serve, telling the client why.""" + + +class WebsocketServerConnection(ServerConnection): + """A server's end of one websocket session, over an accepted ``WebSocket``.""" + + def __init__(self, websocket: WebSocket): + self._websocket = websocket + + @property + def peer(self) -> str: + return str(self._websocket.client) + + @property + def query_params(self) -> QueryParams: + return self._websocket.query_params + + async def send(self, message: bytes) -> None: + await self._websocket.send_bytes(message) + + async def receive(self) -> bytes: + try: + return await self._websocket.receive_bytes() + except WebSocketDisconnect as e: + raise PeerDisconnected(str(e)) from e + + async def refuse(self, reason: str) -> None: + await self._websocket.close(code=1008, reason=reason[:100]) diff --git a/pyproject.toml b/pyproject.toml index 281f7640e..80b93a584 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,9 @@ dependencies = [ "starlette>=1.0.0", "jinja2", "fire", + # The gRPC wire, one of the two transports `positronic.offboard` serves. The rig client and the + # server both import it, so it is not an extra. + "grpcio", "httpx", "msgpack", "mujoco", diff --git a/uv.lock b/uv.lock index 8e162a212..ec2e0c472 100644 --- a/uv.lock +++ b/uv.lock @@ -2264,6 +2264,47 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cc/9b/4366ad3e1c0688146c70aa6143584d6a8d88583b9390f106250e25a3d5cd/glfw-2.10.0-py2.py3-none-win_amd64.whl", hash = "sha256:7f787ee8645781f10e8800438ce4357ab38c573ffb191aba380c1e72eba6311c", size = 559423, upload-time = "2026-03-10T17:21:34.766Z" }, ] +[[package]] +name = "grpcio" +version = "1.82.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/90/bc/656b89387d6f4ed7e0686c7b64c2ae7e554a759aa58122c8e5fb99392c32/grpcio-1.82.1.tar.gz", hash = "sha256:707b24abd90fcb1e45bcc080577da1dbf9971d107490589b9539af8e1e77b4b5", size = 13187300, upload-time = "2026-07-08T12:36:16.588Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/26/5b/e5092af97fa671ca279b3e373251af4bf87d5fbda7dc85f6a616899562a7/grpcio-1.82.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:0ddb18a9a9e1f46692b3567ae4abb3f8d117ce6afea48650f8eca06d8ab5d06f", size = 6181472, upload-time = "2026-07-08T12:34:31.009Z" }, + { url = "https://files.pythonhosted.org/packages/c1/8f/18053a3a2ca03d0c2a1b8cc7271e705007a16aa5dae84bac00935c5b1a7f/grpcio-1.82.1-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:cf855b1af246720f567b0ce5d0724d45dfa4188eecc3296a2a69257b11b9e94b", size = 11970995, upload-time = "2026-07-08T12:34:33.603Z" }, + { url = "https://files.pythonhosted.org/packages/5b/7e/21b1acb052876ad00959ec4d1b05fe08607d650bcfa282073bb164c2703c/grpcio-1.82.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ddb30cb13e25bc13cea70ffc69d6d90c49d36ea6c1d4549e6912f70177834cac", size = 6760127, upload-time = "2026-07-08T12:34:36.122Z" }, + { url = "https://files.pythonhosted.org/packages/3e/12/25eef9c245c54f0061317d13a302357fe8ea03bac240b2b02ececcf54da4/grpcio-1.82.1-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:1e822b2774f719c017cbe700b6e47173b6ae290fb84906f52a5a3c2c60b62e1e", size = 7484377, upload-time = "2026-07-08T12:34:38.368Z" }, + { url = "https://files.pythonhosted.org/packages/a0/41/1a348767eb9d9bd7765dc4fa8a01723d3bb386d67f981ee5c6f9c02b8b1c/grpcio-1.82.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5dafb1ece8ed45dee7c738f166ec82e19673221ed5ab8967f72858a4685345b2", size = 6924269, upload-time = "2026-07-08T12:34:40.583Z" }, + { url = "https://files.pythonhosted.org/packages/e4/b9/3aae7a03d34c86ea27988db859a6087c186f6c3f53f9b551e07afd989bfa/grpcio-1.82.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e06503106e7271e0a49fd5a1ac04747f1e47e87d900476db6fe45bc87ee411f4", size = 7531848, upload-time = "2026-07-08T12:34:43.277Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/3c4afa625d0dac9090707966916284c035fc5b2fb3e2c51e156accee6735/grpcio-1.82.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:ff99bc8cafb6a952201c37b995f425e641c93ffa6e072258525feab57290141d", size = 8568217, upload-time = "2026-07-08T12:34:45.502Z" }, + { url = "https://files.pythonhosted.org/packages/20/9c/d8489c628e73e20a3d034e7f66912de7b1acb405f01d388f056a88e47924/grpcio-1.82.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:644ae1b94266ac785330f4590a69e52b6a7eb73029043a02209db81c81397d69", size = 7938771, upload-time = "2026-07-08T12:34:48.323Z" }, + { url = "https://files.pythonhosted.org/packages/4b/b7/0a92cfd1658f3a896d4aa12d4efeb7dd4ddfc723725ae22741a5241ea710/grpcio-1.82.1-cp311-cp311-win32.whl", hash = "sha256:e203d2e19d471630084a16c815616f8211dff21c268ab3c5f5bf38417832e074", size = 4256432, upload-time = "2026-07-08T12:34:50.432Z" }, + { url = "https://files.pythonhosted.org/packages/c7/6a/2872c761b025d9ec74386f22a4a7d59c5a5b00ebf718761b33739ffc45de/grpcio-1.82.1-cp311-cp311-win_amd64.whl", hash = "sha256:0d8299c285fe6cc6a1f56badf8d3bc5078c8d20273ee64bafa3783b4bc29a769", size = 5009633, upload-time = "2026-07-08T12:34:52.67Z" }, + { url = "https://files.pythonhosted.org/packages/dc/88/d1350bf3343a2ed87d801584e40609f6c6bd3087926eeca03de50348cf4a/grpcio-1.82.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:c09bd5fa0d5b1fbd773ec349fe61441c3e4ebf168c229aa7538a820bdfad6a58", size = 6144689, upload-time = "2026-07-08T12:34:55.567Z" }, + { url = "https://files.pythonhosted.org/packages/e6/33/71875cdecd27c24ac1385d4783a09853f01b84a825a36aec2a2bc7d0d080/grpcio-1.82.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:1eae24810720734598e3e6a1a528d5de0f265fe3fc86575e9ecce424b9ec7379", size = 11952034, upload-time = "2026-07-08T12:34:58.128Z" }, + { url = "https://files.pythonhosted.org/packages/82/b2/d9125df3d8a140dec12cc82c05b7deafedeababcff6496f28b2fd5634d10/grpcio-1.82.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a6bd5daf5bde7b24d7ad2cbaf8bf9eac620d96222016bb5e7ddde930dec0673f", size = 6710772, upload-time = "2026-07-08T12:35:01.33Z" }, + { url = "https://files.pythonhosted.org/packages/88/9b/69e2d1627398b964f34437dc476a5aff5a2cc8e7f247d26272b5674b5faf/grpcio-1.82.1-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:1ecfde669cb687ac020d31ff76debe5dc7a62213335f02262eb6625628da1c03", size = 7450677, upload-time = "2026-07-08T12:35:03.926Z" }, + { url = "https://files.pythonhosted.org/packages/e4/e7/8f855ca29c294956122a2a73023655b9b02602d5111dad2b9b00e7631c68/grpcio-1.82.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:011c8badee95734dee8bf05ce3464756a0ac3ebb8d443afd20c0e2b5e4640ad9", size = 6886855, upload-time = "2026-07-08T12:35:06.174Z" }, + { url = "https://files.pythonhosted.org/packages/5f/f0/fa87e85f49925f44c479d07e58b051e69bcfef6b6d5fbc6749d140f6730a/grpcio-1.82.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b85f4564926fb23114d239392bdcae200db1e6179629edd7d7ab0ab89c96a197", size = 7501323, upload-time = "2026-07-08T12:35:08.49Z" }, + { url = "https://files.pythonhosted.org/packages/67/55/2e0b10ae1d3ef9dcc480b91dc2158f4931fc4675d3af0a2836e39b2a744f/grpcio-1.82.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:2c0c8270833395644c3fe6b6a806397955a2bc0538000a19a78b90c05a6c16e0", size = 8536899, upload-time = "2026-07-08T12:35:10.975Z" }, + { url = "https://files.pythonhosted.org/packages/bd/95/a3d8b0431fa221efc51ee39d73595ede74ba82a43b7c4313192e580face2/grpcio-1.82.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2ba199205ff46c7778290fe1673c91ac8e7e45678dd5c86e9e56fa33ec8788f6", size = 7913892, upload-time = "2026-07-08T12:35:13.944Z" }, + { url = "https://files.pythonhosted.org/packages/b8/92/f2651ec704d9852a56faef394775038afba435b50ce82ab2404d119c3355/grpcio-1.82.1-cp312-cp312-win32.whl", hash = "sha256:06127691866e295c14e84a1fb86356dd962254f6abd0da4ca4b001eea9e89438", size = 4240985, upload-time = "2026-07-08T12:35:16.048Z" }, + { url = "https://files.pythonhosted.org/packages/96/4f/a5fe8bf0d0a1b24855f370293075c931f27de4eb55f0f158786095bf3c11/grpcio-1.82.1-cp312-cp312-win_amd64.whl", hash = "sha256:1fa3223a3a2e1db74f4c2b255189eb7ea875dfba56e221d252ee3fc7b204778e", size = 5001580, upload-time = "2026-07-08T12:35:18.689Z" }, + { url = "https://files.pythonhosted.org/packages/1b/3e/496992d08c0aaa11272eb6228dc8ab947da01fe835de243cd00521bce4c4/grpcio-1.82.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:b454a2d97bfab7565683a02345f86bd182ab69fd7c2bdb7414171e7538f266b1", size = 6146068, upload-time = "2026-07-08T12:35:21.365Z" }, + { url = "https://files.pythonhosted.org/packages/e7/8f/f263d6f14fdba6b56cfadd91fd3e158a52682b72c6016d1f8723d435659f/grpcio-1.82.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:3dde70abfc80b3be11de53ba0d601c439e7fb2afd3583ad1788d1146bec92fdc", size = 11948600, upload-time = "2026-07-08T12:35:24.312Z" }, + { url = "https://files.pythonhosted.org/packages/8c/14/3a02e6ee49c2d85bc15eaae321e0e11ab3542cad3c5b2de121ecce0c4296/grpcio-1.82.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f5523099c98c292ea1ae08e617249db760c56a78f8deae879027fe7d1ffbcbf6", size = 6714591, upload-time = "2026-07-08T12:35:27.027Z" }, + { url = "https://files.pythonhosted.org/packages/69/80/58e3738696f48ab7645347b98d8a7f93d10e00e6218388fbfcd6c9310e3d/grpcio-1.82.1-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:5e5c4dc0a59b0f8490a6bdfd6fc8395b9d8ad8a8407c7d67ca7b5bba15c0877f", size = 7454995, upload-time = "2026-07-08T12:35:29.599Z" }, + { url = "https://files.pythonhosted.org/packages/f3/6c/2557c1a889363072fbf2285ecd0e8c44860d4dbd60f017a32537c5b863e2/grpcio-1.82.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c40d94ba820329cc191981bc22fa6f6eed0799c6d921f3c6709521d59d4a2fd7", size = 6888621, upload-time = "2026-07-08T12:35:32.38Z" }, + { url = "https://files.pythonhosted.org/packages/d2/66/907706ccaff1223f1e10fd5b37fc16faead43392fccb4e786e7e390ac141/grpcio-1.82.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4c816180e31e273caaec6f8bd86a8392499d5bbb26f41da44e3dce48bde69095", size = 7505069, upload-time = "2026-07-08T12:35:35.072Z" }, + { url = "https://files.pythonhosted.org/packages/b3/7c/ff97b0d0f635987ee5ec80dfedafa1aad629303745d48e8637d10eec5b80/grpcio-1.82.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:e31fd780b261830720cb70b0fd8f0aa51d49e75a66d7464ad2e31d4b765f2580", size = 8535384, upload-time = "2026-07-08T12:35:37.954Z" }, + { url = "https://files.pythonhosted.org/packages/62/9e/a97fddd970a8d1588cade06eca20443761c1858b0ad6590a5c835aa18062/grpcio-1.82.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9d76152d7c31d7210d4a106e5d8b64da5bba5d6abf11be30e2f7b0a0c59bbcbf", size = 7910707, upload-time = "2026-07-08T12:35:40.797Z" }, + { url = "https://files.pythonhosted.org/packages/20/e4/eaba1517888af483a88d449eb7566f0f7f63446d46f339c5891798435875/grpcio-1.82.1-cp313-cp313-win32.whl", hash = "sha256:38e9dcb5258226fb3282630b31b16a968df52c8c6ad514af540646e0a4578f8a", size = 4240363, upload-time = "2026-07-08T12:35:43.298Z" }, + { url = "https://files.pythonhosted.org/packages/b0/42/66a98d47732e35290bef722f6149fed3709cd4cf61166f6f53a12f417302/grpcio-1.82.1-cp313-cp313-win_amd64.whl", hash = "sha256:3dbfb52c36d9511ac2b8e6c94fdde837b393ae520cc321f52a333a2deedf5a90", size = 5000980, upload-time = "2026-07-08T12:35:46.262Z" }, +] + [[package]] name = "gymnasium" version = "0.29.1" @@ -4898,6 +4939,7 @@ dependencies = [ { name = "dearpygui", marker = "platform_machine != 'aarch64' or (extra == 'extra-10-positronic-lerobot' and extra == 'extra-10-positronic-lerobot-0-3-3') or (extra == 'extra-10-positronic-lerobot' and extra == 'extra-10-positronic-molmoact2') or (extra == 'extra-10-positronic-lerobot' and extra == 'extra-10-positronic-yam') or (extra == 'extra-10-positronic-lerobot-0-3-3' and extra == 'extra-10-positronic-molmoact2') or (extra == 'extra-10-positronic-lerobot-0-3-3' and extra == 'extra-10-positronic-yam') or (extra == 'extra-10-positronic-molmoact2' and extra == 'extra-10-positronic-yam')" }, { name = "fastapi" }, { name = "fire" }, + { name = "grpcio" }, { name = "httpx" }, { name = "jinja2" }, { name = "msgpack" }, @@ -5009,6 +5051,7 @@ requires-dist = [ { name = "fastapi" }, { name = "feetech-servo-sdk", marker = "extra == 'hardware'" }, { name = "fire" }, + { name = "grpcio" }, { name = "httpx" }, { name = "huggingface-hub", marker = "extra == 'dreamzero'" }, { name = "i2rt", marker = "extra == 'yam'", git = "https://github.com/i2rt-robotics/i2rt.git?rev=5d47b358bafb30c65e397f2ece506550a0db4594" }, From c2487b3e6ecab3394ede78171c1eb2af91e76477 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 00:22:21 +0000 Subject: [PATCH 02/46] Open the endpoint-description line on its subject The writing rules refuse a sentence that opens on its reason. Ticket: none - a one-line prose fix the style gate named --- positronic/offboard/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index 832cd3b4b..19443bcc5 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -93,7 +93,7 @@ Rules: Any violation — including an unknown key — fails at connect: the server sends `{"status": "error", "error": ...}` and ends the session before anything moves, and the Python client raises `RuntimeError`. Overrides apply per session, and the `local_stack` declared in the ready handshake reflects them. -Because the whole session configuration fits in the URL, one string is a complete endpoint description: +One string is a complete endpoint description, because the whole session configuration fits in the URL: `--policy=.remote --policy.url='gpu-host:8000?codec.fps=10'` accepts `host`, `host:port`, and full `http(s)`/`ws(s)`/`grpc` URLs — optionally with `/api/v1/session/` — and forwards the query string verbatim. Credentials are the exception and stay a separate `headers` argument, so the URL itself is safe to hand around. From deba93f519b4986f3fe0fb10b86e0553162617a8 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 00:25:38 +0000 Subject: [PATCH 03/46] Tighten two comments the writing gate named Cut the cold-backend comment to the fact a reader acts on, and drop the HTTP/2 rationale the offboard README already carries. Ticket: none - a comment pass on the change in flight --- positronic/offboard/client.py | 9 +++------ positronic/offboard/server.py | 3 +-- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 5c2d15a75..2f479b7c3 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -259,12 +259,9 @@ def new_session(self) -> InferenceSession: # 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), it refuses the gRPC call (``RpcError``), or it accepts the - # connection and then drops or stalls the status handshake inside ``InferenceSession`` - # (``ConnectionClosed``/``PeerDisconnected``/``TimeoutError``). All mean "not ready yet", so retry - # within the deadline instead of letting one kill the run. + # Each of these is a backend that is not ready yet — a timed-out connect, a reset TLS + # handshake, a refused upgrade or gRPC call, a dropped status handshake — so one must not + # kill the run. ``_ConnectRetries`` decides which of them is the endpoint saying no. except ( TimeoutError, ssl.SSLError, diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 84ab72963..c312afd68 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -462,8 +462,7 @@ def serve( codec, source, checkpoint directory — is reached through the pipeline itself (``--pipeline.source.checkpoints_dir=...``), so each of those values has exactly one name. - ``grpc_port`` adds the gRPC wire beside the websocket one. HTTP/2 needs the frames a websocket-only - front drops, so give it a port a client reaches directly. + ``grpc_port`` adds the gRPC wire beside the websocket one (see the offboard README). The bearer token gating the server comes from ``AUTH_TOKEN_ENV`` rather than a flag, which would put a secret in the process arguments; unset serves open. From b73a9333a20b099005844d12024d7ee7e5d8a5c1 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 00:34:07 +0000 Subject: [PATCH 04/46] Bracket an IPv6 host in the address the gRPC wire binds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit gRPC's target syntax needs an IPv6 literal in brackets, so a server hosted on `::` bound `:::` — which gRPC refuses by returning port 0, and the server then started, accepted nothing and said nothing. The bind reports a refusal now instead of starting deaf. Ticket: none - a review finding on an unmerged branch --- positronic/offboard/grpc_wire.py | 11 ++++++++++- positronic/offboard/tests/conftest.py | 3 ++- positronic/offboard/tests/test_grpc_wire.py | 14 ++++++++++++++ 3 files changed, 26 insertions(+), 2 deletions(-) diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index d1b78184a..ad4671dbd 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -167,6 +167,11 @@ def _headers(context: grpc.aio.ServicerContext) -> dict[str, str]: return {key: value for key, value in (context.invocation_metadata() or ()) if isinstance(value, str)} +def _bind_target(host: str, port: int) -> str: + """The address to bind, with an IPv6 literal in the brackets gRPC's target syntax requires.""" + return f'[{host}]:{port}' if ':' in host else f'{host}:{port}' + + async def serve( serve_session: Callable[[GrpcServerConnection], Awaitable[None]], authorized: Callable[[Mapping[str, str]], bool], @@ -194,7 +199,11 @@ async def _serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerC handler = grpc.stream_stream_rpc_method_handler(_serve_one, request_deserializer=None, response_serializer=None) server = grpc.aio.server(options=_MESSAGE_SIZE_OPTIONS) server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler(SERVICE, {METHOD: handler}),)) - bound = server.add_insecure_port(f'{host}:{port}') + bound = server.add_insecure_port(_bind_target(host, port)) + if bound == 0: + # gRPC reports a refused bind by returning port 0, so a server left to start here would + # accept nothing and say nothing. + raise OSError(f'gRPC could not bind {_bind_target(host, port)}') await server.start() logger.info(f'gRPC sessions on {host}:{bound}') return server diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index d1bd5dd29..1b9b71f6e 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -34,7 +34,8 @@ def start_server() -> Generator[StartServer, None, None]: def start(pipeline, *, grpc: bool = False, **server_kwargs) -> tuple[str, int, PolicyServer]: grpc_port = _find_free_port() if grpc else None - server = PolicyServer(pipeline, host='localhost', port=_find_free_port(), grpc_port=grpc_port, **server_kwargs) + host = server_kwargs.pop('host', 'localhost') + server = PolicyServer(pipeline, host=host, port=_find_free_port(), grpc_port=grpc_port, **server_kwargs) uv_server = uvicorn.Server(uvicorn.Config(server.app, host=server.host, port=server.port, log_level='warning')) async def _run(): diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index 5837abaa6..ea65ac1ee 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -172,3 +172,17 @@ def test_a_port_that_never_answers_is_named_at_the_deadline(): client = InferenceClient('grpc://localhost:1', open_timeout=0.2, connect_deadline=0.0) with pytest.raises(TimeoutError, match='grpc://localhost:1'): client.new_session() + + +def test_an_ipv6_host_binds_in_brackets(start_server: StartServer, make_mock_policy): + """gRPC's target syntax brackets an IPv6 literal, so a bare '::1' would bind ':::' and fail.""" + assert grpc_wire._bind_target('::', 9000) == '[::]:9000' + assert grpc_wire._bind_target('0.0.0.0', 9000) == '0.0.0.0:9000' + + policy = make_mock_policy([{'action': [4]}], {'model_name': 'stub'}) + _host, _port, server = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True, host='::1') + session = InferenceClient(f'grpc://[{server.host}]:{server.grpc_port}').new_session() + try: + assert session.infer({'image': 'test'}) == [{'action': [4]}] + finally: + session.close() From 4ab68e4275b993b6b2687ff7dd7210142c9c2d54 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 00:43:58 +0000 Subject: [PATCH 05/46] Pin the frame ceiling for both wires, and close a connection the handshake refuses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three findings from the review round. The 16 MiB ceiling was gRPC's alone and matched uvicorn's `ws_max_size` default by coincidence, so a move in that default would have parted the two wires silently. It lives in `wire` now and both the gRPC options and `uvicorn.Config` are passed it. A session the server refuses in a protocol frame — an unknown model, a session param it rejects — raises past every transport handler, and a gRPC connection holds a reader thread until it is closed. `_open_session` owns one attempt and closes the connection whenever the handshake does not finish. The query test reads `LOCAL_STACK` and `SEQ` rather than spelling them. Ticket: none - review findings on an unmerged branch --- positronic/offboard/client.py | 20 ++++++++++++----- positronic/offboard/grpc_wire.py | 8 ++----- positronic/offboard/server.py | 4 +++- positronic/offboard/tests/conftest.py | 7 +++++- positronic/offboard/tests/test_grpc_wire.py | 25 ++++++++++++++++++++- positronic/offboard/wire.py | 5 +++++ 6 files changed, 55 insertions(+), 14 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 2f479b7c3..133150de2 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -245,16 +245,28 @@ def _connect(self) -> wire.ClientConnection: ) return wire.WebsocketClientConnection(websocket) + def _open_session(self) -> InferenceSession: + """One attempt at a session, closing the connection whenever the handshake does not finish. + + A refusal the server sends as a protocol frame — an unknown model, a session param it rejects + — raises past every transport handler, and a gRPC connection holds a reader thread until it + is closed. + """ + conn = self._connect() + try: + return InferenceSession(conn, infer_timeout=self.infer_timeout) + except BaseException: + conn.close() + raise + def new_session(self) -> InferenceSession: """Creates a new inference session on the model the URL names.""" deadline = time.monotonic() + self.connect_deadline backoff = 1.0 retries = _ConnectRetries() while True: - conn = None try: - conn = self._connect() - return InferenceSession(conn, infer_timeout=self.infer_timeout) + return self._open_session() # ``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: @@ -270,8 +282,6 @@ def new_session(self) -> InferenceSession: grpc.RpcError, wire.PeerDisconnected, ) as e: - if conn is not None: - conn.close() if retries.take(e) is _ConnectOutcome.SURFACE: raise if time.monotonic() >= deadline: diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index ad4671dbd..617c9aedb 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -27,13 +27,9 @@ SESSION_PATH_HEADER = 'positronic-session-path' SESSION_QUERY_HEADER = 'positronic-session-query' -# An observation is a stack of camera frames, and the gRPC default of 4 MiB refuses one. 16 MiB is the -# ceiling uvicorn already gives the websocket wire (``ws_max_size``), so both wires carry the same frame. -_MAX_MESSAGE_BYTES = 16 * 1024 * 1024 - _MESSAGE_SIZE_OPTIONS = [ - ('grpc.max_receive_message_length', _MAX_MESSAGE_BYTES), - ('grpc.max_send_message_length', _MAX_MESSAGE_BYTES), + ('grpc.max_receive_message_length', wire.MAX_MESSAGE_BYTES), + ('grpc.max_send_message_length', wire.MAX_MESSAGE_BYTES), ] # How long ``close`` waits for the server to end the stream, so its own session cleanup runs. diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index c312afd68..d3f9cd472 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -424,7 +424,9 @@ def authorized(headers: Mapping[str, str]) -> bool: def serve(self): async def _run(): await self._startup() - config = uvicorn.Config(self.app, host=self.host, port=self.port, log_level='info') + config = uvicorn.Config( + self.app, host=self.host, port=self.port, log_level='info', ws_max_size=wire.MAX_MESSAGE_BYTES + ) server = uvicorn.Server(config) self._last_activity = time.monotonic() watchdog = None diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 1b9b71f6e..0a7b5d480 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -8,6 +8,7 @@ import pytest import uvicorn +from positronic.offboard import wire from positronic.offboard.server import PolicyServer from positronic.policy import Policy, Session from positronic.policy.executor import Executor @@ -36,7 +37,11 @@ def start(pipeline, *, grpc: bool = False, **server_kwargs) -> tuple[str, int, P grpc_port = _find_free_port() if grpc else None host = server_kwargs.pop('host', 'localhost') server = PolicyServer(pipeline, host=host, port=_find_free_port(), grpc_port=grpc_port, **server_kwargs) - uv_server = uvicorn.Server(uvicorn.Config(server.app, host=server.host, port=server.port, log_level='warning')) + uv_server = uvicorn.Server( + uvicorn.Config( + server.app, host=server.host, port=server.port, log_level='warning', ws_max_size=wire.MAX_MESSAGE_BYTES + ) + ) async def _run(): await server._startup() diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index ea65ac1ee..3414d15cd 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -7,9 +7,11 @@ import pytest from positronic.offboard import grpc_wire, wire +from positronic.offboard import keys as offboard_keys from positronic.offboard.client import InferenceClient, _ConnectRetries from positronic.offboard.server import AUTH_HEADER, PolicyServer, bearer from positronic.offboard.tests.conftest import DictSource, StartServer +from positronic.policy.base import SEQ from positronic.policy.layers import ChunkedSchedule, TemporalStack from positronic.policy.spec import ModelSource, PolicySource, remote @@ -105,7 +107,10 @@ def test_the_query_carries_the_session_params(start_server, make_mock_policy): _host, _port, server = start_server(pipe, grpc=True) session = InferenceClient(grpc_url(server, f'{wire.SESSION_PATH}?offsets=[-0.5, 0.0]')).new_session() try: - assert session.metadata['local_stack']['seq'][0]['args']['offsets_sec'] == [-0.5, 0.0] + stack = session.metadata[offboard_keys.LOCAL_STACK][SEQ] + # `args` and the layer's own constructor keyword are the spec grammar's, written wherever a + # layer renders itself; this reader spells them as the wire carries them. + assert stack[0]['args']['offsets_sec'] == [-0.5, 0.0] finally: session.close() @@ -186,3 +191,21 @@ def test_an_ipv6_host_binds_in_brackets(start_server: StartServer, make_mock_pol assert session.infer({'image': 'test'}) == [{'action': [4]}] finally: session.close() + + +def test_a_refused_handshake_closes_the_connection(both_wires): + """A model the source does not know is refused in a protocol frame, past the transport handlers, + and the gRPC connection behind it holds a reader thread until something closes it.""" + client = InferenceClient(grpc_url(both_wires[0], f'{wire.SESSION_PATH}/unknown-model')) + opened = [] + connect = client._connect + + def record(): + opened.append(connect()) + return opened[-1] + + client._connect = record + with pytest.raises(RuntimeError): + client.new_session() + assert opened, 'the session never opened a connection' + assert opened[0]._closed, 'the refused session left its connection open' diff --git a/positronic/offboard/wire.py b/positronic/offboard/wire.py index cdf770cef..d04a2087c 100644 --- a/positronic/offboard/wire.py +++ b/positronic/offboard/wire.py @@ -15,6 +15,11 @@ # session metadata, so both wires address a model the same way. SESSION_PATH = '/api/v1/session' +# The largest frame a session may carry, on either wire. An observation is a stack of camera frames, +# so the gRPC default of 4 MiB refuses one; uvicorn's own default happens to be this, and passing it +# explicitly is what keeps the two wires equal when that default moves. +MAX_MESSAGE_BYTES = 16 * 1024 * 1024 + class PeerDisconnected(Exception): """The peer ended the session.""" From be5921e411e6857193bd2d2baebeb235a5f361b4 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 07:36:56 +0000 Subject: [PATCH 06/46] Carry a gRPC session through a TLS edge and a front's idle close An authenticated endpoint is served by a TLS-terminating front, so `grpcs://` dials the gRPC port through one: the client verifies the edge against its own roots and the server keeps its plaintext port, which is what every front hands an HTTP/2 stream to. A `grpcs://` URL defaults to port 443 like the other TLS schemes. A front also drops a connection it has read nothing from, and one inference sends nothing until it answers, so the client pings through the silence. Three of gRPC's own defaults stop that working and each is set here: the client sends two pings and stops, it spaces them five minutes apart, and the server answers a 20s ping with `GOAWAY too_many_pings`. A certificate the client's roots do not cover now raises the exception the websocket wire raises, so the connect loop reads it as permanent instead of retrying its whole deadline out. gRPC's readiness future says only that the channel is down, so one call on it fetches the reason. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/client.py | 34 ++-- positronic/offboard/grpc_wire.py | 76 +++++++- positronic/offboard/tests/test_grpc_wire.py | 201 +++++++++++++++++++- 3 files changed, 290 insertions(+), 21 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 133150de2..7447d1aea 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -172,21 +172,22 @@ def take(self, e: Exception) -> _ConnectOutcome: return _ConnectOutcome.RETRY if again else _ConnectOutcome.SURFACE -# The URL scheme that puts a session on the gRPC wire. -_GRPC_SCHEME = 'grpc' +# The URL schemes that put a session on the gRPC wire: plaintext, and behind a TLS edge. +_GRPC_SCHEMES = ('grpc', 'grpcs') class InferenceClient: """The wire connection to one inference server, addressed by one URL. Accepted URL forms: ``host``, ``host:port``, and ``scheme://host[:port][/api/v1/session[/]]``, - each with an optional ``?query``. ``https``/``wss`` enable TLS (bare or ``http``/``ws`` forms don't); the - port defaults to the scheme's own, 443 for TLS and 80 otherwise. Everything the URL says about the - session — the model id it names and the query it carries as session params — reaches the server exactly - as written, so every session opened here serves that model with those params. + each with an optional ``?query``. ``https``/``wss``/``grpcs`` enable TLS (bare or ``http``/``ws``/``grpc`` + forms don't); the port defaults to the scheme's own, 443 for TLS and 80 otherwise. Everything the URL + says about the session — the model id it names and the query it carries as session params — reaches + the server exactly as written, so every session opened here serves that model with those params. - ``grpc://`` names the same session on the gRPC wire, which the server offers on a port of its own. That - port carries sessions alone, so ``list_models`` needs the HTTP URL. + ``grpc://`` names the same session on the gRPC wire, which the server offers on a port of its own, and + ``grpcs://`` names that port behind a TLS edge. Either port carries sessions alone, so ``list_models`` + needs the HTTP URL. ``headers`` carry auth, whether the server checks it or a proxy in front of it does — credentials stay out of the URL, which is meant to be safe to hand around. @@ -206,12 +207,13 @@ def __init__( infer_timeout: float = DEFAULT_INFER_TIMEOUT, ): split = urllib.parse.urlsplit(url if '://' in url else f'//{url}') - if split.scheme not in ('', 'http', 'ws', 'https', 'wss', _GRPC_SCHEME): + if split.scheme not in ('', 'http', 'ws', 'https', 'wss', *_GRPC_SCHEMES): raise ValueError(f'Unsupported scheme {split.scheme!r} in {url!r}') if not split.hostname: raise ValueError(f'No host in {url!r}') - secure = split.scheme in ('https', 'wss') - session_scheme = _GRPC_SCHEME if split.scheme == _GRPC_SCHEME else ('wss' if secure else 'ws') + grpc_wired = split.scheme in _GRPC_SCHEMES + secure = split.scheme in ('https', 'wss', 'grpcs') + session_scheme = split.scheme if grpc_wired else ('wss' if secure else 'ws') http_scheme = 'https' if secure else 'http' default_port = 443 if secure else 80 # urlsplit strips the brackets an IPv6 host needs back in a netloc. @@ -223,7 +225,8 @@ def __init__( query = f'?{split.query}' if split.query else '' self._session_path = _session_path(split.path, url) self._query = split.query - self._grpc_target = f'{host}:{port}' if split.scheme == _GRPC_SCHEME else None + self._grpc_target = f'{host}:{port}' if grpc_wired else None + self._grpc_secure = secure self.session_url = f'{session_scheme}://{netloc}{self._session_path}{query}' self.api_url = None if self._grpc_target else f'{http_scheme}://{netloc}/api/v1' self.headers = dict(headers) if headers else None @@ -235,7 +238,12 @@ def _connect(self) -> wire.ClientConnection: """One session's connection, over the wire the URL names.""" if self._grpc_target is not None: return grpc_wire.GrpcClientConnection( - self._grpc_target, self._session_path, self._query, self.headers, self.open_timeout + self._grpc_target, + self._session_path, + self._query, + self.headers, + self.open_timeout, + secure=self._grpc_secure, ) # A proxy between here and the server closes a connection it has read nothing from, often # after 60s — well inside one ``infer_timeout`` inference, which sends nothing until it diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 617c9aedb..198a69dec 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -6,6 +6,7 @@ import logging import queue +import ssl import threading import urllib.parse from collections.abc import AsyncIterator, Awaitable, Callable, Mapping @@ -32,15 +33,77 @@ ('grpc.max_send_message_length', wire.MAX_MESSAGE_BYTES), ] +# A front between the two ends closes a connection it has read nothing from — the Nebius managed +# ingress after ~90s — and one inference sends nothing until it answers. The client pings through +# that silence, and the server must tolerate the pings: gRPC's own server defaults are a five-minute +# floor and two strikes, which answer a 20s ping with ``GOAWAY too_many_pings``. +_PING_EVERY_MS = 20_000 +_PING_ANSWER_TIMEOUT_MS = 10_000 +_PING_TOLERATED_EVERY_MS = 10_000 + # How long ``close`` waits for the server to end the stream, so its own session cleanup runs. _CLOSE_TIMEOUT_SEC = 5.0 +# A path no handler serves, so asking why a channel is down never opens a session on a server that +# turns out to be up after all. +_PROBE_PATH = f'/{SERVICE}/ChannelProbe' + +# What gRPC's own status details call a certificate the client's roots do not cover. +_UNVERIFIABLE_CERTIFICATE = 'CERTIFICATE_VERIFY_FAILED' + + +def _why_not_ready(channel: grpc.Channel, timeout: float) -> str: + """What gRPC says stopped the channel coming up. Its readiness future carries only that it did not.""" + probe = channel.stream_stream(_PROBE_PATH, request_serializer=None, response_deserializer=None) + try: + next(probe(iter(()), timeout=timeout)) + except grpc.RpcError as e: + return e.details() or '' + except StopIteration: + return '' + return '' + + +def _client_options() -> list[tuple[str, int]]: + return [ + *_MESSAGE_SIZE_OPTIONS, + ('grpc.keepalive_time_ms', _PING_EVERY_MS), + ('grpc.keepalive_timeout_ms', _PING_ANSWER_TIMEOUT_MS), + # Both of gRPC's own client throttles stop the pings during exactly the silent wait they + # exist for: it sends two and stops, and it spaces them five minutes apart. + ('grpc.http2.max_pings_without_data', 0), + ('grpc.http2.min_time_between_pings_ms', _PING_EVERY_MS), + ] + + +def _server_options() -> list[tuple[str, int]]: + return [ + *_MESSAGE_SIZE_OPTIONS, + ('grpc.http2.min_ping_interval_without_data_ms', _PING_TOLERATED_EVERY_MS), + ('grpc.http2.max_ping_strikes', 0), + ] + + +def channel_credentials() -> grpc.ChannelCredentials: + """The roots a ``grpcs://`` channel verifies the edge against: the system's own.""" + return grpc.ssl_channel_credentials() + + +def _channel(target: str, secure: bool) -> grpc.Channel: + options = _client_options() + if secure: + return grpc.secure_channel(target, channel_credentials(), options=options) + return grpc.insecure_channel(target, options=options) + class GrpcClientConnection: """A client's end of one gRPC session. A reader thread drains the response stream into a queue, because the stream itself has no per-message timeout and ``recv`` needs one. + + ``secure`` dials over TLS, which is the shape an authenticated endpoint takes: a TLS edge in front + of the server's plaintext gRPC port. """ def __init__( @@ -50,13 +113,19 @@ def __init__( query: str, headers: Mapping[str, str] | None = None, open_timeout: float = 10.0, + secure: bool = False, ): self._target = target - self._channel = grpc.insecure_channel(target, options=_MESSAGE_SIZE_OPTIONS) + self._channel = _channel(target, secure) try: grpc.channel_ready_future(self._channel).result(timeout=open_timeout) except grpc.FutureTimeoutError: + details = _why_not_ready(self._channel, timeout=open_timeout) self._channel.close() + if _UNVERIFIABLE_CERTIFICATE in details: + # The same exception the websocket wire raises here, so one connect loop reads a + # misconfigured edge as permanent over either wire rather than retrying its deadline out. + raise ssl.SSLCertVerificationError(f'gRPC channel to {target}: {details}') from None raise TimeoutError(f'gRPC channel to {target} is not ready within {open_timeout}s') from None # gRPC metadata keys are lower case, and they are the same header names the websocket wire sends. metadata = tuple((key.lower(), value) for key, value in (headers or {}).items()) + ( @@ -178,6 +247,9 @@ async def serve( ``authorized`` reads the session headers and refuses before the session opens, as the websocket wire refuses the upgrade. + + The port is plaintext. An authenticated deployment puts a TLS edge in front of it, which terminates + TLS and hands this server the HTTP/2 stream, so no shape needs a certificate here. """ async def _serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerContext) -> None: @@ -193,7 +265,7 @@ async def _serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerC await context.abort(grpc.StatusCode.INTERNAL, str(e)) handler = grpc.stream_stream_rpc_method_handler(_serve_one, request_deserializer=None, response_serializer=None) - server = grpc.aio.server(options=_MESSAGE_SIZE_OPTIONS) + server = grpc.aio.server(options=_server_options()) server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler(SERVICE, {METHOD: handler}),)) bound = server.add_insecure_port(_bind_target(host, port)) if bound == 0: diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index 3414d15cd..56a72543a 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -1,10 +1,24 @@ """The gRPC wire: one session runs over it exactly as it runs over the websocket.""" +import asyncio +import datetime +import pathlib +import queue +import ssl +import tempfile +import threading +import time +from collections.abc import Callable, Generator from unittest.mock import ANY, MagicMock import configuronic as cfn import grpc import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.hazmat.primitives.serialization import Encoding, NoEncryption, PrivateFormat +from cryptography.x509.oid import NameOID from positronic.offboard import grpc_wire, wire from positronic.offboard import keys as offboard_keys @@ -141,19 +155,147 @@ def test_the_grpc_wire_refuses_a_session_without_the_token(authed_server, header assert refused.value.code() is grpc.StatusCode.PERMISSION_DENIED -@pytest.mark.parametrize('url', ['grpcs://gpu-host:9000', 'tcp://gpu-host:9000']) -def test_an_unknown_scheme_is_refused(url): +def _self_signed(host: str) -> tuple[bytes, bytes]: + """A certificate and key for ``host``, PEM encoded, valid from yesterday.""" + key = ec.generate_private_key(ec.SECP256R1()) + name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, host)]) + day = datetime.timedelta(days=1) + now = datetime.datetime.now(datetime.UTC) + certificate = ( + x509 + .CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - day) + .not_valid_after(now + day) + .add_extension(x509.SubjectAlternativeName([x509.DNSName(host)]), critical=False) + .sign(key, hashes.SHA256()) + ) + private = key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption()) + return certificate.public_bytes(Encoding.PEM), private + + +async def _copy(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + while chunk := await reader.read(65536): + writer.write(chunk) + await writer.drain() + except OSError: + pass + finally: + writer.close() + + +@pytest.fixture +def tls_edge() -> Generator[Callable[[str, int], tuple[int, bytes]], None, None]: + """Starts a TLS front over a plaintext gRPC port, the shape an authenticated endpoint takes. + + It terminates TLS, selects HTTP/2 over ALPN and copies the bytes on, so the client and the server + speak one h2 connection end to end and the server holds no certificate. Answers the front's own + port and the root to verify it against. + """ + stops: list[tuple[asyncio.AbstractEventLoop, asyncio.Event]] = [] + + def start(backend_host: str, backend_port: int) -> tuple[int, bytes]: + certificate, private = _self_signed('localhost') + started: queue.SimpleQueue = queue.SimpleQueue() + + async def _serve_edge() -> None: + with tempfile.TemporaryDirectory() as keys: + chain, key_file = pathlib.Path(keys, 'chain.pem'), pathlib.Path(keys, 'key.pem') + chain.write_bytes(certificate) + key_file.write_bytes(private) + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(chain, key_file) + context.set_alpn_protocols(['h2']) + + async def _handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + backend_r, backend_w = await asyncio.open_connection(backend_host, backend_port) + await asyncio.gather(_copy(reader, backend_w), _copy(backend_r, writer)) + + edge = await asyncio.start_server(_handle, 'localhost', 0, ssl=context) + stop = asyncio.Event() + started.put((edge.sockets[0].getsockname()[1], asyncio.get_running_loop(), stop)) + async with edge: + await stop.wait() + + threading.Thread(target=asyncio.run, args=(_serve_edge(),), daemon=True).start() + port, loop, stop = started.get(timeout=5.0) + stops.append((loop, stop)) + return port, certificate + + yield start + for loop, stop in stops: + loop.call_soon_threadsafe(stop.set) + + +@pytest.fixture +def edged(tls_edge, monkeypatch) -> Callable[[PolicyServer], str]: + """The ``grpcs://`` URL of a server reached through a TLS edge, with the client trusting its root.""" + + def url(server: PolicyServer) -> str: + port, root = tls_edge(server.host, server.grpc_port) + monkeypatch.setattr(grpc_wire, 'channel_credentials', lambda: grpc.ssl_channel_credentials(root)) + return f'grpcs://localhost:{port}' + + return url + + +def test_a_session_through_a_tls_edge_handshakes_and_infers(both_wires, edged): + server, policy = both_wires + session = InferenceClient(edged(server)).new_session() + try: + assert session.metadata['model_name'] == 'stub' + obs = {'image': 'test'} + assert session.infer(obs) == [{'action': [1, 2, 3]}] + policy._mock_session.assert_called_with(obs, ANY) + finally: + session.close() + + +def test_a_tls_edge_carries_the_bearer_token(authed_server, edged): + session = InferenceClient(edged(authed_server), headers={AUTH_HEADER: bearer(_TOKEN)}).new_session() + try: + assert session.metadata['model_name'] == 'stub' + finally: + session.close() + + +def test_a_tls_edge_session_without_the_token_is_refused(authed_server, edged, monkeypatch): + monkeypatch.setattr(_ConnectRetries, 'MAX_FORBIDDEN_ATTEMPTS', 1) + with pytest.raises(grpc.RpcError) as refused: + InferenceClient(edged(authed_server)).new_session() + assert refused.value.code() is grpc.StatusCode.PERMISSION_DENIED + + +def test_an_unknown_scheme_is_refused(): with pytest.raises(ValueError, match='Unsupported scheme'): - InferenceClient(url) + InferenceClient('tcp://gpu-host:9000') -def test_a_grpc_url_names_the_session_port_alone(): - client = InferenceClient('grpc://gpu-host:9000') - assert client.session_url == 'grpc://gpu-host:9000/api/v1/session' +@pytest.mark.parametrize('url', ['grpc://gpu-host:9000', 'grpcs://gpu-host:9000']) +def test_a_grpc_url_names_the_session_port_alone(url): + client = InferenceClient(url) + assert client.session_url == f'{url}/api/v1/session' with pytest.raises(ValueError, match='gRPC session port'): client.list_models() +@pytest.mark.parametrize( + ('url', 'target', 'secure'), + [ + ('grpc://gpu-host', 'gpu-host:80', False), + ('grpcs://gpu-host', 'gpu-host:443', True), + ('grpcs://gpu-host:9000', 'gpu-host:9000', True), + ], +) +def test_the_scheme_fixes_the_port_and_the_tls(url, target, secure): + client = InferenceClient(url) + assert (client._grpc_target, client._grpc_secure) == (target, secure) + + @pytest.mark.parametrize( ('session_path', 'model_id'), [ @@ -209,3 +351,50 @@ def record(): client.new_session() assert opened, 'the session never opened a connection' assert opened[0]._closed, 'the refused session left its connection open' + + +# Long enough for the client to send more pings than gRPC's own server default tolerates. +_SILENCE_SEC = 8.0 + + +@pytest.fixture +def chatty_client(monkeypatch) -> None: + """Pings often enough that a silence measured in seconds stands in for one measured in minutes.""" + monkeypatch.setattr(grpc_wire, '_PING_EVERY_MS', 500) + + +def _silent_then_infer(server: PolicyServer) -> list[dict]: + session = InferenceClient(grpc_url(server)).new_session() + try: + time.sleep(_SILENCE_SEC) + return session.infer({'image': 'test'}) + finally: + session.close() + + +def test_a_session_answers_after_a_silence_no_frame_crossed(both_wires, chatty_client): + """One inference can outlast a front's idle close, so the wire's own pings hold the stream open.""" + assert _silent_then_infer(both_wires[0]) == [{'action': [1, 2, 3]}] + + +def test_a_server_on_the_grpc_ping_defaults_kills_the_silent_session( + start_server, make_mock_policy, chatty_client, monkeypatch +): + """gRPC's own server defaults answer those pings with ``GOAWAY too_many_pings``.""" + monkeypatch.setattr(grpc_wire, '_server_options', lambda: list(grpc_wire._MESSAGE_SIZE_OPTIONS)) + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + _host, _port, server = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True) + with pytest.raises(grpc.RpcError, match='Too many pings'): + _silent_then_infer(server) + + +def test_a_certificate_the_client_cannot_verify_is_not_retried(both_wires, tls_edge, monkeypatch): + """A root that does not cover the edge is permanent, so it surfaces on the first attempt.""" + port, _root = tls_edge(both_wires[0].host, both_wires[0].grpc_port) + unrelated, _key = _self_signed('localhost') + monkeypatch.setattr(grpc_wire, 'channel_credentials', lambda: grpc.ssl_channel_credentials(unrelated)) + client = InferenceClient(f'grpcs://localhost:{port}', open_timeout=2.0, connect_deadline=20.0) + started = time.monotonic() + with pytest.raises(ssl.SSLCertVerificationError, match=grpc_wire._UNVERIFIABLE_CERTIFICATE): + client.new_session() + assert time.monotonic() - started < 8.0, 'the connect retried a permanent failure' From 6155128fe58a7547b52ebedd6226067cf5348f30 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 07:40:24 +0000 Subject: [PATCH 07/46] Add a wire-only stub server for measuring a session's transport It answers one constant chunk, so a round trip measures the wire and nothing else. `?delay_sec=120` holds an inference open for two minutes, which is what a cold model does to a connection with nothing crossing it meanwhile. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/stub.py | 77 +++++++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 positronic/offboard/stub.py diff --git a/positronic/offboard/stub.py b/positronic/offboard/stub.py new file mode 100644 index 000000000..0e7b6e0a5 --- /dev/null +++ b/positronic/offboard/stub.py @@ -0,0 +1,77 @@ +"""A server with no model: every inference answers the same chunk, so a session measures the wire alone. + +``delay_sec`` is a session param, so ``?delay_sec=120`` holds one inference open for two minutes — +what a cold model does to a connection, with nothing on the wire meanwhile. +""" + +import time +from collections.abc import Mapping +from typing import Any + +import configuronic as cfn + +from pimm.logging import init_logging +from positronic import keys +from positronic.offboard.server import serve +from positronic.policy import Policy, Session +from positronic.policy.base import DelegatingSession, Layer, Runtime +from positronic.policy.layers import ChunkedSchedule +from positronic.policy.spec import Pipeline, PolicySource, remote + +# One action, at the start of the chunk. A served session must answer a trajectory, and this is the +# smallest one that is. +CHUNK = [{keys.ACTION_TIMESTAMP: 0.0}] + + +class StubSession(Session): + def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]]: + return CHUNK + + @property + def meta(self) -> dict[str, Any]: + return {'model_name': 'stub'} + + +class StubPolicy(Policy): + """Answers ``CHUNK``, whatever it is asked.""" + + def new_session(self, context: dict[str, Any] | None = None, rt: Runtime | None = None) -> Session: + return StubSession() + + +class DelayedSession(DelegatingSession): + def __init__(self, inner: Session, delay_sec: float): + super().__init__(inner) + self._delay_sec = delay_sec + + def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] | None: + time.sleep(self._delay_sec) + return super().__call__(obs, time_ns) + + +class Delay(Layer): + """Holds every answer for ``delay_sec``, standing in for a model slow enough to outlast a front's + idle close. A layer rather than an argument of the policy: a session param may tune the pipeline + around the model source, never the source itself. + """ + + def __init__(self, delay_sec: float = 0.0): + self._delay_sec = delay_sec + + def make_session(self, inner: Session) -> Session: + return DelayedSession(inner, self._delay_sec) + + +# One instance for the process: a server compares the source a session param rebuilds against the one +# it launched with, and two sources are equal only when they hold the same policy. +POLICY = StubPolicy() + + +@cfn.config(delay_sec=0.0) +def pipeline(delay_sec: float) -> Pipeline: + return ChunkedSchedule() | remote | Delay(delay_sec) | PolicySource(POLICY) + + +if __name__ == '__main__': + init_logging() + cfn.cli(serve.override(pipeline=pipeline)) From f6028f2b3651f9eae5bd248492828f730650bc29 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 07:40:40 +0000 Subject: [PATCH 08/46] Serve both wires from a Nebius endpoint, and say what its front does The endpoint now exposes port 9000 beside 8000 and prints the `grpcs://` URL for it. The gRPC port is declared as an ordinary HTTP port: that front negotiates HTTP/2 over ALPN and carries a gRPC session end to end, where a port declared `/tcp` gets a `tls://` URL whose front negotiates no ALPN at all and gRPC refuses it. The README records what a round trip costs through such a front, and where the front stops being the faster wire's problem. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 27 ++++++++++++++++++++++----- workflows/nebius/serve.sh | 31 ++++++++++++++++++++++++++++--- 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index 19443bcc5..a80b2d452 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -15,6 +15,7 @@ the same order, so everything below holds on each. |---|---|---| | WebSocket | `ws://host:8000/api/v1/session[/]` | the server's `port`, beside the HTTP routes | | gRPC | `grpc://host:9000/api/v1/session[/]` | the server's `grpc_port`, sessions alone | +| gRPC over TLS | `grpcs://host:443/api/v1/session[/]` | a TLS edge in front of that same `grpc_port` | The WebSocket wire is the default, and a server serves gRPC only when `grpc_port` names a port. A gRPC session is one bidirectional stream of the same frames, so no `.proto` file describes them. @@ -22,9 +23,24 @@ The session path and the query cross as the `positronic-session-path` and `posit metadata, and `Authorization` crosses as the `authorization` metadata. Python's WebSocket stack costs about 30 ms per 846 KiB observation in framing and reassembly, which -gRPC does in about 1 ms. Take the gRPC wire on an endpoint a client reaches directly. A managed HTTPS -front usually translates HTTP into its own protocol and drops the HTTP/2 frame detail gRPC needs, so -an endpoint behind one keeps the WebSocket wire. +gRPC does in about 1 ms, so take the gRPC wire wherever it reaches. + +It reaches through a managed HTTPS front, which is what serves an authenticated endpoint: the front +terminates TLS and the HTTP/2 connection runs end to end, so the server binds a plaintext port and +holds no certificate of its own. The front has to negotiate HTTP/2 over ALPN — check a new one with +`openssl s_client -alpn h2 -connect :443`. On a Nebius Serverless Endpoint that means declaring +the gRPC port as an ordinary HTTP port and dialling its `https://` host as `grpcs://:443`; a +port declared `/tcp` is fronted by a `tls://` URL that negotiates no ALPN, which gRPC refuses with +`Cannot check peer: missing selected ALPN property`. + +Through such an endpoint an 846 KiB observation round-trips in about 6 ms over gRPC against about +60 ms over the WebSocket, and gRPC holds that at 10 Hz, which is 8 MB/s of observation. The front +shapes a session that outruns it: a back-to-back loop settles at about 83 ms a round trip after some +11 MB, and gets its speed back after a minute of quiet. The WebSocket holds its 60 ms throughout, +never being fast enough to be shaped. + +Both wires ping through a silent wait, so a front that drops a connection it has read nothing from — +the managed one after about 90 s — does not cut an inference the model is still working on. `/api/v1/models` is an HTTP route, so it stays on the server's `port`. `InferenceClient.list_models` over a `grpc://` URL says so. @@ -95,7 +111,7 @@ Any violation — including an unknown key — fails at connect: the server send One string is a complete endpoint description, because the whole session configuration fits in the URL: `--policy=.remote --policy.url='gpu-host:8000?codec.fps=10'` accepts `host`, `host:port`, and full -`http(s)`/`ws(s)`/`grpc` URLs — optionally with `/api/v1/session/` — and forwards the query string verbatim. +`http(s)`/`ws(s)`/`grpc(s)` URLs — optionally with `/api/v1/session/` — and forwards the query string verbatim. Credentials are the exception and stay a separate `headers` argument, so the URL itself is safe to hand around. ### Session Flow @@ -261,8 +277,9 @@ from positronic.offboard.client import InferenceClient client = InferenceClient('localhost:8000') # A named model, tuned for every session this client opens # client = InferenceClient('localhost:8000/api/v1/session/model_a?codec.fps=10') -# The same session on the gRPC wire +# The same session on the gRPC wire, on a LAN and behind a TLS edge # client = InferenceClient('grpc://localhost:9000/api/v1/session/model_a') +# client = InferenceClient('grpcs://gpu-host:443/api/v1/session/model_a') session = client.new_session() meta = session.metadata diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 95abdcab8..31d758df5 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -6,6 +6,12 @@ # itself takes ~10-15 min more to finish uv sync and load the model into GPU # memory after the URL appears. # +# Both wires are served: the websocket on port 8000 and gRPC on port 9000. The +# gRPC port is declared as an ordinary HTTP port, because that front negotiates +# HTTP/2 over ALPN and so carries gRPC end to end; a port declared `/tcp` gets a +# tls:// URL whose front negotiates no ALPN at all, which gRPC refuses with +# "Cannot check peer: missing selected ALPN property". +# # That URL carries the id of a tunnel created with the endpoint, so it cannot be # chosen or known in advance, and a delete plus re-create earns a new one even # under the same name. Nothing may hold it across a redeploy. `nebius ai endpoint @@ -90,6 +96,12 @@ case " $* " in *) set -- "$@" "--idle_timeout_min=${NEBIUS_IDLE_TIMEOUT_MIN:-20}" ;; esac +GRPC_PORT=9000 +case " $* " in + *" --grpc_port="*|*" --grpc_port "*) ;; + *) set -- "$@" "--grpc_port=${GRPC_PORT}" ;; +esac + SERVER_ARGS="run --python 3.13 ${EXTRA}python -m positronic.vendors.${VENDOR}.server $*" echo "Creating $VENDOR endpoint '$NAME'..." @@ -101,6 +113,7 @@ nebius ai endpoint create \ --container-command uv \ --args "$SERVER_ARGS" \ --container-port 8000 \ + --container-port "${GRPC_PORT}" \ --platform gpu-h100-sxm \ --preset "$PRESET" \ --working-dir /positronic \ @@ -127,10 +140,11 @@ echo "Waiting for the managed HTTPS URL (typically <1 min)..." URL="" for i in $(seq 1 30); do - # This field also carries bare `IP:port` entries, which serve no TLS and would put the bearer token - # on the wire in cleartext — take the https:// one, and fail rather than fall back. + # Each managed URL names the container port it fronts, so the two wires are told apart by that + # prefix. This field also carries bare `IP:port` entries, which serve no TLS and would put the + # bearer token on the wire in cleartext — take the https:// ones, and fail rather than fall back. URL=$(nebius ai endpoint get "$ID" --format json 2>/dev/null \ - | jq -r '[.status.public_endpoints[]? | select(startswith("https://"))] | first // empty') + | jq -r '[.status.public_endpoints[]? | select(startswith("https://port8000-"))] | first // empty') if [ -n "$URL" ]; then break; fi sleep 10 done @@ -140,10 +154,16 @@ if [ -z "$URL" ]; then exit 1 fi +GRPC_HOST=$(nebius ai endpoint get "$ID" --format json 2>/dev/null \ + | jq -r "[.status.public_endpoints[]? | select(startswith(\"https://port${GRPC_PORT}-\"))] | first // empty" \ + | sed 's|^https://||') +GRPC_URL="grpcs://${GRPC_HOST}:443" + cat < Date: Wed, 9 Sep 2026 07:45:05 +0000 Subject: [PATCH 09/46] Say the front's contract once, where the README already holds it Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/grpc_wire.py | 10 ++++------ workflows/nebius/serve.sh | 6 ++---- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 198a69dec..20219d8b2 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -33,10 +33,8 @@ ('grpc.max_send_message_length', wire.MAX_MESSAGE_BYTES), ] -# A front between the two ends closes a connection it has read nothing from — the Nebius managed -# ingress after ~90s — and one inference sends nothing until it answers. The client pings through -# that silence, and the server must tolerate the pings: gRPC's own server defaults are a five-minute -# floor and two strikes, which answer a 20s ping with ``GOAWAY too_many_pings``. +# A front closes a connection it has read nothing from — ~90s on the Nebius managed ingress — and one +# inference sends nothing until it answers, so the client pings through that silence. _PING_EVERY_MS = 20_000 _PING_ANSWER_TIMEOUT_MS = 10_000 _PING_TOLERATED_EVERY_MS = 10_000 @@ -79,6 +77,7 @@ def _client_options() -> list[tuple[str, int]]: def _server_options() -> list[tuple[str, int]]: return [ *_MESSAGE_SIZE_OPTIONS, + # gRPC's own defaults, a five-minute floor and two strikes, answer a 20s ping with GOAWAY. ('grpc.http2.min_ping_interval_without_data_ms', _PING_TOLERATED_EVERY_MS), ('grpc.http2.max_ping_strikes', 0), ] @@ -248,8 +247,7 @@ async def serve( ``authorized`` reads the session headers and refuses before the session opens, as the websocket wire refuses the upgrade. - The port is plaintext. An authenticated deployment puts a TLS edge in front of it, which terminates - TLS and hands this server the HTTP/2 stream, so no shape needs a certificate here. + The port is plaintext; a TLS edge in front of it is what serves an authenticated endpoint. """ async def _serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerContext) -> None: diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 31d758df5..0d89c7b23 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -7,10 +7,8 @@ # memory after the URL appears. # # Both wires are served: the websocket on port 8000 and gRPC on port 9000. The -# gRPC port is declared as an ordinary HTTP port, because that front negotiates -# HTTP/2 over ALPN and so carries gRPC end to end; a port declared `/tcp` gets a -# tls:// URL whose front negotiates no ALPN at all, which gRPC refuses with -# "Cannot check peer: missing selected ALPN property". +# gRPC port is declared as an ordinary HTTP port and never `/tcp` — the offboard +# README says what each front does to a gRPC session. # # That URL carries the id of a tunnel created with the endpoint, so it cannot be # chosen or known in advance, and a delete plus re-create earns a new one even From a5aa9aa2584a7e932c10c1b48c78f54bc9896c6e Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 07:47:45 +0000 Subject: [PATCH 10/46] Put the subject first in two sentences about the TLS edge Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 2 +- positronic/offboard/grpc_wire.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index a80b2d452..b02f8a3b3 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -25,7 +25,7 @@ metadata, and `Authorization` crosses as the `authorization` metadata. Python's WebSocket stack costs about 30 ms per 846 KiB observation in framing and reassembly, which gRPC does in about 1 ms, so take the gRPC wire wherever it reaches. -It reaches through a managed HTTPS front, which is what serves an authenticated endpoint: the front +It reaches through a managed HTTPS front, which is how an authenticated endpoint is served: the front terminates TLS and the HTTP/2 connection runs end to end, so the server binds a plaintext port and holds no certificate of its own. The front has to negotiate HTTP/2 over ALPN — check a new one with `openssl s_client -alpn h2 -connect :443`. On a Nebius Serverless Endpoint that means declaring diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 20219d8b2..1cf08d3ba 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -247,7 +247,7 @@ async def serve( ``authorized`` reads the session headers and refuses before the session opens, as the websocket wire refuses the upgrade. - The port is plaintext; a TLS edge in front of it is what serves an authenticated endpoint. + The port is plaintext; a TLS edge in front of it serves an authenticated endpoint. """ async def _serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerContext) -> None: From 859fecd097c4e4e0de6f073b296aea8707437ace Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 08:11:00 +0000 Subject: [PATCH 11/46] Read a TLS edge no client can use as permanent, and expose the port asked for MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three review findings. A front that selects no HTTP/2 over ALPN refuses every gRPC client, exactly as an unverifiable certificate does, and both answer `UNAVAILABLE` like a cold backend — so the connect loop reads the status details rather than the code alone, and stops. The connection raises what gRPC blamed instead of a bare timeout, which is what carries those details. A test edge that selects no protocol pins the new case; it binds one address, because gRPC reports the last address it failed on and a second family refusing the connection hides what the first blamed. `serve.sh` exposed port 9000 whatever the caller asked the server to listen on. It now reads the port back out of `--grpc_port` when one is given. The test proxy caught every `OSError`, which hid a broken edge behind the reset an ending session is expected to raise. It now catches the two that end a session. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/client.py | 5 ++- positronic/offboard/grpc_wire.py | 29 ++++++++------ positronic/offboard/tests/test_grpc_wire.py | 44 +++++++++++++++------ workflows/nebius/serve.sh | 10 +++-- 4 files changed, 60 insertions(+), 28 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 7447d1aea..803875806 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -131,7 +131,8 @@ def _refusal(e: Exception) -> _Refusal: """How to read a refused connect, over either wire. Each gRPC code stands for the HTTP status its wire twin answers: ``PERMISSION_DENIED`` for 403, - ``UNAVAILABLE`` for 503, ``RESOURCE_EXHAUSTED`` for 429. + ``UNAVAILABLE`` for 503, ``RESOURCE_EXHAUSTED`` for 429. A TLS edge no client can use answers + ``UNAVAILABLE`` too, exactly as a cold backend does, so its details are what tell them apart. """ if isinstance(e, InvalidStatus): status = e.response.status_code @@ -142,6 +143,8 @@ def _refusal(e: Exception) -> _Refusal: return _Refusal.FINAL # A gRPC error carries its code as a `Call`; anything else says nothing about the server. if isinstance(e, grpc.Call): + if grpc_wire.edge_is_unusable(e.details() or ''): + return _Refusal.FINAL code = e.code() if code is grpc.StatusCode.PERMISSION_DENIED: return _Refusal.FORBIDDEN diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 1cf08d3ba..737ff3701 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -6,7 +6,6 @@ import logging import queue -import ssl import threading import urllib.parse from collections.abc import AsyncIterator, Awaitable, Callable, Mapping @@ -46,20 +45,26 @@ # turns out to be up after all. _PROBE_PATH = f'/{SERVICE}/ChannelProbe' -# What gRPC's own status details call a certificate the client's roots do not cover. -_UNVERIFIABLE_CERTIFICATE = 'CERTIFICATE_VERIFY_FAILED' +# What gRPC's status details call an edge no client can use: a certificate its roots do not cover, +# and a front that selects no HTTP/2 over ALPN. +UNUSABLE_EDGE = ('CERTIFICATE_VERIFY_FAILED', 'missing selected ALPN property') -def _why_not_ready(channel: grpc.Channel, timeout: float) -> str: +def edge_is_unusable(details: str) -> bool: + """Whether a gRPC status blames the TLS edge's own configuration rather than a cold backend.""" + return any(marker in details for marker in UNUSABLE_EDGE) + + +def _connect_refusal(channel: grpc.Channel, timeout: float) -> grpc.RpcError | None: """What gRPC says stopped the channel coming up. Its readiness future carries only that it did not.""" probe = channel.stream_stream(_PROBE_PATH, request_serializer=None, response_deserializer=None) try: next(probe(iter(()), timeout=timeout)) except grpc.RpcError as e: - return e.details() or '' + return e except StopIteration: - return '' - return '' + return None + return None def _client_options() -> list[tuple[str, int]]: @@ -119,12 +124,12 @@ def __init__( try: grpc.channel_ready_future(self._channel).result(timeout=open_timeout) except grpc.FutureTimeoutError: - details = _why_not_ready(self._channel, timeout=open_timeout) + refusal = _connect_refusal(self._channel, timeout=open_timeout) self._channel.close() - if _UNVERIFIABLE_CERTIFICATE in details: - # The same exception the websocket wire raises here, so one connect loop reads a - # misconfigured edge as permanent over either wire rather than retrying its deadline out. - raise ssl.SSLCertVerificationError(f'gRPC channel to {target}: {details}') from None + # An edge that refuses every client is permanent, so raise what gRPC blamed rather than a + # timeout: the connect loop reads the status and stops instead of retrying its deadline out. + if refusal is not None and edge_is_unusable(refusal.details() or ''): + raise refusal from None raise TimeoutError(f'gRPC channel to {target} is not ready within {open_timeout}s') from None # gRPC metadata keys are lower case, and they are the same header names the websocket wire sends. metadata = tuple((key.lower(), value) for key, value in (headers or {}).items()) + ( diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index 56a72543a..da29cacac 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -2,6 +2,7 @@ import asyncio import datetime +import ipaddress import pathlib import queue import ssl @@ -155,6 +156,11 @@ def test_the_grpc_wire_refuses_a_session_without_the_token(authed_server, header assert refused.value.code() is grpc.StatusCode.PERMISSION_DENIED +# The edge answers on one address, not on both families a name resolves to: gRPC reports the last +# address it failed on, so a second leg refusing the connection would hide what the first blamed. +EDGE_HOST = '127.0.0.1' + + def _self_signed(host: str) -> tuple[bytes, bytes]: """A certificate and key for ``host``, PEM encoded, valid from yesterday.""" key = ec.generate_private_key(ec.SECP256R1()) @@ -170,7 +176,7 @@ def _self_signed(host: str) -> tuple[bytes, bytes]: .serial_number(x509.random_serial_number()) .not_valid_before(now - day) .not_valid_after(now + day) - .add_extension(x509.SubjectAlternativeName([x509.DNSName(host)]), critical=False) + .add_extension(x509.SubjectAlternativeName([x509.IPAddress(ipaddress.ip_address(host))]), critical=False) .sign(key, hashes.SHA256()) ) private = key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption()) @@ -182,7 +188,9 @@ async def _copy(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> N while chunk := await reader.read(65536): writer.write(chunk) await writer.drain() - except OSError: + # Whichever end closes first leaves the other half of the pair writing into a dead socket, which + # is how a session ends. Anything else is the edge itself failing and belongs in the test's face. + except (ConnectionResetError, BrokenPipeError): pass finally: writer.close() @@ -194,12 +202,13 @@ def tls_edge() -> Generator[Callable[[str, int], tuple[int, bytes]], None, None] It terminates TLS, selects HTTP/2 over ALPN and copies the bytes on, so the client and the server speak one h2 connection end to end and the server holds no certificate. Answers the front's own - port and the root to verify it against. + port and the root to verify it against. ``alpn=False`` selects no protocol at all, which is what + a front fronting a raw TCP port does. """ stops: list[tuple[asyncio.AbstractEventLoop, asyncio.Event]] = [] - def start(backend_host: str, backend_port: int) -> tuple[int, bytes]: - certificate, private = _self_signed('localhost') + def start(backend_host: str, backend_port: int, alpn: bool = True) -> tuple[int, bytes]: + certificate, private = _self_signed(EDGE_HOST) started: queue.SimpleQueue = queue.SimpleQueue() async def _serve_edge() -> None: @@ -209,13 +218,14 @@ async def _serve_edge() -> None: key_file.write_bytes(private) context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) context.load_cert_chain(chain, key_file) - context.set_alpn_protocols(['h2']) + if alpn: + context.set_alpn_protocols(['h2']) async def _handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: backend_r, backend_w = await asyncio.open_connection(backend_host, backend_port) await asyncio.gather(_copy(reader, backend_w), _copy(backend_r, writer)) - edge = await asyncio.start_server(_handle, 'localhost', 0, ssl=context) + edge = await asyncio.start_server(_handle, EDGE_HOST, 0, ssl=context) stop = asyncio.Event() started.put((edge.sockets[0].getsockname()[1], asyncio.get_running_loop(), stop)) async with edge: @@ -238,7 +248,7 @@ def edged(tls_edge, monkeypatch) -> Callable[[PolicyServer], str]: def url(server: PolicyServer) -> str: port, root = tls_edge(server.host, server.grpc_port) monkeypatch.setattr(grpc_wire, 'channel_credentials', lambda: grpc.ssl_channel_credentials(root)) - return f'grpcs://localhost:{port}' + return f'grpcs://{EDGE_HOST}:{port}' return url @@ -391,10 +401,22 @@ def test_a_server_on_the_grpc_ping_defaults_kills_the_silent_session( def test_a_certificate_the_client_cannot_verify_is_not_retried(both_wires, tls_edge, monkeypatch): """A root that does not cover the edge is permanent, so it surfaces on the first attempt.""" port, _root = tls_edge(both_wires[0].host, both_wires[0].grpc_port) - unrelated, _key = _self_signed('localhost') + unrelated, _key = _self_signed(EDGE_HOST) monkeypatch.setattr(grpc_wire, 'channel_credentials', lambda: grpc.ssl_channel_credentials(unrelated)) - client = InferenceClient(f'grpcs://localhost:{port}', open_timeout=2.0, connect_deadline=20.0) + _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE[0]) + + +def test_an_edge_that_selects_no_alpn_is_not_retried(both_wires, tls_edge, monkeypatch): + """A front fronting a raw TCP port terminates TLS and names no protocol, which gRPC cannot use.""" + port, root = tls_edge(both_wires[0].host, both_wires[0].grpc_port, alpn=False) + monkeypatch.setattr(grpc_wire, 'channel_credentials', lambda: grpc.ssl_channel_credentials(root)) + _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE[1]) + + +def _surfaces_at_once(url: str, blamed: str) -> None: + """Assert a connect to ``url`` fails naming ``blamed``, without spending its retry deadline.""" + client = InferenceClient(url, open_timeout=2.0, connect_deadline=20.0) started = time.monotonic() - with pytest.raises(ssl.SSLCertVerificationError, match=grpc_wire._UNVERIFIABLE_CERTIFICATE): + with pytest.raises(grpc.RpcError, match=blamed): client.new_session() assert time.monotonic() - started < 8.0, 'the connect retried a permanent failure' diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 0d89c7b23..4c3a3bb43 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -94,10 +94,12 @@ case " $* " in *) set -- "$@" "--idle_timeout_min=${NEBIUS_IDLE_TIMEOUT_MIN:-20}" ;; esac -GRPC_PORT=9000 -case " $* " in - *" --grpc_port="*|*" --grpc_port "*) ;; - *) set -- "$@" "--grpc_port=${GRPC_PORT}" ;; +# The endpoint exposes the port the server listens on, so a caller's own --grpc_port decides both. +ARGS=" $* " +case "$ARGS" in + *" --grpc_port="*) GRPC_PORT=${ARGS#*--grpc_port=}; GRPC_PORT=${GRPC_PORT%% *} ;; + *" --grpc_port "*) GRPC_PORT=${ARGS#*--grpc_port }; GRPC_PORT=${GRPC_PORT%% *} ;; + *) GRPC_PORT=9000; set -- "$@" "--grpc_port=${GRPC_PORT}" ;; esac SERVER_ARGS="run --python 3.13 ${EXTRA}python -m positronic.vendors.${VENDOR}.server $*" From 9904200df629e5dbd87dd2921f492bdf1f57ae70 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 08:15:53 +0000 Subject: [PATCH 12/46] State the ping defaults without counting them first Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/grpc_wire.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 737ff3701..a3559ee5a 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -72,8 +72,7 @@ def _client_options() -> list[tuple[str, int]]: *_MESSAGE_SIZE_OPTIONS, ('grpc.keepalive_time_ms', _PING_EVERY_MS), ('grpc.keepalive_timeout_ms', _PING_ANSWER_TIMEOUT_MS), - # Both of gRPC's own client throttles stop the pings during exactly the silent wait they - # exist for: it sends two and stops, and it spaces them five minutes apart. + # Left to itself gRPC sends two pings without data, five minutes apart. ('grpc.http2.max_pings_without_data', 0), ('grpc.http2.min_time_between_pings_ms', _PING_EVERY_MS), ] From 267a3a905f3a2b1e977df73835e56bf7485f6cfd Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 08:21:06 +0000 Subject: [PATCH 13/46] Say why the websocket is not shaped, rather than closing on the sound of it Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index b02f8a3b3..ee28fbbd1 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -37,7 +37,7 @@ Through such an endpoint an 846 KiB observation round-trips in about 6 ms over g 60 ms over the WebSocket, and gRPC holds that at 10 Hz, which is 8 MB/s of observation. The front shapes a session that outruns it: a back-to-back loop settles at about 83 ms a round trip after some 11 MB, and gets its speed back after a minute of quiet. The WebSocket holds its 60 ms throughout, -never being fast enough to be shaped. +below the rate the front shapes at. Both wires ping through a silent wait, so a front that drops a connection it has read nothing from — the managed one after about 90 s — does not cut an inference the model is still working on. From cc26f7d012f24fdd598256121adaac5aa8156c15 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 08:24:51 +0000 Subject: [PATCH 14/46] Leave the front's idle close to the README, and name the constant's job Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/grpc_wire.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index a3559ee5a..7a8330f6c 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -32,8 +32,7 @@ ('grpc.max_send_message_length', wire.MAX_MESSAGE_BYTES), ] -# A front closes a connection it has read nothing from — ~90s on the Nebius managed ingress — and one -# inference sends nothing until it answers, so the client pings through that silence. +# How often the client pings a connection nothing is crossing, so no front reads it as dead. _PING_EVERY_MS = 20_000 _PING_ANSWER_TIMEOUT_MS = 10_000 _PING_TOLERATED_EVERY_MS = 10_000 From 4e9fa7528058b4080cd980b5a5fd952be84ee3c8 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 08:36:56 +0000 Subject: [PATCH 15/46] Refuse a closed gRPC session, and drop the credentials alias MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two review findings. A timed-out inference closes the connection, and the server can still answer inside the five seconds `close` waits for the stream to end. That reply sat in the inbox, where the next `infer` would have read it — one observation's actions against the next observation's state. `send` and `recv` now refuse a closed session. `channel_credentials` wrapped one call and had one caller, so it was a production name held open for the tests. It is inlined, and the tests substitute roots on gRPC's own function. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/grpc_wire.py | 19 ++++++++++------ positronic/offboard/tests/test_grpc_wire.py | 24 ++++++++++++++++++--- 2 files changed, 34 insertions(+), 9 deletions(-) diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 7a8330f6c..03f2837af 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -86,15 +86,11 @@ def _server_options() -> list[tuple[str, int]]: ] -def channel_credentials() -> grpc.ChannelCredentials: - """The roots a ``grpcs://`` channel verifies the edge against: the system's own.""" - return grpc.ssl_channel_credentials() - - def _channel(target: str, secure: bool) -> grpc.Channel: options = _client_options() if secure: - return grpc.secure_channel(target, channel_credentials(), options=options) + # No roots named, so the channel verifies the edge against the system's own. + return grpc.secure_channel(target, grpc.ssl_channel_credentials(), options=options) return grpc.insecure_channel(target, options=options) @@ -158,10 +154,21 @@ def _read(self) -> None: finally: self._responses.cancel() + def _refuse_if_closed(self) -> None: + """Refuse a closed session, whose inbox may still hold a reply that arrived during ``close``. + + A timed-out inference closes here, so reading that reply would pair one observation's actions + with the next observation's state. + """ + if self._closed: + raise wire.PeerDisconnected(f'The session on {self._target} is closed') + def send(self, message: bytes) -> None: + self._refuse_if_closed() self._outbox.put(message) def recv(self, timeout: float | None = None) -> bytes: + self._refuse_if_closed() try: answer = self._inbox.get(timeout=timeout) except queue.Empty: diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index da29cacac..7fc833a46 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -241,13 +241,19 @@ async def _handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> loop.call_soon_threadsafe(stop.set) +def _trust_only(monkeypatch, root: bytes) -> None: + """Verify every channel this test opens against ``root``, in place of the system's own.""" + system_roots = grpc.ssl_channel_credentials + monkeypatch.setattr(grpc, 'ssl_channel_credentials', lambda: system_roots(root)) + + @pytest.fixture def edged(tls_edge, monkeypatch) -> Callable[[PolicyServer], str]: """The ``grpcs://`` URL of a server reached through a TLS edge, with the client trusting its root.""" def url(server: PolicyServer) -> str: port, root = tls_edge(server.host, server.grpc_port) - monkeypatch.setattr(grpc_wire, 'channel_credentials', lambda: grpc.ssl_channel_credentials(root)) + _trust_only(monkeypatch, root) return f'grpcs://{EDGE_HOST}:{port}' return url @@ -402,14 +408,14 @@ def test_a_certificate_the_client_cannot_verify_is_not_retried(both_wires, tls_e """A root that does not cover the edge is permanent, so it surfaces on the first attempt.""" port, _root = tls_edge(both_wires[0].host, both_wires[0].grpc_port) unrelated, _key = _self_signed(EDGE_HOST) - monkeypatch.setattr(grpc_wire, 'channel_credentials', lambda: grpc.ssl_channel_credentials(unrelated)) + _trust_only(monkeypatch, unrelated) _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE[0]) def test_an_edge_that_selects_no_alpn_is_not_retried(both_wires, tls_edge, monkeypatch): """A front fronting a raw TCP port terminates TLS and names no protocol, which gRPC cannot use.""" port, root = tls_edge(both_wires[0].host, both_wires[0].grpc_port, alpn=False) - monkeypatch.setattr(grpc_wire, 'channel_credentials', lambda: grpc.ssl_channel_credentials(root)) + _trust_only(monkeypatch, root) _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE[1]) @@ -420,3 +426,15 @@ def _surfaces_at_once(url: str, blamed: str) -> None: with pytest.raises(grpc.RpcError, match=blamed): client.new_session() assert time.monotonic() - started < 8.0, 'the connect retried a permanent failure' + + +def test_a_timed_out_session_refuses_the_next_inference(both_wires): + """The timeout closes the connection, and the server may answer inside the close's own wait.""" + server, policy = both_wires + policy._mock_session.side_effect = lambda *_: time.sleep(1.0) or [{'action': [1, 2, 3]}] + session = InferenceClient(grpc_url(server), infer_timeout=0.2).new_session() + with pytest.raises(TimeoutError): + session.infer({'image': 'test'}) + # Without the guard this answers the first observation's actions, against the second's state. + with pytest.raises(wire.PeerDisconnected): + session.infer({'image': 'test'}) From bb1d88841e725df9b1c255c88d5bc80a08b5ff48 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 08:41:02 +0000 Subject: [PATCH 16/46] State the session-param constraint without weighing the alternative Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/stub.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/positronic/offboard/stub.py b/positronic/offboard/stub.py index 0e7b6e0a5..faf2372f4 100644 --- a/positronic/offboard/stub.py +++ b/positronic/offboard/stub.py @@ -51,8 +51,7 @@ def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] class Delay(Layer): """Holds every answer for ``delay_sec``, standing in for a model slow enough to outlast a front's - idle close. A layer rather than an argument of the policy: a session param may tune the pipeline - around the model source, never the source itself. + idle close. A session param may tune the pipeline around the model source, never the source itself. """ def __init__(self, delay_sec: float = 0.0): From b4bc5da9e0a578fd12257fa35d6b7f7e95322b5f Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 09:08:43 +0000 Subject: [PATCH 17/46] Refuse a send into an ended stream, and pin the websocket's own ceiling Three review findings. gRPC stops reading the request iterator once a stream ends, so a send after the caller has been told why it ended sat in the outbox while the next `recv` waited out a whole inference timeout on an inbox nothing refills. The connection records the end when it hands out the reason, and refuses a send past it. The websocket client kept the library's 1 MiB receive default while the server, the gRPC client and the gRPC server all carry the shared 16 MiB ceiling, so a result in between closed the session rather than answering it. `serve.sh` printed `grpcs://:443` when the endpoint served no URL for the gRPC port. It fails there instead. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/client.py | 6 +++++- positronic/offboard/grpc_wire.py | 22 ++++++++++----------- positronic/offboard/tests/test_grpc_wire.py | 15 ++++++++++++++ workflows/nebius/serve.sh | 4 ++++ 4 files changed, 35 insertions(+), 12 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 803875806..51db2861f 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -252,7 +252,11 @@ def _connect(self) -> wire.ClientConnection: # after 60s — well inside one ``infer_timeout`` inference, which sends nothing until it # answers. The pings keep it open. websocket = connect( - self.session_url, open_timeout=self.open_timeout, additional_headers=self.headers, ping_interval=20.0 + self.session_url, + open_timeout=self.open_timeout, + additional_headers=self.headers, + ping_interval=20.0, + max_size=wire.MAX_MESSAGE_BYTES, ) return wire.WebsocketClientConnection(websocket) diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 03f2837af..26fb594de 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -133,6 +133,7 @@ def __init__( self._outbox: queue.SimpleQueue[bytes | None] = queue.SimpleQueue() self._inbox: queue.SimpleQueue[bytes | BaseException] = queue.SimpleQueue() self._closed = False + self._ended = False call = self._channel.stream_stream(METHOD_PATH, request_serializer=None, response_deserializer=None) self._responses = call(self._requests(), metadata=metadata) self._reader = threading.Thread(target=self._read, name='grpc-session-reader', daemon=True) @@ -154,26 +155,25 @@ def _read(self) -> None: finally: self._responses.cancel() - def _refuse_if_closed(self) -> None: - """Refuse a closed session, whose inbox may still hold a reply that arrived during ``close``. - - A timed-out inference closes here, so reading that reply would pair one observation's actions - with the next observation's state. - """ - if self._closed: - raise wire.PeerDisconnected(f'The session on {self._target} is closed') - def send(self, message: bytes) -> None: - self._refuse_if_closed() + # Past either of these gRPC has stopped reading the request iterator, so the write would sit in + # the outbox while ``recv`` waited out a whole inference timeout on an inbox nothing refills. + if self._closed or self._ended: + raise wire.PeerDisconnected(f'The session on {self._target} has ended') self._outbox.put(message) def recv(self, timeout: float | None = None) -> bytes: - self._refuse_if_closed() + # A closed session's inbox may hold a reply that arrived during ``close``, which would pair one + # observation's actions with the next observation's state. + if self._closed: + raise wire.PeerDisconnected(f'The session on {self._target} is closed') try: answer = self._inbox.get(timeout=timeout) except queue.Empty: raise TimeoutError(f'No message from {self._target} within {timeout}s') from None if isinstance(answer, BaseException): + # What ended the stream is queued once, so the caller learns why before writes are refused. + self._ended = True raise answer return answer diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index 7fc833a46..cfb9713e8 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -438,3 +438,18 @@ def test_a_timed_out_session_refuses_the_next_inference(both_wires): # Without the guard this answers the first observation's actions, against the second's state. with pytest.raises(wire.PeerDisconnected): session.infer({'image': 'test'}) + + +def test_a_connection_refuses_to_send_once_the_server_ends_the_stream(both_wires): + """gRPC stops reading the request iterator then, so a write would wait out a whole timeout.""" + server, _policy = both_wires + conn = grpc_wire.GrpcClientConnection(f'{server.host}:{server.grpc_port}', f'{wire.SESSION_PATH}/unknown-model', '') + try: + conn.recv(timeout=10.0) + # The server refuses the model in a frame, then ends the stream with that status. + with pytest.raises(grpc.RpcError): + conn.recv(timeout=10.0) + with pytest.raises(wire.PeerDisconnected): + conn.send(b'an observation the stream can no longer carry') + finally: + conn.close() diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 4c3a3bb43..e37c673a3 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -157,6 +157,10 @@ fi GRPC_HOST=$(nebius ai endpoint get "$ID" --format json 2>/dev/null \ | jq -r "[.status.public_endpoints[]? | select(startswith(\"https://port${GRPC_PORT}-\"))] | first // empty" \ | sed 's|^https://||') +if [ -z "$GRPC_HOST" ]; then + echo "The endpoint serves no https:// URL for port ${GRPC_PORT}. Check: nebius ai endpoint get $ID" >&2 + exit 1 +fi GRPC_URL="grpcs://${GRPC_HOST}:443" cat < Date: Wed, 9 Sep 2026 09:15:14 +0000 Subject: [PATCH 18/46] Name the subject first in two comments about an ended session Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/client.py | 2 +- positronic/offboard/grpc_wire.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 51db2861f..ca949bf26 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -132,7 +132,7 @@ def _refusal(e: Exception) -> _Refusal: Each gRPC code stands for the HTTP status its wire twin answers: ``PERMISSION_DENIED`` for 403, ``UNAVAILABLE`` for 503, ``RESOURCE_EXHAUSTED`` for 429. A TLS edge no client can use answers - ``UNAVAILABLE`` too, exactly as a cold backend does, so its details are what tell them apart. + ``UNAVAILABLE`` too, exactly as a cold backend does, so its details tell them apart. """ if isinstance(e, InvalidStatus): status = e.response.status_code diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 26fb594de..25bc8f304 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -156,8 +156,8 @@ def _read(self) -> None: self._responses.cancel() def send(self, message: bytes) -> None: - # Past either of these gRPC has stopped reading the request iterator, so the write would sit in - # the outbox while ``recv`` waited out a whole inference timeout on an inbox nothing refills. + # A write past either of these sits in the outbox while ``recv`` waits out a whole inference + # timeout on an inbox nothing refills: gRPC has stopped reading the request iterator. if self._closed or self._ended: raise wire.PeerDisconnected(f'The session on {self._target} has ended') self._outbox.put(message) From 3b24ad734633a14c65ef78fa9bbd29bbfcc5e886 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 10:22:52 +0000 Subject: [PATCH 19/46] Keep the connect refusal probe inside one attempt's budget `channel_ready_future` and the probe that asks why it failed each took the whole `open_timeout`, so a target that black-holes connection attempts blocked for nearly twice the documented limit before the connect loop saw a status. Both now run against one deadline, with a second reserved for the probe. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/grpc_wire.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 25bc8f304..525492555 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -7,6 +7,7 @@ import logging import queue import threading +import time import urllib.parse from collections.abc import AsyncIterator, Awaitable, Callable, Mapping @@ -44,6 +45,10 @@ # turns out to be up after all. _PROBE_PATH = f'/{SERVICE}/ChannelProbe' +# The slice of one connect attempt's budget the refusal probe may spend. A target that black-holes +# connection attempts answers neither, so both waits must fit inside the caller's ``open_timeout``. +_REFUSAL_PROBE_SEC = 1.0 + # What gRPC's status details call an edge no client can use: a certificate its roots do not cover, # and a front that selects no HTTP/2 over ALPN. UNUSABLE_EDGE = ('CERTIFICATE_VERIFY_FAILED', 'missing selected ALPN property') @@ -115,10 +120,12 @@ def __init__( ): self._target = target self._channel = _channel(target, secure) + deadline = time.monotonic() + open_timeout + ready_timeout = max(0.0, open_timeout - _REFUSAL_PROBE_SEC) try: - grpc.channel_ready_future(self._channel).result(timeout=open_timeout) + grpc.channel_ready_future(self._channel).result(timeout=ready_timeout) except grpc.FutureTimeoutError: - refusal = _connect_refusal(self._channel, timeout=open_timeout) + refusal = _connect_refusal(self._channel, timeout=max(0.0, deadline - time.monotonic())) self._channel.close() # An edge that refuses every client is permanent, so raise what gRPC blamed rather than a # timeout: the connect loop reads the status and stops instead of retrying its deadline out. From d93e06a1c4b50d70c8f47246d78cd7279923105f Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 10:25:07 +0000 Subject: [PATCH 20/46] Type the URL scheme a session opens on The scheme was dispatched as a closed set of strings, and `grpcs` was spelled a second time to recover its TLS half. A scheme added to one tuple and not the other would have left the wire and the security disagreeing. One enum member now carries both, and the parse fails on anything it does not name. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/client.py | 42 ++++++++++++++++++++++++++--------- 1 file changed, 31 insertions(+), 11 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index ca949bf26..8768d1455 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -175,8 +175,28 @@ def take(self, e: Exception) -> _ConnectOutcome: return _ConnectOutcome.RETRY if again else _ConnectOutcome.SURFACE -# The URL schemes that put a session on the gRPC wire: plaintext, and behind a TLS edge. -_GRPC_SCHEMES = ('grpc', 'grpcs') +class _Scheme(Enum): + """A URL scheme a session may open on, and what it settles: which wire, and whether it is TLS.""" + + EMPTY = ('', False, False) + HTTP = ('http', False, False) + WS = ('ws', False, False) + HTTPS = ('https', True, False) + WSS = ('wss', True, False) + GRPC = ('grpc', False, True) + GRPCS = ('grpcs', True, True) + + def __init__(self, text: str, secure: bool, grpc_wired: bool): + self.text = text + self.secure = secure + self.grpc_wired = grpc_wired + + @classmethod + def of(cls, text: str) -> '_Scheme': + for scheme in cls: + if scheme.text == text: + return scheme + raise ValueError(f'Unsupported scheme {text!r}') class InferenceClient: @@ -210,15 +230,15 @@ def __init__( infer_timeout: float = DEFAULT_INFER_TIMEOUT, ): split = urllib.parse.urlsplit(url if '://' in url else f'//{url}') - if split.scheme not in ('', 'http', 'ws', 'https', 'wss', *_GRPC_SCHEMES): - raise ValueError(f'Unsupported scheme {split.scheme!r} in {url!r}') + try: + scheme = _Scheme.of(split.scheme) + except ValueError: + raise ValueError(f'Unsupported scheme {split.scheme!r} in {url!r}') from None if not split.hostname: raise ValueError(f'No host in {url!r}') - grpc_wired = split.scheme in _GRPC_SCHEMES - secure = split.scheme in ('https', 'wss', 'grpcs') - session_scheme = split.scheme if grpc_wired else ('wss' if secure else 'ws') - http_scheme = 'https' if secure else 'http' - default_port = 443 if secure else 80 + session_scheme = scheme.text if scheme.grpc_wired else ('wss' if scheme.secure else 'ws') + http_scheme = 'https' if scheme.secure else 'http' + default_port = 443 if scheme.secure else 80 # urlsplit strips the brackets an IPv6 host needs back in a netloc. host = f'[{split.hostname}]' if ':' in split.hostname else split.hostname port = default_port if split.port is None else split.port @@ -228,8 +248,8 @@ def __init__( query = f'?{split.query}' if split.query else '' self._session_path = _session_path(split.path, url) self._query = split.query - self._grpc_target = f'{host}:{port}' if grpc_wired else None - self._grpc_secure = secure + self._grpc_target = f'{host}:{port}' if scheme.grpc_wired else None + self._grpc_secure = scheme.secure self.session_url = f'{session_scheme}://{netloc}{self._session_path}{query}' self.api_url = None if self._grpc_target else f'{http_scheme}://{netloc}/api/v1' self.headers = dict(headers) if headers else None From 931109aa10a9c6807992742193d9386ba5dba964 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 10:28:36 +0000 Subject: [PATCH 21/46] Name the websocket port once in the serve script The create declared `--container-port 8000` and the poll selected `https://port8000-`, so a port change could leave the reader looking for an endpoint creation never made. Ticket: Positronic-Robotics/internal#1191 #refs --- workflows/nebius/serve.sh | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index e37c673a3..046bdb52a 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -94,6 +94,10 @@ case " $* " in *) set -- "$@" "--idle_timeout_min=${NEBIUS_IDLE_TIMEOUT_MIN:-20}" ;; esac +# The port the websocket wire listens on, named once: the create declares it and the poll below +# selects the managed URL that fronts it. +WS_PORT=8000 + # The endpoint exposes the port the server listens on, so a caller's own --grpc_port decides both. ARGS=" $* " case "$ARGS" in @@ -112,7 +116,7 @@ nebius ai endpoint create \ --image "$IMAGE" \ --container-command uv \ --args "$SERVER_ARGS" \ - --container-port 8000 \ + --container-port "${WS_PORT}" \ --container-port "${GRPC_PORT}" \ --platform gpu-h100-sxm \ --preset "$PRESET" \ @@ -144,7 +148,7 @@ for i in $(seq 1 30); do # prefix. This field also carries bare `IP:port` entries, which serve no TLS and would put the # bearer token on the wire in cleartext — take the https:// ones, and fail rather than fall back. URL=$(nebius ai endpoint get "$ID" --format json 2>/dev/null \ - | jq -r '[.status.public_endpoints[]? | select(startswith("https://port8000-"))] | first // empty') + | jq -r "[.status.public_endpoints[]? | select(startswith(\"https://port${WS_PORT}-\"))] | first // empty") if [ -n "$URL" ]; then break; fi sleep 10 done From 3c26716816f2626e47af3900c1c7d95e53ba13c8 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 10:29:01 +0000 Subject: [PATCH 22/46] Name a remote session for what it is, not for one wire it can take `round_trip` and `_Endpoint` called every session `ws_session`, which stopped being true when a session could open on gRPC. `_Endpoint.close` set its client to `None`, so a field that is never absent read as optional and every use of it needed a grandfathered baseline entry. Sessions own the connections and the client holds only where to open one, so the close has nothing to release and the two entries go. Ticket: Positronic-Robotics/internal#1191 #refs --- .basedpyright/baseline.json | 18 ------ .../offboard/tests/test_remote_policy.py | 56 +++++++++---------- positronic/policy/remote.py | 23 ++++---- 3 files changed, 40 insertions(+), 57 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 220d04b59..22180d033 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -4884,24 +4884,6 @@ } } ], - "./positronic/policy/remote.py": [ - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 38, - "endColumn": 49, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 34, - "endColumn": 45, - "lineCount": 1 - } - } - ], "./positronic/policy/tests/test_golden_pipeline.py": [ { "code": "reportOptionalMemberAccess", diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 060644a38..bf1c2f341 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -26,7 +26,7 @@ CHUNKED_STACK = {'local_stack': {'name': 'chunked_schedule'}} -def _mock_ws_session(metadata=None): +def _mock_session(metadata=None): session = MagicMock() session.metadata = metadata or {} session.infer.return_value = {'action': 'test'} @@ -34,20 +34,20 @@ def _mock_ws_session(metadata=None): def _mock_remote_policy(metadata=None, infer_return=None): - """A RemotePolicy whose wire client is mocked out; returns (policy, mock_ws).""" - mock_ws = _mock_ws_session(metadata) + """A RemotePolicy whose wire client is mocked out; returns (policy, mock_session).""" + mock_session = _mock_session(metadata) if infer_return is not None: - mock_ws.infer.return_value = infer_return + mock_session.infer.return_value = infer_return policy = RemotePolicy('localhost:0') policy._endpoint._client = MagicMock() - policy._endpoint._client.new_session.return_value = mock_ws - return policy, mock_ws + policy._endpoint._client.new_session.return_value = mock_session + return policy, mock_session def _mock_endpoint(metadata=None, infer_return=None): """The bare wire connection, with no declared stack in front of it.""" - policy, mock_ws = _mock_remote_policy(metadata, infer_return) - return policy._endpoint, mock_ws + policy, mock_session = _mock_remote_policy(metadata, infer_return) + return policy._endpoint, mock_session def _make_image(h, w): @@ -316,8 +316,8 @@ def test_remote_session_normalizes_single_dict(open_session): def test_remote_session_passes_through_none(open_session): - endpoint, mock_ws = _mock_endpoint() - mock_ws.infer.return_value = None + endpoint, mock_session = _mock_endpoint() + mock_session.infer.return_value = None session, rt = open_session(endpoint) assert round_trip(session, rt, {}) is None @@ -327,7 +327,7 @@ def test_a_call_while_a_round_trip_is_in_flight_answers_none(open_session): """A session never waits. Every call while the round trip is in flight answers ``None``, and none of them starts a second round trip.""" chunk = [{'a': 1, 'timestamp': 0.0}] - endpoint, mock_ws = _mock_endpoint() + endpoint, mock_session = _mock_endpoint() started, release = threading.Event(), threading.Event() def blocked(obs): @@ -335,13 +335,13 @@ def blocked(obs): assert release.wait(ANSWER_SEC), 'the test never released the round-trip' return chunk - mock_ws.infer.side_effect = blocked + mock_session.infer.side_effect = blocked session, rt = open_session(endpoint) assert session({}, 0) is None assert started.wait(ANSWER_SEC), 'the round-trip never started' assert session({}, 0) is None - assert mock_ws.infer.call_count == 1 + assert mock_session.infer.call_count == 1 release.set() rt.wait(ANSWER_SEC) @@ -360,7 +360,7 @@ def test_opening_a_session_without_a_runtime_is_refused(): def test_cancel_drops_the_chunk_of_the_round_trip_in_flight(open_session): """A cancelled session drops the chunk it waited for, because that chunk applies to a world the cancel says has gone, and it asks for a new one.""" - endpoint, mock_ws = _mock_endpoint(infer_return=[{'a': 1, 'timestamp': 0.0}]) + endpoint, mock_session = _mock_endpoint(infer_return=[{'a': 1, 'timestamp': 0.0}]) session, rt = open_session(endpoint) assert session({}, 0) is None @@ -370,14 +370,14 @@ def test_cancel_drops_the_chunk_of_the_round_trip_in_flight(open_session): assert session({}, 0) is None # the cancelled answer, read and thrown away assert session({}, 0) is None # a round-trip of its own rt.wait(ANSWER_SEC) - assert mock_ws.infer.call_count == 2 + assert mock_session.infer.call_count == 2 def test_a_cancelled_round_trip_still_raises_what_it_failed_with(open_session): """A dropped chunk drops no failure. The session reads a cancelled answer, so a stalled server raises to the caller that asked for the episode.""" - endpoint, mock_ws = _mock_endpoint() - mock_ws.infer.side_effect = TimeoutError('server stalled') + endpoint, mock_session = _mock_endpoint() + mock_session.infer.side_effect = TimeoutError('server stalled') session, rt = open_session(endpoint) assert session({}, 0) is None @@ -391,8 +391,8 @@ def test_a_cancelled_round_trip_still_raises_what_it_failed_with(open_session): def test_a_cancel_dies_with_the_answer_it_was_made_against(open_session): """A cancel ends with the round trip it was made against, even when that round trip fails. A caller that catches the failure and keeps the session gets the next chunk.""" - endpoint, mock_ws = _mock_endpoint(infer_return=[{'a': 1, 'timestamp': 0.0}]) - mock_ws.infer.side_effect = [TimeoutError('server stalled'), [{'a': 1, 'timestamp': 0.0}]] + endpoint, mock_session = _mock_endpoint(infer_return=[{'a': 1, 'timestamp': 0.0}]) + mock_session.infer.side_effect = [TimeoutError('server stalled'), [{'a': 1, 'timestamp': 0.0}]] session, rt = open_session(endpoint) assert session({}, 0) is None @@ -407,14 +407,14 @@ def test_a_cancel_dies_with_the_answer_it_was_made_against(open_session): def test_closing_a_session_with_a_round_trip_in_flight_is_refused(open_session): """A runtime closes before the session it serves. A caller that closes the websocket under a round trip gets an error that names the order, and not a failure on a dead socket.""" - endpoint, mock_ws = _mock_endpoint() + endpoint, mock_session = _mock_endpoint() release = threading.Event() def blocked(obs): assert release.wait(ANSWER_SEC), 'the test never released the round-trip' return None - mock_ws.infer.side_effect = blocked + mock_session.infer.side_effect = blocked session, _rt = open_session(endpoint) assert session({}, 0) is None @@ -460,8 +460,8 @@ def _stamp_encode(image): def test_records_infer_span_when_inference_raises(tmp_path, open_session): """A round trip that raises still records the time it took to fail, and the answer raises it again at the call that reads it.""" - endpoint, mock_ws = _mock_endpoint() - mock_ws.infer.side_effect = TimeoutError('server stalled') + endpoint, mock_session = _mock_endpoint() + mock_session.infer.side_effect = TimeoutError('server stalled') session, rt = open_session(endpoint) with telemetry.bind(tmp_path, telemetry_keys.HARNESS_PROCESS, 'run-infer-raise'): with pytest.raises(TimeoutError): @@ -486,7 +486,7 @@ def test_empty_declaration_fails_before_motion(): def test_declared_stack_built_at_session_open(open_session): """The server-declared local stack runs in front of the connection.""" - policy, mock_ws = _mock_remote_policy(CHUNKED_STACK, infer_return=[{'a': 1, 'timestamp': 0.0}]) + policy, mock_session = _mock_remote_policy(CHUNKED_STACK, infer_return=[{'a': 1, 'timestamp': 0.0}]) session, rt = open_session(policy) assert round_trip(session, rt, {keys.OBS_TIME_NS: 0}, int(1e9)) == [{'a': 1, 'timestamp': 1.0}] @@ -503,19 +503,19 @@ def test_unknown_declared_entry_fails_before_motion(): def test_compression_follows_the_server_declaration(open_session): """A server behind a message-size cap declares ``remote(compress_images=True)`` and the rig obeys.""" - endpoint, mock_ws = _mock_endpoint({'compress_images': True}, infer_return=[]) + endpoint, mock_session = _mock_endpoint({'compress_images': True}, infer_return=[]) session, rt = open_session(endpoint) round_trip(session, rt, {'cam': _make_image(48, 64)}) - assert isinstance(mock_ws.infer.call_args.args[0]['cam'], dict) + assert isinstance(mock_session.infer.call_args.args[0]['cam'], dict) def test_frames_stay_raw_where_the_server_declares_no_compression(open_session): - endpoint, mock_ws = _mock_endpoint({'compress_images': False}, infer_return=[]) + endpoint, mock_session = _mock_endpoint({'compress_images': False}, infer_return=[]) session, rt = open_session(endpoint) round_trip(session, rt, {'cam': _make_image(48, 64)}) - assert isinstance(mock_ws.infer.call_args.args[0]['cam'], np.ndarray) + assert isinstance(mock_session.infer.call_args.args[0]['cam'], np.ndarray) # rules-allow: hardcoded-keys — the command mapping below is spelled the way a server sends it. Reading diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index b6a4c7b8e..d5dda95b5 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -42,7 +42,7 @@ def _prepare_obs(obs: cabc.Mapping[str, Any], compress_images: bool) -> dict[str def round_trip( - ws_session: InferenceSession, obs: cabc.Mapping[str, Any], compress_images: bool + session: InferenceSession, obs: cabc.Mapping[str, Any], compress_images: bool ) -> list[dict[str, Any]] | dict[str, Any]: """One inference over the wire, timed as the ``policy.infer`` span. @@ -53,7 +53,7 @@ def round_trip( prepared = _prepare_obs(obs, compress_images) infer_start_ns = time.time_ns() try: - return ws_session.infer(prepared) + return session.infer(prepared) finally: telemetry.record_span(telemetry_keys.SPAN_POLICY_INFER, infer_start_ns, time.time_ns()) @@ -68,8 +68,8 @@ class RemoteSession(Session): ``compress_images`` comes from what the server declared (see ``RemoteMarker``). """ - def __init__(self, ws_session: InferenceSession, rt: Runtime, compress_images: bool = False): - self._session = ws_session + def __init__(self, session: InferenceSession, rt: Runtime, compress_images: bool = False): + self._session = session self._rt = rt self._compress_images = compress_images self._answer: Answer | None = None @@ -109,7 +109,7 @@ def close(self): in_flight = self._answer is not None and not self._answer.done() logger.info('RemoteSession.close: answer_in_flight=%s', in_flight) assert not in_flight, ( - 'close the runtime serving this session first: the round trip in flight uses the websocket that this closes' + 'close the runtime serving this session first: the round trip in flight uses the connection this closes' ) self._session.close() logger.info('RemoteSession.close: session closed') @@ -128,26 +128,27 @@ def __init__(self, url: str, *, headers: dict[str, str] | None, infer_timeout: f def server_meta(self) -> dict[str, Any]: if self._server_meta is None: - ws_session = self._client.new_session() + session = self._client.new_session() try: - self._server_meta = dict(ws_session.metadata) + self._server_meta = dict(session.metadata) finally: - ws_session.close() + session.close() return self._server_meta def new_session(self, context=None, rt=None) -> RemoteSession: if rt is None: raise ValueError('A remote session runs its inference on a runtime: pass rt to new_session.') compress = bool(self.server_meta().get(offboard_keys.COMPRESS_IMAGES)) - ws_session = self._client.new_session() - return RemoteSession(ws_session, rt, compress_images=compress) + session = self._client.new_session() + return RemoteSession(session, rt, compress_images=compress) @property def functions(self) -> cabc.Mapping[str, cabc.Callable[..., Any]]: return {INFER: round_trip} def close(self): - self._client = None + # Sessions own the connections; the client itself holds only where to open one. + pass class RemotePolicy(Policy): From 030eba28203bd91def28b01656449dbc2230a03d Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 9 Sep 2026 10:37:02 +0000 Subject: [PATCH 23/46] Name the compress-images metadata key once in its test The server declares it under `offboard_keys.COMPRESS_IMAGES` and the rig reads it under the same constant; the test spelled the string itself, so a rename would have left the test asserting on a key nothing writes. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/tests/test_remote_policy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index bf1c2f341..a5a346776 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -439,7 +439,7 @@ def test_infer_span_excludes_client_side_image_preparation(tmp_path, open_sessio """``policy.infer`` is the remote round-trip, so JPEG-encoding the observation stays outside it: folding client CPU work into the span would inflate the inference percentiles and the policy-server capacity estimate the report derives from them.""" - endpoint, _ = _mock_endpoint({'compress_images': True}, infer_return=[]) + endpoint, _ = _mock_endpoint({offboard_keys.COMPRESS_IMAGES: True}, infer_return=[]) session, rt = open_session(endpoint) encoded_at: list[int] = [] @@ -503,7 +503,7 @@ def test_unknown_declared_entry_fails_before_motion(): def test_compression_follows_the_server_declaration(open_session): """A server behind a message-size cap declares ``remote(compress_images=True)`` and the rig obeys.""" - endpoint, mock_session = _mock_endpoint({'compress_images': True}, infer_return=[]) + endpoint, mock_session = _mock_endpoint({offboard_keys.COMPRESS_IMAGES: True}, infer_return=[]) session, rt = open_session(endpoint) round_trip(session, rt, {'cam': _make_image(48, 64)}) @@ -511,7 +511,7 @@ def test_compression_follows_the_server_declaration(open_session): def test_frames_stay_raw_where_the_server_declares_no_compression(open_session): - endpoint, mock_session = _mock_endpoint({'compress_images': False}, infer_return=[]) + endpoint, mock_session = _mock_endpoint({offboard_keys.COMPRESS_IMAGES: False}, infer_return=[]) session, rt = open_session(endpoint) round_trip(session, rt, {'cam': _make_image(48, 64)}) From deffec0c23e29baaafc9a848e935e0f101c70806 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Thu, 10 Sep 2026 08:00:13 +0000 Subject: [PATCH 24/46] Cover the close report on both wires, and place two helpers above their callers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `InferenceSession.close` logs one line whichever wire carries the session, and each wire fills it with what only that protocol can say. A test drives both against one server and reads the pair: the websocket reports close code 1000, and the gRPC end reports that the server ended the stream. The test that ends a stream from the server now asserts the other reading, so a close the peer had already ended is distinguishable from one the server answered — which is the whole point of the line. `_server_options` moves beside `serve`, and `_surfaces_at_once` above the first test that calls it. The three ping constants stay together: the client's ping interval and the server's tolerance of it are one contract, and a reader has to see them side by side to check they agree. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/grpc_wire.py | 23 +++++++------ positronic/offboard/tests/test_grpc_wire.py | 38 ++++++++++++++++----- 2 files changed, 42 insertions(+), 19 deletions(-) diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 525492555..e2c640ec6 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -82,15 +82,6 @@ def _client_options() -> list[tuple[str, int]]: ] -def _server_options() -> list[tuple[str, int]]: - return [ - *_MESSAGE_SIZE_OPTIONS, - # gRPC's own defaults, a five-minute floor and two strikes, answer a 20s ping with GOAWAY. - ('grpc.http2.min_ping_interval_without_data_ms', _PING_TOLERATED_EVERY_MS), - ('grpc.http2.max_ping_strikes', 0), - ] - - def _channel(target: str, secure: bool) -> grpc.Channel: options = _client_options() if secure: @@ -196,7 +187,10 @@ def close(self) -> str: self._channel.close() # The websocket wire reads the same two facts off a close code. A stream the server never ended # means it still holds this session, so the next one's handshake waits on a slot nobody released. - return f'peer had ended the stream {self._ended}, server ended it within {_CLOSE_TIMEOUT_SEC}s {server_ended_stream}' + return ( + f'peer had ended the stream {self._ended}, ' + f'server ended it within {_CLOSE_TIMEOUT_SEC}s {server_ended_stream}' + ) def model_id_of(session_path: str) -> str | None: @@ -253,6 +247,15 @@ def _bind_target(host: str, port: int) -> str: return f'[{host}]:{port}' if ':' in host else f'{host}:{port}' +def _server_options() -> list[tuple[str, int]]: + return [ + *_MESSAGE_SIZE_OPTIONS, + # gRPC's own defaults, a five-minute floor and two strikes, answer a 20s ping with GOAWAY. + ('grpc.http2.min_ping_interval_without_data_ms', _PING_TOLERATED_EVERY_MS), + ('grpc.http2.max_ping_strikes', 0), + ] + + async def serve( serve_session: Callable[[GrpcServerConnection], Awaitable[None]], authorized: Callable[[Mapping[str, str]], bool], diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index cfb9713e8..a7977cbae 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -3,6 +3,7 @@ import asyncio import datetime import ipaddress +import logging import pathlib import queue import ssl @@ -70,6 +71,23 @@ def test_both_wires_answer_one_observation_alike(both_wires): over_grpc.close() +def test_both_wires_report_what_their_close_saw(both_wires, caplog): + """A close the server answered has to read differently from one it never saw, whichever wire carried + the session. The second leaves the server holding the slot, and the next session's handshake waits on + it, so each wire reports the distinction in the terms its own protocol offers.""" + server, _policy = both_wires + over_ws = InferenceClient(f'{server.host}:{server.port}').new_session() + over_grpc = InferenceClient(grpc_url(server)).new_session() + + with caplog.at_level(logging.INFO, logger='positronic.offboard.client'): + over_ws.close() + over_grpc.close() + + ws_report, grpc_report = (r.getMessage() for r in caplog.records if 'InferenceSession.close' in r.getMessage()) + assert 'close code 1000' in ws_report # the server answered the close frame + assert 'server ended it within 5.0s True' in grpc_report + + def test_closing_a_session_ends_it_on_the_server(both_wires): """``close`` half-closes the stream and waits, so the server releases the session before it returns.""" server, _policy = both_wires @@ -404,6 +422,15 @@ def test_a_server_on_the_grpc_ping_defaults_kills_the_silent_session( _silent_then_infer(server) +def _surfaces_at_once(url: str, blamed: str) -> None: + """Assert a connect to ``url`` fails naming ``blamed``, without spending its retry deadline.""" + client = InferenceClient(url, open_timeout=2.0, connect_deadline=20.0) + started = time.monotonic() + with pytest.raises(grpc.RpcError, match=blamed): + client.new_session() + assert time.monotonic() - started < 8.0, 'the connect retried a permanent failure' + + def test_a_certificate_the_client_cannot_verify_is_not_retried(both_wires, tls_edge, monkeypatch): """A root that does not cover the edge is permanent, so it surfaces on the first attempt.""" port, _root = tls_edge(both_wires[0].host, both_wires[0].grpc_port) @@ -419,15 +446,6 @@ def test_an_edge_that_selects_no_alpn_is_not_retried(both_wires, tls_edge, monke _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE[1]) -def _surfaces_at_once(url: str, blamed: str) -> None: - """Assert a connect to ``url`` fails naming ``blamed``, without spending its retry deadline.""" - client = InferenceClient(url, open_timeout=2.0, connect_deadline=20.0) - started = time.monotonic() - with pytest.raises(grpc.RpcError, match=blamed): - client.new_session() - assert time.monotonic() - started < 8.0, 'the connect retried a permanent failure' - - def test_a_timed_out_session_refuses_the_next_inference(both_wires): """The timeout closes the connection, and the server may answer inside the close's own wait.""" server, policy = both_wires @@ -451,5 +469,7 @@ def test_a_connection_refuses_to_send_once_the_server_ends_the_stream(both_wires conn.recv(timeout=10.0) with pytest.raises(wire.PeerDisconnected): conn.send(b'an observation the stream can no longer carry') + # The peer ended the stream, so this close reads differently from one the server answered. + assert 'peer had ended the stream True' in conn.close() finally: conn.close() From f7100442d6a1a2746f638bc3872f58fec953fa9f Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Thu, 10 Sep 2026 08:16:32 +0000 Subject: [PATCH 25/46] Hold a test server's websocket port from the draw to the serve `start_server` drew a port, closed the socket, and passed the number to uvicorn. These tests run in parallel, so another worker could bind that port in between; the loser then failed to bind and the fixture reported `Server failed to start`. It reddened the 3.11 job on this branch. The socket now stays bound and is handed to `uvicorn.Server.serve`, so no window exists. Its family comes from the host, because one test binds `::1`. The gRPC port keeps drawing a number, which gRPC binds itself. That bind cannot report the same collision: gRPC sets `SO_REUSEPORT`, so a second server binds the same port and the kernel shares the connections between them. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/tests/conftest.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 0a7b5d480..b02c9c5d9 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -16,6 +16,21 @@ from positronic.policy.spec import ModelSource, PolicySource, remote +def _bind_free_socket(host: str) -> socket.socket: + """A socket holding a free port on ``host``, to hand to the server that will serve on it. + + Drawing a port and closing the socket loses the port to whoever binds next, and these tests run in + parallel. Staying bound from the draw to the serve is what makes the port ours. The family comes from + ``host`` itself, so an IPv6 test binds an IPv6 socket. + """ + family, _type, _proto, _canon, address = socket.getaddrinfo( + host, 0, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE + )[0] + sock = socket.socket(family, socket.SOCK_STREAM) + sock.bind(address) + return sock + + def _find_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(('', 0)) @@ -36,7 +51,10 @@ def start_server() -> Generator[StartServer, None, None]: def start(pipeline, *, grpc: bool = False, **server_kwargs) -> tuple[str, int, PolicyServer]: grpc_port = _find_free_port() if grpc else None host = server_kwargs.pop('host', 'localhost') - server = PolicyServer(pipeline, host=host, port=_find_free_port(), grpc_port=grpc_port, **server_kwargs) + ws_socket = _bind_free_socket(host) + server = PolicyServer( + pipeline, host=host, port=ws_socket.getsockname()[1], grpc_port=grpc_port, **server_kwargs + ) uv_server = uvicorn.Server( uvicorn.Config( server.app, host=server.host, port=server.port, log_level='warning', ws_max_size=wire.MAX_MESSAGE_BYTES @@ -48,7 +66,7 @@ async def _run(): # Started first, so the websocket port answering means both wires are up. grpc_server = await server._start_grpc() if grpc_port is not None else None try: - await uv_server.serve() + await uv_server.serve(sockets=[ws_socket]) finally: if grpc_server is not None: await grpc_server.stop(grace=None) From 4907e2e42e6743f3b9b60c6a56524181e3044aaa Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Thu, 10 Sep 2026 21:33:08 +0000 Subject: [PATCH 26/46] Keep a short open timeout usable, and bind the gRPC test wire on any free port A fixed one-second probe reservation left the readiness wait nothing when `open_timeout` was at or under it, so a healthy server timed out. The probe now takes half the budget at most, and an UNIMPLEMENTED answer proves the channel is up. `grpc_wire.serve` returns the port it bound and `PolicyServer.grpc_port` becomes that port, so the tests ask for 0 and hold the port through startup instead of drawing one and closing the socket another worker can take. A cancel during `asyncio.to_thread(session, ...)` does not stop the worker, so the session close runs beside a live inference. The server logs it, which turns a silent backend corruption into a grep; guarding the race is a change to the concurrency model and stays open. `_session_path` and `_connect_refusal` move above their only callers, and `_Endpoint.close` goes: it only repeated the inherited no-op. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/client.py | 32 +++++------ positronic/offboard/grpc_wire.py | 61 +++++++++++++-------- positronic/offboard/server.py | 48 ++++++++++------ positronic/offboard/tests/conftest.py | 14 ++--- positronic/offboard/tests/test_grpc_wire.py | 11 ++++ positronic/policy/remote.py | 4 -- 6 files changed, 99 insertions(+), 71 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 8768d1455..f118f1461 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -95,22 +95,6 @@ def close(self): logger.info('InferenceSession.close: %s', self._conn.close()) -def _session_path(path: str, url: str) -> str: - """The session path a URL names: ``/api/v1/session``, plus the model id it addresses, if any. - - A URL naming no model — a bare host, or the endpoint with or without a trailing slash — addresses the - endpoint itself, which serves whatever the server pinned. - """ - if path.rstrip('/') in ('', wire.SESSION_PATH): - return wire.SESSION_PATH - if not path.startswith(f'{wire.SESSION_PATH}/'): - raise ValueError(f'Unexpected path {path!r} in {url!r}; expected {wire.SESSION_PATH}[/]') - # Kept as written, percent-encoding included, so the server decodes exactly the id whoever handed out - # the URL meant: a trailing slash is part of that id, and an id may itself be a path (a HuggingFace - # repo, say), whose own slashes stay separators. - return path - - class _ConnectOutcome(Enum): RETRY = 'retry' SURFACE = 'surface' @@ -199,6 +183,22 @@ def of(cls, text: str) -> '_Scheme': raise ValueError(f'Unsupported scheme {text!r}') +def _session_path(path: str, url: str) -> str: + """The session path a URL names: ``/api/v1/session``, plus the model id it addresses, if any. + + A URL naming no model — a bare host, or the endpoint with or without a trailing slash — addresses the + endpoint itself, which serves whatever the server pinned. + """ + if path.rstrip('/') in ('', wire.SESSION_PATH): + return wire.SESSION_PATH + if not path.startswith(f'{wire.SESSION_PATH}/'): + raise ValueError(f'Unexpected path {path!r} in {url!r}; expected {wire.SESSION_PATH}[/]') + # Kept as written, percent-encoding included, so the server decodes exactly the id whoever handed out + # the URL meant: a trailing slash is part of that id, and an id may itself be a path (a HuggingFace + # repo, say), whose own slashes stay separators. + return path + + class InferenceClient: """The wire connection to one inference server, addressed by one URL. diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index e2c640ec6..7ba6d9c3d 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -45,7 +45,7 @@ # turns out to be up after all. _PROBE_PATH = f'/{SERVICE}/ChannelProbe' -# The slice of one connect attempt's budget the refusal probe may spend. A target that black-holes +# The largest slice of one connect attempt's budget the refusal probe may spend. A target that black-holes # connection attempts answers neither, so both waits must fit inside the caller's ``open_timeout``. _REFUSAL_PROBE_SEC = 1.0 @@ -59,18 +59,6 @@ def edge_is_unusable(details: str) -> bool: return any(marker in details for marker in UNUSABLE_EDGE) -def _connect_refusal(channel: grpc.Channel, timeout: float) -> grpc.RpcError | None: - """What gRPC says stopped the channel coming up. Its readiness future carries only that it did not.""" - probe = channel.stream_stream(_PROBE_PATH, request_serializer=None, response_deserializer=None) - try: - next(probe(iter(()), timeout=timeout)) - except grpc.RpcError as e: - return e - except StopIteration: - return None - return None - - def _client_options() -> list[tuple[str, int]]: return [ *_MESSAGE_SIZE_OPTIONS, @@ -90,6 +78,27 @@ def _channel(target: str, secure: bool) -> grpc.Channel: return grpc.insecure_channel(target, options=options) +def _probe_share(open_timeout: float) -> float: + """What one connect attempt gives the refusal probe, leaving the readiness wait the rest. + + Half at most, so an ``open_timeout`` under ``_REFUSAL_PROBE_SEC`` still waits for a healthy server + instead of going straight to asking why it is down. + """ + return min(_REFUSAL_PROBE_SEC, open_timeout / 2) + + +def _connect_refusal(channel: grpc.Channel, timeout: float) -> grpc.RpcError | None: + """What gRPC says stopped the channel coming up. Its readiness future carries only that it did not.""" + probe = channel.stream_stream(_PROBE_PATH, request_serializer=None, response_deserializer=None) + try: + next(probe(iter(()), timeout=timeout)) + except grpc.RpcError as e: + return e + except StopIteration: + return None + return None + + class GrpcClientConnection: """A client's end of one gRPC session. @@ -112,17 +121,19 @@ def __init__( self._target = target self._channel = _channel(target, secure) deadline = time.monotonic() + open_timeout - ready_timeout = max(0.0, open_timeout - _REFUSAL_PROBE_SEC) try: - grpc.channel_ready_future(self._channel).result(timeout=ready_timeout) + grpc.channel_ready_future(self._channel).result(timeout=open_timeout - _probe_share(open_timeout)) except grpc.FutureTimeoutError: refusal = _connect_refusal(self._channel, timeout=max(0.0, deadline - time.monotonic())) - self._channel.close() - # An edge that refuses every client is permanent, so raise what gRPC blamed rather than a - # timeout: the connect loop reads the status and stops instead of retrying its deadline out. - if refusal is not None and edge_is_unusable(refusal.details() or ''): - raise refusal from None - raise TimeoutError(f'gRPC channel to {target} is not ready within {open_timeout}s') from None + # The probe path is served by no handler, so an UNIMPLEMENTED means the edge carried the call: + # the channel is up, and the readiness wait was short rather than the server absent. + if refusal is None or refusal.code() is not grpc.StatusCode.UNIMPLEMENTED: + self._channel.close() + # An edge that refuses every client is permanent, so raise what gRPC blamed rather than a + # timeout: the connect loop reads the status and stops instead of retrying its deadline out. + if refusal is not None and edge_is_unusable(refusal.details() or ''): + raise refusal from None + raise TimeoutError(f'gRPC channel to {target} is not ready within {open_timeout}s') from None # gRPC metadata keys are lower case, and they are the same header names the websocket wire sends. metadata = tuple((key.lower(), value) for key, value in (headers or {}).items()) + ( (SESSION_PATH_HEADER, session_path), @@ -261,8 +272,10 @@ async def serve( authorized: Callable[[Mapping[str, str]], bool], host: str, port: int, -) -> grpc.aio.Server: - """Start a gRPC server that gives every accepted session to ``serve_session``. +) -> tuple[grpc.aio.Server, int]: + """Start a gRPC server that gives every accepted session to ``serve_session``, and say what it bound. + + A ``port`` of 0 binds any free one, which is the port that comes back. ``authorized`` reads the session headers and refuses before the session opens, as the websocket wire refuses the upgrade. @@ -292,4 +305,4 @@ async def _serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerC raise OSError(f'gRPC could not bind {_bind_target(host, port)}') await server.start() logger.info(f'gRPC sessions on {host}:{bound}') - return server + return server, bound diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index d3f9cd472..4a7e6fe61 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -19,7 +19,7 @@ from starlette.datastructures import QueryParams from positronic.offboard import keys as offboard_keys -from positronic.policy import Policy, Recorder +from positronic.policy import Policy, Recorder, Session from positronic.policy.base import Layer from positronic.policy.executor import blocking from positronic.policy.spec import ModelSource, Pipeline, split @@ -307,6 +307,29 @@ async def grpc_session(self, conn: grpc_wire.GrpcServerConnection): """Serves one gRPC session, on the model the session path names.""" await self._serve_session(conn, grpc_wire.model_id_of(conn.session_path)) + async def _answer_observations(self, conn: wire.ServerConnection, session: Session) -> None: + """Answer every observation the client sends, until it disconnects.""" + while True: + message = await conn.receive() + self._last_activity = time.monotonic() + try: + raw_obs = deserialise(message) + # Plain acquire, not the keepalive helper: the client is awaiting a ``result`` and would + # mis-parse a ``waiting`` message. Its ``infer_timeout`` bounds the wait. + async with self._infer_lock: + try: + # The server's clock is not the rig's. + actions = await asyncio.to_thread(session, raw_obs, time.time_ns()) + except asyncio.CancelledError: + # Cancelling this await does not stop the worker, so the session close runs beside + # a live inference. Logged to give a later wrong answer a cause. + logger.error('Cancelled mid-inference: the worker is still in the backend') + raise + await conn.send(serialise({protocol.RESULT: actions})) + except Exception as e: + logger.error(f'Error processing message: {e}', exc_info=True) + await conn.send(serialise({protocol.ERROR: str(e)})) + async def _serve_session(self, conn: wire.ServerConnection, model_id: str | None): logger.info(f'Connected to {conn.peer} requesting {model_id or "default"}') @@ -355,20 +378,7 @@ async def _serve_session(self, conn: wire.ServerConnection, model_id: str | None await conn.send(serialise({protocol.STATUS: protocol.ServerStatus.READY, protocol.META: meta})) try: - while True: - message = await conn.receive() - self._last_activity = time.monotonic() - try: - raw_obs = deserialise(message) - # Plain acquire, not the keepalive helper: the client is awaiting a ``result`` and - # would mis-parse a ``waiting`` message. Its ``infer_timeout`` bounds the wait. - async with self._infer_lock: - # The server's clock is not the rig's. - actions = await asyncio.to_thread(session, raw_obs, time.time_ns()) - await conn.send(serialise({protocol.RESULT: actions})) - except Exception as e: - logger.error(f'Error processing message: {e}', exc_info=True) - await conn.send(serialise({protocol.ERROR: str(e)})) + await self._answer_observations(conn, session) except wire.PeerDisconnected: logger.info('Client disconnected') @@ -413,13 +423,17 @@ async def _idle_watchdog(self, server: uvicorn.Server): return async def _start_grpc(self) -> grpc.aio.Server: - """Start the gRPC wire on ``grpc_port``, sharing this server's model slot and inference lock.""" + """Start the gRPC wire on ``grpc_port``, sharing this server's model slot and inference lock. + + ``grpc_port`` becomes the port actually bound, so a server asked for any free one names it. + """ assert self.grpc_port is not None def authorized(headers: Mapping[str, str]) -> bool: return self._authorized(headers.get(AUTH_HEADER.lower())) - return await grpc_wire.serve(self.grpc_session, authorized, self.host, self.grpc_port) + server, self.grpc_port = await grpc_wire.serve(self.grpc_session, authorized, self.host, self.grpc_port) + return server def serve(self): async def _run(): diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index b02c9c5d9..7454694f3 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -31,12 +31,6 @@ def _bind_free_socket(host: str) -> socket.socket: return sock -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(('', 0)) - return s.getsockname()[1] - - StartServer = Callable[..., tuple[str, int, PolicyServer]] @@ -44,16 +38,16 @@ def _find_free_port() -> int: def start_server() -> Generator[StartServer, None, None]: """Factory serving pipelines on daemon threads; every started server is stopped and joined at teardown. - ``grpc=True`` also serves the gRPC wire, on a port of its own that ``PolicyServer.grpc_port`` names. + ``grpc=True`` also serves the gRPC wire, on a free port of its own that ``PolicyServer.grpc_port`` + names once the server has bound it. """ running: list[tuple[uvicorn.Server, threading.Thread]] = [] def start(pipeline, *, grpc: bool = False, **server_kwargs) -> tuple[str, int, PolicyServer]: - grpc_port = _find_free_port() if grpc else None host = server_kwargs.pop('host', 'localhost') ws_socket = _bind_free_socket(host) server = PolicyServer( - pipeline, host=host, port=ws_socket.getsockname()[1], grpc_port=grpc_port, **server_kwargs + pipeline, host=host, port=ws_socket.getsockname()[1], grpc_port=0 if grpc else None, **server_kwargs ) uv_server = uvicorn.Server( uvicorn.Config( @@ -64,7 +58,7 @@ def start(pipeline, *, grpc: bool = False, **server_kwargs) -> tuple[str, int, P async def _run(): await server._startup() # Started first, so the websocket port answering means both wires are up. - grpc_server = await server._start_grpc() if grpc_port is not None else None + grpc_server = await server._start_grpc() if server.grpc_port is not None else None try: await uv_server.serve(sockets=[ws_socket]) finally: diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index a7977cbae..b2ac3fa46 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -355,6 +355,17 @@ def test_a_port_that_never_answers_is_named_at_the_deadline(): client.new_session() +def test_an_open_timeout_under_the_probe_budget_still_opens(both_wires): + """The refusal probe takes a share of the budget, so a healthy server answers a short one.""" + server, _policy = both_wires + budget = grpc_wire._REFUSAL_PROBE_SEC / 2 + session = InferenceClient(grpc_url(server), open_timeout=budget, connect_deadline=0.0).new_session() + try: + assert session.infer({'image': 'test'}) == [{'action': [1, 2, 3]}] + finally: + session.close() + + def test_an_ipv6_host_binds_in_brackets(start_server: StartServer, make_mock_policy): """gRPC's target syntax brackets an IPv6 literal, so a bare '::1' would bind ':::' and fail.""" assert grpc_wire._bind_target('::', 9000) == '[::]:9000' diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index d5dda95b5..c666e89cd 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -146,10 +146,6 @@ def new_session(self, context=None, rt=None) -> RemoteSession: def functions(self) -> cabc.Mapping[str, cabc.Callable[..., Any]]: return {INFER: round_trip} - def close(self): - # Sessions own the connections; the client itself holds only where to open one. - pass - class RemotePolicy(Policy): """Policy running against a remote inference server, owning the stack in front of the connection. From 569d4dc63bb6c5d341ffea44bb7e2f2b7c8895eb Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Thu, 10 Sep 2026 22:09:47 +0000 Subject: [PATCH 27/46] Say the port rule without listing the TLS schemes The `InferenceClient` section named `https` and `wss`, so adding `grpcs` left it wrong: the default is 443 whenever the scheme is TLS. The rule cannot go stale as schemes arrive; the wire table above it names which they are. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index ee28fbbd1..6eff26c2c 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -267,7 +267,7 @@ The CLI entry point every vendor server exposes. A vendor binds `pipeline` to ea ### `client.InferenceClient` A Python client for connecting to an inference server. One URL addresses it, in the same forms -`RemotePolicy` accepts: an omitted port is the scheme's own, 443 for `https`/`wss` and 80 otherwise. The URL +`RemotePolicy` accepts: an omitted port is the scheme's own, 443 for a TLS scheme and 80 otherwise. The URL fixes the wire, the model and the session params, so serving another model means another client. ```python From eb26d6272660d6ae74ec6580f95577765e618998 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Thu, 10 Sep 2026 22:19:27 +0000 Subject: [PATCH 28/46] Say which port the serve script fixes and which the caller names The header claimed gRPC on 9000, while `--grpc_port` decides it and 9000 is only the default. "Container port" was the same claim one line down; the websocket port is the one this script fixes. Ticket: Positronic-Robotics/internal#1191 #refs --- workflows/nebius/serve.sh | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 046bdb52a..79c0f9f8d 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -6,9 +6,9 @@ # itself takes ~10-15 min more to finish uv sync and load the model into GPU # memory after the URL appears. # -# Both wires are served: the websocket on port 8000 and gRPC on port 9000. The -# gRPC port is declared as an ordinary HTTP port and never `/tcp` — the offboard -# README says what each front does to a gRPC session. +# Both wires are served: the websocket on 8000, and gRPC on whatever `--grpc_port` +# names. The gRPC port is declared as an ordinary HTTP port and never `/tcp` — the +# offboard README says what each front does to a gRPC session. # # That URL carries the id of a tunnel created with the endpoint, so it cannot be # chosen or known in advance, and a delete plus re-create earns a new one even @@ -22,7 +22,7 @@ # ingress mode strips the WebSocket upgrade headers and so cannot pass inference # sessions at all. # -# Hardcoded: GPU platform, container port. Vendor selects image + uv extra. One +# Hardcoded: GPU platform, websocket port. Vendor selects image + uv extra. One # setting of its own, via env: NEBIUS_PRESET. Everything shared with the other # scripts here lives in common.sh. From 0b9ab74e0bb2227963b5c003e18dd699c3665991 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Thu, 10 Sep 2026 22:23:18 +0000 Subject: [PATCH 29/46] List the gRPC schemes where the eval docs name what a policy URL takes The `--policy.url` section named `http`, `ws` and `wss` and stopped there, so the wire this branch adds was absent from the one doc a user reads to reach a server. The offboard README and `InferenceClient` already list `grpc(s)`; this was the third copy and the only stale one. Ticket: Positronic-Robotics/internal#1191 #refs --- docs/inference.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/inference.md b/docs/inference.md index 71cc4e9f6..21525a06a 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -56,7 +56,7 @@ uv run positronic eval run --eval=.sim.positronic.stack_cubes \ --policy.url='https://gpu-server/api/v1/session/checkpoint-20000?codec.fps=10&local.pad_start=false' ``` -Accepted forms: `host`, `host:port`, and `https://host[:port][/api/v1/session[/]]` (`http`, `ws` and `wss` work too), each with an optional query. `https`/`wss` enable TLS. An omitted port is the scheme's own — 443 for TLS and 80 otherwise — so name the port a server listens on (`:8000` for every vendor server's default). Naming no model id serves the checkpoint the server pinned at startup. +Accepted forms: `host`, `host:port`, and `scheme://host[:port][/api/v1/session[/]]`, each with an optional query. The scheme settles the wire and the TLS: `http`/`https` and `ws`/`wss` take the websocket wire, `grpc`/`grpcs` the gRPC one, and the `s` forms are the TLS ones. An omitted port is the scheme's own — 443 for TLS and 80 otherwise — so name the port a server listens on (`:8000` for every vendor server's websocket default). Naming no model id serves the checkpoint the server pinned at startup. **Credentials stay out of the URL, and out of the command line.** The URL is meant to be safe to paste around, so a token rides a header instead. It stays off the command line too: `save_run_metadata()` writes `sys.argv` beside the run's episodes. Three policy configs build the header: From c84de5cd9155b0517a1b44e7f0ee17b10ca7ed90 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Thu, 10 Sep 2026 22:34:07 +0000 Subject: [PATCH 30/46] Name the session protocol rather than one of its wires in the eval docs The guide opened by calling remote inference WebSocket-only and labelled the offboard README "WebSocket protocol", so its overview contradicted the URL schemes further down. One protocol, two wires, said once in each place. Ticket: Positronic-Robotics/internal#1191 #refs --- docs/inference.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/inference.md b/docs/inference.md index 21525a06a..70cd2ec80 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -1,10 +1,10 @@ # Inference Guide -Deploy trained policies for evaluation and production use. Positronic supports local inference (model loaded on robot/simulator machine) and inference with remote server (model runs on separate GPU server via WebSocket). +Deploy trained policies for evaluation and production use. Positronic supports local inference (model loaded on robot/simulator machine) and inference with remote server (model runs on a separate GPU server, over a websocket or gRPC). ## Inference with Remote Server -Positronic's unified WebSocket protocol connects any hardware to any model (LeRobot, GR00T, OpenPI). The key benefit is running heavy models on powerful GPU hardware (OpenPI needs ~62GB, GR00T ~8GB) separate from the robot/simulator machine. +Positronic's unified session protocol connects any hardware to any model (LeRobot, GR00T, OpenPI); the same frames cross either wire, a websocket or gRPC. The key benefit is running heavy models on powerful GPU hardware (OpenPI needs ~62GB, GR00T ~8GB) separate from the robot/simulator machine. Each server carries a full **policy pipeline** — one chain naming the rig-side stack, the `remote` split marker, the server-side codec, and the model source that loads checkpoints (see `positronic.policy.spec`). The server runs the half right of the marker and declares the half left of it in its handshake; the client builds the declared stack automatically. Vendors ship their pipelines by name, and every name is a server subcommand — `groot-server ee_rot6d_joints` launches that one. The available names are listed in each vendor's README. @@ -119,5 +119,5 @@ Run inference with recording, review in Positronic server, score manually (succe - [Training Workflow](training-workflow.md) – Preparing data and training - [Codecs Guide](codecs.md) – Observation/action encoding -- [Offboard README](../positronic/offboard/README.md) – WebSocket protocol +- [Offboard README](../positronic/offboard/README.md) – the session protocol and both wires - Vendor guides: [OpenPI](../positronic/vendors/openpi/README.md) | [GR00T](../positronic/vendors/gr00t/README.md) | [SmolVLA](../positronic/vendors/lerobot/README.md) | [LeRobot ACT](../positronic/vendors/lerobot_0_3_3/README.md) From a45453fa71e4d34d5fa48c10acd5b11504001769 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Fri, 11 Sep 2026 10:31:57 +0000 Subject: [PATCH 31/46] Take the wires sessions arrive on, so the server names none of them `PolicyServer` held each wire's bind, its accept-time auth and its model-id extraction: a FastAPI app and two routes for the websocket, `grpc_session` and `_start_grpc` for gRPC. `wire.ServerConnection` already covered an open session, so the inference loop was transport-free while opening a session was not, and a third wire cost five edits to the server. `wire.Wire` covers the other half: bind, accept, refuse and stop. Each wire reads its own route for the model a session names and checks its own session headers. `serve` takes the list of them, and the CLI entry point is where the two wires are named. The handshake meta names the port that carried the session. A gRPC session reported the websocket's port before. The test fixture serves through `PolicyServer.serve` rather than rebuilding it, which removed a second copy of the wire lifecycle and gave the idle watchdog its first test. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 12 +- positronic/offboard/grpc_wire.py | 100 +++++++----- positronic/offboard/server.py | 145 +++++++++-------- positronic/offboard/tests/conftest.py | 87 ++++------ positronic/offboard/tests/test_grpc_wire.py | 113 +++++++------ .../offboard/tests/test_remote_policy.py | 2 +- positronic/offboard/tests/test_server.py | 38 +++-- positronic/offboard/wire.py | 150 +++++++++++++++++- 8 files changed, 411 insertions(+), 236 deletions(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index 6eff26c2c..a2527d2b7 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -253,14 +253,22 @@ The one server implementation behind every vendor. It serves a **policy pipeline ```python from positronic.offboard import PolicyServer +from positronic.offboard.wire import WebsocketWire from positronic.policy.spec import PolicySource, remote from positronic.policy.layers import ChunkedSchedule pipeline = ChunkedSchedule() | remote | PolicySource(my_policy) -PolicyServer(pipeline, host='0.0.0.0', port=8000).serve() +server = PolicyServer(pipeline) +server.serve([WebsocketWire('0.0.0.0', 8000, server.api)]) ``` -`PolicySource` serves one ready in-process policy; vendors instead define a `ModelSource` over a checkpoint directory. Passing a `cfn.Config` that builds the pipeline — as the vendor servers do with their named pipelines — enables [session parameters](#session-parameters); an instantiated pipeline serves exactly as launched. `recording_dir` enables the per-session recording taps described above, `grpc_port` adds the gRPC wire, and `idle_timeout_min` shuts the server down after that many minutes without activity. +`serve` takes the wires sessions arrive on, and the server names none of them: each wire binds its own +port, reads its own route for the model a session asks for, and checks its own session headers. Add +`grpc_wire.GrpcWire(host, port)` to the list to serve gRPC beside the websocket. A wire that speaks HTTP +takes `server.api`, the model catalogue, and answers it on the same port it carries sessions on. A wire +asked for port 0 binds any free one and names it in `wire.endpoint`. + +`PolicySource` serves one ready in-process policy; vendors instead define a `ModelSource` over a checkpoint directory. Passing a `cfn.Config` that builds the pipeline — as the vendor servers do with their named pipelines — enables [session parameters](#session-parameters); an instantiated pipeline serves exactly as launched. `recording_dir` enables the per-session recording taps described above, and `idle_timeout_min` ends the server after that many minutes without activity. ### `server.serve` The CLI entry point every vendor server exposes. A vendor binds `pipeline` to each of its named pipelines and lists the results as subcommands, so `-server ` launches one. Only `--host`, `--port`, `--grpc_port`, `--recording_dir` and `--idle_timeout_min` are flags of `serve` itself; everything the served model is — codec, source, checkpoint directory — is reached through the pipeline (`--pipeline.source.checkpoints_dir=...`), which is also where a deployment preset binds it. diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 7ba6d9c3d..fb4942bd5 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -9,7 +9,7 @@ import threading import time import urllib.parse -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from collections.abc import AsyncIterator, Mapping import grpc import grpc.aio @@ -217,15 +217,26 @@ def model_id_of(session_path: str) -> str | None: class GrpcServerConnection(wire.ServerConnection): """A server's end of one gRPC session.""" - def __init__(self, requests: AsyncIterator[bytes], context: grpc.aio.ServicerContext, headers: Mapping[str, str]): + def __init__( + self, + requests: AsyncIterator[bytes], + context: grpc.aio.ServicerContext, + headers: Mapping[str, str], + endpoint: wire.Endpoint, + ): self._requests = requests self._context = context self._headers = headers + self._endpoint = endpoint @property def peer(self) -> str: return self._context.peer() + @property + def endpoint(self) -> wire.Endpoint: + return self._endpoint + @property def session_path(self) -> str: return self._headers.get(SESSION_PATH_HEADER, wire.SESSION_PATH) @@ -267,42 +278,55 @@ def _server_options() -> list[tuple[str, int]]: ] -async def serve( - serve_session: Callable[[GrpcServerConnection], Awaitable[None]], - authorized: Callable[[Mapping[str, str]], bool], - host: str, - port: int, -) -> tuple[grpc.aio.Server, int]: - """Start a gRPC server that gives every accepted session to ``serve_session``, and say what it bound. - - A ``port`` of 0 binds any free one, which is the port that comes back. +class GrpcWire(wire.Wire): + """The gRPC wire: sessions on a port of their own, one bidirectional stream each. - ``authorized`` reads the session headers and refuses before the session opens, as the websocket - wire refuses the upgrade. - - The port is plaintext; a TLS edge in front of it serves an authenticated endpoint. + A ``port`` of 0 binds any free one. The port is plaintext; a TLS edge in front of it serves an + authenticated endpoint. """ - async def _serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerContext) -> None: - headers = _headers(context) - if not authorized(headers): - await context.abort(grpc.StatusCode.PERMISSION_DENIED, 'Invalid or missing bearer token') - try: - await serve_session(GrpcServerConnection(requests, context, headers)) - except Exception as e: - # The session itself reports what it can over the stream; anything reaching here happened - # before or beyond that, so the client learns of it from the status alone. - logger.error(f'Failed gRPC session: {e}', exc_info=True) - await context.abort(grpc.StatusCode.INTERNAL, str(e)) - - handler = grpc.stream_stream_rpc_method_handler(_serve_one, request_deserializer=None, response_serializer=None) - server = grpc.aio.server(options=_server_options()) - server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler(SERVICE, {METHOD: handler}),)) - bound = server.add_insecure_port(_bind_target(host, port)) - if bound == 0: - # gRPC reports a refused bind by returning port 0, so a server left to start here would - # accept nothing and say nothing. - raise OSError(f'gRPC could not bind {_bind_target(host, port)}') - await server.start() - logger.info(f'gRPC sessions on {host}:{bound}') - return server, bound + def __init__(self, host: str, port: int): + self._host = host + self._port = port + self._server: grpc.aio.Server | None = None + self._endpoint: wire.Endpoint | None = None + + @property + def endpoint(self) -> wire.Endpoint: + assert self._endpoint is not None, 'The gRPC wire has not started' + return self._endpoint + + async def start(self, session: wire.SessionHandler, authorized: wire.Authorized) -> None: + async def serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerContext) -> None: + headers = _headers(context) + if not authorized(headers): + await context.abort(grpc.StatusCode.PERMISSION_DENIED, 'Invalid or missing bearer token') + conn = GrpcServerConnection(requests, context, headers, self.endpoint) + try: + await session(conn, model_id_of(conn.session_path)) + except Exception as e: + # The session itself reports what it can over the stream; anything reaching here happened + # before or beyond that, so the client learns of it from the status alone. + logger.error(f'Failed gRPC session: {e}', exc_info=True) + await context.abort(grpc.StatusCode.INTERNAL, str(e)) + + handler = grpc.stream_stream_rpc_method_handler(serve_one, request_deserializer=None, response_serializer=None) + server = grpc.aio.server(options=_server_options()) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler(SERVICE, {METHOD: handler}),)) + bound = server.add_insecure_port(_bind_target(self._host, self._port)) + if bound == 0: + # gRPC reports a refused bind by returning port 0, so a server left to start here would + # accept nothing and say nothing. + raise OSError(f'gRPC could not bind {_bind_target(self._host, self._port)}') + self._server = server + self._endpoint = wire.Endpoint(self._host, bound) + await server.start() + logger.info(f'gRPC sessions on {self._host}:{bound}') + + async def serve(self) -> None: + assert self._server is not None, 'The gRPC wire has not started' + await self._server.wait_for_termination() + + async def stop(self) -> None: + if self._server is not None: + await self._server.stop(grace=None) diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 4a7e6fe61..246d08358 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -7,15 +7,13 @@ import os import time from collections import Counter -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from importlib.metadata import version as _pkg_version from typing import Any import configuronic as cfn -import grpc.aio import pos3 -import uvicorn -from fastapi import Depends, FastAPI, Header, HTTPException, WebSocket, WebSocketException, status +from fastapi import APIRouter, Depends, Header, HTTPException from starlette.datastructures import QueryParams from positronic.offboard import keys as offboard_keys @@ -193,9 +191,10 @@ class PolicyServer: The session flow is: accept → session params → resolve → load via manager → remote-half wrap → reset → inference loop - A session runs over one of two wires (see ``positronic.offboard.wire``): the websocket, on ``port``, - and gRPC, on ``grpc_port``. Both carry the same frames, so the flow above is the same on each. The - HTTP routes stay on ``port``; a ``grpc_port`` of ``None`` serves the websocket alone. + ``serve`` takes the wires sessions arrive on (see ``positronic.offboard.wire``). Each one reads its + own route for the model a session names and checks its own session headers, so the flow above is the + same over every wire and this server names none of them. ``api`` holds the server's own HTTP routes, + which a wire that speaks HTTP serves beside its sessions. On startup (before accepting connections): resolve(None) → load. @@ -207,12 +206,9 @@ class PolicyServer: def __init__( self, pipeline: cfn.Config | Pipeline, - host: str = '0.0.0.0', - port: int = 8000, recording_dir: str | None = None, idle_timeout_min: float | None = None, auth_token: str | None = None, - grpc_port: int | None = None, ): self._pipeline_cfg = pipeline if isinstance(pipeline, cfn.Config) else None self._pipeline = pipeline.instantiate() if isinstance(pipeline, cfn.Config) else pipeline @@ -225,10 +221,6 @@ def __init__( _declared_stack(local) self._source = self._pipeline.source self._manager = PolicyManager(self._source) - self.host = host - self.port = port - self.grpc_port = grpc_port - self.metadata: dict[str, Any] = {offboard_keys.HOST: host, offboard_keys.PORT: port} # Synced once; each session builds its own ``Recorder`` so concurrent streams never mix. self._recording_dir = pos3.sync(recording_dir) if recording_dir else None @@ -240,6 +232,9 @@ def __init__( self._infer_lock = asyncio.Lock() self._default_id: str | None = None + # Set while ``serve`` runs, so ``shutdown`` can reach its loop from another thread. + self._loop: asyncio.AbstractEventLoop | None = None + self._stop: asyncio.Event | None = None # ``None`` serves open, so a broken secret must not reach that path by accident. Empty would read # as open; anything an ``Authorization`` header cannot carry — a newline off the end of a file, a @@ -248,15 +243,15 @@ def __init__( raise ValueError('auth_token must be non-empty printable ASCII without spaces; pass None to serve open') self._auth_token = auth_token - self.app = FastAPI() - http_auth, ws_auth = [Depends(self._require_http_auth)], [Depends(self._require_ws_auth)] - self.app.get('/api/v1/models', dependencies=http_auth)(self.get_models) - self.app.websocket(wire.SESSION_PATH, dependencies=ws_auth)(self.default_session) - # ``:path`` so an id that is itself a path (a HuggingFace repo, say) opens under the name - # ``/api/v1/models`` advertises. - self.app.websocket(f'{wire.SESSION_PATH}/{{model_id:path}}', dependencies=ws_auth)(self.model_session) + self._api = APIRouter() + self._api.get('/api/v1/models', dependencies=[Depends(self._require_http_auth)])(self.get_models) - def _authorized(self, authorization: str | None) -> bool: + @property + def api(self) -> APIRouter: + """The server's own HTTP routes: the model catalogue a client reads before it opens a session.""" + return self._api + + def _token_matches(self, authorization: str | None) -> bool: if self._auth_token is None: return True if authorization is None: @@ -265,15 +260,14 @@ def _authorized(self, authorization: str | None) -> bool: # a non-ASCII ``str``, which would answer a malformed header with a 500 instead of a refusal. return hmac.compare_digest(authorization.encode(), bearer(self._auth_token).encode()) + def _authorized(self, headers: Mapping[str, str]) -> bool: + """Whether session headers carry the bearer token this server gates on. Every wire asks this.""" + return self._token_matches(headers.get(AUTH_HEADER.lower())) + def _require_http_auth(self, authorization: str | None = Header(default=None, alias=AUTH_HEADER)) -> None: - if not self._authorized(authorization): + if not self._token_matches(authorization): raise HTTPException(status_code=401, detail='Invalid or missing bearer token') - async def _require_ws_auth(self, websocket: WebSocket) -> None: - """Rejects before ``accept()``, so an unauthorized peer never reaches the session handshake.""" - if not self._authorized(websocket.headers.get(AUTH_HEADER)): - raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION) - async def get_models(self) -> dict: return {'models': self._source.get_models()} @@ -293,20 +287,6 @@ def _session_pipeline(self, params: dict[str, Any]) -> Pipeline: raise ValueError('Session params must not change the model source; it is fixed at launch') return pipeline - async def default_session(self, websocket: WebSocket): - """Serves the model pinned at startup. Naming a model is the path's job, so every query param here - is a pipeline override.""" - await websocket.accept() - await self._serve_session(wire.WebsocketServerConnection(websocket), None) - - async def model_session(self, websocket: WebSocket, model_id: str): - await websocket.accept() - await self._serve_session(wire.WebsocketServerConnection(websocket), model_id) - - async def grpc_session(self, conn: grpc_wire.GrpcServerConnection): - """Serves one gRPC session, on the model the session path names.""" - await self._serve_session(conn, grpc_wire.model_id_of(conn.session_path)) - async def _answer_observations(self, conn: wire.ServerConnection, session: Session) -> None: """Answer every observation the client sends, until it disconnects.""" while True: @@ -366,8 +346,10 @@ async def _serve_session(self, conn: wire.ServerConnection, model_id: str | None self._infer_lock.release() assert session is not None # Later entries win: per-episode session facts over static ones, the server's own last. + endpoint = conn.endpoint meta = { - **self.metadata, + offboard_keys.HOST: endpoint.host, + offboard_keys.PORT: endpoint.port, **self._source.meta(rid), offboard_keys.CHECKPOINT_ID: rid, **session.meta, @@ -408,60 +390,71 @@ async def _startup(self): logger.info(f'Pinned default checkpoint at startup: {self._default_id}') await self._manager.get_policy(self._default_id) - async def _idle_watchdog(self, server: uvicorn.Server): + async def _idle_watchdog(self): + """Return once no session has touched the server for ``idle_timeout_min``.""" assert self.idle_timeout_min is not None timeout_s = self.idle_timeout_min * 60 poll = min(timeout_s, 30) - while not server.should_exit: + while True: await asyncio.sleep(poll) if self._active_sessions > 0: continue idle = time.monotonic() - self._last_activity if idle >= timeout_s: logger.warning(f'No activity for {idle:.0f}s (idle timeout {timeout_s:.0f}s); shutting down server') - server.should_exit = True return - async def _start_grpc(self) -> grpc.aio.Server: - """Start the gRPC wire on ``grpc_port``, sharing this server's model slot and inference lock. + def serve(self, wires: Sequence[wire.Wire], on_ready: Callable[[], None] | None = None): + """Serve sessions on every wire in ``wires``, until one of them ends or the server goes idle. - ``grpc_port`` becomes the port actually bound, so a server asked for any free one names it. - """ - assert self.grpc_port is not None - - def authorized(headers: Mapping[str, str]) -> bool: - return self._authorized(headers.get(AUTH_HEADER.lower())) + Every wire shares this server's model slot and inference lock, so a session is served the same + whichever one carried it. - server, self.grpc_port = await grpc_wire.serve(self.grpc_session, authorized, self.host, self.grpc_port) - return server + ``on_ready`` runs on the server's own loop once every wire has bound, which is where a caller + that asked for port 0 reads back the port each wire took. + """ - def serve(self): async def _run(): + self._loop, self._stop = asyncio.get_running_loop(), asyncio.Event() await self._startup() - config = uvicorn.Config( - self.app, host=self.host, port=self.port, log_level='info', ws_max_size=wire.MAX_MESSAGE_BYTES - ) - server = uvicorn.Server(config) + for w in wires: + await w.start(self._serve_session, self._authorized) self._last_activity = time.monotonic() - watchdog = None - grpc_server = await self._start_grpc() if self.grpc_port is not None else None + if on_ready is not None: + on_ready() + serving = [asyncio.create_task(w.serve()) for w in wires] + # What ends the server, beside a wire ending on its own: a caller's ``shutdown``, and the + # idle timeout. + ending: list[asyncio.Task] = [asyncio.create_task(self._stop.wait())] if self.idle_timeout_min and self.idle_timeout_min > 0: - watchdog = asyncio.create_task(self._idle_watchdog(server)) + ending.append(asyncio.create_task(self._idle_watchdog())) try: - await server.serve() + done, _still_running = await asyncio.wait(serving + ending, return_when=asyncio.FIRST_COMPLETED) + # A wire that ended on an error raises here, rather than reading as the shutdown this waits for. + for task in done: + task.result() finally: - if watchdog is not None: - watchdog.cancel() - if grpc_server is not None: - await grpc_server.stop(grace=None) + for task in ending: + task.cancel() + for w in wires: + await w.stop() + # Each wire ends the sessions it carries before this returns and the model slot closes. + await asyncio.gather(*serving, return_exceptions=True) try: asyncio.run(_run()) except KeyboardInterrupt: logger.info('Server stopped by user') finally: + self._loop, self._stop = None, None self._manager.close() + def shutdown(self): + """Ask a running ``serve`` to end, from any thread. A server that is not serving ignores it.""" + loop, stop = self._loop, self._stop + if loop is not None and stop is not None: + loop.call_soon_threadsafe(stop.set) + @cfn.config(host='0.0.0.0', port=8000, recording_dir=None, idle_timeout_min=None, grpc_port=None) def serve( @@ -482,13 +475,17 @@ def serve( The bearer token gating the server comes from ``AUTH_TOKEN_ENV`` rather than a flag, which would put a secret in the process arguments; unset serves open. + + This is where the flags name the wires, and the only place either wire is named: ``PolicyServer`` + takes whatever list it is handed. """ - PolicyServer( + server = PolicyServer( pipeline, - host=host, - port=port, recording_dir=recording_dir, idle_timeout_min=idle_timeout_min, auth_token=os.environ.get(AUTH_TOKEN_ENV), - grpc_port=grpc_port, - ).serve() + ) + wires: list[wire.Wire] = [wire.WebsocketWire(host, port, server.api)] + if grpc_port is not None: + wires.append(grpc_wire.GrpcWire(host, grpc_port)) + server.serve(wires) diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 7454694f3..2692c267a 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -1,14 +1,11 @@ -import asyncio -import socket import threading -import time from collections.abc import Callable, Generator, Mapping +from typing import NamedTuple from unittest.mock import MagicMock import pytest -import uvicorn -from positronic.offboard import wire +from positronic.offboard import grpc_wire, wire from positronic.offboard.server import PolicyServer from positronic.policy import Policy, Session from positronic.policy.executor import Executor @@ -16,72 +13,46 @@ from positronic.policy.spec import ModelSource, PolicySource, remote -def _bind_free_socket(host: str) -> socket.socket: - """A socket holding a free port on ``host``, to hand to the server that will serve on it. +class Served(NamedTuple): + """A running server, and the ports its wires took.""" - Drawing a port and closing the socket loses the port to whoever binds next, and these tests run in - parallel. Staying bound from the draw to the serve is what makes the port ours. The family comes from - ``host`` itself, so an IPv6 test binds an IPv6 socket. - """ - family, _type, _proto, _canon, address = socket.getaddrinfo( - host, 0, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE - )[0] - sock = socket.socket(family, socket.SOCK_STREAM) - sock.bind(address) - return sock + host: str + port: int + server: PolicyServer + grpc_port: int | None -StartServer = Callable[..., tuple[str, int, PolicyServer]] +StartServer = Callable[..., Served] @pytest.fixture def start_server() -> Generator[StartServer, None, None]: """Factory serving pipelines on daemon threads; every started server is stopped and joined at teardown. - ``grpc=True`` also serves the gRPC wire, on a free port of its own that ``PolicyServer.grpc_port`` - names once the server has bound it. + Each wire asks for port 0 and holds what it binds, so servers started in parallel never draw the same + port. ``grpc=True`` serves the gRPC wire beside the websocket one. """ - running: list[tuple[uvicorn.Server, threading.Thread]] = [] + running: list[tuple[PolicyServer, threading.Thread]] = [] - def start(pipeline, *, grpc: bool = False, **server_kwargs) -> tuple[str, int, PolicyServer]: + def start(pipeline, *, grpc: bool = False, **server_kwargs) -> Served: host = server_kwargs.pop('host', 'localhost') - ws_socket = _bind_free_socket(host) - server = PolicyServer( - pipeline, host=host, port=ws_socket.getsockname()[1], grpc_port=0 if grpc else None, **server_kwargs - ) - uv_server = uvicorn.Server( - uvicorn.Config( - server.app, host=server.host, port=server.port, log_level='warning', ws_max_size=wire.MAX_MESSAGE_BYTES - ) - ) - - async def _run(): - await server._startup() - # Started first, so the websocket port answering means both wires are up. - grpc_server = await server._start_grpc() if server.grpc_port is not None else None - try: - await uv_server.serve(sockets=[ws_socket]) - finally: - if grpc_server is not None: - await grpc_server.stop(grace=None) - - thread = threading.Thread(target=asyncio.run, args=(_run(),), daemon=True) + server = PolicyServer(pipeline, **server_kwargs) + wires: list[wire.Wire] = [wire.WebsocketWire(host, 0, server.api)] + if grpc: + wires.append(grpc_wire.GrpcWire(host, 0)) + ready = threading.Event() + thread = threading.Thread(target=server.serve, args=(wires, ready.set), daemon=True) thread.start() - running.append((uv_server, thread)) - - deadline = time.time() + 5.0 - while time.time() < deadline: - try: - with socket.create_connection((server.host, server.port), timeout=0.1): - return server.host, server.port, server - except (ConnectionRefusedError, OSError): - time.sleep(0.05) - raise RuntimeError('Server failed to start') + running.append((server, thread)) + if not ready.wait(timeout=10.0): + raise RuntimeError('Server failed to start') + return Served(host, wires[0].endpoint.port, server, wires[1].endpoint.port if grpc else None) yield start - for uv_server, thread in running: - uv_server.should_exit = True - thread.join(timeout=5.0) + for server, thread in running: + server.shutdown() + thread.join(timeout=10.0) + assert not thread.is_alive(), 'the server did not stop when asked' @pytest.fixture @@ -173,7 +144,7 @@ def inference_server(start_server: StartServer, mock_policy: MagicMock) -> tuple Returns: tuple[str, int]: (host, port) """ - host, port, _server = start_server(ChunkedSchedule() | remote | PolicySource(mock_policy)) + host, port, *_ = start_server(ChunkedSchedule() | remote | PolicySource(mock_policy)) return host, port @@ -181,5 +152,5 @@ def inference_server(start_server: StartServer, mock_policy: MagicMock) -> tuple def multi_policy_server( start_server: StartServer, mock_policy_registry: dict[str, MagicMock] ) -> tuple[str, int, dict[str, MagicMock]]: - host, port, _server = start_server(ChunkedSchedule() | remote | DictSource(mock_policy_registry)) + host, port, *_ = start_server(ChunkedSchedule() | remote | DictSource(mock_policy_registry)) return host, port, mock_policy_registry diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index b2ac3fa46..5b9347a4e 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -25,8 +25,8 @@ from positronic.offboard import grpc_wire, wire from positronic.offboard import keys as offboard_keys from positronic.offboard.client import InferenceClient, _ConnectRetries -from positronic.offboard.server import AUTH_HEADER, PolicyServer, bearer -from positronic.offboard.tests.conftest import DictSource, StartServer +from positronic.offboard.server import AUTH_HEADER, bearer +from positronic.offboard.tests.conftest import DictSource, Served, StartServer from positronic.policy.base import SEQ from positronic.policy.layers import ChunkedSchedule, TemporalStack from positronic.policy.spec import ModelSource, PolicySource, remote @@ -34,21 +34,21 @@ _TOKEN = 'test-secret-token' -def grpc_url(server: PolicyServer, path: str = '') -> str: - return f'grpc://{server.host}:{server.grpc_port}{path}' +def grpc_url(served: Served, path: str = '') -> str: + return f'grpc://{served.host}:{served.grpc_port}{path}' @pytest.fixture -def both_wires(start_server: StartServer, make_mock_policy) -> tuple[PolicyServer, MagicMock]: +def both_wires(start_server: StartServer, make_mock_policy) -> tuple[Served, MagicMock]: """A server offering both wires over one policy.""" policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - _host, _port, server = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True) - return server, policy + served = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True) + return served, policy def test_a_grpc_session_handshakes_and_infers(both_wires): - server, policy = both_wires - session = InferenceClient(grpc_url(server)).new_session() + served, policy = both_wires + session = InferenceClient(grpc_url(served)).new_session() try: assert session.metadata['model_name'] == 'stub' obs = {'image': 'test'} @@ -58,26 +58,43 @@ def test_a_grpc_session_handshakes_and_infers(both_wires): session.close() +def _apart_from_the_endpoint(meta: dict) -> dict: + return {key: value for key, value in meta.items() if key not in (offboard_keys.HOST, offboard_keys.PORT)} + + def test_both_wires_answer_one_observation_alike(both_wires): - server, _policy = both_wires + served, _policy = both_wires obs = {'image': 'test'} - over_ws = InferenceClient(f'{server.host}:{server.port}').new_session() - over_grpc = InferenceClient(grpc_url(server)).new_session() + over_ws = InferenceClient(f'{served.host}:{served.port}').new_session() + over_grpc = InferenceClient(grpc_url(served)).new_session() try: - assert over_grpc.metadata == over_ws.metadata + assert _apart_from_the_endpoint(over_grpc.metadata) == _apart_from_the_endpoint(over_ws.metadata) assert over_grpc.infer(obs) == over_ws.infer(obs) finally: over_ws.close() over_grpc.close() +def test_each_wire_names_its_own_port_in_the_meta(both_wires): + """The two wires bind two ports, and a session reads back the one that carried it.""" + served, _policy = both_wires + over_ws = InferenceClient(f'{served.host}:{served.port}').new_session() + over_grpc = InferenceClient(grpc_url(served)).new_session() + try: + assert over_ws.metadata[offboard_keys.PORT] == served.port + assert over_grpc.metadata[offboard_keys.PORT] == served.grpc_port + finally: + over_ws.close() + over_grpc.close() + + def test_both_wires_report_what_their_close_saw(both_wires, caplog): """A close the server answered has to read differently from one it never saw, whichever wire carried the session. The second leaves the server holding the slot, and the next session's handshake waits on it, so each wire reports the distinction in the terms its own protocol offers.""" - server, _policy = both_wires - over_ws = InferenceClient(f'{server.host}:{server.port}').new_session() - over_grpc = InferenceClient(grpc_url(server)).new_session() + served, _policy = both_wires + over_ws = InferenceClient(f'{served.host}:{served.port}').new_session() + over_grpc = InferenceClient(grpc_url(served)).new_session() with caplog.at_level(logging.INFO, logger='positronic.offboard.client'): over_ws.close() @@ -90,16 +107,16 @@ def test_both_wires_report_what_their_close_saw(both_wires, caplog): def test_closing_a_session_ends_it_on_the_server(both_wires): """``close`` half-closes the stream and waits, so the server releases the session before it returns.""" - server, _policy = both_wires - session = InferenceClient(grpc_url(server)).new_session() - assert server._active_sessions == 1 + served, _policy = both_wires + session = InferenceClient(grpc_url(served)).new_session() + assert served.server._active_sessions == 1 session.close() - assert server._active_sessions == 0 + assert served.server._active_sessions == 0 def test_a_failed_inference_reaches_the_client_as_an_exception(both_wires): - server, policy = both_wires - session = InferenceClient(grpc_url(server)).new_session() + served, policy = both_wires + session = InferenceClient(grpc_url(served)).new_session() try: policy._mock_session.side_effect = RuntimeError('no such joint') with pytest.raises(RuntimeError, match='no such joint'): @@ -111,9 +128,9 @@ def test_a_failed_inference_reaches_the_client_as_an_exception(both_wires): def test_a_session_that_cannot_open_reaches_the_client_as_an_exception(start_server, make_mock_policy): """A model the source refuses fails in the handshake, before the session serves anything.""" policies = {'alpha': make_mock_policy([{'action': [1]}], {'model_name': 'alpha'})} - _host, _port, server = start_server(ChunkedSchedule() | remote | DictSource(policies), grpc=True) + served = start_server(ChunkedSchedule() | remote | DictSource(policies), grpc=True) with pytest.raises(RuntimeError, match='Unknown model'): - InferenceClient(grpc_url(server, f'{wire.SESSION_PATH}/beta')).new_session() + InferenceClient(grpc_url(served, f'{wire.SESSION_PATH}/beta')).new_session() def test_the_session_path_names_the_model(start_server, make_mock_policy): @@ -121,8 +138,8 @@ def test_the_session_path_names_the_model(start_server, make_mock_policy): 'alpha': make_mock_policy([{'action': ['alpha']}], {'model_name': 'alpha'}), 'beta': make_mock_policy([{'action': ['beta']}], {'model_name': 'beta'}), } - _host, _port, server = start_server(ChunkedSchedule() | remote | DictSource(policies), grpc=True) - session = InferenceClient(grpc_url(server, f'{wire.SESSION_PATH}/beta')).new_session() + served = start_server(ChunkedSchedule() | remote | DictSource(policies), grpc=True) + session = InferenceClient(grpc_url(served, f'{wire.SESSION_PATH}/beta')).new_session() try: assert session.metadata['model_name'] == 'beta' assert session.infer({'obs': 'beta'}) == [{'action': ['beta']}] @@ -137,8 +154,8 @@ def _tunable_pipe(source: ModelSource, offsets: tuple[float, ...] = (-0.1, 0.0)) def test_the_query_carries_the_session_params(start_server, make_mock_policy): policies = {'alpha': make_mock_policy([{'action': ['alpha']}], {'model_name': 'alpha'})} pipe = cfn.Config(_tunable_pipe, source=cfn.Config(DictSource, policies=policies)) - _host, _port, server = start_server(pipe, grpc=True) - session = InferenceClient(grpc_url(server, f'{wire.SESSION_PATH}?offsets=[-0.5, 0.0]')).new_session() + served = start_server(pipe, grpc=True) + session = InferenceClient(grpc_url(served, f'{wire.SESSION_PATH}?offsets=[-0.5, 0.0]')).new_session() try: stack = session.metadata[offboard_keys.LOCAL_STACK][SEQ] # `args` and the layer's own constructor keyword are the spec grammar's, written wherever a @@ -149,10 +166,10 @@ def test_the_query_carries_the_session_params(start_server, make_mock_policy): @pytest.fixture -def authed_server(start_server: StartServer, make_mock_policy) -> PolicyServer: +def authed_server(start_server: StartServer, make_mock_policy) -> Served: policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - _host, _port, server = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True, auth_token=_TOKEN) - return server + served = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True, auth_token=_TOKEN) + return served def test_the_grpc_wire_gates_on_the_bearer_token(authed_server): @@ -266,11 +283,11 @@ def _trust_only(monkeypatch, root: bytes) -> None: @pytest.fixture -def edged(tls_edge, monkeypatch) -> Callable[[PolicyServer], str]: +def edged(tls_edge, monkeypatch) -> Callable[[Served], str]: """The ``grpcs://`` URL of a server reached through a TLS edge, with the client trusting its root.""" - def url(server: PolicyServer) -> str: - port, root = tls_edge(server.host, server.grpc_port) + def url(served: Served) -> str: + port, root = tls_edge(served.host, served.grpc_port) _trust_only(monkeypatch, root) return f'grpcs://{EDGE_HOST}:{port}' @@ -278,8 +295,8 @@ def url(server: PolicyServer) -> str: def test_a_session_through_a_tls_edge_handshakes_and_infers(both_wires, edged): - server, policy = both_wires - session = InferenceClient(edged(server)).new_session() + served, policy = both_wires + session = InferenceClient(edged(served)).new_session() try: assert session.metadata['model_name'] == 'stub' obs = {'image': 'test'} @@ -357,9 +374,9 @@ def test_a_port_that_never_answers_is_named_at_the_deadline(): def test_an_open_timeout_under_the_probe_budget_still_opens(both_wires): """The refusal probe takes a share of the budget, so a healthy server answers a short one.""" - server, _policy = both_wires + served, _policy = both_wires budget = grpc_wire._REFUSAL_PROBE_SEC / 2 - session = InferenceClient(grpc_url(server), open_timeout=budget, connect_deadline=0.0).new_session() + session = InferenceClient(grpc_url(served), open_timeout=budget, connect_deadline=0.0).new_session() try: assert session.infer({'image': 'test'}) == [{'action': [1, 2, 3]}] finally: @@ -372,8 +389,8 @@ def test_an_ipv6_host_binds_in_brackets(start_server: StartServer, make_mock_pol assert grpc_wire._bind_target('0.0.0.0', 9000) == '0.0.0.0:9000' policy = make_mock_policy([{'action': [4]}], {'model_name': 'stub'}) - _host, _port, server = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True, host='::1') - session = InferenceClient(f'grpc://[{server.host}]:{server.grpc_port}').new_session() + served = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True, host='::1') + session = InferenceClient(f'grpc://[{served.host}]:{served.grpc_port}').new_session() try: assert session.infer({'image': 'test'}) == [{'action': [4]}] finally: @@ -408,8 +425,8 @@ def chatty_client(monkeypatch) -> None: monkeypatch.setattr(grpc_wire, '_PING_EVERY_MS', 500) -def _silent_then_infer(server: PolicyServer) -> list[dict]: - session = InferenceClient(grpc_url(server)).new_session() +def _silent_then_infer(served: Served) -> list[dict]: + session = InferenceClient(grpc_url(served)).new_session() try: time.sleep(_SILENCE_SEC) return session.infer({'image': 'test'}) @@ -428,9 +445,9 @@ def test_a_server_on_the_grpc_ping_defaults_kills_the_silent_session( """gRPC's own server defaults answer those pings with ``GOAWAY too_many_pings``.""" monkeypatch.setattr(grpc_wire, '_server_options', lambda: list(grpc_wire._MESSAGE_SIZE_OPTIONS)) policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - _host, _port, server = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True) + served = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True) with pytest.raises(grpc.RpcError, match='Too many pings'): - _silent_then_infer(server) + _silent_then_infer(served) def _surfaces_at_once(url: str, blamed: str) -> None: @@ -459,9 +476,9 @@ def test_an_edge_that_selects_no_alpn_is_not_retried(both_wires, tls_edge, monke def test_a_timed_out_session_refuses_the_next_inference(both_wires): """The timeout closes the connection, and the server may answer inside the close's own wait.""" - server, policy = both_wires + served, policy = both_wires policy._mock_session.side_effect = lambda *_: time.sleep(1.0) or [{'action': [1, 2, 3]}] - session = InferenceClient(grpc_url(server), infer_timeout=0.2).new_session() + session = InferenceClient(grpc_url(served), infer_timeout=0.2).new_session() with pytest.raises(TimeoutError): session.infer({'image': 'test'}) # Without the guard this answers the first observation's actions, against the second's state. @@ -471,8 +488,8 @@ def test_a_timed_out_session_refuses_the_next_inference(both_wires): def test_a_connection_refuses_to_send_once_the_server_ends_the_stream(both_wires): """gRPC stops reading the request iterator then, so a write would wait out a whole timeout.""" - server, _policy = both_wires - conn = grpc_wire.GrpcClientConnection(f'{server.host}:{server.grpc_port}', f'{wire.SESSION_PATH}/unknown-model', '') + served, _policy = both_wires + conn = grpc_wire.GrpcClientConnection(f'{served.host}:{served.grpc_port}', f'{wire.SESSION_PATH}/unknown-model', '') try: conn.recv(timeout=10.0) # The server refuses the model in a frame, then ends the stream with that status. diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index a5a346776..96a90a9e2 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -527,7 +527,7 @@ def test_a_command_crossing_a_live_websocket_arrives_typed(start_server, make_mo pose = [0.4, 0.0, 0.6, 1, 0, 0, 0, 1, 0, 0, 0, 1] # translation + a 3x3 rotation, the wire's own layout wire_action = [{keys.ROBOT_COMMAND: {'type': 'cartesian_pos', 'pose': pose}, 'timestamp': 0.0}] served = make_mock_policy(wire_action, {'model_name': 'm'}) - host, port, _ = start_server(ChunkedSchedule() | remote | PolicySource(served)) + host, port, *_ = start_server(ChunkedSchedule() | remote | PolicySource(served)) session, rt = open_session(RemotePolicy(f'{host}:{port}')) actions = round_trip(session, rt, {keys.OBS_TIME_NS: 0}) diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 6bad1fd82..d45227287 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -1,5 +1,6 @@ import os import socket +import threading import time import urllib.parse from collections.abc import Callable, Generator @@ -14,6 +15,7 @@ from positronic import keys from positronic.offboard import keys as offboard_keys +from positronic.offboard import wire from positronic.offboard.client import InferenceClient, InferenceSession, _ConnectRetries from positronic.offboard.protocol import deserialise from positronic.offboard.server import AUTH_HEADER, AUTH_TOKEN_ENV, PolicyServer, bearer @@ -26,6 +28,9 @@ from positronic.policy.layers import ChunkedSchedule, TemporalStack from positronic.policy.spec import ModelSource, PolicySource, inline, remote +# Short enough to keep the idle test quick, long enough that a loaded box still reaches the first poll. +_A_MOMENT_IDLE = 0.5 + class _StubSource(ModelSource): """Serves one ready policy under any requested id, so route-supplied checkpoints resolve as-is.""" @@ -47,10 +52,21 @@ def meta(self, model_id: str) -> dict[str, Any]: return {'type': 'stub'} +def test_an_idle_server_stops_itself(make_mock_policy): + """The idle watchdog ends every wire it is serving on, so ``serve`` returns with nobody asking it to.""" + server = PolicyServer( + ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {})), idle_timeout_min=_A_MOMENT_IDLE / 60 + ) + serving = threading.Thread(target=server.serve, args=([wire.WebsocketWire('localhost', 0, server.api)],)) + serving.start() + serving.join(timeout=_A_MOMENT_IDLE * 20) + assert not serving.is_alive(), 'the idle watchdog left the server running' + + @pytest.fixture def stub_server(start_server, make_mock_policy) -> tuple[str, int, PolicyServer, MagicMock]: policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, server = start_server(ChunkedSchedule() | remote | _StubSource(policy)) + host, port, server, _ = start_server(ChunkedSchedule() | remote | _StubSource(policy)) return host, port, server, policy @@ -127,7 +143,7 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) def test_latest_checkpoint_pinned_once_at_startup(start_server, make_mock_policy): source = _LatestSource(make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'})) - host, port, _server = start_server(ChunkedSchedule() | remote | source) + host, port, *_ = start_server(ChunkedSchedule() | remote | source) # A newer checkpoint lands after startup (e.g. a training job writes it)... source.latest = '200' client = InferenceClient(f'{host}:{port}') @@ -156,7 +172,7 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) def test_load_progress_frames_reach_the_client(start_server, make_mock_policy): policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, _server = start_server(ChunkedSchedule() | remote | _ProgressSource(policy)) + host, port, *_ = start_server(ChunkedSchedule() | remote | _ProgressSource(policy)) # Requesting a non-pinned id forces a load inside the handshake; the source's progress # callbacks must arrive as ``loading`` frames before ``ready``. ws = connect(f'ws://{host}:{port}/api/v1/session/other') @@ -185,7 +201,7 @@ def meta(self): @pytest.fixture def codec_server(start_server, make_mock_policy) -> tuple[str, int, MagicMock]: policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, _server = start_server(ChunkedSchedule() | remote | _IdentityCodec() | _StubSource(policy)) + host, port, *_ = start_server(ChunkedSchedule() | remote | _IdentityCodec() | _StubSource(policy)) return host, port, policy @@ -224,7 +240,7 @@ def test_a_backend_that_cannot_answer_its_warmup_raises_and_still_ends_its_sessi def test_local_stack_declared_in_handshake(start_server, make_mock_policy): stub = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) pipeline = ChunkedSchedule() | remote | _IdentityCodec() | _StubSource(stub) - host, port, _server = start_server(pipeline) + host, port, *_ = start_server(pipeline) client = InferenceClient(f'{host}:{port}') session = client.new_session() try: @@ -276,7 +292,7 @@ def test_in_process_equals_remote_for_same_pipeline(start_server, open_session): def pipeline(): return ChunkedSchedule() | remote | ActionTimestamp(fps=10.0) | PolicySource(_ScriptedPolicy()) - host, port, _server = start_server(pipeline()) + host, port, *_ = start_server(pipeline()) remote_session, rt = open_session(RemotePolicy(f'{host}:{port}')) local_session, local_rt = open_session(inline(pipeline())) @@ -312,7 +328,7 @@ def _param_session(host: str, port: int, query: list[tuple[str, str]]) -> Infere def param_server(start_server, make_mock_policy) -> Generator[tuple[str, int], None, None]: stub = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) pipe_cfg = cfn.Config(_tunable_pipe, source=cfn.Config(_StubSource, policy=stub)) - host, port, _server = start_server(pipe_cfg) + host, port, *_ = start_server(pipe_cfg) yield host, port @@ -342,7 +358,7 @@ def _fps_pipe(source: ModelSource, fps: float = 10.0): def test_session_param_retunes_the_served_remote_half(start_server): pipe_cfg = cfn.Config(_fps_pipe, source=cfn.Config(PolicySource, policy=_ScriptedPolicy())) - host, port, _server = start_server(pipe_cfg) + host, port, *_ = start_server(pipe_cfg) # The wire carries the server-side half's output: relative timestamps spaced 1/fps. default_session = _param_session(host, port, []) @@ -407,7 +423,7 @@ def test_source_touching_session_param_rejected(param_server): def test_plain_pipe_server_rejects_session_params(start_server, make_mock_policy): stub = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, _server = start_server(_tunable_pipe(_StubSource(stub))) + host, port, *_ = start_server(_tunable_pipe(_StubSource(stub))) with pytest.raises(RuntimeError, match='config-launched'): _param_session(host, port, [('pad_start', 'false')]) @@ -433,7 +449,7 @@ def authed_endpoint(start_server, make_mock_policy) -> tuple[str, str]: if _LIVE_ENDPOINT: return _LIVE_ENDPOINT, os.environ[AUTH_TOKEN_ENV] policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, _server = start_server(ChunkedSchedule() | remote | _StubSource(policy), auth_token=_TOKEN) + host, port, *_ = start_server(ChunkedSchedule() | remote | _StubSource(policy), auth_token=_TOKEN) return f'{host}:{port}', _TOKEN @@ -517,7 +533,7 @@ def test_a_non_ascii_authorization_header_is_refused_rather_than_crashing(start_ """A header carries bytes, and Starlette hands them over latin-1 decoded, so a peer can put a non-ASCII ``str`` in front of the token comparison.""" policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, _server = start_server(ChunkedSchedule() | remote | _StubSource(policy), auth_token=_TOKEN) + host, port, *_ = start_server(ChunkedSchedule() | remote | _StubSource(policy), auth_token=_TOKEN) with socket.create_connection((host, port), timeout=5.0) as sock: sock.sendall( b'GET /api/v1/models HTTP/1.1\r\nHost: localhost\r\n' diff --git a/positronic/offboard/wire.py b/positronic/offboard/wire.py index d04a2087c..f4879544a 100644 --- a/positronic/offboard/wire.py +++ b/positronic/offboard/wire.py @@ -1,13 +1,16 @@ -"""The transports one session runs over, and the two ends of an open one. +"""The transports a session runs over, the two ends of an open one, and the server's end of a wire. A wire carries the ``protocol`` frames as opaque bytes and reads none of them, so the handshake and the inference loop read the same over every wire. ``grpc_wire`` holds the gRPC one. """ import abc -from typing import Protocol +import socket +from collections.abc import Awaitable, Callable, Mapping +from typing import NamedTuple, Protocol -from fastapi import WebSocket, WebSocketDisconnect +import uvicorn +from fastapi import APIRouter, Depends, FastAPI, WebSocket, WebSocketDisconnect, WebSocketException, status from starlette.datastructures import QueryParams from websockets.sync.connection import Connection @@ -20,11 +23,22 @@ # explicitly is what keeps the two wires equal when that default moves. MAX_MESSAGE_BYTES = 16 * 1024 * 1024 +# How long ``WebsocketWire.stop`` lets an open session finish before it cuts the connection. Left to +# itself uvicorn waits for ever, so a session mid-inference would hold the whole server open. +STOP_GRACE_SEC = 2 + class PeerDisconnected(Exception): """The peer ended the session.""" +class Endpoint(NamedTuple): + """Where a wire serves.""" + + host: str + port: int + + class ClientConnection(Protocol): """A client's end of one open session.""" @@ -73,6 +87,11 @@ class ServerConnection(abc.ABC): def peer(self) -> str: """Whom this session serves, for the log.""" + @property + @abc.abstractmethod + def endpoint(self) -> Endpoint: + """Where the wire that accepted this session serves.""" + @property @abc.abstractmethod def query_params(self) -> QueryParams: @@ -93,13 +112,18 @@ async def refuse(self, reason: str) -> None: class WebsocketServerConnection(ServerConnection): """A server's end of one websocket session, over an accepted ``WebSocket``.""" - def __init__(self, websocket: WebSocket): + def __init__(self, websocket: WebSocket, endpoint: Endpoint): self._websocket = websocket + self._endpoint = endpoint @property def peer(self) -> str: return str(self._websocket.client) + @property + def endpoint(self) -> Endpoint: + return self._endpoint + @property def query_params(self) -> QueryParams: return self._websocket.query_params @@ -115,3 +139,121 @@ async def receive(self) -> bytes: async def refuse(self, reason: str) -> None: await self._websocket.close(code=1008, reason=reason[:100]) + + +# What a wire hands the server for each session it accepts: the connection, and the model the route +# names, which is ``None`` where the route names the model the server pinned. +SessionHandler = Callable[[ServerConnection, str | None], Awaitable[None]] + +# Whether the session headers carry a credential the server accepts. Header names are lower case. +Authorized = Callable[[Mapping[str, str]], bool] + + +class Wire(abc.ABC): + """One transport that sessions arrive on. + + A wire reads its own route for the model a session names, and refuses an unauthorized peer before + the session opens. So a server hands every wire one ``SessionHandler`` and serves them all alike. + """ + + @property + @abc.abstractmethod + def endpoint(self) -> Endpoint: + """Where this wire serves. The port is bound, and so known, once ``start`` returns.""" + + @abc.abstractmethod + async def start(self, session: SessionHandler, authorized: Authorized) -> None: + """Bind, and give every accepted session to ``session``. Raises when the port is not free.""" + + @abc.abstractmethod + async def serve(self) -> None: + """Carry sessions until ``stop``, or until the wire ends for its own reason.""" + + @abc.abstractmethod + async def stop(self) -> None: + """End the wire, and every session on it.""" + + +def _listening_socket(host: str, port: int) -> socket.socket: + """A socket bound on ``host``, where a ``port`` of 0 takes any free one. + + The family comes from ``host`` itself, so an IPv6 host binds an IPv6 socket. Binding here rather + than inside uvicorn is what names the port before the wire serves, and holds it from then on. + """ + family, kind, proto, _canonical, address = socket.getaddrinfo( + host, port, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE + )[0] + sock = socket.socket(family, kind, proto) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(address) + # Listening here rather than at the first accept is what makes the port answer from the moment + # ``start`` returns: the kernel queues a connect that beats the serving loop to it. + sock.listen() + return sock + + +class WebsocketWire(Wire): + """The websocket wire: a session upgrades on ``SESSION_PATH``, and ``api`` answers on the same port. + + One uvicorn serves both, so the routes a client reads a model catalogue from sit on the endpoint it + opens sessions on. + """ + + def __init__(self, host: str, port: int, api: APIRouter): + self._host = host + self._port = port + self._api = api + self._socket: socket.socket | None = None + self._server: uvicorn.Server | None = None + self._endpoint: Endpoint | None = None + + @property + def endpoint(self) -> Endpoint: + assert self._endpoint is not None, 'The websocket wire has not started' + return self._endpoint + + async def start(self, session: SessionHandler, authorized: Authorized) -> None: + self._socket = _listening_socket(self._host, self._port) + self._endpoint = Endpoint(self._host, self._socket.getsockname()[1]) + app = FastAPI() + app.include_router(self._api) + self._route_sessions(app, session, authorized) + config = uvicorn.Config( + app, + host=self._host, + port=self._endpoint.port, + log_level='info', + ws_max_size=MAX_MESSAGE_BYTES, + timeout_graceful_shutdown=STOP_GRACE_SEC, + ) + self._server = uvicorn.Server(config) + + def _route_sessions(self, app: FastAPI, session: SessionHandler, authorized: Authorized) -> None: + async def require_auth(websocket: WebSocket) -> None: + """Refuses before ``accept()``, so an unauthorized peer never reaches the session handshake.""" + if not authorized(websocket.headers): + raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION) + + async def serve_pinned_model(websocket: WebSocket) -> None: + """Serves the model the server pinned. Naming a model is the path's job, so every query param + here is a pipeline override.""" + await websocket.accept() + await session(WebsocketServerConnection(websocket, self.endpoint), None) + + async def serve_named_model(websocket: WebSocket, model_id: str) -> None: + await websocket.accept() + await session(WebsocketServerConnection(websocket, self.endpoint), model_id) + + auth = [Depends(require_auth)] + app.websocket(SESSION_PATH, dependencies=auth)(serve_pinned_model) + # ``:path`` so an id that is itself a path (a HuggingFace repo, say) opens under the name the + # model catalogue advertises. + app.websocket(f'{SESSION_PATH}/{{model_id:path}}', dependencies=auth)(serve_named_model) + + async def serve(self) -> None: + assert self._server is not None and self._socket is not None, 'The websocket wire has not started' + await self._server.serve(sockets=[self._socket]) + + async def stop(self) -> None: + if self._server is not None: + self._server.should_exit = True From 14a2c3861e47fc696a24dabe20d15a9c6aa9c888 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Fri, 11 Sep 2026 10:40:24 +0000 Subject: [PATCH 32/46] Report every wire that failed, and put two constants beside their users `serve` gathered the wire tasks with `return_exceptions=True` and read none of the results, so a wire that failed as it stopped was discarded without a line in the log. One failure now raises out of `serve` and every other one is logged at ERROR. `STOP_GRACE_SEC` moves into `WebsocketWire`, which is its only user, and the idle test's own constant moves above the test that reads it. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/server.py | 15 ++++++--- positronic/offboard/tests/test_server.py | 39 ++++++++++++++++++++++-- positronic/offboard/wire.py | 10 +++--- 3 files changed, 51 insertions(+), 13 deletions(-) diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 246d08358..1beb8827b 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -429,17 +429,22 @@ async def _run(): if self.idle_timeout_min and self.idle_timeout_min > 0: ending.append(asyncio.create_task(self._idle_watchdog())) try: - done, _still_running = await asyncio.wait(serving + ending, return_when=asyncio.FIRST_COMPLETED) - # A wire that ended on an error raises here, rather than reading as the shutdown this waits for. - for task in done: - task.result() + await asyncio.wait(serving + ending, return_when=asyncio.FIRST_COMPLETED) finally: for task in ending: task.cancel() for w in wires: await w.stop() # Each wire ends the sessions it carries before this returns and the model slot closes. - await asyncio.gather(*serving, return_exceptions=True) + outcomes = await asyncio.gather(*serving, return_exceptions=True) + + failed = [(w, e) for w, e in zip(wires, outcomes, strict=True) if isinstance(e, Exception)] + # Only one failure can reach the caller, so the rest are reported here or nowhere. + for w, error in failed[1:]: + logger.error(f'{type(w).__name__} also failed: {error}', exc_info=error) + if failed: + # A wire that ended on an error raises, rather than reading as the shutdown this waits for. + raise failed[0][1] try: asyncio.run(_run()) diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index d45227287..f4fee83a6 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -1,3 +1,5 @@ +import asyncio +import logging import os import socket import threading @@ -28,9 +30,6 @@ from positronic.policy.layers import ChunkedSchedule, TemporalStack from positronic.policy.spec import ModelSource, PolicySource, inline, remote -# Short enough to keep the idle test quick, long enough that a loaded box still reaches the first poll. -_A_MOMENT_IDLE = 0.5 - class _StubSource(ModelSource): """Serves one ready policy under any requested id, so route-supplied checkpoints resolve as-is.""" @@ -52,6 +51,40 @@ def meta(self, model_id: str) -> dict[str, Any]: return {'type': 'stub'} +# Short enough to keep this test quick, long enough that a loaded box still reaches the first poll. +_A_MOMENT_IDLE = 0.5 + + +class _FailingWire(wire.Wire): + """Serves for ``after`` seconds, then falls over.""" + + def __init__(self, after: float): + self._after = after + + @property + def endpoint(self) -> wire.Endpoint: + return wire.Endpoint('localhost', 0) + + async def start(self, session: wire.SessionHandler, authorized: wire.Authorized) -> None: + pass + + async def serve(self) -> None: + await asyncio.sleep(self._after) + raise RuntimeError(f'the {self._after}s wire fell over') + + async def stop(self) -> None: + pass + + +def test_a_failing_wire_reaches_the_caller_and_the_rest_are_logged(make_mock_policy, caplog): + """No wire ends in silence: one failure raises out of ``serve``, and every other one is logged.""" + server = PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {}))) + with caplog.at_level(logging.ERROR, logger='positronic.offboard.server'): + with pytest.raises(RuntimeError, match='the 0.05s wire fell over'): + server.serve([_FailingWire(0.05), _FailingWire(0.1)]) + assert any('the 0.1s wire fell over' in record.getMessage() for record in caplog.records) + + def test_an_idle_server_stops_itself(make_mock_policy): """The idle watchdog ends every wire it is serving on, so ``serve`` returns with nobody asking it to.""" server = PolicyServer( diff --git a/positronic/offboard/wire.py b/positronic/offboard/wire.py index f4879544a..b15e63254 100644 --- a/positronic/offboard/wire.py +++ b/positronic/offboard/wire.py @@ -23,10 +23,6 @@ # explicitly is what keeps the two wires equal when that default moves. MAX_MESSAGE_BYTES = 16 * 1024 * 1024 -# How long ``WebsocketWire.stop`` lets an open session finish before it cuts the connection. Left to -# itself uvicorn waits for ever, so a session mid-inference would hold the whole server open. -STOP_GRACE_SEC = 2 - class PeerDisconnected(Exception): """The peer ended the session.""" @@ -199,6 +195,10 @@ class WebsocketWire(Wire): opens sessions on. """ + # How long ``stop`` lets an open session finish before it cuts the connection. Left to itself + # uvicorn waits for ever, so a session mid-inference would hold the whole server open. + STOP_GRACE_SEC = 2 + def __init__(self, host: str, port: int, api: APIRouter): self._host = host self._port = port @@ -224,7 +224,7 @@ async def start(self, session: SessionHandler, authorized: Authorized) -> None: port=self._endpoint.port, log_level='info', ws_max_size=MAX_MESSAGE_BYTES, - timeout_graceful_shutdown=STOP_GRACE_SEC, + timeout_graceful_shutdown=self.STOP_GRACE_SEC, ) self._server = uvicorn.Server(config) From 71609909a1a2da28fbb4d7152cb75fdc2ba73816 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Fri, 11 Sep 2026 10:56:22 +0000 Subject: [PATCH 33/46] Say what `PolicyServer` owns, and leave the wire contract to `Wire` The class docstring restated what a wire reads and checks, which `wire.Wire` and the offboard README both already say. It names what this server does with a wire instead. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/server.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 1beb8827b..5f7d4c700 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -191,10 +191,9 @@ class PolicyServer: The session flow is: accept → session params → resolve → load via manager → remote-half wrap → reset → inference loop - ``serve`` takes the wires sessions arrive on (see ``positronic.offboard.wire``). Each one reads its - own route for the model a session names and checks its own session headers, so the flow above is the - same over every wire and this server names none of them. ``api`` holds the server's own HTTP routes, - which a wire that speaks HTTP serves beside its sessions. + ``serve`` takes the wires sessions arrive on, and names none of them (see + ``positronic.offboard.wire``, which states what a wire owes a server). ``api`` holds this server's + own HTTP routes, which a wire that speaks HTTP serves beside its sessions. On startup (before accepting connections): resolve(None) → load. From 9c001385c0b9302fdbb713e5ccef5ee45ac36421 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Fri, 11 Sep 2026 11:12:04 +0000 Subject: [PATCH 34/46] Drive the vendor server tests through the session handler `PolicyServer.default_session` and `model_session` are gone, and the two lerobot vendor suites still called them. Both skip without their vendor package, so a local run reported them as skipped and CI failed on them. Each now builds the connection the websocket wire builds and calls `_serve_session`, which is the same handler the wire reaches. They are the only two callers the removal left. Two comments in `_listening_socket` said `is what` where the verb does the work. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/wire.py | 6 +++--- positronic/vendors/lerobot/tests/test_server.py | 9 +++++++-- positronic/vendors/lerobot_0_3_3/tests/test_server.py | 9 +++++++-- 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/positronic/offboard/wire.py b/positronic/offboard/wire.py index b15e63254..7b6c11b93 100644 --- a/positronic/offboard/wire.py +++ b/positronic/offboard/wire.py @@ -174,7 +174,7 @@ def _listening_socket(host: str, port: int) -> socket.socket: """A socket bound on ``host``, where a ``port`` of 0 takes any free one. The family comes from ``host`` itself, so an IPv6 host binds an IPv6 socket. Binding here rather - than inside uvicorn is what names the port before the wire serves, and holds it from then on. + than inside uvicorn names the port before the wire serves, and holds it from then on. """ family, kind, proto, _canonical, address = socket.getaddrinfo( host, port, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE @@ -182,8 +182,8 @@ def _listening_socket(host: str, port: int) -> socket.socket: sock = socket.socket(family, kind, proto) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.bind(address) - # Listening here rather than at the first accept is what makes the port answer from the moment - # ``start`` returns: the kernel queues a connect that beats the serving loop to it. + # Listening here rather than at the first accept makes the port answer from the moment ``start`` + # returns: the kernel queues a connect that beats the serving loop to it. sock.listen() return sock diff --git a/positronic/vendors/lerobot/tests/test_server.py b/positronic/vendors/lerobot/tests/test_server.py index 6cd5e4562..2556c83f2 100644 --- a/positronic/vendors/lerobot/tests/test_server.py +++ b/positronic/vendors/lerobot/tests/test_server.py @@ -4,6 +4,7 @@ from fastapi import WebSocketDisconnect from starlette.datastructures import QueryParams +from positronic.offboard import wire from positronic.offboard.protocol import deserialise from positronic.offboard.server import PolicyServer from positronic.policy.layers import ChunkedSchedule @@ -39,6 +40,10 @@ async def close(self, **kwargs): self.events.append('close') await self._close(**kwargs) + def as_connection(self) -> wire.WebsocketServerConnection: + """What the websocket wire hands the server for one session it has accepted.""" + return wire.WebsocketServerConnection(self, wire.Endpoint('localhost', 8000)) + @pytest.mark.asyncio async def test_lerobot_server_uses_configured_checkpoint(monkeypatch): @@ -60,7 +65,7 @@ async def fake_get_policy(checkpoint_id: str, websocket=None): await server._startup() websocket = _DummyWebSocket() - await server.default_session(websocket) + await server._serve_session(websocket.as_connection(), None) assert requested['checkpoint_id'] == '42' assert websocket.events == ['send_bytes'] @@ -95,7 +100,7 @@ async def test_lerobot_server_reports_unknown_checkpoint_id(monkeypatch): server._manager.get_policy.reset_mock() websocket = _DummyWebSocket() - await server.model_session(websocket, '42') + await server._serve_session(websocket.as_connection(), '42') assert websocket.events == ['send_bytes', 'close'] error_response = deserialise(websocket._send_bytes.await_args.args[0]) diff --git a/positronic/vendors/lerobot_0_3_3/tests/test_server.py b/positronic/vendors/lerobot_0_3_3/tests/test_server.py index 65165ca24..5e06e6cea 100644 --- a/positronic/vendors/lerobot_0_3_3/tests/test_server.py +++ b/positronic/vendors/lerobot_0_3_3/tests/test_server.py @@ -4,6 +4,7 @@ from fastapi import WebSocketDisconnect from starlette.datastructures import QueryParams +from positronic.offboard import wire from positronic.offboard.protocol import deserialise from positronic.policy.executor import blocking from positronic.policy.layers import ChunkedSchedule @@ -53,6 +54,10 @@ async def close(self, **kwargs): self.events.append('close') await self._close(**kwargs) + def as_connection(self) -> wire.WebsocketServerConnection: + """What the websocket wire hands the server for one session it has accepted.""" + return wire.WebsocketServerConnection(self, wire.Endpoint('localhost', 8000)) + def test_handshake_metadata_does_not_depend_on_the_factory(monkeypatch): """A factory's whole contract is returning a policy, so a plain one carrying no extra attributes @@ -96,7 +101,7 @@ async def fake_get_policy(checkpoint_id: str, websocket=None): await server._startup() websocket = _DummyWebSocket() - await server.default_session(websocket) + await server._serve_session(websocket.as_connection(), None) assert requested['checkpoint_id'] == '42' ready = deserialise(websocket._send_bytes.await_args_list[0].args[0]) @@ -134,7 +139,7 @@ async def test_lerobot_server_reports_unknown_checkpoint_id(monkeypatch): server._manager.get_policy.reset_mock() websocket = _DummyWebSocket() - await server.model_session(websocket, '42') + await server._serve_session(websocket.as_connection(), '42') assert websocket.events == ['send_bytes', 'close'] error_payload = websocket._send_bytes.await_args.args[0] From 43536c7aa062d0060017bcb4f613ea588573f3e8 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Fri, 11 Sep 2026 11:18:44 +0000 Subject: [PATCH 35/46] Stop the wires that bound when a later one cannot `serve` started every wire before it entered the cleanup `try`, so a `GrpcWire` that bound stayed bound when the `WebsocketWire` after it found its port taken: a live listener in front of a server nothing serves. Startup runs inside the `try` now, and the cleanup stops the wires that actually started. The status-update section said the frames prevent a WebSocket keepalive timeout. They hold the client's 30s-per-message handshake open, which is the same on either wire. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 2 +- positronic/offboard/server.py | 32 ++++++++++++++---------- positronic/offboard/tests/test_server.py | 28 +++++++++++++++++++++ 3 files changed, 48 insertions(+), 14 deletions(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index a2527d2b7..920a96265 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -161,7 +161,7 @@ This metadata tells the client: #### 2. Status Updates (Long Model Loading) -Some models may take a long time to load (e.g., OpenPI and GR00T can take 120-300s). The server sends periodic status updates during loading to prevent WebSocket keepalive timeouts: +Some models may take a long time to load (e.g., OpenPI and GR00T can take 120-300s). The client gives the handshake 30s per message, so the server sends periodic status updates during loading. This holds on either wire: ```json { diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 5f7d4c700..c96fb796a 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -416,28 +416,34 @@ def serve(self, wires: Sequence[wire.Wire], on_ready: Callable[[], None] | None async def _run(): self._loop, self._stop = asyncio.get_running_loop(), asyncio.Event() await self._startup() - for w in wires: - await w.start(self._serve_session, self._authorized) - self._last_activity = time.monotonic() - if on_ready is not None: - on_ready() - serving = [asyncio.create_task(w.serve()) for w in wires] - # What ends the server, beside a wire ending on its own: a caller's ``shutdown``, and the - # idle timeout. - ending: list[asyncio.Task] = [asyncio.create_task(self._stop.wait())] - if self.idle_timeout_min and self.idle_timeout_min > 0: - ending.append(asyncio.create_task(self._idle_watchdog())) + # A wire binds when it starts, so one that started is stopped even where a later one cannot + # bind and nothing ever serves. + started: list[wire.Wire] = [] + serving: list[asyncio.Task] = [] + ending: list[asyncio.Task] = [] try: + for w in wires: + await w.start(self._serve_session, self._authorized) + started.append(w) + self._last_activity = time.monotonic() + if on_ready is not None: + on_ready() + serving = [asyncio.create_task(w.serve()) for w in started] + # What ends the server, beside a wire ending on its own: a caller's ``shutdown``, and + # the idle timeout. + ending = [asyncio.create_task(self._stop.wait())] + if self.idle_timeout_min and self.idle_timeout_min > 0: + ending.append(asyncio.create_task(self._idle_watchdog())) await asyncio.wait(serving + ending, return_when=asyncio.FIRST_COMPLETED) finally: for task in ending: task.cancel() - for w in wires: + for w in started: await w.stop() # Each wire ends the sessions it carries before this returns and the model slot closes. outcomes = await asyncio.gather(*serving, return_exceptions=True) - failed = [(w, e) for w, e in zip(wires, outcomes, strict=True) if isinstance(e, Exception)] + failed = [(w, e) for w, e in zip(started, outcomes, strict=True) if isinstance(e, Exception)] # Only one failure can reach the caller, so the rest are reported here or nowhere. for w, error in failed[1:]: logger.error(f'{type(w).__name__} also failed: {error}', exc_info=error) diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index f4fee83a6..babc5c053 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -60,6 +60,7 @@ class _FailingWire(wire.Wire): def __init__(self, after: float): self._after = after + self.stopped = False @property def endpoint(self) -> wire.Endpoint: @@ -72,10 +73,37 @@ async def serve(self) -> None: await asyncio.sleep(self._after) raise RuntimeError(f'the {self._after}s wire fell over') + async def stop(self) -> None: + self.stopped = True + + +class _UnbindableWire(wire.Wire): + """A wire whose port is taken.""" + + @property + def endpoint(self) -> wire.Endpoint: + raise AssertionError('it never bound') + + async def start(self, session: wire.SessionHandler, authorized: wire.Authorized) -> None: + raise OSError('that port is taken') + + async def serve(self) -> None: + raise AssertionError('it never served') + async def stop(self) -> None: pass +def test_a_wire_that_cannot_bind_stops_the_ones_that_did(make_mock_policy): + """A wire binds when it starts, so a startup that gives up does not leave an earlier one holding a + port against a server nothing is serving.""" + server = PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {}))) + bound = _FailingWire(_A_MOMENT_IDLE) + with pytest.raises(OSError, match='that port is taken'): + server.serve([bound, _UnbindableWire()]) + assert bound.stopped, 'the wire that had bound was left holding its port' + + def test_a_failing_wire_reaches_the_caller_and_the_rest_are_logged(make_mock_policy, caplog): """No wire ends in silence: one failure raises out of ``serve``, and every other one is logged.""" server = PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {}))) From 55bb7f4a5bd701ce5e8b2bb9ea3fd71ad30b27f8 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Fri, 11 Sep 2026 21:04:57 +0000 Subject: [PATCH 36/46] Close the websocket wire's socket when startup rolls back A wire binds its listening socket in `start`, and uvicorn takes it over only when `serve` runs. A startup that rolls back stops a wire that bound but never served, so `stop` set uvicorn's `should_exit` on a server that never ran and the socket kept its port. Close the socket in `stop` where the wire never served. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/tests/test_server.py | 11 +++++++++++ positronic/offboard/wire.py | 6 ++++++ 2 files changed, 17 insertions(+) diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index babc5c053..52e270a2d 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -104,6 +104,17 @@ def test_a_wire_that_cannot_bind_stops_the_ones_that_did(make_mock_policy): assert bound.stopped, 'the wire that had bound was left holding its port' +def test_a_websocket_wire_releases_its_port_when_startup_rolls_back(make_mock_policy): + """A ``WebsocketWire`` binds a real socket when it starts, so a startup that rolls back frees the + port it took, not only the stub wires that own no socket.""" + server = PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {}))) + bound = wire.WebsocketWire('localhost', 0, server.api) + with pytest.raises(OSError, match='that port is taken'): + server.serve([bound, _UnbindableWire()]) + # A leaked listener would still hold the port, so binding a fresh socket to it would raise. + wire._listening_socket('localhost', bound.endpoint.port).close() + + def test_a_failing_wire_reaches_the_caller_and_the_rest_are_logged(make_mock_policy, caplog): """No wire ends in silence: one failure raises out of ``serve``, and every other one is logged.""" server = PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {}))) diff --git a/positronic/offboard/wire.py b/positronic/offboard/wire.py index 7b6c11b93..30069e23e 100644 --- a/positronic/offboard/wire.py +++ b/positronic/offboard/wire.py @@ -206,6 +206,7 @@ def __init__(self, host: str, port: int, api: APIRouter): self._socket: socket.socket | None = None self._server: uvicorn.Server | None = None self._endpoint: Endpoint | None = None + self._served = False @property def endpoint(self) -> Endpoint: @@ -252,8 +253,13 @@ async def serve_named_model(websocket: WebSocket, model_id: str) -> None: async def serve(self) -> None: assert self._server is not None and self._socket is not None, 'The websocket wire has not started' + self._served = True await self._server.serve(sockets=[self._socket]) async def stop(self) -> None: if self._server is not None: self._server.should_exit = True + # uvicorn releases the socket as it shuts down. A startup that rolls back stops a wire that + # bound but never served, so uvicorn never runs; close the socket here to free the port. + if self._socket is not None and not self._served: + self._socket.close() From 811c04c98f59d190a35cd57e9729b53d5445ce76 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 09:30:23 +0000 Subject: [PATCH 37/46] Move the websocket wire into its own module `wire.py` keeps the protocol every wire implements. `websocket_wire.py` holds the websocket wire and the two ends of a websocket session, beside `grpc_wire.py`. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 2 +- positronic/offboard/client.py | 4 +- positronic/offboard/server.py | 4 +- positronic/offboard/tests/conftest.py | 4 +- positronic/offboard/tests/test_server.py | 10 +- positronic/offboard/websocket_wire.py | 158 ++++++++++++++++++ positronic/offboard/wire.py | 154 +---------------- .../vendors/lerobot/tests/test_server.py | 6 +- .../lerobot_0_3_3/tests/test_server.py | 6 +- 9 files changed, 177 insertions(+), 171 deletions(-) create mode 100644 positronic/offboard/websocket_wire.py diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index 920a96265..9b7a249e9 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -253,7 +253,7 @@ The one server implementation behind every vendor. It serves a **policy pipeline ```python from positronic.offboard import PolicyServer -from positronic.offboard.wire import WebsocketWire +from positronic.offboard.websocket_wire import WebsocketWire from positronic.policy.spec import PolicySource, remote from positronic.policy.layers import ChunkedSchedule diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index f118f1461..05418f399 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -11,7 +11,7 @@ from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidStatus from websockets.sync.client import connect -from . import grpc_wire, protocol, wire +from . import grpc_wire, protocol, websocket_wire, wire from .protocol import deserialise, serialise, typed_commands logger = logging.getLogger(__name__) @@ -278,7 +278,7 @@ def _connect(self) -> wire.ClientConnection: ping_interval=20.0, max_size=wire.MAX_MESSAGE_BYTES, ) - return wire.WebsocketClientConnection(websocket) + return websocket_wire.WebsocketClientConnection(websocket) def _open_session(self) -> InferenceSession: """One attempt at a session, closing the connection whenever the handshake does not finish. diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index c96fb796a..2c12e8f7f 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -22,7 +22,7 @@ from positronic.policy.executor import blocking from positronic.policy.spec import ModelSource, Pipeline, split -from . import grpc_wire, protocol, wire +from . import grpc_wire, protocol, websocket_wire, wire from .protocol import deserialise, serialise logger = logging.getLogger(__name__) @@ -495,7 +495,7 @@ def serve( idle_timeout_min=idle_timeout_min, auth_token=os.environ.get(AUTH_TOKEN_ENV), ) - wires: list[wire.Wire] = [wire.WebsocketWire(host, port, server.api)] + wires: list[wire.Wire] = [websocket_wire.WebsocketWire(host, port, server.api)] if grpc_port is not None: wires.append(grpc_wire.GrpcWire(host, grpc_port)) server.serve(wires) diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 2692c267a..e4e23bedd 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -5,7 +5,7 @@ import pytest -from positronic.offboard import grpc_wire, wire +from positronic.offboard import grpc_wire, websocket_wire, wire from positronic.offboard.server import PolicyServer from positronic.policy import Policy, Session from positronic.policy.executor import Executor @@ -37,7 +37,7 @@ def start_server() -> Generator[StartServer, None, None]: def start(pipeline, *, grpc: bool = False, **server_kwargs) -> Served: host = server_kwargs.pop('host', 'localhost') server = PolicyServer(pipeline, **server_kwargs) - wires: list[wire.Wire] = [wire.WebsocketWire(host, 0, server.api)] + wires: list[wire.Wire] = [websocket_wire.WebsocketWire(host, 0, server.api)] if grpc: wires.append(grpc_wire.GrpcWire(host, 0)) ready = threading.Event() diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 52e270a2d..f1498312f 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -17,13 +17,13 @@ from positronic import keys from positronic.offboard import keys as offboard_keys -from positronic.offboard import wire +from positronic.offboard import websocket_wire, wire from positronic.offboard.client import InferenceClient, InferenceSession, _ConnectRetries from positronic.offboard.protocol import deserialise from positronic.offboard.server import AUTH_HEADER, AUTH_TOKEN_ENV, PolicyServer, bearer from positronic.offboard.server_utils import warmup from positronic.offboard.tests.conftest import round_trip -from positronic.offboard.wire import WebsocketClientConnection +from positronic.offboard.websocket_wire import WebsocketClientConnection from positronic.policy import Codec, Policy, RemotePolicy, Session from positronic.policy.base import Runtime from positronic.policy.codec import ActionTimestamp @@ -108,11 +108,11 @@ def test_a_websocket_wire_releases_its_port_when_startup_rolls_back(make_mock_po """A ``WebsocketWire`` binds a real socket when it starts, so a startup that rolls back frees the port it took, not only the stub wires that own no socket.""" server = PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {}))) - bound = wire.WebsocketWire('localhost', 0, server.api) + bound = websocket_wire.WebsocketWire('localhost', 0, server.api) with pytest.raises(OSError, match='that port is taken'): server.serve([bound, _UnbindableWire()]) # A leaked listener would still hold the port, so binding a fresh socket to it would raise. - wire._listening_socket('localhost', bound.endpoint.port).close() + websocket_wire._listening_socket('localhost', bound.endpoint.port).close() def test_a_failing_wire_reaches_the_caller_and_the_rest_are_logged(make_mock_policy, caplog): @@ -129,7 +129,7 @@ def test_an_idle_server_stops_itself(make_mock_policy): server = PolicyServer( ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {})), idle_timeout_min=_A_MOMENT_IDLE / 60 ) - serving = threading.Thread(target=server.serve, args=([wire.WebsocketWire('localhost', 0, server.api)],)) + serving = threading.Thread(target=server.serve, args=([websocket_wire.WebsocketWire('localhost', 0, server.api)],)) serving.start() serving.join(timeout=_A_MOMENT_IDLE * 20) assert not serving.is_alive(), 'the idle watchdog left the server running' diff --git a/positronic/offboard/websocket_wire.py b/positronic/offboard/websocket_wire.py new file mode 100644 index 000000000..ed7e5b780 --- /dev/null +++ b/positronic/offboard/websocket_wire.py @@ -0,0 +1,158 @@ +"""The websocket wire, and the two ends of a websocket session.""" + +import socket + +import uvicorn +from fastapi import APIRouter, Depends, FastAPI, WebSocket, WebSocketDisconnect, WebSocketException, status +from starlette.datastructures import QueryParams +from websockets.sync.connection import Connection + +from . import wire + + +class WebsocketClientConnection: + """A client's end of one websocket session.""" + + def __init__(self, websocket: Connection): + self._websocket = websocket + + def send(self, message: bytes) -> None: + self._websocket.send(message) + + def recv(self, timeout: float | None = None) -> bytes: + message = self._websocket.recv(timeout=timeout) + assert isinstance(message, bytes), f'A frame is bytes, and this one is {type(message).__name__}' + return message + + def close(self) -> str: + state_before_close = self._websocket.state.name + self._websocket.close() + # A close that times out still reaches CLOSED locally; only the close code says the server answered. + return f'state {state_before_close} -> {self._websocket.state.name}, close code {self._websocket.close_code}' + + +class WebsocketServerConnection(wire.ServerConnection): + """A server's end of one websocket session, over an accepted ``WebSocket``.""" + + def __init__(self, websocket: WebSocket, endpoint: wire.Endpoint): + self._websocket = websocket + self._endpoint = endpoint + + @property + def peer(self) -> str: + return str(self._websocket.client) + + @property + def endpoint(self) -> wire.Endpoint: + return self._endpoint + + @property + def query_params(self) -> QueryParams: + return self._websocket.query_params + + async def send(self, message: bytes) -> None: + await self._websocket.send_bytes(message) + + async def receive(self) -> bytes: + try: + return await self._websocket.receive_bytes() + except WebSocketDisconnect as e: + raise wire.PeerDisconnected(str(e)) from e + + async def refuse(self, reason: str) -> None: + await self._websocket.close(code=1008, reason=reason[:100]) + + +def _listening_socket(host: str, port: int) -> socket.socket: + """A socket bound on ``host``, where a ``port`` of 0 takes any free one. + + The family comes from ``host`` itself, so an IPv6 host binds an IPv6 socket. Binding here rather + than inside uvicorn names the port before the wire serves, and holds it from then on. + """ + family, kind, proto, _canonical, address = socket.getaddrinfo( + host, port, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE + )[0] + sock = socket.socket(family, kind, proto) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(address) + # Listening here rather than at the first accept makes the port answer from the moment ``start`` + # returns: the kernel queues a connect that beats the serving loop to it. + sock.listen() + return sock + + +class WebsocketWire(wire.Wire): + """The websocket wire: a session upgrades on ``wire.SESSION_PATH``, and ``api`` answers on the same port. + + One uvicorn serves both, so the routes a client reads a model catalogue from sit on the endpoint it + opens sessions on. + """ + + # How long ``stop`` lets an open session finish before it cuts the connection. Left to itself + # uvicorn waits for ever, so a session mid-inference would hold the whole server open. + STOP_GRACE_SEC = 2 + + def __init__(self, host: str, port: int, api: APIRouter): + self._host = host + self._port = port + self._api = api + self._socket: socket.socket | None = None + self._server: uvicorn.Server | None = None + self._endpoint: wire.Endpoint | None = None + self._served = False + + @property + def endpoint(self) -> wire.Endpoint: + assert self._endpoint is not None, 'The websocket wire has not started' + return self._endpoint + + async def start(self, session: wire.SessionHandler, authorized: wire.Authorized) -> None: + self._socket = _listening_socket(self._host, self._port) + self._endpoint = wire.Endpoint(self._host, self._socket.getsockname()[1]) + app = FastAPI() + app.include_router(self._api) + self._route_sessions(app, session, authorized) + config = uvicorn.Config( + app, + host=self._host, + port=self._endpoint.port, + log_level='info', + ws_max_size=wire.MAX_MESSAGE_BYTES, + timeout_graceful_shutdown=self.STOP_GRACE_SEC, + ) + self._server = uvicorn.Server(config) + + def _route_sessions(self, app: FastAPI, session: wire.SessionHandler, authorized: wire.Authorized) -> None: + async def require_auth(websocket: WebSocket) -> None: + """Refuses before ``accept()``, so an unauthorized peer never reaches the session handshake.""" + if not authorized(websocket.headers): + raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION) + + async def serve_pinned_model(websocket: WebSocket) -> None: + """Serves the model the server pinned. Naming a model is the path's job, so every query param + here is a pipeline override.""" + await websocket.accept() + await session(WebsocketServerConnection(websocket, self.endpoint), None) + + async def serve_named_model(websocket: WebSocket, model_id: str) -> None: + await websocket.accept() + await session(WebsocketServerConnection(websocket, self.endpoint), model_id) + + auth = [Depends(require_auth)] + app.websocket(wire.SESSION_PATH, dependencies=auth)(serve_pinned_model) + # ``:path`` so an id that is itself a path (a HuggingFace repo, say) opens under the name the + # model catalogue advertises. + app.websocket(f'{wire.SESSION_PATH}/{{model_id:path}}', dependencies=auth)(serve_named_model) + + async def serve(self) -> None: + assert self._server is not None and self._socket is not None, 'The websocket wire has not started' + self._served = True + await self._server.serve(sockets=[self._socket]) + + async def stop(self) -> None: + if self._server is not None: + self._server.should_exit = True + # uvicorn releases the socket as it shuts down. A startup that rolls back stops a wire that + # bound but never served, so uvicorn never runs; close the socket here to free the port. + if self._socket is not None and not self._served: + self._socket.close() diff --git a/positronic/offboard/wire.py b/positronic/offboard/wire.py index 30069e23e..a62fcd005 100644 --- a/positronic/offboard/wire.py +++ b/positronic/offboard/wire.py @@ -1,18 +1,14 @@ """The transports a session runs over, the two ends of an open one, and the server's end of a wire. A wire carries the ``protocol`` frames as opaque bytes and reads none of them, so the handshake and -the inference loop read the same over every wire. ``grpc_wire`` holds the gRPC one. +the inference loop read the same over every wire. ``websocket_wire`` and ``grpc_wire`` hold the two. """ import abc -import socket from collections.abc import Awaitable, Callable, Mapping from typing import NamedTuple, Protocol -import uvicorn -from fastapi import APIRouter, Depends, FastAPI, WebSocket, WebSocketDisconnect, WebSocketException, status from starlette.datastructures import QueryParams -from websockets.sync.connection import Connection # The route a session opens on. The websocket wire puts it in the URL; the gRPC wire names it in the # session metadata, so both wires address a model the same way. @@ -54,27 +50,6 @@ def close(self) -> str: ... -class WebsocketClientConnection: - """A client's end of one websocket session.""" - - def __init__(self, websocket: Connection): - self._websocket = websocket - - def send(self, message: bytes) -> None: - self._websocket.send(message) - - def recv(self, timeout: float | None = None) -> bytes: - message = self._websocket.recv(timeout=timeout) - assert isinstance(message, bytes), f'A frame is bytes, and this one is {type(message).__name__}' - return message - - def close(self) -> str: - state_before_close = self._websocket.state.name - self._websocket.close() - # A close that times out still reaches CLOSED locally; only the close code says the server answered. - return f'state {state_before_close} -> {self._websocket.state.name}, close code {self._websocket.close_code}' - - class ServerConnection(abc.ABC): """A server's end of one open session.""" @@ -105,38 +80,6 @@ async def refuse(self, reason: str) -> None: """End a session the server cannot serve, telling the client why.""" -class WebsocketServerConnection(ServerConnection): - """A server's end of one websocket session, over an accepted ``WebSocket``.""" - - def __init__(self, websocket: WebSocket, endpoint: Endpoint): - self._websocket = websocket - self._endpoint = endpoint - - @property - def peer(self) -> str: - return str(self._websocket.client) - - @property - def endpoint(self) -> Endpoint: - return self._endpoint - - @property - def query_params(self) -> QueryParams: - return self._websocket.query_params - - async def send(self, message: bytes) -> None: - await self._websocket.send_bytes(message) - - async def receive(self) -> bytes: - try: - return await self._websocket.receive_bytes() - except WebSocketDisconnect as e: - raise PeerDisconnected(str(e)) from e - - async def refuse(self, reason: str) -> None: - await self._websocket.close(code=1008, reason=reason[:100]) - - # What a wire hands the server for each session it accepts: the connection, and the model the route # names, which is ``None`` where the route names the model the server pinned. SessionHandler = Callable[[ServerConnection, str | None], Awaitable[None]] @@ -168,98 +111,3 @@ async def serve(self) -> None: @abc.abstractmethod async def stop(self) -> None: """End the wire, and every session on it.""" - - -def _listening_socket(host: str, port: int) -> socket.socket: - """A socket bound on ``host``, where a ``port`` of 0 takes any free one. - - The family comes from ``host`` itself, so an IPv6 host binds an IPv6 socket. Binding here rather - than inside uvicorn names the port before the wire serves, and holds it from then on. - """ - family, kind, proto, _canonical, address = socket.getaddrinfo( - host, port, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE - )[0] - sock = socket.socket(family, kind, proto) - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - sock.bind(address) - # Listening here rather than at the first accept makes the port answer from the moment ``start`` - # returns: the kernel queues a connect that beats the serving loop to it. - sock.listen() - return sock - - -class WebsocketWire(Wire): - """The websocket wire: a session upgrades on ``SESSION_PATH``, and ``api`` answers on the same port. - - One uvicorn serves both, so the routes a client reads a model catalogue from sit on the endpoint it - opens sessions on. - """ - - # How long ``stop`` lets an open session finish before it cuts the connection. Left to itself - # uvicorn waits for ever, so a session mid-inference would hold the whole server open. - STOP_GRACE_SEC = 2 - - def __init__(self, host: str, port: int, api: APIRouter): - self._host = host - self._port = port - self._api = api - self._socket: socket.socket | None = None - self._server: uvicorn.Server | None = None - self._endpoint: Endpoint | None = None - self._served = False - - @property - def endpoint(self) -> Endpoint: - assert self._endpoint is not None, 'The websocket wire has not started' - return self._endpoint - - async def start(self, session: SessionHandler, authorized: Authorized) -> None: - self._socket = _listening_socket(self._host, self._port) - self._endpoint = Endpoint(self._host, self._socket.getsockname()[1]) - app = FastAPI() - app.include_router(self._api) - self._route_sessions(app, session, authorized) - config = uvicorn.Config( - app, - host=self._host, - port=self._endpoint.port, - log_level='info', - ws_max_size=MAX_MESSAGE_BYTES, - timeout_graceful_shutdown=self.STOP_GRACE_SEC, - ) - self._server = uvicorn.Server(config) - - def _route_sessions(self, app: FastAPI, session: SessionHandler, authorized: Authorized) -> None: - async def require_auth(websocket: WebSocket) -> None: - """Refuses before ``accept()``, so an unauthorized peer never reaches the session handshake.""" - if not authorized(websocket.headers): - raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION) - - async def serve_pinned_model(websocket: WebSocket) -> None: - """Serves the model the server pinned. Naming a model is the path's job, so every query param - here is a pipeline override.""" - await websocket.accept() - await session(WebsocketServerConnection(websocket, self.endpoint), None) - - async def serve_named_model(websocket: WebSocket, model_id: str) -> None: - await websocket.accept() - await session(WebsocketServerConnection(websocket, self.endpoint), model_id) - - auth = [Depends(require_auth)] - app.websocket(SESSION_PATH, dependencies=auth)(serve_pinned_model) - # ``:path`` so an id that is itself a path (a HuggingFace repo, say) opens under the name the - # model catalogue advertises. - app.websocket(f'{SESSION_PATH}/{{model_id:path}}', dependencies=auth)(serve_named_model) - - async def serve(self) -> None: - assert self._server is not None and self._socket is not None, 'The websocket wire has not started' - self._served = True - await self._server.serve(sockets=[self._socket]) - - async def stop(self) -> None: - if self._server is not None: - self._server.should_exit = True - # uvicorn releases the socket as it shuts down. A startup that rolls back stops a wire that - # bound but never served, so uvicorn never runs; close the socket here to free the port. - if self._socket is not None and not self._served: - self._socket.close() diff --git a/positronic/vendors/lerobot/tests/test_server.py b/positronic/vendors/lerobot/tests/test_server.py index 2556c83f2..2078612f1 100644 --- a/positronic/vendors/lerobot/tests/test_server.py +++ b/positronic/vendors/lerobot/tests/test_server.py @@ -4,7 +4,7 @@ from fastapi import WebSocketDisconnect from starlette.datastructures import QueryParams -from positronic.offboard import wire +from positronic.offboard import websocket_wire, wire from positronic.offboard.protocol import deserialise from positronic.offboard.server import PolicyServer from positronic.policy.layers import ChunkedSchedule @@ -40,9 +40,9 @@ async def close(self, **kwargs): self.events.append('close') await self._close(**kwargs) - def as_connection(self) -> wire.WebsocketServerConnection: + def as_connection(self) -> websocket_wire.WebsocketServerConnection: """What the websocket wire hands the server for one session it has accepted.""" - return wire.WebsocketServerConnection(self, wire.Endpoint('localhost', 8000)) + return websocket_wire.WebsocketServerConnection(self, wire.Endpoint('localhost', 8000)) @pytest.mark.asyncio diff --git a/positronic/vendors/lerobot_0_3_3/tests/test_server.py b/positronic/vendors/lerobot_0_3_3/tests/test_server.py index 5e06e6cea..ed522ff8c 100644 --- a/positronic/vendors/lerobot_0_3_3/tests/test_server.py +++ b/positronic/vendors/lerobot_0_3_3/tests/test_server.py @@ -4,7 +4,7 @@ from fastapi import WebSocketDisconnect from starlette.datastructures import QueryParams -from positronic.offboard import wire +from positronic.offboard import websocket_wire, wire from positronic.offboard.protocol import deserialise from positronic.policy.executor import blocking from positronic.policy.layers import ChunkedSchedule @@ -54,9 +54,9 @@ async def close(self, **kwargs): self.events.append('close') await self._close(**kwargs) - def as_connection(self) -> wire.WebsocketServerConnection: + def as_connection(self) -> websocket_wire.WebsocketServerConnection: """What the websocket wire hands the server for one session it has accepted.""" - return wire.WebsocketServerConnection(self, wire.Endpoint('localhost', 8000)) + return websocket_wire.WebsocketServerConnection(self, wire.Endpoint('localhost', 8000)) def test_handshake_metadata_does_not_depend_on_the_factory(monkeypatch): From da23946ac330e89649175fae88f57f8cc0fef8b0 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 09:50:10 +0000 Subject: [PATCH 38/46] Cut the comments and docs of the wires to what holds Each comment and docstring the gRPC wire adds states one fact, in one or two lines. A sentence that gave a reason, weighed an alternative, or told the reader how to read a fact is gone. A docstring that repeated its test name is gone. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 50 ++++++------- positronic/offboard/client.py | 33 ++++----- positronic/offboard/grpc_wire.py | 80 ++++++++++----------- positronic/offboard/server.py | 33 ++++----- positronic/offboard/stub.py | 18 ++--- positronic/offboard/tests/conftest.py | 4 +- positronic/offboard/tests/test_grpc_wire.py | 60 +++++++--------- positronic/offboard/tests/test_server.py | 14 ++-- positronic/offboard/websocket_wire.py | 33 ++++----- positronic/offboard/wire.py | 27 ++++--- workflows/nebius/serve.sh | 17 +++-- 11 files changed, 163 insertions(+), 206 deletions(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index 9b7a249e9..daa99a85f 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -9,7 +9,7 @@ The unified protocol is built to enable ANY hardware to connect to ANY model. Al ### Wires The protocol is a sequence of msgpack frames, and two wires carry them. Both carry the same frames in -the same order, so everything below holds on each. +the same order; everything below holds on each. | Wire | URL | Port | |---|---|---| @@ -18,32 +18,32 @@ the same order, so everything below holds on each. | gRPC over TLS | `grpcs://host:443/api/v1/session[/]` | a TLS edge in front of that same `grpc_port` | The WebSocket wire is the default, and a server serves gRPC only when `grpc_port` names a port. A -gRPC session is one bidirectional stream of the same frames, so no `.proto` file describes them. +gRPC session is one bidirectional stream of the same frames; no `.proto` file describes them. The session path and the query cross as the `positronic-session-path` and `positronic-session-query` metadata, and `Authorization` crosses as the `authorization` metadata. -Python's WebSocket stack costs about 30 ms per 846 KiB observation in framing and reassembly, which -gRPC does in about 1 ms, so take the gRPC wire wherever it reaches. +Python's WebSocket stack spends about 30 ms per 846 KiB observation on framing and reassembly; gRPC +spends about 1 ms. Take the gRPC wire wherever it reaches. -It reaches through a managed HTTPS front, which is how an authenticated endpoint is served: the front -terminates TLS and the HTTP/2 connection runs end to end, so the server binds a plaintext port and -holds no certificate of its own. The front has to negotiate HTTP/2 over ALPN — check a new one with -`openssl s_client -alpn h2 -connect :443`. On a Nebius Serverless Endpoint that means declaring -the gRPC port as an ordinary HTTP port and dialling its `https://` host as `grpcs://:443`; a -port declared `/tcp` is fronted by a `tls://` URL that negotiates no ALPN, which gRPC refuses with +It reaches through a managed HTTPS front, which is how an authenticated endpoint is served. The front +terminates TLS, and the HTTP/2 connection runs end to end; the server binds a plaintext port and holds +no certificate. The front must select HTTP/2 over ALPN. Check a new front with +`openssl s_client -alpn h2 -connect :443`. On a Nebius Serverless Endpoint, declare the gRPC +port as an ordinary HTTP port and dial its `https://` host as `grpcs://:443`. A port declared +`/tcp` gets a `tls://` URL that selects no ALPN protocol, and gRPC refuses it with `Cannot check peer: missing selected ALPN property`. -Through such an endpoint an 846 KiB observation round-trips in about 6 ms over gRPC against about -60 ms over the WebSocket, and gRPC holds that at 10 Hz, which is 8 MB/s of observation. The front -shapes a session that outruns it: a back-to-back loop settles at about 83 ms a round trip after some -11 MB, and gets its speed back after a minute of quiet. The WebSocket holds its 60 ms throughout, -below the rate the front shapes at. +Through such an endpoint an 846 KiB observation round-trips in about 6 ms over gRPC and about 60 ms +over the WebSocket. gRPC holds 6 ms at 10 Hz, which is 8 MB/s of observation. The front shapes a +session that sends faster: a back-to-back loop settles at about 83 ms a round trip after some 11 MB, +and returns to 6 ms after a minute of quiet. The WebSocket holds its 60 ms throughout, below the rate +the front shapes at. -Both wires ping through a silent wait, so a front that drops a connection it has read nothing from — -the managed one after about 90 s — does not cut an inference the model is still working on. +Both wires ping through a silent wait. A front drops a connection it reads nothing from (the managed +front after about 90 s), and the pings keep an inference open through that wait. -`/api/v1/models` is an HTTP route, so it stays on the server's `port`. `InferenceClient.list_models` -over a `grpc://` URL says so. +`/api/v1/models` is an HTTP route and stays on the server's `port`. `InferenceClient.list_models` +refuses a `grpc://` URL. ### Authentication @@ -161,7 +161,7 @@ This metadata tells the client: #### 2. Status Updates (Long Model Loading) -Some models may take a long time to load (e.g., OpenPI and GR00T can take 120-300s). The client gives the handshake 30s per message, so the server sends periodic status updates during loading. This holds on either wire: +Some models may take a long time to load (e.g., OpenPI and GR00T can take 120-300s). The client gives the handshake 30 s per message; the server sends status updates during loading, on either wire: ```json { @@ -262,11 +262,11 @@ server = PolicyServer(pipeline) server.serve([WebsocketWire('0.0.0.0', 8000, server.api)]) ``` -`serve` takes the wires sessions arrive on, and the server names none of them: each wire binds its own -port, reads its own route for the model a session asks for, and checks its own session headers. Add -`grpc_wire.GrpcWire(host, port)` to the list to serve gRPC beside the websocket. A wire that speaks HTTP -takes `server.api`, the model catalogue, and answers it on the same port it carries sessions on. A wire -asked for port 0 binds any free one and names it in `wire.endpoint`. +`serve` takes the wires sessions arrive on. Each wire binds its own port, reads its own route for the +model a session asks for, and checks its own session headers. Add `grpc_wire.GrpcWire(host, port)` to +the list to serve gRPC beside the websocket. A wire that speaks HTTP takes `server.api`, the model +catalogue, and answers it on the port it carries sessions on. A wire asked for port 0 binds any free +one and names it in `wire.endpoint`. `PolicySource` serves one ready in-process policy; vendors instead define a `ModelSource` over a checkpoint directory. Passing a `cfn.Config` that builds the pipeline — as the vendor servers do with their named pipelines — enables [session parameters](#session-parameters); an instantiated pipeline serves exactly as launched. `recording_dir` enables the per-session recording taps described above, and `idle_timeout_min` ends the server after that many minutes without activity. diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 05418f399..965718d46 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -103,20 +103,19 @@ class _ConnectOutcome(Enum): class _Refusal(Enum): """What a refused connect says about the server.""" - COLD = 'cold' # still coming up; retry to the deadline + COLD = 'cold' # a backend still starting; retry to the deadline FORBIDDEN = 'forbidden' # a cold backend, or a refused credential; a few attempts, then surface - FINAL = 'final' # the endpoint is saying no; surface at once + FINAL = 'final' # a permanent refusal; surface at once _COLD_GRPC_CODES = (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.DEADLINE_EXCEEDED) def _refusal(e: Exception) -> _Refusal: - """How to read a refused connect, over either wire. + """What a refused connect says about the server, over either wire. - Each gRPC code stands for the HTTP status its wire twin answers: ``PERMISSION_DENIED`` for 403, - ``UNAVAILABLE`` for 503, ``RESOURCE_EXHAUSTED`` for 429. A TLS edge no client can use answers - ``UNAVAILABLE`` too, exactly as a cold backend does, so its details tell them apart. + ``PERMISSION_DENIED`` reads as 403, ``UNAVAILABLE`` as 503, ``RESOURCE_EXHAUSTED`` as 429. A TLS + edge no client can use answers ``UNAVAILABLE`` too; its details tell it from a cold backend. """ if isinstance(e, InvalidStatus): status = e.response.status_code @@ -139,8 +138,8 @@ def _refusal(e: Exception) -> _Refusal: class _ConnectRetries: """The retry policy over one ``new_session``'s connect attempts. - A refusal both wires answer for a cold backend and for a refused credential — HTTP 403, gRPC - ``PERMISSION_DENIED`` — gets a few attempts rather than the whole ``connect_deadline``. + A 403 or a ``PERMISSION_DENIED`` means a cold backend or a refused credential, and gets + ``MAX_FORBIDDEN_ATTEMPTS`` attempts. """ MAX_FORBIDDEN_ATTEMPTS = 3 @@ -208,9 +207,8 @@ class InferenceClient: says about the session — the model id it names and the query it carries as session params — reaches the server exactly as written, so every session opened here serves that model with those params. - ``grpc://`` names the same session on the gRPC wire, which the server offers on a port of its own, and - ``grpcs://`` names that port behind a TLS edge. Either port carries sessions alone, so ``list_models`` - needs the HTTP URL. + ``grpc://`` opens the session on the gRPC wire, on the server's own gRPC port; ``grpcs://`` reaches that + port through a TLS edge. That port carries sessions alone: ``list_models`` needs the HTTP URL. ``headers`` carry auth, whether the server checks it or a proxy in front of it does — credentials stay out of the URL, which is meant to be safe to hand around. @@ -281,11 +279,10 @@ def _connect(self) -> wire.ClientConnection: return websocket_wire.WebsocketClientConnection(websocket) def _open_session(self) -> InferenceSession: - """One attempt at a session, closing the connection whenever the handshake does not finish. + """One attempt at a session. The connection closes when the handshake does not finish. - A refusal the server sends as a protocol frame — an unknown model, a session param it rejects - — raises past every transport handler, and a gRPC connection holds a reader thread until it - is closed. + A refusal sent as a protocol frame (an unknown model, a rejected session param) raises past every + transport handler, and a gRPC connection holds a reader thread until it is closed. """ conn = self._connect() try: @@ -306,9 +303,9 @@ def new_session(self) -> InferenceSession: # 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 - # Each of these is a backend that is not ready yet — a timed-out connect, a reset TLS - # handshake, a refused upgrade or gRPC call, a dropped status handshake — so one must not - # kill the run. ``_ConnectRetries`` decides which of them is the endpoint saying no. + # Each of these can be a backend that is not ready: a timed-out connect, a reset TLS handshake, a + # refused upgrade or gRPC call, a dropped status handshake. ``_ConnectRetries`` tells a permanent + # refusal apart. except ( TimeoutError, ssl.SSLError, diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index fb4942bd5..09d2ea24d 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -1,7 +1,6 @@ -"""The gRPC wire: one bidirectional stream per session, carrying the same ``protocol`` frames. +"""The gRPC wire: one bidirectional stream per session, which carries the ``protocol`` frames. -The stream is untyped bytes on both sides, so there is no protobuf schema and no generated code: a -generic handler with no serialiser hands each frame over as it arrived. +The stream is untyped bytes on both sides. There is no protobuf schema and no generated code. """ import logging @@ -24,7 +23,7 @@ METHOD = 'Session' METHOD_PATH = f'/{SERVICE}/{METHOD}' -# What the websocket wire says in the URL, said here in the session metadata. +# The session path and the query cross as metadata; the websocket wire carries them in the URL. SESSION_PATH_HEADER = 'positronic-session-path' SESSION_QUERY_HEADER = 'positronic-session-query' @@ -33,29 +32,28 @@ ('grpc.max_send_message_length', wire.MAX_MESSAGE_BYTES), ] -# How often the client pings a connection nothing is crossing, so no front reads it as dead. +# How often the client pings an idle connection. A front drops a connection it reads nothing from. _PING_EVERY_MS = 20_000 _PING_ANSWER_TIMEOUT_MS = 10_000 _PING_TOLERATED_EVERY_MS = 10_000 -# How long ``close`` waits for the server to end the stream, so its own session cleanup runs. +# How long ``close`` waits for the server to end the stream and release the session. _CLOSE_TIMEOUT_SEC = 5.0 -# A path no handler serves, so asking why a channel is down never opens a session on a server that -# turns out to be up after all. +# A path no handler serves: a probe of it opens no session on a server that is up. _PROBE_PATH = f'/{SERVICE}/ChannelProbe' -# The largest slice of one connect attempt's budget the refusal probe may spend. A target that black-holes -# connection attempts answers neither, so both waits must fit inside the caller's ``open_timeout``. +# The largest share of one connect attempt's budget the refusal probe may spend. Both waits fit inside +# the caller's ``open_timeout``: a target that drops every connect answers neither. _REFUSAL_PROBE_SEC = 1.0 -# What gRPC's status details call an edge no client can use: a certificate its roots do not cover, -# and a front that selects no HTTP/2 over ALPN. +# The status details of an edge no client can use: a certificate the roots do not cover, and a front +# that selects no HTTP/2 over ALPN. UNUSABLE_EDGE = ('CERTIFICATE_VERIFY_FAILED', 'missing selected ALPN property') def edge_is_unusable(details: str) -> bool: - """Whether a gRPC status blames the TLS edge's own configuration rather than a cold backend.""" + """Whether a gRPC status blames the TLS edge's own configuration.""" return any(marker in details for marker in UNUSABLE_EDGE) @@ -64,7 +62,7 @@ def _client_options() -> list[tuple[str, int]]: *_MESSAGE_SIZE_OPTIONS, ('grpc.keepalive_time_ms', _PING_EVERY_MS), ('grpc.keepalive_timeout_ms', _PING_ANSWER_TIMEOUT_MS), - # Left to itself gRPC sends two pings without data, five minutes apart. + # The gRPC default sends two pings without data, five minutes apart. ('grpc.http2.max_pings_without_data', 0), ('grpc.http2.min_time_between_pings_ms', _PING_EVERY_MS), ] @@ -73,22 +71,21 @@ def _client_options() -> list[tuple[str, int]]: def _channel(target: str, secure: bool) -> grpc.Channel: options = _client_options() if secure: - # No roots named, so the channel verifies the edge against the system's own. + # No roots named: the channel verifies the edge against the system's own roots. return grpc.secure_channel(target, grpc.ssl_channel_credentials(), options=options) return grpc.insecure_channel(target, options=options) def _probe_share(open_timeout: float) -> float: - """What one connect attempt gives the refusal probe, leaving the readiness wait the rest. + """The share of one connect attempt the refusal probe gets; the readiness wait gets the rest. - Half at most, so an ``open_timeout`` under ``_REFUSAL_PROBE_SEC`` still waits for a healthy server - instead of going straight to asking why it is down. + Half at most: an ``open_timeout`` under ``_REFUSAL_PROBE_SEC`` still waits for a healthy server. """ return min(_REFUSAL_PROBE_SEC, open_timeout / 2) def _connect_refusal(channel: grpc.Channel, timeout: float) -> grpc.RpcError | None: - """What gRPC says stopped the channel coming up. Its readiness future carries only that it did not.""" + """What gRPC says stopped the channel. The readiness future says only that the channel is not ready.""" probe = channel.stream_stream(_PROBE_PATH, request_serializer=None, response_deserializer=None) try: next(probe(iter(()), timeout=timeout)) @@ -102,11 +99,8 @@ def _connect_refusal(channel: grpc.Channel, timeout: float) -> grpc.RpcError | N class GrpcClientConnection: """A client's end of one gRPC session. - A reader thread drains the response stream into a queue, because the stream itself has no - per-message timeout and ``recv`` needs one. - - ``secure`` dials over TLS, which is the shape an authenticated endpoint takes: a TLS edge in front - of the server's plaintext gRPC port. + A reader thread drains the response stream into a queue: the stream has no per-message timeout, and + ``recv`` needs one. ``secure`` dials over TLS, to a TLS edge in front of the server's plaintext port. """ def __init__( @@ -125,16 +119,15 @@ def __init__( grpc.channel_ready_future(self._channel).result(timeout=open_timeout - _probe_share(open_timeout)) except grpc.FutureTimeoutError: refusal = _connect_refusal(self._channel, timeout=max(0.0, deadline - time.monotonic())) - # The probe path is served by no handler, so an UNIMPLEMENTED means the edge carried the call: - # the channel is up, and the readiness wait was short rather than the server absent. + # An ``UNIMPLEMENTED`` from the probe path means the channel is up: the readiness wait was too short. if refusal is None or refusal.code() is not grpc.StatusCode.UNIMPLEMENTED: self._channel.close() - # An edge that refuses every client is permanent, so raise what gRPC blamed rather than a - # timeout: the connect loop reads the status and stops instead of retrying its deadline out. + # An edge that refuses every client is permanent, and the connect loop retries a ``TimeoutError`` + # to its deadline. if refusal is not None and edge_is_unusable(refusal.details() or ''): raise refusal from None raise TimeoutError(f'gRPC channel to {target} is not ready within {open_timeout}s') from None - # gRPC metadata keys are lower case, and they are the same header names the websocket wire sends. + # gRPC metadata keys are lower case; the header names are the websocket wire's. metadata = tuple((key.lower(), value) for key, value in (headers or {}).items()) + ( (SESSION_PATH_HEADER, session_path), (SESSION_QUERY_HEADER, query), @@ -154,7 +147,7 @@ def _requests(self): yield message def _read(self) -> None: - """Drain the response stream into the inbox, ending it with what stopped it.""" + """Drain the response stream into the inbox, and end the inbox with what stopped the stream.""" try: for message in self._responses: self._inbox.put(message) @@ -165,15 +158,15 @@ def _read(self) -> None: self._responses.cancel() def send(self, message: bytes) -> None: - # A write past either of these sits in the outbox while ``recv`` waits out a whole inference - # timeout on an inbox nothing refills: gRPC has stopped reading the request iterator. + # gRPC stops reading the request iterator once the stream ends, and a write then sits in the outbox + # until ``recv`` times out. if self._closed or self._ended: raise wire.PeerDisconnected(f'The session on {self._target} has ended') self._outbox.put(message) def recv(self, timeout: float | None = None) -> bytes: - # A closed session's inbox may hold a reply that arrived during ``close``, which would pair one - # observation's actions with the next observation's state. + # A reply that arrived during ``close`` sits in the inbox, and would pair one observation's actions + # with the next observation. if self._closed: raise wire.PeerDisconnected(f'The session on {self._target} is closed') try: @@ -181,7 +174,7 @@ def recv(self, timeout: float | None = None) -> bytes: except queue.Empty: raise TimeoutError(f'No message from {self._target} within {timeout}s') from None if isinstance(answer, BaseException): - # What ended the stream is queued once, so the caller learns why before writes are refused. + # What ended the stream is queued once. The caller reads it before ``send`` refuses a write. self._ended = True raise answer return answer @@ -191,13 +184,13 @@ def close(self) -> str: return 'already closed' self._closed = True self._outbox.put(None) - # The half-close ends the server's session, and the server then ends the stream. Waiting for - # that lets the server release its model slot; closing the channel now would cut it short. + # The half-close ends the server's session, and the server then ends the stream. A channel closed + # before that cuts the server's cleanup short. self._reader.join(timeout=_CLOSE_TIMEOUT_SEC) server_ended_stream = not self._reader.is_alive() self._channel.close() - # The websocket wire reads the same two facts off a close code. A stream the server never ended - # means it still holds this session, so the next one's handshake waits on a slot nobody released. + # A stream the server never ended means the server still holds this session, and the next session's + # handshake waits on its slot. return ( f'peer had ended the stream {self._ended}, ' f'server ended it within {_CLOSE_TIMEOUT_SEC}s {server_ended_stream}' @@ -205,7 +198,7 @@ def close(self) -> str: def model_id_of(session_path: str) -> str | None: - """The model a session path names, or ``None`` where it names the model the server pinned.""" + """The model a session path names, or ``None`` for the model the server pinned.""" prefix = f'{wire.SESSION_PATH}/' if session_path == wire.SESSION_PATH: return None @@ -305,8 +298,8 @@ async def serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerCo try: await session(conn, model_id_of(conn.session_path)) except Exception as e: - # The session itself reports what it can over the stream; anything reaching here happened - # before or beyond that, so the client learns of it from the status alone. + # The session reports its own errors over the stream. One that reaches here reaches the client + # as the status alone. logger.error(f'Failed gRPC session: {e}', exc_info=True) await context.abort(grpc.StatusCode.INTERNAL, str(e)) @@ -315,8 +308,7 @@ async def serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerCo server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler(SERVICE, {METHOD: handler}),)) bound = server.add_insecure_port(_bind_target(self._host, self._port)) if bound == 0: - # gRPC reports a refused bind by returning port 0, so a server left to start here would - # accept nothing and say nothing. + # gRPC reports a refused bind as port 0, and a server started on it accepts nothing and says nothing. raise OSError(f'gRPC could not bind {_bind_target(self._host, self._port)}') self._server = server self._endpoint = wire.Endpoint(self._host, bound) diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 2c12e8f7f..5f0506c7d 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -191,9 +191,8 @@ class PolicyServer: The session flow is: accept → session params → resolve → load via manager → remote-half wrap → reset → inference loop - ``serve`` takes the wires sessions arrive on, and names none of them (see - ``positronic.offboard.wire``, which states what a wire owes a server). ``api`` holds this server's - own HTTP routes, which a wire that speaks HTTP serves beside its sessions. + ``serve`` takes the wires sessions arrive on (``positronic.offboard.wire``). ``api`` holds this + server's own HTTP routes, for a wire that speaks HTTP. On startup (before accepting connections): resolve(None) → load. @@ -231,7 +230,7 @@ def __init__( self._infer_lock = asyncio.Lock() self._default_id: str | None = None - # Set while ``serve`` runs, so ``shutdown`` can reach its loop from another thread. + # Set while ``serve`` runs; ``shutdown`` reaches the loop from another thread. self._loop: asyncio.AbstractEventLoop | None = None self._stop: asyncio.Event | None = None @@ -260,7 +259,7 @@ def _token_matches(self, authorization: str | None) -> bool: return hmac.compare_digest(authorization.encode(), bearer(self._auth_token).encode()) def _authorized(self, headers: Mapping[str, str]) -> bool: - """Whether session headers carry the bearer token this server gates on. Every wire asks this.""" + """Whether the session headers carry the bearer token this server gates on.""" return self._token_matches(headers.get(AUTH_HEADER.lower())) def _require_http_auth(self, authorization: str | None = Header(default=None, alias=AUTH_HEADER)) -> None: @@ -300,8 +299,8 @@ async def _answer_observations(self, conn: wire.ServerConnection, session: Sessi # The server's clock is not the rig's. actions = await asyncio.to_thread(session, raw_obs, time.time_ns()) except asyncio.CancelledError: - # Cancelling this await does not stop the worker, so the session close runs beside - # a live inference. Logged to give a later wrong answer a cause. + # A cancelled await does not stop the worker, and the session close runs beside a live + # inference. The log gives a later wrong answer a cause. logger.error('Cancelled mid-inference: the worker is still in the backend') raise await conn.send(serialise({protocol.RESULT: actions})) @@ -406,18 +405,14 @@ async def _idle_watchdog(self): def serve(self, wires: Sequence[wire.Wire], on_ready: Callable[[], None] | None = None): """Serve sessions on every wire in ``wires``, until one of them ends or the server goes idle. - Every wire shares this server's model slot and inference lock, so a session is served the same - whichever one carried it. - - ``on_ready`` runs on the server's own loop once every wire has bound, which is where a caller - that asked for port 0 reads back the port each wire took. + Every wire shares this server's model slot and inference lock. ``on_ready`` runs on the server's + own loop once every wire has bound; a caller that asked for port 0 reads the port there. """ async def _run(): self._loop, self._stop = asyncio.get_running_loop(), asyncio.Event() await self._startup() - # A wire binds when it starts, so one that started is stopped even where a later one cannot - # bind and nothing ever serves. + # A wire binds when it starts, and a started wire is stopped even when a later one cannot bind. started: list[wire.Wire] = [] serving: list[asyncio.Task] = [] ending: list[asyncio.Task] = [] @@ -429,8 +424,7 @@ async def _run(): if on_ready is not None: on_ready() serving = [asyncio.create_task(w.serve()) for w in started] - # What ends the server, beside a wire ending on its own: a caller's ``shutdown``, and - # the idle timeout. + # What else ends the server: a caller's ``shutdown``, and the idle timeout. ending = [asyncio.create_task(self._stop.wait())] if self.idle_timeout_min and self.idle_timeout_min > 0: ending.append(asyncio.create_task(self._idle_watchdog())) @@ -444,11 +438,11 @@ async def _run(): outcomes = await asyncio.gather(*serving, return_exceptions=True) failed = [(w, e) for w, e in zip(started, outcomes, strict=True) if isinstance(e, Exception)] - # Only one failure can reach the caller, so the rest are reported here or nowhere. + # Only one failure can raise; the rest are logged here or nowhere. for w, error in failed[1:]: logger.error(f'{type(w).__name__} also failed: {error}', exc_info=error) if failed: - # A wire that ended on an error raises, rather than reading as the shutdown this waits for. + # A wire that ended on an error raises; a silent return reads as a shutdown. raise failed[0][1] try: @@ -485,9 +479,6 @@ def serve( The bearer token gating the server comes from ``AUTH_TOKEN_ENV`` rather than a flag, which would put a secret in the process arguments; unset serves open. - - This is where the flags name the wires, and the only place either wire is named: ``PolicyServer`` - takes whatever list it is handed. """ server = PolicyServer( pipeline, diff --git a/positronic/offboard/stub.py b/positronic/offboard/stub.py index faf2372f4..7e3f7617e 100644 --- a/positronic/offboard/stub.py +++ b/positronic/offboard/stub.py @@ -1,7 +1,7 @@ -"""A server with no model: every inference answers the same chunk, so a session measures the wire alone. +"""A server with no model: every inference answers the same chunk, and a session measures the wire alone. -``delay_sec`` is a session param, so ``?delay_sec=120`` holds one inference open for two minutes — -what a cold model does to a connection, with nothing on the wire meanwhile. +``delay_sec`` is a session param: ``?delay_sec=120`` holds one inference open for two minutes, with +nothing on the wire meanwhile. """ import time @@ -18,8 +18,7 @@ from positronic.policy.layers import ChunkedSchedule from positronic.policy.spec import Pipeline, PolicySource, remote -# One action, at the start of the chunk. A served session must answer a trajectory, and this is the -# smallest one that is. +# The smallest trajectory a served session can answer: one action, at the start of the chunk. CHUNK = [{keys.ACTION_TIMESTAMP: 0.0}] @@ -50,8 +49,9 @@ def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] class Delay(Layer): - """Holds every answer for ``delay_sec``, standing in for a model slow enough to outlast a front's - idle close. A session param may tune the pipeline around the model source, never the source itself. + """Holds every answer for ``delay_sec``, in place of a model slow enough to outlast a front's idle close. + + A session param can tune a layer and cannot change the model source. """ def __init__(self, delay_sec: float = 0.0): @@ -61,8 +61,8 @@ def make_session(self, inner: Session) -> Session: return DelayedSession(inner, self._delay_sec) -# One instance for the process: a server compares the source a session param rebuilds against the one -# it launched with, and two sources are equal only when they hold the same policy. +# One instance for the process. A server refuses a session param that rebuilds a different source, and +# two ``PolicySource``s are equal only over one policy object. POLICY = StubPolicy() diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index e4e23bedd..6d2865806 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -29,8 +29,8 @@ class Served(NamedTuple): def start_server() -> Generator[StartServer, None, None]: """Factory serving pipelines on daemon threads; every started server is stopped and joined at teardown. - Each wire asks for port 0 and holds what it binds, so servers started in parallel never draw the same - port. ``grpc=True`` serves the gRPC wire beside the websocket one. + Each wire asks for port 0, and servers started in parallel never draw the same port. ``grpc=True`` + serves the gRPC wire beside the websocket one. """ running: list[tuple[PolicyServer, threading.Thread]] = [] diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index 5b9347a4e..327a85bf2 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -1,4 +1,4 @@ -"""The gRPC wire: one session runs over it exactly as it runs over the websocket.""" +"""The gRPC wire: a session runs over it as it runs over the websocket.""" import asyncio import datetime @@ -40,7 +40,7 @@ def grpc_url(served: Served, path: str = '') -> str: @pytest.fixture def both_wires(start_server: StartServer, make_mock_policy) -> tuple[Served, MagicMock]: - """A server offering both wires over one policy.""" + """A server that offers both wires over one policy.""" policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) served = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True) return served, policy @@ -76,7 +76,6 @@ def test_both_wires_answer_one_observation_alike(both_wires): def test_each_wire_names_its_own_port_in_the_meta(both_wires): - """The two wires bind two ports, and a session reads back the one that carried it.""" served, _policy = both_wires over_ws = InferenceClient(f'{served.host}:{served.port}').new_session() over_grpc = InferenceClient(grpc_url(served)).new_session() @@ -89,9 +88,7 @@ def test_each_wire_names_its_own_port_in_the_meta(both_wires): def test_both_wires_report_what_their_close_saw(both_wires, caplog): - """A close the server answered has to read differently from one it never saw, whichever wire carried - the session. The second leaves the server holding the slot, and the next session's handshake waits on - it, so each wire reports the distinction in the terms its own protocol offers.""" + """The server holds the slot of a session whose close it never saw, and the next handshake waits on it.""" served, _policy = both_wires over_ws = InferenceClient(f'{served.host}:{served.port}').new_session() over_grpc = InferenceClient(grpc_url(served)).new_session() @@ -106,7 +103,7 @@ def test_both_wires_report_what_their_close_saw(both_wires, caplog): def test_closing_a_session_ends_it_on_the_server(both_wires): - """``close`` half-closes the stream and waits, so the server releases the session before it returns.""" + """``close`` returns after the server has released the session.""" served, _policy = both_wires session = InferenceClient(grpc_url(served)).new_session() assert served.server._active_sessions == 1 @@ -158,8 +155,6 @@ def test_the_query_carries_the_session_params(start_server, make_mock_policy): session = InferenceClient(grpc_url(served, f'{wire.SESSION_PATH}?offsets=[-0.5, 0.0]')).new_session() try: stack = session.metadata[offboard_keys.LOCAL_STACK][SEQ] - # `args` and the layer's own constructor keyword are the spec grammar's, written wherever a - # layer renders itself; this reader spells them as the wire carries them. assert stack[0]['args']['offsets_sec'] == [-0.5, 0.0] finally: session.close() @@ -182,8 +177,8 @@ def test_the_grpc_wire_gates_on_the_bearer_token(authed_server): @pytest.mark.parametrize('header', [None, bearer('wrong'), _TOKEN]) def test_the_grpc_wire_refuses_a_session_without_the_token(authed_server, header, monkeypatch): - # A refused credential and a cold backend answer alike, so the client spends attempts on it; one - # is enough to see the refusal. + # A refused credential answers like a cold backend, and the client retries it; one attempt shows the + # refusal. monkeypatch.setattr(_ConnectRetries, 'MAX_FORBIDDEN_ATTEMPTS', 1) headers = None if header is None else {AUTH_HEADER: header} with pytest.raises(grpc.RpcError) as refused: @@ -191,8 +186,8 @@ def test_the_grpc_wire_refuses_a_session_without_the_token(authed_server, header assert refused.value.code() is grpc.StatusCode.PERMISSION_DENIED -# The edge answers on one address, not on both families a name resolves to: gRPC reports the last -# address it failed on, so a second leg refusing the connection would hide what the first blamed. +# An address, and no name that resolves to two families: gRPC reports the last address it failed on, +# and a refused second family would hide what the first blamed. EDGE_HOST = '127.0.0.1' @@ -223,8 +218,8 @@ async def _copy(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> N while chunk := await reader.read(65536): writer.write(chunk) await writer.drain() - # Whichever end closes first leaves the other half of the pair writing into a dead socket, which - # is how a session ends. Anything else is the edge itself failing and belongs in the test's face. + # The end that closes first leaves the other half of the pair writing into a dead socket, which is how + # a session ends. Any other error is the edge's own, and fails the test. except (ConnectionResetError, BrokenPipeError): pass finally: @@ -233,12 +228,11 @@ async def _copy(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> N @pytest.fixture def tls_edge() -> Generator[Callable[[str, int], tuple[int, bytes]], None, None]: - """Starts a TLS front over a plaintext gRPC port, the shape an authenticated endpoint takes. + """Starts a TLS front over a plaintext gRPC port, as an authenticated endpoint is served. - It terminates TLS, selects HTTP/2 over ALPN and copies the bytes on, so the client and the server - speak one h2 connection end to end and the server holds no certificate. Answers the front's own - port and the root to verify it against. ``alpn=False`` selects no protocol at all, which is what - a front fronting a raw TCP port does. + The front terminates TLS, selects HTTP/2 over ALPN and copies the bytes on. It answers its own port + and the root to verify it against. ``alpn=False`` selects no protocol, as a front over a raw TCP + port does. """ stops: list[tuple[asyncio.AbstractEventLoop, asyncio.Event]] = [] @@ -284,7 +278,7 @@ def _trust_only(monkeypatch, root: bytes) -> None: @pytest.fixture def edged(tls_edge, monkeypatch) -> Callable[[Served], str]: - """The ``grpcs://`` URL of a server reached through a TLS edge, with the client trusting its root.""" + """The ``grpcs://`` URL of a server reached through a TLS edge; the client trusts the edge's root.""" def url(served: Served) -> str: port, root = tls_edge(served.host, served.grpc_port) @@ -366,14 +360,13 @@ def test_a_path_outside_the_session_route_is_refused(): def test_a_port_that_never_answers_is_named_at_the_deadline(): - """Nothing listens on port 1, so the channel never becomes ready and the connect deadline passes.""" + """Nothing listens on port 1; the channel never becomes ready.""" client = InferenceClient('grpc://localhost:1', open_timeout=0.2, connect_deadline=0.0) with pytest.raises(TimeoutError, match='grpc://localhost:1'): client.new_session() def test_an_open_timeout_under_the_probe_budget_still_opens(both_wires): - """The refusal probe takes a share of the budget, so a healthy server answers a short one.""" served, _policy = both_wires budget = grpc_wire._REFUSAL_PROBE_SEC / 2 session = InferenceClient(grpc_url(served), open_timeout=budget, connect_deadline=0.0).new_session() @@ -384,7 +377,7 @@ def test_an_open_timeout_under_the_probe_budget_still_opens(both_wires): def test_an_ipv6_host_binds_in_brackets(start_server: StartServer, make_mock_policy): - """gRPC's target syntax brackets an IPv6 literal, so a bare '::1' would bind ':::' and fail.""" + """A bare '::1' binds as ':::', which gRPC refuses.""" assert grpc_wire._bind_target('::', 9000) == '[::]:9000' assert grpc_wire._bind_target('0.0.0.0', 9000) == '0.0.0.0:9000' @@ -398,8 +391,8 @@ def test_an_ipv6_host_binds_in_brackets(start_server: StartServer, make_mock_pol def test_a_refused_handshake_closes_the_connection(both_wires): - """A model the source does not know is refused in a protocol frame, past the transport handlers, - and the gRPC connection behind it holds a reader thread until something closes it.""" + """A refusal in a protocol frame raises past the transport handlers, and the connection holds a reader + thread until it is closed.""" client = InferenceClient(grpc_url(both_wires[0], f'{wire.SESSION_PATH}/unknown-model')) opened = [] connect = client._connect @@ -421,7 +414,7 @@ def record(): @pytest.fixture def chatty_client(monkeypatch) -> None: - """Pings often enough that a silence measured in seconds stands in for one measured in minutes.""" + """Pings every 500 ms, and a silence of seconds stands in for one of minutes.""" monkeypatch.setattr(grpc_wire, '_PING_EVERY_MS', 500) @@ -435,7 +428,7 @@ def _silent_then_infer(served: Served) -> list[dict]: def test_a_session_answers_after_a_silence_no_frame_crossed(both_wires, chatty_client): - """One inference can outlast a front's idle close, so the wire's own pings hold the stream open.""" + """The wire's own pings hold the stream open through an inference that outlasts a front's idle close.""" assert _silent_then_infer(both_wires[0]) == [{'action': [1, 2, 3]}] @@ -451,7 +444,7 @@ def test_a_server_on_the_grpc_ping_defaults_kills_the_silent_session( def _surfaces_at_once(url: str, blamed: str) -> None: - """Assert a connect to ``url`` fails naming ``blamed``, without spending its retry deadline.""" + """Assert that a connect to ``url`` fails, names ``blamed``, and spends no retry deadline.""" client = InferenceClient(url, open_timeout=2.0, connect_deadline=20.0) started = time.monotonic() with pytest.raises(grpc.RpcError, match=blamed): @@ -460,7 +453,6 @@ def _surfaces_at_once(url: str, blamed: str) -> None: def test_a_certificate_the_client_cannot_verify_is_not_retried(both_wires, tls_edge, monkeypatch): - """A root that does not cover the edge is permanent, so it surfaces on the first attempt.""" port, _root = tls_edge(both_wires[0].host, both_wires[0].grpc_port) unrelated, _key = _self_signed(EDGE_HOST) _trust_only(monkeypatch, unrelated) @@ -468,7 +460,7 @@ def test_a_certificate_the_client_cannot_verify_is_not_retried(both_wires, tls_e def test_an_edge_that_selects_no_alpn_is_not_retried(both_wires, tls_edge, monkeypatch): - """A front fronting a raw TCP port terminates TLS and names no protocol, which gRPC cannot use.""" + """A front over a raw TCP port terminates TLS and names no ALPN protocol, and gRPC refuses it.""" port, root = tls_edge(both_wires[0].host, both_wires[0].grpc_port, alpn=False) _trust_only(monkeypatch, root) _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE[1]) @@ -481,13 +473,13 @@ def test_a_timed_out_session_refuses_the_next_inference(both_wires): session = InferenceClient(grpc_url(served), infer_timeout=0.2).new_session() with pytest.raises(TimeoutError): session.infer({'image': 'test'}) - # Without the guard this answers the first observation's actions, against the second's state. + # The late answer is the first observation's actions. with pytest.raises(wire.PeerDisconnected): session.infer({'image': 'test'}) def test_a_connection_refuses_to_send_once_the_server_ends_the_stream(both_wires): - """gRPC stops reading the request iterator then, so a write would wait out a whole timeout.""" + """gRPC stops reading the request iterator, and a write waits out a whole timeout.""" served, _policy = both_wires conn = grpc_wire.GrpcClientConnection(f'{served.host}:{served.grpc_port}', f'{wire.SESSION_PATH}/unknown-model', '') try: @@ -497,7 +489,7 @@ def test_a_connection_refuses_to_send_once_the_server_ends_the_stream(both_wires conn.recv(timeout=10.0) with pytest.raises(wire.PeerDisconnected): conn.send(b'an observation the stream can no longer carry') - # The peer ended the stream, so this close reads differently from one the server answered. + # The close report says the peer ended the stream. assert 'peer had ended the stream True' in conn.close() finally: conn.close() diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index f1498312f..35e522adc 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -51,12 +51,12 @@ def meta(self, model_id: str) -> dict[str, Any]: return {'type': 'stub'} -# Short enough to keep this test quick, long enough that a loaded box still reaches the first poll. +# Short enough for a quick test, long enough that a loaded box reaches the first poll. _A_MOMENT_IDLE = 0.5 class _FailingWire(wire.Wire): - """Serves for ``after`` seconds, then falls over.""" + """Serves for ``after`` seconds, then raises.""" def __init__(self, after: float): self._after = after @@ -95,8 +95,7 @@ async def stop(self) -> None: def test_a_wire_that_cannot_bind_stops_the_ones_that_did(make_mock_policy): - """A wire binds when it starts, so a startup that gives up does not leave an earlier one holding a - port against a server nothing is serving.""" + """A wire binds when it starts, and a startup that gives up frees the port an earlier wire took.""" server = PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {}))) bound = _FailingWire(_A_MOMENT_IDLE) with pytest.raises(OSError, match='that port is taken'): @@ -105,13 +104,12 @@ def test_a_wire_that_cannot_bind_stops_the_ones_that_did(make_mock_policy): def test_a_websocket_wire_releases_its_port_when_startup_rolls_back(make_mock_policy): - """A ``WebsocketWire`` binds a real socket when it starts, so a startup that rolls back frees the - port it took, not only the stub wires that own no socket.""" + """A ``WebsocketWire`` binds a real socket when it starts, and a startup that rolls back frees it.""" server = PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {}))) bound = websocket_wire.WebsocketWire('localhost', 0, server.api) with pytest.raises(OSError, match='that port is taken'): server.serve([bound, _UnbindableWire()]) - # A leaked listener would still hold the port, so binding a fresh socket to it would raise. + # A leaked listener holds the port, and a fresh bind to it raises. websocket_wire._listening_socket('localhost', bound.endpoint.port).close() @@ -125,7 +123,7 @@ def test_a_failing_wire_reaches_the_caller_and_the_rest_are_logged(make_mock_pol def test_an_idle_server_stops_itself(make_mock_policy): - """The idle watchdog ends every wire it is serving on, so ``serve`` returns with nobody asking it to.""" + """The idle watchdog ends every wire, and ``serve`` returns with no ``shutdown`` call.""" server = PolicyServer( ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {})), idle_timeout_min=_A_MOMENT_IDLE / 60 ) diff --git a/positronic/offboard/websocket_wire.py b/positronic/offboard/websocket_wire.py index ed7e5b780..d64b7dd18 100644 --- a/positronic/offboard/websocket_wire.py +++ b/positronic/offboard/websocket_wire.py @@ -64,32 +64,24 @@ async def refuse(self, reason: str) -> None: def _listening_socket(host: str, port: int) -> socket.socket: - """A socket bound on ``host``, where a ``port`` of 0 takes any free one. - - The family comes from ``host`` itself, so an IPv6 host binds an IPv6 socket. Binding here rather - than inside uvicorn names the port before the wire serves, and holds it from then on. - """ + """A listening socket bound on ``host``, where a ``port`` of 0 takes any free one.""" family, kind, proto, _canonical, address = socket.getaddrinfo( host, port, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE )[0] sock = socket.socket(family, kind, proto) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.bind(address) - # Listening here rather than at the first accept makes the port answer from the moment ``start`` - # returns: the kernel queues a connect that beats the serving loop to it. + # The port answers from the moment ``start`` returns: the kernel queues a connect that arrives before + # the serving loop runs. sock.listen() return sock class WebsocketWire(wire.Wire): - """The websocket wire: a session upgrades on ``wire.SESSION_PATH``, and ``api`` answers on the same port. - - One uvicorn serves both, so the routes a client reads a model catalogue from sit on the endpoint it - opens sessions on. - """ + """The websocket wire: a session upgrades on ``wire.SESSION_PATH``, and ``api`` answers on the same port.""" - # How long ``stop`` lets an open session finish before it cuts the connection. Left to itself - # uvicorn waits for ever, so a session mid-inference would hold the whole server open. + # How long ``stop`` lets an open session finish before it cuts the connection. The uvicorn default + # waits for ever, and a session mid-inference holds the whole server open. STOP_GRACE_SEC = 2 def __init__(self, host: str, port: int, api: APIRouter): @@ -124,13 +116,12 @@ async def start(self, session: wire.SessionHandler, authorized: wire.Authorized) def _route_sessions(self, app: FastAPI, session: wire.SessionHandler, authorized: wire.Authorized) -> None: async def require_auth(websocket: WebSocket) -> None: - """Refuses before ``accept()``, so an unauthorized peer never reaches the session handshake.""" + """Refuse before ``accept()``. An unauthorized peer never reaches the session handshake.""" if not authorized(websocket.headers): raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION) async def serve_pinned_model(websocket: WebSocket) -> None: - """Serves the model the server pinned. Naming a model is the path's job, so every query param - here is a pipeline override.""" + """Serve the model the server pinned. The path names a model; every query param is a pipeline override.""" await websocket.accept() await session(WebsocketServerConnection(websocket, self.endpoint), None) @@ -140,8 +131,8 @@ async def serve_named_model(websocket: WebSocket, model_id: str) -> None: auth = [Depends(require_auth)] app.websocket(wire.SESSION_PATH, dependencies=auth)(serve_pinned_model) - # ``:path`` so an id that is itself a path (a HuggingFace repo, say) opens under the name the - # model catalogue advertises. + # ``:path``: a model id can itself be a path (a HuggingFace repo), and opens under the name the + # catalogue advertises. app.websocket(f'{wire.SESSION_PATH}/{{model_id:path}}', dependencies=auth)(serve_named_model) async def serve(self) -> None: @@ -152,7 +143,7 @@ async def serve(self) -> None: async def stop(self) -> None: if self._server is not None: self._server.should_exit = True - # uvicorn releases the socket as it shuts down. A startup that rolls back stops a wire that - # bound but never served, so uvicorn never runs; close the socket here to free the port. + # uvicorn releases the socket when it shuts down. A wire that bound but never served has no + # uvicorn to release it. if self._socket is not None and not self._served: self._socket.close() diff --git a/positronic/offboard/wire.py b/positronic/offboard/wire.py index a62fcd005..f40392512 100644 --- a/positronic/offboard/wire.py +++ b/positronic/offboard/wire.py @@ -1,7 +1,7 @@ -"""The transports a session runs over, the two ends of an open one, and the server's end of a wire. +"""The transports a session runs over, and the two ends of one open session. -A wire carries the ``protocol`` frames as opaque bytes and reads none of them, so the handshake and -the inference loop read the same over every wire. ``websocket_wire`` and ``grpc_wire`` hold the two. +A wire carries the ``protocol`` frames as opaque bytes and reads none of them. ``websocket_wire`` and +``grpc_wire`` hold the two wires. """ import abc @@ -10,13 +10,11 @@ from starlette.datastructures import QueryParams -# The route a session opens on. The websocket wire puts it in the URL; the gRPC wire names it in the -# session metadata, so both wires address a model the same way. +# The route a session opens on: in the URL on the websocket wire, in the session metadata on the gRPC wire. SESSION_PATH = '/api/v1/session' -# The largest frame a session may carry, on either wire. An observation is a stack of camera frames, -# so the gRPC default of 4 MiB refuses one; uvicorn's own default happens to be this, and passing it -# explicitly is what keeps the two wires equal when that default moves. +# The largest frame a session may carry, on either wire. An observation is a stack of camera frames, and +# the gRPC default of 4 MiB refuses one. MAX_MESSAGE_BYTES = 16 * 1024 * 1024 @@ -43,9 +41,8 @@ def recv(self, timeout: float | None = None) -> bytes: def close(self) -> str: """Close this end, and report what the wire saw, for the log. - A peer that answered the close leaves a different trace from one that had already gone while the - server still held the session, and the second is what strands the next session's handshake. Only - the wire can tell the two apart, and each says it in its own terms. + The report says whether the peer answered the close, in the wire's own terms. A server that still + holds a session strands the next session's handshake, and only the wire can see that. """ ... @@ -77,11 +74,11 @@ async def receive(self) -> bytes: @abc.abstractmethod async def refuse(self, reason: str) -> None: - """End a session the server cannot serve, telling the client why.""" + """End a session the server cannot serve, and tell the client why.""" # What a wire hands the server for each session it accepts: the connection, and the model the route -# names, which is ``None`` where the route names the model the server pinned. +# names, or ``None`` for the model the server pinned. SessionHandler = Callable[[ServerConnection, str | None], Awaitable[None]] # Whether the session headers carry a credential the server accepts. Header names are lower case. @@ -92,13 +89,13 @@ class Wire(abc.ABC): """One transport that sessions arrive on. A wire reads its own route for the model a session names, and refuses an unauthorized peer before - the session opens. So a server hands every wire one ``SessionHandler`` and serves them all alike. + the session opens. """ @property @abc.abstractmethod def endpoint(self) -> Endpoint: - """Where this wire serves. The port is bound, and so known, once ``start`` returns.""" + """Where this wire serves. The port is known once ``start`` returns.""" @abc.abstractmethod async def start(self, session: SessionHandler, authorized: Authorized) -> None: diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 79c0f9f8d..bb2a4129f 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -6,9 +6,9 @@ # itself takes ~10-15 min more to finish uv sync and load the model into GPU # memory after the URL appears. # -# Both wires are served: the websocket on 8000, and gRPC on whatever `--grpc_port` -# names. The gRPC port is declared as an ordinary HTTP port and never `/tcp` — the -# offboard README says what each front does to a gRPC session. +# Both wires are served: the websocket on 8000, and gRPC on the port `--grpc_port` +# names. The gRPC port is declared as an ordinary HTTP port; a `/tcp` port gets a +# front gRPC refuses. The offboard README says what each front does to a session. # # That URL carries the id of a tunnel created with the endpoint, so it cannot be # chosen or known in advance, and a delete plus re-create earns a new one even @@ -94,11 +94,10 @@ case " $* " in *) set -- "$@" "--idle_timeout_min=${NEBIUS_IDLE_TIMEOUT_MIN:-20}" ;; esac -# The port the websocket wire listens on, named once: the create declares it and the poll below -# selects the managed URL that fronts it. +# The websocket port: the create declares it, and the poll below selects the managed URL that fronts it. WS_PORT=8000 -# The endpoint exposes the port the server listens on, so a caller's own --grpc_port decides both. +# The endpoint exposes the port the server listens on; a caller's own --grpc_port names both. ARGS=" $* " case "$ARGS" in *" --grpc_port="*) GRPC_PORT=${ARGS#*--grpc_port=}; GRPC_PORT=${GRPC_PORT%% *} ;; @@ -144,9 +143,9 @@ echo "Waiting for the managed HTTPS URL (typically <1 min)..." URL="" for i in $(seq 1 30); do - # Each managed URL names the container port it fronts, so the two wires are told apart by that - # prefix. This field also carries bare `IP:port` entries, which serve no TLS and would put the - # bearer token on the wire in cleartext — take the https:// ones, and fail rather than fall back. + # Each managed URL names the container port it fronts, and that prefix tells the two wires apart. + # This field also carries bare `IP:port` entries, which serve no TLS and would put the bearer + # token on the wire in cleartext: take the https:// ones, and fail with no fallback. URL=$(nebius ai endpoint get "$ID" --format json 2>/dev/null \ | jq -r "[.status.public_endpoints[]? | select(startswith(\"https://port${WS_PORT}-\"))] | first // empty") if [ -n "$URL" ]; then break; fi From 816cdc0afa00c1623791abfabb5cdbb2e6214c07 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 09:54:24 +0000 Subject: [PATCH 39/46] Cut two client comments to two lines and give `serve.sh` a usage block The inference guide states the GPU split as a fact. Ticket: Positronic-Robotics/internal#1191 #refs --- docs/inference.md | 2 +- positronic/offboard/client.py | 10 ++++------ workflows/nebius/serve.sh | 4 ++++ 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/docs/inference.md b/docs/inference.md index 70cd2ec80..db02c56fe 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -4,7 +4,7 @@ Deploy trained policies for evaluation and production use. Positronic supports l ## Inference with Remote Server -Positronic's unified session protocol connects any hardware to any model (LeRobot, GR00T, OpenPI); the same frames cross either wire, a websocket or gRPC. The key benefit is running heavy models on powerful GPU hardware (OpenPI needs ~62GB, GR00T ~8GB) separate from the robot/simulator machine. +Positronic's unified session protocol connects any hardware to any model (LeRobot, GR00T, OpenPI); the same frames cross either wire, a websocket or gRPC. A heavy model (OpenPI needs ~62GB, GR00T ~8GB) runs on GPU hardware separate from the robot/simulator machine. Each server carries a full **policy pipeline** — one chain naming the rig-side stack, the `remote` split marker, the server-side codec, and the model source that loads checkpoints (see `positronic.policy.spec`). The server runs the half right of the marker and declares the half left of it in its handshake; the client builds the declared stack automatically. Vendors ship their pipelines by name, and every name is a server subcommand — `groot-server ee_rot6d_joints` launches that one. The available names are listed in each vendor's README. diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 965718d46..488f3894c 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -192,9 +192,8 @@ def _session_path(path: str, url: str) -> str: return wire.SESSION_PATH if not path.startswith(f'{wire.SESSION_PATH}/'): raise ValueError(f'Unexpected path {path!r} in {url!r}; expected {wire.SESSION_PATH}[/]') - # Kept as written, percent-encoding included, so the server decodes exactly the id whoever handed out - # the URL meant: a trailing slash is part of that id, and an id may itself be a path (a HuggingFace - # repo, say), whose own slashes stay separators. + # Kept as written, percent-encoding included: a trailing slash is part of the id, and an id that is + # itself a path (a HuggingFace repo) keeps its slashes as separators. return path @@ -266,9 +265,8 @@ def _connect(self) -> wire.ClientConnection: self.open_timeout, secure=self._grpc_secure, ) - # A proxy between here and the server closes a connection it has read nothing from, often - # after 60s — well inside one ``infer_timeout`` inference, which sends nothing until it - # answers. The pings keep it open. + # A proxy closes a connection it has read nothing from, often after 60 s, and one inference sends + # nothing until it answers. The pings keep it open. websocket = connect( self.session_url, open_timeout=self.open_timeout, diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index bb2a4129f..7953947fc 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -1,6 +1,10 @@ #!/usr/bin/env bash # Submit a Nebius Serverless Endpoint running a vendor inference server. # +# Usage +# bash workflows/nebius/serve.sh [server args...] +# NEBIUS_PRESET=8gpu-128vcpu-1600gb bash workflows/nebius/serve.sh dreamzero dz-server ee --num_gpus=8 +# # The endpoint gets no public IP: Nebius fronts every HTTP container port with a # managed https:// URL, which is what this polls for and prints. The container # itself takes ~10-15 min more to finish uv sync and load the model into GPU From c01b1d3cfa30a881b71b847b4d4c8c4ba7ac6240 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 09:57:29 +0000 Subject: [PATCH 40/46] State the README's protocol claims as facts Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index daa99a85f..2e9304ed2 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -4,7 +4,7 @@ This package implements the protocol and utilities for offboard policy inference ## Protocol v1 -The unified protocol is built to enable ANY hardware to connect to ANY model. All Positronic inference servers (LeRobot, GR00T, OpenPI) implement this protocol, allowing a single `.remote` policy client to work across all vendors. +The protocol connects any hardware to any model. All Positronic inference servers (LeRobot, GR00T, OpenPI) implement it, so a single `.remote` policy client works across all vendors. ### Wires @@ -244,7 +244,7 @@ uv run positronic eval run --eval=.sim.positronic.stack_cubes \ **Server-side recording:** Servers accept an optional `recording_dir`. When set, each session writes a rerun `.rrd` file that taps both sides of the codec: `raw` captures the obs/action at the wire boundary, and `inference` captures the encoded observation and raw model output. -**Python Client:** We provide a Python client (`positronic.offboard.client.InferenceClient`) that handles the protocol automatically. While the API is currently in alpha and may change, we'll do our best to maintain backward compatibility for the inference client. +**Python Client:** A Python client (`positronic.offboard.client.InferenceClient`) handles the protocol. The API is in alpha and may change. ## Classes From 1a9a9d1d8d460834315fd7eac7c49dfa4c75fce8 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 10:58:54 +0000 Subject: [PATCH 41/46] Rename `UNUSABLE_EDGE_DETAILS` and cut the inherited docstrings to what holds The name says what the tuple holds, and the comment above it goes. Five passages predate the pull request and sit in files it touches. Each one carried design rationale or ran far past the docstring budget. The rationale is in the offboard README and in `workflows/nebius/README.md`, next to the subject it explains. The passages keep each constraint and each footgun: - `PolicyServer`: what it serves, the two halves of the marker, the fixed source, session params, and the pinned default checkpoint. - `serve`: the token comes from the environment, a flag exposes a secret. - `InferenceClient`: the URL forms, the gRPC port, the headers, the timeouts. - `serve.sh`: the managed URL changes on re-create, and the `--auth token` ingress strips the WebSocket upgrade headers. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/client.py | 25 +++++++---------- positronic/offboard/grpc_wire.py | 6 ++-- positronic/offboard/server.py | 31 ++++++--------------- positronic/offboard/tests/test_grpc_wire.py | 4 +-- workflows/nebius/serve.sh | 16 ++++------- 5 files changed, 28 insertions(+), 54 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 488f3894c..6561cd030 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -200,21 +200,16 @@ def _session_path(path: str, url: str) -> str: class InferenceClient: """The wire connection to one inference server, addressed by one URL. - Accepted URL forms: ``host``, ``host:port``, and ``scheme://host[:port][/api/v1/session[/]]``, - each with an optional ``?query``. ``https``/``wss``/``grpcs`` enable TLS (bare or ``http``/``ws``/``grpc`` - forms don't); the port defaults to the scheme's own, 443 for TLS and 80 otherwise. Everything the URL - says about the session — the model id it names and the query it carries as session params — reaches - the server exactly as written, so every session opened here serves that model with those params. - - ``grpc://`` opens the session on the gRPC wire, on the server's own gRPC port; ``grpcs://`` reaches that - port through a TLS edge. That port carries sessions alone: ``list_models`` needs the HTTP URL. - - ``headers`` carry auth, whether the server checks it or a proxy in front of it does — credentials stay - 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. + The URL is ``host``, ``host:port`` or ``scheme://host[:port][/api/v1/session[/]]``, each with + an optional ``?query``. ``https``/``wss``/``grpcs`` enable TLS; the other schemes do not. The port + defaults to 443 with TLS and to 80 without. The model id and the query reach the server as written, + and every session opened here carries them. ``grpc://`` opens the session on the server's gRPC port; + ``grpcs://`` reaches that port through a TLS edge. That port carries sessions alone: ``list_models`` + needs the HTTP URL. + + ``headers`` carry the credentials; the URL carries none. ``open_timeout`` bounds one TCP/TLS handshake, + ``connect_deadline`` the retries until a cold backend answers, and ``infer_timeout`` one inference + round trip. """ def __init__( diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 09d2ea24d..28c106cc5 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -47,14 +47,12 @@ # the caller's ``open_timeout``: a target that drops every connect answers neither. _REFUSAL_PROBE_SEC = 1.0 -# The status details of an edge no client can use: a certificate the roots do not cover, and a front -# that selects no HTTP/2 over ALPN. -UNUSABLE_EDGE = ('CERTIFICATE_VERIFY_FAILED', 'missing selected ALPN property') +UNUSABLE_EDGE_DETAILS = ('CERTIFICATE_VERIFY_FAILED', 'missing selected ALPN property') def edge_is_unusable(details: str) -> bool: """Whether a gRPC status blames the TLS edge's own configuration.""" - return any(marker in details for marker in UNUSABLE_EDGE) + return any(marker in details for marker in UNUSABLE_EDGE_DETAILS) def _client_options() -> list[tuple[str, int]]: diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 5f0506c7d..d9766f1ba 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -177,28 +177,13 @@ class PolicyServer: """Serves a policy pipeline: one layer chain with a ``remote`` marker, closed by a ``ModelSource`` (see ``positronic.policy.spec``). - The half right of the marker wraps the model here; the half left of it is published as the - ``local_stack`` spec in the ``ready`` handshake for the rig to build, alongside the marker's own - wire settings. The source is the only model loader and is fixed at launch. + The half right of the marker wraps the model here. The half left of it goes to the rig as the + ``local_stack`` spec in the ``ready`` handshake, with the marker's own wire settings. The source is + the only model loader and is fixed at launch. - When ``pipeline`` is a ``cfn.Config``, query params on the session URL become dotted - overrides into the pipeline config (e.g. ``?codec.fps=10``), applied and instantiated per session. - Values must be JSON literals (unparseable values pass through as strings) and are applied with - ``Config.override_data``, so a param can tune an argument but never name a Python object to - import; params that change the model source are rejected too. A server built from an - already-instantiated ``Pipeline`` rejects all session params. - - The session flow is: - accept → session params → resolve → load via manager → remote-half wrap → reset → inference loop - - ``serve`` takes the wires sessions arrive on (``positronic.offboard.wire``). ``api`` holds this - server's own HTTP routes, for a wire that speaks HTTP. - - On startup (before accepting connections): resolve(None) → load. - - The default checkpoint is resolved once, at startup, and pinned for every request that names no - explicit one — a running server never switches to a newer checkpoint that lands later. A request - for /api/v1/session/{model_id} still loads that one on demand. + A ``cfn.Config`` pipeline takes session params as dotted overrides (``?codec.fps=10``; the offboard + README states the rules). An instantiated ``Pipeline`` refuses every session param. The default + checkpoint is resolved at startup and pinned; a session that names a model id loads that one. """ def __init__( @@ -477,8 +462,8 @@ def serve( ``grpc_port`` adds the gRPC wire beside the websocket one (see the offboard README). - The bearer token gating the server comes from ``AUTH_TOKEN_ENV`` rather than a flag, which would put - a secret in the process arguments; unset serves open. + The bearer token comes from ``AUTH_TOKEN_ENV``; a flag would put a secret in the process arguments. + Unset serves open. """ server = PolicyServer( pipeline, diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index 327a85bf2..bf165ac2e 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -456,14 +456,14 @@ def test_a_certificate_the_client_cannot_verify_is_not_retried(both_wires, tls_e port, _root = tls_edge(both_wires[0].host, both_wires[0].grpc_port) unrelated, _key = _self_signed(EDGE_HOST) _trust_only(monkeypatch, unrelated) - _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE[0]) + _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE_DETAILS[0]) def test_an_edge_that_selects_no_alpn_is_not_retried(both_wires, tls_edge, monkeypatch): """A front over a raw TCP port terminates TLS and names no ALPN protocol, and gRPC refuses it.""" port, root = tls_edge(both_wires[0].host, both_wires[0].grpc_port, alpn=False) _trust_only(monkeypatch, root) - _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE[1]) + _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE_DETAILS[1]) def test_a_timed_out_session_refuses_the_next_inference(both_wires): diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 7953947fc..39e0486d0 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -14,17 +14,13 @@ # names. The gRPC port is declared as an ordinary HTTP port; a `/tcp` port gets a # front gRPC refuses. The offboard README says what each front does to a session. # -# That URL carries the id of a tunnel created with the endpoint, so it cannot be -# chosen or known in advance, and a delete plus re-create earns a new one even -# under the same name. Nothing may hold it across a redeploy. `nebius ai endpoint -# stop`/`start` keeps it where `stop.sh` would not; a URL that survives re-create -# needs a standalone `nebius tunnel` and its agent in the container. See the -# README's "The managed URL is assigned, not chosen". +# The managed URL is assigned, never chosen, and a delete plus re-create of the +# same name gets a new one; `stop.sh` deletes, `nebius ai endpoint stop`/`start` +# keeps the URL. See the README's "The managed URL is assigned, not chosen". # -# The server is gated on a bearer token (AUTH_TOKEN, from MysteryBox). Auth is -# in-process rather than `nebius ai endpoint create --auth token`, because that -# ingress mode strips the WebSocket upgrade headers and so cannot pass inference -# sessions at all. +# The server is gated on a bearer token (AUTH_TOKEN, from MysteryBox). Auth stays +# in-process: `nebius ai endpoint create --auth token` strips the WebSocket +# upgrade headers and passes no inference session. # # Hardcoded: GPU platform, websocket port. Vendor selects image + uv extra. One # setting of its own, via env: NEBIUS_PRESET. Everything shared with the other From c9e66e28624e4337fd8d428b011b86186a020869 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 10:59:46 +0000 Subject: [PATCH 42/46] State where the served model's values are named Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/server.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index d9766f1ba..1aca6bd8b 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -456,9 +456,8 @@ def serve( ): """The CLI entry point every vendor server exposes: bind ``pipeline``, and the commands are configs of this. - Only the sockets and the recording taps are flags of their own; everything the served model is — - codec, source, checkpoint directory — is reached through the pipeline itself - (``--pipeline.source.checkpoints_dir=...``), so each of those values has exactly one name. + Only the sockets and the recording taps are flags of their own. The codec, the source and the checkpoint + directory are reached through the pipeline (``--pipeline.source.checkpoints_dir=...``), each under one name. ``grpc_port`` adds the gRPC wire beside the websocket one (see the offboard README). From 9738f0c6a92f810ee58fe59acfe486a66140b6c3 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 11:15:56 +0000 Subject: [PATCH 43/46] Raise a wire-neutral refusal from each wire's dial `wire.Refusal` names what a refused connect says about the server, and `wire.ConnectRefused` carries it with the library error as the cause. `websocket_wire.dial` and `grpc_wire.dial` open a client's connection and classify what their library refuses. The gRPC wire also classifies a status that ends the stream before its first message, and the websocket client end raises `PeerDisconnected` for a closed connection. `InferenceClient.new_session` catches `ConnectRefused`, `TimeoutError` and `PeerDisconnected`, spends the refusal against the retry budget, and keeps the connect deadline. It raises `ConnectRefused` where it raised the library exception. `client.py` imports neither `grpc` nor `websockets`. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/README.md | 5 + positronic/offboard/client.py | 103 ++++-------------- positronic/offboard/grpc_wire.py | 95 ++++++++++------ positronic/offboard/tests/test_grpc_wire.py | 38 +++++-- .../offboard/tests/test_remote_policy.py | 68 +++++++----- positronic/offboard/tests/test_server.py | 30 ++++- positronic/offboard/websocket_wire.py | 47 +++++++- positronic/offboard/wire.py | 23 +++- 8 files changed, 252 insertions(+), 157 deletions(-) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index 2e9304ed2..f623b2963 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -294,6 +294,11 @@ meta = session.metadata action = session.infer(observation) ``` +`new_session` retries a cold backend until `connect_deadline`, and raises `TimeoutError` when it stays +cold. A refusal that no retry clears raises `wire.ConnectRefused`, whose `refusal` says what the server +answered: `FORBIDDEN` for a refused credential, `FINAL` for a permanent refusal. `new_session` raises no +exception of the websocket or gRPC library. + ## Vendor Implementations Every vendor ships a `ModelSource` plus named pipelines and serves them through the one `PolicyServer`: diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 6561cd030..a6fdd0ecd 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -1,15 +1,10 @@ import logging -import ssl import time import urllib.parse from enum import Enum -from http import HTTPStatus from typing import Any -import grpc import httpx -from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidStatus -from websockets.sync.client import connect from . import grpc_wire, protocol, websocket_wire, wire from .protocol import deserialise, serialise, typed_commands @@ -100,46 +95,11 @@ class _ConnectOutcome(Enum): SURFACE = 'surface' -class _Refusal(Enum): - """What a refused connect says about the server.""" - - COLD = 'cold' # a backend still starting; retry to the deadline - FORBIDDEN = 'forbidden' # a cold backend, or a refused credential; a few attempts, then surface - FINAL = 'final' # a permanent refusal; surface at once - - -_COLD_GRPC_CODES = (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.DEADLINE_EXCEEDED) - - -def _refusal(e: Exception) -> _Refusal: - """What a refused connect says about the server, over either wire. - - ``PERMISSION_DENIED`` reads as 403, ``UNAVAILABLE`` as 503, ``RESOURCE_EXHAUSTED`` as 429. A TLS - edge no client can use answers ``UNAVAILABLE`` too; its details tell it from a cold backend. - """ - if isinstance(e, InvalidStatus): - status = e.response.status_code - if status == HTTPStatus.FORBIDDEN: - return _Refusal.FORBIDDEN - if status >= HTTPStatus.INTERNAL_SERVER_ERROR or status == HTTPStatus.TOO_MANY_REQUESTS: - return _Refusal.COLD - return _Refusal.FINAL - # A gRPC error carries its code as a `Call`; anything else says nothing about the server. - if isinstance(e, grpc.Call): - if grpc_wire.edge_is_unusable(e.details() or ''): - return _Refusal.FINAL - code = e.code() - if code is grpc.StatusCode.PERMISSION_DENIED: - return _Refusal.FORBIDDEN - return _Refusal.COLD if code in _COLD_GRPC_CODES else _Refusal.FINAL - return _Refusal.COLD - - class _ConnectRetries: """The retry policy over one ``new_session``'s connect attempts. - A 403 or a ``PERMISSION_DENIED`` means a cold backend or a refused credential, and gets - ``MAX_FORBIDDEN_ATTEMPTS`` attempts. + A ``FORBIDDEN`` refusal means a cold backend or a refused credential, and gets ``MAX_FORBIDDEN_ATTEMPTS`` + attempts. """ MAX_FORBIDDEN_ATTEMPTS = 3 @@ -147,14 +107,13 @@ class _ConnectRetries: def __init__(self) -> None: self._forbidden_attempts = 0 - def take(self, e: Exception) -> _ConnectOutcome: + def take(self, refusal: wire.Refusal) -> _ConnectOutcome: """Spend a refused connect against the budget.""" - refusal = _refusal(e) - if refusal is _Refusal.FORBIDDEN: + if refusal is wire.Refusal.FORBIDDEN: self._forbidden_attempts += 1 again = self._forbidden_attempts < self.MAX_FORBIDDEN_ATTEMPTS else: - again = refusal is _Refusal.COLD + again = refusal is wire.Refusal.COLD return _ConnectOutcome.RETRY if again else _ConnectOutcome.SURFACE @@ -252,7 +211,7 @@ def __init__( def _connect(self) -> wire.ClientConnection: """One session's connection, over the wire the URL names.""" if self._grpc_target is not None: - return grpc_wire.GrpcClientConnection( + return grpc_wire.dial( self._grpc_target, self._session_path, self._query, @@ -260,16 +219,7 @@ def _connect(self) -> wire.ClientConnection: self.open_timeout, secure=self._grpc_secure, ) - # A proxy closes a connection it has read nothing from, often after 60 s, and one inference sends - # nothing until it answers. The pings keep it open. - websocket = connect( - self.session_url, - open_timeout=self.open_timeout, - additional_headers=self.headers, - ping_interval=20.0, - max_size=wire.MAX_MESSAGE_BYTES, - ) - return websocket_wire.WebsocketClientConnection(websocket) + return websocket_wire.dial(self.session_url, self.headers, self.open_timeout) def _open_session(self) -> InferenceSession: """One attempt at a session. The connection closes when the handshake does not finish. @@ -285,37 +235,30 @@ def _open_session(self) -> InferenceSession: raise def new_session(self) -> InferenceSession: - """Creates a new inference session on the model the URL names.""" + """Creates a new inference session on the model the URL names. + + Raises ``wire.ConnectRefused`` when the wire refuses the session and no retry clears it. + """ deadline = time.monotonic() + self.connect_deadline backoff = 1.0 retries = _ConnectRetries() while True: try: return self._open_session() - # ``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 - # Each of these can be a backend that is not ready: a timed-out connect, a reset TLS handshake, a - # refused upgrade or gRPC call, a dropped status handshake. ``_ConnectRetries`` tells a permanent - # refusal apart. - except ( - TimeoutError, - ssl.SSLError, - ConnectionClosed, - InvalidHandshake, - grpc.RpcError, - wire.PeerDisconnected, - ) as e: - if retries.take(e) is _ConnectOutcome.SURFACE: - raise - if time.monotonic() >= deadline: - 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) - backoff = min(backoff * 2, 30.0) + except wire.ConnectRefused as e: + refusal, not_ready = e.refusal, e + # A status handshake the server did not finish: a backend that is not ready. + except (TimeoutError, wire.PeerDisconnected) as e: + refusal, not_ready = wire.Refusal.COLD, e except OSError as e: raise type(e)(f'{e} (connecting to {self.session_url})') from e + if retries.take(refusal) is _ConnectOutcome.SURFACE: + raise not_ready + if time.monotonic() >= deadline: + raise TimeoutError(f'{not_ready} (connecting to {self.session_url})') from not_ready + logger.info('Server not ready (cold start?): %s; retrying in %.0fs', not_ready, backoff) + time.sleep(backoff) + backoff = min(backoff * 2, 30.0) def list_models(self) -> list[str]: """List available models from the server.""" diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index 28c106cc5..ba292c72d 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -47,12 +47,25 @@ # the caller's ``open_timeout``: a target that drops every connect answers neither. _REFUSAL_PROBE_SEC = 1.0 -UNUSABLE_EDGE_DETAILS = ('CERTIFICATE_VERIFY_FAILED', 'missing selected ALPN property') +# Status details that blame the TLS edge's own configuration. No client can use such an edge. +_UNUSABLE_EDGE_DETAILS = ('CERTIFICATE_VERIFY_FAILED', 'missing selected ALPN property') +_COLD_CODES = (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.DEADLINE_EXCEEDED) -def edge_is_unusable(details: str) -> bool: - """Whether a gRPC status blames the TLS edge's own configuration.""" - return any(marker in details for marker in UNUSABLE_EDGE_DETAILS) + +def _refusal(status: grpc.RpcError) -> wire.Refusal: + """What a status that ended a call before it opened says about the server. + + ``PERMISSION_DENIED`` reads as 403, ``UNAVAILABLE`` as 503, ``RESOURCE_EXHAUSTED`` as 429. An unusable + edge answers ``UNAVAILABLE`` too; its details tell it from a cold backend. + """ + details = status.details() or '' + if any(marker in details for marker in _UNUSABLE_EDGE_DETAILS): + return wire.Refusal.FINAL + code = status.code() + if code is grpc.StatusCode.PERMISSION_DENIED: + return wire.Refusal.FORBIDDEN + return wire.Refusal.COLD if code in _COLD_CODES else wire.Refusal.FINAL def _client_options() -> list[tuple[str, int]]: @@ -94,46 +107,56 @@ def _connect_refusal(channel: grpc.Channel, timeout: float) -> grpc.RpcError | N return None +def _ready_channel(target: str, secure: bool, open_timeout: float) -> grpc.Channel: + """A channel to ``target`` that is ready. Raises ``wire.ConnectRefused`` when it is not within ``open_timeout``.""" + channel = _channel(target, secure) + deadline = time.monotonic() + open_timeout + try: + grpc.channel_ready_future(channel).result(timeout=open_timeout - _probe_share(open_timeout)) + except grpc.FutureTimeoutError as not_ready: + refusal = _connect_refusal(channel, timeout=max(0.0, deadline - time.monotonic())) + # An ``UNIMPLEMENTED`` from the probe path means the channel is up: the readiness wait was too short. + if refusal is not None and refusal.code() is grpc.StatusCode.UNIMPLEMENTED: + return channel + channel.close() + if refusal is None: + message = f'gRPC channel to {target} is not ready within {open_timeout}s' + raise wire.ConnectRefused(wire.Refusal.COLD, message) from not_ready + raise wire.ConnectRefused(_refusal(refusal), str(refusal)) from refusal + return channel + + +def dial( + target: str, session_path: str, query: str, headers: Mapping[str, str] | None, open_timeout: float, secure: bool +) -> 'GrpcClientConnection': + """A client's end of one session on ``target``. Raises ``wire.ConnectRefused`` when the channel does not open. + + ``secure`` dials over TLS, to a TLS edge in front of the server's plaintext port. + """ + channel = _ready_channel(target, secure, open_timeout) + # gRPC metadata keys are lower case; the header names are the websocket wire's. + metadata = tuple((key.lower(), value) for key, value in (headers or {}).items()) + ( + (SESSION_PATH_HEADER, session_path), + (SESSION_QUERY_HEADER, query), + ) + return GrpcClientConnection(channel, target, metadata) + + class GrpcClientConnection: - """A client's end of one gRPC session. + """A client's end of one gRPC session, over a ready ``channel``. A reader thread drains the response stream into a queue: the stream has no per-message timeout, and - ``recv`` needs one. ``secure`` dials over TLS, to a TLS edge in front of the server's plaintext port. + ``recv`` needs one. """ - def __init__( - self, - target: str, - session_path: str, - query: str, - headers: Mapping[str, str] | None = None, - open_timeout: float = 10.0, - secure: bool = False, - ): + def __init__(self, channel: grpc.Channel, target: str, metadata: tuple[tuple[str, str], ...]): self._target = target - self._channel = _channel(target, secure) - deadline = time.monotonic() + open_timeout - try: - grpc.channel_ready_future(self._channel).result(timeout=open_timeout - _probe_share(open_timeout)) - except grpc.FutureTimeoutError: - refusal = _connect_refusal(self._channel, timeout=max(0.0, deadline - time.monotonic())) - # An ``UNIMPLEMENTED`` from the probe path means the channel is up: the readiness wait was too short. - if refusal is None or refusal.code() is not grpc.StatusCode.UNIMPLEMENTED: - self._channel.close() - # An edge that refuses every client is permanent, and the connect loop retries a ``TimeoutError`` - # to its deadline. - if refusal is not None and edge_is_unusable(refusal.details() or ''): - raise refusal from None - raise TimeoutError(f'gRPC channel to {target} is not ready within {open_timeout}s') from None - # gRPC metadata keys are lower case; the header names are the websocket wire's. - metadata = tuple((key.lower(), value) for key, value in (headers or {}).items()) + ( - (SESSION_PATH_HEADER, session_path), - (SESSION_QUERY_HEADER, query), - ) + self._channel = channel self._outbox: queue.SimpleQueue[bytes | None] = queue.SimpleQueue() self._inbox: queue.SimpleQueue[bytes | BaseException] = queue.SimpleQueue() self._closed = False self._ended = False + self._received = False call = self._channel.stream_stream(METHOD_PATH, request_serializer=None, response_deserializer=None) self._responses = call(self._requests(), metadata=metadata) self._reader = threading.Thread(target=self._read, name='grpc-session-reader', daemon=True) @@ -174,7 +197,11 @@ def recv(self, timeout: float | None = None) -> bytes: if isinstance(answer, BaseException): # What ended the stream is queued once. The caller reads it before ``send`` refuses a write. self._ended = True + # A status before any message crossed is the server refusing the call. + if isinstance(answer, grpc.RpcError) and not self._received: + raise wire.ConnectRefused(_refusal(answer), str(answer)) from answer raise answer + self._received = True return answer def close(self) -> str: diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index bf165ac2e..4c52d3337 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -181,9 +181,9 @@ def test_the_grpc_wire_refuses_a_session_without_the_token(authed_server, header # refusal. monkeypatch.setattr(_ConnectRetries, 'MAX_FORBIDDEN_ATTEMPTS', 1) headers = None if header is None else {AUTH_HEADER: header} - with pytest.raises(grpc.RpcError) as refused: + with pytest.raises(wire.ConnectRefused) as refused: InferenceClient(grpc_url(authed_server), headers=headers).new_session() - assert refused.value.code() is grpc.StatusCode.PERMISSION_DENIED + assert refused.value.refusal is wire.Refusal.FORBIDDEN # An address, and no name that resolves to two families: gRPC reports the last address it failed on, @@ -310,9 +310,29 @@ def test_a_tls_edge_carries_the_bearer_token(authed_server, edged): def test_a_tls_edge_session_without_the_token_is_refused(authed_server, edged, monkeypatch): monkeypatch.setattr(_ConnectRetries, 'MAX_FORBIDDEN_ATTEMPTS', 1) - with pytest.raises(grpc.RpcError) as refused: + with pytest.raises(wire.ConnectRefused) as refused: InferenceClient(edged(authed_server)).new_session() - assert refused.value.code() is grpc.StatusCode.PERMISSION_DENIED + assert refused.value.refusal is wire.Refusal.FORBIDDEN + + +@pytest.mark.parametrize( + ('code', 'details', 'refusal'), + [ + (grpc.StatusCode.PERMISSION_DENIED, 'Invalid or missing bearer token', wire.Refusal.FORBIDDEN), + (grpc.StatusCode.UNAVAILABLE, 'connection refused', wire.Refusal.COLD), + (grpc.StatusCode.RESOURCE_EXHAUSTED, '', wire.Refusal.COLD), + (grpc.StatusCode.DEADLINE_EXCEEDED, '', wire.Refusal.COLD), + (grpc.StatusCode.UNAVAILABLE, 'Cannot check peer: missing selected ALPN property', wire.Refusal.FINAL), + (grpc.StatusCode.UNAVAILABLE, 'CERTIFICATE_VERIFY_FAILED', wire.Refusal.FINAL), + (grpc.StatusCode.UNIMPLEMENTED, '', wire.Refusal.FINAL), + (grpc.StatusCode.INTERNAL, '', wire.Refusal.FINAL), + ], +) +def test_a_status_that_refuses_the_call_reads_as_its_http_status_does(code, details, refusal): + status = MagicMock() + status.code.return_value = code + status.details.return_value = details + assert grpc_wire._refusal(status) is refusal def test_an_unknown_scheme_is_refused(): @@ -447,8 +467,9 @@ def _surfaces_at_once(url: str, blamed: str) -> None: """Assert that a connect to ``url`` fails, names ``blamed``, and spends no retry deadline.""" client = InferenceClient(url, open_timeout=2.0, connect_deadline=20.0) started = time.monotonic() - with pytest.raises(grpc.RpcError, match=blamed): + with pytest.raises(wire.ConnectRefused, match=blamed) as refused: client.new_session() + assert refused.value.refusal is wire.Refusal.FINAL assert time.monotonic() - started < 8.0, 'the connect retried a permanent failure' @@ -456,14 +477,14 @@ def test_a_certificate_the_client_cannot_verify_is_not_retried(both_wires, tls_e port, _root = tls_edge(both_wires[0].host, both_wires[0].grpc_port) unrelated, _key = _self_signed(EDGE_HOST) _trust_only(monkeypatch, unrelated) - _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE_DETAILS[0]) + _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire._UNUSABLE_EDGE_DETAILS[0]) def test_an_edge_that_selects_no_alpn_is_not_retried(both_wires, tls_edge, monkeypatch): """A front over a raw TCP port terminates TLS and names no ALPN protocol, and gRPC refuses it.""" port, root = tls_edge(both_wires[0].host, both_wires[0].grpc_port, alpn=False) _trust_only(monkeypatch, root) - _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire.UNUSABLE_EDGE_DETAILS[1]) + _surfaces_at_once(f'grpcs://{EDGE_HOST}:{port}', grpc_wire._UNUSABLE_EDGE_DETAILS[1]) def test_a_timed_out_session_refuses_the_next_inference(both_wires): @@ -481,7 +502,8 @@ def test_a_timed_out_session_refuses_the_next_inference(both_wires): def test_a_connection_refuses_to_send_once_the_server_ends_the_stream(both_wires): """gRPC stops reading the request iterator, and a write waits out a whole timeout.""" served, _policy = both_wires - conn = grpc_wire.GrpcClientConnection(f'{served.host}:{served.grpc_port}', f'{wire.SESSION_PATH}/unknown-model', '') + target = f'{served.host}:{served.grpc_port}' + conn = grpc_wire.dial(target, f'{wire.SESSION_PATH}/unknown-model', '', None, 10.0, secure=False) try: conn.recv(timeout=10.0) # The server refuses the model in a frame, then ends the stream with that status. diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 96a90a9e2..11eacd035 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -1,17 +1,14 @@ import threading import time -from http import HTTPStatus from unittest.mock import MagicMock, patch import numpy as np import pytest -from websockets.datastructures import Headers -from websockets.exceptions import InvalidStatus -from websockets.http11 import Response from positronic import keys, telemetry, telemetry_keys from positronic.drivers.roboarm import command from positronic.offboard import keys as offboard_keys +from positronic.offboard import wire from positronic.offboard.client import DEFAULT_INFER_TIMEOUT, InferenceClient, _ConnectRetries from positronic.offboard.tests.conftest import ANSWER_SEC, round_trip from positronic.policy import RemotePolicy @@ -94,7 +91,7 @@ def test_headers_stored_and_copied(self): 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.websocket_wire.connect') as mock_connect, patch('positronic.offboard.client.InferenceSession') as mock_session_cls, ): client = InferenceClient('localhost:8000', headers=headers) @@ -108,7 +105,7 @@ def test_new_session_passes_additional_headers(self): def test_new_session_without_headers_passes_none(self): with ( - patch('positronic.offboard.client.connect') as mock_connect, + patch('positronic.offboard.websocket_wire.connect') as mock_connect, patch('positronic.offboard.client.InferenceSession'), ): client = InferenceClient('localhost:8000') @@ -198,7 +195,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.websocket_wire.connect') as mock_connect, patch('positronic.offboard.client.InferenceSession'), ): client = InferenceClient('localhost:8000/api/v1/session/10000?fps=10') @@ -210,57 +207,68 @@ def test_every_session_dials_the_session_url(self): assert call.args[0] == client.session_url == 'ws://localhost:8000/api/v1/session/10000?fps=10' -def _refused(status: HTTPStatus) -> InvalidStatus: - return InvalidStatus(Response(status, 'refused', Headers())) +def _refused(refusal: wire.Refusal) -> wire.ConnectRefused: + return wire.ConnectRefused(refusal, 'refused') -class TestNewSessionRetriesRefusedUpgrades: - """Which non-101 upgrade responses are a backend still coming up, and which are the endpoint saying no.""" +class TestNewSessionRetriesRefusedConnects: + """Which refusals are a backend still coming up, and which are the endpoint saying no.""" - def test_a_403_retries_and_the_session_that_follows_is_returned(self): + def test_a_forbidden_refusal_retries_and_the_session_that_follows_is_returned(self): with ( patch( - 'positronic.offboard.client.connect', side_effect=[_refused(HTTPStatus.FORBIDDEN), MagicMock()] - ) as mock_connect, + 'positronic.offboard.websocket_wire.dial', side_effect=[_refused(wire.Refusal.FORBIDDEN), MagicMock()] + ) as mock_dial, patch('positronic.offboard.client.InferenceSession') as mock_session_cls, patch('positronic.offboard.client.time.sleep'), ): session = InferenceClient('localhost:8000').new_session() - assert mock_connect.call_count == 2 + assert mock_dial.call_count == 2 assert session is mock_session_cls.return_value - def test_a_403_gives_up_once_its_attempts_are_spent(self): + def test_a_forbidden_refusal_gives_up_once_its_attempts_are_spent(self): with ( patch( - 'positronic.offboard.client.connect', - side_effect=[_refused(HTTPStatus.FORBIDDEN)] * (_ConnectRetries.MAX_FORBIDDEN_ATTEMPTS + 5), - ) as mock_connect, + 'positronic.offboard.websocket_wire.dial', + side_effect=[_refused(wire.Refusal.FORBIDDEN)] * (_ConnectRetries.MAX_FORBIDDEN_ATTEMPTS + 5), + ) as mock_dial, patch('positronic.offboard.client.InferenceSession'), patch('positronic.offboard.client.time.sleep'), - pytest.raises(InvalidStatus), + pytest.raises(wire.ConnectRefused), ): InferenceClient('localhost:8000').new_session() - assert mock_connect.call_count == _ConnectRetries.MAX_FORBIDDEN_ATTEMPTS + assert mock_dial.call_count == _ConnectRetries.MAX_FORBIDDEN_ATTEMPTS - @pytest.mark.parametrize('status', [HTTPStatus.UNAUTHORIZED, HTTPStatus.NOT_FOUND]) - def test_a_refusal_that_no_warm_up_clears_is_raised_at_once(self, status): + def test_a_final_refusal_is_raised_at_once(self): with ( - patch('positronic.offboard.client.connect', side_effect=_refused(status)) as mock_connect, + patch('positronic.offboard.websocket_wire.dial', side_effect=_refused(wire.Refusal.FINAL)) as mock_dial, patch('positronic.offboard.client.InferenceSession'), patch('positronic.offboard.client.time.sleep'), - pytest.raises(InvalidStatus), + pytest.raises(wire.ConnectRefused) as refused, ): InferenceClient('localhost:8000').new_session() - assert mock_connect.call_count == 1 + assert mock_dial.call_count == 1 + assert refused.value.refusal is wire.Refusal.FINAL + + def test_a_cold_refusal_retries_to_the_deadline(self): + with ( + patch('positronic.offboard.websocket_wire.dial', side_effect=_refused(wire.Refusal.COLD)) as mock_dial, + patch('positronic.offboard.client.InferenceSession'), + patch('positronic.offboard.client.time.sleep'), + pytest.raises(TimeoutError, match='ws://localhost:8000'), + ): + InferenceClient('localhost:8000', connect_deadline=0.0).new_session() + + assert mock_dial.call_count == 1 def test_each_session_opens_on_a_full_budget(self): - """A client that spent 403s opening one session still gets all of them for the next.""" - one_session = [_refused(HTTPStatus.FORBIDDEN)] * (_ConnectRetries.MAX_FORBIDDEN_ATTEMPTS - 1) + [MagicMock()] + """A client that spent forbidden refusals opening one session still gets all of them for the next.""" + one_session = [_refused(wire.Refusal.FORBIDDEN)] * (_ConnectRetries.MAX_FORBIDDEN_ATTEMPTS - 1) + [MagicMock()] with ( - patch('positronic.offboard.client.connect', side_effect=one_session * 2) as mock_connect, + patch('positronic.offboard.websocket_wire.dial', side_effect=one_session * 2) as mock_dial, patch('positronic.offboard.client.InferenceSession'), patch('positronic.offboard.client.time.sleep'), ): @@ -268,7 +276,7 @@ def test_each_session_opens_on_a_full_budget(self): client.new_session() client.new_session() - assert mock_connect.call_count == 2 * len(one_session) + assert mock_dial.call_count == 2 * len(one_session) def test_remote_policy_hands_the_url_and_headers_to_the_client(): diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 35e522adc..faec3f915 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -6,13 +6,16 @@ import time import urllib.parse from collections.abc import Callable, Generator +from http import HTTPStatus from typing import Any -from unittest.mock import ANY, MagicMock +from unittest.mock import ANY, MagicMock, patch import configuronic as cfn import httpx import pytest +from websockets.datastructures import Headers from websockets.exceptions import InvalidStatus +from websockets.http11 import Response from websockets.sync.client import connect from positronic import keys @@ -539,12 +542,35 @@ def test_auth_rejects_requests_without_the_token(authed_endpoint, make_header, m url, token = authed_endpoint header = make_header(token) client = InferenceClient(url, headers=None if header is None else {AUTH_HEADER: header}) - with pytest.raises(InvalidStatus): + with pytest.raises(wire.ConnectRefused) as refused: client.new_session() + assert refused.value.refusal is wire.Refusal.FORBIDDEN with pytest.raises(httpx.HTTPStatusError): client.list_models() +@pytest.mark.parametrize( + ('status', 'refusal'), + [ + (HTTPStatus.FORBIDDEN, wire.Refusal.FORBIDDEN), + (HTTPStatus.TOO_MANY_REQUESTS, wire.Refusal.COLD), + (HTTPStatus.SERVICE_UNAVAILABLE, wire.Refusal.COLD), + (HTTPStatus.BAD_GATEWAY, wire.Refusal.COLD), + (HTTPStatus.UNAUTHORIZED, wire.Refusal.FINAL), + (HTTPStatus.NOT_FOUND, wire.Refusal.FINAL), + ], +) +def test_a_non_101_answer_to_the_upgrade_says_what_the_server_is(status, refusal): + refused_upgrade = InvalidStatus(Response(status, 'refused', Headers())) + with ( + patch('positronic.offboard.websocket_wire.connect', side_effect=refused_upgrade), + pytest.raises(wire.ConnectRefused) as refused, + ): + websocket_wire.dial('ws://localhost:8000/api/v1/session', None, 1.0) + assert refused.value.refusal is refusal + assert refused.value.__cause__ is refused_upgrade + + @pytest.mark.endpoint def test_auth_accepts_the_token(authed_endpoint): url, token = authed_endpoint diff --git a/positronic/offboard/websocket_wire.py b/positronic/offboard/websocket_wire.py index d64b7dd18..8acfdc8ef 100644 --- a/positronic/offboard/websocket_wire.py +++ b/positronic/offboard/websocket_wire.py @@ -1,10 +1,15 @@ """The websocket wire, and the two ends of a websocket session.""" import socket +import ssl +from collections.abc import Mapping +from http import HTTPStatus import uvicorn from fastapi import APIRouter, Depends, FastAPI, WebSocket, WebSocketDisconnect, WebSocketException, status from starlette.datastructures import QueryParams +from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidStatus +from websockets.sync.client import connect from websockets.sync.connection import Connection from . import wire @@ -17,10 +22,16 @@ def __init__(self, websocket: Connection): self._websocket = websocket def send(self, message: bytes) -> None: - self._websocket.send(message) + try: + self._websocket.send(message) + except ConnectionClosed as e: + raise wire.PeerDisconnected(str(e)) from e def recv(self, timeout: float | None = None) -> bytes: - message = self._websocket.recv(timeout=timeout) + try: + message = self._websocket.recv(timeout=timeout) + except ConnectionClosed as e: + raise wire.PeerDisconnected(str(e)) from e assert isinstance(message, bytes), f'A frame is bytes, and this one is {type(message).__name__}' return message @@ -31,6 +42,38 @@ def close(self) -> str: return f'state {state_before_close} -> {self._websocket.state.name}, close code {self._websocket.close_code}' +def _status_refusal(status_code: int) -> wire.Refusal: + """What a non-101 answer to the upgrade says about the server.""" + if status_code == HTTPStatus.FORBIDDEN: + return wire.Refusal.FORBIDDEN + if status_code >= HTTPStatus.INTERNAL_SERVER_ERROR or status_code == HTTPStatus.TOO_MANY_REQUESTS: + return wire.Refusal.COLD + return wire.Refusal.FINAL + + +def dial(url: str, headers: Mapping[str, str] | None, open_timeout: float) -> WebsocketClientConnection: + """A client's end of one session on ``url``. Raises ``wire.ConnectRefused`` when the upgrade does not open.""" + try: + # A proxy closes a connection it has read nothing from, often after 60 s, and one inference sends + # nothing until it answers. The pings keep it open. + websocket = connect( + url, + open_timeout=open_timeout, + additional_headers=headers, + ping_interval=20.0, + max_size=wire.MAX_MESSAGE_BYTES, + ) + except InvalidStatus as e: + raise wire.ConnectRefused(_status_refusal(e.response.status_code), str(e)) from e + except ssl.SSLCertVerificationError as e: + raise wire.ConnectRefused(wire.Refusal.FINAL, str(e)) from e + # A timed-out connect, a reset TLS handshake, a refused upgrade, a dropped handshake: a backend that is + # not ready. + except (TimeoutError, ssl.SSLError, ConnectionClosed, InvalidHandshake) as e: + raise wire.ConnectRefused(wire.Refusal.COLD, str(e)) from e + return WebsocketClientConnection(websocket) + + class WebsocketServerConnection(wire.ServerConnection): """A server's end of one websocket session, over an accepted ``WebSocket``.""" diff --git a/positronic/offboard/wire.py b/positronic/offboard/wire.py index f40392512..bc82071bc 100644 --- a/positronic/offboard/wire.py +++ b/positronic/offboard/wire.py @@ -6,6 +6,7 @@ import abc from collections.abc import Awaitable, Callable, Mapping +from enum import Enum from typing import NamedTuple, Protocol from starlette.datastructures import QueryParams @@ -22,6 +23,22 @@ class PeerDisconnected(Exception): """The peer ended the session.""" +class Refusal(Enum): + """What a refused connect says about the server.""" + + COLD = 'cold' # a backend still starting; retry to the deadline + FORBIDDEN = 'forbidden' # a cold backend, or a refused credential; a few attempts, then surface + FINAL = 'final' # a permanent refusal; surface at once + + +class ConnectRefused(Exception): + """A wire could not open a session. The library error that refused it is the cause.""" + + def __init__(self, refusal: Refusal, message: str): + super().__init__(message) + self.refusal = refusal + + class Endpoint(NamedTuple): """Where a wire serves.""" @@ -35,7 +52,11 @@ class ClientConnection(Protocol): def send(self, message: bytes) -> None: ... def recv(self, timeout: float | None = None) -> bytes: - """The next message. Raises ``TimeoutError`` when none arrives in time.""" + """The next message. + + Raises ``TimeoutError`` when none arrives in time, ``PeerDisconnected`` once the server ends the + session, and ``ConnectRefused`` when the server refuses the session before its first message. + """ ... def close(self) -> str: From 0f3f5b43f327b1ab86f28a7c6fe434e0081722a4 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 12:17:10 +0000 Subject: [PATCH 44/46] Report a lost peer when a status ends a started gRPC session A status that ends the stream before its first message is the server refusing the call. A status after one is a lost peer, and the wire raised the `grpc.RpcError` itself. `InferenceClient.new_session` retries `ConnectRefused`, `TimeoutError` and `PeerDisconnected`, so that error escaped a cold start the websocket wire retries: the websocket client end raises `PeerDisconnected` for a closed connection whether or not a frame crossed. The gRPC wire now raises it too, with the status as the cause. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/grpc_wire.py | 9 +++++--- positronic/offboard/tests/test_grpc_wire.py | 24 ++++++++++++++++----- 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/positronic/offboard/grpc_wire.py b/positronic/offboard/grpc_wire.py index ba292c72d..88a3b9c67 100644 --- a/positronic/offboard/grpc_wire.py +++ b/positronic/offboard/grpc_wire.py @@ -197,9 +197,12 @@ def recv(self, timeout: float | None = None) -> bytes: if isinstance(answer, BaseException): # What ended the stream is queued once. The caller reads it before ``send`` refuses a write. self._ended = True - # A status before any message crossed is the server refusing the call. - if isinstance(answer, grpc.RpcError) and not self._received: - raise wire.ConnectRefused(_refusal(answer), str(answer)) from answer + if isinstance(answer, grpc.RpcError): + # A status before any message crossed is the server refusing the call. A status after one is a + # lost peer, which is what the other wire reports and what the connect retry reads as cold. + if not self._received: + raise wire.ConnectRefused(_refusal(answer), str(answer)) from answer + raise wire.PeerDisconnected(f'{self._target} ended the session: {answer}') from answer raise answer self._received = True return answer diff --git a/positronic/offboard/tests/test_grpc_wire.py b/positronic/offboard/tests/test_grpc_wire.py index 4c52d3337..c3a785491 100644 --- a/positronic/offboard/tests/test_grpc_wire.py +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -455,11 +455,11 @@ def test_a_session_answers_after_a_silence_no_frame_crossed(both_wires, chatty_c def test_a_server_on_the_grpc_ping_defaults_kills_the_silent_session( start_server, make_mock_policy, chatty_client, monkeypatch ): - """gRPC's own server defaults answer those pings with ``GOAWAY too_many_pings``.""" + """gRPC's own server defaults answer those pings with ``GOAWAY too_many_pings``, and the session is lost.""" monkeypatch.setattr(grpc_wire, '_server_options', lambda: list(grpc_wire._MESSAGE_SIZE_OPTIONS)) policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) served = start_server(ChunkedSchedule() | remote | PolicySource(policy), grpc=True) - with pytest.raises(grpc.RpcError, match='Too many pings'): + with pytest.raises(wire.PeerDisconnected, match='Too many pings'): _silent_then_infer(served) @@ -499,15 +499,29 @@ def test_a_timed_out_session_refuses_the_next_inference(both_wires): session.infer({'image': 'test'}) -def test_a_connection_refuses_to_send_once_the_server_ends_the_stream(both_wires): - """gRPC stops reading the request iterator, and a write waits out a whole timeout.""" +def test_a_status_after_the_first_frame_surfaces_as_a_lost_peer(both_wires): + """A stream that ends after frames have crossed raises a lost peer, which the connect retry reads as cold.""" served, _policy = both_wires target = f'{served.host}:{served.grpc_port}' conn = grpc_wire.dial(target, f'{wire.SESSION_PATH}/unknown-model', '', None, 10.0, secure=False) try: conn.recv(timeout=10.0) # The server refuses the model in a frame, then ends the stream with that status. - with pytest.raises(grpc.RpcError): + with pytest.raises(wire.PeerDisconnected) as gone: + conn.recv(timeout=10.0) + assert isinstance(gone.value.__cause__, grpc.RpcError) + finally: + conn.close() + + +def test_a_connection_refuses_to_send_once_the_server_ends_the_stream(both_wires): + """``send`` raises as soon as the terminal status is read, and the write never reaches the outbox.""" + served, _policy = both_wires + target = f'{served.host}:{served.grpc_port}' + conn = grpc_wire.dial(target, f'{wire.SESSION_PATH}/unknown-model', '', None, 10.0, secure=False) + try: + conn.recv(timeout=10.0) + with pytest.raises(wire.PeerDisconnected): conn.recv(timeout=10.0) with pytest.raises(wire.PeerDisconnected): conn.send(b'an observation the stream can no longer carry') From b884a13fb20b5e31e9b1a430e8dd4f9564cb2881 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 12:17:33 +0000 Subject: [PATCH 45/46] Log an unexpected error delivery failure at ERROR A failed session tells the client before it cleans up. That send reaches a peer which may already be gone, and one blanket catch recorded every outcome at DEBUG. A peer that has gone raises `PeerDisconnected` and stays at DEBUG. Any other failure names the peer at ERROR, so a delivery path that breaks for another reason is not silent. Ticket: Positronic-Robotics/internal#1191 #refs --- positronic/offboard/server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 1aca6bd8b..9180a51ca 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -352,8 +352,10 @@ async def _serve_session(self, conn: wire.ServerConnection, model_id: str | None try: await conn.send(serialise({protocol.STATUS: protocol.ServerStatus.ERROR, protocol.ERROR: str(e)})) await conn.refuse(str(e)) + except wire.PeerDisconnected: + logger.debug('The client was gone before the error reached it', exc_info=True) except Exception: - logger.debug('Failed to send error to client', exc_info=True) + logger.error(f'Failed to tell {conn.peer} its session failed: {e}', exc_info=True) finally: self._active_sessions = max(0, self._active_sessions - 1) self._last_activity = time.monotonic() From 735741f0b349d28b1a067246d3e611034c17730e Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Sat, 12 Sep 2026 12:17:57 +0000 Subject: [PATCH 46/46] Serve both wires in the offboard docs `PolicyServer` takes no `host` or `port`, and `serve` takes the wires it serves on, so the custom-model example in `docs/connect-your-model.md` raised `TypeError`. It now builds the server, then serves it a `WebsocketWire` over `server.api`. The offboard README says every wire carries the same frames, while its session endpoints were named and shown for the websocket alone. The headings name the path, and the examples show both schemes. Ticket: Positronic-Robotics/internal#1191 #refs --- docs/connect-your-model.md | 4 +++- positronic/offboard/README.md | 9 ++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/docs/connect-your-model.md b/docs/connect-your-model.md index 7b36a65cb..a98dd56bb 100644 --- a/docs/connect-your-model.md +++ b/docs/connect-your-model.md @@ -175,6 +175,7 @@ Implement a `Policy`, close a pipeline over it with `PolicySource`, and hand the ```python from positronic.drivers.roboarm import command from positronic.offboard import PolicyServer +from positronic.offboard.websocket_wire import WebsocketWire from positronic.policy import Policy, Session from positronic.policy.spec import PolicySource, remote from positronic.policy.layers import ChunkedSchedule, StopOnFault @@ -209,7 +210,8 @@ class MyPolicy(Policy): pipeline = StopOnFault() | ChunkedSchedule() | remote | PolicySource(MyPolicy(load_my_model())) -PolicyServer(pipeline, host='0.0.0.0', port=8000).serve() +server = PolicyServer(pipeline) +server.serve([WebsocketWire('0.0.0.0', 8000, server.api)]) ``` The pipeline reads left to right: everything left of the `remote` marker is the client-side stack the server declares in its handshake (here the standard `StopOnFault` and `ChunkedSchedule`); everything right of it runs on the server. `PolicySource` is the pipeline's terminal — a model source that serves one already-built policy. diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index f623b2963..9186bbebe 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -74,16 +74,19 @@ curl http://localhost:8000/api/v1/models Use this to discover which models are available before connecting. -#### `WS /api/v1/session` +#### `/api/v1/session` Establishes an inference session with the **default** model — the checkpoint pinned at server startup (the configured one, or the latest available at that moment). -#### `WS /api/v1/session/{model_id}` +#### `/api/v1/session/{model_id}` Establishes an inference session with a **specific** model. **Example:** - `ws://localhost:8000/api/v1/session` → Default model +- `grpc://localhost:9000/api/v1/session` → Default model, over gRPC - `ws://localhost:8000/api/v1/session/10000` → Model 10000 -- `ws://localhost:8000/api/v1/session/20000` → Model 20000 +- `grpc://localhost:9000/api/v1/session/10000` → Model 10000, over gRPC + +Each wire from the table above takes the same path; only the scheme and the port change. The id is everything after the prefix, slashes included, so a source may advertise one that is itself a path: `ws://localhost:8000/api/v1/session/GEAR-Dreams/DreamZero-DROID` serves that HuggingFace checkpoint. Anything else