Report what the server spent, beside the answer it sends

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:

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):

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)}))

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):

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)

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'