From 1e474316a7266067d323b41a85bfe101dcb749fb Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Tue, 8 Sep 2026 20:08:11 +0000 Subject: [PATCH] Measure what a served request costs us, with no model behind it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A probe that stands the real server up on loopback behind a model that answers instantly, replays a recorded episode through the rig-side stack and the real client, and divides each round trip into the phases the server reports. On one DROID episode, 25-frame stacks of two cameras bounded to 1024x288 (846 KiB on the wire), a request costs 57 ms and none of it is inference: 27 ms decoding, 30 ms in the websocket transport. Every figure is proportional to the payload — a 5-frame stack costs a fifth of it. The same measurement found uvicorn's default websocket implementation reassembling a received message twice as slowly as its sans-io one, so the server names the faster one. Ticket: none - a probe, and the one-line server change it measured --- positronic/offboard/server.py | 6 +- positronic/offboard/serving_cost.py | 264 ++++++++++++++++++ positronic/offboard/tests/conftest.py | 6 +- .../offboard/tests/test_remote_policy.py | 6 +- .../offboard/tests/test_serving_cost.py | 54 ++++ positronic/policy/remote.py | 4 +- 6 files changed, 332 insertions(+), 8 deletions(-) create mode 100644 positronic/offboard/serving_cost.py create mode 100644 positronic/offboard/tests/test_serving_cost.py diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index 582f4979b..5216e1bda 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -33,6 +33,10 @@ AUTH_HEADER = 'Authorization' +# uvicorn's default ('websockets') reassembles an 846 KiB observation in 58 ms, against 29 ms here +# (measured by positronic/offboard/serving_cost.py). +WS_IMPL = 'websockets-sansio' + def bearer(token: str) -> str: """The ``AUTH_HEADER`` value carrying ``token``.""" @@ -445,7 +449,7 @@ async def _idle_watchdog(self, server: uvicorn.Server): def serve(self): async def _run(): await self._startup() - config = uvicorn.Config(self.app, host=self.host, port=self.port, log_level='info') + config = uvicorn.Config(self.app, host=self.host, port=self.port, log_level='info', ws=WS_IMPL) server = uvicorn.Server(config) self._last_activity = time.monotonic() watchdog = None diff --git a/positronic/offboard/serving_cost.py b/positronic/offboard/serving_cost.py new file mode 100644 index 000000000..153e5e1a6 --- /dev/null +++ b/positronic/offboard/serving_cost.py @@ -0,0 +1,264 @@ +"""What one inference costs the serving path itself, with no model behind it. + +Replays a recorded episode against a real ``PolicyServer`` on loopback whose model answers a fixed +chunk instantly, so every millisecond reported is serving cost, divided by the phases the server +reports. The default stack is the one a DROID endpoint declares: 25 frames of two cameras and the +arm's pose, bounded to 1024x288, JPEG-encoded per frame and re-queried every 24 rows. + +Usage + uv run --locked python -m positronic.offboard.serving_cost \\ + --dataset.path= --requests=20 + ... --compress_images=False # send raw stacks instead of per-frame JPEG + ... --frames=25 --rate_hz=15 --width=1024 --height=288 --chunk_rows=24 --out=rows.json +""" + +import asyncio +import json +import socket +import threading +import time +from collections.abc import Iterable, Iterator, Sequence +from pathlib import Path +from typing import Any + +import configuronic as cfn +import numpy as np +import pos3 +import uvicorn + +import positronic.cfg.ds +from pimm.logging import init_logging +from positronic import keys +from positronic.dataset.dataset import Dataset +from positronic.dataset.episode import Episode +from positronic.offboard import protocol +from positronic.offboard.client import InferenceClient, InferenceSession +from positronic.offboard.server import WS_IMPL, PolicyServer +from positronic.policy.base import DelegatingPolicy, DelegatingSession, Layer, Policy, Session +from positronic.policy.codec import RestrictImageSize +from positronic.policy.layers import ChunkedSchedule, StopOnFault, TemporalStack +from positronic.policy.remote import prepare_obs +from positronic.policy.spec import PolicySource, remote + + +class InstantChunk(Policy): + """The model the probe serves: it answers a fixed chunk of ``rows`` and does no work at all.""" + + def __init__(self, rows: int, period_s: float): + self.chunk = [ + { + keys.ACTION_TIMESTAMP: row * period_s, + keys.TARGET_JOINTS: np.zeros(7, dtype=np.float32), + keys.TARGET_GRIP: np.float32(0.0), + } + for row in range(rows) + ] + + def new_session(self, context: dict[str, Any] | None = None, rt=None) -> Session: + return InstantChunk._Session(self.chunk) + + class _Session(Session): + def __init__(self, chunk: list[dict[str, Any]]): + self._chunk = chunk + + def __call__(self, obs, time_ns): + return self._chunk + + +class CapturingWire(DelegatingPolicy): + """Stands where the wire stands during capture: answers like the server and keeps what it was sent. + + It collects the messages the rig would have put on the wire, and the replay sends those. + """ + + def __init__(self, inner: Policy): + super().__init__(inner) + self.sent: list[dict[str, Any]] = [] + + def new_session(self, context: dict[str, Any] | None = None, rt=None) -> Session: + return CapturingWire._Session(self._inner.new_session(context, rt), self.sent) + + class _Session(DelegatingSession): + def __init__(self, inner: Session, sent: list[dict[str, Any]]): + super().__init__(inner) + self._sent = sent + + def __call__(self, obs, time_ns): + self._sent.append(dict(obs)) + return super().__call__(obs, time_ns) + + +def rig_stack(cameras: Sequence[str], frames: int, rate_hz: float, width: int, height: int) -> Layer: + """The rig-side half a DROID endpoint declares, with the stack depth and image bound the caller names.""" + offsets = tuple(-(frames - 1 - step) / rate_hz for step in range(frames)) + stacked = (*cameras, keys.EE_POSE, keys.GRIP) + return ( + StopOnFault() + | TemporalStack(stacked, offsets) + | ChunkedSchedule() + | RestrictImageSize(width=width, height=height) + ) + + +# What a rig observation carries beside its cameras. The episode holds much more — the arm's URDF, its +# meshes, every recorded command — and none of that crosses the wire. +STATE_KEYS = (keys.JOINTS, keys.JOINT_VEL, keys.EE_POSE, keys.GRIP, keys.ROBOT_STATUS) + + +def observations(episode: Episode, cameras: Sequence[str], rate_hz: float) -> Iterator[dict[str, Any]]: + """The episode as the harness hands it to the stack: one observation per control tick.""" + period_ns = int(1e9 / rate_hz) + for ts in range(episode.start_ts, episode.last_ts + 1, period_ns): + sample = episode.time[ts] + obs = {key: sample[key] for key in (*STATE_KEYS, *cameras) if key in sample} + if keys.TASK in sample: + obs[keys.TASK] = sample[keys.TASK] + yield {**obs, keys.OBS_TIME_NS: ts, keys.WALL_TIME_NS: ts} + + +def capture(ticks: Iterable[dict[str, Any]], stack: Layer, model: Policy, requests: int) -> list[dict[str, Any]]: + """Run the rig-side stack over ``ticks`` and collect the first ``requests`` payloads it sends.""" + wire = CapturingWire(model) + session = stack.wrap(wire).new_session() + try: + for obs in ticks: + session(obs, obs[keys.OBS_TIME_NS]) + if len(wire.sent) >= requests: + break + finally: + session.close() + return wire.sent + + +def serve(pipeline) -> tuple[uvicorn.Server, threading.Thread, int]: + """Serve ``pipeline`` on a free loopback port, and hand back what stops it.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as free_port: + free_port.bind(('127.0.0.1', 0)) + port = free_port.getsockname()[1] + policy_server = PolicyServer(pipeline, host='127.0.0.1', port=port) + served = uvicorn.Server( + uvicorn.Config(policy_server.app, host='127.0.0.1', port=port, log_level='warning', ws=WS_IMPL) + ) + + async def run(): + await policy_server._startup() + await served.serve() + + thread = threading.Thread(target=asyncio.run, args=(run(),), daemon=True) + thread.start() + deadline = time.time() + 30.0 + while time.time() < deadline: + try: + with socket.create_connection(('127.0.0.1', port), timeout=0.2): + return served, thread, port + except OSError: + time.sleep(0.05) + raise RuntimeError('the probe server never came up') + + +def replay(session: InferenceSession, payloads: list[dict[str, Any]], compress_images: bool) -> list[dict[str, float]]: + """Send each payload, and report what its round trip cost beside what the server reports spending.""" + rows = [] + for obs in payloads: + started = time.perf_counter() + prepared = prepare_obs(obs, compress_images) + encoded = time.perf_counter() + # The same pack ``infer`` does next, timed on its own so the round trip below divides. + message = protocol.serialise(prepared) + packed = time.perf_counter() + session.infer(prepared) + answered = time.perf_counter() + served = dict(session.served_timing) + round_trip_ms = (answered - packed) * 1000.0 + pack_ms = (packed - encoded) * 1000.0 + rows.append({ + 'wire_kib': len(message) / 1024.0, + 'prepare_ms': (encoded - started) * 1000.0, + 'pack_ms': pack_ms, + 'round_trip_ms': round_trip_ms, + **served, + # What the round trip spends outside the server's own span: the socket both ways, and the + # websocket receive the server pays before ``served_ms`` opens. + 'transfer_ms': round_trip_ms - pack_ms - served.get(protocol.TIMING_SERVED, 0.0), + }) + return rows + + +def report(rows: list[dict[str, float]]) -> str: + columns = list(rows[0]) + label_width = max(len(column) for column in columns) + lines = [f'{"":<{label_width}} {"median":>9} {"p95":>9} {"max":>9}'] + for column in columns: + values = np.array([row[column] for row in rows]) + lines.append( + f'{column:<{label_width}} {np.median(values):9.1f} {np.percentile(values, 95):9.1f} {values.max():9.1f}' + ) + return '\n'.join(lines) + + +@cfn.config( + dataset=positronic.cfg.ds.local, + episode=0, + requests=20, + frames=25, + rate_hz=15.0, + width=1024, + height=288, + cameras=(keys.WRIST_IMAGE, keys.EXTERIOR_IMAGE), + chunk_rows=24, + compress_images=True, + out=None, +) +def main( + dataset: Dataset, + episode: int, + requests: int, + frames: int, + rate_hz: float, + width: int, + height: int, + cameras: Sequence[str], + chunk_rows: int, + compress_images: bool, + out: str | None, +): + model = InstantChunk(chunk_rows, 1.0 / rate_hz) + stack = rig_stack(cameras, frames, rate_hz, width, height) + + chosen = dataset[episode] + assert isinstance(chosen, Episode), 'name one episode, not a slice of them' + payloads = capture(observations(chosen, cameras, rate_hz), stack, model, requests) + if not payloads: + raise ValueError(f'episode {episode} is shorter than one {chunk_rows}-row chunk; nothing was sent') + print(f'captured {len(payloads)} payload(s) off episode {episode}') + + served, thread, port = serve(stack | remote(compress_images=compress_images) | PolicySource(model)) + try: + session = InferenceClient(f'ws://127.0.0.1:{port}').new_session() + try: + replay(session, payloads[:1], compress_images) # warm up, so no first touch is timed + rows = replay(session, payloads, compress_images) + finally: + session.close() + finally: + served.should_exit = True + thread.join(timeout=10.0) + + print( + f'\n{len(rows)} requests, {frames} frames x {len(cameras)} cameras, bound {width}x{height}, ' + f'compress_images={compress_images}, no model behind the server\n' + ) + print(report(rows)) + if out is not None: + Path(out).write_text(json.dumps(rows, indent=1)) + print(f'\nper-request rows -> {out}') + + +@pos3.with_mirror() +def _internal_main(): + init_logging() + cfn.cli(main) + + +if __name__ == '__main__': + _internal_main() diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 54dfcbb7a..22f410936 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -8,7 +8,7 @@ import pytest import uvicorn -from positronic.offboard.server import PolicyServer +from positronic.offboard.server import WS_IMPL, PolicyServer from positronic.policy import Policy, Session from positronic.policy.executor import Executor from positronic.policy.layers import ChunkedSchedule @@ -31,7 +31,9 @@ def start_server() -> Generator[StartServer, None, None]: 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')) + uv_server = uvicorn.Server( + uvicorn.Config(server.app, host=server.host, port=server.port, log_level='warning', ws=WS_IMPL) + ) async def _run(): await server._startup() diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index add6375c3..537ebf388 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -17,7 +17,7 @@ from positronic.policy import RemotePolicy from positronic.policy.codec import ActionHorizon from positronic.policy.layers import ChunkedSchedule -from positronic.policy.remote import _prepare_obs +from positronic.policy.remote import prepare_obs from positronic.policy.spec import PolicySource, remote # These fixtures stand in for a server, so they spell the handshake fields rather than importing the @@ -59,12 +59,12 @@ class TestPrepareObs: def test_images_pass_through_untouched_by_default(self): obs = {'cam': _make_image(480, 640), 'state': np.array([1.0])} - prepared = _prepare_obs(obs, compress_images=False) + prepared = prepare_obs(obs, compress_images=False) assert prepared.keys() == obs.keys() assert all(prepared[key] is value for key, value in obs.items()) def test_compression_reaches_nested_images(self): - result = _prepare_obs( + result = prepare_obs( { 'cam': _make_image(48, 64), 'video': {'wrist': _make_image(48, 64)}, diff --git a/positronic/offboard/tests/test_serving_cost.py b/positronic/offboard/tests/test_serving_cost.py new file mode 100644 index 000000000..061a68a27 --- /dev/null +++ b/positronic/offboard/tests/test_serving_cost.py @@ -0,0 +1,54 @@ +import numpy as np + +from positronic import keys +from positronic.offboard import protocol +from positronic.offboard.client import InferenceClient +from positronic.offboard.serving_cost import InstantChunk, capture, replay, rig_stack +from positronic.policy.spec import PolicySource, remote + +CAMERAS = (keys.WRIST_IMAGE, keys.EXTERIOR_IMAGE) + + +def _ticks(count: int, period_ns: int = 66_666_666): + """A stand-in for the harness: one moving frame per camera per control tick.""" + rng = np.random.default_rng(0) + for tick in range(count): + frame = rng.integers(0, 255, (48, 64, 3), dtype=np.uint8) + yield { + keys.EE_POSE: np.zeros(7), + keys.GRIP: 0.0, + keys.ROBOT_STATUS: 0, + keys.OBS_TIME_NS: 1_000_000_000 + tick * period_ns, + **dict.fromkeys(CAMERAS, frame), + } + + +def test_replay_divides_a_round_trip_into_the_phases_the_server_reports(start_server): + stack = rig_stack(CAMERAS, frames=3, rate_hz=15.0, width=64, height=48) + model = InstantChunk(rows=2, period_s=1 / 15.0) + payloads = capture(_ticks(12), stack, model, requests=2) + assert payloads, 'the stack sent nothing' + + host, port, _ = start_server(stack | remote(compress_images=True) | PolicySource(model)) + session = InferenceClient(f'ws://{host}:{port}').new_session() + try: + rows = replay(session, payloads, compress_images=True) + finally: + session.close() + + assert len(rows) == len(payloads) + for row in rows: + assert row['wire_kib'] > 0 + assert row[protocol.TIMING_SERVED] >= row[protocol.TIMING_DECODE] + assert row['round_trip_ms'] >= row[protocol.TIMING_SERVED] + + +def test_a_captured_payload_carries_one_stack_per_stacked_key(): + stack = rig_stack(CAMERAS, frames=3, rate_hz=15.0, width=64, height=48) + payloads = capture(_ticks(12), stack, InstantChunk(rows=2, period_s=1 / 15.0), requests=1) + + sent = payloads[0] + for camera in CAMERAS: + assert sent[camera].shape == (3, 48, 64, 3) + assert sent[keys.EE_POSE].shape == (3, 7) + assert sent[keys.GRIP].shape == (3,) diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index c29a700a9..48a41a837 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -32,7 +32,7 @@ def _prepare_value(value: Any) -> Any: return value -def _prepare_obs(obs: cabc.Mapping[str, Any], compress_images: bool) -> dict[str, Any]: +def prepare_obs(obs: cabc.Mapping[str, Any], compress_images: bool) -> dict[str, Any]: if not compress_images: return dict(obs) return {key: _prepare_value(value) for key, value in obs.items()} @@ -48,7 +48,7 @@ def round_trip( encode is not inference. """ with telemetry.span(telemetry_keys.SPAN_POLICY_PREPARE): - prepared = _prepare_obs(obs, compress_images) + prepared = prepare_obs(obs, compress_images) infer_start_ns = time.time_ns() try: return ws_session.infer(prepared)