diff --git a/positronic/cli/eval/tests/test_timing_report.py b/positronic/cli/eval/tests/test_timing_report.py index d577b25ee..e8d4f96ef 100644 --- a/positronic/cli/eval/tests/test_timing_report.py +++ b/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 diff --git a/positronic/cli/eval/timing_report.py b/positronic/cli/eval/timing_report.py index 78f138c52..0b9ea3e52 100644 --- a/positronic/cli/eval/timing_report.py +++ b/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)) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 98369c4f7..1859a9222 100644 --- a/positronic/offboard/client.py +++ b/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: diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 9f2f6201e..add6375c3 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/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(): diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 78c0f6bb2..564907400 100644 --- a/positronic/policy/codec.py +++ b/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 diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 2a2ae2811..44af25e78 100644 --- a/positronic/policy/harness.py +++ b/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) diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index 530039e06..3e7c895b2 100644 --- a/positronic/policy/layers.py +++ b/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() diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index b6a4c7b8e..002ebce99 100644 --- a/positronic/policy/remote.py +++ b/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) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index b05bed4c8..92cf5cf3d 100644 --- a/positronic/policy/tests/test_harness.py +++ b/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 diff --git a/positronic/telemetry.py b/positronic/telemetry.py index fe308b029..33346de9e 100644 --- a/positronic/telemetry.py +++ b/positronic/telemetry.py @@ -35,8 +35,9 @@ import socket import threading import time +import uuid from collections.abc import Callable, Generator, Iterator -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from contextvars import ContextVar from pathlib import Path from typing import TYPE_CHECKING, Any, NamedTuple @@ -52,6 +53,8 @@ ATTR_PROCESS_NAME, ATTR_PROCESS_PID, ATTR_RUN_ID, + ENV_RUN_ID, + ENV_TELEMETRY_DIR, SPANS_SUFFIX, ) @@ -149,10 +152,10 @@ def _encode_attrs(attrs: dict[str, Any]) -> dict[str, Any]: @contextmanager -def bind(out_dir: Path | str, process: str, run_id: str) -> Generator['TracerProvider', None, None]: - """Provider lifecycle for one process's telemetry: stream spans to ``.spans.jsonl`` under a - resource block carrying this process's identity, and register the provider so ``span`` records. The batch - processor is flushed and shut down on exit — an abrupt exit would otherwise lose its queued tail.""" +def _bind_to(path: Path, process: str, run_id: str) -> Generator['TracerProvider', None, None]: + """Provider lifecycle for one process's telemetry: stream spans to ``path`` under a resource block + carrying this process's identity, and register the provider so ``span`` records. The batch processor is + flushed and shut down on exit — an abrupt exit would otherwise lose its queued tail.""" global _provider try: # the OTel SDK and its file exporter ship in the optional `telemetry` extra from opentelemetry.exporter.otlp.json.file import FileSpanExporter # noqa: PLC0415 @@ -162,7 +165,6 @@ def bind(out_dir: Path | str, process: str, run_id: str) -> Generator['TracerPro from opentelemetry.sdk.trace.sampling import ALWAYS_ON # noqa: PLC0415 except ImportError as error: raise RuntimeError(_MISSING_EXTRA) from error - path = spans_path(out_dir, process) path.parent.mkdir(parents=True, exist_ok=True) resource = Resource.create({ ATTR_RUN_ID: run_id, @@ -189,6 +191,29 @@ def bind(out_dir: Path | str, process: str, run_id: str) -> Generator['TracerPro _provider = None +@contextmanager +def bind(out_dir: Path | str, process: str, run_id: str) -> Generator['TracerProvider', None, None]: + """``_bind_to`` for a run's output directory, which holds the spans at ``.spans.jsonl``.""" + with _bind_to(spans_path(out_dir, process), process, run_id) as provider: + yield provider + + +def bind_from_env(process: str): + """Bind ``process``'s sidecar from the telemetry environment, for a binary that is not the eval CLI. + + The directory turns recording on, and ``process`` names the file, as it does for the env server. This + mints a run id when the environment sets none. Every record holds the run id in its resource block, and + the reduce keys an episode by it. Two runs that share one file therefore stay apart. + + Inert while the directory is unset, and while a provider is already bound. + """ + directory = os.environ.get(ENV_TELEMETRY_DIR) + if directory is None or _provider is not None: + return nullcontext() + run_id = os.environ.get(ENV_RUN_ID) or uuid.uuid4().hex + return _bind_to(Path(directory) / f'{process}{SPANS_SUFFIX}', process, run_id) + + def force_flush() -> None: """Flush the batch processor's queue so a crash after this point loses no already-ended span; the owner of a long-running span flushes as it closes, so a later crash loses at most that span's tail. Inert while @@ -287,8 +312,12 @@ def _seal_truncated_line(path: Path) -> None: class SpanRec(NamedTuple): """One parsed span: hex ``span_id``/``parent_id`` (``parent_id`` is ``None`` for a root), wall-clock - epoch-ns bounds, the flat attribute map, and the recording process's name (the ``process.name`` resource - attribute every sidecar writer stamps — ``''`` when a file carries none).""" + epoch-ns bounds, the flat attribute map, and the recording process's name and run id (the ``process.name`` + and ``run.id`` resource attributes every sidecar writer stamps — ``''`` when a file carries none). + + A reduce reads a telemetry directory whole, so it needs the run id to tell two runs in it apart: a run + that opens no ``eval.pass`` span leaves every episode a root, and no other field separates them. + """ name: str start_ns: int @@ -297,6 +326,7 @@ class SpanRec(NamedTuple): span_id: str parent_id: str | None process: str = '' + run_id: str = '' def _decode_value(value: dict[str, Any]) -> Any: @@ -332,6 +362,7 @@ def read_spans(path: Path | str) -> Iterator[SpanRec]: for resource_spans in doc.get('resourceSpans', []): resource_attrs = _decode_attrs(resource_spans.get('resource', {}).get('attributes', [])) process = str(resource_attrs.get(ATTR_PROCESS_NAME, '')) + run_id = str(resource_attrs.get(ATTR_RUN_ID, '')) for scope_spans in resource_spans.get('scopeSpans', []): for span_data in scope_spans.get('spans', []): yield SpanRec( @@ -342,6 +373,7 @@ def read_spans(path: Path | str) -> Iterator[SpanRec]: span_id=span_data['spanId'], parent_id=span_data.get('parentSpanId') or None, process=process, + run_id=run_id, ) diff --git a/positronic/telemetry_keys.py b/positronic/telemetry_keys.py index 08deed8b2..2a09c4056 100644 --- a/positronic/telemetry_keys.py +++ b/positronic/telemetry_keys.py @@ -26,12 +26,28 @@ SPAN_POLICY_INFER = 'policy.infer' SPAN_RECORD_IO = 'record.io' +# The rig-side stack between the harness and the wire. `policy.call` opens on every control tick, +# including the ones a scheduling layer answers without inferring. +SPAN_POLICY_CALL = 'policy.call' +SPAN_POLICY_STACK = 'policy.stack' +SPAN_POLICY_ENCODE = 'policy.encode' +SPAN_POLICY_PREPARE = 'policy.prepare' +SPAN_WIRE_SEND = 'wire.send' +SPAN_WIRE_RECV = 'wire.recv' + ATTR_EPISODE_INDEX = 'episode.index' ATTR_EPISODE_STEPS = 'episode.steps' ATTR_EPISODE_VIRTUAL_S = 'episode.virtual_s' ATTR_EPISODE_PARTIAL = 'episode.partial' ATTR_PASS_FAILED = 'pass.failed' +# 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' +# Whether a `policy.call` came back with a trajectory. The tick that STARTS a round trip answers False +# too, that trip being asynchronous; it is the tick carrying a `policy.encode`. +ATTR_POLICY_ANSWERED = 'policy.answered' + # The harness process's sidecar name — the discriminator between client-side spans (episode, client env.step) # and an env server's own file, which reduces rely on. HARNESS_PROCESS = 'harness' diff --git a/positronic/tests/test_telemetry.py b/positronic/tests/test_telemetry.py index e451ce226..0b1e06e08 100644 --- a/positronic/tests/test_telemetry.py +++ b/positronic/tests/test_telemetry.py @@ -14,7 +14,7 @@ from opentelemetry.sdk.trace.sampling import ALWAYS_OFF from positronic import telemetry -from positronic.simulator.env_server.telemetry import ENV_PROCESS +from positronic.simulator.env_server.telemetry import ENV_PROCESS, ENV_RUN_ID, ENV_TELEMETRY_DIR from positronic.telemetry_keys import HARNESS_PROCESS, SPAN_EVAL_PASS @@ -173,6 +173,12 @@ def test_unbound_span_is_inert(tmp_path): assert not (tmp_path / 'telemetry').exists() +def _resource_attrs(path): + """The resource block a sidecar stamps on every span, from its first line.""" + line = json.loads(path.read_text().splitlines()[0]) + return telemetry._decode_attrs(line['resourceSpans'][0]['resource']['attributes']) + + def test_resource_carries_process_identity(tmp_path): """Every span document's resource block names the run and the writing process, so a sidecar identifies itself without a second file.""" @@ -180,8 +186,7 @@ def test_resource_carries_process_identity(tmp_path): with telemetry.span('probe'): pass - line = json.loads((telemetry.spans_path(tmp_path, ENV_PROCESS)).read_text().splitlines()[0]) - attrs = telemetry._decode_attrs(line['resourceSpans'][0]['resource']['attributes']) + attrs = _resource_attrs(telemetry.spans_path(tmp_path, ENV_PROCESS)) assert attrs[telemetry.ATTR_RUN_ID] == 'run-1' assert attrs[telemetry.ATTR_PROCESS_NAME] == ENV_PROCESS assert attrs[telemetry.ATTR_PROCESS_PID] == os.getpid() @@ -440,3 +445,60 @@ def test_readers_tolerate_truncated_final_line(tmp_path): stats_path.write_text('{"t_ns": 1, "gpus": []}\n{"t_ns": 2, "cpu_sy') stats = list(telemetry.read_stats(stats_path)) assert [sample[telemetry.STAT_T_NS] for sample in stats] == [1] + + +def test_bind_from_env_is_inert_without_the_env_vars(tmp_path, monkeypatch): + monkeypatch.delenv(ENV_TELEMETRY_DIR, raising=False) + monkeypatch.delenv(ENV_RUN_ID, raising=False) + with telemetry.bind_from_env(HARNESS_PROCESS): + with telemetry.span('client'): + pass + assert not (tmp_path / telemetry.TELEMETRY_SUBDIR).exists() + + +def _run_id(path): + return _resource_attrs(path)[telemetry.ATTR_RUN_ID] + + +def _env_sidecars(tmp_path): + """The sidecars ``bind_from_env`` wrote under a telemetry dir.""" + return sorted((tmp_path / telemetry.TELEMETRY_SUBDIR).glob(f'*{telemetry.SPANS_SUFFIX}')) + + +def test_bind_from_env_mints_a_run_id_when_none_is_given(tmp_path, monkeypatch): + """A run id left unset is minted, so every record stamps one for the reduce to group by.""" + monkeypatch.setenv(ENV_TELEMETRY_DIR, str(tmp_path / telemetry.TELEMETRY_SUBDIR)) + monkeypatch.delenv(ENV_RUN_ID, raising=False) + with telemetry.bind_from_env(HARNESS_PROCESS): + with telemetry.span('client'): + pass + (path,) = _env_sidecars(tmp_path) + assert _run_id(path) + + +def test_bind_from_env_records_under_the_process_it_names(tmp_path, monkeypatch): + """An attended rollout sets two env vars and gets the sidecar the eval CLI writes.""" + monkeypatch.setenv(ENV_TELEMETRY_DIR, str(tmp_path / telemetry.TELEMETRY_SUBDIR)) + monkeypatch.setenv(ENV_RUN_ID, 'rollout-1') + with telemetry.bind_from_env(HARNESS_PROCESS): + with telemetry.span('client'): + pass + (path,) = _env_sidecars(tmp_path) + assert _spans_by_name(path)['client'].process == HARNESS_PROCESS + assert _run_id(path) == 'rollout-1' + + +def test_bind_from_env_defers_to_a_provider_already_bound(tmp_path, monkeypatch): + """The CLI owns the lifecycle under ``eval run --timing``, and keeps it.""" + monkeypatch.setenv(ENV_TELEMETRY_DIR, str(tmp_path / 'ignored' / telemetry.TELEMETRY_SUBDIR)) + monkeypatch.setenv(ENV_RUN_ID, 'from-env') + with telemetry.bind(tmp_path, HARNESS_PROCESS, 'from-the-cli'): + with telemetry.bind_from_env(HARNESS_PROCESS): + with telemetry.span('client'): + pass + with telemetry.span('after'): + pass + path = telemetry.spans_path(tmp_path, HARNESS_PROCESS) + assert set(_spans_by_name(path)) == {'client', 'after'} + assert _run_id(path) == 'from-the-cli' + assert not (tmp_path / 'ignored').exists() diff --git a/pyproject.toml b/pyproject.toml index 56ebf19ea..7ea7851e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -107,7 +107,7 @@ yam = [ "i2rt", # i2rt YAM arm: joint-space position-PD + gravity-comp over CAN; not on PyPI, see [tool.uv.sources] ] telemetry = [ - # `positronic eval run --timing` only: the OTel SDK + exporter that write the span sidecars, and the + # Recording telemetry only: the OTel SDK + exporter that write the span sidecars, and the # machine-load sampler's probes. "nvidia-ml-py", "opentelemetry-exporter-otlp-json-file>=0.65b0",