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/docs/connect-your-model.md b/docs/connect-your-model.md index 7b36a65cb..609826ce2 100644 --- a/docs/connect-your-model.md +++ b/docs/connect-your-model.md @@ -1,6 +1,6 @@ # Connect Your Model -Positronic lets any robot run any policy over one WebSocket protocol. A trained model runs as a server; the robot — or a simulator — runs a client that streams observations to it and executes the actions it returns. This guide explains how that split works and how to plug in your own model. +Positronic lets any robot run any policy over one protocol, which a WebSocket or a gRPC wire carries. A trained model runs as a server; the robot — or a simulator — runs a client that streams observations to it and executes the actions it returns. This guide explains how that split works and how to plug in your own model. **What you need:** [uv](https://docs.astral.sh/uv/) and a clone of the repo (`git clone git@github.com:Positronic-Robotics/positronic.git`). Docker is optional — it is only a convenient way to get a vendor model's Python dependencies; the server itself is an ordinary webserver you can also run from a checkout. @@ -90,7 +90,7 @@ Four small concepts make up the API. You meet them whether you use a built-in se ## The wire format -This is the concrete data crossing the WebSocket. Every message is [msgpack](https://msgpack.org/) with numpy array support (see [Serialization](#serialization)). +This is the concrete data a wire carries. Every message is [msgpack](https://msgpack.org/) with numpy array support (see [Serialization](#serialization)). ### Observation (client → server) @@ -166,7 +166,7 @@ The client side can record too: `--output_dir` saves the full episode as a Posit ## Implement your own server -To connect a custom model you implement this WebSocket protocol. The full low-level spec — endpoints, handshake, status messages — is in the [Offboard README](../positronic/offboard/README.md); the rest of this section shows the shortcut for Positronic-based servers. +To connect a custom model you implement this protocol. The full low-level spec — endpoints, handshake, status messages — is in the [Offboard README](../positronic/offboard/README.md); the rest of this section shows the shortcut for Positronic-based servers. ### Ready, in-process models @@ -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/docs/inference.md b/docs/inference.md index 71cc4e9f6..db02c56fe 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. 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. @@ -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: @@ -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) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index ddf7140e1..9186bbebe 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -4,13 +4,53 @@ 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 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 + +The protocol is a sequence of msgpack frames, and two wires carry them. Both carry the same frames in +the same order; 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 | +| 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; 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 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; 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 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. 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 and stays on the server's `port`. `InferenceClient.list_models` +refuses a `grpc://` URL. ### 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. @@ -34,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 @@ -67,14 +110,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: +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)` 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. -### WebSocket Flow +### Session Flow #### 1. Handshake Upon connection, the server sends a ready packet with metadata: @@ -121,7 +164,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 30 s per message; the server sends status updates during loading, on either wire: ```json { @@ -202,9 +245,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:** A Python client (`positronic.offboard.client.InferenceClient`) handles the protocol. The API is in alpha and may change. ## Classes @@ -213,22 +256,30 @@ The one server implementation behind every vendor. It serves a **policy pipeline ```python from positronic.offboard import PolicyServer +from positronic.offboard.websocket_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, and `idle_timeout_min` shuts the server down after that many minutes without activity. +`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. ### `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. +`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 positronic.offboard.client import InferenceClient @@ -237,12 +288,20 @@ 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, 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 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 98369c4f7..a6fdd0ecd 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -1,17 +1,12 @@ import logging -import ssl import time import urllib.parse from enum import Enum -from http import HTTPStatus from typing import Any 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, websocket_wire, wire from .protocol import deserialise, serialise, typed_commands logger = logging.getLogger(__name__) @@ -23,8 +18,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 +32,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 +68,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,31 +87,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, - ) - - -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 ('', '/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[/]') - # 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 + logger.info('InferenceSession.close: %s', self._conn.close()) class _ConnectOutcome(Enum): @@ -125,8 +98,8 @@ class _ConnectOutcome(Enum): 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 ``FORBIDDEN`` refusal means a cold backend or a refused credential, and gets ``MAX_FORBIDDEN_ATTEMPTS`` + attempts. """ MAX_FORBIDDEN_ATTEMPTS = 3 @@ -134,34 +107,68 @@ 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.""" - if not isinstance(e, InvalidStatus): - return _ConnectOutcome.RETRY - status = e.response.status_code - if status == HTTPStatus.FORBIDDEN: + if refusal is wire.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 wire.Refusal.COLD return _ConnectOutcome.RETRY if again else _ConnectOutcome.SURFACE +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}') + + +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: 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 + + 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. + 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 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. + ``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__( @@ -174,14 +181,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'): - 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}') - secure = split.scheme in ('https', 'wss') - ws_scheme = '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 @@ -189,55 +197,73 @@ 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 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 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.dial( + self._grpc_target, + self._session_path, + self._query, + self.headers, + self.open_timeout, + secure=self._grpc_secure, + ) + 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. + + 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: + 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.""" + """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: - ws = 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) - # ``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() - 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) + return self._open_session() + 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.""" + 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..0604aa84e --- /dev/null +++ b/positronic/offboard/grpc_wire.py @@ -0,0 +1,357 @@ +"""The gRPC wire: one bidirectional stream per session, which carries the ``protocol`` frames. + +The stream is untyped bytes on both sides. There is no protobuf schema and no generated code. +""" + +import logging +import queue +import threading +import time +import urllib.parse +from collections.abc import AsyncIterator, 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}' + +# 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' + +_MESSAGE_SIZE_OPTIONS = [ + ('grpc.max_receive_message_length', wire.MAX_MESSAGE_BYTES), + ('grpc.max_send_message_length', wire.MAX_MESSAGE_BYTES), +] + +# 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 and release the session. +_CLOSE_TIMEOUT_SEC = 5.0 + +# 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 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 + +# 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') + +# Status details carrying an authoritative answer that the host has no address. A resolver that timed out +# says something else, and stays cold: a name can start resolving, where a misspelt one never does. +_NO_SUCH_HOST_DETAILS = ('Domain name not found', 'DNS server returned answer with no data') + +_COLD_CODES = (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.DEADLINE_EXCEEDED) + + +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. Two refusals + no wait fixes answer ``UNAVAILABLE`` too — an unusable edge, and a host with no address — and their + details tell them from a cold backend. + """ + details = status.details() or '' + if any(marker in details for marker in _UNUSABLE_EDGE_DETAILS + _NO_SUCH_HOST_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]]: + return [ + *_MESSAGE_SIZE_OPTIONS, + ('grpc.keepalive_time_ms', _PING_EVERY_MS), + ('grpc.keepalive_timeout_ms', _PING_ANSWER_TIMEOUT_MS), + # 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), + ] + + +def _channel(target: str, secure: bool) -> grpc.Channel: + options = _client_options() + if secure: + # 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: + """The share of one connect attempt the refusal probe gets; the readiness wait gets the rest. + + 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. 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)) + except grpc.RpcError as e: + return e + except StopIteration: + return None + 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, 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. + """ + + def __init__(self, channel: grpc.Channel, target: str, metadata: tuple[tuple[str, str], ...]): + self._target = target + 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) + 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, and end the inbox with what stopped the stream.""" + 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: + # 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 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: + 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. The caller reads it before ``send`` refuses a write. + self._ended = True + 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 + + 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. 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() + # 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}' + ) + + +def model_id_of(session_path: str) -> str | None: + """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 + 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], + 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) + + @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)} + + +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}' + + +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), + ] + + +class GrpcWire(wire.Wire): + """The gRPC wire: sessions on a port of their own, one bidirectional stream each. + + A ``port`` of 0 binds any free one. The port is plaintext; a TLS edge in front of it serves an + authenticated endpoint. + """ + + 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 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)) + + 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 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) + 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 fbe7d33b7..34facceab 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -7,23 +7,22 @@ import os import time from collections import Counter -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from importlib.metadata import version as _pkg_version from typing import Any import configuronic as cfn import pos3 -import uvicorn -from fastapi import Depends, FastAPI, Header, HTTPException, WebSocket, WebSocketDisconnect, WebSocketException, status +from fastapi import APIRouter, Depends, Header, HTTPException 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 -from . import protocol +from . import grpc_wire, protocol, websocket_wire, wire from .protocol import deserialise, serialise logger = logging.getLogger(__name__) @@ -38,7 +37,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 +48,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 +67,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 +76,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 +93,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 +102,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 @@ -183,32 +177,18 @@ 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 websocket 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: - accept → session params → resolve → load via manager → remote-half wrap → reset → inference loop - - 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__( 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, @@ -224,9 +204,6 @@ def __init__( _declared_stack(local) self._source = self._pipeline.source self._manager = PolicyManager(self._source) - self.host = host - self.port = 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 @@ -238,6 +215,9 @@ def __init__( self._infer_lock = asyncio.Lock() self._default_id: str | None = None + # Set while ``serve`` runs; ``shutdown`` reaches the 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 @@ -246,15 +226,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('/api/v1/session', 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._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: @@ -263,15 +243,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 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: - 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()} @@ -291,30 +270,44 @@ 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 self._serve_session(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"}') + 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: + # 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})) + 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"}') 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,15 +322,17 @@ 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: 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, @@ -345,35 +340,22 @@ 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() - 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 websocket.send_bytes(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 self._answer_observations(conn, session) + 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 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() @@ -393,59 +375,111 @@ 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 - def serve(self): + @staticmethod + def _raise_first_wire_failure(started: Sequence[wire.Wire], outcomes: Sequence[Any]): + """Raise the first wire that ended on an error, and log every other one.""" + failed = [(w, e) for w, e in zip(started, outcomes, strict=True) if isinstance(e, Exception)] + # 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; a silent return reads as a shutdown. + raise failed[0][1] + + 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. ``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. + """ + if not wires: + raise ValueError('wires must hold at least one wire; a server with none binds nothing and answers nobody') + 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') - server = uvicorn.Server(config) - self._last_activity = time.monotonic() - watchdog = None - if self.idle_timeout_min and self.idle_timeout_min > 0: - watchdog = asyncio.create_task(self._idle_watchdog(server)) + # 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] = [] try: - await server.serve() + 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 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())) + await asyncio.wait(serving + ending, return_when=asyncio.FIRST_COMPLETED) finally: - if watchdog is not None: - watchdog.cancel() + for task in ending: + task.cancel() + 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) + + self._raise_first_wire_failure(started, outcomes) try: asyncio.run(_run()) except KeyboardInterrupt: logger.info('Server stopped by user') finally: + self._loop, self._stop = None, None 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): + 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( + 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 — - 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). - 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. """ - 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), - ).serve() + ) + 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/stub.py b/positronic/offboard/stub.py new file mode 100644 index 000000000..7e3f7617e --- /dev/null +++ b/positronic/offboard/stub.py @@ -0,0 +1,76 @@ +"""A server with no model: every inference answers the same chunk, and a session measures the wire alone. + +``delay_sec`` is a session param: ``?delay_sec=120`` holds one inference open for two minutes, 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 + +# The smallest trajectory a served session can answer: one action, at the start of the chunk. +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``, 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): + 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 refuses a session param that rebuilds a different source, and +# two ``PolicySource``s are equal only over one policy object. +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)) diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 54dfcbb7a..6d2865806 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -1,13 +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 grpc_wire, websocket_wire, wire from positronic.offboard.server import PolicyServer from positronic.policy import Policy, Session from positronic.policy.executor import Executor @@ -15,45 +13,46 @@ from positronic.policy.spec import ModelSource, PolicySource, remote -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(('', 0)) - return s.getsockname()[1] +class Served(NamedTuple): + """A running server, and the ports its wires took.""" + 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.""" - 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) - 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() + """Factory serving pipelines on daemon threads; every started server is stopped and joined at teardown. - thread = threading.Thread(target=asyncio.run, args=(_run(),), daemon=True) + 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]] = [] + + def start(pipeline, *, grpc: bool = False, **server_kwargs) -> Served: + host = server_kwargs.pop('host', 'localhost') + server = PolicyServer(pipeline, **server_kwargs) + wires: list[wire.Wire] = [websocket_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 @@ -104,7 +103,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]): @@ -145,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 @@ -153,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 new file mode 100644 index 000000000..2cd9c7583 --- /dev/null +++ b/positronic/offboard/tests/test_grpc_wire.py @@ -0,0 +1,546 @@ +"""The gRPC wire: a session runs over it as it runs over the websocket.""" + +import asyncio +import datetime +import ipaddress +import logging +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 +from positronic.offboard.client import InferenceClient, _ConnectRetries +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 + +_TOKEN = 'test-secret-token' + + +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[Served, MagicMock]: + """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 + + +def test_a_grpc_session_handshakes_and_infers(both_wires): + served, policy = both_wires + session = InferenceClient(grpc_url(served)).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 _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): + served, _policy = both_wires + obs = {'image': 'test'} + over_ws = InferenceClient(f'{served.host}:{served.port}').new_session() + over_grpc = InferenceClient(grpc_url(served)).new_session() + try: + 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): + 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): + """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() + + 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`` 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 + session.close() + assert served.server._active_sessions == 0 + + +def test_a_failed_inference_reaches_the_client_as_an_exception(both_wires): + 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'): + 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'})} + served = start_server(ChunkedSchedule() | remote | DictSource(policies), grpc=True) + with pytest.raises(RuntimeError, match='Unknown model'): + InferenceClient(grpc_url(served, 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'}), + } + 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']}] + 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)) + 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] + assert stack[0]['args']['offsets_sec'] == [-0.5, 0.0] + finally: + session.close() + + +@pytest.fixture +def authed_server(start_server: StartServer, make_mock_policy) -> Served: + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + 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): + 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 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(wire.ConnectRefused) as refused: + InferenceClient(grpc_url(authed_server), headers=headers).new_session() + 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, +# and a refused second family 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()) + 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.IPAddress(ipaddress.ip_address(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() + # 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: + 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, as an authenticated endpoint is served. + + 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]] = [] + + 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: + 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) + 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, EDGE_HOST, 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) + + +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[[Served], str]: + """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) + _trust_only(monkeypatch, root) + return f'grpcs://{EDGE_HOST}:{port}' + + return url + + +def test_a_session_through_a_tls_edge_handshakes_and_infers(both_wires, edged): + served, policy = both_wires + session = InferenceClient(edged(served)).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(wire.ConnectRefused) as refused: + InferenceClient(edged(authed_server)).new_session() + 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.UNAVAILABLE, + 'address lookup failed for gpu-host:443: Domain name not found', + wire.Refusal.FINAL, + ), + ( + grpc.StatusCode.UNAVAILABLE, + 'address lookup failed for gpu-host:443: DNS server returned answer with no data', + wire.Refusal.FINAL, + ), + ( + grpc.StatusCode.UNAVAILABLE, + 'address lookup failed for gpu-host:443: Timeout while contacting DNS servers', + wire.Refusal.COLD, + ), + (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(): + with pytest.raises(ValueError, match='Unsupported scheme'): + InferenceClient('tcp://gpu-host:9000') + + +@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'), + [ + (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; 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): + 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() + 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): + """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' + + policy = make_mock_policy([{'action': [4]}], {'model_name': 'stub'}) + 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: + session.close() + + +def test_a_refused_handshake_closes_the_connection(both_wires): + """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 + + 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' + + +# 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 every 500 ms, and a silence of seconds stands in for one of minutes.""" + monkeypatch.setattr(grpc_wire, '_PING_EVERY_MS', 500) + + +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'}) + finally: + session.close() + + +def test_a_session_answers_after_a_silence_no_frame_crossed(both_wires, chatty_client): + """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]}] + + +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``, 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(wire.PeerDisconnected, match='Too many pings'): + _silent_then_infer(served) + + +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(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' + + +def test_a_certificate_the_client_cannot_verify_is_not_retried(both_wires, tls_edge, monkeypatch): + 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]) + + +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]) + + +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.""" + served, policy = both_wires + policy._mock_session.side_effect = lambda *_: time.sleep(1.0) or [{'action': [1, 2, 3]}] + session = InferenceClient(grpc_url(served), infer_timeout=0.2).new_session() + with pytest.raises(TimeoutError): + session.infer({'image': 'test'}) + # The late answer is the first observation's actions. + with pytest.raises(wire.PeerDisconnected): + session.infer({'image': 'test'}) + + +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(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') + # 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_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 9f2f6201e..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 @@ -26,7 +23,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 +31,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): @@ -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) @@ -102,11 +99,13 @@ 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 ( - 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') @@ -196,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') @@ -208,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'), ): @@ -266,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(): @@ -314,8 +324,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 @@ -325,7 +335,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): @@ -333,13 +343,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) @@ -358,7 +368,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 @@ -368,14 +378,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 @@ -389,8 +399,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 @@ -405,14 +415,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 @@ -437,7 +447,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] = [] @@ -458,8 +468,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): @@ -484,7 +494,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}] @@ -501,19 +511,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({offboard_keys.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({offboard_keys.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 @@ -525,7 +535,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 2bf3bea0a..59f0200c5 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -1,24 +1,32 @@ +import asyncio +import logging import os import socket +import threading 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 from positronic.offboard import keys as offboard_keys +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.websocket_wire import WebsocketClientConnection from positronic.policy import Codec, Policy, RemotePolicy, Session from positronic.policy.base import Runtime from positronic.policy.codec import ActionTimestamp @@ -46,10 +54,99 @@ def meta(self, model_id: str) -> dict[str, Any]: return {'type': 'stub'} +# 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 raises.""" + + def __init__(self, after: float): + self._after = after + self.stopped = False + + @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: + 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_server_with_no_wire_refuses_to_serve(make_mock_policy): + """A server that binds nothing answers nobody, so it raises instead of reporting itself ready.""" + server = PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {}))) + with pytest.raises(ValueError, match='at least one wire'): + server.serve([], on_ready=lambda: pytest.fail('it reported ready with no wire bound')) + + +def test_a_wire_that_cannot_bind_stops_the_ones_that_did(make_mock_policy): + """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'): + server.serve([bound, _UnbindableWire()]) + 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, 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 holds the port, and a fresh bind to it raises. + 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): + """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, and ``serve`` returns with no ``shutdown`` call.""" + server = PolicyServer( + ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {})), idle_timeout_min=_A_MOMENT_IDLE / 60 + ) + 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' + + @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 @@ -126,7 +223,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}') @@ -155,7 +252,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') @@ -184,7 +281,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 @@ -223,7 +320,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: @@ -275,7 +372,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())) @@ -304,14 +401,14 @@ 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 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 @@ -341,7 +438,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, []) @@ -359,7 +456,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 @@ -405,7 +503,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')]) @@ -431,7 +529,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 @@ -451,12 +549,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 @@ -484,7 +605,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() @@ -513,7 +636,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/websocket_wire.py b/positronic/offboard/websocket_wire.py new file mode 100644 index 000000000..8acfdc8ef --- /dev/null +++ b/positronic/offboard/websocket_wire.py @@ -0,0 +1,192 @@ +"""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 + + +class WebsocketClientConnection: + """A client's end of one websocket session.""" + + def __init__(self, websocket: Connection): + self._websocket = websocket + + def send(self, message: bytes) -> None: + try: + self._websocket.send(message) + except ConnectionClosed as e: + raise wire.PeerDisconnected(str(e)) from e + + def recv(self, timeout: float | None = None) -> bytes: + 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 + + 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}' + + +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``.""" + + 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 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) + # 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.""" + + # 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): + 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: + """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: + """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) + + 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``: 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: + 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 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 new file mode 100644 index 000000000..bc82071bc --- /dev/null +++ b/positronic/offboard/wire.py @@ -0,0 +1,131 @@ +"""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. ``websocket_wire`` and +``grpc_wire`` hold the two wires. +""" + +import abc +from collections.abc import Awaitable, Callable, Mapping +from enum import Enum +from typing import NamedTuple, Protocol + +from starlette.datastructures import QueryParams + +# 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, and +# the gRPC default of 4 MiB refuses one. +MAX_MESSAGE_BYTES = 16 * 1024 * 1024 + + +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.""" + + host: str + port: int + + +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, ``PeerDisconnected`` once the server ends the + session, and ``ConnectRefused`` when the server refuses the session before its first message. + """ + ... + + def close(self) -> str: + """Close this end, and report what the wire saw, for the log. + + 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. + """ + ... + + +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 endpoint(self) -> Endpoint: + """Where the wire that accepted this session serves.""" + + @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, and tell the client why.""" + + +# What a wire hands the server for each session it accepts: the connection, and the model the route +# 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. +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. + """ + + @property + @abc.abstractmethod + def endpoint(self) -> Endpoint: + """Where this wire serves. The port is 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.""" diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index b6a4c7b8e..c666e89cd 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,27 +128,24 @@ 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 - class RemotePolicy(Policy): """Policy running against a remote inference server, owning the stack in front of the connection. diff --git a/positronic/vendors/lerobot/tests/test_server.py b/positronic/vendors/lerobot/tests/test_server.py index 6cd5e4562..2078612f1 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 websocket_wire, 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) -> websocket_wire.WebsocketServerConnection: + """What the websocket wire hands the server for one session it has accepted.""" + return websocket_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..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,6 +4,7 @@ from fastapi import WebSocketDisconnect from starlette.datastructures import QueryParams +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 @@ -53,6 +54,10 @@ async def close(self, **kwargs): self.events.append('close') await self._close(**kwargs) + def as_connection(self) -> websocket_wire.WebsocketServerConnection: + """What the websocket wire hands the server for one session it has accepted.""" + return websocket_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] 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/utilities/validate_server.py b/utilities/validate_server.py index 560b02f50..f5e071784 100644 --- a/utilities/validate_server.py +++ b/utilities/validate_server.py @@ -71,8 +71,9 @@ def main( ): """Validate an inference server by iterating all available models and running inference for each. - ``url`` names the server, in any form ``InferenceClient`` takes; a gated one also needs its bearer - token exported as ``AUTH_TOKEN``. + ``url`` names the server, in any form ``InferenceClient`` takes except ``grpc://`` and ``grpcs://``: + this lists the models first, and the gRPC port carries sessions alone. A gated server also needs its + bearer token exported as ``AUTH_TOKEN``. Example: 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" }, diff --git a/workflows/nebius/README.md b/workflows/nebius/README.md index 2abf2c448..ad04ccc91 100644 --- a/workflows/nebius/README.md +++ b/workflows/nebius/README.md @@ -210,13 +210,15 @@ an empty `wandb/` placeholder. SmolVLA matches the same layout; OpenPI and GR00T checkpoint shapes (see each vendor's README under `positronic/vendors/`). Live WandB metrics flow to your account directly via the API key — they aren't synced to S3. -## Serve a checkpoint as an HTTP endpoint +## Serve a checkpoint as an endpoint `serve.sh` creates a [Nebius Serverless Endpoint](https://docs.nebius.com/serverless/endpoints/manage) running `python -m positronic.vendors..server` on H100. The endpoint gets no public IP: -Nebius fronts the container's port 8000 with a managed `https://` URL, which terminates TLS and -is the contact address. That URL survives endpoint stop/start; deleting an endpoint retires it, -so a re-created one of the same name gets a new URL. Supported vendors: `lerobot_0_3_3`, +Nebius fronts each container port with its own managed `https://` URL, which terminates TLS and is +the contact address. The server listens on two — port 8000 for the websocket wire and port 9000 for +the gRPC one — so the endpoint returns two URLs. `--grpc_port=` moves the second one. Both +survive endpoint stop/start; deleting an endpoint retires them, so a re-created one of the same name +gets new ones. Supported vendors: `lerobot_0_3_3`, `lerobot`, `openpi`, `gr00t`. Every endpoint is gated on a bearer token — see [Authenticated inference](#authenticated-inference) @@ -254,8 +256,10 @@ bash workflows/nebius/serve.sh gr00t groot-server ee_rot6d_rel \ --pipeline.source.checkpoints_dir=s3:///checkpoints/groot// ``` -`serve.sh` blocks until the managed URL appears (typically <1 min), then prints a banner with -that URL, the endpoint ID, and the commands to follow logs and tear down. The container takes +`serve.sh` blocks until the managed URLs appear (typically <1 min), then prints a banner with both +of them, the endpoint ID, and the commands to follow logs and tear down. A rig points at either wire: +the `https://` URL for the websocket, and for gRPC the port-9000 host dialled as `grpcs://:443`, +which the banner prints ready to paste. The container takes another ~10–15 min to finish `uv sync` and load the model into GPU memory; once `INFO Started server process` appears in `nebius ai endpoint logs`, sanity-check with (`AUTH_TOKEN` loaded as in [Authenticated inference](#authenticated-inference)): @@ -281,15 +285,16 @@ When you're done, `stop.sh` deletes the endpoint: bash workflows/nebius/stop.sh my-act-demo ``` -Deleting retires the managed URL, and a re-created endpoint of the same name gets a new one — so anything -holding it, a robot config or an eval job, breaks on redeploy. To keep the URL, use `nebius ai endpoint +Deleting retires the managed URLs, and a re-created endpoint of the same name gets new ones — so anything +holding one, a robot config or an eval job, breaks on redeploy. To keep the URL, use `nebius ai endpoint stop ` instead: it releases the compute too, and `start` resumes on the same URL. -### The managed URL is assigned, not chosen +### A managed URL is assigned, not chosen -It belongs to [a tunnel](https://docs.nebius.com/tunnels/overview) Nebius creates with the endpoint — -`https://port8000-.tunnel.applications..nebius.cloud`. No flag sets it and nothing -derives it, which is why `serve.sh` polls `status.public_endpoints` to learn it. +Each belongs to [a tunnel](https://docs.nebius.com/tunnels/overview) Nebius creates with the endpoint — +`https://port-.tunnel.applications..nebius.cloud`. The port prefix is +what tells the two wires apart. No flag sets a URL and nothing derives one, which is why `serve.sh` polls +`status.public_endpoints` to learn them. A URL that outlives the endpoint needs a tunnel of your own (`nebius tunnel create`) with its agent in the container, which also names the host (`services.name`, up to 20 lowercase alphanumerics — `phail` rather diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 95abdcab8..39e0486d0 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -1,24 +1,28 @@ #!/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 # memory after the URL appears. # -# 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". +# 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. +# +# 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, 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. @@ -90,6 +94,17 @@ case " $* " in *) set -- "$@" "--idle_timeout_min=${NEBIUS_IDLE_TIMEOUT_MIN:-20}" ;; esac +# 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; a caller's own --grpc_port names 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 $*" echo "Creating $VENDOR endpoint '$NAME'..." @@ -100,7 +115,8 @@ 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" \ --working-dir /positronic \ @@ -127,10 +143,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, 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://"))] | first // empty') + | jq -r "[.status.public_endpoints[]? | select(startswith(\"https://port${WS_PORT}-\"))] | first // empty") if [ -n "$URL" ]; then break; fi sleep 10 done @@ -140,10 +157,20 @@ 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://||') +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 <