diff --git a/docs/inference.md b/docs/inference.md index 71cc4e9f6..e8ea75737 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -58,6 +58,8 @@ uv run positronic eval run --eval=.sim.positronic.stack_cubes \ 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. +**A Unix socket reaches a server on the same machine.** `--uds /run/policy.sock` binds that socket path in place of a host and a port. `--policy.url=unix:///run/policy.sock` dials it, over no network. A model id and session params follow the socket path as they follow a host: `unix:///run/policy.sock/api/v1/session/10000?codec.fps=10`. Use this carrier for a policy process that runs beside the harness and has no network interface of its own. + **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: - `.authed_remote` — a bearer token read from `AUTH_TOKEN`, which it raises about when that is unset. Every endpoint [`workflows/nebius/serve.sh`](../workflows/nebius/README.md) creates is gated this way, whether the server checks the token itself or a proxy in front of it does. diff --git a/docs/training-workflow.md b/docs/training-workflow.md index 62752f7dc..04230a267 100644 --- a/docs/training-workflow.md +++ b/docs/training-workflow.md @@ -219,6 +219,7 @@ cd docker && docker compose run --rm --service-ports openpi-server ee \ | `--pipeline.ee_frame` | OpenPI only: the EE frame the checkpoint speaks, relative to the rig's `default` | `None` | | `--port` | Server port | `8000` (default) | | `--host` | Server host | `0.0.0.0` (default, binds to all interfaces) | +| `--uds` | Unix socket path to bind in place of `--host`/`--port`, for a client on the same machine | `/run/policy.sock` | The subcommand picks the pipeline and `--pipeline.` reaches anywhere inside it, so every value the served model is built from has exactly one name. The same paths are the per-session query params on the client's `--policy.url` (see the [Inference Guide](inference.md)), except `source.*`, which is fixed at launch. diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index ddf7140e1..ba96c5753 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -74,6 +74,10 @@ Because the whole session configuration fits in the URL, one string is a complet `http(s)`/`ws(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. +A `unix://` URL reaches a server on the same machine over a Unix socket, which needs no network: the server +binds the path with `--uds`, and `unix:///run/policy.sock[/api/v1/session[/]][?query]` dials it. The +socket path runs to the first `/api/v1` segment; everything after it is the URL path the server reads. + ### WebSocket Flow #### 1. Handshake @@ -223,7 +227,7 @@ PolicyServer(pipeline, host='0.0.0.0', port=8000).serve() `PolicySource` serves one ready in-process policy; vendors instead define a `ModelSource` over a checkpoint directory. Passing a `cfn.Config` that builds the pipeline — as the vendor servers do with their named pipelines — enables [session parameters](#session-parameters); an instantiated pipeline serves exactly as launched. `recording_dir` enables the per-session recording taps described above, and `idle_timeout_min` shuts the server down after that many minutes without activity. ### `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`, `--uds`, `--recording_dir` and `--idle_timeout_min` are flags of `serve` itself (`--uds` binds a Unix socket path in place of `--host`/`--port`); 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 diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 98369c4f7..fd50330c8 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -1,14 +1,18 @@ import logging +import os +import re import ssl +import stat import time import urllib.parse from enum import Enum +from functools import partial 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.client import connect, unix_connect from websockets.sync.connection import Connection from . import protocol @@ -117,6 +121,23 @@ def _session_path(path: str, url: str) -> str: return path +def _socket_and_path(split: urllib.parse.SplitResult, url: str) -> tuple[str, str]: + """The socket path a ``unix://`` URL names, decoded, and the URL path left over for the server. + + The split runs over the encoded path, so an escaped ``/api/v1`` cannot be read as the marker. + Decoding follows, and it resolves every escape: ``%2F`` becomes a separator like any other, so a + socket path cannot hold a directory whose own name carries a slash. Only the socket path is + decoded, because it names a file; the URL path reaches the server as written, so a model id + carries its own escapes. + """ + if split.netloc or not split.path.startswith('/'): + raise ValueError(f'Socket path must be absolute in {url!r}; write unix:///path/to.sock') + marker = re.search(r'/api/v1(?=/|$)', split.path) + if marker is None: + return urllib.parse.unquote(split.path), '' + return urllib.parse.unquote(split.path[: marker.start()]), split.path[marker.start() :] + + class _ConnectOutcome(Enum): RETRY = 'retry' SURFACE = 'surface' @@ -131,8 +152,21 @@ class _ConnectRetries: MAX_FORBIDDEN_ATTEMPTS = 3 - def __init__(self) -> None: + def __init__(self, connect_deadline: float, url: str) -> None: self._forbidden_attempts = 0 + self._deadline = time.monotonic() + connect_deadline + self._backoff = 1.0 + self._url = url + + def wait_or_surface(self, e: Exception) -> None: + """Spend one refused connect against the budget, or let it surface. Call it from the handler.""" + if self.take(e) is _ConnectOutcome.SURFACE: + raise + if time.monotonic() >= self._deadline: + raise TimeoutError(f'{e} (connecting to {self._url})') from e + logger.info('Server not ready (cold start?): %s; retrying in %.0fs', e, self._backoff) + time.sleep(self._backoff) + self._backoff = min(self._backoff * 2, 30.0) def take(self, e: Exception) -> _ConnectOutcome: """Spend a refused connect against the budget.""" @@ -156,6 +190,11 @@ class InferenceClient: session — the model id it names and the query it carries as session params — reaches the server exactly as written, so every session opened here serves that model with those params. + ``unix://[/api/v1/session[/]][?query]`` reaches a server on the same + machine over a Unix domain socket, which needs no network. The socket path runs to the first + ``/api/v1`` segment, so ``unix:///run/policy.sock`` is the default session and + ``unix:///run/policy.sock/api/v1/session/10000?fps=10`` names a model and a param. TLS does not apply. + ``headers`` carry auth, whether the server checks it or a proxy in front of it does — credentials stay out of the URL, which is meant to be safe to hand around. @@ -174,45 +213,74 @@ def __init__( infer_timeout: float = DEFAULT_INFER_TIMEOUT, ): split = urllib.parse.urlsplit(url if '://' in url else f'//{url}') - if split.scheme not in ('', 'http', 'ws', 'https', 'wss'): + if split.scheme not in ('', 'http', 'ws', 'https', 'wss', 'unix'): raise ValueError(f'Unsupported scheme {split.scheme!r} in {url!r}') - if not split.hostname: - raise ValueError(f'No host in {url!r}') secure = split.scheme in ('https', 'wss') + if split.scheme == 'unix': + uds, path = _socket_and_path(split, url) + # A socket path is not a host. The server reads the path and the query alone, so the + # handshake asks for them under a host that stands in for the socket. + netloc = 'localhost' + else: + uds = None + if not split.hostname: + raise ValueError(f'No host in {url!r}') + path = split.path + default_port = 443 if 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 + netloc = host if port == default_port else f'{host}:{port}' ws_scheme = 'wss' if secure else 'ws' http_scheme = 'https' if secure else 'http' - default_port = 443 if secure else 80 - # urlsplit strips the brackets an IPv6 host needs back in a netloc. - host = f'[{split.hostname}]' if ':' in split.hostname else split.hostname - port = default_port if split.port is None else split.port - netloc = host if port == default_port else f'{host}:{port}' # 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}' + session_path = _session_path(path, url) + self.uds = uds + # The URL the websocket handshake asks for, and the TCP address to dial when there is no socket. + self._ws_uri = f'{ws_scheme}://{netloc}{session_path}{query}' + # What an error names. Over a socket the stand-in host would not say which socket failed. + self.session_url = self._ws_uri if uds is None else f'unix://{uds}{session_path}{query}' self.api_url = 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 _socket_may_still_appear(self, e: OSError) -> bool: + """Whether a failed dial is a co-located server that has not bound its socket yet. + + Only an absent path and a refusal can mean that; every other ``OSError`` is settled, and + waiting for it spends the whole deadline on an answer that will not change. A refusal then + reads the path, which tells a restarting server from a path naming something that is not a + socket. + """ + assert self.uds is not None + if not isinstance(e, (FileNotFoundError, ConnectionRefusedError)): + return False + try: + return stat.S_ISSOCK(os.stat(self.uds).st_mode) + except FileNotFoundError: + return True + except OSError: + return False + def new_session(self) -> InferenceSession: """Creates a new inference session on the model the URL names.""" - deadline = time.monotonic() + self.connect_deadline - backoff = 1.0 - retries = _ConnectRetries() + retries = _ConnectRetries(self.connect_deadline, self.session_url) 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, + dial = ( + partial(connect, self._ws_uri) + if self.uds is None + else partial(unix_connect, self.uds, uri=self._ws_uri) ) + ws = dial(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. @@ -226,18 +294,18 @@ def new_session(self) -> InferenceSession: 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) + retries.wait_or_surface(e) except OSError as e: - raise type(e)(f'{e} (connecting to {self.session_url})') from e + if ws is not None: + ws.close() + if self.uds is None or not self._socket_may_still_appear(e): + raise type(e)(f'{e} (connecting to {self.session_url})') from e + retries.wait_or_surface(e) def list_models(self) -> list[str]: """List available models from the server.""" - response = httpx.get(f'{self.api_url}/models', headers=self.headers) + transport = None if self.uds is None else httpx.HTTPTransport(uds=self.uds) + with httpx.Client(transport=transport) as client: + response = client.get(f'{self.api_url}/models', headers=self.headers) response.raise_for_status() return response.json()['models'] diff --git a/positronic/offboard/keys.py b/positronic/offboard/keys.py index a9af0a26f..137cd8a99 100644 --- a/positronic/offboard/keys.py +++ b/positronic/offboard/keys.py @@ -4,6 +4,8 @@ # what the rig builds and obeys — the local stack spec, image compression, the positronic version it runs. HOST = 'host' PORT = 'port' +# The socket path a server bound instead of a host and a port. One of the two pairs is present, never both. +UDS = 'uds' CHECKPOINT_ID = 'checkpoint_id' LOCAL_STACK = 'local_stack' COMPRESS_IMAGES = 'compress_images' diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index fbe7d33b7..1f15bbef9 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -1,10 +1,13 @@ """The inference server: serves a policy pipeline (see ``positronic.policy.spec``) over the offboard protocol.""" import asyncio +import errno import hmac import json import logging import os +import socket +import stat import time from collections import Counter from collections.abc import Callable @@ -202,6 +205,9 @@ class PolicyServer: 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. + + ``uds`` binds a Unix socket path instead of ``host:port``, which serves a client on the same machine + over no network. A client reaches it with a ``unix://`` URL. """ def __init__( @@ -212,6 +218,7 @@ def __init__( recording_dir: str | None = None, idle_timeout_min: float | None = None, auth_token: str | None = None, + uds: str | None = None, ): self._pipeline_cfg = pipeline if isinstance(pipeline, cfn.Config) else None self._pipeline = pipeline.instantiate() if isinstance(pipeline, cfn.Config) else pipeline @@ -226,7 +233,11 @@ def __init__( self._manager = PolicyManager(self._source) self.host = host self.port = port - self.metadata: dict[str, Any] = {offboard_keys.HOST: host, offboard_keys.PORT: port} + self.uds = uds + # Where the server listens, as the handshake reports it. A socket path is not a host, so it has its own key. + self.metadata: dict[str, Any] = ( + {offboard_keys.HOST: host, offboard_keys.PORT: port} if uds is None else {offboard_keys.UDS: uds} + ) # 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 @@ -407,20 +418,80 @@ async def _idle_watchdog(self, server: uvicorn.Server): server.should_exit = True return + # The probe bounds its wait, and reads a wait that runs out as a live server: a server whose + # backlog is full holds a connect open, and an unbounded one would stall startup. + LIVE_SOCKET_PROBE_SEC = 1.0 + + @staticmethod + def _is_stale_socket(path: str) -> bool: + """Whether ``path`` is a socket no server answers on, so replacing it takes nothing from anybody. + + A live socket, a probe that runs out of time against a full backlog, and a path that holds + something other than a socket are none of them stale. + """ + try: + if not stat.S_ISSOCK(os.stat(path).st_mode): + return False + except FileNotFoundError: + return False + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: + probe.settimeout(PolicyServer.LIVE_SOCKET_PROBE_SEC) + try: + probe.connect(path) + except ConnectionRefusedError: + return True + except OSError: + return False + return False + + @staticmethod + def claim_socket_path(path: str) -> socket.socket: + """Bind and listen on ``path``, and return the socket, or refuse a path something already holds. + + The bind is the claim, so two servers starting together cannot both take one path: the loser's + bind fails. A probe follows it only to tell a stale file from a live server. Serve the returned + socket by its descriptor: a server handed the path instead binds again, and unlinks this claim. + """ + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + try: + sock.bind(path) + except OSError as taken: + if taken.errno != errno.EADDRINUSE: + raise + if not PolicyServer._is_stale_socket(path): + raise OSError(errno.EADDRINUSE, f'{path!r} is already in use') from None + os.unlink(path) + sock.bind(path) + # The mode is the deployment's, through its umask: widening it here would open the socket + # to every local account that can reach the directory. + sock.listen() + except BaseException: + sock.close() + raise + return sock + def serve(self): async def _run(): 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)) + sock, watchdog = None, None try: + if self.uds is not None: + sock = self.claim_socket_path(self.uds) + fd = None if sock is None else sock.fileno() + config = uvicorn.Config(self.app, host=self.host, port=self.port, fd=fd, log_level='info') + server = uvicorn.Server(config) + self._last_activity = time.monotonic() + if self.idle_timeout_min and self.idle_timeout_min > 0: + watchdog = asyncio.create_task(self._idle_watchdog(server)) await server.serve() finally: if watchdog is not None: watchdog.cancel() + if sock is not None: + # The file stays: a successor reads it as stale, where an unlink here could take a + # path that successor has already claimed. + sock.close() try: asyncio.run(_run()) @@ -430,14 +501,23 @@ async def _run(): self._manager.close() -@cfn.config(host='0.0.0.0', port=8000, recording_dir=None, idle_timeout_min=None) -def serve(pipeline: cfn.Config, host: str, port: int, recording_dir: str | None, idle_timeout_min: float | None): +@cfn.config(host='0.0.0.0', port=8000, recording_dir=None, idle_timeout_min=None, uds=None) +def serve( + pipeline: cfn.Config, + host: str, + port: int, + recording_dir: str | None, + idle_timeout_min: float | None, + uds: str | 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. + ``--uds`` binds that Unix socket path and leaves ``host`` and ``port`` unused. + The bearer token gating the server comes from ``AUTH_TOKEN_ENV`` rather than a flag, which would put a secret in the process arguments; unset serves open. """ @@ -448,4 +528,5 @@ def serve(pipeline: cfn.Config, host: str, port: int, recording_dir: str | None, recording_dir=recording_dir, idle_timeout_min=idle_timeout_min, auth_token=os.environ.get(AUTH_TOKEN_ENV), + uds=uds, ).serve() diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 54dfcbb7a..92e13fb3b 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -1,5 +1,7 @@ import asyncio +import os import socket +import tempfile import threading import time from collections.abc import Callable, Generator, Mapping @@ -21,39 +23,90 @@ def _find_free_port() -> int: return s.getsockname()[1] +RunningServers = list[tuple[uvicorn.Server, threading.Thread]] + StartServer = Callable[..., tuple[str, int, PolicyServer]] +StartUnixServer = Callable[..., PolicyServer] + @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 running_servers() -> Generator[RunningServers, None, None]: + """Every server a test started, stopped and joined at teardown.""" + running: RunningServers = [] + yield running + for uv_server, thread in running: + uv_server.should_exit = True + thread.join(timeout=5.0) + + +def _serve_in_background(server: PolicyServer, config: uvicorn.Config, running: RunningServers) -> None: + uv_server = uvicorn.Server(config) + + async def _run(): + await server._startup() + await uv_server.serve() + + thread = threading.Thread(target=asyncio.run, args=(_run(),), daemon=True) + thread.start() + running.append((uv_server, thread)) + + +def _wait_until_it_accepts(dial: Callable[[], None]) -> None: + deadline = time.time() + 5.0 + while time.time() < deadline: + try: + dial() + return + except OSError: + time.sleep(0.05) + raise RuntimeError('Server failed to start') + + +@pytest.fixture +def socket_path() -> Generator[str, None, None]: + """A path for a Unix socket, short enough for the 104-byte limit that ``tmp_path`` can pass.""" + with tempfile.TemporaryDirectory(dir='/tmp') as directory: + yield os.path.join(directory, 's.sock') + + +@pytest.fixture +def start_server(running_servers: RunningServers) -> StartServer: + """Factory serving pipelines on daemon threads.""" 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')) + config = uvicorn.Config(server.app, host=server.host, port=server.port, log_level='warning') + _serve_in_background(server, config, running_servers) + _wait_until_it_accepts(lambda: socket.create_connection((server.host, server.port), timeout=0.1).close()) + return server.host, server.port, server + + return start - async def _run(): - await server._startup() - await uv_server.serve() - thread = threading.Thread(target=asyncio.run, args=(_run(),), daemon=True) - thread.start() - running.append((uv_server, thread)) +def _dial_unix(path: str) -> None: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: + sock.settimeout(0.1) + sock.connect(path) - 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') + +@pytest.fixture +def start_unix_server(running_servers: RunningServers) -> Generator[StartUnixServer, None, None]: + """Factory serving pipelines on a Unix socket, as ``PolicyServer.serve`` claims one.""" + claimed: list[socket.socket] = [] + + def start(pipeline, uds: str, **server_kwargs) -> PolicyServer: + server = PolicyServer(pipeline, uds=uds, **server_kwargs) + sock = PolicyServer.claim_socket_path(uds) + claimed.append(sock) + config = uvicorn.Config(server.app, fd=sock.fileno(), log_level='warning') + _serve_in_background(server, config, running_servers) + _wait_until_it_accepts(lambda: _dial_unix(uds)) + return server yield start - for uv_server, thread in running: - uv_server.should_exit = True - thread.join(timeout=5.0) + for sock in claimed: + sock.close() @pytest.fixture diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 9f2f6201e..d33bcd214 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -117,7 +117,8 @@ def test_new_session_without_headers_passes_none(self): def test_list_models_passes_headers(self): headers = {'Modal-Key': 'k', 'Modal-Secret': 's'} - with patch('positronic.offboard.client.httpx.get') as mock_get: + with patch('positronic.offboard.client.httpx.Client') as mock_client: + mock_get = mock_client.return_value.__enter__.return_value.get mock_get.return_value.json.return_value = {'models': ['m1']} client = InferenceClient('localhost:8000', headers=headers) @@ -127,7 +128,8 @@ def test_list_models_passes_headers(self): assert mock_get.call_args.kwargs['headers'] == headers def test_list_models_without_headers_passes_none(self): - with patch('positronic.offboard.client.httpx.get') as mock_get: + with patch('positronic.offboard.client.httpx.Client') as mock_client: + mock_get = mock_client.return_value.__enter__.return_value.get mock_get.return_value.json.return_value = {'models': []} client = InferenceClient('localhost:8000') client.list_models() @@ -207,6 +209,55 @@ def test_every_session_dials_the_session_url(self): for call in mock_connect.call_args_list: assert call.args[0] == client.session_url == 'ws://localhost:8000/api/v1/session/10000?fps=10' + def test_a_socket_path_alone_is_the_default_session(self): + client = InferenceClient('unix:///run/policy.sock') + assert client.uds == '/run/policy.sock' + assert client.session_url == 'unix:///run/policy.sock/api/v1/session' + assert client.api_url == 'http://localhost/api/v1' + + def test_a_socket_path_ends_at_the_api_segment(self): + client = InferenceClient('unix:///run/policy.sock/api/v1/session/10000?fps=10') + assert client.uds == '/run/policy.sock' + assert client.session_url == 'unix:///run/policy.sock/api/v1/session/10000?fps=10' + + def test_a_query_on_a_bare_socket_path_rides_along(self): + client = InferenceClient('unix:///run/policy.sock?fps=10') + assert client.uds == '/run/policy.sock' + assert client.session_url == 'unix:///run/policy.sock/api/v1/session?fps=10' + + def test_the_api_marker_is_a_whole_segment(self): + """A socket under a directory whose name only starts with the marker is still the whole path.""" + client = InferenceClient('unix:///run/api/v1x/policy.sock') + assert client.uds == '/run/api/v1x/policy.sock' + assert client.session_url == 'unix:///run/api/v1x/policy.sock/api/v1/session' + + def test_the_socket_path_is_decoded_and_the_session_path_is_not(self): + """The socket path names a file, so its escapes are resolved; the model id reaches the server + as written, which is how an id carrying its own escapes survives.""" + client = InferenceClient('unix:///run/a%20b%25c/policy.sock/api/v1/session/s3%3A//ckpt') + assert client.uds == '/run/a b%c/policy.sock' + assert client.session_url == 'unix:///run/a b%c/policy.sock/api/v1/session/s3%3A//ckpt' + + def test_an_escaped_separator_is_a_separator_once_decoded(self): + """The decode resolves every escape, so no socket path can hold a slash inside one name.""" + client = InferenceClient('unix:///run/odd%2Fname/policy.sock') + assert client.uds == '/run/odd/name/policy.sock' + + def test_a_relative_socket_path_rejected(self): + with pytest.raises(ValueError, match='absolute'): + InferenceClient('unix://policy.sock') + + def test_a_unix_url_dials_the_socket_and_asks_for_the_url_path(self): + with ( + patch('positronic.offboard.client.unix_connect') as mock_connect, + patch('positronic.offboard.client.InferenceSession'), + ): + InferenceClient('unix:///run/policy.sock/api/v1/session/10000?fps=10').new_session() + + mock_connect.assert_called_once() + assert mock_connect.call_args.args[0] == '/run/policy.sock' + assert mock_connect.call_args.kwargs['uri'] == 'ws://localhost/api/v1/session/10000?fps=10' + def _refused(status: HTTPStatus) -> InvalidStatus: return InvalidStatus(Response(status, 'refused', Headers())) diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 2bf3bea0a..27ce2f44a 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -1,10 +1,14 @@ +import errno import os +import pathlib import socket +import stat +import threading import time import urllib.parse from collections.abc import Callable, Generator from typing import Any -from unittest.mock import ANY, MagicMock +from unittest.mock import ANY, MagicMock, patch import configuronic as cfn import httpx @@ -232,6 +236,207 @@ def test_local_stack_declared_in_handshake(start_server, make_mock_policy): session.close() +@pytest.fixture +def unix_stub_server(start_unix_server, socket_path, make_mock_policy) -> tuple[str, MagicMock]: + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + start_unix_server(ChunkedSchedule() | remote | _StubSource(policy), socket_path) + return socket_path, policy + + +def test_a_pipeline_served_over_a_unix_socket(unix_stub_server): + socket_path, policy = unix_stub_server + client = InferenceClient(f'unix://{socket_path}') + + assert client.list_models() == ['stub'] + session = client.new_session() + try: + assert session.metadata['model_name'] == 'stub' + assert session.metadata[offboard_keys.LOCAL_STACK] == {'name': 'chunked_schedule'} + assert session.metadata[offboard_keys.UDS] == socket_path + assert offboard_keys.HOST not in session.metadata + assert offboard_keys.PORT not in session.metadata + + 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_unix_url_carries_the_model_id_past_the_socket_path(unix_stub_server): + socket_path, _policy = unix_stub_server + + session = InferenceClient(f'unix://{socket_path}/api/v1/session/10000').new_session() + try: + assert session.metadata[offboard_keys.CHECKPOINT_ID] == '10000' + finally: + session.close() + + +def test_a_socket_path_carrying_url_escapes_is_dialled_as_a_filename(start_unix_server, socket_path, make_mock_policy): + """``urlsplit`` leaves the path encoded, so a directory holding a space or a percent would + otherwise be dialled as a file that does not exist.""" + odd = pathlib.Path(socket_path).parent / 'a b%c' + odd.mkdir() + uds = str(odd / 's.sock') + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + start_unix_server(ChunkedSchedule() | remote | _StubSource(policy), uds) + + client = InferenceClient(f'unix://{urllib.parse.quote(uds)}') + + assert client.uds == uds + assert client.list_models() == ['stub'] + session = client.new_session() + try: + assert session.infer({'obs': 'data'}) == [{'action': [1, 2, 3]}] + finally: + session.close() + + +@pytest.mark.timeout(60.0) +def test_a_client_waits_for_a_socket_the_server_has_not_bound_yet(start_unix_server, socket_path, make_mock_policy): + """``serve`` binds only once the model has loaded, so a co-located client starting beside its + server finds no socket at all for that interval.""" + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + pipeline = ChunkedSchedule() | remote | _StubSource(policy) + late = threading.Timer(1.5, lambda: start_unix_server(pipeline, socket_path)) + late.start() + + try: + session = InferenceClient(f'unix://{socket_path}', connect_deadline=30.0).new_session() + finally: + late.join() + try: + assert session.metadata['model_name'] == 'stub' + assert session.infer({'obs': 'data'}) == [{'action': [1, 2, 3]}] + finally: + session.close() + + +@pytest.mark.timeout(60.0) +def test_a_client_waits_for_a_server_restarting_over_the_socket_it_left( + start_unix_server, socket_path, make_mock_policy +): + """A bound path whose server has gone refuses the dial, and the successor binds over it. The wait + covers that restart as it covers a first start.""" + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as gone: + gone.bind(socket_path) + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + pipeline = ChunkedSchedule() | remote | _StubSource(policy) + late = threading.Timer(1.5, lambda: start_unix_server(pipeline, socket_path)) + late.start() + + try: + session = InferenceClient(f'unix://{socket_path}', connect_deadline=30.0).new_session() + finally: + late.join() + try: + assert session.infer({'obs': 'data'}) == [{'action': [1, 2, 3]}] + finally: + session.close() + + +def test_a_dial_this_process_broke_fails_at_once_over_a_live_socket(unix_stub_server): + """A descriptor limit is this process's own, so no server appearing clears it. The socket is live + and the path says so, which is exactly when reading the path alone would wait out the deadline.""" + socket_path, _policy = unix_stub_server + started = time.monotonic() + + with patch('positronic.offboard.client.unix_connect') as dial: + dial.side_effect = OSError(errno.EMFILE, 'Too many open files') + with pytest.raises(OSError) as refusal: + InferenceClient(f'unix://{socket_path}', connect_deadline=30.0).new_session() + + assert 'Too many open files' in str(refusal.value) + assert time.monotonic() - started < 5.0 + + +def test_a_dial_at_a_path_holding_something_that_is_not_a_socket_fails_at_once(socket_path): + """No waiting clears a wrong path. Which errno says so differs by platform, so this asserts the + connect deadline goes unspent.""" + pathlib.Path(socket_path).write_text('not a socket') + started = time.monotonic() + + with pytest.raises(OSError): + InferenceClient(f'unix://{socket_path}', connect_deadline=30.0).new_session() + + assert time.monotonic() - started < 5.0 + + +def test_a_server_binds_over_the_socket_an_earlier_run_left(start_unix_server, socket_path, make_mock_policy): + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as stale: + stale.bind(socket_path) + + start_unix_server(ChunkedSchedule() | remote | _StubSource(policy), socket_path) + + assert InferenceClient(f'unix://{socket_path}').list_models() == ['stub'] + + +def test_a_path_that_is_not_a_socket_is_refused_and_left_alone(socket_path): + """A wrong ``uds`` is refused rather than emptied of a file nobody meant to lose.""" + path = pathlib.Path(socket_path) + path.write_text('not a socket') + + with pytest.raises(OSError) as refusal: + PolicyServer.claim_socket_path(socket_path) + + assert refusal.value.errno == errno.EADDRINUSE + assert path.read_text() == 'not a socket' + + +@pytest.mark.timeout(30.0) +def test_a_second_claim_on_one_path_is_refused_and_the_first_goes_on_serving(socket_path): + """The bind is the claim, so two servers starting on one absent path cannot both pass it.""" + held = PolicyServer.claim_socket_path(socket_path) + try: + with pytest.raises(OSError) as refusal: + PolicyServer.claim_socket_path(socket_path) + assert refusal.value.errno == errno.EADDRINUSE + assert socket_path in str(refusal.value) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(socket_path) + assert held.accept()[0].close() is None + finally: + held.close() + + +def test_a_claimed_socket_keeps_the_mode_the_umask_gives(socket_path): + """A restrictive umask is the deployment's choice, and widening it would open the socket to every + local account that can reach the directory.""" + previous = os.umask(0o077) + try: + sock = PolicyServer.claim_socket_path(socket_path) + finally: + os.umask(previous) + try: + assert stat.S_IMODE(os.stat(socket_path).st_mode) & 0o077 == 0 + finally: + sock.close() + + +@pytest.mark.timeout(30.0) +def test_a_server_refuses_a_socket_a_live_server_listens_on(socket_path, make_mock_policy): + """``asyncio.create_unix_server`` unlinks the file it finds, so only a refusal here keeps the + address with the server that owns it.""" + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + server = PolicyServer(ChunkedSchedule() | remote | _StubSource(policy), uds=socket_path) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as live: + live.bind(socket_path) + live.listen() + + with pytest.raises(OSError) as refusal: + server.serve() + assert refusal.value.errno == errno.EADDRINUSE + assert socket_path in str(refusal.value) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(socket_path) + assert live.accept()[0].close() is None + + def test_pipeline_with_no_rig_side_half_refused_at_startup(make_mock_policy): """Nothing left of the marker leaves the rig nothing to run, so the server refuses to serve it.""" stub = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) diff --git a/positronic/vendors/gr00t/README.md b/positronic/vendors/gr00t/README.md index 8276a81d7..8729b3e31 100644 --- a/positronic/vendors/gr00t/README.md +++ b/positronic/vendors/gr00t/README.md @@ -128,6 +128,7 @@ cd docker && docker compose run --rm --service-ports groot-server ee_rot6d_joint | `--pipeline.source.checkpoints_dir` | Experiment directory (contains `checkpoint-N` folders) | Required | `~/checkpoints/groot/my_task_v1/` | | `--pipeline.source.checkpoint` | Specific checkpoint ID | Latest | `10000`, `50000` | | `--port` | Server port | `8000` | `8001` | +| `--uds` | Unix socket path to bind in place of `--host`/`--port` | `None` | `/run/policy.sock` | | `--pipeline.source.modality_config` | Override the pipeline's paired modality config | Paired | `ee_rot6d_q` | **Session parameters:** a client can tune the served pipeline per connection via query params on the diff --git a/positronic/vendors/lerobot/README.md b/positronic/vendors/lerobot/README.md index 4fbee6ae0..0bf6ad5a6 100644 --- a/positronic/vendors/lerobot/README.md +++ b/positronic/vendors/lerobot/README.md @@ -109,6 +109,7 @@ cd docker && docker compose run --rm --service-ports lerobot-server ee \ | `--pipeline.source.device` | Torch device the policy runs on | Auto-detected | `cuda`, `mps`, `cpu` | | `--port` | Server port | `8000` | `8001` | | `--host` | Server host | `0.0.0.0` | Binds to all interfaces | +| `--uds` | Unix socket path to bind in place of `--host`/`--port` | `None` | `/run/policy.sock` | | `--recording_dir` | Directory for server-side inference recordings | `None` | `s3://inference/...` | | `--idle_timeout_min` | Shut down after this many idle minutes | `None` | `30` | diff --git a/positronic/vendors/lerobot_0_3_3/README.md b/positronic/vendors/lerobot_0_3_3/README.md index 79dafee7f..91847d9f2 100644 --- a/positronic/vendors/lerobot_0_3_3/README.md +++ b/positronic/vendors/lerobot_0_3_3/README.md @@ -106,6 +106,7 @@ cd docker && docker compose run --rm --service-ports lerobot-0_3_3-server ee \ | `--pipeline.source.model_type` | Names what the factory builds, for the handshake metadata | `act` | `diffusion` | | `--port` | Server port | `8000` | `8001` | | `--host` | Server host | `0.0.0.0` | Binds to all interfaces | +| `--uds` | Unix socket path to bind in place of `--host`/`--port` | `None` | `/run/policy.sock` | | `--recording_dir` | Directory for server-side inference recordings | `None` | `s3://inference/...` | | `--idle_timeout_min` | Shut down after this many idle minutes | `None` | `30` | diff --git a/positronic/vendors/molmoact2/README.md b/positronic/vendors/molmoact2/README.md index 8835be723..640d2a971 100644 --- a/positronic/vendors/molmoact2/README.md +++ b/positronic/vendors/molmoact2/README.md @@ -32,8 +32,9 @@ uv run --python 3.13 --extra molmoact2 python -m positronic.vendors.molmoact2.se The server serves a named policy pipeline — the codec plus the HuggingFace model source. MolmoAct2 ships one pipeline, `droid`, which is the default subcommand. The codec lives server-side, so clients send raw -observations and receive decoded joint commands. `--host`, `--port`, `--recording_dir` and -`--idle_timeout_min` are the server's flags; the model is reached through the pipeline +observations and receive decoded joint commands. `--host`, `--port`, `--uds`, `--recording_dir` and +`--idle_timeout_min` are the server's flags, and `--uds` binds a Unix socket path in place of +`--host`/`--port`; the model is reached through the pipeline (`--pipeline.source.hf_repo`, `.device_map`, `.norm_tag`, `.num_steps`), with defaults in [`server.py`](./server.py). Sanity-check once warm: diff --git a/positronic/vendors/openpi/README.md b/positronic/vendors/openpi/README.md index 32b0c55ce..64c7f3de5 100644 --- a/positronic/vendors/openpi/README.md +++ b/positronic/vendors/openpi/README.md @@ -140,6 +140,7 @@ emits absolute `JointPosition` chunks executed at RoboLab's leaderboard cadence - `--pipeline.source.checkpoint`: (Optional) Specific checkpoint step to load. If omitted, loads the latest checkpoint - `--pipeline.source.config_name`: (Optional) OpenPI config name; overrides the pipeline's pairing (base pipelines use `pi05_positronic_lowmem`) - `--port`: (Optional) Port to serve on (default: 8000) +- `--uds`: (Optional) Unix socket path to bind in place of `--host`/`--port`, for a client on the same machine - `--pipeline.source.openpi_ws_port`: (Optional) Internal port for OpenPI subprocess (default: 8001) - `--recording_dir`: (Optional) Directory for server-side `.rrd` recordings (local or S3) - `--idle_timeout_min`: (Optional) Shut down after this many minutes without activity