diff --git a/positronic/eval_timing.py b/positronic/eval_timing.py index 6a04f0a63..78fb4e8e5 100644 --- a/positronic/eval_timing.py +++ b/positronic/eval_timing.py @@ -224,7 +224,7 @@ def record_env_phases(physics_s: float, render_s: float, server_s: float) -> Non timer.add_env_phases(physics_s, render_s, server_s) -def _start_gpu_sampler(out_dir: Path) -> subprocess.Popen | None: +def start_gpu_sampler(out_dir: Path) -> subprocess.Popen | None: """Background ``nvidia-smi dmon`` writing this box's util+memory to ``gpu_dmon.log``. ``None`` when no ``nvidia-smi`` is on PATH (a CPU dev box) — GPU telemetry is then simply absent, not @@ -237,7 +237,7 @@ def _start_gpu_sampler(out_dir: Path) -> subprocess.Popen | None: log_path = out_dir / GPU_LOG_FILENAME log_path.unlink(missing_ok=True) if shutil.which('nvidia-smi') is None: - logger.info('EvalTimer: no nvidia-smi on PATH; skipping GPU sampling') + logger.info('no nvidia-smi on PATH; skipping GPU sampling') return None # Sample only the GPU this eval runs on — the first CUDA-visible device, else device 0. Left unpinned, # dmon logs every visible GPU and ``_parse_dmon`` would average idle/unrelated devices into the numbers. @@ -263,7 +263,7 @@ def bind(out_dir: Path) -> Iterator[EvalTimer]: """Bind a fresh timer (and a GPU sampler) for the enclosed run, then flush ``timing.jsonl`` on exit.""" timer = EvalTimer(out_dir) token = _ACTIVE.set(timer) - sampler = _start_gpu_sampler(out_dir) + sampler = start_gpu_sampler(out_dir) try: yield timer finally: diff --git a/positronic/eval_timing_shim.py b/positronic/eval_timing_shim.py new file mode 100644 index 000000000..25e626b6e --- /dev/null +++ b/positronic/eval_timing_shim.py @@ -0,0 +1,514 @@ +"""Emit ``positronic eval timing-report``-compatible telemetry from a rollout loop positronic does not run. + +Some sims (MolmoSpaces is the first) run their own rollout loop and know nothing about positronic's dataset +format, yet a sizing pass wants the same pass-level wall-clock report ``positronic eval timing-report`` +produces for native evals. This shim is the adapter: the harness wraps it around the rollout loop, marks each +phase (policy wait, env step, reset, record IO) and each episode's outcome, and the shim writes both halves +the report needs — the ``timing.jsonl`` wall-clock records AND the recorded-dataset stubs (block/episode dirs +with ``meta.json`` + ``static.json`` + a timestamped signal) the report joins them against by ``episode_uid``. +Optional ``nvidia-smi dmon`` sampling folds in GPU utilisation and peak VRAM, and is simply absent on a CPU +box. A lightweight host sampler additionally logs per-second CPU (``/proc/stat``) and RAM (``/proc/meminfo``) +to ``host_stats.log`` on any box, GPU or CPU; the reducer ignores that sidecar (it reads only +``timing.jsonl``, ``gpu_dmon.log`` and the recorded episode dirs), so it is a pass-time telemetry artifact, +not a join input. + +The emitted layout is exactly what the reducer consumes:: + + out_dir/ + timing.jsonl # one EpisodeTiming JSON object per rollout + gpu_dmon.log # optional nvidia-smi dmon -s um log (sim box) + host_stats.log # optional per-interval host CPU/RAM samples (any box) + 000000000000/ # 12-digit block dir = (episode_id // 1000) * 1000 + 000000000000/ # 12-digit episode dir = episode_id + meta.json # uid (the join key), created_ts_ns, duration_ns, size_mb + static.json # eval.success / eval.terminated / eval.scored + rollout.parquet # one timestamped signal, so duration derives on its own too + +Usage:: + + shim = TimingShim(out_dir, task='pick_cube') + with shim.run(): + for trial in range(n_trials): + with shim.episode(trial=trial, sim_duration_s=n_steps * control_dt) as ep: + with ep.reset(): + obs = env.reset() + for _ in range(n_steps): + with ep.policy(): + action = policy(obs) + with ep.env_step(): + obs, info = env.step(action) + ep.add_env_phases(info['physics_s'], info['render_s'], info['server_s']) + with ep.record_io(): + recorder.write(obs) + ep.set_outcome(success=bool(info['success']), terminated=True, scored=True) + +Every ``ep.()`` block is a context manager that times its body; the ``add_*`` methods are the +equivalent callbacks for phases whose wall time the harness already measured. ``run()`` and ``episode()`` +own the pass and per-episode lifecycles; both are optional wrappers over the explicit ``start_gpu`` / +``close`` and ``begin_episode`` / ``finish_episode`` calls beneath them. +""" + +import contextlib +import json +import logging +import shutil +import subprocess +import threading +import time +import uuid +from collections.abc import Callable, Iterator +from dataclasses import asdict, dataclass, field +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq + +from positronic.dataset.local_dataset import UNFINISHED_MARKER +from positronic.eval_timing import GPU_LOG_FILENAME, TIMING_FILENAME, EpisodeTiming, start_gpu_sampler + +logger = logging.getLogger(__name__) + +HOST_LOG_FILENAME = 'host_stats.log' +SIGNAL_FILENAME = 'rollout.parquet' + + +@dataclass +class Episode: + """One in-flight rollout: running phase sums plus the outcome and join facts for its recorded stub. + + ``sim_duration_s`` is the virtual time the rollout advanced (``n_steps * control_dt`` for a fixed-step + loop), which becomes the recorded duration the reducer divides by wall time for the real-time factor. + ``size_bytes`` is the real on-disk cost of one rollout's recording; left ``None``, the reducer measures + the stub dir instead (near zero). ``success`` / ``terminated`` / ``scored`` become the ``eval.*`` verdicts + the reducer reads for the success rate. + """ + + task: str + trial: int + uid: str + sim_duration_s: float | None = None + size_bytes: int | None = None + success: bool | None = None + terminated: bool | None = None + scored: bool = False + reset_s: float = 0.0 + env_step_s: float = 0.0 + policy_wait_s: float = 0.0 + record_io_s: float = 0.0 + env_physics_s: float = 0.0 + env_render_s: float = 0.0 + env_server_s: float = 0.0 + infer_ms: list[float] = field(default_factory=list) + _wall_start: float = 0.0 + _discarded: bool = False + + def add_reset(self, seconds: float) -> None: + self.reset_s += seconds + + def add_env_step(self, seconds: float) -> None: + self.env_step_s += seconds + + def add_infer(self, seconds: float) -> None: + """One policy round-trip: the whole of ``policy_wait_s`` and one entry in the latency distribution.""" + self.policy_wait_s += seconds + self.infer_ms.append(seconds * 1000.0) + + def add_record_io(self, seconds: float) -> None: + self.record_io_s += seconds + + def add_env_phases(self, physics_s: float, render_s: float, server_s: float) -> None: + """One env step's server-reported decomposition: physics substeps, rendering, whole in-step wall.""" + self.env_physics_s += physics_s + self.env_render_s += render_s + self.env_server_s += server_s + + def reset(self) -> contextlib.AbstractContextManager[None]: + return _timed(self.add_reset) + + def env_step(self) -> contextlib.AbstractContextManager[None]: + return _timed(self.add_env_step) + + def policy(self) -> contextlib.AbstractContextManager[None]: + return _timed(self.add_infer) + + def record_io(self) -> contextlib.AbstractContextManager[None]: + return _timed(self.add_record_io) + + def set_outcome(self, *, success: bool | None = None, terminated: bool | None = None, scored: bool = False) -> None: + """Record the eval verdicts that drive the success rate. + + A meaningful rate needs either ``success`` (taken as-is) or ``scored=True`` with ``terminated`` (a + no-success termination then counts as a failure). Leaving all three unset keeps the episode out of the + rate, and the run still reduces. + """ + self.success = success + self.terminated = terminated + self.scored = scored + + def discard(self) -> None: + """Drop this rollout without recording it — an abort that should not weigh on the pass.""" + self._discarded = True + + def to_timing(self, wall_s: float, finished_at: float) -> EpisodeTiming: + measured = self.reset_s + self.env_step_s + self.policy_wait_s + self.record_io_s + return EpisodeTiming( + task=self.task, + trial=self.trial, + episode_uid=self.uid, + wall_s=wall_s, + reset_s=self.reset_s, + env_step_s=self.env_step_s, + policy_wait_s=self.policy_wait_s, + record_io_s=self.record_io_s, + overhead_s=max(wall_s - measured, 0.0), + infer_ms=[round(ms, 3) for ms in self.infer_ms], + finished_at=finished_at, + env_physics_s=self.env_physics_s, + env_render_s=self.env_render_s, + env_server_s=self.env_server_s, + ) + + def eval_static(self) -> dict[str, bool]: + """The ``eval.*`` verdicts to write into ``static.json``, omitting the ones left unset.""" + static: dict[str, bool] = {} + if self.success is not None: + static['eval.success'] = self.success + if self.terminated is not None: + static['eval.terminated'] = self.terminated + if self.scored: + static['eval.scored'] = True + return static + + +class TimingShim: + """Collects one ``EpisodeTiming`` per rollout, appending each to ``timing.jsonl`` as it seals, and writes + the recorded dataset dir the reducer consumes. + + The harness records episodes strictly in sequence — one is in flight at a time — so episode ids are a + simple counter and the recorded-stub layout mirrors positronic's block/episode numbering. Records are + flushed per episode, not buffered, so a preempted or killed pass keeps the timings of the episodes it + completed. + """ + + def __init__( + self, + out_dir: str | Path, + *, + task: str | None = None, + sample_gpu: bool = True, + sample_host: bool = True, + write_episodes: bool = True, + ) -> None: + self._out_dir = Path(out_dir) + self._task = task + self._sample_gpu = sample_gpu + self._sample_host = sample_host + self._write_episodes = write_episodes + self._count = 0 + self._next_episode_id = 0 + self._timing_started = False + self._sampler: subprocess.Popen | None = None + self._host_sampler: _HostSampler | None = None + + @contextlib.contextmanager + def run(self) -> Iterator['TimingShim']: + """Bind the pass: create the output dir, clear stale sidecars, and start GPU + host sampling. + + A retry reusing the same ``--out`` dir must not mix a prior attempt's artifacts into this pass: the + reducer joins any episode dirs it finds and auto-reads ``gpu_dmon.log``, and a leftover + ``host_stats.log`` would be mistaken for current telemetry. So stale block dirs, ``timing.jsonl``, the + GPU log and the host log are cleared up front — the GPU and host logs even with their sampling off, + else a prior pass's samples would be folded into this one. Each episode appends its record to + ``timing.jsonl`` as it seals, so a pass killed mid-run keeps the timings it already completed. + """ + self._out_dir.mkdir(parents=True, exist_ok=True) + for stale in self._out_dir.glob('[0-9]' * 12): + shutil.rmtree(stale) + self._prepare_timing_file() + (self._out_dir / GPU_LOG_FILENAME).unlink(missing_ok=True) + (self._out_dir / HOST_LOG_FILENAME).unlink(missing_ok=True) + self._next_episode_id = 0 + self.start_gpu() + self.start_host() + try: + yield self + finally: + self.close() + + def start_gpu(self) -> None: + if self._sample_gpu: + self._sampler = start_gpu_sampler(self._out_dir) + + def start_host(self) -> None: + """Start the host CPU/RAM sampler; a no-op when host sampling is off or ``/proc`` is unreadable.""" + if not self._sample_host: + return + sampler = _HostSampler(self._out_dir) + if sampler.start(): + self._host_sampler = sampler + + def begin_episode( + self, + trial: int, + task: str | None = None, + *, + uid: str | None = None, + sim_duration_s: float | None = None, + size_bytes: int | None = None, + ) -> Episode: + """Open a new in-flight rollout. ``uid`` defaults to a fresh id and is reused as the join key.""" + resolved_task = task if task is not None else self._task + if resolved_task is None: + raise ValueError('task is required: pass it to TimingShim(task=...) or episode(task=...)') + episode = Episode( + task=resolved_task, + trial=trial, + uid=uid if uid is not None else uuid.uuid4().hex, + sim_duration_s=sim_duration_s, + size_bytes=size_bytes, + ) + episode._wall_start = time.perf_counter() + return episode + + def finish_episode(self, episode: Episode) -> None: + """Seal an in-flight rollout: write its recorded stub and append its timing record to disk. + + The record is appended and flushed as the episode seals, so a pass killed mid-run keeps every + completed rollout's timing instead of losing an in-memory buffer. + """ + if episode._discarded: + return + # Freeze the rollout wall and its seal timestamp on the same side of the stub write, so the reducer's + # start = finished_at - wall_s stays exact and the shim's own stub I/O falls outside the episode + # window (the rollout never ran it) instead of skewing wall_s and finished_at apart. + wall_s = time.perf_counter() - episode._wall_start + finished_at = time.time() + episode_id = self._next_episode_id + self._next_episode_id += 1 + if self._write_episodes: + _write_episode_dir(self._out_dir, episode_id, episode) + self._append_timing(episode.to_timing(wall_s, finished_at)) + + @contextlib.contextmanager + def episode( + self, + trial: int, + task: str | None = None, + *, + uid: str | None = None, + sim_duration_s: float | None = None, + size_bytes: int | None = None, + ) -> Iterator[Episode]: + """Time one rollout end to end. A propagating exception aborts it (not recorded), like ``discard``.""" + episode = self.begin_episode(trial, task, uid=uid, sim_duration_s=sim_duration_s, size_bytes=size_bytes) + try: + yield episode + except BaseException: + episode.discard() + raise + finally: + self.finish_episode(episode) + + def close(self) -> Path: + """Stop GPU + host sampling and return the ``timing.jsonl`` path; records flushed as each episode sealed. + + A pass that sealed no episodes — or that ran without ``run()`` — still leaves an empty, valid file. + """ + self._stop_gpu() + self._stop_host() + if not self._timing_started: + self._prepare_timing_file() + path = self._out_dir / TIMING_FILENAME + logger.info(f'TimingShim: wrote {self._count} episode timings to {path}') + return path + + def _prepare_timing_file(self) -> None: + """Truncate ``timing.jsonl`` for a fresh pass; each finished episode then appends one line as it seals.""" + self._out_dir.mkdir(parents=True, exist_ok=True) + (self._out_dir / TIMING_FILENAME).write_text('') + self._timing_started = True + self._count = 0 + + def _append_timing(self, record: EpisodeTiming) -> None: + """Append one sealed record to ``timing.jsonl``, truncating a stale file lazily on the first write.""" + if not self._timing_started: + self._prepare_timing_file() + with (self._out_dir / TIMING_FILENAME).open('a') as f: + f.write(json.dumps(asdict(record)) + '\n') + self._count += 1 + + def _stop_gpu(self) -> None: + if self._sampler is None: + return + self._sampler.terminate() + try: + self._sampler.wait(timeout=5) + except subprocess.TimeoutExpired: + self._sampler.kill() + self._sampler = None + + def _stop_host(self) -> None: + if self._host_sampler is None: + return + self._host_sampler.stop() + self._host_sampler = None + + +@contextlib.contextmanager +def _timed(sink: Callable[[float], None]) -> Iterator[None]: + """Feed the wall duration of the enclosed block to ``sink``.""" + start = time.perf_counter() + try: + yield + finally: + sink(time.perf_counter() - start) + + +def _write_episode_dir(root: Path, episode_id: int, episode: Episode) -> Path: + """Write the recorded-episode stub the reducer joins against, in positronic's block/episode layout.""" + block_dir = root / f'{(episode_id // 1000) * 1000:012d}' + ep_dir = block_dir / f'{episode_id:012d}' + ep_dir.mkdir(parents=True, exist_ok=True) + # The dir is visible before its files land; the marker keeps a preemption mid-write from + # presenting a half-written stub as a finished episode (same contract as LocalDataset's writer). + unfinished = ep_dir / UNFINISHED_MARKER + unfinished.write_text('unfinished', encoding='utf-8') + + created_ts_ns = time.time_ns() + duration_ns = int((episode.sim_duration_s or 0.0) * 1e9) + meta: dict[str, object] = {'uid': episode.uid, 'created_ts_ns': created_ts_ns, 'duration_ns': duration_ns} + if episode.size_bytes is not None: + meta['size_mb'] = episode.size_bytes / (1024 * 1024) + (ep_dir / 'meta.json').write_text(json.dumps(meta)) + (ep_dir / 'static.json').write_text(json.dumps(episode.eval_static())) + _write_signal(ep_dir / SIGNAL_FILENAME, created_ts_ns, duration_ns) + unfinished.unlink() + return ep_dir + + +def _write_signal(path: Path, start_ts_ns: int, duration_ns: int) -> None: + """Write a two-sample timestamped signal so the episode's duration derives from data as well as meta.""" + end_ts_ns = start_ts_ns + max(duration_ns, 1) + table = pa.table({ + 'timestamp': pa.array([start_ts_ns, end_ts_ns], type=pa.int64()), + 'value': pa.array([0.0, 1.0], type=pa.float64()), + }) + pq.write_table(table, path) + + +def _read_proc_stat() -> tuple[tuple[int, int], list[tuple[int, int]]] | None: + """``((idle, total), [(idle, total), ...])`` cumulative jiffies for the whole CPU and each core. + + ``idle`` folds in ``iowait`` (the two idle columns); ``total`` sums every column. ``None`` when + ``/proc/stat`` is unreadable — a box without procfs — so host sampling degrades to absent, not an error. + """ + try: + text = Path('/proc/stat').read_text() + except OSError: + return None + overall: tuple[int, int] | None = None + cores: list[tuple[int, int]] = [] + for line in text.splitlines(): + if not line.startswith('cpu'): + break # the cpu/cpuN lines lead /proc/stat; the rest (intr, ctxt, …) is irrelevant here + parts = line.split() + nums = [int(x) for x in parts[1:]] + idle = nums[3] + (nums[4] if len(nums) > 4 else 0) + pair = (idle, sum(nums)) + if parts[0] == 'cpu': + overall = pair + elif parts[0][3:].isdigit(): + cores.append(pair) + if overall is None: + return None + return overall, cores + + +def _read_meminfo() -> tuple[int, int] | None: + """``(mem_total_kb, mem_available_kb)`` from ``/proc/meminfo``, or ``None`` when unreadable.""" + try: + text = Path('/proc/meminfo').read_text() + except OSError: + return None + total_kb: int | None = None + avail_kb: int | None = None + for line in text.splitlines(): + key, _, rest = line.partition(':') + if key == 'MemTotal': + total_kb = int(rest.split()[0]) + elif key == 'MemAvailable': + avail_kb = int(rest.split()[0]) + if total_kb is not None and avail_kb is not None: + break + if total_kb is None or avail_kb is None: + return None + return total_kb, avail_kb + + +def _cpu_busy_pct(prev: tuple[int, int], cur: tuple[int, int]) -> float: + """Percent busy over the interval between two ``(idle, total)`` jiffie snapshots.""" + idle_delta = cur[0] - prev[0] + total_delta = cur[1] - prev[1] + if total_delta <= 0: + return 0.0 + return round(100.0 * (1.0 - idle_delta / total_delta), 2) + + +class _HostSampler: + """Daemon thread logging per-interval host CPU (``/proc/stat``) and RAM (``/proc/meminfo``) samples. + + Writes ``host_stats.log`` — one JSON object per interval carrying monotonic + wall timestamps, overall and + per-core CPU utilisation (busy fraction over that interval), and memory total/used/available in GiB. Works + on any Linux box, GPU or CPU. A missing ``/proc`` at start yields no log (``start`` returns ``False``); a + read that begins failing mid-pass simply stops emitting, so the pass is never taken down by its telemetry. + """ + + def __init__(self, out_dir: Path, *, interval_s: float = 1.0) -> None: + self._log_path = out_dir / HOST_LOG_FILENAME + self._interval_s = interval_s + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, name='timingshim-host', daemon=True) + self._baseline: tuple[tuple[int, int], list[tuple[int, int]]] | None = None + self._fh = None + + def start(self) -> bool: + """Capture the CPU baseline and launch the sampler thread; ``False`` when ``/proc/stat`` is unreadable.""" + self._baseline = _read_proc_stat() + if self._baseline is None: + logger.info('/proc/stat unavailable; skipping host CPU/RAM sampling') + return False + self._log_path.unlink(missing_ok=True) + self._fh = self._log_path.open('w') + self._thread.start() + return True + + def _run(self) -> None: + prev_overall, prev_cores = self._baseline + # Wait one interval before the first sample so every logged line is a real delta off the baseline. + while not self._stop.wait(self._interval_s): + stat = _read_proc_stat() + mem = _read_meminfo() + if stat is None or mem is None: + continue + overall, cores = stat + total_kb, avail_kb = mem + sample = { + 't_mono': round(time.monotonic(), 3), + 't_wall': round(time.time(), 3), + 'cpu_pct': _cpu_busy_pct(prev_overall, overall), + 'per_core_pct': [_cpu_busy_pct(p, c) for p, c in zip(prev_cores, cores, strict=False)], + 'mem_total_gb': round(total_kb / (1024 * 1024), 3), + 'mem_used_gb': round((total_kb - avail_kb) / (1024 * 1024), 3), + 'mem_avail_gb': round(avail_kb / (1024 * 1024), 3), + } + self._fh.write(json.dumps(sample) + '\n') + self._fh.flush() + prev_overall, prev_cores = overall, cores + + def stop(self) -> None: + """Signal the thread, join it, and close the log.""" + self._stop.set() + self._thread.join(timeout=self._interval_s + 5.0) + if self._fh is not None: + self._fh.close() + self._fh = None diff --git a/positronic/tests/test_eval_timing_shim.py b/positronic/tests/test_eval_timing_shim.py new file mode 100644 index 000000000..a239f5699 --- /dev/null +++ b/positronic/tests/test_eval_timing_shim.py @@ -0,0 +1,344 @@ +"""Tests for ``eval_timing_shim``: the emitted ``timing.jsonl`` + episode stubs feed ``EpisodeTiming`` and +reduce cleanly through the real ``positronic eval timing-report`` CLI.""" + +import json +import os +import shutil +import subprocess +import time +from dataclasses import fields +from pathlib import Path + +import pytest + +from positronic.dataset.local_dataset import UNFINISHED_MARKER +from positronic.eval_timing import GPU_LOG_FILENAME, TIMING_FILENAME, EpisodeTiming +from positronic.eval_timing_shim import ( + HOST_LOG_FILENAME, + SIGNAL_FILENAME, + TimingShim, + _cpu_busy_pct, + _HostSampler, + _read_meminfo, + _read_proc_stat, +) + +REPO_ROOT = Path(__file__).resolve().parents[2] +SECONDS_FIELDS = [f.name for f in fields(EpisodeTiming) if f.name.endswith('_s')] + + +def _drive_episode(ep, n_steps: int) -> None: + """Run one synthetic rollout through the shim's phase hooks with small real sleeps for a realistic split.""" + with ep.reset(): + time.sleep(0.002) + for _ in range(n_steps): + with ep.policy(): + time.sleep(0.005) + with ep.env_step(): + time.sleep(0.004) + ep.add_env_phases(physics_s=0.002, render_s=0.001, server_s=0.0035) + with ep.record_io(): + time.sleep(0.001) + + +def _build_synthetic_run(out_dir: Path, *, sample_gpu: bool = False) -> dict: + """Emit a full run: four scored episodes (two success, one fail, one unscored) plus a discarded abort. + + Returns the facts the assertions need — the expected recorded-episode count and success rate. + """ + outcomes = [ + {'success': True, 'terminated': True, 'scored': True}, + {'success': True, 'terminated': True, 'scored': True}, + {'success': False, 'terminated': True, 'scored': True}, + {'success': None, 'terminated': None, 'scored': False}, # unscored: excluded from the rate + ] + shim = TimingShim(out_dir, task='pick_cube', sample_gpu=sample_gpu) + with shim.run(): + for trial, outcome in enumerate(outcomes): + with shim.episode(trial=trial, sim_duration_s=5.0, size_bytes=2 * 1024 * 1024) as ep: + _drive_episode(ep, n_steps=3) + ep.set_outcome(**outcome) + # An abort mid-rollout must not be recorded, in timing.jsonl or as an episode dir. + with pytest.raises(RuntimeError): + with shim.episode(trial=99, sim_duration_s=5.0) as ep: + _drive_episode(ep, n_steps=1) + raise RuntimeError('injected abort') + return {'episodes': len(outcomes), 'expected_success_rate': 2 / 3} + + +def _read_timing(out_dir: Path) -> list[dict]: + lines = (out_dir / TIMING_FILENAME).read_text().splitlines() + return [json.loads(line) for line in lines if line.strip()] + + +def test_timing_jsonl_feeds_episode_timing(tmp_path): + facts = _build_synthetic_run(tmp_path) + rows = _read_timing(tmp_path) + + # The discarded episode is absent; only the four completed rollouts are recorded. + assert len(rows) == facts['episodes'] + + for row in rows: + # The line round-trips into the dataclass exactly as the reducer decodes it — no missing/extra keys. + record = EpisodeTiming(**row) + assert isinstance(record.task, str) + assert isinstance(record.trial, int) and not isinstance(record.trial, bool) + assert isinstance(record.episode_uid, str) and record.episode_uid + for name in SECONDS_FIELDS: + value = getattr(record, name) + assert isinstance(value, float), f'{name} must be float, got {type(value)}' + assert value >= 0.0, f'{name} negative' + assert all(isinstance(x, float) for x in record.infer_ms) + assert len(record.infer_ms) == 3 # one entry per policy call + # Epoch seconds at seal — a real recent wall clock, not a relative offset. + assert isinstance(record.finished_at, float) and record.finished_at > 1.7e9 + + # Producer invariants the reducer's derived metrics assume. + assert record.policy_wait_s == pytest.approx(sum(record.infer_ms) / 1000.0, abs=1e-3) + measured = record.reset_s + record.env_step_s + record.policy_wait_s + record.record_io_s + assert record.overhead_s == pytest.approx(max(record.wall_s - measured, 0.0), abs=1e-6) + # env_step server decomposition stays within the client-observed env-step wall. + assert record.env_server_s <= record.env_step_s + 1e-9 + assert record.env_physics_s + record.env_render_s <= record.env_server_s + 1e-9 + + +def test_dataset_dir_layout_joins_by_uid(tmp_path): + _build_synthetic_run(tmp_path) + rows = _read_timing(tmp_path) + recorded_uids = set() + for meta_path in sorted(tmp_path.rglob('meta.json')): + meta = json.loads(meta_path.read_text()) + assert 'uid' in meta and 'duration_ns' in meta + # 12-digit block/episode names, as load_all_datasets requires. + assert meta_path.parent.name.isdigit() and len(meta_path.parent.name) == 12 + assert meta_path.parent.parent.name.isdigit() and len(meta_path.parent.parent.name) == 12 + assert (meta_path.parent / 'static.json').exists() + assert (meta_path.parent / SIGNAL_FILENAME).exists() + recorded_uids.add(meta['uid']) + # Every timing record's join key has a matching recorded episode (the reducer fails otherwise). + assert {row['episode_uid'] for row in rows} == recorded_uids + assert len(recorded_uids) == len(rows) + + +def test_preempted_stub_write_leaves_unfinished_marker(tmp_path, monkeypatch): + # A run preempted after the episode dir becomes visible but before its files land must not present a + # finished episode (regression: a half-written static.json failed timing-report on the whole dataset). + # The stub carries LocalDataset's marker until the last file lands; a sealed write leaves no marker. + import positronic.eval_timing_shim as shim_mod + + shim = TimingShim(tmp_path, task='pick_cube', sample_gpu=False, sample_host=False) + ep = shim.begin_episode(trial=0, sim_duration_s=1.0) + _drive_episode(ep, n_steps=1) + ep.set_outcome(success=True, terminated=True, scored=True) + + def preempted(path, start_ts_ns, duration_ns): + raise RuntimeError('preempted mid-write') + + monkeypatch.setattr(shim_mod, '_write_signal', preempted) + with pytest.raises(RuntimeError): + shim.finish_episode(ep) + markers = list(tmp_path.rglob(UNFINISHED_MARKER)) + assert len(markers) == 1 # the half-written stub is identifiable as unfinished + + monkeypatch.undo() + ep2 = shim.begin_episode(trial=1, sim_duration_s=1.0) + _drive_episode(ep2, n_steps=1) + ep2.set_outcome(success=True, terminated=True, scored=True) + shim.finish_episode(ep2) + assert list(tmp_path.rglob(UNFINISHED_MARKER)) == markers # the sealed episode carries no marker + + +def test_finished_episodes_flush_to_disk_before_close(tmp_path): + # A pass killed mid-run (no close(), no run() wrapper) must still leave every finished rollout's timing + # on disk: records append as each episode seals rather than buffering until close(). A stale timing.jsonl + # from a prior pass is truncated lazily on the first flush. + (tmp_path / TIMING_FILENAME).write_text('stale line from a prior pass\n') + shim = TimingShim(tmp_path, task='pick_cube', sample_gpu=False, sample_host=False) + uids = [] + for trial in range(3): + ep = shim.begin_episode(trial=trial, sim_duration_s=5.0) + _drive_episode(ep, n_steps=2) + ep.set_outcome(success=True, terminated=True, scored=True) + shim.finish_episode(ep) + uids.append(ep.uid) + # The record is durable the moment the episode seals — no close() has run yet. + assert len(_read_timing(tmp_path)) == trial + 1 + + rows = _read_timing(tmp_path) + assert [row['episode_uid'] for row in rows] == uids # stale line gone, three fresh records in seal order + + +def test_stub_write_stays_outside_reconstructed_episode_window(tmp_path, monkeypatch): + # The reducer reconstructs an episode's start as finished_at - wall_s. Both must be frozen on the same + # side of the stub write, else slow-filesystem stub I/O shifts the reconstructed start. Force a slow + # stub write and assert the reconstruction still lands on the real rollout start, not 0.2s later. + import positronic.eval_timing_shim as shim_mod + + real_write = shim_mod._write_episode_dir + + def slow_write(root, episode_id, episode): + time.sleep(0.2) + return real_write(root, episode_id, episode) + + monkeypatch.setattr(shim_mod, '_write_episode_dir', slow_write) + + shim = TimingShim(tmp_path, task='pick_cube', sample_gpu=False, sample_host=False) + t_begin = time.time() + ep = shim.begin_episode(trial=0, sim_duration_s=5.0) + _drive_episode(ep, n_steps=1) + ep.set_outcome(success=True, terminated=True, scored=True) + shim.finish_episode(ep) # the 0.2s stub write happens after wall_s and finished_at are captured + + row = _read_timing(tmp_path)[0] + reconstructed_start = row['finished_at'] - row['wall_s'] + assert reconstructed_start == pytest.approx(t_begin, abs=0.1) # unshifted by the 0.2s stub I/O + assert row['wall_s'] < 0.15 # the rollout wall excludes the stub write entirely + + +def test_run_clears_stale_gpu_log_with_sampling_off(tmp_path): + # The reducer auto-reads any gpu_dmon.log in the dataset dir, so a prior GPU pass's log must not + # survive into a fresh sample_gpu=False pass. + (tmp_path / GPU_LOG_FILENAME).write_text('stale dmon samples\n') + _build_synthetic_run(tmp_path, sample_gpu=False) + assert not (tmp_path / GPU_LOG_FILENAME).exists() + + +def test_run_clears_stale_host_log_with_sampling_off(tmp_path): + # Symmetric with the GPU log: a prior pass's host_stats.log must not survive a sample_host=False pass, + # else stale CPU/RAM samples sit next to the fresh timing.jsonl and are mistaken for current telemetry. + (tmp_path / HOST_LOG_FILENAME).write_text('stale host samples\n') + shim = TimingShim(tmp_path, task='pick_cube', sample_gpu=False, sample_host=False) + with shim.run(): + pass + assert not (tmp_path / HOST_LOG_FILENAME).exists() + + +def test_gpu_sampling_degrades_gracefully(tmp_path): + _build_synthetic_run(tmp_path, sample_gpu=True) + gpu_log = tmp_path / GPU_LOG_FILENAME + if shutil.which('nvidia-smi') is None: + assert not gpu_log.exists(), 'no GPU log should be written on a box without nvidia-smi' + else: + assert gpu_log.exists() + + +HOST_SAMPLE_KEYS = {'t_mono', 't_wall', 'cpu_pct', 'per_core_pct', 'mem_total_gb', 'mem_used_gb', 'mem_avail_gb'} + + +def test_host_proc_readers_parse_this_box(): + """The /proc readers return well-formed data on a real Linux box (no psutil, procfs only).""" + stat = _read_proc_stat() + if stat is None: + pytest.skip('/proc/stat unavailable on this platform') + overall, cores = stat + assert len(overall) == 2 and all(isinstance(v, int) for v in overall) + assert len(cores) == os.cpu_count() # one (idle, total) pair per core + assert all(idle <= total for idle, total in [overall, *cores]) + + total_kb, avail_kb = _read_meminfo() + assert total_kb > 0 and 0 < avail_kb <= total_kb + + +def test_cpu_busy_pct_bounds_and_zero_interval(): + # Half the added jiffies were idle -> 50% busy over the interval. + assert _cpu_busy_pct((100, 200), (150, 300)) == pytest.approx(50.0) + # All added jiffies idle -> 0% busy; all busy -> 100%. + assert _cpu_busy_pct((100, 200), (200, 300)) == pytest.approx(0.0) + assert _cpu_busy_pct((100, 200), (100, 300)) == pytest.approx(100.0) + # No elapsed jiffies (two identical snapshots) can't divide by zero. + assert _cpu_busy_pct((100, 200), (100, 200)) == 0.0 + + +def test_host_sampler_writes_monotone_samples_and_shuts_down_clean(tmp_path): + sampler = _HostSampler(tmp_path, interval_s=0.05) + if not sampler.start(): + pytest.skip('/proc/stat unavailable on this platform') + time.sleep(0.3) + sampler.stop() + + # Thread is joined and the file handle closed on stop() — no lingering daemon, no open fd. + assert not sampler._thread.is_alive() + assert sampler._fh is None + + lines = (tmp_path / HOST_LOG_FILENAME).read_text().splitlines() + assert len(lines) >= 2, f'expected several samples at 50ms over 300ms, got {len(lines)}' + + prev_mono = None + for line in lines: + sample = json.loads(line) + assert set(sample) == HOST_SAMPLE_KEYS, f'unexpected keys: {set(sample) ^ HOST_SAMPLE_KEYS}' + assert 0.0 <= sample['cpu_pct'] <= 100.0 + assert len(sample['per_core_pct']) == os.cpu_count() + assert all(0.0 <= pct <= 100.0 for pct in sample['per_core_pct']) + assert sample['mem_total_gb'] > 0.0 + assert sample['mem_used_gb'] >= 0.0 + assert sample['mem_avail_gb'] >= 0.0 + assert sample['mem_used_gb'] + sample['mem_avail_gb'] == pytest.approx(sample['mem_total_gb'], abs=5e-3) + # Monotonic timestamps strictly increase across the interval-spaced samples. + if prev_mono is not None: + assert sample['t_mono'] > prev_mono + prev_mono = sample['t_mono'] + + +def test_shim_run_lifecycle_emits_host_log(tmp_path): + """The pass lifecycle (run/close) starts and stops the host sampler, leaving a host_stats.log behind.""" + shim = TimingShim(tmp_path, task='pick_cube', sample_gpu=False) + with shim.run(): + if shim._host_sampler is None: + pytest.skip('/proc/stat unavailable on this platform') + time.sleep(0.15) + assert shim._host_sampler is None # stopped by close() + assert (tmp_path / HOST_LOG_FILENAME).exists() + + +def test_shim_host_sampling_can_be_disabled(tmp_path): + shim = TimingShim(tmp_path, task='pick_cube', sample_gpu=False, sample_host=False) + with shim.run(): + assert shim._host_sampler is None + assert not (tmp_path / HOST_LOG_FILENAME).exists() + + +def _write_gpu_log(out_dir: Path) -> None: + """A minimal ``nvidia-smi dmon -s um`` shaped log so the reducer's column-driven GPU parse is exercised.""" + (out_dir / GPU_LOG_FILENAME).write_text( + '#Date Time gpu sm mem enc dec fb\n' + '20260721 12:00:00 0 55 10 0 0 8000\n' + '20260721 12:00:01 0 65 12 0 0 9000\n' + ) + + +def test_reducer_end_to_end(tmp_path): + facts = _build_synthetic_run(tmp_path, sample_gpu=False) + _write_gpu_log(tmp_path) # a box without nvidia-smi has no dmon log; inject one to exercise the GPU parse + # The host sampler's sidecar is a pass-time artifact the reducer must ignore, not join on. Drop a + # realistic one at the dataset root and assert the reduce is unchanged by its presence. + (tmp_path / HOST_LOG_FILENAME).write_text( + '{"t_mono": 1.0, "t_wall": 1000.0, "cpu_pct": 12.5, "per_core_pct": [10.0, 15.0], ' + '"mem_total_gb": 31.1, "mem_used_gb": 14.2, "mem_avail_gb": 16.9}\n' + '{"t_mono": 2.0, "t_wall": 1001.0, "cpu_pct": 33.0, "per_core_pct": [30.0, 36.0], ' + '"mem_total_gb": 31.1, "mem_used_gb": 15.0, "mem_avail_gb": 16.1}\n' + ) + + result = subprocess.run( + ['uv', 'run', '--locked', 'positronic', 'eval', 'timing-report', '--dataset_dir', str(tmp_path)], + cwd=REPO_ROOT, + capture_output=True, + text=True, + timeout=300, + ) + assert result.returncode == 0, f'timing-report failed:\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}' + + summary = json.loads((tmp_path / 'timing_summary.json').read_text()) + assert summary['episodes'] == facts['episodes'] + assert summary['infer_calls'] == facts['episodes'] * 3 + assert summary['success_rate'] == pytest.approx(facts['expected_success_rate']) + assert summary['real_time_factor'] > 0.0 + assert summary['infer_p50_ms'] > 0.0 + # env server decomposition was recorded, so the split is present, not null. + assert summary['env_step_split'] is not None + # bytes/rollout comes from the stamped size_bytes (2 MiB), joined from meta.json. + assert summary['mean_bytes_per_rollout'] == pytest.approx(2 * 1024 * 1024) + # the injected dmon log parses: mean sm util over {55,65}=60, peak fb 9000 MiB. + assert summary['gpu']['sim'] is not None + assert summary['gpu']['sim']['mean_util_pct'] == pytest.approx(60.0) + assert summary['gpu']['sim']['peak_vram_gb'] == pytest.approx(9000 / 1024.0)