diff --git a/docs/inference.md b/docs/inference.md index e8ea75737..532ea1bd6 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -60,6 +60,8 @@ Accepted forms: `host`, `host:port`, and `https://host[:port][/api/v1/session[/< **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. +Such a server also takes each observation's images through shared memory, which keeps 8 MB of frames out of the message; `--frame_ring=false` turns that off. [`positronic/offboard/README.md`](../positronic/offboard/README.md) states the contract. + **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/positronic/offboard/README.md b/positronic/offboard/README.md index ba96c5753..21e966f58 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -78,6 +78,35 @@ A `unix://` URL reaches a server on the same machine over a Unix socket, which n 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. +#### The frame ring + +A server on a Unix socket carries each observation's images through shared memory instead of the message. +It declares `frame_ring` in the ready handshake, with this session's id as the value. A client that dialled a +`unix://` URL then creates a ring, hands the descriptor over, and sends a reference in place of every image. +A client that ignores the declaration keeps sending whole images, and so does every client over TCP. + +- **The ring is a sealed `memfd`.** The client maps it writable, seals it with + `F_SEAL_SHRINK | F_SEAL_GROW | F_SEAL_FUTURE_WRITE`, and only then hands the descriptor over. The server maps + it read-only and cannot write it, resize it, or punch a hole in it. That holds whatever the server's code + does, so a server that runs untrusted code gets the frames and no way to change them. +- **The descriptor travels beside the session socket.** The server binds a second `AF_UNIX` socket at the + session socket's path plus `.frames`, of type `SOCK_SEQPACKET`, and the client dials the same suffix on the + path it dialled. Each side builds that path from the socket path it already holds, so a bind mount that gives + the two processes different names for one directory still lands them on the same socket. The client sends the + descriptor with `SCM_RIGHTS`, names the session id from the handshake, and waits for the server to map it. +- **A ring holds four slots.** One round trip is in flight at a time, so the writer returns to a slot four + inferences later, and a server that still reads an earlier observation reads the bytes written for it. Each + slot carries a sequence number before its payload and one after it; a reader that finds either one different + from the reference refuses that observation rather than serving other pixels under it. +- **A larger frame grows the ring.** The client creates a bigger one and hands it over before it sends any + reference to it. The server keeps every mapping it was handed, so a view it built earlier stays readable. +- **The views are read-only.** Code that writes an observation's image in place raises; a codec that resizes or + copies is unaffected. + +`--frame_ring=false` on the server keeps every image in the message. A server declares no ring where the +kernel seals no `memfd` — a macOS server, or Linux before 5.1 — or where the session socket's path plus +`.frames` is longer than a Unix socket address may be. + ### WebSocket Flow #### 1. Handshake @@ -100,7 +129,8 @@ Upon connection, the server sends a ready packet with metadata: {"name": "restrict_image_size", "args": {"width": 224, "height": 224}} ]}, "compress_images": false, - "positronic_version": "0.2.1" + "positronic_version": "0.2.1", + "frame_ring": "9f2c1ab4e7d05613" } } ``` @@ -122,6 +152,8 @@ This metadata tells the client: - `compress_images` — the `remote` marker's own wire setting: whether the rig JPEG-encodes frames before sending, for an endpoint behind a proxy with a message-size cap - `positronic_version` — the server's positronic version, for diagnosing declaration mismatches +- `frame_ring` — this session's id, present when the server takes images through shared memory (see above); + absent when it does not #### 2. Status Updates (Long Model Loading) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index fd50330c8..9181ad599 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -15,6 +15,8 @@ from websockets.sync.client import connect, unix_connect from websockets.sync.connection import Connection +from . import frame_ring as frames +from . import keys as offboard_keys from . import protocol from .protocol import deserialise, serialise, typed_commands @@ -27,10 +29,24 @@ class InferenceSession: - def __init__(self, websocket: Connection, infer_timeout: float = DEFAULT_INFER_TIMEOUT): + """One session on one server, and the observations it sends. + + ``uds`` is the Unix socket the session was dialled over, which says the server runs on this host. + A server that also declares a frame ring then gets every image through shared memory, and the + message carries a reference in place of each one. + """ + + def __init__(self, websocket: Connection, infer_timeout: float = DEFAULT_INFER_TIMEOUT, uds: str | None = None): self._websocket = websocket self._infer_timeout = infer_timeout self._metadata = self._handshake() + self._frames = self._frame_writer(uds) + + def _frame_writer(self, uds: str | None) -> frames.FrameWriter | None: + session_id = self._metadata.get(offboard_keys.FRAME_RING) + if uds is None or session_id is None or not frames.SUPPORTED: + return None + return frames.FrameWriter(frames.channel_path(uds), session_id) def _handshake(self, timeout_per_message: float = 30.0) -> dict[str, Any]: """Receive status updates until server is ready. @@ -72,7 +88,7 @@ def infer(self, obs: dict[str, Any]) -> Any: arrays/scalars, and no arbitrary Python objects. The result is whatever the server's session returned — canonically a list of action dicts, but a bare dict or ``None`` too. """ - serialised = serialise(obs) + serialised = serialise(obs if self._frames is None else self._frames.pack(obs)) logger.debug('Size of serialised obs: %1.f KiB', len(serialised) / 1024) self._websocket.send(serialised) @@ -94,6 +110,8 @@ def infer(self, obs: dict[str, Any]) -> Any: return typed_commands(response[protocol.RESULT]) def close(self): + if self._frames is not None: + self._frames.close() 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. @@ -281,7 +299,7 @@ def new_session(self) -> InferenceSession: 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) + return InferenceSession(ws, infer_timeout=self.infer_timeout, uds=self.uds) # ``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: diff --git a/positronic/offboard/frame_ring.py b/positronic/offboard/frame_ring.py new file mode 100644 index 000000000..fbfaf2092 --- /dev/null +++ b/positronic/offboard/frame_ring.py @@ -0,0 +1,347 @@ +"""A shared-memory ring that carries an observation's frames to a server on the same host. + +The client creates the ring, seals it, and hands its descriptor over an ``AF_UNIX`` socket beside the +session socket. Each inference writes the images into a slot and sends a reference for each, which +the server reads through a mapping the seals let it neither write nor resize. +""" + +import contextlib +import fcntl +import logging +import mmap +import os +import socket +import threading +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np + +from positronic.utils import serialization + +logger = logging.getLogger(__name__) + +# The longest path ``sockaddr_un`` carries, in bytes: ``sun_path`` is 108 on Linux, and one holds the NUL. +MAX_SOCKET_PATH = 107 + +# The suffix of the descriptor socket, which each side derives from its own session socket path. +SOCKET_SUFFIX = '.frames' + +# Slots per ring, so the writer never takes a slot the server may still read. +SLOTS = 4 + +# The slot header: two sequence numbers, then padding that puts every payload on a 64-byte boundary. +_HEADER_BYTES = 64 +_ALIGN = 64 + +# The envelope one image travels in when the ring carries its bytes: the slot and the sequence number +# that say which write it belongs to, and where the array sits inside that slot. +_RING = b'__ring__' +_SLOT = b'slot' +_SEQ = b'seq' +_OFFSET = b'offset' +_SHAPE = b'shape' +_DTYPE = b'dtype' + +# The fields of a handover, which the client sends with the descriptor. +_SESSION = 'session' +_SLOTS = 'slots' +_SLOT_BYTES = 'slot_bytes' + +_ACK = b'\x01' +_HANDOVER_BYTES = 4096 + + +class TornFrame(RuntimeError): + """The slot a reference names holds another write, so the pixels under it belong elsewhere.""" + + +def channel_path(socket_path: str) -> str: + """The descriptor socket that belongs to the session socket at ``socket_path``.""" + return socket_path + SOCKET_SUFFIX + + +def _aligned(nbytes: int) -> int: + return -(-nbytes // _ALIGN) * _ALIGN + + +# The seal numbers from ``linux/fcntl.h``: a Python built against other headers exports none of them. +# ``F_SEAL_FUTURE_WRITE`` (Linux 5.1) spares the writer's own mapping, which ``F_SEAL_WRITE`` cannot. +_F_ADD_SEALS = 1033 +_F_SEAL_SHRINK = 0x0002 +_F_SEAL_GROW = 0x0004 +_F_SEAL_FUTURE_WRITE = 0x0010 +_SEALS = _F_SEAL_SHRINK | _F_SEAL_GROW | _F_SEAL_FUTURE_WRITE + + +def _ring_is_supported() -> bool: + """Whether this host builds a sealed ring: macOS has no ``memfd_create``, Linux before 5.1 no seal.""" + if not hasattr(os, 'memfd_create'): + return False + try: + fd = os.memfd_create('positronic-frames-probe', os.MFD_CLOEXEC | os.MFD_ALLOW_SEALING) + except OSError as refused: + logger.warning('No frame ring: this host refuses memfd_create (%s)', refused) + return False + try: + fcntl.fcntl(fd, _F_ADD_SEALS, _SEALS) + except OSError as refused: + logger.warning('No frame ring: this host refuses the memfd seals (%s)', refused) + return False + finally: + os.close(fd) + return True + + +# A server that cannot build a ring declares none, and every image stays in the message. +SUPPORTED = _ring_is_supported() + + +class FrameRing: + """One sealed ring of ``slots`` slots, each holding ``slot_bytes`` of image data. + + The constructor maps it writable before it seals it, which is the only order the seals allow. + """ + + def __init__(self, slot_bytes: int, slots: int = SLOTS): + self.slots = slots + self.slot_bytes = slot_bytes + self._stride = _HEADER_BYTES + _aligned(slot_bytes) + self.fd = os.memfd_create('positronic-frames', os.MFD_CLOEXEC | os.MFD_ALLOW_SEALING) + try: + os.ftruncate(self.fd, self._stride * slots) + self._map = mmap.mmap(self.fd, self._stride * slots, mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE) + fcntl.fcntl(self.fd, _F_ADD_SEALS, _SEALS) + except BaseException: + os.close(self.fd) + raise + self._seq = 0 + + def write(self, arrays: Sequence[np.ndarray]) -> list[dict[bytes, Any]]: + """Copy ``arrays`` into the next slot and return one reference each.""" + self._seq += 1 + slot = self._seq % self.slots + counters = np.ndarray(2, dtype=np.uint64, buffer=self._map, offset=slot * self._stride) + counters[0] = self._seq + payload = slot * self._stride + _HEADER_BYTES + offset = 0 + references = [] + for array in arrays: + destination = np.ndarray(array.shape, dtype=array.dtype, buffer=self._map, offset=payload + offset) + np.copyto(destination, array) + references.append({ + _RING: True, + _SLOT: slot, + _SEQ: self._seq, + _OFFSET: offset, + _SHAPE: list(array.shape), + _DTYPE: array.dtype.str, + }) + offset += _aligned(array.nbytes) + counters[1] = self._seq + return references + + def close(self) -> None: + self._map.close() + os.close(self.fd) + + +def _detach_images(value: Any, found: list[tuple[dict[bytes, Any], np.ndarray]]) -> Any: + """``value`` with an empty reference in place of every image, each paired with its array in ``found``.""" + if serialization.is_image(value): + reference: dict[bytes, Any] = {} + found.append((reference, value)) + return reference + if isinstance(value, Mapping): + return {key: _detach_images(item, found) for key, item in value.items()} + if isinstance(value, list | tuple): + return type(value)(_detach_images(item, found) for item in value) + return value + + +# How long a handover waits for the server to map the ring. +HANDOVER_TIMEOUT_SEC = 10.0 + + +class FrameWriter: + """The client's half: one ring per session, handed to the server and grown when a frame outgrows it. + + ``pack`` returns the observation with a reference in place of every image. + """ + + def __init__(self, channel: str, session_id: str): + self._channel = channel + self._session_id = session_id + self._ring: FrameRing | None = None + + def pack(self, obs: Mapping[str, Any]) -> dict[str, Any]: + found: list[tuple[dict[bytes, Any], np.ndarray]] = [] + packed = _detach_images(obs, found) + if not found: + return packed + arrays = [array for _reference, array in found] + ring = self._ring_for(sum(_aligned(array.nbytes) for array in arrays)) + for (reference, _array), written in zip(found, ring.write(arrays), strict=True): + reference.update(written) + return packed + + def _ring_for(self, slot_bytes: int) -> FrameRing: + if self._ring is not None and self._ring.slot_bytes >= slot_bytes: + return self._ring + ring = FrameRing(slot_bytes) + try: + self._hand_over(ring) + except Exception: + ring.close() + raise + if self._ring is not None: + # A mapping outlives the descriptor it was made from, so this frees the ring only here. + self._ring.close() + self._ring = ring + return ring + + def _hand_over(self, ring: FrameRing) -> None: + header = serialization.serialise({_SESSION: self._session_id, _SLOTS: ring.slots, _SLOT_BYTES: ring.slot_bytes}) + with socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET) as sock: + sock.settimeout(HANDOVER_TIMEOUT_SEC) + sock.connect(self._channel) + socket.send_fds(sock, [header], [ring.fd]) + if sock.recv(len(_ACK)) != _ACK: + raise RuntimeError(f'The server at {self._channel} did not map the frame ring it was handed') + + def close(self) -> None: + if self._ring is not None: + self._ring.close() + self._ring = None + + +class MappedRing: + """The server's half: the ring mapped read-only, and the array a reference names. + + Nothing unmaps the ring. Every view holds a reference to the mapping, so it lives until the last + view is gone; ``mmap.close`` under a live view leaves that view pointing at memory the process no + longer owns. + """ + + def __init__(self, fd: int, slots: int, slot_bytes: int): + self._slots = slots + self._slot_bytes = slot_bytes + self._stride = _HEADER_BYTES + _aligned(slot_bytes) + self._map = mmap.mmap(fd, self._stride * slots, mmap.MAP_SHARED, mmap.PROT_READ) + + def array(self, reference: Mapping[bytes, Any]) -> np.ndarray: + """A read-only view of the image ``reference`` names, over the ring's own pages.""" + slot, seq, offset = reference[_SLOT], reference[_SEQ], reference[_OFFSET] + dtype = np.dtype(reference[_DTYPE]) + shape = tuple(reference[_SHAPE]) + nbytes = dtype.itemsize * int(np.prod(shape)) + if not 0 <= slot < self._slots or offset < 0 or offset + nbytes > self._slot_bytes: + raise ValueError(f'A frame reference names slot {slot} at {offset}+{nbytes}, outside the ring') + counters = np.ndarray(2, dtype=np.uint64, buffer=self._map, offset=slot * self._stride) + if counters[0] != seq or counters[1] != seq: + raise TornFrame(f'Slot {slot} holds write {counters[0]}..{counters[1]}, and the reference names {seq}') + return np.ndarray(shape, dtype=dtype, buffer=self._map, offset=slot * self._stride + _HEADER_BYTES + offset) + + +class FrameChannel: + """The socket that carries ring descriptors to this server, and the rings each session holds. + + A thread accepts each handover, maps the ring read-only and answers, under the session id the + server put in its ready handshake. The caller claims ``path`` and hands the socket to ``start``. + """ + + def __init__(self, path: str): + self.path = path + self._rings: dict[str, MappedRing | None] = {} + self._lock = threading.Lock() + self._socket: socket.socket | None = None + self._thread: threading.Thread | None = None + self._closing = False + + def start(self, sock: socket.socket) -> None: + """Serve handovers on ``sock``, which the caller already bound to ``path`` and listened on.""" + self._socket = sock + self._thread = threading.Thread(target=self._accept_forever, name='frame-channel', daemon=True) + self._thread.start() + logger.info('Frame ring channel listening on %s', self.path) + + def open_session(self, session_id: str) -> None: + with self._lock: + self._rings[session_id] = None + + def close_session(self, session_id: str) -> None: + with self._lock: + self._rings.pop(session_id, None) + + def resolve(self, session_id: str, obs: Any) -> Any: + """``obs`` with a read-only view in place of every frame reference it carries.""" + if isinstance(obs, Mapping): + if _RING in obs: + return self._ring(session_id).array(obs) + return {key: self.resolve(session_id, item) for key, item in obs.items()} + if isinstance(obs, list | tuple): + return type(obs)(self.resolve(session_id, item) for item in obs) + return obs + + def _ring(self, session_id: str) -> MappedRing: + with self._lock: + ring = self._rings.get(session_id) + if ring is None: + raise RuntimeError(f'Session {session_id} sent a frame reference before it handed over a ring') + return ring + + def _accept_forever(self) -> None: + assert self._socket is not None + while not self._closing: + try: + connection, _address = self._socket.accept() + except OSError: + if not self._closing: + logger.exception('The frame ring channel stopped accepting') + return + with connection: + if self._closing: + return + try: + self._take_ring(connection) + # One bad handover must not take the channel down. + except Exception: + logger.exception('A frame ring handover failed') + + def _take_ring(self, connection: socket.socket) -> None: + connection.settimeout(HANDOVER_TIMEOUT_SEC) + message, fds, _flags, _address = socket.recv_fds(connection, _HANDOVER_BYTES, 1) + if not fds: + raise RuntimeError('A frame ring handover carried no descriptor') + try: + header = serialization.deserialise(message) + ring = MappedRing(fds[0], header[_SLOTS], header[_SLOT_BYTES]) + finally: + os.close(fds[0]) + session_id = header[_SESSION] + with self._lock: + # A ring this replaces stays mapped while any view over it lives, so dropping it here frees + # only what nobody reads. + open_here = session_id in self._rings + if open_here: + self._rings[session_id] = ring + if not open_here: + raise RuntimeError(f'A frame ring arrived for session {session_id}, which is not open here') + connection.send(_ACK) + + def close(self) -> None: + self._closing = True + if self._socket is not None: + # Closing a socket another thread waits in ``accept`` on does not wake that thread, so + # knock first. A refused knock means it is not waiting, which is the state this wants. + with contextlib.suppress(OSError), socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET) as knock: + knock.settimeout(HANDOVER_TIMEOUT_SEC) + knock.connect(self.path) + self._socket.close() + self._socket = None + if self._thread is not None: + # A handover under way holds the thread for its own deadline, so the join waits that long. + self._thread.join(timeout=HANDOVER_TIMEOUT_SEC) + self._thread = None + with self._lock: + self._rings.clear() diff --git a/positronic/offboard/keys.py b/positronic/offboard/keys.py index 137cd8a99..ac409f10b 100644 --- a/positronic/offboard/keys.py +++ b/positronic/offboard/keys.py @@ -10,3 +10,5 @@ LOCAL_STACK = 'local_stack' COMPRESS_IMAGES = 'compress_images' POSITRONIC_VERSION = 'positronic_version' +# The id of this session, present when the server takes frames through a shared-memory ring. +FRAME_RING = 'frame_ring' diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 1f15bbef9..88cfac46c 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -6,6 +6,7 @@ import json import logging import os +import secrets import socket import stat import time @@ -20,8 +21,9 @@ from fastapi import Depends, FastAPI, Header, HTTPException, WebSocket, WebSocketDisconnect, WebSocketException, status from starlette.datastructures import QueryParams +from positronic.offboard import frame_ring as frames 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 @@ -208,6 +210,9 @@ class PolicyServer: ``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. + + ``frame_ring`` takes each observation's images through shared memory, which a Unix socket makes + possible; ``positronic.offboard.frame_ring`` states the contract. """ def __init__( @@ -219,6 +224,7 @@ def __init__( idle_timeout_min: float | None = None, auth_token: str | None = None, uds: str | None = None, + frame_ring: bool = True, ): self._pipeline_cfg = pipeline if isinstance(pipeline, cfn.Config) else None self._pipeline = pipeline.instantiate() if isinstance(pipeline, cfn.Config) else pipeline @@ -238,6 +244,13 @@ def __init__( self.metadata: dict[str, Any] = ( {offboard_keys.HOST: host, offboard_keys.PORT: port} if uds is None else {offboard_keys.UDS: uds} ) + # A ring rides beside the Unix socket, so it needs a kernel that seals a memfd and a companion + # path short enough to bind. + channel = frames.channel_path(uds) if uds is not None and frame_ring and frames.SUPPORTED else None + if channel is not None and len(os.fsencode(channel)) > frames.MAX_SOCKET_PATH: + logger.warning('No frame ring: the companion socket path %r is too long to bind', channel) + channel = None + self._frames = frames.FrameChannel(channel) if channel is not None else None # 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 @@ -318,6 +331,11 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): self._last_activity = time.monotonic() policy: Policy | None = None session = None + # The id a frame ring is handed over under. It names this session and nothing else, so a ring + # reaches the session that declared it. + session_id = secrets.token_hex(8) + if self._frames is not None: + self._frames.open_session(session_id) try: pipeline = self._session_pipeline(_session_params(websocket.query_params)) local, border, remote_half = split(pipeline) @@ -356,25 +374,11 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): offboard_keys.COMPRESS_IMAGES: border.compress_images, offboard_keys.POSITRONIC_VERSION: _pkg_version('positronic'), } + if self._frames is not None: + meta[offboard_keys.FRAME_RING] = session_id await websocket.send_bytes(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: - logger.info('Client disconnected') + await self._answer_until_disconnect(websocket, session, session_id) except Exception as e: logger.error(f'Failed session: {e}', exc_info=True) @@ -388,6 +392,8 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): finally: self._active_sessions = max(0, self._active_sessions - 1) self._last_activity = time.monotonic() + if self._frames is not None: + self._frames.close_session(session_id) try: if session is not None: # Both ends of a session's life touch the backend — close does a reset round-trip — so @@ -399,7 +405,31 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): if policy is not None: await self._manager.release_session() + async def _answer_until_disconnect(self, websocket: WebSocket, session: Session, session_id: str) -> None: + """Answer one observation at a time until the client goes away.""" + try: + while True: + message = await websocket.receive_bytes() + self._last_activity = time.monotonic() + try: + raw_obs = deserialise(message) + if self._frames is not None: + raw_obs = self._frames.resolve(session_id, raw_obs) + # 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: + logger.info('Client disconnected') + async def _startup(self): + if self._frames is not None: + self._frames.start(self.claim_socket_path(self._frames.path, socket.SOCK_SEQPACKET)) self._default_id = self._source.resolve(None) logger.info(f'Pinned default checkpoint at startup: {self._default_id}') await self._manager.get_policy(self._default_id) @@ -423,8 +453,8 @@ async def _idle_watchdog(self, server: uvicorn.Server): 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. + def _is_stale_socket(path: str, kind: int) -> bool: + """Whether ``path`` is a socket of type ``kind`` no server answers on, so replacing it takes nothing. 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. @@ -434,7 +464,7 @@ def _is_stale_socket(path: str) -> bool: return False except FileNotFoundError: return False - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: + with socket.socket(socket.AF_UNIX, kind) as probe: probe.settimeout(PolicyServer.LIVE_SOCKET_PROBE_SEC) try: probe.connect(path) @@ -445,21 +475,21 @@ def _is_stale_socket(path: str) -> bool: return False @staticmethod - def claim_socket_path(path: str) -> socket.socket: + def claim_socket_path(path: str, kind: int = socket.SOCK_STREAM) -> 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) + sock = socket.socket(socket.AF_UNIX, kind) try: try: sock.bind(path) except OSError as taken: if taken.errno != errno.EADDRINUSE: raise - if not PolicyServer._is_stale_socket(path): + if not PolicyServer._is_stale_socket(path, kind): raise OSError(errno.EADDRINUSE, f'{path!r} is already in use') from None os.unlink(path) sock.bind(path) @@ -498,10 +528,12 @@ async def _run(): except KeyboardInterrupt: logger.info('Server stopped by user') finally: + if self._frames is not None: + self._frames.close() self._manager.close() -@cfn.config(host='0.0.0.0', port=8000, recording_dir=None, idle_timeout_min=None, uds=None) +@cfn.config(host='0.0.0.0', port=8000, recording_dir=None, idle_timeout_min=None, uds=None, frame_ring=True) def serve( pipeline: cfn.Config, host: str, @@ -509,6 +541,7 @@ def serve( recording_dir: str | None, idle_timeout_min: float | None, uds: str | None, + frame_ring: bool, ): """The CLI entry point every vendor server exposes: bind ``pipeline``, and the commands are configs of this. @@ -516,7 +549,8 @@ def serve( codec, source, checkpoint directory — is reached through the pipeline itself (``--pipeline.source.checkpoints_dir=...``), so each of those values has exactly one name. - ``--uds`` binds that Unix socket path and leaves ``host`` and ``port`` unused. + ``--uds`` binds that Unix socket path and leaves ``host`` and ``port`` unused. ``--frame_ring=false`` + keeps every image in the message. 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. @@ -529,4 +563,5 @@ def serve( idle_timeout_min=idle_timeout_min, auth_token=os.environ.get(AUTH_TOKEN_ENV), uds=uds, + frame_ring=frame_ring, ).serve() diff --git a/positronic/offboard/tests/bench_frame_ring.py b/positronic/offboard/tests/bench_frame_ring.py new file mode 100644 index 000000000..ff5d42853 --- /dev/null +++ b/positronic/offboard/tests/bench_frame_ring.py @@ -0,0 +1,143 @@ +"""Measure what one inference costs when the frames ride a shared-memory ring, and when they do not. + +The server runs in this process on a Unix socket, and the served session reads every pixel it is +given, as a codec does. So each arm pays for touching the frames, and the boundary costs the +difference between them. + +Usage + uv run --locked python -m positronic.offboard.tests.bench_frame_ring + uv run --locked python -m positronic.offboard.tests.bench_frame_ring --calls 100 --cameras 3 +""" + +import argparse +import asyncio +import statistics +import tempfile +import threading +import time +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import uvicorn + +from positronic import keys +from positronic.offboard import client +from positronic.offboard import keys as offboard_keys +from positronic.offboard.client import InferenceClient +from positronic.offboard.server import PolicyServer +from positronic.policy import Policy, Session +from positronic.policy.layers import ChunkedSchedule +from positronic.policy.spec import PolicySource, remote + +HD720 = (720, 1280, 3) + + +class _ReadEveryPixel(Session): + """A session that reads each frame once and answers, so both arms pay the same read.""" + + def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]]: + for value in obs.values(): + if isinstance(value, np.ndarray): + int(value.sum()) + return [{keys.ACTION_TIMESTAMP: 0.0}] + + @property + def meta(self) -> dict[str, Any]: + return {} + + def close(self) -> None: + pass + + +class _StubPolicy(Policy): + def new_session(self, context=None, rt=None) -> Session: + return _ReadEveryPixel() + + @property + def functions(self): + return {} + + def close(self) -> None: + pass + + +def _serve(socket_path: str, frame_ring: bool) -> tuple[PolicyServer, uvicorn.Server, threading.Thread]: + server = PolicyServer( + ChunkedSchedule() | remote | PolicySource(_StubPolicy()), uds=socket_path, frame_ring=frame_ring + ) + PolicyServer.claim_socket_path(socket_path) + uv_server = uvicorn.Server(uvicorn.Config(server.app, uds=socket_path, log_level='error')) + + async def run(): + await server._startup() + await uv_server.serve() + + thread = threading.Thread(target=asyncio.run, args=(run(),), daemon=True) + thread.start() + while not uv_server.started: + time.sleep(0.02) + return server, uv_server, thread + + +def _measure(socket_path: str, frame_ring: bool, calls: int, cameras: int) -> tuple[list[float], int, bool]: + server, uv_server, thread = _serve(socket_path, frame_ring) + sent: list[int] = [] + packer = client.serialise + client.serialise = lambda obj: _record(packer(obj), sent) + try: + session = InferenceClient(f'unix://{socket_path}').new_session() + declared = offboard_keys.FRAME_RING in session.metadata + obs: dict[str, Any] = { + f'image.{i}': np.random.default_rng(i).integers(0, 256, HD720, dtype=np.uint8) for i in range(cameras) + } + obs[keys.GRIP] = 0.5 + for _ in range(5): + session.infer(obs) + times = [] + for _ in range(calls): + start = time.perf_counter() + session.infer(obs) + times.append((time.perf_counter() - start) * 1e3) + message = max(sent) + session.close() + finally: + client.serialise = packer + uv_server.should_exit = True + thread.join(timeout=5.0) + if server._frames is not None: + server._frames.close() + return times, message, declared + + +def _record(message: bytes, sent: list[int]) -> bytes: + sent.append(len(message)) + return message + + +def _report(name: str, times: list[float], message: int, declared: bool) -> None: + ordered = sorted(times) + p95 = ordered[min(len(ordered) - 1, int(0.95 * len(ordered)))] + print( + f'{name:<12} median {statistics.median(times):7.2f} ms p95 {p95:7.2f} ms ' + f'min {ordered[0]:7.2f} ms message {message / 1e6:6.2f} MB declared {declared}' + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--calls', type=int, default=50) + parser.add_argument('--cameras', type=int, default=3) + args = parser.parse_args() + + frame = np.prod(HD720) * args.cameras + print(f'{args.cameras} frames of {HD720} per inference, {frame / 1e6:.1f} MB raw, {args.calls} calls') + with tempfile.TemporaryDirectory(dir='/tmp') as directory: + for name, frame_ring in (('message', False), ('ring', True)): + path = str(Path(directory) / f'{name}.sock') + _report(name, *_measure(path, frame_ring, args.calls, args.cameras)) + + +if __name__ == '__main__': + main() diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 92e13fb3b..65eed10db 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -23,7 +23,7 @@ def _find_free_port() -> int: return s.getsockname()[1] -RunningServers = list[tuple[uvicorn.Server, threading.Thread]] +RunningServers = list[tuple[uvicorn.Server, threading.Thread, PolicyServer]] StartServer = Callable[..., tuple[str, int, PolicyServer]] @@ -35,9 +35,12 @@ 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: + for uv_server, thread, server in running: uv_server.should_exit = True thread.join(timeout=5.0) + # A server started here never runs through ``serve``, so nothing else closes its frame channel. + if server._frames is not None: + server._frames.close() def _serve_in_background(server: PolicyServer, config: uvicorn.Config, running: RunningServers) -> None: @@ -49,7 +52,7 @@ async def _run(): thread = threading.Thread(target=asyncio.run, args=(_run(),), daemon=True) thread.start() - running.append((uv_server, thread)) + running.append((uv_server, thread, server)) def _wait_until_it_accepts(dial: Callable[[], None]) -> None: diff --git a/positronic/offboard/tests/test_offboard.py b/positronic/offboard/tests/test_offboard.py index 6967f70a3..1c49ffbad 100644 --- a/positronic/offboard/tests/test_offboard.py +++ b/positronic/offboard/tests/test_offboard.py @@ -1,3 +1,10 @@ +import ctypes +import ctypes.util +import errno +import mmap +import os +import socket +from collections.abc import Generator from types import MappingProxyType from unittest.mock import ANY @@ -14,8 +21,10 @@ to_wire, ) from positronic.geom import Rotation, Transform3D +from positronic.offboard import frame_ring from positronic.offboard.client import InferenceClient from positronic.offboard.protocol import deserialise, serialise, typed_commands +from positronic.offboard.server import PolicyServer from positronic.utils.serialization import encode_jpeg @@ -269,3 +278,133 @@ def test_a_command_crossing_the_wire_arrives_typed_from_either_shape(self): assert isinstance(from_envelope, CartesianPosition) and isinstance(from_bare, CartesianPosition) np.testing.assert_allclose(from_bare.pose.translation, from_envelope.pose.translation, atol=1e-6) + + +SESSION_ID = 'session-under-test' + +pytestmark_ring = pytest.mark.skipif( + not frame_ring.SUPPORTED, reason='a frame ring needs memfd_create, which Linux has and macOS has not' +) + + +@pytest.fixture +def frame_channel(socket_path: str) -> Generator[frame_ring.FrameChannel, None, None]: + channel = frame_ring.FrameChannel(frame_ring.channel_path(socket_path)) + channel.start(PolicyServer.claim_socket_path(channel.path, socket.SOCK_SEQPACKET)) + channel.open_session(SESSION_ID) + yield channel + channel.close() + + +@pytest.fixture +def frame_writer(socket_path: str, frame_channel) -> Generator[frame_ring.FrameWriter, None, None]: + writer = frame_ring.FrameWriter(frame_ring.channel_path(socket_path), SESSION_ID) + yield writer + writer.close() + + +_PIXELS = np.random.default_rng(0) + + +def _image(height: int = 48, width: int = 64) -> np.ndarray: + return _PIXELS.integers(0, 256, (height, width, 3), dtype=np.uint8) + + +def _over_the_wire(obs): + """``obs`` as the server reads it, so a reference is tested through msgpack rather than beside it.""" + return deserialise(serialise(obs)) + + +@pytestmark_ring +def test_a_ring_carries_every_image_of_an_observation_byte_identical(frame_channel, frame_writer): + image, stack = _image(), np.stack([_image(), _image()]) + obs = {'image.left': image, keys.GRIP: 0.5, 'nested': {'frames': [image, stack]}} + + served = frame_channel.resolve(SESSION_ID, _over_the_wire(frame_writer.pack(obs))) + + np.testing.assert_array_equal(served['image.left'], image) + np.testing.assert_array_equal(served['nested']['frames'][0], image) + np.testing.assert_array_equal(served['nested']['frames'][1], stack) + assert served[keys.GRIP] == 0.5 + + +@pytestmark_ring +def test_a_packed_observation_leaves_the_pixels_out_of_the_message(frame_writer): + image = _image(720, 1280) + + message = serialise(frame_writer.pack({'image.left': image})) + + assert len(message) < image.nbytes // 100 + + +@pytestmark_ring +def test_an_observation_with_no_image_creates_no_ring(frame_writer): + assert frame_writer.pack({keys.GRIP: 0.5}) == {keys.GRIP: 0.5} + assert frame_writer._ring is None + + +@pytestmark_ring +def test_a_sealed_ring_refuses_a_write_a_resize_and_a_hole(): + ring = frame_ring.FrameRing(1024, slots=2) + try: + with pytest.raises(PermissionError): + mmap.mmap(ring.fd, 1024, mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE) + with pytest.raises(PermissionError): + os.write(ring.fd, b'x') + with pytest.raises(PermissionError): + os.ftruncate(ring.fd, 1 << 20) + libc = ctypes.CDLL(ctypes.util.find_library('c'), use_errno=True) + # FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE, from linux/falloc.h. + assert libc.fallocate(ring.fd, 0x01 | 0x02, ctypes.c_int64(0), ctypes.c_int64(1024)) == -1 + assert ctypes.get_errno() == errno.EPERM + finally: + ring.close() + + +@pytestmark_ring +def test_a_slot_written_over_before_it_is_read_is_refused(frame_channel, frame_writer): + stale = _over_the_wire(frame_writer.pack({'image.left': _image()})) + for _ in range(frame_ring.SLOTS): + frame_writer.pack({'image.left': _image()}) + + with pytest.raises(frame_ring.TornFrame): + frame_channel.resolve(SESSION_ID, stale) + + +@pytestmark_ring +def test_a_reference_that_points_outside_the_ring_is_refused(frame_channel, frame_writer): + packed = _over_the_wire(frame_writer.pack({'image.left': _image()})) + packed['image.left'][frame_ring._OFFSET] = 1 << 30 + + with pytest.raises(ValueError, match='outside the ring'): + frame_channel.resolve(SESSION_ID, packed) + + +@pytestmark_ring +def test_a_reference_from_a_session_that_handed_over_no_ring_is_refused(frame_channel): + frame_channel.open_session('another-session') + + with pytest.raises(RuntimeError, match='before it handed over a ring'): + frame_channel.resolve('another-session', {frame_ring._RING: True}) + + +@pytestmark_ring +def test_a_larger_frame_grows_the_ring_and_leaves_the_earlier_view_readable(frame_channel, frame_writer): + small, large = _image(8, 8), _image(64, 64) + + first = frame_channel.resolve(SESSION_ID, _over_the_wire(frame_writer.pack({'image.left': small}))) + second = frame_channel.resolve(SESSION_ID, _over_the_wire(frame_writer.pack({'image.left': large}))) + + np.testing.assert_array_equal(second['image.left'], large) + np.testing.assert_array_equal(first['image.left'], small) + + +@pytestmark_ring +def test_closing_a_channel_that_waits_for_a_handover_ends_its_thread(socket_path): + channel = frame_ring.FrameChannel(frame_ring.channel_path(socket_path)) + channel.start(PolicyServer.claim_socket_path(channel.path, socket.SOCK_SEQPACKET)) + accepting = channel._thread + + channel.close() + + assert accepting is not None and not accepting.is_alive() diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index d33bcd214..1a52584b1 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -102,7 +102,9 @@ def test_new_session_passes_additional_headers(self): mock_connect.assert_called_once() assert mock_connect.call_args.kwargs['additional_headers'] == headers - mock_session_cls.assert_called_once_with(mock_connect.return_value, infer_timeout=DEFAULT_INFER_TIMEOUT) + mock_session_cls.assert_called_once_with( + mock_connect.return_value, infer_timeout=DEFAULT_INFER_TIMEOUT, uds=None + ) def test_new_session_without_headers_passes_none(self): with ( diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 27ce2f44a..4912848dd 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -12,11 +12,13 @@ import configuronic as cfn import httpx +import numpy as np import pytest from websockets.exceptions import InvalidStatus from websockets.sync.client import connect from positronic import keys +from positronic.offboard import client, frame_ring from positronic.offboard import keys as offboard_keys from positronic.offboard.client import InferenceClient, InferenceSession, _ConnectRetries from positronic.offboard.protocol import deserialise @@ -726,3 +728,92 @@ def test_a_non_ascii_authorization_header_is_refused_rather_than_crashing(start_ ) status = sock.recv(64).split(b' ')[1] assert status == b'401' + + +def _messages_the_client_sends(monkeypatch: pytest.MonkeyPatch) -> list[bytes]: + """Every observation message ``InferenceSession.infer`` puts on the wire.""" + sent: list[bytes] = [] + packer = client.serialise + + def record(obj): + message = packer(obj) + sent.append(message) + return message + + monkeypatch.setattr(client, 'serialise', record) + return sent + + +def _frame() -> np.ndarray: + return np.random.default_rng(0).integers(0, 256, (240, 320, 3), dtype=np.uint8) + + +@pytest.mark.skipif(not frame_ring.SUPPORTED, reason='a frame ring needs memfd_create, which macOS has not') +def test_a_unix_session_carries_its_frames_through_shared_memory(unix_stub_server, monkeypatch): + socket_path, policy = unix_stub_server + image = _frame() + sent = _messages_the_client_sends(monkeypatch) + + session = InferenceClient(f'unix://{socket_path}').new_session() + try: + assert offboard_keys.FRAME_RING in session.metadata + assert session.infer({'image.left': image, keys.GRIP: 0.5}) == [{'action': [1, 2, 3]}] + finally: + session.close() + + served = policy._mock_session.call_args.args[0] + np.testing.assert_array_equal(served['image.left'], image) + assert served[keys.GRIP] == 0.5 + assert max(len(message) for message in sent) < image.nbytes // 100 + + +@pytest.mark.skipif(not frame_ring.SUPPORTED, reason='a frame ring needs memfd_create, which macOS has not') +def test_a_ring_hands_the_server_a_view_it_cannot_write(unix_stub_server): + socket_path, policy = unix_stub_server + + session = InferenceClient(f'unix://{socket_path}').new_session() + try: + session.infer({'image.left': _frame()}) + finally: + session.close() + + served = policy._mock_session.call_args.args[0]['image.left'] + assert not served.flags.writeable + + +def test_a_unix_session_with_the_ring_off_carries_its_frames_in_the_message( + start_unix_server, socket_path, make_mock_policy, monkeypatch +): + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + start_unix_server(ChunkedSchedule() | remote | _StubSource(policy), socket_path, frame_ring=False) + image = _frame() + sent = _messages_the_client_sends(monkeypatch) + + session = InferenceClient(f'unix://{socket_path}').new_session() + try: + assert offboard_keys.FRAME_RING not in session.metadata + session.infer({'image.left': image}) + finally: + session.close() + + np.testing.assert_array_equal(policy._mock_session.call_args.args[0]['image.left'], image) + assert max(len(message) for message in sent) > image.nbytes + + +def test_a_server_on_a_port_declares_no_frame_ring(stub_server): + host, port, _server, _policy = stub_server + + session = InferenceClient(f'{host}:{port}').new_session() + try: + assert offboard_keys.FRAME_RING not in session.metadata + finally: + session.close() + + +def test_a_companion_path_too_long_to_bind_declares_no_frame_ring(make_mock_policy): + policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) + uds = '/tmp/' + 'p' * (frame_ring.MAX_SOCKET_PATH - len('/tmp/')) + + server = PolicyServer(ChunkedSchedule() | remote | _StubSource(policy), uds=uds) + + assert server._frames is None diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 78c0f6bb2..8319fbfd6 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -27,6 +27,7 @@ from positronic.drivers.roboarm.models import DEFAULT_FRAME from positronic.policy.base import PAR, SEQ, DelegatingSession, Layer, Session, _ComposedLayer from positronic.utils import merge_dicts +from positronic.utils.serialization import is_image _QUAT = geom.Rotation.Representation.QUAT @@ -467,8 +468,7 @@ def encode(self, data): return {key: self._restrict(key, value) for key, value in data.items()} def _restrict(self, key: str, value: Any) -> Any: - # Codecs nest images inside dicts and lists (e.g. GR00T), so recurse to reach every image array. - if isinstance(value, np.ndarray) and value.ndim in (3, 4) and value.shape[-1] == 3: + if is_image(value): # A TemporalStack emits a (T, H, W, 3) stack, so bound each frame rather than the stack's first axis. if value.ndim == 4: return np.stack([_scaled(frame, self._width, self._height) for frame in value]) diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index b6a4c7b8e..a53c25afc 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -3,7 +3,6 @@ import time from typing import Any -import numpy as np import pos3 from positronic import telemetry, telemetry_keys @@ -11,7 +10,7 @@ from positronic.offboard.client import DEFAULT_INFER_TIMEOUT, InferenceClient, InferenceSession from positronic.policy import keys as policy_keys from positronic.utils import flatten_dict -from positronic.utils.serialization import encode_jpeg +from positronic.utils.serialization import encode_jpeg, is_image from .base import Answer, Layer, Policy, Runtime, Session from .recording import Recorder @@ -24,8 +23,7 @@ def _prepare_value(value: Any) -> Any: - # Codecs nest images inside dicts and lists (e.g. GR00T), so recurse to reach every image array. - if isinstance(value, np.ndarray) and value.ndim in (3, 4) and value.shape[-1] == 3: + if is_image(value): # A raw HD frame — especially a (T, H, W, 3) stack — can exceed a proxy's websocket message cap. return encode_jpeg(value) if isinstance(value, cabc.Mapping): diff --git a/positronic/utils/serialization.py b/positronic/utils/serialization.py index 09948c22a..c8c50f403 100644 --- a/positronic/utils/serialization.py +++ b/positronic/utils/serialization.py @@ -35,6 +35,11 @@ _JPEG_QUALITY = 90 +def is_image(value: Any) -> bool: + """True for a value the wire treats as an image: an ``(H, W, 3)`` frame or a ``(T, H, W, 3)`` stack.""" + return isinstance(value, np.ndarray) and value.ndim in (3, 4) and value.shape[-1] == 3 + + def encode_jpeg(image: np.ndarray) -> dict[bytes, Any]: """JPEG-encode a single ``(H, W, 3)`` image or a ``(T, H, W, 3)`` stack to a compact wire marker.