From 0a0c1ef882c7dff9932efc584787a0a2fac3320c Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Tue, 8 Sep 2026 17:43:40 +0000 Subject: [PATCH] Report what the server spent, beside the answer it sends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A client times its own round trip and cannot see inside it: network and the server's own work arrive as one number. The server now reports what it spent — decode, queue, inference, encode, and the span bracketing them — in the answer's `timing` block, and the client stamps those onto the `policy.infer` span that timed the trip. Subtracting one from the other leaves the network, and both sides measured only their own elapsed time, so no clock has to agree with any other. A server that reports nothing leaves the round trip undivided, which is what a deployment that has not taken this change does. Ticket: Positronic-Robotics/internal#1168 #refs --- positronic/offboard/client.py | 4 +++ positronic/offboard/protocol.py | 12 +++++++ positronic/offboard/server.py | 47 +++++++++++++++++++++---- positronic/policy/remote.py | 4 ++- positronic/policy/tests/test_harness.py | 1 + positronic/telemetry_keys.py | 4 +++ 6 files changed, 65 insertions(+), 7 deletions(-) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index c9596e8a5..c3d5c8a47 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -28,6 +28,9 @@ class InferenceSession: def __init__(self, websocket: Connection, infer_timeout: float = DEFAULT_INFER_TIMEOUT): self._websocket = websocket self._infer_timeout = infer_timeout + # What the server reported spending on the last inference, for the caller that times the round + # trip. Empty against a server that reports nothing, which leaves that round trip undivided. + self.served_timing: dict[str, float] = {} self._metadata = self._handshake() def _handshake(self, timeout_per_message: float = 30.0) -> dict[str, Any]: @@ -90,6 +93,7 @@ def infer(self, obs: dict[str, Any]) -> Any: f'No inference response within {self._infer_timeout}s — server stalled or connection half-open' ) from None response = deserialise(received) + self.served_timing = response.get(protocol.TIMING) or {} if isinstance(response, dict) else {} logger.debug('Size of deserialised response: %1.f KiB', len(response) / 1024) if isinstance(response, dict) and protocol.ERROR in response: diff --git a/positronic/offboard/protocol.py b/positronic/offboard/protocol.py index 20a76f4be..c30809645 100644 --- a/positronic/offboard/protocol.py +++ b/positronic/offboard/protocol.py @@ -24,6 +24,18 @@ META = 'meta' RESULT = 'result' ERROR = 'error' +# What the server spent on one inference, beside the ``RESULT`` it answers with: durations in +# milliseconds on the server's own clock. A server that sends none leaves the round trip undivided. +TIMING = 'timing' + +# The phases ``TIMING`` reports. `SERVED` brackets the other three: it opens on the observation +# arriving and closes as the answer goes back, so a client's round trip minus `SERVED` is network. +TIMING_SERVED = 'served_ms' +TIMING_DECODE = 'decode_ms' +TIMING_INFER = 'infer_ms' +TIMING_ENCODE = 'encode_ms' +# Time the observation waited for the inference slot, inside `SERVED` — a queue rather than compute. +TIMING_QUEUED = 'queued_ms' class ServerStatus(StrEnum): diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index fbe7d33b7..582f4979b 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -7,7 +7,8 @@ import os import time from collections import Counter -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from importlib.metadata import version as _pkg_version from typing import Any @@ -179,6 +180,31 @@ def _declared_stack(local: Layer | None) -> dict[str, Any]: return local.to_spec() +class _ServedTiming: + """What one inference cost the server, in milliseconds on the server's own clock. + + Every figure is a duration. ``served_ms`` opens when the observation arrives and brackets the + phases inside it. + """ + + def __init__(self) -> None: + self._opened = time.perf_counter() + self._phases: dict[str, float] = {} + + @contextmanager + def phase(self, name: str) -> Iterator[None]: + started = time.perf_counter() + try: + yield + finally: + self._phases[name] = (time.perf_counter() - started) * 1000.0 + + def report(self) -> dict[str, float]: + """The phases closed so far, under the span bracketing them. FOOTGUN: read while encoding, so + it carries no encode of its own — the answer's serialisation is outside every figure here.""" + return {protocol.TIMING_SERVED: (time.perf_counter() - self._opened) * 1000.0, **self._phases} + + class PolicyServer: """Serves a policy pipeline: one layer chain with a ``remote`` marker, closed by a ``ModelSource`` (see ``positronic.policy.spec``). @@ -352,13 +378,22 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): message = await websocket.receive_bytes() self._last_activity = time.monotonic() try: - raw_obs = deserialise(message) + timing = _ServedTiming() + with timing.phase(protocol.TIMING_DECODE): + 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})) + with timing.phase(protocol.TIMING_QUEUED): + await self._infer_lock.acquire() + try: + with timing.phase(protocol.TIMING_INFER): + # The server's clock is not the rig's. + actions = await asyncio.to_thread(session, raw_obs, time.time_ns()) + finally: + self._infer_lock.release() + with timing.phase(protocol.TIMING_ENCODE): + answer = serialise({protocol.RESULT: actions, protocol.TIMING: timing.report()}) + await websocket.send_bytes(answer) except Exception as e: logger.error(f'Error processing message: {e}', exc_info=True) await websocket.send_bytes(serialise({protocol.ERROR: str(e)})) diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 9a5e72548..c29a700a9 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -53,7 +53,9 @@ def round_trip( try: return ws_session.infer(prepared) finally: - telemetry.record_span(telemetry_keys.SPAN_POLICY_INFER, infer_start_ns, time.time_ns()) + # Stamped on the span that timed the round trip, which is what a reduce subtracts them from. + served = {f'{telemetry_keys.ATTR_SERVED_PREFIX}{k}': v for k, v in ws_session.served_timing.items()} + telemetry.record_span(telemetry_keys.SPAN_POLICY_INFER, infer_start_ns, time.time_ns(), **served) class RemoteSession(Session): diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 61a0ce367..2fb85f792 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -185,6 +185,7 @@ class _FakeInferenceSession(InferenceSession): def __init__(self, action: list[dict[str, Any]], wall_sec: float = 0.0) -> None: self._action = action self._wall_sec = wall_sec + self.served_timing: dict[str, float] = {} def infer(self, obs: dict[str, Any]) -> list[dict[str, Any]]: time.sleep(self._wall_sec) diff --git a/positronic/telemetry_keys.py b/positronic/telemetry_keys.py index 2a09c4056..5f31e11f8 100644 --- a/positronic/telemetry_keys.py +++ b/positronic/telemetry_keys.py @@ -41,6 +41,10 @@ ATTR_EPISODE_PARTIAL = 'episode.partial' ATTR_PASS_FAILED = 'pass.failed' +# What the server reported spending, stamped on `policy.infer` under this prefix — `served.infer_ms` +# and its siblings, straight from the answer's own `timing` block. +ATTR_SERVED_PREFIX = 'served.' + # Which codec a `policy.encode` span timed, and how many bytes the observation took on the wire. ATTR_CODEC = 'codec' ATTR_WIRE_BYTES = 'wire.bytes'