Time the rig-side stack, and bind telemetry on a real run

positronic/cli/eval/tests/test_timing_report.py

@@ -15,6 +15,7 @@
 from positronic.simulator.env_server.telemetry import ENV_PROCESS
 from positronic.telemetry import (
     ATTR_PROCESS_NAME,
+    ATTR_RUN_ID,
     GPU_INDEX,
     GPU_MEM_USED_B,
     GPU_PROC_MEM_B,
@@ -41,7 +42,7 @@
 _S = 1_000_000_000  # seconds -> ns
 
 
-def _span(name, start_s, end_s, span_id, parent_id=None, attrs=None, process=HARNESS_PROCESS):
+def _span(name, start_s, end_s, span_id, parent_id=None, attrs=None, process=HARNESS_PROCESS, run_id=''):
     encoded = {
         'traceId': '0' * 32,
         'spanId': span_id,
@@ -52,14 +53,10 @@ def _span(name, start_s, end_s, span_id, parent_id=None, attrs=None, process=HAR
     }
     if parent_id is not None:
         encoded['parentSpanId'] = parent_id
-    return {
-        'resourceSpans': [
-            {
-                'resource': {'attributes': [{'key': ATTR_PROCESS_NAME, 'value': {'stringValue': process}}]},
-                'scopeSpans': [{'spans': [encoded]}],
-            }
-        ]
-    }
+    resource = [{'key': ATTR_PROCESS_NAME, 'value': {'stringValue': process}}]
+    if run_id:
+        resource.append({'key': ATTR_RUN_ID, 'value': {'stringValue': run_id}})
+    return {'resourceSpans': [{'resource': {'attributes': resource}, 'scopeSpans': [{'spans': [encoded]}]}]}
 
 
 def _write_lines(path, docs):
@@ -249,6 +246,29 @@ def test_two_killed_runs_in_one_directory_get_a_window_each(tmp_path):
     assert report.wall_split.between_episodes == pytest.approx(0.0)
 
 
+def test_two_attended_runs_in_one_sidecar_get_a_window_each(tmp_path):
+    """A process names its sidecar, so two attended runs against one telemetry directory write one file. An
+    attended rollout opens no ``eval.pass`` span, so its episodes are roots and the parent they share says
+    nothing about which run wrote them. The run id does — and without it one window spans both runs and
+    reports 1040 s of wall for 80 s of work."""
+    telemetry_dir = tmp_path / TELEMETRY_SUBDIR
+    telemetry_dir.mkdir()
+    _write_lines(
+        telemetry_dir / f'{HARNESS_PROCESS}{SPANS_SUFFIX}',
+        [
+            _span(SPAN_EPISODE, start, end, f'ep-{run}', attrs={ATTR_EPISODE_VIRTUAL_S: 20.0}, run_id=run)
+            for run, (start, end) in (('rik-0', (0, 40)), ('rik-1', (1000, 1040)))
+        ],
+    )
+
+    report = _build_report(_read_spans_dir(telemetry_dir), [], policy_gpu=None)
+
+    assert report.episodes == 2
+    assert report.window is WallWindow.W_EPISODES
+    assert report.wall_s == pytest.approx(80.0)
+    assert report.wall_split.between_episodes == pytest.approx(0.0)
+
+
 def test_telemetry_with_neither_a_pass_nor_an_episode_names_what_is_missing(tmp_path):
     """Spans that carry no closed pass and no episode leave nothing to reduce. The refusal must say so: a run
     killed before its first episode finished did record telemetry, so blaming a missing ``--timing`` sends the

positronic/cli/eval/timing_report.py

@@ -16,6 +16,7 @@
 from dataclasses import asdict, dataclass, fields
 from enum import StrEnum
 from pathlib import Path
+from typing import NamedTuple
 
 import configuronic as cfn
 import numpy as np
@@ -440,19 +441,28 @@ def in_episode(ts_ns: int) -> bool:
     )
 
 
-def _episode_windows(episodes: list[SpanRec]) -> dict[str | None, tuple[int, int]]:
-    """One wall window per run whose ``eval.pass`` span never closed, keyed by the pass span the episodes name
-    as their parent — from that run's first episode start to its last episode end.
+class _WindowKey(NamedTuple):
+    """The run, and the pass span within it, that one wall window covers.
 
-    Grouping by parent is what keeps two killed runs appended to one directory apart: each contributes its own
+    The parent alone does not identify a run: an attended rollout opens no ``eval.pass`` span, so every one of
+    its episodes is a root, and two such runs appended to one directory share the ``None`` parent.
+    """
+
+    run_id: str
+    parent_id: str | None
+
+
+def _episode_windows(episodes: list[SpanRec]) -> dict[_WindowKey, tuple[int, int]]:
+    """One wall window per run whose ``eval.pass`` span never closed, from that run's first episode start to
+    its last episode end.
+
+    Grouping by run and parent keeps two runs appended to one file apart: each contributes its own
     window, so the dead wall between them falls outside both, exactly as the gap between two pass spans does.
     """
-    by_parent: dict[str | None, list[SpanRec]] = defaultdict(list)
+    by_window: dict[_WindowKey, list[SpanRec]] = defaultdict(list)
     for episode in episodes:
-        by_parent[episode.parent_id].append(episode)
-    return {
-        parent: (min(e.start_ns for e in group), max(e.end_ns for e in group)) for parent, group in by_parent.items()
-    }
+        by_window[_WindowKey(episode.run_id, episode.parent_id)].append(episode)
+    return {key: (min(e.start_ns for e in group), max(e.end_ns for e in group)) for key, group in by_window.items()}
 
 
 def _build_report(spans: list[SpanRec], stats: list[dict], policy_gpu: GpuSummary | None) -> PassReport:
@@ -469,7 +479,7 @@ def _build_report(spans: list[SpanRec], stats: list[dict], policy_gpu: GpuSummar
     passes = [p for p in spans if p.name == SPAN_EVAL_PASS]
     if passes:
         window = WallWindow.W_PASS
-        windows = {p.span_id: (p.start_ns, p.end_ns) for p in passes}
+        windows = {_WindowKey(p.run_id, p.span_id): (p.start_ns, p.end_ns) for p in passes}
     else:
         window = WallWindow.W_EPISODES
         windows = _episode_windows(all_episodes)
@@ -501,12 +511,12 @@ def _build_report(spans: list[SpanRec], stats: list[dict], policy_gpu: GpuSummar
         logger.warning(
             '%d episode(s) did not run to completion (a failed pass?); their finished phases are included', partial
         )
-    orphans = sum(e.parent_id not in windows for e in all_episodes)
+    orphans = sum(_WindowKey(e.run_id, e.parent_id) not in windows for e in all_episodes)
     if orphans:
         logger.warning('%d episode(s) belong to no completed pass (a killed run?); excluded from the report', orphans)
     # Only episodes belonging to a window reduce: where one pass completed, a killed earlier run's episodes are
     # in the same reused directory but under no window of their own, and would inflate every normalised figure.
-    episodes = [e for e in all_episodes if e.parent_id in windows]
+    episodes = [e for e in all_episodes if _WindowKey(e.run_id, e.parent_id) in windows]
     timings = [_episode_timing(e, children) for e in episodes]
 
     episode_wall_sum = float(sum(t.wall_s for t in timings))

positronic/offboard/client.py

@@ -11,6 +11,8 @@
 from websockets.sync.client import connect
 from websockets.sync.connection import Connection
 
+from positronic import telemetry, telemetry_keys
+
 from . import protocol
 from .protocol import deserialise, serialise, typed_commands
 
@@ -70,10 +72,15 @@ def infer(self, obs: dict[str, Any]) -> Any:
         """
         serialised = serialise(obs)
         logger.debug('Size of serialised obs: %1.f KiB', len(serialised) / 1024)
+        # The pair reads as the uplink and then the wait the server's own time sits inside: each span
+        # holds the socket alone. A send outlasting its own bytes is an uplink too slow for the payload.
+        wire_bytes = {telemetry_keys.ATTR_WIRE_BYTES: len(serialised)}
+        with telemetry.span(telemetry_keys.SPAN_WIRE_SEND, **wire_bytes):
+            self._websocket.send(serialised)
 
-        self._websocket.send(serialised)
         try:
-            response = deserialise(self._websocket.recv(timeout=self._infer_timeout))
+            with telemetry.span(telemetry_keys.SPAN_WIRE_RECV):
+                received = self._websocket.recv(timeout=self._infer_timeout)
         except TimeoutError:
             # The observation is in flight but unanswered; the server's late response would sit in the socket and
             # the next ``recv`` would pair it with a future observation. Close so the desynced session can't be
@@ -82,6 +89,7 @@ def infer(self, obs: dict[str, Any]) -> Any:
             raise TimeoutError(
                 f'No inference response within {self._infer_timeout}s — server stalled or connection half-open'
             ) from None
+        response = deserialise(received)
         logger.debug('Size of deserialised response: %1.f KiB', len(response) / 1024)
 
         if isinstance(response, dict) and protocol.ERROR in response:

positronic/offboard/tests/test_remote_policy.py

@@ -423,14 +423,14 @@ def blocked(obs):
 
 
 def test_records_infer_span_without_scheduling_layer(tmp_path, open_session):
-    """The ``policy.infer`` span is recorded at the remote inference boundary itself, not by a layer in
-    front of it."""
+    """The remote inference boundary records ``policy.infer``, and the preparation before it records
+    ``policy.prepare``."""
     endpoint, _ = _mock_endpoint(infer_return=[{'a': 1, 'timestamp': 0.0}])
     session, rt = open_session(endpoint)
     with telemetry.bind(tmp_path, telemetry_keys.HARNESS_PROCESS, 'run-infer-span'):
         assert round_trip(session, rt, {keys.OBS_TIME_NS: 0}) is not None
     spans = list(telemetry.read_spans(telemetry.spans_path(tmp_path, telemetry_keys.HARNESS_PROCESS)))
-    assert [s.name for s in spans] == [telemetry_keys.SPAN_POLICY_INFER]
+    assert {s.name for s in spans} == {telemetry_keys.SPAN_POLICY_PREPARE, telemetry_keys.SPAN_POLICY_INFER}
 
 
 def test_infer_span_excludes_client_side_image_preparation(tmp_path, open_session):
@@ -449,10 +449,12 @@ def _stamp_encode(image):
         with telemetry.bind(tmp_path, telemetry_keys.HARNESS_PROCESS, 'run-infer-prep'):
             round_trip(session, rt, {'cam': _make_image(48, 64), keys.OBS_TIME_NS: 0})
 
-    (span,) = telemetry.read_spans(telemetry.spans_path(tmp_path, telemetry_keys.HARNESS_PROCESS))
-    assert span.name == telemetry_keys.SPAN_POLICY_INFER
+    spans = {s.name: s for s in telemetry.read_spans(telemetry.spans_path(tmp_path, telemetry_keys.HARNESS_PROCESS))}
     assert encoded_at, 'the observation carried an image to compress'
-    assert span.start_ns >= encoded_at[-1]  # every encode finishes before the span opens, not inside it
+    # Every encode finishes before the infer span opens, and falls inside the span that does measure it.
+    assert spans[telemetry_keys.SPAN_POLICY_INFER].start_ns >= encoded_at[-1]
+    prepare = spans[telemetry_keys.SPAN_POLICY_PREPARE]
+    assert prepare.start_ns <= encoded_at[0] and prepare.end_ns >= encoded_at[-1]
 
 
 def test_records_infer_span_when_inference_raises(tmp_path, open_session):
@@ -465,7 +467,7 @@ def test_records_infer_span_when_inference_raises(tmp_path, open_session):
         with pytest.raises(TimeoutError):
             round_trip(session, rt, {keys.OBS_TIME_NS: 0})
     spans = list(telemetry.read_spans(telemetry.spans_path(tmp_path, telemetry_keys.HARNESS_PROCESS)))
-    assert [s.name for s in spans] == [telemetry_keys.SPAN_POLICY_INFER]
+    assert telemetry_keys.SPAN_POLICY_INFER in {s.name for s in spans}
 
 
 def test_missing_declaration_fails_before_motion():

positronic/policy/codec.py

@@ -17,7 +17,7 @@
 import numpy as np
 from PIL import Image as PilImage
 
-from positronic import geom
+from positronic import geom, telemetry, telemetry_keys
 from positronic import keys as obs_keys
 from positronic.dataset.transforms import Elementwise, lazy_sequence
 from positronic.dataset.transforms.episode import Derive, EpisodeTransform, FromValue, Group, Identity
@@ -115,7 +115,9 @@ def __init__(self, inner: Session, codec: 'Codec'):
         self._codec = codec
 
     def __call__(self, obs, time_ns):
-        encoded = self._codec.encode(obs)
+        codec_name = {telemetry_keys.ATTR_CODEC: type(self._codec).__name__}
+        with telemetry.span(telemetry_keys.SPAN_POLICY_ENCODE, **codec_name):
+            encoded = self._codec.encode(obs)
         action = self._inner(encoded, time_ns)
         if action is None:
             return None

positronic/policy/harness.py

@@ -89,7 +89,16 @@ def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None:
         # A call that joins work already in flight keeps its anchor, so the trial pays for that work one time.
         if not self._rollout.rt.in_flight:
             self._t0_ns, self._wall_t0 = now_ns, time.monotonic()
-        return self._rollout.session(frozen_view(self._owned(obs)), now_ns)
+        # FOOTGUN: recorded, not entered. Entering it re-parents ``policy.infer``, which the pass report
+        # reads off the episode's own children.
+        call_start_ns = time.time_ns()
+        trajectory = None
+        try:
+            trajectory = self._rollout.session(frozen_view(self._owned(obs)), now_ns)
+        finally:
+            answered = {telemetry_keys.ATTR_POLICY_ANSWERED: trajectory is not None}
+            telemetry.record_span(telemetry_keys.SPAN_POLICY_CALL, call_start_ns, time.time_ns(), **answered)
+        return trajectory
 
     def wait(self, should_stop: pimm.SignalReceiver[bool]) -> None:
         """Wait for the function in flight, for as long as the trial charges the loop for it."""
@@ -405,7 +414,7 @@ def _trial_terminal(self, done: pimm.Message[dict] | None, clock: pimm.Clock) ->
             return {eval_keys.TERMINATED: False}
         return None
 
-    def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]:
+    def _guarded(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]:
         try:
             yield from self._run(should_stop, clock)
         except BaseException as exc:
@@ -422,6 +431,12 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p
             for call in self.perform_task.incoming():
                 call.set_exception(pimm.calls.HandlerStopped())
 
+    def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]:
+        # FOOTGUN: outside the handler that seals the episode span — leaving this scope shuts the provider
+        # down, and a span ended after that is dropped. Inert unless the env vars are set.
+        with telemetry.bind_from_env(telemetry_keys.HARNESS_PROCESS):
+            yield from self._guarded(should_stop, clock)
+
     def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]:
         while not should_stop.value:
             call = next(self.perform_task.incoming(), None)

positronic/policy/layers.py

@@ -9,7 +9,7 @@
 
 import numpy as np
 
-from positronic import keys
+from positronic import keys, telemetry, telemetry_keys
 from positronic.drivers.roboarm import RobotStatus
 from positronic.policy.base import DelegatingSession, Layer, Session
 
@@ -180,8 +180,11 @@ def __init__(self, inner: Session, keys: tuple[str, ...], offsets_sec: tuple[flo
 
         def __call__(self, obs, time_ns):
             now = _obs_time(obs)
-            self._buffer.append(now, {k: obs[k] for k in self._keys})
-            return self._inner({**obs, **self._buffer.sample(now)}, time_ns)
+            # Every tick pays the whole of this, the ones the scheduling layer below answers included.
+            with telemetry.span(telemetry_keys.SPAN_POLICY_STACK):
+                self._buffer.append(now, {k: obs[k] for k in self._keys})
+                stacked = self._buffer.sample(now)
+            return self._inner({**obs, **stacked}, time_ns)
 
         def cancel(self):
             self._buffer.reset()

positronic/policy/remote.py

@@ -50,7 +50,8 @@ def round_trip(
     stack must not run on the thread that calls the session. The span starts after it, because that
     encode is not inference.
     """
-    prepared = _prepare_obs(obs, compress_images)
+    with telemetry.span(telemetry_keys.SPAN_POLICY_PREPARE):
+        prepared = _prepare_obs(obs, compress_images)
     infer_start_ns = time.time_ns()
     try:
         return ws_session.infer(prepared)

positronic/policy/tests/test_harness.py

@@ -29,6 +29,7 @@
 from positronic.policy.harness import POLL_PERIOD_SEC, Harness, Rollout, _EpisodeInference
 from positronic.policy.layers import ChunkedSchedule, StopOnFault
 from positronic.policy.remote import INFER, RemoteSession, round_trip
+from positronic.simulator.env_server.telemetry import ENV_RUN_ID, ENV_TELEMETRY_DIR
 from positronic.tests.testing_coutils import EpisodeCaller, ManualDriver, RecordingEmitter, drive_scheduler, drive_until
 
 POLL_PERIOD_NS = round(POLL_PERIOD_SEC * 1e9)
@@ -1938,6 +1939,37 @@ def test_an_inference_outliving_its_episode_parents_to_it(world, tmp_path):
 
 
 @pytest.mark.timeout(3.0)
+def test_seal_exports_when_the_harness_owns_the_provider(world, tmp_path, monkeypatch):
+    """The seal runs while the provider it writes to is still bound, so the sealed span is exported."""
+    monkeypatch.setenv(ENV_TELEMETRY_DIR, str(tmp_path / telemetry.TELEMETRY_SUBDIR))
+    monkeypatch.setenv(ENV_RUN_ID, 'run-crash')
+
+    policy = StubPolicy()
+    scene = pimm.calls.ControlSystemHandler[Any, None](Passive())
+    harness = Harness(make_embodiment(prepare_handlers={eval_keys.SCENE: scene}))
+    wire_call(world, harness.prepare[eval_keys.SCENE], scene)
+    harness.ds_command._bind(RecordingEmitter())
+    task = Task(
+        instruction_source='stack',
+        timeout_sec=10.0,
+        prepare_args={eval_keys.SCENE: {}},
+        meta={eval_keys.TRIAL_INDEX: 0},
+    )
+    _ask(world, harness, policy, task)
+    stop = SimpleNamespace(value=False)
+    clock = _ManualClock()
+
+    with pytest.raises(RuntimeError, match='reset boom'):
+        for _ in harness.run(cast(pimm.SignalReceiver, stop), cast(pimm.Clock, clock)):
+            for call in scene.incoming():
+                call.set_exception(RuntimeError('reset boom'))
+
+    path = telemetry.spans_path(tmp_path, telemetry_keys.HARNESS_PROCESS)
+    episodes = [s for s in telemetry.read_spans(path) if s.name == telemetry_keys.SPAN_EPISODE]
+    assert len(episodes) == 1
+    assert episodes[0].attrs.get(telemetry_keys.ATTR_EPISODE_PARTIAL) is True
+
+
 def test_failed_pass_seals_open_episode_span(world, tmp_path):
     """A ``reset`` raising after the episode span was opened must seal that span before the
     provider flushes on exit. Ending it is what exports it at all: an unended span never leaves the batch