|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 3 | +# SPDX-License-Identifier: Apache-2.0 |
| 4 | + |
| 5 | +"""CLI entry point for the -vvv trace dashboard. |
| 6 | +
|
| 7 | +Reads fixed-size binary frames from the FIFO opened by |
| 8 | +:func:`inference_endpoint.utils.trace.bootstrap` and renders the |
| 9 | +dashboard via ``rich.Live``. Dashboard aggregation lives in |
| 10 | +:mod:`inference_endpoint.utils.trace_dashboard` so it can be unit |
| 11 | +tested without standing up a TUI. |
| 12 | +
|
| 13 | +Linux only: timestamps are compared across processes and rely on |
| 14 | +``CLOCK_MONOTONIC`` being system-wide (per ``man 7 time``). |
| 15 | +""" |
| 16 | + |
| 17 | +# ruff: noqa: I001 |
| 18 | +# The pre-commit ruff hook is pinned to v0.3.3 (see |
| 19 | +# .pre-commit-config.yaml's "TODO: sync rev with ruff version"), which |
| 20 | +# does not auto-detect `inference_endpoint` as a first-party package |
| 21 | +# and therefore disagrees with the project's local ruff (v0.15.8) on |
| 22 | +# import order in this file. File-level noqa keeps both versions quiet |
| 23 | +# until the rev is synced. |
| 24 | +from __future__ import annotations |
| 25 | + |
| 26 | +import argparse |
| 27 | +import fcntl |
| 28 | +import json |
| 29 | +import os |
| 30 | +import sys |
| 31 | +import threading |
| 32 | +import time |
| 33 | + |
| 34 | +import zmq |
| 35 | +from inference_endpoint.async_utils.services.metrics_aggregator.snapshot import ( |
| 36 | + MetricsSnapshotCodec, |
| 37 | + snapshot_to_dict, |
| 38 | +) |
| 39 | +from inference_endpoint.core.record import BATCH_TOPIC, TOPIC_FRAME_SIZE |
| 40 | +from inference_endpoint.utils.trace import ( |
| 41 | + FRAME_SIZE, |
| 42 | + metrics_addr_path, |
| 43 | + snapshot_sidecar_path, |
| 44 | +) |
| 45 | +from inference_endpoint.utils.trace_dashboard import ( |
| 46 | + DASHBOARD_THEME, |
| 47 | + READ_CHUNK, |
| 48 | + REFRESH_HZ, |
| 49 | + Dashboard, |
| 50 | +) |
| 51 | +from rich.console import Console |
| 52 | +from rich.live import Live |
| 53 | + |
| 54 | + |
| 55 | +def _try_load_snapshot(path: str) -> dict | None: |
| 56 | + """Best-effort read of the loadgen snapshot sidecar. Returns None |
| 57 | + if the file is missing or transiently mid-rename (atomic write may |
| 58 | + briefly produce a half-rename window that json.load tolerates).""" |
| 59 | + try: |
| 60 | + with open(path) as f: |
| 61 | + return json.load(f) |
| 62 | + except (OSError, json.JSONDecodeError): |
| 63 | + return None |
| 64 | + |
| 65 | + |
| 66 | +# End-of-run exit policy for the reader thread. |
| 67 | +# |
| 68 | +# The authoritative final snapshot (state=="complete") is written by the |
| 69 | +# parent's trace.teardown() *just before* it closes its FIFO write fd, and |
| 70 | +# every worker closes its write fd at process exit. So true FIFO EOF (all |
| 71 | +# writers closed) is the guaranteed signal that the complete snapshot is |
| 72 | +# already on disk — preferring EOF makes the FINAL panel deterministic. |
| 73 | +# |
| 74 | +# But workers can take tens of seconds to drain their ZMQ queues before |
| 75 | +# exiting, so we can't wait for EOF unconditionally or the dashboard hangs |
| 76 | +# long past the end of the run. Compromise: once lifecycle frames go quiet, |
| 77 | +# keep reading — letting true EOF trigger an immediate clean exit — until |
| 78 | +# this generous cap bounds a wedged-worker hang. After the reader exits, |
| 79 | +# main() polls the sidecar for the complete snapshot. |
| 80 | +_IDLE_EXIT_CAP_S = 30.0 |
| 81 | + |
| 82 | + |
| 83 | +class _FrameReader(threading.Thread): |
| 84 | + """Blocking read loop; ingests whole frames into the Dashboard. |
| 85 | +
|
| 86 | + Uses ``select`` with a short timeout so the thread can check for an |
| 87 | + idle-exit condition — the FIFO may not reach EOF until all 24+ worker |
| 88 | + processes have drained their ZMQ queues and exited, which can be tens |
| 89 | + of seconds after the benchmark has finished. |
| 90 | + """ |
| 91 | + |
| 92 | + def __init__(self, fd: int, dash: Dashboard) -> None: |
| 93 | + super().__init__(daemon=True, name="trace-reader") |
| 94 | + self._fd = fd |
| 95 | + self._dash = dash |
| 96 | + self._pending = bytearray() |
| 97 | + self._eof = threading.Event() |
| 98 | + |
| 99 | + @property |
| 100 | + def eof(self) -> bool: |
| 101 | + return self._eof.is_set() |
| 102 | + |
| 103 | + def run(self) -> None: |
| 104 | + import select |
| 105 | + |
| 106 | + try: |
| 107 | + while True: |
| 108 | + ready, _, _ = select.select([self._fd], [], [], 0.5) |
| 109 | + if ready: |
| 110 | + try: |
| 111 | + chunk = os.read(self._fd, READ_CHUNK) |
| 112 | + except OSError: |
| 113 | + return |
| 114 | + if not chunk: |
| 115 | + return # true EOF — all writers closed |
| 116 | + self._pending.extend(chunk) |
| 117 | + whole = (len(self._pending) // FRAME_SIZE) * FRAME_SIZE |
| 118 | + if whole: |
| 119 | + self._dash.ingest_frames(bytes(self._pending[:whole])) |
| 120 | + del self._pending[:whole] |
| 121 | + # Check lifecycle idle time regardless of whether LOOP_LAG |
| 122 | + # frames are still arriving — workers emit LOOP_LAG every |
| 123 | + # 0.3 s even after the run ends, so frame-arrival time is |
| 124 | + # not a reliable proxy for end-of-run. We keep reading so |
| 125 | + # true EOF (the parent's post-teardown fd close) wins and |
| 126 | + # guarantees the complete snapshot is on disk; this cap only |
| 127 | + # bounds a wedged-worker hang where EOF never arrives. |
| 128 | + d = self._dash |
| 129 | + if (d.is_done or d.is_tail) and d.lifecycle_idle_s >= _IDLE_EXIT_CAP_S: |
| 130 | + return |
| 131 | + finally: |
| 132 | + self._eof.set() |
| 133 | + |
| 134 | + |
| 135 | +class _MetricsSubReader(threading.Thread): |
| 136 | + """Opens a SUB straight to the aggregator's metrics PUB and feeds the |
| 137 | + dashboard fresh LOADGEN snapshots. |
| 138 | +
|
| 139 | + The aggregator is a separate process that publishes every tick (and the |
| 140 | + terminal COMPLETE frame), so it is immune to the main benchmark loop's |
| 141 | + saturation — unlike the sidecar, which is written from the main proc's |
| 142 | + starved in-process subscriber. Best-effort: any failure leaves the |
| 143 | + sidecar fallback in place. Sets ``delivered`` once a snapshot lands so |
| 144 | + main() can stop polling the sidecar. |
| 145 | + """ |
| 146 | + |
| 147 | + def __init__(self, addr_path: str, dash: Dashboard) -> None: |
| 148 | + super().__init__(daemon=True, name="metrics-sub") |
| 149 | + self._addr_path = addr_path |
| 150 | + self._dash = dash |
| 151 | + # Not ``_stop``: that shadows threading.Thread._stop and breaks join(). |
| 152 | + self._stopping = threading.Event() |
| 153 | + self.delivered = threading.Event() |
| 154 | + |
| 155 | + def stop(self) -> None: |
| 156 | + self._stopping.set() |
| 157 | + |
| 158 | + def _wait_for_addr(self) -> str | None: |
| 159 | + # The parent writes the addr only after the aggregator PUB binds, |
| 160 | + # which is well after the dashboard spawns — poll for it. |
| 161 | + deadline = time.monotonic() + 30.0 |
| 162 | + while time.monotonic() < deadline and not self._stopping.is_set(): |
| 163 | + try: |
| 164 | + with open(self._addr_path) as f: |
| 165 | + addr = f.read().strip() |
| 166 | + if addr: |
| 167 | + return addr |
| 168 | + except OSError: |
| 169 | + pass # addr file not published yet — retry until the deadline |
| 170 | + time.sleep(0.2) |
| 171 | + return None |
| 172 | + |
| 173 | + def run(self) -> None: |
| 174 | + addr = self._wait_for_addr() |
| 175 | + if addr is None: |
| 176 | + return # never published — main() keeps polling the sidecar |
| 177 | + ctx = zmq.Context.instance() |
| 178 | + sock = ctx.socket(zmq.SUB) |
| 179 | + sock.setsockopt(zmq.SUBSCRIBE, b"") |
| 180 | + sock.setsockopt(zmq.CONFLATE, 1) # only the freshest snapshot |
| 181 | + sock.setsockopt(zmq.RCVTIMEO, 500) |
| 182 | + sock.setsockopt(zmq.LINGER, 0) |
| 183 | + codec = MetricsSnapshotCodec() |
| 184 | + try: |
| 185 | + sock.connect(addr) |
| 186 | + while not self._stopping.is_set(): |
| 187 | + try: |
| 188 | + raw = sock.recv() |
| 189 | + except zmq.Again: |
| 190 | + continue # no snapshot within RCVTIMEO — re-check stop, retry |
| 191 | + except zmq.ZMQError: |
| 192 | + return |
| 193 | + if raw[:TOPIC_FRAME_SIZE] == BATCH_TOPIC: |
| 194 | + continue # metrics snapshots are never batched |
| 195 | + payload = raw[TOPIC_FRAME_SIZE:] if len(raw) > TOPIC_FRAME_SIZE else raw |
| 196 | + try: |
| 197 | + snap = snapshot_to_dict(codec.decode(payload)) |
| 198 | + except Exception: # noqa: BLE001 — telemetry, never crash |
| 199 | + continue |
| 200 | + self._dash.attach_loadgen_snapshot(snap) |
| 201 | + self.delivered.set() |
| 202 | + finally: |
| 203 | + sock.close(0) |
| 204 | + |
| 205 | + |
| 206 | +# Match the producers' write-fd request in trace.py. Capped by |
| 207 | +# /proc/sys/fs/pipe-max-size; request fails above that → default kept. |
| 208 | +_F_SETPIPE_SZ = getattr(fcntl, "F_SETPIPE_SZ", 1031) |
| 209 | +_KERNEL_PIPE_BUF = 64 * 1024 * 1024 |
| 210 | + |
| 211 | + |
| 212 | +def _open_trace_input(pipe_path: str | None) -> int: |
| 213 | + if pipe_path: |
| 214 | + fd = os.open(pipe_path, os.O_RDONLY) |
| 215 | + try: |
| 216 | + fcntl.fcntl(fd, _F_SETPIPE_SZ, _KERNEL_PIPE_BUF) |
| 217 | + except OSError: |
| 218 | + pass # F_SETPIPE_SZ is best-effort — keep the kernel default |
| 219 | + return fd |
| 220 | + return sys.stdin.fileno() |
| 221 | + |
| 222 | + |
| 223 | +def main() -> int: |
| 224 | + parser = argparse.ArgumentParser(description=__doc__) |
| 225 | + parser.add_argument( |
| 226 | + "--trace-pipe", |
| 227 | + help="FIFO path to read binary trace frames from (default: stdin).", |
| 228 | + ) |
| 229 | + args = parser.parse_args() |
| 230 | + |
| 231 | + dash = Dashboard() |
| 232 | + # Dashboard renders to stderr because the parent benchmark process |
| 233 | + # has redirected its own stdout/stderr to a log file (see trace.bootstrap). |
| 234 | + console = Console(file=sys.stderr, force_terminal=True, theme=DASHBOARD_THEME) |
| 235 | + reader = _FrameReader(_open_trace_input(args.trace_pipe), dash) |
| 236 | + reader.start() |
| 237 | + |
| 238 | + # Convention paths keyed on the parent's pid (the main proc). |
| 239 | + snap_path = snapshot_sidecar_path(os.getppid()) |
| 240 | + # Primary LOADGEN feed: a SUB straight to the aggregator PUB (fresh). |
| 241 | + # Falls back to the main-proc-written sidecar if the SUB never connects. |
| 242 | + sub_reader = _MetricsSubReader(metrics_addr_path(os.getppid()), dash) |
| 243 | + sub_reader.start() |
| 244 | + |
| 245 | + # screen=True uses the alternate-screen buffer so updates redraw |
| 246 | + # cleanly without scrollback noise. When Live() exits the alt |
| 247 | + # screen is torn down — to keep the final frame visible we capture |
| 248 | + # it BEFORE leaving the context and print it to the normal buffer |
| 249 | + # afterward. |
| 250 | + final_frame = None |
| 251 | + with Live( |
| 252 | + dash.render(), |
| 253 | + console=console, |
| 254 | + refresh_per_second=REFRESH_HZ, |
| 255 | + screen=True, |
| 256 | + transient=False, |
| 257 | + ) as live: |
| 258 | + while not reader.eof: |
| 259 | + # SUB is the fresh primary; only fall back to the sidecar |
| 260 | + # until/unless the SUB has delivered at least one snapshot. |
| 261 | + if not sub_reader.delivered.is_set(): |
| 262 | + snap = _try_load_snapshot(snap_path) |
| 263 | + if snap is not None: |
| 264 | + dash.attach_loadgen_snapshot(snap) |
| 265 | + live.update(dash.render()) |
| 266 | + time.sleep(1.0 / REFRESH_HZ) |
| 267 | + # Wait for the terminal snapshot before the closing frame. teardown() |
| 268 | + # writes it to the sidecar before closing the FIFO (reliably on disk |
| 269 | + # by EOF); the SUB can miss the terminal PUB frame, so poll the sidecar |
| 270 | + # as the source of truth. |
| 271 | + _FINAL_SNAP_WAIT_S = 20.0 |
| 272 | + deadline = time.monotonic() + _FINAL_SNAP_WAIT_S |
| 273 | + while time.monotonic() < deadline and not dash.has_terminal_loadgen: |
| 274 | + s = _try_load_snapshot(snap_path) |
| 275 | + if s is not None: |
| 276 | + dash.attach_loadgen_snapshot(s, force=True) |
| 277 | + time.sleep(0.1) |
| 278 | + # Join the SUB before the final render so no late snapshot attaches. |
| 279 | + sub_reader.stop() |
| 280 | + sub_reader.join(timeout=1.0) |
| 281 | + # Wait elapsed with no terminal snapshot → the aggregator never |
| 282 | + # finalized (crash / OOM / drain past the wait). Flag it so the |
| 283 | + # closing frame says "unavailable", not a perpetual "finalizing…". |
| 284 | + if not dash.has_terminal_loadgen: |
| 285 | + dash.mark_final_unavailable() |
| 286 | + # Bypass the per-tick fold-defer window so the final render |
| 287 | + # captures COMPLETE frames that landed within the last 300 ms |
| 288 | + # — otherwise they sit queued and the closing frame shows |
| 289 | + # stale stage histograms / verdict. |
| 290 | + dash.flush_pending_folds() |
| 291 | + final_frame = dash.render() |
| 292 | + live.update(final_frame) |
| 293 | + # Now we're back on the normal screen — print the last snapshot |
| 294 | + # so the user can see the totals + verdict after the run ends. |
| 295 | + if final_frame is not None: |
| 296 | + console.print(final_frame) |
| 297 | + console.print("[dim]── trace finished ──[/dim]") |
| 298 | + return 0 |
| 299 | + |
| 300 | + |
| 301 | +if __name__ == "__main__": |
| 302 | + sys.exit(main()) |
0 commit comments