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'