diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index b09c2c3a03..3683039528 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -80,6 +80,12 @@ def init( if self._is_main_rank: init_tracking(args, primary=False) + # every rank, not just main: per-GPU dashboard phase lanes need each + # rank's Timer events (no-op unless --use-miles-dashboard) + from miles.dashboard.hooks import register_train_actor + + register_train_actor(args) + unsupported = {"train_actor", "train_log_probs"} & set(args.profile_target) if unsupported and args.use_pytorch_profiler: raise NotImplementedError( diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 51e44d5e61..942f24f863 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -53,6 +53,18 @@ def _to_local_gpu_id(physical_gpu_id: int) -> int: ) +def _gpu_uuids(gpu_ids: list[int]) -> list: + """Best-effort NVML UUIDs so the dashboard can reconcile GPU index + spaces across processes; None entries when NVML is unavailable.""" + try: + import pynvml + + pynvml.nvmlInit() + return [str(pynvml.nvmlDeviceGetUUID(pynvml.nvmlDeviceGetHandleByIndex(i))) for i in gpu_ids] + except Exception: + return [None] * len(gpu_ids) + + def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process: from sglang.srt.entrypoints.http_server import launch_server @@ -126,6 +138,29 @@ def __init__( self.sglang_overrides = sglang_overrides or {} self.num_gpus_per_engine = num_gpus_per_engine + def get_topology_info(self) -> dict: + """Placement facts for the dashboard timeline (called after init()). + + ``base_gpu_id`` is node-physical (the local remap for the sglang + server happens later in ``_compute_server_args``), so these ids line + up with NVML device order used by the dashboard's GPU sampler. + """ + from miles.utils.misc import get_current_node_ip + + if self.base_gpu_id is None: # external engines: placement unknown + gpu_ids = [] + else: + gpus_on_node = min(self.num_gpus_per_engine, self.args.num_gpus_per_node) + gpu_ids = list(range(self.base_gpu_id, self.base_gpu_id + gpus_on_node)) + return dict( + url=f"http://{self.server_host}:{self.server_port}", + node_ip=get_current_node_ip(), + gpu_ids=gpu_ids, + gpu_uuids=_gpu_uuids(gpu_ids), + worker_type=self.worker_type, + node_rank=self.node_rank, + ) + def init( self, dist_init_addr, diff --git a/miles/dashboard/README.md b/miles/dashboard/README.md new file mode 100644 index 0000000000..3affd6ca10 --- /dev/null +++ b/miles/dashboard/README.md @@ -0,0 +1,63 @@ +# miles dashboard + +Training-dynamics and efficiency dashboards for miles runs, built on top of +`--dump-details` plus a live telemetry collector. + +## Collect + +```bash +python train.py ... \ + --dump-details /path/to/dump \ + --use-miles-dashboard \ + --use-rollout-entropy # optional: per-token entropy in the token view +``` + +`--use-miles-dashboard` adds a named collector actor on the driver node, one +NVML sampler per GPU node, per-rank phase interval sinks on the existing +`Timer` instrumentation, and an sglang engine-metric scraper. Everything is +appended under `{dump-details}/dashboard/` (append-only JSONL). Overhead on +the training path is a few milliseconds per step; all pushes are +fire-and-forget and a dead collector never affects training. + +## View + +```bash +pip install -e .[dashboard] # fastapi/uvicorn; already in the training image +python -m miles.dashboard.serve --dump-details /path/to/dump [--follow] [--port 7788] +``` + +Works on any machine that can see the directory (login node over NFS, the +training node itself); `--follow` tails a still-running job. Typical remote +usage is an SSH port-forward. Runs recorded *without* `--use-miles-dashboard` +still get the full training-dynamics views (metrics fall back to dump-derived +aggregates under the `dump/` namespace); only the timeline is absent. + +Views: + +- **metrics** — every logged metric plus `dump/*` per-step aggregates; click + a rollout-axis point to drill into that step. +- **timeline** — per-GPU lanes: phase band (rollout / train / update_weights, + hatched `train_wait` idle), NVML utilization, sglang engine overlay + (running requests, throughput, KV usage), bubble strip with click-to-zoom. +- **step drill-down** — per-trajectory table and scatter, GRPO group + degeneracy (`zero_std`), generation-time columns, eval tab. +- **token view** — per-token importance ratio / entropy / advantage strips + with loss-masked regions dimmed. + +## Develop + +```bash +python -m miles.dashboard.serve --demo # generated demo data, no cluster needed +python -m pytest tests/fast/dashboard/ -q +MILES_DASHBOARD_REALDATA_DIR=/path/to/real/dump python -m pytest tests/fast/dashboard/ -q +``` + +Design and implementation-plan documents live with the dashboard workspace +(`miles-dashboard-design.md`); the short version of the architecture: + +``` +producers (Timer sinks / rollout hooks / NVML samplers / sglang scraper) + -> DashboardCollector (named actor, driver node) [writes JSONL streams] +dump_details .pt files (existing dump path) [written by training] + -> serve.py: MetricStore + DumpReader -> FastAPI -> static SPA +``` diff --git a/miles/dashboard/__init__.py b/miles/dashboard/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/miles/dashboard/args.py b/miles/dashboard/args.py new file mode 100644 index 0000000000..36b9223e60 --- /dev/null +++ b/miles/dashboard/args.py @@ -0,0 +1,92 @@ +"""CLI arguments and configuration plumbing for the miles dashboard.""" + +from __future__ import annotations + +import logging + +from miles.dashboard.collector import CollectorConfig +from miles.dashboard.sglang_scraper import DEFAULT_METRIC_WHITELIST + +logger = logging.getLogger(__name__) + +# curated subset of args persisted into meta.json for the dashboard header +_SNAPSHOT_KEYS = ( + "wandb_group", + "colocate", + "num_gpus_per_node", + "actor_num_nodes", + "actor_num_gpus_per_node", + "rollout_num_gpus", + "rollout_num_gpus_per_engine", + "rollout_batch_size", + "n_samples_per_prompt", + "hf_checkpoint", +) + + +def add_dashboard_arguments(parser) -> None: + group = parser.add_argument_group("miles dashboard") + group.add_argument( + "--use-miles-dashboard", + action="store_true", + default=False, + help="Collect dashboard telemetry (phases, GPU util, engine metrics) under {dump-details}/dashboard/. " + "Requires --dump-details. View with `python -m miles.dashboard.serve`.", + ) + group.add_argument("--dashboard-flush-interval", type=float, default=5.0, help="collector disk flush cadence (s)") + group.add_argument("--dashboard-gpu-sample-interval", type=float, default=1.0, help="NVML sampling cadence (s)") + group.add_argument("--dashboard-sglang-scrape-interval", type=float, default=2.0, help="engine scrape cadence (s)") + group.add_argument( + "--dashboard-sglang-scrape-mode", + type=str, + choices=["auto", "router", "direct"], + default="auto", + help="auto scrapes {router}/engine_metrics, or each engine's /metrics under --use-miles-router", + ) + group.add_argument( + "--dashboard-sglang-metrics", + type=str, + default=None, + help="comma-separated override of the scraped sglang metric whitelist", + ) + group.add_argument( + "--dashboard-forward-prometheus", + action="store_true", + default=False, + help="also push dashboard gauges to the --use-prometheus collector for external Grafana", + ) + + +def validate_dashboard_args(args) -> None: + if not args.use_miles_dashboard: + return + assert args.dump_details is not None, ( + "--use-miles-dashboard writes telemetry under {dump-details}/dashboard/ and the " + "trajectory views read the rollout/train dumps, so --dump-details is required" + ) + if not args.use_rollout_entropy: + logger.warning( + "--use-miles-dashboard without --use-rollout-entropy: per-token entropy " + "will be missing from the dashboard token view" + ) + + +def collector_config_from_args(args, *, start_ts: float) -> CollectorConfig: + if args.dashboard_sglang_metrics is not None: + whitelist = tuple(m for m in args.dashboard_sglang_metrics.split(",") if m) + assert whitelist, f"empty --dashboard-sglang-metrics: {args.dashboard_sglang_metrics!r}" + else: + whitelist = DEFAULT_METRIC_WHITELIST + snapshot = {key: getattr(args, key) for key in _SNAPSHOT_KEYS if hasattr(args, key)} + return CollectorConfig( + dashboard_dir=f"{args.dump_details}/dashboard", + run_name=args.wandb_group or "miles-run", + start_ts=start_ts, + args_snapshot=snapshot, + flush_interval_seconds=args.dashboard_flush_interval, + gpu_sample_interval_seconds=args.dashboard_gpu_sample_interval, + scrape_interval_seconds=args.dashboard_sglang_scrape_interval, + scrape_mode=args.dashboard_sglang_scrape_mode, + metric_whitelist=whitelist, + forward_prometheus=args.dashboard_forward_prometheus, + ) diff --git a/miles/dashboard/backend.py b/miles/dashboard/backend.py new file mode 100644 index 0000000000..4ee08dd570 --- /dev/null +++ b/miles/dashboard/backend.py @@ -0,0 +1,171 @@ +"""Dashboard tracking-backend internals and Ray lifecycle glue. + +Call matrix of ``init_dashboard`` (design doc §6.1), reached through the +``MilesDashboardBackend`` adapter in ``tracking_utils``: + +- ``train.py`` driver (``primary=True``): builds the config, creates the + named ``DashboardCollector`` actor pinned to the driver node (same pattern + as ``prometheus_utils``), starts its flush loop, and spawns one + ``GpuSampler`` actor per GPU node. +- train actor main rank (``primary=False``): resolves the named actor so + ``dashboard_log`` works there. +- rollout manager (``primary=False, router_addr=...``): additionally attaches + a Timer phase sink for ``rollout`` / ``eval_rollout`` events. The router + itself is registered later through ``hooks.register_router`` — at + init_tracking time the router has not started yet. + +All pushes are fire-and-forget; when the collector cannot be resolved the +module warns once and every later call is a no-op — the dashboard must never +take a training run down with it. +""" + +from __future__ import annotations + +import logging +import time + +from miles.dashboard import hooks +from miles.dashboard.args import collector_config_from_args +from miles.dashboard.collector import COLLECTOR_ACTOR_NAME +from miles.dashboard.logging_utils import RateLimitedWarner +from miles.dashboard.store import MetricsRecord, Role + +logger = logging.getLogger(__name__) + +GET_ACTOR_TIMEOUT_SECONDS = 60.0 +GET_ACTOR_INTERVAL_SECONDS = 2.0 + +_handle = None +_is_primary = False +_resolution_failed = False +_warner = RateLimitedWarner(logger) + + +def init_dashboard(args, *, primary: bool = True, router_addr: str | None = None, **kwargs) -> None: + global _handle, _is_primary + import ray + + if primary: + _is_primary = True + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + from miles.dashboard.collector import DashboardCollector + + config = collector_config_from_args(args, start_ts=time.time()) + _handle = ( + ray.remote(DashboardCollector) + .options( + name=COLLECTOR_ACTOR_NAME, + # driver node, hard-pinned: the dump volume and the run's other + # observability live there, and driver death ends the job anyway + scheduling_strategy=NodeAffinitySchedulingStrategy( + node_id=ray.get_runtime_context().get_node_id(), soft=False + ), + ) + .remote(config, prometheus_handle_factory=_prometheus_factory if config.forward_prometheus else None) + ) + ray.get(_handle.ping.remote()) + _handle.start.remote() + logger.info( + "miles dashboard: telemetry -> %s | view live: python -m miles.dashboard.serve " + "--dump-details %s --follow --port 7788", + config.dashboard_dir, + args.dump_details, + ) + return + + if resolve_collector() is None: + return + if router_addr is not None: + # the kwarg marks the rollout-manager process. It cannot be used for + # scraping: init_tracking runs before start_rollout_servers, so the + # address is still "http://None:None" here — the real registration + # happens via hooks.register_router once the router is up. + hooks.attach_phase_sink(_handle, Role.ROLLOUT_MANAGER) + hooks.attach_trajectory_sink(_handle) + + +def resolve_collector(): + """This process's collector handle, resolving the named actor on first + use (train ranks never run init_dashboard themselves). Returns None — + permanently, after one warning — when the actor cannot be found.""" + global _handle, _resolution_failed + if _handle is not None or _resolution_failed: + return _handle + import ray + + deadline = time.monotonic() + GET_ACTOR_TIMEOUT_SECONDS + while True: + try: + _handle = ray.get_actor(COLLECTOR_ACTOR_NAME) + return _handle + except ValueError: + if time.monotonic() >= deadline: + logger.warning( + "dashboard collector actor not found after %.0fs; telemetry from this process is disabled", + GET_ACTOR_TIMEOUT_SECONDS, + ) + _resolution_failed = True + return None + time.sleep(GET_ACTOR_INTERVAL_SECONDS) + + +def current_collector(): + """The already-resolved handle, or None. Never blocks (hot paths).""" + return _handle + + +def dashboard_log(metrics: dict, *, step: int | None = None, step_key: str | None = None) -> None: + if _handle is None: + return + try: + _handle.push_metrics.remote( + MetricsRecord(ts=time.time(), step_key=step_key, step=step, metrics=_scalars_only(metrics)) + ) + except Exception: + _warner.warn("dashboard metric push failed; dropping this log call") + + +def finish_dashboard() -> None: + global _handle, _is_primary, _resolution_failed + hooks.detach_and_flush() + if _handle is not None and _is_primary: + import ray + + try: + # synchronous: the final flush must land before the driver exits + ray.get(_handle.shutdown.remote(), timeout=30) + except Exception: + logger.warning("dashboard collector shutdown incomplete", exc_info=True) + ray.kill(_handle) + _handle = None + _is_primary = False + _resolution_failed = False + + +# ------------------------------ ray helpers --------------------------------- + + +def _prometheus_factory(): + """Resolves the miles prometheus collector from inside the dashboard + collector actor (its module-global handle is per-process).""" + import ray + + from miles.utils.tracking_utils.prometheus_utils import _COLLECTOR_ACTOR_NAME + + try: + return ray.get_actor(_COLLECTOR_ACTOR_NAME) + except ValueError: + return None + + +def _scalars_only(metrics: dict) -> dict: + """The JSONL store persists plain scalars; tracking payloads can carry + tensors or numpy values, which are converted or dropped here.""" + out = {} + for key, value in metrics.items(): + if isinstance(value, (int, float, str, bool)): + out[key] = value + elif hasattr(type(value), "__float__"): # numpy/torch zero-dim scalars + out[key] = float(value) + return out diff --git a/miles/dashboard/collector.py b/miles/dashboard/collector.py new file mode 100644 index 0000000000..60bcc5d406 --- /dev/null +++ b/miles/dashboard/collector.py @@ -0,0 +1,362 @@ +"""The dashboard ingest hub. + +Producers (Timer sinks on every rank, the rollout manager's hooks, per-node +GPU samplers) push records here; the collector buffers them, appends to the +JSONL streams under ``{dump_details}/dashboard/`` on a flush cadence, runs +the sglang scraper thread once a router is registered, and optionally +forwards a latest-value snapshot to the existing Prometheus collector for +external Grafana. + +This class is deliberately Ray-free: the backend glue (``backend.py``) wraps +it in a named Ray actor pinned to the driver node and spawns the per-node +sampler actors — the collector itself only ever sees plain method calls, so +every behavior here is unit-testable. Producers call in fire-and-forget +style; nothing in the training path ever waits on this class. + +Failure policy: if the disk write fails (disk full, NFS hiccup) the error is +logged LOUDLY on every flush attempt — never masked — while ingestion keeps +running with bounded buffers (oldest records dropped past the cap) so a disk +problem cannot OOM the driver. +""" + +from __future__ import annotations + +import logging +import sys +import threading +import time +from dataclasses import dataclass, field +from typing import Any, ClassVar + +from miles.dashboard.sglang_scraper import DEFAULT_METRIC_WHITELIST, ScrapeMode, SglangScraper +from miles.dashboard.store import ( + EngineInfo, + EngineSample, + GpuSample, + Meta, + MetricsRecord, + MetricStore, + PhaseEvent, + Record, + TopologySnapshot, + TrajectoryEvent, +) + +logger = logging.getLogger(__name__) + +COLLECTOR_ACTOR_NAME = "miles_dashboard_collector" + + +class _SelfGpuPush: + """GpuSampler sink pushing back into this collector actor (the sampler + runs on its own node; the handle crosses the process boundary).""" + + def __init__(self, handle): + self._handle = handle + + def __call__(self, node: str, batch: list[GpuSample]) -> None: + self._handle.push_gpu_samples.remote(node, batch) + + +def _default_list_gpu_nodes() -> list[tuple[str, str]]: + # no initialized ray in this process = not running as the collector actor + # (unit tests, tooling): nothing to reconcile, and never pay the import + ray = sys.modules.get("ray") + if ray is None or not ray.is_initialized(): + return [] + return [ + (node["NodeID"], node["NodeManagerAddress"]) + for node in ray.nodes() + if node.get("Alive") and node.get("Resources", {}).get("GPU", 0) > 0 + ] + + +def _default_spawn_sampler(node_id: str, node_ip: str, interval: float): + """Spawn + start a GpuSampler pinned to one node; None when NVML is + unavailable there (recorded so the node is not retried every tick).""" + import ray + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + from miles.dashboard.gpu_sampler import GpuSampler + + handle = ( + ray.remote(GpuSampler) + .options( + num_cpus=0, + scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=node_id, soft=False), + ) + .remote(_SelfGpuPush(ray.get_runtime_context().current_actor), node=node_ip, interval=interval) + ) + if ray.get(handle.start.remote()): + return handle + ray.kill(handle) # NVML unavailable; the sampler already warned + return None + + +def _default_kill_sampler(handle) -> None: + import ray + + ray.kill(handle) + + +# test seams; production always uses the defaults +_list_gpu_nodes = _default_list_gpu_nodes +_spawn_sampler = _default_spawn_sampler +_kill_sampler = _default_kill_sampler + + +@dataclass +class CollectorConfig: + dashboard_dir: str # {dump_details}/dashboard + run_name: str + start_ts: float + args_snapshot: dict[str, Any] = field(default_factory=dict) + flush_interval_seconds: float = 5.0 + gpu_sample_interval_seconds: float = 1.0 + scrape_interval_seconds: float = 2.0 + scrape_mode: str = "auto" # "auto" or a ScrapeMode value; auto resolves at set_router() + metric_whitelist: tuple[str, ...] = DEFAULT_METRIC_WHITELIST + forward_prometheus: bool = False + + +class DashboardCollector: + # bounded ingest buffers: past this many buffered records per stream the + # oldest are dropped (only reachable when flushing to disk keeps failing) + MAX_BUFFERED_PER_STREAM: ClassVar[int] = 500_000 + + def __init__( + self, + config: CollectorConfig, + *, + prometheus_handle_factory=None, # () -> handle with .update.remote(dict), or None + scraper_http_get=None, # test hook, forwarded to SglangScraper + ): + self.config = config + self._store = MetricStore(config.dashboard_dir) + self._store.write_meta(Meta(run_name=config.run_name, start_ts=config.start_ts, args=config.args_snapshot)) + self._lock = threading.Lock() + self._dropped_since_flush = 0 + self._last_topology: TopologySnapshot | None = None + self._scraper: SglangScraper | None = None + self._scraper_http_get = scraper_http_get + self._prometheus_handle_factory = prometheus_handle_factory + # latest-value caches for the Prometheus forwarding snapshot; kept + # separately because store buffers empty out on every flush + self._latest_gpu: dict[tuple[str, int], GpuSample] = {} + self._latest_running_reqs: dict[str, float] = {} + self._latest_phase_seconds: dict[str, float] = {} + self._scraped_engine_addrs: set[str] = set() + self._actor_engine_addrs: set[str] = set() + self._stop_event = threading.Event() + self._flush_thread: threading.Thread | None = None + # per-node GPU samplers, keyed by ray NodeID and reconciled on every + # flush tick: nodes joining AFTER startup (late rollout nodes, node + # restarts) get a sampler too — a one-shot spawn at init cannot. + # None marks a node where NVML is unavailable (never retried). + self._samplers: dict[str, Any] = {} + + # ------------------------------ lifecycle ------------------------------- + + def ping(self) -> bool: + return True + + def start(self) -> None: + assert self._flush_thread is None, "collector already started" + self._flush_thread = threading.Thread(target=self._run_flush_loop, name="dashboard-flush", daemon=True) + self._flush_thread.start() + + def shutdown(self) -> None: + if self._scraper is not None: + self._scraper.stop() + self._stop_event.set() + if self._flush_thread is not None: + self._flush_thread.join(timeout=self.config.flush_interval_seconds + 1) + for handle in self._samplers.values(): + if handle is not None: + _kill_sampler(handle) + self._samplers.clear() + self.flush() + + # ------------------------------- ingestion ------------------------------ + + def push_metrics(self, record: MetricsRecord) -> None: + self._append(record) + + def push_phases(self, batch: list[PhaseEvent]) -> None: + for event in batch: + self._append(event) + + def push_trajectories(self, batch: list[TrajectoryEvent]) -> None: + for event in batch: + self._append(event) + + def push_gpu_samples(self, node: str, batch: list[GpuSample]) -> None: + for sample in batch: + self._append(sample) + + def update_topology(self, snapshot: TopologySnapshot) -> None: + # an addr the actor path ever registered must not be resurrected as + # "external" after a restart retires it — scrape memory is a fallback + # for engines the actor path NEVER knew, not a liveness source + self._actor_engine_addrs.update(e.addr for e in snapshot.engines if e.worker_type != "external") + snapshot = self._with_external_engines(snapshot) + with self._lock: + if self._last_topology is not None and self._last_topology.engines == snapshot.engines: + return # steady-state re-registration; only changes are recorded + self._last_topology = snapshot + self._append(snapshot) + + def _with_external_engines(self, snapshot: TopologySnapshot) -> TopologySnapshot: + """Merge in engines known only from scraping: externally launched + sglang servers have no miles engine actor, so actor registration + never sees them — but every scraped sample carries their addr. Node + is the addr host; GPU placement is unknown (gpus=[]), which the + frontend resolves by node match.""" + covered = {engine.addr for engine in snapshot.engines} | self._actor_engine_addrs + synthetic = [ + EngineInfo(addr=addr, worker_type="external", engine_rank=-1, gpus=[], gpu_uuids=[]) + for addr in sorted(self._scraped_engine_addrs) + if addr not in covered + ] + if not synthetic: + return snapshot + return TopologySnapshot(ts=snapshot.ts, engines=snapshot.engines + synthetic) + + def _sync_external_topology(self) -> None: + """Flush-loop step (outside the ingest lock): fold newly scraped + engine addrs into the topology. _update_latest only RECORDS addrs — + it runs under self._lock, and update_topology takes the same lock.""" + with self._lock: + covered = set() if self._last_topology is None else {e.addr for e in self._last_topology.engines} + missing = self._scraped_engine_addrs - covered - self._actor_engine_addrs + base = [] if self._last_topology is None else list(self._last_topology.engines) + if missing: + self.update_topology(TopologySnapshot(ts=time.time(), engines=base)) + + def set_router(self, router_addr: str, *, use_miles_router: bool) -> None: + """Register the sglang router and start (or re-point) the scraper.""" + if self.config.scrape_mode == "auto": + mode = ScrapeMode.DIRECT if use_miles_router else ScrapeMode.ROUTER + else: + mode = ScrapeMode(self.config.scrape_mode) + # never hold the lock while stopping a scraper: its thread may be + # blocked on the same lock inside the _append sink (deadlock) + with self._lock: + previous = self._scraper + if previous is not None and previous.router_addr == router_addr and previous.mode == mode: + return + self._scraper = None + if previous is not None: + previous.stop() + kwargs = dict( + mode=mode, + router_addr=router_addr, + engine_addrs=self._current_engine_addrs, + interval=self.config.scrape_interval_seconds, + whitelist=self.config.metric_whitelist, + ) + if self._scraper_http_get is not None: + kwargs["http_get"] = self._scraper_http_get + scraper = SglangScraper(self._append, **kwargs) + with self._lock: + self._scraper = scraper + scraper.start() + logger.info("dashboard scraper started in %s mode against %s", mode, router_addr) + + def _current_engine_addrs(self) -> list[str]: + with self._lock: + if self._last_topology is None: + return [] + return [engine.addr for engine in self._last_topology.engines] + + def _append(self, record: Record) -> None: + with self._lock: + if self._store.buffered_count(record.stream) >= self.MAX_BUFFERED_PER_STREAM: + self._dropped_since_flush += self._store.drop_oldest_buffered(record.stream) + self._store.append(record) + self._update_latest(record) + + def _update_latest(self, record: Record) -> None: + if isinstance(record, GpuSample): + self._latest_gpu[(record.node, record.gpu)] = record + elif isinstance(record, EngineSample): + self._scraped_engine_addrs.add(record.addr) + if record.metric == "sglang_num_running_reqs": + self._latest_running_reqs[record.addr] = record.value + elif isinstance(record, PhaseEvent): + self._latest_phase_seconds[record.name] = record.t1 - record.t0 + + # -------------------------------- flushing ------------------------------ + + def _run_flush_loop(self) -> None: + while not self._stop_event.is_set(): + self._stop_event.wait(self.config.flush_interval_seconds) + try: + self._reconcile_samplers() + self._sync_external_topology() + except Exception: + logger.exception("dashboard sampler reconcile failed; will retry next tick") + self.flush() + + def _reconcile_samplers(self) -> None: + """Diff alive GPU nodes against owned samplers; spawn the missing. + + Runs on the flush cadence (a ray.nodes() call is cheap), so late- + joining nodes start reporting util within one flush interval. A node + restart changes its NodeID, which reads as gone + new -> respawn.""" + alive = _list_gpu_nodes() + alive_ids = {node_id for node_id, _ in alive} + for node_id in list(self._samplers): + if node_id not in alive_ids: + del self._samplers[node_id] # node gone; actor died with it + for node_id, node_ip in alive: + if node_id not in self._samplers: + self._samplers[node_id] = _spawn_sampler(node_id, node_ip, self.config.gpu_sample_interval_seconds) + + def flush(self) -> None: + with self._lock: + if self._dropped_since_flush: + logger.error( + "dashboard collector dropped %d records since the last flush " + "(buffers past the cap — is the disk full?)", + self._dropped_since_flush, + ) + self._dropped_since_flush = 0 + try: + self._store.flush() + except OSError: + # deliberately loud on EVERY failed flush: a disk problem must + # surface, not be masked; buffers stay bounded via _append + logger.exception("dashboard flush to %s failed; records stay buffered", self.config.dashboard_dir) + return + if self.config.forward_prometheus and self._prometheus_handle_factory is not None: + handle = self._prometheus_handle_factory() + if handle is not None: + handle.update.remote(self._prometheus_snapshot()) + + # -------------------------- prometheus forwarding ----------------------- + + def _prometheus_snapshot(self) -> dict[str, float]: + """Latest-value gauges for external Grafana. Keys avoid characters the + prometheus collector's sanitizer does not handle ('.', ':').""" + + def safe(text: str) -> str: + return text.replace("http://", "").replace(".", "_").replace(":", "_") + + with self._lock: + snapshot = { + f"dashboard/gpu_{safe(node)}_{gpu}_util": float(sample.util) + for (node, gpu), sample in self._latest_gpu.items() + } + snapshot |= { + f"dashboard/gpu_{safe(node)}_{gpu}_mem_mb": float(sample.mem_mb) + for (node, gpu), sample in self._latest_gpu.items() + } + snapshot |= { + f"dashboard/engine_{safe(addr)}_running_reqs": value + for addr, value in self._latest_running_reqs.items() + } + snapshot |= { + f"dashboard/phase_{name}_seconds": seconds for name, seconds in self._latest_phase_seconds.items() + } + return snapshot diff --git a/miles/dashboard/dump_reader.py b/miles/dashboard/dump_reader.py new file mode 100644 index 0000000000..400d30776a --- /dev/null +++ b/miles/dashboard/dump_reader.py @@ -0,0 +1,501 @@ +"""Lazy reader for ``--dump-details`` directories: discovery, loading, join. + +``rollout_data/{rid}.pt`` holds the full sample batch of one rollout step; +``train_data/{rid}_{rank}.pt`` holds that rank's DP shard with per-token +tensors and ``sample_indices`` mapping each row back to ``Sample.index``. +``load_joined()`` reunites the two: every rollout sample plus (for train +dumps) its per-token training-side row, deduplicated across TP-duplicate +rank files. + +Files being written concurrently by a live run are handled in two layers: +``rollout_ids()`` hides files younger than ``MIN_AGE_SECONDS`` unless their +train companion already exists, and a ``torch.load`` failure on a fresh file +raises :class:`DumpStillWriting` (the server maps it to HTTP 503) instead of +surfacing as corruption. +""" + +from __future__ import annotations + +import json +import time +from collections import OrderedDict +from dataclasses import dataclass +from pathlib import Path +from typing import Any, ClassVar + +import polars as pl +import torch + +from miles.utils.types import Sample + + +class DumpStillWriting(Exception): + """A dump file exists but cannot be used yet (``torch.save`` in progress).""" + + +# metric columns of DumpReader.step_aggregates(), i.e. the dump-derived L0 +# series (consumed by the server's metric catalog as "dump/") +STEP_AGGREGATE_METRICS = ( + "reward_mean", + "reward_std", + "response_length_mean", + "truncated_frac", + "zero_std_group_frac", + "mean_abs_lp_diff", + "mean_entropy", + "mixed_version_frac", +) + + +def _min_numeric_version(versions: list[str] | None) -> int | None: + numeric = [int(v) for v in versions or [] if str(v).isdigit()] + return min(numeric) if numeric else None + + +def _tool_call_count(sample: Sample) -> int | None: + """Tool messages in a chat-style prompt; None for plain-string prompts + (single-turn math runs have no message structure to count).""" + if not isinstance(sample.prompt, list): + return None + return sum(1 for message in sample.prompt if isinstance(message, dict) and message.get("role") == "tool") + + +@dataclass +class RolloutIds: + train: list[int] + eval: list[int] + + +@dataclass +class TrainRow: + """Per-sample slice of one rank's train dump. + + Required columns fail loudly when absent; columns that legitimately depend + on run configuration (entropy needs ``--use-rollout-entropy``, + ``ref_log_probs`` needs a KL term, ...) are ``None`` when not dumped. + + ``raw_reward`` is passed in separately by the caller: unlike every other + column it is stored batch-global (full batch, rollout order — see the + "splited at train side" block in ``split_train_data_by_dp``), so indexing + it by shard row would silently misattribute rewards. + """ + + sample_index: int + rank: int + tokens: torch.Tensor + response_length: int + total_length: int + reward: float + loss_mask: torch.Tensor + log_probs: torch.Tensor | None # absent when the run does not dump them + rollout_log_probs: torch.Tensor | None + ref_log_probs: torch.Tensor | None + entropy: torch.Tensor | None + ref_entropy: torch.Tensor | None + advantages: torch.Tensor | None + returns: torch.Tensor | None + raw_reward: Any + truncated: int | None + weight_versions: list[str] | None + + @classmethod + def from_columns(cls, columns: dict, row: int, *, rank: int, raw_reward) -> TrainRow: + def optional(key: str): + values = columns.get(key) + return None if values is None else values[row] + + return cls( + sample_index=columns["sample_indices"][row], + rank=rank, + tokens=columns["tokens"][row], + response_length=columns["response_lengths"][row], + total_length=columns["total_lengths"][row], + reward=columns["rewards"][row], + loss_mask=columns["loss_masks"][row], + log_probs=optional("log_probs"), + rollout_log_probs=optional("rollout_log_probs"), + ref_log_probs=optional("ref_log_probs"), + entropy=optional("entropy"), + ref_entropy=optional("ref_entropy"), + advantages=optional("advantages"), + returns=optional("returns"), + raw_reward=raw_reward, + truncated=optional("truncated"), + weight_versions=optional("weight_versions"), + ) + + +@dataclass +class JoinedRollout: + rollout_id: int + evaluation: bool + samples: list[Sample] + train_rows: dict[int, TrainRow] # keyed by Sample.index; empty for eval dumps + + @property + def train_coverage(self) -> float: + return len(self.train_rows) / len(self.samples) if self.samples else 0.0 + + +class DumpReader: + # A fresh rollout file is only trusted once its train companion exists + # (written strictly after it) or it has stopped changing for this long. + MIN_AGE_SECONDS: ClassVar[float] = 10.0 + # torch.load failures on files younger than this are "still being written"; + # on older files they are real corruption and propagate. + FRESH_SECONDS: ClassVar[float] = 60.0 + + # bump to invalidate summary parquet caches when their columns change + SUMMARY_VERSION: ClassVar[int] = 2 # v2: staleness/turns/tool columns + + def __init__(self, dump_dir: Path | str, *, cache_dir: Path | str | None = None, tensor_lru: int = 2): + self.dump_dir = Path(dump_dir) + self.rollout_dir = self.dump_dir / "rollout_data" + self.train_dir = self.dump_dir / "train_data" + self.cache_dir = Path(cache_dir) if cache_dir is not None else self.dump_dir / "dashboard" / "cache" + self.tensor_lru = tensor_lru + self._joined_cache: OrderedDict[tuple[int, bool], JoinedRollout] = OrderedDict() + self._trajectory_cache: OrderedDict[tuple[int, bool], dict[int, dict]] = OrderedDict() + self._tokenizer = None + self._tokenizer_loaded = False + + def rollout_ids(self) -> RolloutIds: + ids = RolloutIds(train=[], eval=[]) + if not self.rollout_dir.is_dir(): + return ids + now = time.time() + for path in self.rollout_dir.glob("*.pt"): + evaluation = path.stem.startswith("eval_") + rollout_id = int(path.stem.removeprefix("eval_")) + if self._visible(path, rollout_id, evaluation=evaluation, now=now): + (ids.eval if evaluation else ids.train).append(rollout_id) + ids.train.sort() + ids.eval.sort() + return ids + + def load_joined(self, rollout_id: int, *, evaluation: bool = False) -> JoinedRollout: + name = f"eval_{rollout_id}.pt" if evaluation else f"{rollout_id}.pt" + pack = self._torch_load(self.rollout_dir / name) + assert pack["rollout_id"] == rollout_id, f"{pack['rollout_id']=} != {rollout_id=} in {name}" + samples = [Sample.from_dict(data) for data in pack["samples"]] + # raw_reward is stored batch-global, so it is indexed by the sample's + # position in the rollout dump, not by shard row. + batch_position = {s.index: i for i, s in enumerate(samples)} + + train_rows: dict[int, TrainRow] = {} + if not evaluation: + for path in self._train_paths(rollout_id): + rank_pack = self._torch_load(path) + columns = rank_pack["rollout_data"] + raw_rewards = columns.get("raw_reward") + if raw_rewards is not None: + assert len(raw_rewards) == len(samples), ( + f"{path}: raw_reward must be batch-global " + f"(expected {len(samples)} entries, got {len(raw_rewards)})" + ) + for row, sample_index in enumerate(columns["sample_indices"]): + assert ( + sample_index in batch_position + ), f"{path} references sample_index {sample_index} absent from the rollout dump" + if (existing := train_rows.get(sample_index)) is not None: + # TP-duplicate rank carrying the same DP shard. + assert existing.response_length == columns["response_lengths"][row], ( + f"rank {rank_pack['rank']} disagrees with rank {existing.rank} " + f"on sample {sample_index} in rollout {rollout_id}" + ) + continue + train_rows[sample_index] = TrainRow.from_columns( + columns, + row, + rank=rank_pack["rank"], + raw_reward=None if raw_rewards is None else raw_rewards[batch_position[sample_index]], + ) + + return JoinedRollout(rollout_id=rollout_id, evaluation=evaluation, samples=samples, train_rows=train_rows) + + def joined(self, rollout_id: int, *, evaluation: bool = False) -> JoinedRollout: + """LRU-cached :meth:`load_joined`. A completed rollout's dumps never + change, so entries stay valid; the LRU (``tensor_lru`` ids resident) + bounds memory since one id holds every per-token tensor of its step.""" + key = (rollout_id, evaluation) + if key in self._joined_cache: + self._joined_cache.move_to_end(key) + return self._joined_cache[key] + result = self.load_joined(rollout_id, evaluation=evaluation) + self._joined_cache[key] = result + while len(self._joined_cache) > self.tensor_lru: + self._joined_cache.popitem(last=False) + return result + + @property + def tokenizer(self): + """Tokenizer persisted by the run's data source, or None if absent.""" + if not self._tokenizer_loaded: + self._tokenizer_loaded = True + tokenizer_dir = self.dump_dir / "tokenizer" + if tokenizer_dir.is_dir(): + from miles.utils.processing_utils import load_tokenizer + + self._tokenizer = load_tokenizer(str(tokenizer_dir)) + return self._tokenizer + + # ------------------------------- L1 views ------------------------------- + + def summary(self, rollout_id: int, *, evaluation: bool = False) -> pl.DataFrame: + """Per-sample summary table (one row per Sample), parquet-cached under + ``cache_dir`` and invalidated on source mtime or SUMMARY_VERSION change.""" + stem = f"rollout_{'eval_' if evaluation else ''}{rollout_id}" + cache_path = self.cache_dir / f"{stem}.parquet" + sources_path = self.cache_dir / f"{stem}.sources.json" + sources = self._source_stamps(rollout_id, evaluation=evaluation) + if cache_path.exists() and sources_path.exists() and json.loads(sources_path.read_text()) == sources: + return pl.read_parquet(cache_path) + + joined = self.joined(rollout_id, evaluation=evaluation) + df = pl.DataFrame([self._summary_row(s, joined.train_rows.get(s.index)) for s in joined.samples], strict=False) + self.cache_dir.mkdir(parents=True, exist_ok=True) + df.write_parquet(cache_path) + sources_path.write_text(json.dumps(sources)) + return df + + def groups(self, rollout_id: int, *, evaluation: bool = False) -> pl.DataFrame: + """Per-GRPO-group aggregates; ``zero_std`` flags degenerate groups + (all samples got the same reward, so advantages vanish).""" + reward_column = "reward" if evaluation else "raw_reward" + return ( + self.summary(rollout_id, evaluation=evaluation) + .group_by("group_index") + .agg( + n=pl.len(), + reward_mean=pl.col(reward_column).mean(), + reward_std=pl.col(reward_column).std(), + response_length_mean=pl.col("response_length").mean(), + truncated_frac=pl.col("truncated").cast(pl.Float64).mean(), + ) + .with_columns(zero_std=pl.col("reward_std").fill_null(0.0) <= 1e-12) + .sort("group_index") + ) + + def step_aggregates(self) -> pl.DataFrame: + """Dump-derived per-step series: the L0 fallback when no metrics.jsonl + exists. First call computes (and parquet-caches) every step's summary.""" + rows = [] + for rollout_id in self.rollout_ids().train: + df = self.summary(rollout_id) + groups = self.groups(rollout_id) + rows.append( + dict( + rollout_id=rollout_id, + n_samples=df.height, + reward_mean=df["raw_reward"].mean(), + reward_std=df["raw_reward"].std(), + response_length_mean=df["response_length"].mean(), + truncated_frac=df["truncated"].cast(pl.Float64).mean(), + zero_std_group_frac=groups["zero_std"].cast(pl.Float64).mean(), + mean_abs_lp_diff=df["mean_abs_lp_diff"].mean(), + mean_entropy=df["mean_entropy"].mean(), + mixed_version_frac=df["mixed_version"].cast(pl.Float64).mean(), + ) + ) + return pl.DataFrame(rows, strict=False) + + def trajectory_messages(self, rollout_id: int, sample_index: int, *, evaluation: bool = False) -> dict: + """Sidecar row for one sample; missing file or sample raises (-> 404), + which is how the frontend learns the run recorded no conversation.""" + key = (rollout_id, evaluation) + if key not in self._trajectory_cache: + name = f"eval_{rollout_id}.jsonl" if evaluation else f"{rollout_id}.jsonl" + with open(self.dump_dir / "trajectory" / name) as f: + rows = {row["sample_index"]: row for row in map(json.loads, f)} + self._trajectory_cache[key] = rows + while len(self._trajectory_cache) > 4: + self._trajectory_cache.popitem(last=False) + self._trajectory_cache.move_to_end(key) + rows = self._trajectory_cache[key] + if sample_index not in rows: + raise KeyError(f"sample {sample_index} has no recorded conversation in rollout {rollout_id}") + return rows[sample_index] + + # ------------------------------- L2 view -------------------------------- + + def tokens( + self, rollout_id: int, sample_index: int, *, start: int = 0, end: int | None = None, evaluation: bool = False + ) -> dict: + """Per-token payload for one sample over token positions [start, end). + + Token ids/text cover the whole requested slice; per-token stat arrays + cover only its overlap with the response region (stat ``i`` maps to + token position ``prompt_len + a + i``). ``response_offset`` is the + index within the returned token slice where the response begins. + """ + joined = self.joined(rollout_id, evaluation=evaluation) + sample = next((s for s in joined.samples if s.index == sample_index), None) + if sample is None: + raise KeyError(f"unknown sample_index {sample_index} in rollout {rollout_id}") + row = joined.train_rows.get(sample_index) + + total = len(sample.tokens) + prompt_len = total - sample.response_length + start = max(0, start) + end = total if end is None else min(end, total) + if start >= end: + raise ValueError(f"empty token range [{start}, {end}) for total_len={total}") + a = max(0, start - prompt_len) + b = max(0, end - prompt_len) + + def response_slice(values) -> list[float] | None: + return None if values is None else [float(v) for v in values[a:b]] + + token_ids = [int(t) for t in sample.tokens[start:end]] + lp_diff = ( + row.log_probs - row.rollout_log_probs + if row is not None and row.log_probs is not None and row.rollout_log_probs is not None + else None + ) + return dict( + rollout_id=rollout_id, + sample_index=sample_index, + evaluation=evaluation, + total_len=total, + prompt_len=prompt_len, + start=start, + end=end, + response_offset=min(len(token_ids), max(0, prompt_len - start)), + token_ids=token_ids, + token_text=self._decode_tokens(token_ids), + rollout_log_probs=( + response_slice(sample.rollout_log_probs) + if sample.rollout_log_probs is not None + else response_slice(row.rollout_log_probs) if row is not None else None + ), + loss_mask=None if row is None else [int(v) for v in row.loss_mask[a:b]], + train_log_probs=None if row is None or row.log_probs is None else response_slice(row.log_probs), + ref_log_probs=None if row is None else response_slice(row.ref_log_probs), + lp_diff=response_slice(lp_diff), + imp_ratio=None if lp_diff is None else response_slice(lp_diff.exp()), + entropy=None if row is None else response_slice(row.entropy), + ref_entropy=None if row is None else response_slice(row.ref_entropy), + advantages=None if row is None else response_slice(row.advantages), + returns=None if row is None else response_slice(row.returns), + ) + + # ------------------------------- internals ------------------------------ + + def _decode_tokens(self, token_ids: list[int]) -> list[str] | None: + if self.tokenizer is None: + return None + return [self.tokenizer.decode([token_id]) for token_id in token_ids] + + def _source_stamps(self, rollout_id: int, *, evaluation: bool) -> dict: + rollout_path = self.rollout_dir / (f"eval_{rollout_id}.pt" if evaluation else f"{rollout_id}.pt") + paths = [rollout_path] + ([] if evaluation else self._train_paths(rollout_id)) + return {"_summary_version": self.SUMMARY_VERSION, **{p.name: p.stat().st_mtime for p in paths}} + + def _summary_row(self, sample: Sample, row: TrainRow | None) -> dict: + spec = sample.spec_info + cache_info = sample.prefix_cache_info + entry = dict( + sample_index=sample.index, + group_index=sample.group_index, + status=sample.status.value, + remove_sample=sample.remove_sample, + response_length=sample.response_length, + total_length=len(sample.tokens), + reward=float(sample.reward) if isinstance(sample.reward, (int, float)) else None, + weight_version=sample.weight_versions[-1] if sample.weight_versions else None, + weight_version_min=_min_numeric_version(sample.weight_versions), + mixed_version=len(set(sample.weight_versions)) > 1 if sample.weight_versions else None, + turns=len(sample.weight_versions) if sample.weight_versions else None, + tool_calls=_tool_call_count(sample), + non_generation_time=sample.non_generation_time, + spec_accept_rate=( + spec.spec_accept_token_num / spec.spec_draft_token_num if spec.spec_draft_token_num else None + ), + prefix_cache_hit_rate=( + cache_info.cached_tokens / cache_info.total_prompt_tokens if cache_info.total_prompt_tokens else None + ), + ) + if row is None: + train_columns = dict.fromkeys( + ( + "raw_reward", + "normalized_reward", + "dumped_rank", + "mean_entropy", + "max_entropy", + "ref_entropy_mean", + "mean_abs_lp_diff", + "max_abs_lp_diff", + "mean_imp_ratio", + "adv_mean", + "adv_std", + "return_mean", + ) + ) + train_columns["truncated"] = sample.status == Sample.Status.TRUNCATED + return entry | train_columns + + mask = row.loss_mask > 0 + lp_diff = ( + None if row.log_probs is None or row.rollout_log_probs is None else row.log_probs - row.rollout_log_probs + ) + entropy = _masked(row.entropy, mask) + abs_diff = _masked(None if lp_diff is None else lp_diff.abs(), mask) + advantages = _masked(row.advantages, mask) + return entry | dict( + raw_reward=None if row.raw_reward is None else float(row.raw_reward), + normalized_reward=float(row.reward), + truncated=bool(row.truncated) if row.truncated is not None else sample.status == Sample.Status.TRUNCATED, + dumped_rank=row.rank, + mean_entropy=_mean(entropy), + max_entropy=_max(entropy), + ref_entropy_mean=_mean(_masked(row.ref_entropy, mask)), + mean_abs_lp_diff=_mean(abs_diff), + max_abs_lp_diff=_max(abs_diff), + mean_imp_ratio=_mean(_masked(None if lp_diff is None else lp_diff.exp(), mask)), + adv_mean=_mean(advantages), + adv_std=_std(advantages), + return_mean=_mean(_masked(row.returns, mask)), + ) + + def _train_paths(self, rollout_id: int) -> list[Path]: + return sorted(self.train_dir.glob(f"{rollout_id}_*.pt"), key=lambda p: int(p.stem.rsplit("_", 1)[1])) + + def _visible(self, path: Path, rollout_id: int, *, evaluation: bool, now: float) -> bool: + if now - path.stat().st_mtime > self.MIN_AGE_SECONDS: + return True + return not evaluation and (self.train_dir / f"{rollout_id}_0.pt").exists() + + def _torch_load(self, path: Path): + try: + return torch.load(path, weights_only=False, map_location="cpu") + except FileNotFoundError: + raise + except Exception as e: + if time.time() - path.stat().st_mtime < self.FRESH_SECONDS: + raise DumpStillWriting(str(path)) from e + raise + + +# ---------------------- masked per-token statistics ------------------------- + + +def _masked(values: torch.Tensor | None, mask: torch.Tensor) -> torch.Tensor | None: + """Loss-masked positions (tool outputs, removed samples) are excluded from + all summary statistics; an empty selection yields None, not NaN.""" + if values is None: + return None + selected = values[mask] + return selected.float() if selected.numel() else None + + +def _mean(values: torch.Tensor | None) -> float | None: + return None if values is None else float(values.mean()) + + +def _max(values: torch.Tensor | None) -> float | None: + return None if values is None else float(values.max()) + + +def _std(values: torch.Tensor | None) -> float | None: + return None if values is None else float(values.std()) diff --git a/miles/dashboard/gpu_sampler.py b/miles/dashboard/gpu_sampler.py new file mode 100644 index 0000000000..5a2e935101 --- /dev/null +++ b/miles/dashboard/gpu_sampler.py @@ -0,0 +1,127 @@ +"""Per-node GPU utilization sampler feeding the dashboard timeline. + +One instance runs per GPU node (the collector spawns it as a Ray actor with +``NodeAffinitySchedulingStrategy``; the class itself is plain Python and +unit-testable). A daemon thread samples every NVML device on the node at +``interval`` seconds — physical device order, independent of +``CUDA_VISIBLE_DEVICES`` — buffers locally, and hands batches to the injected +``push(node, batch)`` callable (the collector wraps its own Ray handle) every +``FLUSH_INTERVAL_SECONDS``, so there is roughly one RPC per node per flush +rather than one per sample. + +Degradation: NVML being unavailable (no pynvml, driver mismatch) disables the +sampler with a single warning — the timeline just lacks the util band. A +device that fails mid-run (e.g. during a GPU reset) is skipped for that tick +with rate-limited warnings; the other devices keep reporting. +""" + +from __future__ import annotations + +import logging +import threading +import time +from collections.abc import Callable +from typing import ClassVar + +from miles.dashboard.logging_utils import RateLimitedWarner +from miles.dashboard.store import GpuSample + +logger = logging.getLogger(__name__) + + +class GpuSampler: + FLUSH_INTERVAL_SECONDS: ClassVar[float] = 5.0 + + def __init__( + self, + push: Callable[[str, list[GpuSample]], None], + *, + node: str, + interval: float = 1.0, + nvml=None, + ): + assert interval > 0, f"{interval=}" + self._push = push + self.node = node + self.interval = interval + self._nvml = nvml + self._handles: list = [] + self._uuids: list[str] = [] + self._buffer: list[GpuSample] = [] + self._buffer_lock = threading.Lock() + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + self._warner = RateLimitedWarner(logger) + self.available = self._init_nvml() + + def _init_nvml(self) -> bool: + try: + if self._nvml is None: + import pynvml + + self._nvml = pynvml + self._nvml.nvmlInit() + count = self._nvml.nvmlDeviceGetCount() + self._handles = [self._nvml.nvmlDeviceGetHandleByIndex(i) for i in range(count)] + self._uuids = [str(self._nvml.nvmlDeviceGetUUID(handle)) for handle in self._handles] + return True + except Exception as e: + logger.warning("NVML unavailable on %s (%s); GPU utilization will not be collected", self.node, e) + return False + + # ------------------------------ lifecycle ------------------------------- + + def gpu_uuids(self) -> list[str]: + return list(self._uuids) + + def start(self) -> bool: + """Begin sampling; returns False (and stays inert) when NVML is unavailable.""" + if not self.available: + return False + assert self._thread is None, "sampler already started" + self._thread = threading.Thread(target=self._run, name="dashboard-gpu-sampler", daemon=True) + self._thread.start() + return True + + def stop(self) -> None: + self._stop_event.set() + if self._thread is not None: + self._thread.join(timeout=self.interval + self.FLUSH_INTERVAL_SECONDS) + self.flush() + + def _run(self) -> None: + next_flush = time.monotonic() + self.FLUSH_INTERVAL_SECONDS + while not self._stop_event.is_set(): + self.sample_once(time.time()) + if time.monotonic() >= next_flush: + self.flush() + next_flush = time.monotonic() + self.FLUSH_INTERVAL_SECONDS + self._stop_event.wait(self.interval) + + # -------------------------------- sampling ------------------------------ + + def sample_once(self, ts: float) -> int: + """Sample every device once into the buffer. Returns the sample count.""" + if not self.available: + return 0 + count = 0 + for gpu, handle in enumerate(self._handles): + try: + util = int(self._nvml.nvmlDeviceGetUtilizationRates(handle).gpu) + mem_mb = int(self._nvml.nvmlDeviceGetMemoryInfo(handle).used) >> 20 + power_w = int(self._nvml.nvmlDeviceGetPowerUsage(handle)) // 1000 + except Exception: + self._warner.warn(f"NVML read failed for gpu {gpu} on {self.node}; skipping this tick") + continue + with self._buffer_lock: + self._buffer.append( + GpuSample(ts=ts, node=self.node, gpu=gpu, util=util, mem_mb=mem_mb, power_w=power_w) + ) + count += 1 + return count + + def flush(self) -> None: + with self._buffer_lock: + batch, self._buffer = self._buffer, [] + if batch: + self._push(self.node, batch) diff --git a/miles/dashboard/hooks.py b/miles/dashboard/hooks.py new file mode 100644 index 0000000000..8be5ea66a7 --- /dev/null +++ b/miles/dashboard/hooks.py @@ -0,0 +1,349 @@ +"""Process-side hooks feeding the dashboard collector. + +These are the functions the (tiny) core wiring calls: a Timer event sink on +every training rank, and engine-topology registration from the rollout +manager. Every entry point is a no-op when the dashboard is disabled or the +collector cannot be reached, and every sink swallows its own exceptions with +rate-limited warnings — a deliberate exception to fail-loud (design doc +§6.3): observability must never be able to kill a training step. +""" + +from __future__ import annotations + +import logging +import threading +import time +from dataclasses import dataclass + +from miles.dashboard.logging_utils import RateLimitedWarner +from miles.dashboard.store import EngineInfo, PhaseEvent, Role, TopologySnapshot, TrajectoryEvent, TrajectoryEventKind +from miles.utils.lifecycle import TrajectoryLifecycle +from miles.utils.timer import Timer + +logger = logging.getLogger(__name__) + +BATCH_MAX_EVENTS = 64 +BATCH_MAX_SECONDS = 2.0 + + +@dataclass +class _Identity: + node: str + gpus: list[int] + rank: int + + +def _default_resolve_identity() -> _Identity: + import ray + import torch.distributed as dist + + from miles.utils.misc import get_current_node_ip + + return _Identity( + node=get_current_node_ip(), + gpus=[int(gpu) for gpu in ray.get_gpu_ids()], + rank=dist.get_rank() if dist.is_initialized() else -1, + ) + + +def _default_ray_get(refs: list): + import ray + + return ray.get(refs) + + +# test seams; production always uses the defaults +_resolve_identity = _default_resolve_identity +_ray_get = _default_ray_get + + +class PhaseSink: + """Timer event sink: stamps process identity onto intervals and pushes + them to the collector in batches (fire-and-forget).""" + + def __init__(self, handle, role: str): + self._handle = handle + self.role = role + self._buffer: list[PhaseEvent] = [] + self._lock = threading.Lock() + self._last_flush = time.monotonic() + self._identity: _Identity | None = None + self._warner = RateLimitedWarner(logger) + + def begin(self, name: str, t0: float) -> None: + """Timer.start notification: push an OPEN interval immediately (no + batching — starts are rare and while-open visibility is the point). + The closing event supersedes it on the read side.""" + try: + with self._lock: + if self._identity is None or (self._identity.rank < 0 and self.role == Role.TRAIN): + self._identity = _resolve_identity() + identity = self._identity + event = PhaseEvent( + name=name, + t0=t0, + t1=PhaseEvent.OPEN_T1, + node=identity.node, + gpus=identity.gpus, + rank=identity.rank, + role=self.role, + ) + self._handle.push_phases.remote([event]) + except Exception: + self._warner.warn("dashboard phase sink failed; dropping events") + + def __call__(self, name: str, t0: float, t1: float) -> None: + try: + with self._lock: + # lazy: torch.distributed is usually not initialized yet when + # the sink attaches; re-resolve until a real rank appears + if self._identity is None or (self._identity.rank < 0 and self.role == Role.TRAIN): + self._identity = _resolve_identity() + identity = self._identity + self._buffer.append( + PhaseEvent( + name=name, + t0=t0, + t1=t1, + node=identity.node, + gpus=identity.gpus, + rank=identity.rank, + role=self.role, + ) + ) + batch = self._take_batch_if_due() + if batch: + self._handle.push_phases.remote(batch) + except Exception: + self._warner.warn("dashboard phase sink failed; dropping events") + + def _take_batch_if_due(self) -> list[PhaseEvent] | None: + if len(self._buffer) < BATCH_MAX_EVENTS and time.monotonic() - self._last_flush < BATCH_MAX_SECONDS: + return None + batch, self._buffer = self._buffer, [] + self._last_flush = time.monotonic() + return batch + + def flush(self) -> None: + try: + with self._lock: + batch, self._buffer = self._buffer, [] + if batch: + self._handle.push_phases.remote(batch) + except Exception: + self._warner.warn("dashboard phase sink flush failed; dropping events") + + +class TrajectorySink: + """Per-sample lifecycle events from the rollout process, batched to the + collector (fire-and-forget). Core code reaches it via + ``miles.dashboard.lifecycle.trajectory_sink()`` — duck-typed so the core + call sites never import this module.""" + + def __init__(self, handle): + self._handle = handle + self._buffer: list[TrajectoryEvent] = [] + self._lock = threading.Lock() + self._last_flush = time.monotonic() + self._warner = RateLimitedWarner(logger) + + # ------------------------- core-facing emitters ------------------------- + + def attempt_start(self, sample) -> None: + self._emit(TrajectoryEventKind.ATTEMPT_START, sample) + + def gen_start(self, sample) -> None: + self._emit(TrajectoryEventKind.GEN_START, sample) + + def attempt_end(self, result) -> None: + """``result`` is the Sample or the list of per-turn Samples a generate + function returned; per-turn samples share the identity and carry their + own lifecycle metadata segments (agentic endpoint path).""" + samples = result if isinstance(result, list) else [result] + for sample in samples: + value = (sample.metadata or {}).get("lifecycle") + # one dict per turn-sample; a TITO-merged sample carries the list + for segment in value if isinstance(value, list) else [value] if value else []: + gap_end = segment.get("req_ts") or segment.get("t0") + if segment.get("prev_t1") is not None and gap_end is not None: + # the gap between chat calls is agent-side work (tool + # execution, environment steps) — design §18.3. req_ts + # (server-edge arrival) bounds it exactly; without it the + # gen start approximates and absorbs engine queueing + self.tool_span(sample, segment["prev_t1"], gap_end, turn=segment["turn"], detail="agent gap") + if segment.get("t0") is not None: + self._emit(TrajectoryEventKind.GEN_START, sample, ts=segment["t0"], turn=segment["turn"]) + self._emit(TrajectoryEventKind.GEN_END, sample, ts=segment["t1"], turn=segment["turn"]) + last = samples[-1] + status = last.status.value if getattr(last, "status", None) is not None else "" + self._emit(TrajectoryEventKind.ATTEMPT_END, last, detail=status) + + def gen_span(self, sample, t0: float, t1: float, turn: int, detail: str = "") -> None: + self._emit(TrajectoryEventKind.GEN_START, sample, ts=t0, turn=turn) + self._emit(TrajectoryEventKind.GEN_END, sample, ts=t1, turn=turn, detail=detail) + + def tool_span(self, sample, t0: float, t1: float, turn: int, detail: str = "") -> None: + self._emit(TrajectoryEventKind.TOOL_START, sample, ts=t0, turn=turn, detail=detail) + self._emit(TrajectoryEventKind.TOOL_END, sample, ts=t1, turn=turn, detail=detail) + + # ------------------------------ internals ------------------------------- + + def _emit(self, kind: str, sample, *, ts: float | None = None, turn: int = -1, detail: str = "") -> None: + try: + versions = getattr(sample, "weight_versions", None) or [] + event = TrajectoryEvent( + ts=time.time() if ts is None else ts, + kind=kind, + sample_index=sample.index if sample.index is not None else -1, + group_index=sample.group_index if sample.group_index is not None else -1, + turn=turn, + weight_version=str(versions[-1]) if versions else "", + detail=detail, + ) + with self._lock: + self._buffer.append(event) + batch = self._take_batch_if_due() + if batch: + self._handle.push_trajectories.remote(batch) + except Exception: + self._warner.warn("dashboard trajectory sink failed; dropping events") + + def _take_batch_if_due(self) -> list[TrajectoryEvent] | None: + if len(self._buffer) < BATCH_MAX_EVENTS and time.monotonic() - self._last_flush < BATCH_MAX_SECONDS: + return None + batch, self._buffer = self._buffer, [] + self._last_flush = time.monotonic() + return batch + + def flush(self) -> None: + try: + with self._lock: + batch, self._buffer = self._buffer, [] + if batch: + self._handle.push_trajectories.remote(batch) + except Exception: + self._warner.warn("dashboard trajectory sink flush failed; dropping events") + + +_phase_sink: PhaseSink | None = None +_engines_fingerprint: tuple | None = None +_warner = RateLimitedWarner(logger) + + +def attach_phase_sink(handle, role: str) -> None: + global _phase_sink + if _phase_sink is not None: + return # one sink per process + _phase_sink = PhaseSink(handle, role) + Timer().event_sinks.append(_phase_sink) + + +def attach_trajectory_sink(handle) -> None: + seam = TrajectoryLifecycle() + if seam.sink is not None: + return # one sink per process + seam.sink = TrajectorySink(handle) + + +def detach_and_flush() -> None: + global _phase_sink, _engines_fingerprint + if _phase_sink is not None: + if _phase_sink in Timer().event_sinks: + Timer().event_sinks.remove(_phase_sink) + _phase_sink.flush() + _phase_sink = None + _engines_fingerprint = None + seam = TrajectoryLifecycle() + if seam.sink is not None: + seam.sink.flush() + seam.sink = None + + +def register_train_actor(args) -> None: + """Called by TrainActor.init on EVERY rank — init_tracking only runs on + the megatron main rank, but per-GPU phase lanes need every rank's Timer.""" + if not args.use_miles_dashboard: + return + from miles.dashboard import backend + + handle = backend.resolve_collector() + if handle is None: + return + attach_phase_sink(handle, Role.TRAIN) + + +def register_router(args) -> None: + """Called by the rollout manager AFTER start_rollout_servers: only then are + ``args.sglang_router_ip/port`` filled in. init_tracking runs earlier in + __init__, so the backend cannot register the router at init time.""" + from miles.dashboard import backend + + handle = backend.current_collector() + if handle is None: + return + # a None ip here is a wiring-order bug, not runtime flakiness: fail loud + assert args.sglang_router_ip is not None, "register_router must run after start_rollout_servers" + try: + handle.set_router.remote( + f"http://{args.sglang_router_ip}:{args.sglang_router_port}", + use_miles_router=args.use_miles_router, + ) + except Exception: + _warner.warn("dashboard router registration failed; engine metrics will be missing") + + +def register_engines(servers) -> None: + """Called at the top of every RolloutManager.generate(): pushes an engine + topology snapshot whenever the set of engine actors changed (startup, + fault-tolerance recovery). Steady state costs one local tuple compare.""" + global _engines_fingerprint + from miles.dashboard import backend + + handle = backend.current_collector() + if handle is None: + return + try: + chunks = _alive_engine_chunks(servers) + fingerprint = tuple(id(engine.actor_handle) for chunk in chunks for engine in chunk) + if fingerprint == _engines_fingerprint: + return + infos = _ray_get([engine.actor_handle.get_topology_info.remote() for chunk in chunks for engine in chunk]) + handle.update_topology.remote(TopologySnapshot(ts=time.time(), engines=_group_engines(chunks, infos))) + _engines_fingerprint = fingerprint + except Exception: + _warner.warn("dashboard engine registration failed; topology may be stale") + + +def _alive_engine_chunks(servers) -> list[list]: + """Multi-node engines occupy ``nodes_per_engine`` consecutive entries of + ``group.all_engines``; only the first (master) owns the router-visible + URL. Chunks with any dead member are skipped until recovery completes.""" + chunks = [] + for server in servers.values(): + for group in server.server_groups: + stride = group.nodes_per_engine + engines = group.all_engines + for i in range(0, len(engines), stride): + chunk = engines[i : i + stride] + if all(engine.is_allocated and engine.is_alive for engine in chunk): + chunks.append(chunk) + return chunks + + +def _group_engines(chunks: list[list], infos: list[dict]) -> list[EngineInfo]: + engines = [] + index = 0 + for engine_rank, chunk in enumerate(chunks): + chunk_infos = infos[index : index + len(chunk)] + index += len(chunk) + master = chunk_infos[0] + engines.append( + EngineInfo( + addr=master["url"], + worker_type=master["worker_type"], + engine_rank=engine_rank, + gpus=[[info["node_ip"], gpu] for info in chunk_infos for gpu in info["gpu_ids"]], + gpu_uuids=[uuid for info in chunk_infos for uuid in info["gpu_uuids"]], + ) + ) + return engines diff --git a/miles/dashboard/logging_utils.py b/miles/dashboard/logging_utils.py new file mode 100644 index 0000000000..59dd346874 --- /dev/null +++ b/miles/dashboard/logging_utils.py @@ -0,0 +1,30 @@ +"""Logging helpers shared by the dashboard's background collectors.""" + +from __future__ import annotations + +import logging +import time + + +class RateLimitedWarner: + """Warn at most once per ``interval_seconds``. + + Background loops (scraper ticks, NVML sampling) hit the same failure + thousands of times during an outage; observability must neither spam the + training logs nor die silently, so the first warning goes through and + repeats are suppressed for the window. + """ + + def __init__(self, logger: logging.Logger, *, interval_seconds: float = 300.0): + self._logger = logger + self.interval_seconds = interval_seconds + self._last_warn = float("-inf") + + def warn(self, message: str) -> None: + now = time.monotonic() + if now - self._last_warn >= self.interval_seconds: + self._last_warn = now + self._logger.warning("%s (further warnings suppressed for %.0fs)", message, self.interval_seconds) + + def reset_window_for_test(self) -> None: + self._last_warn = float("-inf") diff --git a/miles/dashboard/serve.py b/miles/dashboard/serve.py new file mode 100644 index 0000000000..a0468fe799 --- /dev/null +++ b/miles/dashboard/serve.py @@ -0,0 +1,82 @@ +"""Standalone dashboard server over a ``--dump-details`` directory. + + python -m miles.dashboard.serve --dump-details /path/to/dump_details \\ + [--host 0.0.0.0] [--port 7788] [--follow] [--tensor-lru 2] [--cache-dir DIR] + +Works on any machine that can see the directory (login node, laptop over +NFS, the training node itself). ``--follow`` tails the JSONL telemetry +streams every few seconds for quasi-live viewing of a running job; dump +files are re-discovered per request either way. Typical remote usage is +behind an SSH port-forward. +""" + +from __future__ import annotations + +import argparse +import threading +import time +from pathlib import Path + +import uvicorn + +from miles.dashboard.dump_reader import DumpReader +from miles.dashboard.server import make_app +from miles.dashboard.store import MetricStore + +FOLLOW_INTERVAL_SECONDS = 2.0 + + +def make_demo_dir(target: Path) -> Path: + """Populate ``target`` with generated demo data (dumps + telemetry). + + Uses the dummy generators from the test suite, so it requires a miles + repo checkout (they are deliberately not shipped in the wheel).""" + try: + from tests.fast.dashboard.dummy_dump import dump_dummy_run + from tests.fast.dashboard.dummy_telemetry import dump_dummy_telemetry + except ImportError as e: + raise SystemExit("--demo needs the dummy generators under tests/; run from a miles repo checkout") from e + dump_dummy_run(target, steps=3, num_prompts=8, n_samples_per_prompt=4, max_response_len=48) + dump_dummy_telemetry(target, steps=3, samples_per_step=32) + return target + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--dump-details", default=None, help="the run's --dump-details directory") + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=7788) + parser.add_argument("--follow", action="store_true", help="tail telemetry streams of a still-running job") + parser.add_argument("--tensor-lru", type=int, default=2, help="rollout steps kept resident in tensor memory") + parser.add_argument("--cache-dir", default=None, help="summary cache dir (default: /dashboard/cache)") + parser.add_argument("--demo", action="store_true", help="serve generated demo data (needs a repo checkout)") + args = parser.parse_args(argv) + + if args.demo: + import tempfile + + dump_dir = make_demo_dir(Path(tempfile.mkdtemp(prefix="miles_dashboard_demo_"))) + else: + assert args.dump_details is not None, "--dump-details is required (or use --demo)" + dump_dir = Path(args.dump_details) + assert dump_dir.is_dir(), f"--dump-details directory not found: {dump_dir}" + + store = MetricStore.load(dump_dir / "dashboard") + reader = DumpReader(dump_dir, cache_dir=args.cache_dir, tensor_lru=args.tensor_lru) + app = make_app(store, reader, follow=args.follow) + + if args.follow: + # Append-only streams + GIL-atomic list appends make concurrent reads + # from request handlers safe; a reader may just miss the newest records. + def _tail() -> None: + while True: + time.sleep(FOLLOW_INTERVAL_SECONDS) + store.follow() + + threading.Thread(target=_tail, daemon=True, name="dashboard-follow").start() + + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/miles/dashboard/server.py b/miles/dashboard/server.py new file mode 100644 index 0000000000..85947d1dea --- /dev/null +++ b/miles/dashboard/server.py @@ -0,0 +1,256 @@ +"""HTTP API of the miles dashboard. + +``make_app(store, reader)`` wires the two read-side data sources — the +:class:`MetricStore` (JSONL telemetry under ``{dump_details}/dashboard/``) +and the :class:`DumpReader` (rollout/train ``.pt`` dumps) — into the REST +API consumed by the SPA. The server is strictly read-only over files on +disk; live viewing is the same app with a follow loop tailing the store +(see ``serve.py``). + +Metric catalog: keys from ``metrics.jsonl`` are served as-is; dump-derived +per-step aggregates are namespaced as ``dump/`` so the L0 view works +even for runs where the dashboard backend was not enabled during training. + +Error mapping: ``DumpStillWriting`` -> 503 (client retries), +``FileNotFoundError``/``KeyError`` -> 404, ``ValueError`` -> 400. +""" + +from __future__ import annotations + +import json +import math +from contextlib import contextmanager +from pathlib import Path + +from fastapi import FastAPI, HTTPException, Query +from fastapi.responses import FileResponse, Response +from fastapi.staticfiles import StaticFiles + +from miles.dashboard.dump_reader import STEP_AGGREGATE_METRICS, DumpReader, DumpStillWriting +from miles.dashboard.store import MetricStore, Stream + +DUMP_METRIC_PREFIX = "dump/" + +_STATIC_DIR = Path(__file__).parent / "static" + + +@contextmanager +def _translate_errors(): + try: + yield + except DumpStillWriting as e: + raise HTTPException(status_code=503, detail=f"dump still being written: {e}") from e + except (FileNotFoundError, KeyError) as e: + raise HTTPException(status_code=404, detail=str(e)) from e + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) from e + + +def _json_safe(value): + """NaN/inf are not valid JSON; per-token GPU outputs can contain them.""" + if isinstance(value, float) and not math.isfinite(value): + return None + if isinstance(value, list): + return [_json_safe(v) for v in value] + if isinstance(value, dict): + return {k: _json_safe(v) for k, v in value.items()} + return value + + +def _table(df) -> dict: + return dict(columns=df.columns, rows=[_json_safe(row) for row in df.to_dicts()]) + + +def make_app(store: MetricStore, reader: DumpReader, *, follow: bool = False) -> FastAPI: + app = FastAPI(title="miles dashboard", docs_url=None, redoc_url=None) + + @app.get("/api/meta") + def meta(): + ids = reader.rollout_ids() + # dump-derived aggregates are the L0 fallback for dump-only dirs; a + # run with a telemetry stream never advertises them (they torch.load + # raw sample dumps — the /api/metrics handler still serves explicit + # requests, but the catalog stays wandb-shaped) + offline = ids.train and not store.records[Stream.METRICS] + dump_keys = [DUMP_METRIC_PREFIX + column for column in STEP_AGGREGATE_METRICS] if offline else [] + return dict( + mode="follow" if follow else "static", + run_name=store.meta.run_name if store.meta else None, + start_ts=store.meta.start_ts if store.meta else None, + time_range=store.time_range(), + rollout_ids=dict(train=ids.train, eval=ids.eval), + metric_keys=store.metric_keys() + dump_keys, + engine_metric_keys=store.engine_metric_names(), + step_keys=store.step_keys(), + capabilities=dict( + has_metrics=store.has_stream(Stream.METRICS), + has_tokenizer=(reader.dump_dir / "tokenizer").is_dir(), + has_timeline=store.has_stream(Stream.PHASES) or store.has_stream(Stream.GPU_UTIL), + has_engine_series=store.has_stream(Stream.ENGINE_SERIES), + max_window_s=MetricStore.MAX_WINDOW_S, + ), + ) + + def _check_window(t0: float | None, t1: float | None) -> None: + if t0 is None or t1 is None: + return + if t1 <= t0: + raise ValueError(f"bad window: {t1=} <= {t0=}") + if t1 - t0 > MetricStore.MAX_WINDOW_S: + raise ValueError(f"window {t1 - t0:.0f}s exceeds max_window_s {MetricStore.MAX_WINDOW_S:.0f}") + + # ------------------------------ timeline -------------------------------- + + @app.get("/api/rollout/{rollout_id}/trajectories") + def rollout_trajectories(rollout_id: int, sample_index: int | None = None): + """Batch anatomy: the consuming step's samples resolved to their + lifecycle lanes. The event scan is capped at one viewport (4 h) before + the consume anchor (design §18.5); empty lanes = run predates the + trajectory probes. ``sample_index`` narrows to one sample (the L2 + page's own lane) without touching the step summary.""" + with _translate_errors(): + indices = ( + {sample_index} + if sample_index is not None + else set(reader.summary(rollout_id)["sample_index"].to_list()) + ) + consume = next((b["ts"] for b in store.bubbles() if b["step"] == rollout_id), None) + if consume is None: + window = store.time_range() + consume = window[1] if window else None + t0 = consume - MetricStore.MAX_WINDOW_S if consume is not None else None + lanes = store.trajectory_lanes(t0=t0, t1=consume, sample_indices=indices) + return _json_safe(dict(lanes=lanes, consume_ts=consume, t0=t0)) + + @app.get("/api/timeline/topology") + def timeline_topology(): + return dict(lanes=store.lanes(), windows=store.topology_windows()) + + @app.get("/api/timeline/phases") + def timeline_phases(t0: float | None = None, t1: float | None = None, lanes: str | None = None): + with _translate_errors(): + _check_window(t0, t1) + return dict(phases=store.phases_by_lane(t0=t0, t1=t1, lanes=store.resolve_lanes(lanes))) + + @app.get("/api/timeline/gpu") + def timeline_gpu( + t0: float | None = None, t1: float | None = None, max_points: int = 2000, lanes: str | None = None + ): + with _translate_errors(): + _check_window(t0, t1) + if max_points < 2: + raise ValueError(f"{max_points=} must be >= 2") + return dict(lanes=store.gpu_series(t0=t0, t1=t1, max_points=max_points, lanes=store.resolve_lanes(lanes))) + + @app.get("/api/timeline/heatmap") + def timeline_heatmap( + metric: str = "util", + t0: float | None = None, + t1: float | None = None, + x_buckets: int = 1200, + lanes: str | None = None, + ): + """Binary rank carpet: [4-byte LE header length][header JSON][uint8 + matrix, row-major] — one byte per (lane, time bucket) cell.""" + with _translate_errors(): + _check_window(t0, t1) + if not 2 <= x_buckets <= 4000: + raise ValueError(f"{x_buckets=} out of range [2, 4000]") + result = store.heatmap(metric, t0=t0, t1=t1, x_buckets=x_buckets, lanes=store.resolve_lanes(lanes)) + values = result.pop("values") + header = json.dumps(_json_safe(result)).encode() + return Response( + content=len(header).to_bytes(4, "little") + header + values, + media_type="application/octet-stream", + ) + + @app.get("/api/timeline/outliers") + def timeline_outliers(criterion: str, t0: float | None = None, t1: float | None = None, top_k: int = 16): + with _translate_errors(): + if not 1 <= top_k <= 256: + raise ValueError(f"{top_k=} out of range [1, 256]") + return dict(outliers=store.outliers(criterion, t0=t0, t1=t1, top_k=top_k)) + + @app.get("/api/timeline/engine_series") + def timeline_engine_series(metric: str, t0: float | None = None, t1: float | None = None, max_points: int = 2000): + with _translate_errors(): + _check_window(t0, t1) + if max_points < 2: + raise ValueError(f"{max_points=} must be >= 2") + return dict(series=store.engine_series(metric, t0=t0, t1=t1, max_points=max_points)) + + @app.get("/api/timeline/bubbles") + def timeline_bubbles(): + return dict(bubbles=store.bubbles()) + + @app.get("/api/metrics") + def metrics(keys: str, x: str = "rollout/step", t0: float | None = None, t1: float | None = None): + with _translate_errors(): + requested = [k for k in keys.split(",") if k] + if not requested: + raise ValueError("keys must be a non-empty comma-separated list") + store_keys = [k for k in requested if not k.startswith(DUMP_METRIC_PREFIX)] + series = store.metric_series(store_keys, x_key=x, t0=t0, t1=t1) + dump_keys = [k for k in requested if k.startswith(DUMP_METRIC_PREFIX)] + if dump_keys: + aggregates = reader.step_aggregates() + steps = aggregates["rollout_id"].to_list() if aggregates.height else [] + for key in dump_keys: + column = key.removeprefix(DUMP_METRIC_PREFIX) + if column not in STEP_AGGREGATE_METRICS: + raise ValueError(f"unknown dump metric {key!r}") + points = [ + (step, value) + for step, value in zip(steps, aggregates[column].to_list(), strict=True) + if value is not None + ] + series[key] = dict(x=[p[0] for p in points], y=[_json_safe(p[1]) for p in points], ts=[]) + return _json_safe(series) + + @app.get("/api/rollout/{rollout_id}/summary") + def rollout_summary(rollout_id: int, evaluation: bool = Query(False, alias="eval")): + with _translate_errors(): + df = reader.summary(rollout_id, evaluation=evaluation) + return dict(rollout_id=rollout_id, evaluation=evaluation, **_table(df)) + + @app.get("/api/rollout/{rollout_id}/groups") + def rollout_groups(rollout_id: int, evaluation: bool = Query(False, alias="eval")): + with _translate_errors(): + df = reader.groups(rollout_id, evaluation=evaluation) + return dict(rollout_id=rollout_id, evaluation=evaluation, **_table(df)) + + @app.get("/api/rollout/{rollout_id}/sample/{sample_index}/messages") + def sample_messages(rollout_id: int, sample_index: int, evaluation: bool = Query(False, alias="eval")): + with _translate_errors(): + return reader.trajectory_messages(rollout_id, sample_index, evaluation=evaluation) + + @app.get("/api/rollout/{rollout_id}/sample/{sample_index}/tokens") + def sample_tokens( + rollout_id: int, + sample_index: int, + start: int = 0, + end: int | None = None, + evaluation: bool = Query(False, alias="eval"), + ): + with _translate_errors(): + payload = reader.tokens(rollout_id, sample_index, start=start, end=end, evaluation=evaluation) + return _json_safe(payload) + + if _STATIC_DIR.is_dir(): + app.mount("/static", StaticFiles(directory=_STATIC_DIR), name="static") + + @app.middleware("http") + async def _no_stale_frontend(request, call_next): + # without Cache-Control browsers cache heuristically and keep + # serving a STALE SPA after a stack upgrade; no-cache forces a + # revalidation (cheap 304s) so the frontend always matches serve + response = await call_next(request) + if request.url.path == "/" or request.url.path.startswith("/static"): + response.headers["Cache-Control"] = "no-cache" + return response + + @app.get("/", include_in_schema=False) + def index(): + return FileResponse(_STATIC_DIR / "index.html") + + return app diff --git a/miles/dashboard/sglang_scraper.py b/miles/dashboard/sglang_scraper.py new file mode 100644 index 0000000000..7b8ad78dc7 --- /dev/null +++ b/miles/dashboard/sglang_scraper.py @@ -0,0 +1,173 @@ +"""Scraper turning sglang engine metrics into dashboard ``EngineSample`` records. + +Two fetch modes (design doc §6.6): + +- ``router``: one ``GET {router_addr}/engine_metrics`` per tick. The sgl + router fans the request out to every registered worker and labels each + sample with ``worker_addr`` (sgl-model-gateway ``get_engine_metrics``). +- ``direct``: concurrent ``GET {addr}/metrics`` against every engine address + from the topology. Needed under ``--use-miles-router``, whose router has no + aggregation endpoint; here the scraper attaches the address itself. + +Metric names are normalized (``sglang:num_running_reqs`` from a raw engine +equals ``sglang_num_running_reqs`` from the router aggregation), filtered by +a whitelist, and reduced to the label subset the dashboard uses. Records are +handed to a sink callable (the collector passes ``store.append``), so this +module has no collector or Ray dependency and is fully unit-testable. + +Failure policy: a failed tick (or a failed engine in direct mode) leaves a +gap in the series — values are never fabricated — and warnings are +rate-limited to one per ``WARN_INTERVAL_SECONDS`` (pattern from +``examples/random_async/random_async_sglang_metrics.py``). +""" + +from __future__ import annotations + +import logging +import threading +import time +from collections.abc import Callable, Iterable +from concurrent.futures import ThreadPoolExecutor + +try: + from enum import StrEnum +except ImportError: + from backports.strenum import StrEnum + +from miles.dashboard.logging_utils import RateLimitedWarner +from miles.dashboard.store import EngineSample + +logger = logging.getLogger(__name__) + +DEFAULT_METRIC_WHITELIST = ( + "sglang_num_running_reqs", + "sglang_num_queue_reqs", + "sglang_gen_throughput", + "sglang_token_usage", + "sglang_cache_hit_rate", + # PD disaggregation; these families simply don't exist when PD is off + "sglang_num_prefill_prealloc_queue_reqs", + "sglang_num_prefill_inflight_queue_reqs", + "sglang_num_decode_prealloc_queue_reqs", + "sglang_num_decode_transfer_queue_reqs", + "sglang_kv_transfer_speed_gb_s", + "sglang_kv_transfer_latency_ms", +) + +# label subset passed through to EngineSample.labels (e.g. PD prefill/decode) +KEPT_LABELS = ("engine_type",) + + +class ScrapeMode(StrEnum): + ROUTER = "router" + DIRECT = "direct" + + +def _default_http_get(url: str, timeout: float) -> str: + import requests + + response = requests.get(url, timeout=timeout) + response.raise_for_status() + return response.text + + +class SglangScraper: + + def __init__( + self, + sink: Callable[[EngineSample], None], + *, + mode: ScrapeMode, + router_addr: str | None = None, + engine_addrs: Callable[[], Iterable[str]] | None = None, + interval: float = 2.0, + timeout: float = 5.0, + whitelist: Iterable[str] = DEFAULT_METRIC_WHITELIST, + http_get: Callable[[str, float], str] = _default_http_get, + ): + if mode == ScrapeMode.ROUTER: + assert router_addr is not None, "router mode needs router_addr" + else: + assert engine_addrs is not None, "direct mode needs an engine_addrs provider" + self._sink = sink + self.mode = mode + self.router_addr = router_addr + self.engine_addrs = engine_addrs + self.interval = interval + self.timeout = timeout + self.whitelist = frozenset(whitelist) + self._http_get = http_get + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + self._warner = RateLimitedWarner(logger) + + # ------------------------------ lifecycle ------------------------------- + + def start(self) -> None: + assert self._thread is None, "scraper already started" + self._thread = threading.Thread(target=self._run, name="dashboard-sglang-scraper", daemon=True) + self._thread.start() + + def stop(self) -> None: + self._stop_event.set() + if self._thread is not None: + self._thread.join(timeout=self.timeout + 1) + + def _run(self) -> None: + while not self._stop_event.is_set(): + started = time.time() + try: + self.scrape_once(started) + except Exception: + self._warn("sglang metrics scrape failed; leaving a gap") + self._stop_event.wait(max(0.5, self.interval - (time.time() - started))) + + # ------------------------------- one tick ------------------------------- + + def scrape_once(self, now: float) -> int: + """Fetch, parse and sink one round of samples. Returns the count.""" + count = 0 + for addr, text in self._fetch(): + count += self._ingest(addr, text, now) + return count + + def _fetch(self) -> list[tuple[str | None, str]]: + if self.mode == ScrapeMode.ROUTER: + # addr None: the router aggregation carries worker_addr labels + return [(None, self._http_get(f"{self.router_addr}/engine_metrics", self.timeout))] + + addrs = list(self.engine_addrs()) + + def fetch_one(addr: str) -> tuple[str | None, str] | None: + try: + return addr, self._http_get(f"{addr}/metrics", self.timeout) + except Exception: + self._warn(f"scrape of engine {addr} failed; skipping this tick") + return None + + with ThreadPoolExecutor(max_workers=min(8, max(1, len(addrs)))) as pool: + return [result for result in pool.map(fetch_one, addrs) if result is not None] + + def _ingest(self, addr: str | None, text: str, now: float) -> int: + from prometheus_client.parser import text_string_to_metric_families + + count = 0 + for family in text_string_to_metric_families(text): + # raw engines export "sglang:x", the router aggregation "sglang_x" + name = family.name.replace(":", "_") + if name not in self.whitelist: + continue + for sample in family.samples: + sample_addr = addr if addr is not None else sample.labels.get("worker_addr") + if sample_addr is None: + self._warn(f"router sample for {name} lacks worker_addr; dropping") + continue + labels = {k: v for k, v in sample.labels.items() if k in KEPT_LABELS} + self._sink( + EngineSample(ts=now, addr=sample_addr, metric=name, labels=labels, value=float(sample.value)) + ) + count += 1 + return count + + def _warn(self, message: str) -> None: + self._warner.warn(message) diff --git a/miles/dashboard/static/anatomy.js b/miles/dashboard/static/anatomy.js new file mode 100644 index 0000000000..554fbddfed --- /dev/null +++ b/miles/dashboard/static/anatomy.js @@ -0,0 +1,242 @@ +import { el, fmtNum } from "./app.js"; +import { hideTooltip, showTooltip } from "./charts.js"; + +// batch anatomy swimlanes (design §18.4, visual contract: batch_anatomy_demo) +const COLORS = { + gen: "#1d9e75", + tool: "#d97706", + attempt: "#e7eaef", // queue/underlay: idle recedes, activity pops + consume: "#2f6feb", + text: "#24292f", + muted: "#667080", + stale: "#d85a30", +}; +const M_TOP = 18; + +// sequential cyclic ring for weight versions: v0 is the familiar gen green +// and each +1 steps 25° around a 320° ring that skips the tool-wait amber +// band — neighbouring versions read as neighbouring colors (dark green → +// light green → cyan → …) and the ring wraps back to the start every ~13 +const versionColor = (v) => `hsl(${55 + ((105 + v * 25) % 320)}, 69%, 37%)`; + +// spans carry the version stamped by the sink; older dumps may lack it on a +// per-span basis, so fall back to the lane's (usually single) version +const segmentVersion = (segment, lane) => { + const v = Number(segment.weight_version); + return Number.isFinite(v) && segment.weight_version !== "" ? v : (lane.versions[0] ?? 0); +}; +// two densities, same philosophy as the rank carpet (§15): full detail only +// while every row can carry text; above that the pane compresses to a carpet +// and the sort chips + tooltip + table below carry identification +const DETAIL_MAX = 48; +const MAX_PANE_PX = 640; + +// createAnatomy renders the per-step trajectory swimlanes: one row per sample, +// x = wall clock from the earliest lifecycle event to the consume anchor. +// rowsByIndex supplies the dump-side columns (versions/turns/tools/reward). +export function createAnatomy({ lanes, consumeTs, rowsByIndex, onClickSample }) { + const detailed = lanes.length <= DETAIL_MAX; + const ROW = detailed ? 13 : lanes.length <= 512 ? 4 : 2; + const M_LEFT = detailed ? 96 : 64; + const M_RIGHT = detailed ? 64 : 14; + + const dumpRow = (lane) => rowsByIndex.get(lane.sample_index) ?? {}; + const SORTS = { + submit: (l) => l.first_ts, + staleness: (l) => -(l.versions.length > 1 ? l.versions.at(-1) - l.versions[0] : 0), + "wall span": (l) => -(l.last_ts - l.first_ts), + reward: (l) => dumpRow(l).raw_reward ?? dumpRow(l).reward ?? 0, + turns: (l) => -(dumpRow(l).turns ?? 0), + }; + let sortKey = "submit"; + let order = [...lanes]; + const resort = () => { + order = [...lanes].sort((a, b) => SORTS[sortKey](a) - SORTS[sortKey](b) || a.first_ts - b.first_ts); + }; + + const canvas = el("canvas", { class: "anatomy" }); + const wrap = el("div", { class: "anatomywrap" }, [canvas]); + const sortRow = el("div", { class: "controls" }); + const renderSort = () => { + sortRow.replaceChildren( + el("span", { class: "muted" }, [`${lanes.length} trajectories · wheel = zoom · drag = pan · sort`]), + ...Object.keys(SORTS).map((key) => + el( + "button", + { + class: key === sortKey ? "active" : "", + onclick: () => { + sortKey = key; + renderSort(); + resort(); + draw(); + }, + }, + [key], + ), + ), + ); + }; + const single = lanes.length === 1; // the L2 page's own lane: no sorting to do + const panel = el("div", { class: "panel" }, [ + el("h3", {}, [ + single + ? `sample s${lanes[0].sample_index} lifecycle — generated, waited, tool-called` + : "batch anatomy — when each sample generated, waited, tool-called", + ]), + ...(single ? [] : [sortRow]), + wrap, + el("div", { class: "legend" }, [ + legendSwatch(COLORS.gen, "generating (hue steps with weight version)"), + legendSwatch(COLORS.tool, "tool wait"), + legendSwatch(COLORS.attempt, "queue/attempt"), + el("span", { style: `color: ${COLORS.consume}` }, ["│ consume"]), + ]), + ]); + + const T0 = Math.min(...lanes.map((l) => l.first_ts)); + const T1 = consumeTs ?? Math.max(...lanes.map((l) => l.last_ts)); + // zoomable time window, same idiom as the timeline lane view: wheel = + // zoom at cursor, drag = pan, double-click = reset. Purely visual — all + // segments are already client-side, no refetch. + let v0 = T0; + let v1 = T1; + const X = (t, w) => M_LEFT + ((t - v0) / Math.max(v1 - v0, 1e-9)) * (w - M_LEFT - M_RIGHT); + const tAt = (clientX, rect) => v0 + ((clientX - rect.left - M_LEFT) / (rect.width - M_LEFT - M_RIGHT)) * (v1 - v0); + + function draw() { + const height = M_TOP + order.length * ROW + 6; + canvas.style.height = `${height}px`; + wrap.style.maxHeight = `${Math.min(height, MAX_PANE_PX)}px`; + const dpr = window.devicePixelRatio || 1; + const rect = canvas.getBoundingClientRect(); + canvas.width = rect.width * dpr; + canvas.height = rect.height * dpr; + const ctx = canvas.getContext("2d"); + ctx.scale(dpr, dpr); + const W = rect.width; + ctx.font = "10.5px ui-monospace, monospace"; + + ctx.fillStyle = COLORS.muted; + const span = v1 - v0; + const tick = [1, 5, 10, 30, 60, 120, 300, 600, 1800, 3600].find((s) => span / s <= 8) || 7200; + for (let t = Math.ceil((v0 - T0) / tick) * tick; t <= v1 - T0; t += tick) { + ctx.fillText(`+${Math.floor(t / 60)}:${String(Math.round(t) % 60).padStart(2, "0")}`, X(T0 + t, W) - 10, 10); + } + + const labelEvery = detailed ? 1 : Math.ceil(order.length / 40); + order.forEach((lane, i) => { + const y = M_TOP + i * ROW; + const row = dumpRow(lane); + if (i % labelEvery === 0) { + ctx.fillStyle = COLORS.text; + ctx.fillText(`s${lane.sample_index}`, 8, y + Math.min(ROW, 10)); + } + + const padAttempt = detailed ? 4 : ROW > 2 ? 1 : 0; + const padSegment = detailed ? 2 : 0; + for (const attempt of lane.attempts) { + if ((attempt.t1 ?? T1) < v0 || (attempt.t0 ?? T0) > v1) continue; + ctx.fillStyle = COLORS.attempt; + const x0 = X(attempt.t0 ?? T0, W); + ctx.fillRect(x0, y + padAttempt, Math.max(X(attempt.t1 ?? T1, W) - x0, 1), ROW - 2 * padAttempt); + } + for (const segment of lane.segments) { + if ((segment.t1 ?? T1) < v0 || (segment.t0 ?? T0) > v1) continue; + ctx.fillStyle = segment.kind === "gen" ? versionColor(segmentVersion(segment, lane)) : COLORS.tool; + const x0 = X(segment.t0 ?? T0, W); + ctx.fillRect(x0, y + padSegment, Math.max(X(segment.t1 ?? T1, W) - x0, 1.5), ROW - 2 * padSegment); + } + if (detailed) { + const info = [ + row.turns !== undefined && row.turns !== null ? `${row.turns}t` : null, + row.tool_calls ? `${row.tool_calls}x` : null, + row.raw_reward !== undefined && row.raw_reward !== null ? `r=${fmtNum(row.raw_reward)}` : null, + ] + .filter(Boolean) + .join("·"); + ctx.fillStyle = COLORS.muted; + ctx.fillText(info, W - M_RIGHT + 6, y + 10); + } + }); + + if (consumeTs !== null && consumeTs !== undefined && consumeTs >= v0 && consumeTs <= v1) { + const x = X(consumeTs, W); + ctx.strokeStyle = COLORS.consume; + ctx.lineWidth = 1.5; + ctx.beginPath(); + ctx.moveTo(x, M_TOP - 6); + ctx.lineTo(x, M_TOP + order.length * ROW); + ctx.stroke(); + ctx.lineWidth = 1; + } + } + + canvas.onwheel = (ev) => { + ev.preventDefault(); + const rect = canvas.getBoundingClientRect(); + const pivot = tAt(ev.clientX, rect); + const factor = Math.exp(ev.deltaY * 0.002); + v0 = Math.max(T0, pivot + (v0 - pivot) * factor); + v1 = Math.min(T1, pivot + (v1 - pivot) * factor); + if (v1 - v0 < 1) v1 = v0 + 1; + draw(); + }; + canvas.ondblclick = () => ((v0 = T0), (v1 = T1), draw()); + let dragFrom = null; // {x, v0, v1}; a <4px move on release is a click + canvas.onmousedown = (ev) => (dragFrom = { x: ev.clientX, v0, v1 }); + canvas.onmouseup = (ev) => { + const moved = dragFrom && Math.abs(ev.clientX - dragFrom.x) >= 4; + dragFrom = null; + if (moved) return; + const rect = canvas.getBoundingClientRect(); + const i = Math.floor((ev.clientY - rect.top - M_TOP) / ROW); + if (i >= 0 && i < order.length) onClickSample(order[i].sample_index); + }; + + canvas.onmousemove = (ev) => { + const rect = canvas.getBoundingClientRect(); + if (dragFrom) { + const dt = ((dragFrom.x - ev.clientX) / (rect.width - M_LEFT - M_RIGHT)) * (dragFrom.v1 - dragFrom.v0); + const span = dragFrom.v1 - dragFrom.v0; + v0 = Math.max(T0, Math.min(dragFrom.v0 + dt, T1 - span)); + v1 = v0 + span; + draw(); + return; + } + const i = Math.floor((ev.clientY - rect.top - M_TOP) / ROW); + const x = ev.clientX - rect.left; + if (i < 0 || i >= order.length || x < M_LEFT || x > rect.width - M_RIGHT) { + hideTooltip(); + return; + } + const lane = order[i]; + const t = tAt(ev.clientX, rect); + const segment = lane.segments.find((s) => (s.t0 ?? T0) <= t && t < (s.t1 ?? T1)); + const versions = lane.versions.length > 1 ? ` v${lane.versions[0]}–v${lane.versions.at(-1)}` : ""; + const lines = [`s${lane.sample_index} +${fmtNum(t - T0)}s ${lane.status}${versions}`]; + if (segment) { + const version = segment.weight_version ? ` · v${segment.weight_version}` : ""; + const turn = segment.turn > 0 ? `turn ${segment.turn} ` : ""; + lines.push(`${turn}${segment.kind === "gen" ? "generating" : "tool wait"}${version}`); + } else if (lane.attempts.some((a) => (a.t0 ?? T0) <= t && t < (a.t1 ?? T1))) { + lines.push("queued / waiting"); + } + showTooltip(ev.clientX, ev.clientY, lines.join("\n")); + }; + canvas.onmouseleave = () => { + hideTooltip(); + dragFrom = null; // leaving the canvas cancels a pan + }; + + renderSort(); + resort(); + queueMicrotask(draw); + window.addEventListener("resize", draw); + return panel; +} + +function legendSwatch(color, label) { + const swatch = el("span", { class: "bar", style: `width: 12px; background: ${color}` }); + return el("span", { style: "display: inline-flex; gap: 4px; align-items: center" }, [swatch, label]); +} diff --git a/miles/dashboard/static/api.js b/miles/dashboard/static/api.js new file mode 100644 index 0000000000..c1e278031c --- /dev/null +++ b/miles/dashboard/static/api.js @@ -0,0 +1,45 @@ +const RETRIES_503 = 3; + +async function fetchOk(path, params) { + const url = new URL(path, location.origin); + for (const [k, v] of Object.entries(params)) { + if (v !== undefined && v !== null) url.searchParams.set(k, v); + } + for (let attempt = 0; ; attempt++) { + const res = await fetch(url); + if (res.status === 503 && attempt < RETRIES_503) { + await new Promise((r) => setTimeout(r, 1500)); + continue; + } + if (!res.ok) { + let detail = ""; + try { + detail = (await res.json()).detail ?? ""; + } catch { + /* non-json error body */ + } + throw new Error(`HTTP ${res.status}: ${detail || url.pathname}`); + } + return res; + } +} + +export async function api(path, params = {}) { + return (await fetchOk(path, params)).json(); +} + +// framed binary endpoints (/api/timeline/heatmap): +// [4-byte LE header length][header JSON][uint8 matrix, row-major] +export async function apiBinary(path, params = {}) { + const buf = await (await fetchOk(path, params)).arrayBuffer(); + const headerLen = new DataView(buf).getUint32(0, true); + const header = JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 4, headerLen))); + return { ...header, values: new Uint8Array(buf, 4 + headerLen) }; +} + +let metaCache = null; + +export async function getMeta() { + if (metaCache === null) metaCache = await api("/api/meta"); + return metaCache; +} diff --git a/miles/dashboard/static/app.js b/miles/dashboard/static/app.js new file mode 100644 index 0000000000..f545daca8d --- /dev/null +++ b/miles/dashboard/static/app.js @@ -0,0 +1,108 @@ +import { getMeta } from "./api.js"; +import { renderMetrics } from "./views_metrics.js"; +import { renderRollout } from "./views_rollout.js"; +import { renderTimeline } from "./views_timeline.js"; +import { renderTokens } from "./views_tokens.js"; + +// tiny DOM builder: el("div", {class: "x", onclick: fn}, [children|strings]) +export function el(tag, attrs = {}, children = []) { + const node = document.createElement(tag); + for (const [k, v] of Object.entries(attrs)) { + if (k.startsWith("on")) node[k] = v; + else if (v !== null && v !== undefined) node.setAttribute(k, v); + } + for (const child of children) { + node.append(child); + } + return node; +} + +// views with background work (follow-mode auto refresh) register a cleanup +// here; render() runs it before switching views so intervals never leak +let activeViewCleanup = null; + +export function setViewCleanup(cleanup) { + activeViewCleanup = cleanup; +} + +export const fmtNum = (v) => { + if (v === null || v === undefined) return "—"; + if (typeof v === "boolean") return v ? "✓" : ""; + if (typeof v !== "number") return String(v); + if (Number.isInteger(v)) return String(v); + const a = Math.abs(v); + if (a >= 1e5 || (a < 1e-3 && a > 0)) return v.toExponential(2); + return v.toFixed(3); +}; + +function parseRoute() { + const [path, query] = (location.hash.slice(1) || "/").split("?"); + const segments = path.split("/").filter(Boolean); + const params = new URLSearchParams(query || ""); + if (segments[0] === "timeline") { + return { view: "timeline", lanes: params.get("lanes") }; + } + if (segments[0] === "rollout" && segments.length >= 2) { + const rolloutId = Number(segments[1]); + const evaluation = params.get("eval") === "1"; + if (segments[2] === "sample" && segments.length === 4) { + return { view: "tokens", rolloutId, sampleIndex: Number(segments[3]), evaluation }; + } + return { view: "rollout", rolloutId, evaluation }; + } + return { view: "metrics" }; +} + +function crumbs(route, meta) { + const nav = (label, href, active) => el("a", { class: `nav${active ? " active" : ""}`, href }, [label]); + const parts = [nav("Metrics", "#/", route.view === "metrics")]; + if (meta.capabilities.has_timeline) { + parts.push(nav("Compute Utilization", "#/timeline", route.view === "timeline")); + } + // the per-step data view is a top-level destination, not a hidden + // click-through from chart points; land on the newest train step + const latest = meta.rollout_ids.train.at(-1); + if (latest !== undefined) { + parts.push(nav("Data / Trajectories", `#/rollout/${latest}`, route.view === "rollout" || route.view === "tokens")); + } + if (route.view === "rollout" || route.view === "tokens") { + const evalSuffix = route.evaluation ? "?eval=1" : ""; + parts.push( + el("span", { class: "crumb" }, [ + "› ", + el("a", { href: `#/rollout/${route.rolloutId}${evalSuffix}` }, [ + `${route.evaluation ? "eval " : ""}step ${route.rolloutId}`, + ]), + ]), + ); + } + if (route.view === "tokens") { + parts.push(el("span", { class: "crumb" }, [`› sample ${route.sampleIndex}`])); + } + document.getElementById("crumbs").replaceChildren(...parts); +} + +async function render() { + if (activeViewCleanup) { + activeViewCleanup(); + activeViewCleanup = null; + } + const route = parseRoute(); + const view = document.getElementById("view"); + view.replaceChildren(el("p", { class: "muted" }, ["loading…"])); + try { + const meta = await getMeta(); + crumbs(route, meta); + document.getElementById("runinfo").textContent = + `${meta.run_name ?? "unnamed run"} · ${meta.mode}` + (meta.capabilities.has_metrics ? "" : " · dump-derived metrics"); + if (route.view === "metrics") await renderMetrics(view, meta); + else if (route.view === "timeline") await renderTimeline(view, meta, route); + else if (route.view === "rollout") await renderRollout(view, meta, route); + else await renderTokens(view, meta, route); + } catch (err) { + view.replaceChildren(el("div", { class: "error" }, [String(err)])); + } +} + +window.addEventListener("hashchange", render); +render(); diff --git a/miles/dashboard/static/carpet.js b/miles/dashboard/static/carpet.js new file mode 100644 index 0000000000..ae92e902ed --- /dev/null +++ b/miles/dashboard/static/carpet.js @@ -0,0 +1,242 @@ +import { apiBinary } from "./api.js"; +import { el, fmtNum } from "./app.js"; +import { hideTooltip, showTooltip } from "./charts.js"; + +// margins match views_timeline so the carpet and the lane view share an x domain +const M_LEFT = 96; +const M_RIGHT = 14; +const X_BUCKETS = 1200; +const METRICS = ["util", "phase", "mem_mb", "power_w", "lifecycle"]; +const LIFECYCLE_COLORS = { queue: "#cfd6e0", generating: "#1d9e75", tool_wait: "#d97706" }; +const MAX_LABELED_ROWS = 48; + +// sequential ramp for magnitude metrics: one hue, light -> dark +const RAMP_LO = [238, 241, 246]; +const RAMP_HI = [22, 64, 143]; +const NO_DATA = [245, 246, 248]; // page bg: absence recedes +const MUTED = "#667080"; +const TEXT = "#24292f"; +const ACCENT = "#2f6feb"; + +// carpet-only wake_up shift: the lane view's #ba7517 is deutan-indistinguishable +// from update_weights (dE 5.9) without position/width context, so the carpet +// moves it toward yellow (dE 20.7 against update_weights, validator-checked) +const CARPET_WAKE_UP = "#d4a017"; + +const hexToRgb = (hex) => [1, 3, 5].map((i) => parseInt(hex.slice(i, i + 2), 16)); + +// createCarpet returns {root, refresh({t0?, t1?}), redraw, destroy}. The carpet +// is the overview over ALL lanes and the row-brush selector; opts: +// phaseColors lane-view palette (wake_up overridden here) +// runStart() -> run start ts, so tooltips share the lane view's +m:ss origin +// selectedKeys() -> Set("node:gpu") for the left-edge selection markers +// onBrush(keys) brushed lane keys (click or row drag); caller decides add/remove +export function createCarpet(opts) { + const phaseColors = { ...opts.phaseColors, wake_up: CARPET_WAKE_UP }; + let metric = METRICS[0]; + let data = null; // last heatmap response + let range = {}; + let brush = null; // {r0, r1} while dragging + + const canvas = el("canvas", { class: "carpet" }); + const scaleLabel = el("span", { class: "muted", style: "font-size: 12px" }); + const chipsRow = el("div", { class: "controls" }); + const root = el("div", { class: "panel" }, [chipsRow, canvas]); + + const renderChips = () => { + chipsRow.replaceChildren( + el("span", { class: "muted" }, ["carpet"]), + ...METRICS.map((m) => + el( + "button", + { + class: m === metric ? "active" : "", + onclick: () => { + metric = m; + renderChips(); + refresh(range); + }, + }, + [m === "lifecycle" ? "traj" : m], + ), + ), + scaleLabel, + el("span", { class: "muted" }, ["click gpu = add/remove · drag = brush"]), + ); + }; + + async function refresh(newRange = range) { + range = newRange; + data = await apiBinary("/api/timeline/heatmap", { + metric, + x_buckets: X_BUCKETS, + t0: range.t0, + t1: range.t1, + }); + redraw(); + } + + const rowH = () => { + const n = data.rows.length; + return n <= 16 ? 14 : n <= 64 ? 8 : n <= 256 ? 4 : 2; + }; + + function redraw() { + if (!data) return; + const rows = data.rows.length; + const h = Math.max(rows * (rows ? rowH() : 1), 24); + canvas.style.height = `${h}px`; + const dpr = window.devicePixelRatio || 1; + const rect = canvas.getBoundingClientRect(); + canvas.width = rect.width * dpr; + canvas.height = rect.height * dpr; + const ctx = canvas.getContext("2d"); + ctx.scale(dpr, dpr); + ctx.font = "11px ui-monospace, monospace"; + if (!rows) { + ctx.fillStyle = MUTED; + ctx.fillText("waiting for timeline data…", M_LEFT, 16); + return; + } + const plotW = rect.width - M_LEFT - M_RIGHT; + + // 256-entry color lookup: palette ids for phase, ramp for magnitudes + const lut = new Array(256).fill(NO_DATA); + if (data.palette) { + data.palette.forEach((name, i) => { + if (i > 0) lut[i] = hexToRgb(LIFECYCLE_COLORS[name] ?? phaseColors[name] ?? "#98a1b0"); + }); + } else { + for (let v = 0; v < 256; v++) { + lut[v] = RAMP_LO.map((lo, c) => Math.round(lo + ((RAMP_HI[c] - lo) * v) / 255)); + } + lut[0] = NO_DATA; + } + + const img = new ImageData(data.x_buckets, rows); + for (let i = 0; i < data.values.length; i++) { + const [r, g, b] = lut[data.values[i]]; + img.data[i * 4] = r; + img.data[i * 4 + 1] = g; + img.data[i * 4 + 2] = b; + img.data[i * 4 + 3] = 255; + } + const off = document.createElement("canvas"); + off.width = data.x_buckets; + off.height = rows; + off.getContext("2d").putImageData(img, 0, 0); + ctx.imageSmoothingEnabled = false; + ctx.drawImage(off, 0, 0, data.x_buckets, rows, M_LEFT, 0, plotW, h); + + const sampleRows = data.rows[0]?.sample_index !== undefined; + if (!sampleRows) { + // node boundaries (solid) and train/rollout role changes (dashed) + ctx.strokeStyle = MUTED; + const roleKey = (row) => row.roles.join(","); + for (let r = 1; r < rows; r++) { + const prev = data.rows[r - 1]; + const cur = data.rows[r]; + if (prev.node !== cur.node) ctx.setLineDash([]); + else if (roleKey(prev) !== roleKey(cur)) ctx.setLineDash([3, 3]); + else continue; + ctx.beginPath(); + ctx.moveTo(M_LEFT, r * rowH()); + ctx.lineTo(rect.width - M_RIGHT, r * rowH()); + ctx.stroke(); + } + ctx.setLineDash([]); + } + + const labelEvery = Math.ceil(rows / MAX_LABELED_ROWS); + ctx.fillStyle = TEXT; + for (let r = 0; r < rows; r += labelEvery) { + const label = sampleRows ? `s${data.rows[r].sample_index}` : `g${data.rows[r].index}`; + ctx.fillText(label, 8, r * rowH() + Math.min(rowH(), 11)); + } + + if (!sampleRows) { + // left-edge markers for lanes in the current selection + const selected = opts.selectedKeys(); + ctx.fillStyle = ACCENT; + for (let r = 0; r < rows; r++) { + if (selected.has(`${data.rows[r].node}:${data.rows[r].gpu}`)) { + ctx.fillRect(M_LEFT - 6, r * rowH(), 3, rowH()); + } + } + } + + if (brush) { + const [a, b] = [Math.min(brush.r0, brush.r1), Math.max(brush.r0, brush.r1)]; + ctx.fillStyle = "rgba(47, 111, 235, 0.18)"; + ctx.fillRect(M_LEFT, a * rowH(), plotW, (b - a + 1) * rowH()); + } + + scaleLabel.textContent = data.scale + ? `0–${fmtNum(data.scale.max)}` + : data.rows_total > data.rows.length + ? `showing ${data.rows.length}/${data.rows_total} trajectories` + : ""; + } + + const rowAt = (clientY) => { + const rect = canvas.getBoundingClientRect(); + return Math.max(0, Math.min(data.rows.length - 1, Math.floor((clientY - rect.top) / rowH()))); + }; + + canvas.onmousedown = (ev) => { + if (!data || !data.rows.length) return; + if (data.rows[0].sample_index !== undefined) return; // trajectories: nothing to select + const r = rowAt(ev.clientY); + brush = { r0: r, r1: r }; + redraw(); + }; + const onMouseUp = () => { + if (!brush) return; + const [a, b] = [Math.min(brush.r0, brush.r1), Math.max(brush.r0, brush.r1)]; + const keys = data.rows.slice(a, b + 1).map((row) => `${row.node}:${row.gpu}`); + brush = null; + redraw(); + opts.onBrush(keys); + }; + window.addEventListener("mouseup", onMouseUp); + + canvas.onmousemove = (ev) => { + if (!data || !data.rows.length) return; + if (brush) { + brush.r1 = rowAt(ev.clientY); + redraw(); + return; + } + const rect = canvas.getBoundingClientRect(); + const x = ev.clientX - rect.left; + if (x < M_LEFT || x > rect.width - M_RIGHT) { + hideTooltip(); + return; + } + const r = rowAt(ev.clientY); + const row = data.rows[r]; + const col = Math.min( + data.x_buckets - 1, + Math.floor(((x - M_LEFT) / (rect.width - M_LEFT - M_RIGHT)) * data.x_buckets), + ); + const value = data.values[r * data.x_buckets + col]; + const t = data.t0 + ((col + 0.5) / data.x_buckets) * (data.t1 - data.t0); + const identity = + row.sample_index !== undefined + ? `s${row.sample_index} (group ${row.group_index})` + : `g${row.index} ${row.node}:${row.gpu} (${row.roles.join("+") || "?"})`; + const lines = [`${identity} +${fmtNum(t - opts.runStart())}s`]; + if (data.palette) lines.push(`phase: ${data.palette[value] || "—"}`); + else lines.push(`${metric}: ${fmtNum((value / 255) * data.scale.max)}`); + showTooltip(ev.clientX, ev.clientY, lines.join("\n")); + }; + canvas.onmouseleave = hideTooltip; + + renderChips(); + return { + root, + refresh, + redraw, + destroy: () => window.removeEventListener("mouseup", onMouseUp), + }; +} diff --git a/miles/dashboard/static/charts.js b/miles/dashboard/static/charts.js new file mode 100644 index 0000000000..3d47618238 --- /dev/null +++ b/miles/dashboard/static/charts.js @@ -0,0 +1,407 @@ +// Minimal dependency-free canvas charts (line + scatter) with point picking. + +const MARGIN = { left: 52, right: 14, top: 10, bottom: 24 }; + +function setupCanvas(canvas) { + const dpr = window.devicePixelRatio || 1; + const rect = canvas.getBoundingClientRect(); + canvas.width = rect.width * dpr; + canvas.height = rect.height * dpr; + const ctx = canvas.getContext("2d"); + ctx.scale(dpr, dpr); + return { ctx, width: rect.width, height: rect.height }; +} + +function niceTicks(min, max, target = 5) { + if (!(max > min)) { + max = min + (Math.abs(min) || 1); + } + const span = max - min; + const step0 = Math.pow(10, Math.floor(Math.log10(span / target))); + const step = [1, 2, 5, 10].map((m) => m * step0).find((s) => span / s <= target) || step0 * 10; + const ticks = []; + for (let t = Math.ceil(min / step) * step; t <= max + 1e-9 * span; t += step) ticks.push(t); + return ticks; +} + +const fmt = (v) => { + if (v === 0) return "0"; + const a = Math.abs(v); + if (a >= 1e5 || a < 1e-3) return v.toExponential(1); + if (a >= 100) return v.toFixed(0); + if (a >= 1) return String(Math.round(v * 100) / 100); + return String(Math.round(v * 1000) / 1000); +}; + +// points: [{x, y, label?, flag?}]; opts: {line: bool, onClick(point), color} +// canvas._zoom = {x: [x0,x1]|null, y: [y0,y1]|null} — drag-selected domains; +// a null axis autoscales (y re-fits the points inside the x window) +const MAX_RAW_POINTS = 700; + +export function drawChart(canvas, points, opts = {}) { + const zoom = canvas._zoom ?? null; + let pts = zoom?.x ? points.filter((p) => p.x >= zoom.x[0] && p.x <= zoom.x[1]) : points; + let bucketed = false; + if (opts.line !== false && pts.length > MAX_RAW_POINTS) { + bucketed = true; + const size = Math.ceil(pts.length / MAX_RAW_POINTS); + const buckets = []; + for (let i = 0; i < pts.length; i += size) { + const bucket = pts.slice(i, i + size); + const y = bucket.reduce((sum, p) => sum + p.y, 0) / bucket.length; + buckets.push({ + x: bucket[0].x, + y, + label: `${fmt(bucket[0].x)}–${fmt(bucket.at(-1).x)}\nmean = ${fmt(y)} (${bucket.length} pts)`, + }); + } + pts = buckets; + } + const { ctx, width, height } = setupCanvas(canvas); + const css = getComputedStyle(document.documentElement); + const colText = css.getPropertyValue("--muted").trim(); + const colBorder = css.getPropertyValue("--border").trim(); + const colMain = opts.color || css.getPropertyValue("--accent").trim(); + const colFlag = css.getPropertyValue("--bad").trim(); + ctx.clearRect(0, 0, width, height); + + const plotW = width - MARGIN.left - MARGIN.right; + const plotH = height - MARGIN.top - MARGIN.bottom; + ctx.font = "11px ui-monospace, monospace"; + if (!points.length) { + ctx.fillStyle = colText; + ctx.fillText("no data", MARGIN.left + plotW / 2 - 20, MARGIN.top + plotH / 2); + return; + } + + let xMin, xMax; + if (zoom?.x) { + [xMin, xMax] = zoom.x; + } else { + const xs = pts.map((p) => p.x); + [xMin, xMax] = [Math.min(...xs), Math.max(...xs)]; + } + if (xMin === xMax) [xMin, xMax] = [xMin - 0.5, xMax + 0.5]; + let yMin, yMax; + if (zoom?.y) { + [yMin, yMax] = zoom.y; + } else { + const ys = pts.map((p) => p.y); + [yMin, yMax] = ys.length ? [Math.min(...ys), Math.max(...ys)] : [0, 1]; + if (yMin === yMax) [yMin, yMax] = [yMin - 0.5, yMax + 0.5]; + const yPad = (yMax - yMin) * 0.08; + yMin -= yPad; + yMax += yPad; + } + const X = (v) => MARGIN.left + ((v - xMin) / (xMax - xMin)) * plotW; + const Y = (v) => MARGIN.top + plotH - ((v - yMin) / (yMax - yMin)) * plotH; + + ctx.strokeStyle = colBorder; + ctx.fillStyle = colText; + ctx.lineWidth = 1; + for (const t of niceTicks(yMin, yMax)) { + ctx.beginPath(); + ctx.moveTo(MARGIN.left, Y(t)); + ctx.lineTo(width - MARGIN.right, Y(t)); + ctx.stroke(); + ctx.fillText(fmt(t), 4, Y(t) + 3); + } + for (const t of niceTicks(xMin, xMax, 8)) { + ctx.fillText(fmt(t), X(t) - 8, height - 6); + } + + // marks are clipped to the plot area so zoomed-out data cannot paint over + // the margins; the full point list is drawn so lines run to the box edges + ctx.save(); + ctx.beginPath(); + ctx.rect(MARGIN.left, MARGIN.top, plotW, plotH); + ctx.clip(); + if (opts.line !== false) { + const linePts = bucketed ? pts : points; + ctx.strokeStyle = colMain; + ctx.lineWidth = 1.5; + ctx.beginPath(); + linePts.forEach((p, i) => (i ? ctx.lineTo(X(p.x), Y(p.y)) : ctx.moveTo(X(p.x), Y(p.y)))); + ctx.stroke(); + } + if (!bucketed) { + for (const p of points) { + ctx.fillStyle = p.flag ? colFlag : colMain; + ctx.beginPath(); + ctx.arc(X(p.x), Y(p.y), opts.line !== false ? 3 : 3.5, 0, Math.PI * 2); + ctx.fill(); + } + } + ctx.restore(); + + const nearest = (ev) => { + const rect = canvas.getBoundingClientRect(); + const mx = ev.clientX - rect.left; + const my = ev.clientY - rect.top; + let best = null; + let bestDist = 20; + for (const p of pts) { + const d = Math.hypot(X(p.x) - mx, Y(p.y) - my); + if (d < bestDist) { + bestDist = d; + best = p; + } + } + return best; + }; + const hover = (ev) => { + const p = nearest(ev); + canvas.style.cursor = p && opts.onClick ? "pointer" : "crosshair"; + if (p) { + showTooltip(ev.clientX, ev.clientY, p.label ?? `${fmt(p.x)}, ${fmt(p.y)}`); + } else { + hideTooltip(); + } + }; + installDragZoom(canvas, { + plotLeft: MARGIN.left, + plotRight: width - MARGIN.right, + plotTop: MARGIN.top, + plotBottom: height - MARGIN.bottom, + invX: (px) => xMin + ((px - MARGIN.left) / plotW) * (xMax - xMin), + invY: (py) => yMax - ((py - MARGIN.top) / plotH) * (yMax - yMin), + redraw: () => drawChart(canvas, points, opts), + hover, + onClick: + opts.onClick && + ((ev) => { + const p = nearest(ev); + if (p) opts.onClick(p); + }), + }); +} + +// an axis participates in a zoom only when dragged at least this many px, so +// a horizontal-ish sweep zooms x and re-autoscales y instead of pinning y to +// the few pixels the hand wobbled +const AXIS_MIN_DRAG = 10; + +// drag a box on the plot = zoom into the selected x/y ranges; double-click +// resets. Handlers are re-installed on every render so invX/invY match the +// pixels on screen; drag state and the zoom live on the canvas element +// because redraw() re-renders (and re-installs handlers) mid-drag. move/up +// listen on window: a real drag routes events to whatever element is under +// the pointer, so canvas-local handlers would lose any drag that strays +// outside the canvas band and never complete the zoom. +function installDragZoom(canvas, { plotLeft, plotRight, plotTop, plotBottom, invX, invY, redraw, hover, onClick }) { + const pos = (ev) => { + const rect = canvas.getBoundingClientRect(); + return { + x: Math.max(plotLeft, Math.min(plotRight, ev.clientX - rect.left)), + y: Math.max(plotTop, Math.min(plotBottom, ev.clientY - rect.top)), + }; + }; + canvas.onmousemove = (ev) => { + if (canvas._dragAnchor == null) hover(ev); + }; + canvas.onmousedown = (ev) => { + ev.preventDefault(); // no native text selection mid-drag + canvas._dragAnchor = pos(ev); + hideTooltip(); + const onMove = (mv) => { + const cur = pos(mv); + redraw(); + const a = canvas._dragAnchor; + const ctx = canvas.getContext("2d"); + ctx.fillStyle = "rgba(47, 111, 235, 0.15)"; + // full-height band while the drag is x-only; a box once it has y extent + const boxY = Math.abs(cur.y - a.y) >= AXIS_MIN_DRAG; + ctx.fillRect( + Math.min(a.x, cur.x), + boxY ? Math.min(a.y, cur.y) : plotTop, + Math.abs(cur.x - a.x), + boxY ? Math.abs(cur.y - a.y) : plotBottom - plotTop, + ); + }; + const onUp = (up) => { + window.removeEventListener("mousemove", onMove); + window.removeEventListener("mouseup", onUp); + const a = canvas._dragAnchor; + canvas._dragAnchor = null; + const cur = pos(up); + const dx = Math.abs(cur.x - a.x); + const dy = Math.abs(cur.y - a.y); + if (dx < 4 && dy < 4) { + redraw(); + if (onClick) onClick(up); + return; + } + const prev = canvas._zoom ?? { x: null, y: null }; + canvas._zoom = { + x: dx >= AXIS_MIN_DRAG ? [invX(Math.min(a.x, cur.x)), invX(Math.max(a.x, cur.x))] : prev.x, + y: + dy >= AXIS_MIN_DRAG + ? [invY(Math.max(a.y, cur.y)), invY(Math.min(a.y, cur.y))] + : dx >= AXIS_MIN_DRAG + ? null // x-only sweep: let y re-fit the new window + : prev.y, + }; + redraw(); + }; + window.addEventListener("mousemove", onMove); + window.addEventListener("mouseup", onUp); + }; + canvas.onmouseleave = hideTooltip; + canvas.ondblclick = () => { + canvas._zoom = null; + redraw(); + }; +} + +// ------------------------------ tooltip ------------------------------------- + +let tooltipEl = null; + +export function showTooltip(clientX, clientY, text) { + if (!tooltipEl) { + tooltipEl = document.createElement("div"); + tooltipEl.id = "tooltip"; + document.body.appendChild(tooltipEl); + } + tooltipEl.textContent = text; + tooltipEl.style.display = "block"; + const pad = 12; + const w = tooltipEl.offsetWidth; + const x = clientX + pad + w > window.innerWidth ? clientX - w - pad : clientX + pad; + tooltipEl.style.left = `${x}px`; + tooltipEl.style.top = `${Math.min(clientY + pad, window.innerHeight - tooltipEl.offsetHeight - pad)}px`; +} + +export function hideTooltip() { + if (tooltipEl) tooltipEl.style.display = "none"; +} + +// diverging color for values around a center (e.g. imp_ratio around 1) +export function divergingColor(t) { + // t in [-1, 1]: blue -> transparent -> red + const a = Math.min(1, Math.abs(t)) * 0.85; + return t >= 0 ? `rgba(224, 96, 96, ${a})` : `rgba(78, 161, 255, ${a})`; +} + +// sequential color for magnitudes in [0, 1] +export function sequentialColor(t) { + return `rgba(70, 194, 142, ${Math.min(1, Math.max(0, t)) * 0.85})`; +} + + +// fixed categorical order, assigned by sorted-series index (stable across +// refreshes because the store sorts by addr); cycles past 10 series +export const SERIES_COLORS = [ + "#2f6feb", "#e8590c", "#12b886", "#d6336c", "#7048e8", + "#f59f00", "#0c8599", "#5c940d", "#a61e4d", "#495057", +]; + +// seriesList: [{label, ts: [], value: []}] — one thin line per engine on a +// shared scale, one color per engine; hover names the engine +export function drawMultiLine(canvas, seriesList, opts = {}) { + const zoom = canvas._zoom ?? null; + // opts.hidden: legend-unchecked engines render as empty (drop out of the + // y-scale too); opts.colorIndex(label) pins colors to a page-wide order + const shown = opts.hidden ? seriesList.map((s) => (opts.hidden.has(s.label) ? { ...s, ts: [], value: [] } : s)) : seriesList; + const visible = zoom?.x + ? shown.map((s) => { + const ts = []; + const value = []; + s.ts.forEach((t, i) => { + if (t >= zoom.x[0] && t <= zoom.x[1]) { + ts.push(t); + value.push(s.value[i]); + } + }); + return { ...s, ts, value }; + }) + : shown; + const { ctx, width, height } = setupCanvas(canvas); + const css = getComputedStyle(document.documentElement); + const colText = css.getPropertyValue("--muted").trim(); + const colBorder = css.getPropertyValue("--border").trim(); + ctx.clearRect(0, 0, width, height); + const plotW = width - MARGIN.left - MARGIN.right; + const plotH = height - MARGIN.top - MARGIN.bottom; + ctx.font = "11px ui-monospace, monospace"; + + const alive = visible.filter((s) => s.ts.length); + if (!alive.length && !zoom) { + ctx.fillStyle = colText; + ctx.fillText("no data", MARGIN.left + plotW / 2 - 20, MARGIN.top + plotH / 2); + return; + } + // x origin stays the run-wide first sample so +m:ss labels keep meaning under zoom + const withData = seriesList.filter((s) => s.ts.length); + const x0 = withData.length ? Math.min(...withData.map((s) => s.ts[0])) : 0; + const [xMin, xMax] = zoom?.x ?? [x0, alive.length ? Math.max(...alive.map((s) => s.ts.at(-1))) : x0 + 1]; + let yMin, yMax; + if (zoom?.y) { + [yMin, yMax] = zoom.y; + } else { + [yMin, yMax] = alive.length + ? [Math.min(...alive.map((s) => Math.min(...s.value))), Math.max(...alive.map((s) => Math.max(...s.value)))] + : [0, 1]; + } + if (yMin === yMax) [yMin, yMax] = [yMin - 0.5, yMax + 0.5]; + const X = (t) => MARGIN.left + ((t - xMin) / Math.max(xMax - xMin, 1e-9)) * plotW; + const Y = (v) => MARGIN.top + (1 - (v - yMin) / (yMax - yMin)) * plotH; + + ctx.strokeStyle = colBorder; + ctx.fillStyle = colText; + for (const tick of niceTicks(yMin, yMax, 4)) { + ctx.beginPath(); + ctx.moveTo(MARGIN.left, Y(tick)); + ctx.lineTo(width - MARGIN.right, Y(tick)); + ctx.stroke(); + ctx.fillText(fmt(tick), 6, Y(tick) + 4); + } + for (const tick of niceTicks(xMin - x0, xMax - x0, 6)) { + const rel = Math.round(tick); + ctx.fillText(`+${Math.floor(rel / 60)}:${String(rel % 60).padStart(2, "0")}`, X(x0 + tick) - 12, height - 6); + } + + ctx.save(); + ctx.beginPath(); + ctx.rect(MARGIN.left, MARGIN.top, plotW, plotH); + ctx.clip(); + ctx.lineWidth = 1; + visible.forEach((s, i) => { + if (!s.ts.length) return; + // index in the page-wide engine order, not the surviving subset: color follows the engine + const idx = opts.colorIndex ? opts.colorIndex(s.label) : i; + ctx.strokeStyle = SERIES_COLORS[idx % SERIES_COLORS.length]; + ctx.beginPath(); + s.ts.forEach((t, j) => (j ? ctx.lineTo(X(t), Y(s.value[j])) : ctx.moveTo(X(t), Y(s.value[j])))); + ctx.stroke(); + }); + ctx.restore(); + + const hover = (ev) => { + const rect = canvas.getBoundingClientRect(); + const mx = ev.clientX - rect.left; + const my = ev.clientY - rect.top; + let best = null; + for (const s of alive) { + for (let i = 0; i < s.ts.length; i++) { + const d = (X(s.ts[i]) - mx) ** 2 + (Y(s.value[i]) - my) ** 2; + if (!best || d < best.d) best = { d, s, i }; + } + } + if (!best || best.d > 30 ** 2) { + hideTooltip(); + return; + } + showTooltip(ev.clientX, ev.clientY, `${best.s.label}\n+${fmt(best.s.ts[best.i] - x0)}s = ${fmt(best.s.value[best.i])}`); + }; + installDragZoom(canvas, { + plotLeft: MARGIN.left, + plotRight: width - MARGIN.right, + plotTop: MARGIN.top, + plotBottom: height - MARGIN.bottom, + invX: (px) => xMin + ((px - MARGIN.left) / plotW) * Math.max(xMax - xMin, 1e-9), + invY: (py) => yMax - ((py - MARGIN.top) / plotH) * (yMax - yMin), + redraw: () => drawMultiLine(canvas, seriesList, opts), + hover, + }); +} diff --git a/miles/dashboard/static/conversation.js b/miles/dashboard/static/conversation.js new file mode 100644 index 0000000000..ba1e940a88 --- /dev/null +++ b/miles/dashboard/static/conversation.js @@ -0,0 +1,91 @@ +import { el, fmtNum } from "./app.js"; + +const CLOSED_BY_DEFAULT = new Set(["system", "tool"]); + +// row: one trajectory-sidecar entry {status, reward, messages: [...]} +// (OpenAI message shape: role/content plus optional reasoning_content, +// tool_calls, name) +export function renderConversation(row) { + const chips = [el("span", { class: "chip" }, [row.status])]; + if (row.reward !== null && row.reward !== undefined) { + chips.push(el("span", { class: "chip" }, [`reward ${fmtNum(row.reward)}`])); + } + return el("div", { class: "panel" }, [ + el("h3", {}, ["conversation"]), + el("div", { class: "controls" }, chips), + ...row.messages.map(messageCard), + ]); +} + +// display-level split of raw assistant text (engine multi_turn path records +// the model's verbatim emission): blocks, markup, and +// trailing special tokens; unknown markup stays visible as-is +function splitRawAssistant(text) { + const thinking = []; + let rest = text.replace(/([\s\S]*?)<\/think>/g, (_, body) => { + thinking.push(body.trim()); + return ""; + }); + const calls = []; + rest = rest.replace(/\s*([\s\S]*?)\s*<\/tool_call>/g, (_, body) => { + calls.push(body); + return ""; + }); + rest = rest.replace(/(?:<\|[^|>]+\|>\s*)+$/, "").trim(); + return { thinking: thinking.join("\n\n"), calls, rest }; +} + +function toolCallBlock(name, args) { + let pretty = args; + try { + pretty = JSON.stringify(typeof args === "string" ? JSON.parse(args) : args, null, 2); + } catch { + /* keep raw args */ + } + return el("div", {}, [el("span", { class: "chip" }, [name]), el("pre", { class: "msg-text" }, [pretty])]); +} + +function messageCard(message, index) { + const role = message.role ?? "?"; + const body = []; + let content = message.content; + let rawCalls = []; + if (role === "assistant" && typeof content === "string" && !message.reasoning_content && !message.tool_calls) { + const split = splitRawAssistant(content); + if (split.thinking) body.push(el("div", { class: "msg-thinking" }, [split.thinking])); + content = split.rest; + rawCalls = split.calls; + } + if (message.reasoning_content) { + body.push(el("div", { class: "msg-thinking" }, [message.reasoning_content])); + } + if (content) { + body.push(el("pre", { class: "msg-text" }, [typeof content === "string" ? content : JSON.stringify(content, null, 2)])); + } + for (const raw of rawCalls) { + let name = "tool_call"; + let args = raw; + try { + const parsed = JSON.parse(raw); + name = parsed.name ?? name; + args = parsed.arguments ?? parsed; + } catch { + /* unparseable call body: show verbatim */ + } + body.push(toolCallBlock(name, args)); + } + for (const call of message.tool_calls ?? []) { + const fn = call.function ?? call; + body.push(toolCallBlock(fn.name ?? "tool", fn.arguments ?? "")); + } + const bodyWrap = el("div", { class: `msg-body${CLOSED_BY_DEFAULT.has(role) ? " closed" : ""}` }, body); + const header = el( + "div", + { class: "msg-head", onclick: () => bodyWrap.classList.toggle("closed") }, + [ + el("span", { class: `msg-role role-${role}` }, [role + (message.name ? ` · ${message.name}` : "")]), + el("span", { class: "muted" }, [`#${index}`]), + ], + ); + return el("div", { class: "msg" }, [header, bodyWrap]); +} diff --git a/miles/dashboard/static/index.html b/miles/dashboard/static/index.html new file mode 100644 index 0000000000..dc36a4becd --- /dev/null +++ b/miles/dashboard/static/index.html @@ -0,0 +1,18 @@ + + + + + + miles dashboard + + + +
+ miles dashboard + + +
+

loading…

+ + + diff --git a/miles/dashboard/static/style.css b/miles/dashboard/static/style.css new file mode 100644 index 0000000000..61de14b330 --- /dev/null +++ b/miles/dashboard/static/style.css @@ -0,0 +1,118 @@ +:root { + --bg: #f5f6f8; + --panel: #ffffff; + --border: #dde1e8; + --text: #24292f; + --muted: #667080; + --accent: #2f6feb; + --good: #1a9e6e; + --bad: #d1442e; + font-size: 14px; +} +* { box-sizing: border-box; } +body { margin: 0; background: var(--bg); color: var(--text); font-family: -apple-system, "Segoe UI", sans-serif; } +a { color: var(--accent); text-decoration: none; } +.muted { color: var(--muted); } + +#topbar { + display: flex; align-items: stretch; gap: 22px; padding: 0 18px; + background: #1d2433; border-bottom: 1px solid #2a3346; + position: sticky; top: 0; z-index: 10; min-height: 46px; +} +.brand { font-weight: 700; color: #fff; font-size: 16px; align-self: center; } +#crumbs { flex: 1; display: flex; align-items: stretch; gap: 8px; } +#crumbs a.nav { + display: flex; align-items: center; padding: 0 22px; + color: #d7deea; font-weight: 700; font-size: 15px; +} +#crumbs a.nav:hover { background: rgba(255, 255, 255, 0.14); color: #fff; } +#crumbs a.nav.active { color: #fff; box-shadow: inset 0 -2px 0 var(--accent); } +#crumbs .crumb { align-self: center; color: #8fa0b8; font-size: 13px; } +#crumbs .crumb a { color: #aebedd; } +#runinfo { color: #8fa0b8; font-size: 12px; align-self: center; } + +#view { padding: 16px 18px 60px; max-width: 1400px; margin: 0 auto; } + +.panel { background: var(--panel); border: 1px solid var(--border); border-radius: 6px; padding: 12px; margin-bottom: 14px; } +.panel h3 { margin: 0 0 8px; font-size: 13px; font-weight: 600; color: var(--muted); text-transform: uppercase; letter-spacing: 0.04em; } +.row { display: flex; gap: 14px; flex-wrap: wrap; align-items: flex-start; } +.row > .panel { flex: 1; min-width: 380px; } + +.controls { display: flex; gap: 8px; align-items: center; flex-wrap: wrap; margin-bottom: 10px; } +.controls input[type="text"], .controls input[type="number"], .controls select { + background: var(--bg); color: var(--text); border: 1px solid var(--border); border-radius: 4px; padding: 4px 8px; +} +.controls input[type="number"] { width: 90px; } +button, .btn { + background: var(--bg); color: var(--text); border: 1px solid var(--border); + border-radius: 4px; padding: 4px 10px; cursor: pointer; font-size: 13px; +} +button:hover { border-color: var(--accent); } +button.active { border-color: var(--accent); color: var(--accent); } + +.keylist { max-height: 300px; overflow-y: auto; display: flex; flex-direction: column; gap: 2px; } +.keylist label { display: block; padding: 2px 4px; cursor: pointer; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; } +.keylist label:hover { background: var(--bg); } + +canvas.chart { width: 100%; height: 220px; display: block; cursor: crosshair; } +.chart-title { font-size: 13px; color: var(--text); margin: 0 0 4px; font-family: ui-monospace, monospace; } + +table.data { border-collapse: collapse; width: 100%; font-size: 13px; font-variant-numeric: tabular-nums; } +table.data th, table.data td { padding: 4px 8px; border-bottom: 1px solid var(--border); text-align: right; white-space: nowrap; } +table.data th:first-child, table.data td:first-child { text-align: left; } +table.data th { color: var(--muted); cursor: pointer; user-select: none; position: sticky; top: 0; background: var(--panel); } +table.data th.sorted { color: var(--accent); } +table.data tbody tr:hover { background: #eef2f8; cursor: pointer; } +table.data tr.flagged td { color: var(--bad); } +.tablewrap { max-height: 480px; overflow: auto; } + +.tabs { display: flex; gap: 4px; margin-bottom: 10px; } +.tokens { line-height: 1.9; font-family: ui-monospace, monospace; font-size: 13px; word-break: break-all; } +.tok { border-radius: 2px; padding: 1px 0; position: relative; } +.tok.prompt { color: var(--muted); } +.tok.masked { opacity: 0.35; } +.tok:hover { outline: 1px solid var(--accent); } +#tooltip { + position: fixed; pointer-events: none; background: #2b3038; color: #f0f2f5; border: 1px solid #2b3038; + border-radius: 4px; padding: 6px 9px; font-size: 12px; z-index: 100; display: none; + font-family: ui-monospace, monospace; white-space: pre; +} +.legend { display: flex; align-items: center; gap: 8px; font-size: 12px; color: var(--muted); } +.legend .bar { width: 140px; height: 10px; border-radius: 2px; } +.statgrid { display: grid; grid-template-columns: repeat(auto-fill, minmax(150px, 1fr)); gap: 8px; margin-bottom: 10px; } +.stat { background: var(--bg); border: 1px solid var(--border); border-radius: 4px; padding: 6px 10px; } +.stat .v { font-size: 16px; font-weight: 600; } +.stat .k { font-size: 11px; color: var(--muted); } +.error { color: var(--bad); padding: 12px; } + +canvas.timeline { width: 100%; display: block; cursor: crosshair; } +canvas.carpet { width: 100%; display: block; cursor: crosshair; } +canvas.anatomy { width: 100%; display: block; cursor: pointer; } +.anatomywrap { overflow-y: auto; } +.chip { + display: inline-flex; align-items: center; gap: 4px; background: var(--bg); + border: 1px solid var(--border); border-radius: 12px; padding: 2px 8px; + font-size: 12px; font-family: ui-monospace, monospace; +} +.chip button { border: none; background: none; padding: 0 2px; font-size: 12px; color: var(--muted); } +.chip button:hover { color: var(--bad); border: none; } +.warn { color: #9a6700; font-size: 12px; } +.bubblestrip { display: flex; gap: 3px; align-items: center; margin-bottom: 10px; flex-wrap: wrap; } +.bubble { + min-width: 26px; text-align: center; padding: 2px 4px; border: 1px solid var(--border); + border-radius: 3px; font-size: 12px; cursor: pointer; font-variant-numeric: tabular-nums; +} +.bubble:hover { border-color: var(--accent); } + +.msg { border: 1px solid var(--border); border-radius: 6px; margin-bottom: 8px; } +.msg-head { display: flex; gap: 8px; padding: 6px 10px; cursor: pointer; align-items: baseline; } +.msg-head:hover { background: var(--bg); } +.msg-role { font-weight: 600; font-size: 12px; text-transform: uppercase; letter-spacing: 0.03em; } +.role-user { color: var(--accent); } +.role-assistant { color: var(--good); } +.role-tool { color: #d97706; } +.role-system { color: var(--muted); } +.msg-body { padding: 0 10px 8px; } +.msg-body.closed { display: none; } +.msg-text { white-space: pre-wrap; word-break: break-word; font-size: 13px; margin: 4px 0; max-height: 320px; overflow-y: auto; background: var(--bg); border-radius: 4px; padding: 8px; } +.msg-thinking { white-space: pre-wrap; font-style: italic; color: var(--muted); font-size: 12px; max-height: 200px; overflow-y: auto; border-left: 3px solid var(--border); padding: 6px 10px; margin: 4px 0; } diff --git a/miles/dashboard/static/views_metrics.js b/miles/dashboard/static/views_metrics.js new file mode 100644 index 0000000000..dc7c36e44e --- /dev/null +++ b/miles/dashboard/static/views_metrics.js @@ -0,0 +1,222 @@ +import { api } from "./api.js"; +import { el, setViewCleanup } from "./app.js"; +import { drawChart, drawMultiLine, SERIES_COLORS } from "./charts.js"; + +// wandb-shaped L0: metric names come verbatim from the tracking fan-out, so +// categories are simply the key prefixes. One category at a time (wandb +// sidebar idiom); every key of the active category renders as a line chart. +// "dump" only ever appears for dump-only dirs — the server does not +// advertise it while a telemetry stream exists. +const CATEGORY_ORDER = ["rollout", "perf", "train", "eval", "dump"]; + +const categoryOf = (key) => key.split("/")[0]; + +function axisOf(key) { + if (key.startsWith("dump/")) return "dump"; + if (key.startsWith("train/")) return "train/step"; + if (key.startsWith("eval/")) return "eval/step"; + return "rollout/step"; +} + +function orderedCategories(meta) { + const present = [...new Set(meta.metric_keys.map(categoryOf))]; + const cats = [ + ...CATEGORY_ORDER.filter((c) => present.includes(c)), + ...present.filter((c) => !CATEGORY_ORDER.includes(c) && c !== "dump").sort(), + ]; + // grafana-parity category: scraped engine metrics, x = wall clock + if (meta.engine_metric_keys?.length) cats.splice(cats.includes("dump") ? cats.length - 1 : cats.length, 0, "sglang"); + return cats; +} + +export async function renderMetrics(view, meta) { + const cats = orderedCategories(meta); + let active = sessionStorage.getItem("metricsCategory"); + if (!cats.includes(active)) active = cats[0] ?? null; + + const chartsPanel = el("div", { + style: "flex: 3; min-width: 500px; display: grid; grid-template-columns: repeat(2, minmax(0, 500px)); gap: 0 14px; align-content: start", + }); + const filterInput = el("input", { type: "text", placeholder: "filter…" }); + const catList = el("div", { class: "keylist" }); + + const renderCategories = () => { + catList.replaceChildren( + ...cats.map((cat) => + el( + "button", + { + class: cat === active ? "active" : "", + style: "text-align: left", + onclick: () => { + active = cat; + sessionStorage.setItem("metricsCategory", cat); + renderCategories(); + buildPanels(); + }, + }, + [cat], + ), + ), + ); + }; + + const activeKeys = () => { + const needle = filterInput.value.toLowerCase(); + const pool = active === "sglang" ? meta.engine_metric_keys : meta.metric_keys.filter((k) => categoryOf(k) === active); + return pool.filter((k) => k.toLowerCase().includes(needle)); + }; + + // slots persist across refreshes; a chart repaints ONLY when its data + // changed (signature check), so a live page sits visually still between + // real updates instead of flickering on every poll + const slots = new Map(); // key -> {canvas, status, signature, lastSeries?} + let epoch = 0; + + // sglang legend: one checkbox per engine, all on by default; unchecking + // hides that engine's line in every chart (and drops it from the y-scale) + const engines = []; // sorted union of scraped addrs, fixes the color order page-wide + const hiddenEngines = new Set(); + const legendPanel = el("div", { class: "panel", style: "flex: 0 0 240px; min-width: 240px; display: none" }); + const engineOpts = () => ({ hidden: hiddenEngines, colorIndex: (label) => engines.indexOf(label) }); + const redrawEngineCharts = () => { + for (const slot of slots.values()) { + if (slot.lastSeries) drawMultiLine(slot.canvas, slot.lastSeries, engineOpts()); + } + }; + const renderLegend = () => { + legendPanel.replaceChildren( + el("h3", {}, ["engines"]), + ...engines.map((addr, i) => + el( + "label", + { style: "display: flex; align-items: center; gap: 6px; padding: 2px 0; cursor: pointer; font-size: 12px" }, + [ + el("input", { + type: "checkbox", + ...(hiddenEngines.has(addr) ? {} : { checked: "" }), + onchange: (ev) => { + if (ev.target.checked) hiddenEngines.delete(addr); + else hiddenEngines.add(addr); + redrawEngineCharts(); + }, + }), + el("span", { + style: `display: inline-block; width: 14px; height: 3px; background: ${SERIES_COLORS[i % SERIES_COLORS.length]}`, + }), + addr.replace(/^https?:\/\//, ""), + ], + ), + ), + ); + }; + + function buildPanels() { + slots.clear(); + legendPanel.style.display = active === "sglang" ? "" : "none"; + const keys = activeKeys(); + if (!keys.length) { + chartsPanel.replaceChildren(el("p", { class: "muted" }, [active ? "no metrics here" : "no metrics logged"])); + return; + } + chartsPanel.replaceChildren( + ...keys.map((key) => { + const canvas = el("canvas", { class: "chart" }); + const status = el("p", { class: "muted" }, ["loading…"]); + slots.set(key, { canvas, status, signature: null }); + return el("div", { class: "panel" }, [el("p", { class: "chart-title" }, [key]), status, canvas]); + }), + ); + refresh(); + } + + function refresh() { + const current = ++epoch; + if (active === "sglang") { + for (const [key, slot] of slots.entries()) { + api("/api/timeline/engine_series", { metric: key, max_points: 1000 }) + .then(({ series }) => { + if (current !== epoch) return; + const signature = series.map((s) => `${s.addr}:${s.ts.length}:${s.ts.at(-1)}`).join("|"); + if (signature === slot.signature) return; + slot.signature = signature; + slot.status.remove(); + slot.lastSeries = series.map((s) => ({ label: s.addr, ts: s.ts, value: s.value })); + let grew = false; + for (const s of series) { + if (!engines.includes(s.addr)) { + engines.push(s.addr); + grew = true; + } + } + if (grew) { + engines.sort(); + renderLegend(); + redrawEngineCharts(); // sort may have shifted the color order + } else { + drawMultiLine(slot.canvas, slot.lastSeries, engineOpts()); + } + }) + .catch((err) => { + if (current === epoch && slot.signature === null) slot.status.textContent = String(err); + }); + } + return; + } + const byAxis = new Map(); + for (const key of slots.keys()) { + const axis = axisOf(key); + if (!byAxis.has(axis)) byAxis.set(axis, []); + byAxis.get(axis).push(key); + } + for (const [axis, keys] of byAxis.entries()) { + api("/api/metrics", { keys: keys.join(","), x: axis === "dump" ? "rollout/step" : axis }) + .then((series) => { + if (current !== epoch) return; // category/filter changed meanwhile + for (const key of keys) { + const slot = slots.get(key); + if (!slot) continue; + const s = series[key] ?? { x: [], y: [] }; + const signature = `${s.x.length}:${s.x.at(-1)}:${s.y.at(-1)}`; + if (signature === slot.signature) continue; // unchanged: no repaint + slot.signature = signature; + slot.status.remove(); + drawChart( + slot.canvas, + s.x.map((x, i) => ({ x, y: s.y[i], label: `step ${x}\n${key} = ${s.y[i]}` })), + ); + } + }) + .catch((err) => { + if (current !== epoch) return; + for (const key of keys) { + const slot = slots.get(key); + if (slot && slot.signature === null) slot.status.textContent = String(err); + } + }); + } + } + + filterInput.oninput = buildPanels; + + view.replaceChildren( + el("div", { class: "row" }, [ + el("div", { class: "panel", style: "flex: 0 0 200px; min-width: 200px" }, [ + el("h3", {}, ["metrics"]), + el("div", { class: "controls" }, [filterInput]), + catList, + ]), + chartsPanel, + legendPanel, + ]), + ); + renderCategories(); + buildPanels(); + + if (meta.mode === "follow") { + // refresh is fire-and-forget and idempotent: unchanged charts skip the + // repaint entirely, so the live page sits still between real updates + const intervalId = setInterval(refresh, 5000); + setViewCleanup(() => clearInterval(intervalId)); + } +} diff --git a/miles/dashboard/static/views_rollout.js b/miles/dashboard/static/views_rollout.js new file mode 100644 index 0000000000..94faaa8753 --- /dev/null +++ b/miles/dashboard/static/views_rollout.js @@ -0,0 +1,221 @@ +import { createAnatomy } from "./anatomy.js"; +import { api } from "./api.js"; +import { el, fmtNum } from "./app.js"; +import { drawChart } from "./charts.js"; + +const DEFAULT_COLUMNS = [ + "sample_index", + "group_index", + "raw_reward", + "reward", + "response_length", + "truncated", + "versions", + "turns", + "tool_calls", + "mean_abs_lp_diff", + "mean_imp_ratio", + "mean_entropy", + "adv_mean", + "non_generation_time", +]; + +// staleness span rendered as one sortable-ish string column; mixed-version +// samples are exactly what --use-tis corrects, so they must not blend in +function versionSpan(row) { + if (row.weight_version === null || row.weight_version === undefined) return null; + const lo = row.weight_version_min; + return row.mixed_version ? `v${lo}–v${row.weight_version}` : `v${row.weight_version}`; +} + +function sortableTable(rows, columns, { onRowClick, flagRow, sortState }) { + const wrap = el("div", { class: "tablewrap" }); + const render = () => { + const { column, desc } = sortState; + const sorted = [...rows].sort((a, b) => { + const [va, vb] = [a[column], b[column]]; + if (va === null) return 1; + if (vb === null) return -1; + return (va > vb ? 1 : va < vb ? -1 : 0) * (desc ? -1 : 1); + }); + const head = el("tr", {}, columns.map((c) => + el("th", { + class: c === column ? "sorted" : "", + onclick: () => { + sortState.desc = sortState.column === c ? !sortState.desc : true; + sortState.column = c; + render(); + }, + }, [c === column ? `${c} ${desc ? "▾" : "▴"}` : c]), + )); + const body = sorted.map((row) => + el("tr", { + class: flagRow && flagRow(row) ? "flagged" : "", + onclick: onRowClick ? () => onRowClick(row) : null, + }, columns.map((c) => el("td", {}, [fmtNum(row[c])]))), + ); + wrap.replaceChildren(el("table", { class: "data" }, [el("thead", {}, [head]), el("tbody", {}, body)])); + }; + render(); + return wrap; +} + +function statBox(label, value) { + return el("div", { class: "stat" }, [el("div", { class: "v" }, [fmtNum(value)]), el("div", { class: "k" }, [label])]); +} + +export async function renderRollout(view, meta, route) { + const { rolloutId, evaluation } = route; + const [summary, groups] = await Promise.all([ + api(`/api/rollout/${rolloutId}/summary`, { eval: evaluation }), + api(`/api/rollout/${rolloutId}/groups`, { eval: evaluation }), + ]); + const rows = summary.rows; + for (const row of rows) row.versions = versionSpan(row); + // lifecycle lanes exist only for runs with the trajectory probes (PR §18); + // older dumps or eval steps: no pane, table only + let trajectories = { lanes: [], consume_ts: null }; + if (!evaluation) { + try { + trajectories = await api(`/api/rollout/${rolloutId}/trajectories`); + } catch { + /* endpoint absent or no events: keep the table-only layout */ + } + } + const rewardKey = evaluation ? "reward" : "raw_reward"; + + // -------- header controls: prev/next step, train/eval toggle -------- + const ids = evaluation ? meta.rollout_ids.eval : meta.rollout_ids.train; + const position = ids.indexOf(rolloutId); + const goto = (id, toEval) => (location.hash = `#/rollout/${id}${toEval ? "?eval=1" : ""}`); + // type a step number + Enter to jump straight to it + const jumpInput = el("input", { + type: "number", + value: String(rolloutId), + style: "width: 64px", + onkeydown: (ev) => { + if (ev.key === "Enter") goto(Number(ev.target.value), evaluation); + }, + }); + const controls = el("div", { class: "controls" }, [ + el("button", { onclick: () => position > 0 && goto(ids[position - 1], evaluation) }, ["◀ prev"]), + el("span", {}, [`${evaluation ? "eval" : "train"} step`]), + jumpInput, + el("span", {}, [`(${position + 1}/${ids.length})`]), + el("button", { onclick: () => position < ids.length - 1 && goto(ids[position + 1], evaluation) }, ["next ▶"]), + ]); + if (!evaluation && meta.rollout_ids.eval.includes(rolloutId)) { + controls.append(el("button", { onclick: () => goto(rolloutId, true) }, ["eval view"])); + } + if (evaluation && meta.rollout_ids.train.includes(rolloutId)) { + controls.append(el("button", { onclick: () => goto(rolloutId, false) }, ["train view"])); + } + + // -------- headline stats -------- + const rewards = rows.map((r) => r[rewardKey]).filter((v) => v !== null); + const mean = (vs) => (vs.length ? vs.reduce((a, b) => a + b, 0) / vs.length : null); + const zeroStdGroups = groups.rows.filter((g) => g.zero_std).length; + const stats = el("div", { class: "statgrid" }, [ + statBox("samples", rows.length), + statBox("reward mean", mean(rewards)), + statBox("truncated frac", mean(rows.map((r) => (r.truncated ? 1 : 0)))), + statBox("zero-std groups", `${zeroStdGroups}/${groups.rows.length}`), + statBox("mixed-version frac", mean(rows.map((r) => (r.mixed_version === null ? null : +r.mixed_version)).filter((v) => v !== null))), + // staleness = trainer version at consume − generation version; the engine + // counter is 1 after the startup push and +1 per train step, so the + // trainer holds v(step+1) when consuming step N + statBox( + "avg staleness", + evaluation + ? null + : mean(rows.map((r) => (r.weight_version_min == null ? null : rolloutId + 1 - r.weight_version_min)).filter((v) => v !== null)), + ), + statBox("mean |lp diff|", mean(rows.map((r) => r.mean_abs_lp_diff).filter((v) => v !== null))), + statBox("mean entropy", mean(rows.map((r) => r.mean_entropy).filter((v) => v !== null))), + ]); + + // -------- tabs -------- + const openTokens = (row) => + (location.hash = `#/rollout/${rolloutId}/sample/${row.sample_index}${evaluation ? "?eval=1" : ""}`); + + const samplesTab = () => { + const scatter = el("canvas", { class: "chart", style: "height: 260px" }); + queueMicrotask(() => + drawChart( + scatter, + rows + .filter((r) => r[rewardKey] !== null) + .map((r) => ({ + x: r.response_length, + y: r[rewardKey], + flag: Boolean(r.truncated), + label: `sample ${r.sample_index}\nreward=${fmtNum(r[rewardKey])} len=${r.response_length}` + + (r.mean_abs_lp_diff !== null ? `\n|lp diff|=${fmtNum(r.mean_abs_lp_diff)}` : ""), + row: r, + })), + { line: false, onClick: (p) => openTokens(p.row) }, + ), + ); + const allColumns = el("input", { type: "checkbox" }); + const tableHolder = el("div", {}); + const renderTable = () => { + const columns = allColumns.checked + ? summary.columns + : DEFAULT_COLUMNS.filter((c) => summary.columns.includes(c)); + tableHolder.replaceChildren( + sortableTable(rows, columns, { + onRowClick: openTokens, + flagRow: (r) => r.remove_sample, + sortState: { column: "sample_index", desc: false }, + }), + ); + }; + allColumns.onchange = renderTable; + renderTable(); + const panels = []; + if (trajectories.lanes.length) { + panels.push( + createAnatomy({ + lanes: trajectories.lanes, + consumeTs: trajectories.consume_ts, + rowsByIndex: new Map(rows.map((r) => [r.sample_index, r])), + onClickSample: (index) => openTokens({ sample_index: index }), + }), + ); + } + return [ + ...panels, + el("div", { class: "panel" }, [el("h3", {}, [`${rewardKey} vs response length (red = truncated)`]), scatter]), + el("div", { class: "panel" }, [ + el("h3", {}, ["samples — click a row for the token view"]), + el("div", { class: "controls" }, [el("label", {}, [allColumns, " all columns"])]), + tableHolder, + ]), + ]; + }; + + const groupsTab = () => [ + el("div", { class: "panel" }, [ + el("h3", {}, ["GRPO groups (red = zero reward std → no gradient signal)"]), + sortableTable(groups.rows, groups.columns, { + flagRow: (g) => g.zero_std, + sortState: { column: "group_index", desc: false }, + }), + ]), + ]; + + const tabBody = el("div", {}); + const tabButtons = el("div", { class: "tabs" }); + const tabs = { samples: samplesTab, groups: groupsTab }; + const selectTab = (name) => { + tabBody.replaceChildren(...tabs[name]()); + tabButtons.replaceChildren( + ...Object.keys(tabs).map((t) => + el("button", { class: t === name ? "active" : "", onclick: () => selectTab(t) }, [t]), + ), + ); + }; + selectTab("samples"); + + view.replaceChildren(controls, stats, tabButtons, tabBody); +} diff --git a/miles/dashboard/static/views_timeline.js b/miles/dashboard/static/views_timeline.js new file mode 100644 index 0000000000..7b5cbf8a55 --- /dev/null +++ b/miles/dashboard/static/views_timeline.js @@ -0,0 +1,652 @@ +import { api } from "./api.js"; +import { el, fmtNum, setViewCleanup } from "./app.js"; +import { createCarpet } from "./carpet.js"; +import { hideTooltip, showTooltip } from "./charts.js"; + +const PHASE_COLORS = { + initialize: "#b8c4d8", + rollout: "#1d9e75", + eval_rollout: "#5dcaa5", + actor_train: "#7f77dd", + train_log_probs: "#afa9ec", + log_probs: "#afa9ec", + ref_log_probs: "#9f97e0", + data_preprocess: "#8b95a5", + train_wait: "#e4e7ec", + update_weights: "#d85a30", + ref_model_update: "#b98156", + save_model: "#8a6d3b", + sleep: "#d6dbe3", + wake_up: "#ba7517", +}; +const DEFAULT_PHASE_COLOR = "#98a1b0"; +const OVERLAY_METRICS = [ + "sglang_num_running_reqs", + "sglang_gen_throughput", + "sglang_token_usage", + "sglang_cache_hit_rate", +]; +const OVERLAY_COLOR = "#d97706"; +const MEM_COLOR = "#1a9e6e"; +const UTIL_STROKE = "rgba(47, 111, 235, 0.9)"; +const UTIL_FILL = "rgba(47, 111, 235, 0.14)"; +const TEXT = "#24292f"; +const MUTED = "#667080"; +const GRID = "#e3e7ee"; + +const FOLLOW_REDRAW_MS = 5000; +const LANE_CAP = 32; +const DEFAULT_LANES = 8; // no-selection default: evenly spaced across the cluster +const QUICK_PICKS = [ + { label: "pick: lowest util", criterion: "lowest_util" }, + { label: "pick: slowest update_weights", criterion: "slowest_phase:update_weights" }, +]; +const QUICK_PICK_K = 8; + +const LANE_H = 66; +const PHASE_H = 14; +const UTIL_H = 40; +const M_LEFT = 96; +const M_TOP = 24; +const M_RIGHT = 14; + +export async function renderTimeline(view, meta, route) { + // mutable data state, reloaded on every follow-mode refresh + let selection = route?.lanes || null; // lane-selection grammar string + let autoDefault = !selection; // seed the spaced default once topology is known + let resolvedKeys = null; // full "node:gpu" set the grammar resolved to (null = no pool) + let capped = 0; // lanes hidden by LANE_CAP + let lanes = []; + let windows = []; + let phasesByLane = new Map(); + let bubbles = []; + let gpu = {}; + let engineSeries = []; + let T0 = 0; + let T1 = 1; + let v0 = 0; + let v1 = 1; + let haveData = false; + let overlayMetric = meta.capabilities.has_engine_series ? OVERLAY_METRICS[0] : null; + let showMem = false; + let multiNode = false; + const laneKey = (l) => `${l.node}:${l.gpu}`; + // hard viewport cap (design §17): the lane view never displays more than + // MAXW at once, so every data fetch is O(window), never O(run) + const MAXW = meta.capabilities.max_window_s ?? Infinity; + let allLanes = []; // full topology lane list (seeds the spaced default selection) + let fetched = null; // {t0, t1} bounds the loaded phases/gpu/engine data covers + + // T0/T1 come from meta.time_range (an O(1) edge-stamp read server-side), + // not from scanning fetched data. Keeps the trailing window pinned to the + // growing run unless the user has panned away. + function applyRange(range) { + if (!range) { + haveData = false; + return; + } + const trailing = !haveData || v1 >= T1 - 1e-6; + const span = haveData ? v1 - v0 : Math.min(MAXW, range[1] - range[0] || 1); + [T0, T1] = range; + if (trailing) { + v1 = T1; + v0 = Math.max(T0, T1 - Math.min(span, MAXW)); + } + if (v1 - v0 < 1) v0 = Math.max(T0, v1 - 1); + haveData = true; + } + + async function loadData() { + const [topology, bubblesRes] = await Promise.all([ + api("/api/timeline/topology"), + api("/api/timeline/bubbles"), + ]); + allLanes = topology.lanes; + windows = topology.windows; + bubbles = bubblesRes.bubbles; + if (!haveData) { + lanes = []; + gpu = {}; + phasesByLane = new Map(); + return; + } + + // fetch a margin beyond the viewport so small pans redraw without a reload + const span = v1 - v0; + const margin = Math.max(0, Math.min(span / 4, (MAXW - span) / 2)); + const f0 = Math.max(T0, v0 - margin); + const f1 = Math.min(T1, v1 + margin); + // first sight of the topology: DEFAULT_LANES evenly spaced across the + // global index range (8 gpus -> all, 64 -> g0,g8,...,g56) becomes a REAL + // selection — visible chips the user removes/extends/overwrites, no + // parallel "default budget" state to reason about + if (autoDefault && allLanes.length > 0) { + autoDefault = false; + if (allLanes.length > DEFAULT_LANES) { + selection = Array.from( + { length: DEFAULT_LANES }, + (_, i) => `g:${allLanes[Math.floor((i * allLanes.length) / DEFAULT_LANES)].index}`, + ).join(", "); + history.replaceState(null, "", `#/timeline?lanes=${encodeURIComponent(selection)}`); + } + } + const [phasesRes, gpuRes] = await Promise.all([ + api("/api/timeline/phases", { t0: f0, t1: f1, lanes: selection }), + api("/api/timeline/gpu", { t0: f0, t1: f1, max_points: 4000, lanes: selection }), + ]); + engineSeries = overlayMetric + ? (await api("/api/timeline/engine_series", { metric: overlayMetric, t0: f0, t1: f1, max_points: 4000 })) + .series + : []; + fetched = { t0: f0, t1: f1 }; + lanes = allLanes; + resolvedKeys = null; + if (selection) { + // the filtered responses reveal which lanes the grammar resolved to + const keys = new Set(Object.keys(gpuRes.lanes)); + for (const p of phasesRes.phases) keys.add(`${p.node}:${p.gpu}`); + lanes = lanes.filter((l) => keys.has(laneKey(l))); + resolvedKeys = new Set(lanes.map(laneKey)); + } + const total = lanes.length; + if (lanes.length > LANE_CAP) lanes = lanes.slice(0, LANE_CAP); + capped = Math.max(0, total - lanes.length); + gpu = gpuRes.lanes; + multiNode = new Set(lanes.map((l) => l.node)).size > 1; + phasesByLane = new Map(lanes.map((l) => [laneKey(l), []])); + for (const p of phasesRes.phases) phasesByLane.get(`${p.node}:${p.gpu}`)?.push(p); + } + + // pan/zoom past the fetched bounds: debounced windowed reload + let reloadTimer = null; + function afterViewChange() { + draw(); + syncCarpet(); + if (!fetched || v0 < fetched.t0 - 1e-6 || v1 > fetched.t1 + 1e-6) { + clearTimeout(reloadTimer); + reloadTimer = setTimeout(async () => { + try { + await loadData(); + } catch { + return; // transient fetch failure — the next interaction retries + } + renderAll(); + }, 300); + } + } + + // external engines (no miles actor) have unknown GPU placement (gpus=[]): + // fall back to host identity — the engine's addr host IS its node + const engineHost = (addr) => addr.split("//").pop().split(":")[0]; + const engineAt = (node, gpuId, t) => { + for (const w of windows) { + if (t >= w.t0 && (w.t1 === null || t < w.t1)) { + for (const e of w.engines) { + const match = e.gpus.length + ? e.gpus.some(([n, g]) => n === node && g === gpuId) + : engineHost(e.addr) === node; + if (match) return e.addr; + } + } + } + return null; + }; + + // ------------------------------- toolbar ---------------------------------- + const toolbar = el("div", { class: "controls" }); + const overlayScale = el("span", { style: `color: ${OVERLAY_COLOR}; font-size: 12px` }); + const followBadge = el("span", { class: "muted" }); + const renderToolbar = () => { + const chips = OVERLAY_METRICS.map((m) => + el( + "button", + { + class: m === overlayMetric ? "active" : "", + onclick: async () => { + overlayMetric = m === overlayMetric ? null : m; + renderToolbar(); + await loadData(); + renderAll(); + }, + }, + [m.replace("sglang_", "")], + ), + ); + const memBtn = el( + "button", + { class: showMem ? "active" : "", onclick: () => ((showMem = !showMem), renderToolbar(), draw()) }, + ["mem"], + ); + toolbar.replaceChildren( + el("span", { class: "muted" }, ["overlay"]), + ...(meta.capabilities.has_engine_series ? chips : [el("span", { class: "muted" }, ["(no engine series)"])]), + memBtn, + el( + "button", + { onclick: () => ((v1 = T1), (v0 = Math.max(T0, T1 - MAXW)), afterViewChange()) }, + ["latest window"], + ), + overlayScale, + followBadge, + el("span", { class: "muted" }, [`wheel = zoom · drag = pan · window ≤ ${Math.round(MAXW / 3600)}h`]), + ); + }; + + // ----------------------------- lane selection ----------------------------- + const selRow = el("div", { class: "controls" }); + const selError = el("span", { class: "error", style: "padding: 0" }); + const selTerms = () => (selection ? selection.split(",").map((s) => s.trim()).filter(Boolean) : []); + + async function setSelection(grammar) { + selection = grammar || null; + history.replaceState(null, "", selection ? `#/timeline?lanes=${encodeURIComponent(selection)}` : "#/timeline"); + selError.textContent = ""; + renderSelection(); + try { + await loadData(); + } catch (err) { + selError.textContent = String(err); + return; + } + renderAll(); + carpet.redraw(); // selection markers + } + + function renderSelection() { + const input = el("input", { + type: "text", + placeholder: "add: g:0-31 · rank:0-7 · node: · gpu:: · engine: · role:train", + style: "flex: 1; min-width: 280px", + }); + input.onkeydown = (ev) => { + if (ev.key === "Enter" && input.value.trim()) { + setSelection([...selTerms(), input.value.trim()].join(", ")); + } + }; + selRow.replaceChildren( + el("span", { class: "muted" }, ["lanes"]), + ...selTerms().map((term) => + el("span", { class: "chip" }, [ + term, + el( + "button", + { onclick: () => setSelection(selTerms().filter((t) => t !== term).join(", ")) }, + ["×"], + ), + ]), + ), + ...(selection ? [] : [el("span", { class: "muted" }, ["all"])]), + input, + ...QUICK_PICKS.map(({ label, criterion }) => + el( + "button", + { + onclick: async () => { + const res = await api("/api/timeline/outliers", { criterion, top_k: QUICK_PICK_K }); + if (res.outliers.length) setSelection(res.outliers.map((o) => `gpu:${o.node}:${o.gpu}`).join(", ")); + }, + }, + [label], + ), + ), + ...(selection ? [el("button", { onclick: () => setSelection(null) }, ["clear"])] : []), + ...(capped + ? [el("span", { class: "warn" }, [`${capped + lanes.length} lanes; showing first ${lanes.length}`])] + : []), + selError, + ); + } + + // -------------------------------- carpet ---------------------------------- + const carpet = createCarpet({ + phaseColors: PHASE_COLORS, + runStart: () => T0, + selectedKeys: () => resolvedKeys ?? new Set(), + // toggle semantics: brushed lanes join the detailed pool; a brush that only + // covers already-pooled lanes removes them instead + onBrush: (keys) => { + const pool = resolvedKeys ?? new Set(); + if (keys.every((k) => pool.has(k))) { + const kept = [...pool].filter((k) => !keys.includes(k)); + setSelection(kept.map((k) => `gpu:${k}`).join(", ")); + } else { + const added = keys.filter((k) => !pool.has(k)).map((k) => `gpu:${k}`); + setSelection([...selTerms(), ...added].join(", ")); + } + }, + }); + // the carpet tracks the lane view's zoom window; before any data exists it + // asks for the (empty) server default + const carpetRange = () => (haveData ? { t0: v0, t1: v1 } : {}); + let carpetTimer = null; + const syncCarpet = () => { + clearTimeout(carpetTimer); + carpetTimer = setTimeout(() => carpet.refresh(carpetRange()).catch(() => {}), 250); + }; + + // ----------------------------- bubble strip ------------------------------- + const bubbleStrip = el("div", { class: "bubblestrip" }); + const renderBubbles = () => { + if (!bubbles.length) { + bubbleStrip.replaceChildren(); + return; + } + const worst = Math.max(...bubbles.map((b) => b.wait_ratio ?? 0), 0.01); + bubbleStrip.replaceChildren(el("span", { class: "muted" }, ["wait ratio / step: "])); + for (const b of bubbles) { + const cell = el("span", { class: "bubble" }, [String(b.step)]); + const ratio = b.wait_ratio ?? 0; + cell.style.background = `rgba(224, 96, 96, ${(0.85 * ratio) / worst})`; + cell.onclick = () => { + if (b.step_time !== null) { + v1 = b.ts; + v0 = Math.max(b.ts - b.step_time, b.ts - MAXW); + afterViewChange(); + } + }; + cell.onmousemove = (ev) => + showTooltip( + ev.clientX, + ev.clientY, + `step ${b.step}\nwait_ratio = ${fmtNum(ratio)}\nstep_time = ${fmtNum(b.step_time)}s`, + ); + cell.onmouseleave = hideTooltip; + bubbleStrip.append(cell); + } + }; + + // -------------------------------- canvas ---------------------------------- + const canvas = el("canvas", { class: "timeline" }); + const overlayMax = () => Math.max(...engineSeries.flatMap((s) => s.value), 1e-9); + const memMax = () => Math.max(...Object.values(gpu).flatMap((s) => s.mem_mb), 1); + + function draw() { + canvas.style.height = `${M_TOP + Math.max(lanes.length, 1) * LANE_H + 8}px`; + const dpr = window.devicePixelRatio || 1; + const rect = canvas.getBoundingClientRect(); + canvas.width = rect.width * dpr; + canvas.height = rect.height * dpr; + const ctx = canvas.getContext("2d"); + ctx.scale(dpr, dpr); + const width = rect.width; + const plotW = width - M_LEFT - M_RIGHT; + ctx.font = "11px ui-monospace, monospace"; + if (!haveData) { + ctx.fillStyle = MUTED; + ctx.fillText("waiting for timeline data…", M_LEFT, M_TOP + 20); + overlayScale.textContent = ""; + return; + } + const X = (t) => M_LEFT + ((t - v0) / (v1 - v0)) * plotW; + + ctx.fillStyle = MUTED; + ctx.strokeStyle = GRID; + const span = v1 - v0; + const tickStep = [1, 2, 5, 10, 30, 60, 120, 300, 600, 1800, 3600].find((s) => span / s <= 10) || 7200; + for (let t = Math.ceil((v0 - T0) / tickStep) * tickStep + T0; t <= v1; t += tickStep) { + const rel = Math.round(t - T0); + const label = `+${Math.floor(rel / 60)}:${String(rel % 60).padStart(2, "0")}`; + ctx.fillText(label, X(t) - 12, 12); + ctx.beginPath(); + ctx.moveTo(X(t), M_TOP - 6); + ctx.lineTo(X(t), M_TOP + lanes.length * LANE_H); + ctx.stroke(); + } + + lanes.forEach((lane, i) => { + const key = laneKey(lane); + const yPhase = M_TOP + i * LANE_H + 4; + const yUtilTop = yPhase + PHASE_H + 4; + const yUtilBot = yUtilTop + UTIL_H; + + ctx.fillStyle = TEXT; + ctx.fillText(`g${lane.index}`, 8, yUtilTop + 4); + if (multiNode) { + ctx.fillStyle = MUTED; + ctx.fillText(`${lane.node}:${lane.gpu}`, 8, yUtilTop + 18); + } + ctx.strokeStyle = GRID; + ctx.beginPath(); + ctx.moveTo(M_LEFT, yUtilBot); + ctx.lineTo(width - M_RIGHT, yUtilBot); + ctx.stroke(); + + for (const p of phasesByLane.get(key) ?? []) { + if (p.t1 < v0 || p.t0 > v1) continue; + const x0 = Math.max(X(p.t0), M_LEFT); + const x1 = Math.min(X(p.t1), width - M_RIGHT); + ctx.fillStyle = PHASE_COLORS[p.name] ?? DEFAULT_PHASE_COLOR; + ctx.fillRect(x0, yPhase, Math.max(x1 - x0, 1), PHASE_H); + if (p.name === "train_wait") { + // hatch idle explicitly: bubbles must be visible, not blank + ctx.strokeStyle = MUTED; + ctx.save(); + ctx.beginPath(); + ctx.rect(x0, yPhase, x1 - x0, PHASE_H); + ctx.clip(); + for (let x = x0 - PHASE_H; x < x1; x += 5) { + ctx.beginPath(); + ctx.moveTo(x, yPhase + PHASE_H); + ctx.lineTo(x + PHASE_H, yPhase); + ctx.stroke(); + } + ctx.restore(); + } + } + + const series = gpu[key]; + if (series && series.ts.length) { + const yOf = (u) => yUtilBot - (u / 100) * UTIL_H; + ctx.beginPath(); + let started = false; + for (let j = 0; j < series.ts.length; j++) { + const t = series.ts[j]; + if (t < v0 || t > v1) continue; + const x = X(t); + const y = yOf(series.util[j]); + started ? ctx.lineTo(x, y) : ctx.moveTo(x, y); + started = true; + } + ctx.strokeStyle = UTIL_STROKE; + ctx.lineWidth = 1; + ctx.stroke(); + ctx.lineTo(X(Math.min(series.ts[series.ts.length - 1], v1)), yUtilBot); + ctx.lineTo(X(Math.max(series.ts[0], v0)), yUtilBot); + ctx.closePath(); + ctx.fillStyle = UTIL_FILL; + ctx.fill(); + + if (showMem) { + const scale = memMax(); + ctx.beginPath(); + started = false; + for (let j = 0; j < series.ts.length; j++) { + const t = series.ts[j]; + if (t < v0 || t > v1) continue; + const y = yUtilBot - (series.mem_mb[j] / scale) * UTIL_H; + started ? ctx.lineTo(X(t), y) : ctx.moveTo(X(t), y); + started = true; + } + ctx.strokeStyle = MEM_COLOR; + ctx.stroke(); + } + } + + if (overlayMetric && engineSeries.length) { + const scale = overlayMax(); + for (const w of windows) { + const w0 = Math.max(w.t0, v0); + const w1 = Math.min(w.t1 ?? v1, v1); + if (w0 >= w1) continue; + const addr = engineAt(lane.node, lane.gpu, (w0 + w1) / 2); + const s = engineSeries.find((es) => es.addr === addr); + if (!s) continue; + ctx.beginPath(); + let on = false; + for (let j = 0; j < s.ts.length; j++) { + const t = s.ts[j]; + if (t < w0 || t > w1) continue; + const y = yUtilBot - (s.value[j] / scale) * UTIL_H; + on ? ctx.lineTo(X(t), y) : ctx.moveTo(X(t), y); + on = true; + } + ctx.strokeStyle = OVERLAY_COLOR; + ctx.lineWidth = 1.6; + ctx.stroke(); + ctx.lineWidth = 1; + } + } + }); + + overlayScale.textContent = + overlayMetric && engineSeries.length + ? `${overlayMetric.replace("sglang_", "")} scale 0–${fmtNum(overlayMax())}` + : ""; + } + + // freshness: the newest data timestamp (server-side T1), not the browser + // fetch time — a stalled collector must READ as stale despite live redraws + const renderFreshness = () => { + const stamp = haveData ? new Date(T1 * 1000).toLocaleTimeString() : "…"; + followBadge.textContent = meta.mode === "follow" ? `live · data → ${stamp}` : `data → ${stamp}`; + }; + + function renderAll() { + renderFreshness(); + renderSelection(); + renderBubbles(); + renderLegend(); + draw(); + } + + // ------------------------------ interactions ------------------------------ + const timeAt = (clientX) => { + const rect = canvas.getBoundingClientRect(); + return v0 + ((clientX - rect.left - M_LEFT) / (rect.width - M_LEFT - M_RIGHT)) * (v1 - v0); + }; + canvas.onwheel = (ev) => { + if (!haveData) return; + ev.preventDefault(); + const pivot = timeAt(ev.clientX); + const factor = Math.exp(ev.deltaY * 0.002); + v0 = Math.max(T0, pivot + (v0 - pivot) * factor); + v1 = Math.min(T1, pivot + (v1 - pivot) * factor); + if (v1 - v0 > MAXW) { + // zoom-out stops at the viewport cap, anchored at the cursor + const anchor = (pivot - v0) / (v1 - v0); + v0 = Math.max(T0, pivot - anchor * MAXW); + v1 = v0 + MAXW; + } + afterViewChange(); + }; + let dragFrom = null; + canvas.onmousedown = (ev) => (dragFrom = { x: ev.clientX, v0, v1 }); + const onMouseUp = () => (dragFrom = null); + window.addEventListener("mouseup", onMouseUp); + canvas.ondblclick = () => ((v1 = T1), (v0 = Math.max(T0, T1 - MAXW)), afterViewChange()); + canvas.onmousemove = (ev) => { + if (!haveData) return; + if (dragFrom) { + const rect = canvas.getBoundingClientRect(); + const dt = ((dragFrom.x - ev.clientX) / (rect.width - M_LEFT - M_RIGHT)) * (dragFrom.v1 - dragFrom.v0); + const span = dragFrom.v1 - dragFrom.v0; + v0 = Math.max(T0, Math.min(dragFrom.v0 + dt, T1 - span)); + v1 = v0 + span; + afterViewChange(); + return; + } + const rect = canvas.getBoundingClientRect(); + const laneIdx = Math.floor((ev.clientY - rect.top - M_TOP) / LANE_H); + if (laneIdx < 0 || laneIdx >= lanes.length || ev.clientX - rect.left < M_LEFT) { + hideTooltip(); + return; + } + const lane = lanes[laneIdx]; + const key = laneKey(lane); + const t = timeAt(ev.clientX); + const lines = [`g${lane.index} ${key} +${fmtNum(t - T0)}s`]; + const phase = (phasesByLane.get(key) ?? []).find((p) => p.t0 <= t && t < p.t1); + if (phase) lines.push(`phase: ${phase.name}${phase.rank >= 0 ? ` (rank ${phase.rank})` : ""}`); + const series = gpu[key]; + if (series && series.ts.length) { + let best = 0; + for (let j = 0; j < series.ts.length; j++) { + if (Math.abs(series.ts[j] - t) < Math.abs(series.ts[best] - t)) best = j; + } + lines.push(`util: ${series.util[best]}% mem: ${fmtNum(series.mem_mb[best] / 1024)}G`); + } + if (overlayMetric) { + const addr = engineAt(lane.node, lane.gpu, t); + const s = engineSeries.find((es) => es.addr === addr); + if (addr && s && s.ts.length) { + let best = 0; + for (let j = 0; j < s.ts.length; j++) if (Math.abs(s.ts[j] - t) < Math.abs(s.ts[best] - t)) best = j; + lines.push(`${overlayMetric.replace("sglang_", "")}: ${fmtNum(s.value[best])}`, `engine: ${addr}`); + } + } + showTooltip(ev.clientX, ev.clientY, lines.join("\n")); + }; + canvas.onmouseleave = hideTooltip; + + // -------------------------------- legend ---------------------------------- + const legendPanel = el("div", { class: "panel" }); + const renderLegend = () => { + const present = new Set(); + for (const items of phasesByLane.values()) for (const p of items) present.add(p.name); + legendPanel.replaceChildren( + el("h3", {}, ["legend"]), + el("div", { class: "legend" }, [ + ...[...present].sort().map((name) => { + const swatch = el("span", { + class: "bar", + style: `width: 12px; background: ${PHASE_COLORS[name] ?? DEFAULT_PHASE_COLOR}`, + }); + return el("span", { style: "display: inline-flex; gap: 4px; align-items: center" }, [swatch, name]); + }), + el("span", { style: `color: ${OVERLAY_COLOR}` }, ["— engine overlay"]), + el("span", { style: `color: ${UTIL_STROKE}` }, ["— gpu util"]), + ]), + ); + }; + + view.replaceChildren(toolbar, selRow, carpet.root, bubbleStrip, el("div", { class: "panel" }, [canvas]), legendPanel); + renderToolbar(); + applyRange(meta.time_range); + await Promise.all([loadData(), carpet.refresh(carpetRange())]); + renderAll(); + carpet.redraw(); // selection markers need the loaded lane list + const onResize = () => { + draw(); + carpet.redraw(); + }; + window.addEventListener("resize", onResize); + + // ------------------------- follow-mode auto refresh ----------------------- + let refreshing = false; + let intervalId = null; + if (meta.mode === "follow") { + intervalId = setInterval(async () => { + if (refreshing) return; + refreshing = true; + try { + // fresh global range (cheap edge-stamp read); getMeta() is cached + applyRange((await api("/api/meta")).time_range); + await loadData(); + renderAll(); + await carpet.refresh(carpetRange()); + } catch { + // transient fetch failure (server restart, network blip): keep polling + } finally { + refreshing = false; + } + }, FOLLOW_REDRAW_MS); + } + setViewCleanup(() => { + if (intervalId !== null) clearInterval(intervalId); + clearTimeout(carpetTimer); + clearTimeout(reloadTimer); + carpet.destroy(); + window.removeEventListener("resize", onResize); + window.removeEventListener("mouseup", onMouseUp); + }); +} diff --git a/miles/dashboard/static/views_tokens.js b/miles/dashboard/static/views_tokens.js new file mode 100644 index 0000000000..b288c153bd --- /dev/null +++ b/miles/dashboard/static/views_tokens.js @@ -0,0 +1,249 @@ +import { createAnatomy } from "./anatomy.js"; +import { api } from "./api.js"; +import { el, fmtNum } from "./app.js"; +import { divergingColor, drawChart, hideTooltip, sequentialColor, showTooltip } from "./charts.js"; +import { renderConversation } from "./conversation.js"; + +const WINDOW_SIZES = [256, 1024, 4096]; + +// stat -> how to color a value: diverging stats define center/scale via values +const STATS = { + imp_ratio: { color: (v) => divergingColor(Math.log(Math.max(v, 1e-6))) }, + lp_diff: { diverging: true }, + advantages: { diverging: true }, + returns: { diverging: true }, + entropy: { sequential: true }, + ref_entropy: { sequential: true }, + train_log_probs: { sequential: true, negate: true }, + rollout_log_probs: { sequential: true, negate: true }, + ref_log_probs: { sequential: true, negate: true }, +}; + +function colorFor(stat, values) { + const spec = STATS[stat]; + if (spec.color) return spec.color; + const finite = values.filter((v) => v !== null && Number.isFinite(v)); + if (spec.diverging) { + const scale = Math.max(...finite.map(Math.abs), 1e-9); + return (v) => divergingColor(v / scale); + } + const transformed = spec.negate ? finite.map((v) => -v) : finite; + const [lo, hi] = [Math.min(...transformed), Math.max(...transformed)]; + return (v) => sequentialColor(((spec.negate ? -v : v) - lo) / Math.max(hi - lo, 1e-9)); +} + +export async function renderTokens(view, meta, route) { + const { rolloutId, sampleIndex, evaluation } = route; + view.replaceChildren(el("p", { class: "muted" }, ["loading sample…"])); + + // cheap panels first: the lifecycle lane (telemetry) and the conversation + // (trajectory sidecar); the token machinery costs a full dump load plus + // detokenize, so it only starts when its tab is opened + const panels = []; + if (!evaluation) { + try { + const trajectories = await api(`/api/rollout/${rolloutId}/trajectories`, { sample_index: sampleIndex }); + if (trajectories.lanes.length) { + panels.push( + createAnatomy({ + lanes: trajectories.lanes, + consumeTs: trajectories.consume_ts, + rowsByIndex: new Map(), + onClickSample: () => {}, + }), + ); + } + } catch { + /* endpoint absent or no events: token view stands alone */ + } + } + + let conversationRow = null; + try { + conversationRow = await api(`/api/rollout/${rolloutId}/sample/${sampleIndex}/messages`, { eval: evaluation }); + } catch (err) { + if (!String(err).includes("404")) throw err; // 404 = run recorded no conversation + } + + const tokensPane = el("div"); + let tokensStarted = false; + const startTokens = () => { + if (tokensStarted) return; + tokensStarted = true; + tokensPane.replaceChildren( + el("p", { class: "muted" }, [ + "loading tokens… the first open of a step detokenizes its whole dump and can take several minutes", + ]), + ); + loadTokensPane(tokensPane, rolloutId, sampleIndex, evaluation).catch((err) => { + tokensPane.replaceChildren(el("div", { class: "error" }, [String(err)])); + }); + }; + + if (conversationRow === null) { + view.replaceChildren(...panels, tokensPane); + startTokens(); + return; + } + + const conversationPane = renderConversation(conversationRow); + const tabs = el("div", { class: "tabs" }); + const body = el("div"); + const select = (name) => { + tabs.replaceChildren( + ...["conversation", "tokens"].map((tab) => + el("button", { class: tab === name ? "active" : "", onclick: () => select(tab) }, [tab]), + ), + ); + body.replaceChildren(name === "conversation" ? conversationPane : tokensPane); + if (name === "tokens") startTokens(); + }; + view.replaceChildren(...panels, tabs, body); + select("conversation"); +} + +async function loadTokensPane(root, rolloutId, sampleIndex, evaluation) { + const probe = await api(`/api/rollout/${rolloutId}/sample/${sampleIndex}/tokens`, { + start: 0, + end: 1, + eval: evaluation, + }); + const total = probe.total_len; + + let windowSize = total <= 4096 ? total : 1024; + let start = 0; + let stat = "imp_ratio"; + let chartMetric = "imp_ratio"; + let chartData = null; // full-range payload: the chart spans the whole response + + const controls = el("div", { class: "controls" }); + const strip = el("div", { class: "panel" }); + const chartPanel = el("div", { class: "panel" }); + const chartCanvas = el("canvas", { class: "chart" }); // persistent: zoom survives metric switches + + function renderChart() { + const available = Object.keys(STATS).filter((s) => chartData[s] !== null && chartData[s] !== undefined); + if (!available.length) { + chartPanel.replaceChildren(); + return; + } + if (!available.includes(chartMetric)) chartMetric = available[0]; + const chips = available.map((s) => + el( + "button", + { + class: s === chartMetric ? "active" : "", + onclick: () => { + chartMetric = s; + if (chartCanvas._zoom) chartCanvas._zoom = { x: chartCanvas._zoom.x, y: null }; + renderChart(); + }, + }, + [s], + ), + ); + chartPanel.replaceChildren( + el("h3", {}, ["per-token metrics"]), + el("div", { class: "controls" }, [ + ...chips, + el("span", { class: "muted" }, ["drag = zoom to token range · double-click = reset"]), + ]), + chartCanvas, + ); + const values = chartData[chartMetric]; + const points = []; + values.forEach((v, i) => { + if (v === null) return; + const pos = chartData.prompt_len + i; + points.push({ x: pos, y: v, label: `pos ${pos}\n${chartMetric} = ${fmtNum(v)}` }); + }); + queueMicrotask(() => drawChart(chartCanvas, points)); + } + + async function loadChart() { + chartData = await api(`/api/rollout/${rolloutId}/sample/${sampleIndex}/tokens`, { + start: 0, + end: total, + eval: evaluation, + }); + renderChart(); + } + + async function load() { + const payload = await api(`/api/rollout/${rolloutId}/sample/${sampleIndex}/tokens`, { + start, + end: Math.min(start + windowSize, total), + eval: evaluation, + }); + const available = Object.keys(STATS).filter((s) => payload[s] !== null && payload[s] !== undefined); + if (!available.includes(stat)) stat = available[0]; + + // ---- controls ---- + const statSelect = el("select", {}, available.map((s) => + Object.assign(el("option", { value: s }, [s]), { selected: s === stat }), + )); + statSelect.onchange = () => { + stat = statSelect.value; + load(); + }; + const sizeSelect = el("select", {}, WINDOW_SIZES.filter((w) => w < total).concat([total]).map((w) => + Object.assign(el("option", { value: w }, [w === total ? `all ${total}` : String(w)]), { + selected: w === windowSize, + }), + )); + sizeSelect.onchange = () => { + windowSize = Number(sizeSelect.value); + load(); + }; + const startInput = el("input", { type: "number", value: start, min: 0, max: total - 1 }); + startInput.onchange = () => { + start = Math.max(0, Math.min(Number(startInput.value), total - 1)); + load(); + }; + controls.replaceChildren( + el("button", { onclick: () => ((start = Math.max(0, start - windowSize)), load()) }, ["◀"]), + el("span", {}, [`tokens ${payload.start}–${payload.end} of ${total} (prompt ${payload.prompt_len})`]), + el("button", { onclick: () => ((start = Math.min(total - 1, start + windowSize)), load()) }, ["▶"]), + el("span", { class: "muted" }, [" window"]), + sizeSelect, + el("span", { class: "muted" }, [" start"]), + startInput, + el("span", { class: "muted" }, [" color by"]), + statSelect, + ); + + // ---- token strip ---- + const values = payload[stat] ?? []; + const color = values.length ? colorFor(stat, values) : () => "transparent"; + const spans = payload.token_ids.map((tokenId, i) => { + const respPos = i - payload.response_offset; // index into stat arrays + const inResponse = respPos >= 0 && respPos < (payload.train_log_probs ?? payload.rollout_log_probs ?? []).length; + const text = payload.token_text ? payload.token_text[i] : `·${tokenId}`; + const masked = inResponse && payload.loss_mask !== null && payload.loss_mask[respPos] === 0; + const value = inResponse && values.length ? values[respPos] : null; + const span = el("span", { class: `tok ${inResponse ? "" : "prompt"} ${masked ? "masked" : ""}` }, [ + text.replaceAll("\n", "⏎\n"), + ]); + if (value !== null && value !== undefined && !masked) span.style.background = color(value); + span.onmousemove = (ev) => { + const lines = [`#${payload.start + i} id=${tokenId} ${JSON.stringify(text)}`]; + if (!inResponse) lines.push("(prompt)"); + else { + for (const s of available) lines.push(`${s} = ${fmtNum(payload[s][respPos])}`); + if (masked) lines.push("loss_mask = 0"); + } + showTooltip(ev.clientX, ev.clientY, lines.join("\n")); + }; + span.onmouseleave = hideTooltip; + return span; + }); + strip.replaceChildren( + el("h3", {}, [`tokens — colored by ${stat}${payload.token_text ? "" : " (no tokenizer: ids shown)"}`]), + el("div", { class: "tokens" }, spans), + ); + + } + + root.replaceChildren(controls, strip, chartPanel); + await Promise.all([load(), loadChart()]); +} diff --git a/miles/dashboard/store.py b/miles/dashboard/store.py new file mode 100644 index 0000000000..931a180014 --- /dev/null +++ b/miles/dashboard/store.py @@ -0,0 +1,1232 @@ +"""File-backed store for miles dashboard time series. + +The live collector buffers records via ``append()`` and persists them with +``flush()`` as append-only JSONL streams under ``{dump_details}/dashboard/``. +The dashboard server reads the same directory with ``MetricStore.load()`` and +picks up appended lines incrementally with ``follow()`` — a plain byte-offset +tail, correct because every stream is append-only. + +Each stream holds one record type (one JSON object per line); see the +``Record`` subclasses for the schemas. + +The two high-rate streams (``gpu_util``, ``engine_series``) are held in memory +as columnar polars frames instead of dataclass lists (design doc §17): +~16 B/row instead of ~600 B, vectorized parsing and numpy queries. + +High-rate streams (those two plus ``phases``) write hourly partition files +``{stream}/{YYYYMMDD_HH}.jsonl`` (phases keyed by END hour) and parse lazily: +open reads nothing from them; a windowed query parses only the hours it +touches, through an LRU block cache with per-file byte-offset tail refresh. +A v1 flat ``{stream}.jsonl`` still reads as one "legacy" partition. Lane +metadata for selection resolution is folded into ``lane_catalog.json`` at +flush time so it never needs a stream scan. +""" + +from __future__ import annotations + +import calendar +import io +import json +import os +import time +from collections import OrderedDict +from collections.abc import Callable +from dataclasses import asdict, dataclass, replace +from pathlib import Path +from typing import Any, ClassVar + +try: + from enum import StrEnum +except ImportError: + from backports.strenum import StrEnum + +import numpy as np +import polars as pl + + +class Stream(StrEnum): + METRICS = "metrics" + PHASES = "phases" + TOPOLOGY = "topology" + GPU_UTIL = "gpu_util" + ENGINE_SERIES = "engine_series" + TRAJECTORIES = "trajectories" + + @property + def filename(self) -> str: + return f"{self.value}.jsonl" + + +@dataclass +class Record: + stream: ClassVar[Stream] + + def to_dict(self) -> dict: + return asdict(self) + + @classmethod + def from_dict(cls, data: dict) -> Record: + return cls(**data) + + def timestamps(self) -> tuple[float, ...]: + return (self.ts,) + + +@dataclass +class MetricsRecord(Record): + """One ``tracking_utils.log()`` payload.""" + + stream: ClassVar[Stream] = Stream.METRICS + ts: float + step_key: str | None + step: int | None + metrics: dict[str, Any] + + +class Role(StrEnum): + TRAIN = "train" + ROLLOUT_MANAGER = "rollout_manager" + DERIVED = "derived" # synthesized by the read side, not observed + + +# phase name synthesized per lane for [meta.start_ts, first observed event) +INITIALIZE_PHASE = "initialize" + + +@dataclass +class PhaseEvent(Record): + """One completed timer interval on one process (Timer sink event). + + ``rollout_manager`` events carry no GPUs (the manager is a driver-side + process); the timeline queries expand them onto rollout-engine lanes.""" + + stream: ClassVar[Stream] = Stream.PHASES + name: str + t0: float + t1: float + node: str + gpus: list[int] + rank: int + role: str # a Role value + + # t1 sentinel: the interval was still running when this event was written + # (Timer.start emits it so long phases are visible before they end); the + # closing event with the real t1 supersedes it on the read side + OPEN_T1: ClassVar[float] = -1.0 + + @property + def open(self) -> bool: + return self.t1 == self.OPEN_T1 + + def timestamps(self) -> tuple[float, ...]: + return (self.t0,) if self.open else (self.t0, self.t1) + + +@dataclass +class EngineInfo: + addr: str + worker_type: str + engine_rank: int + gpus: list[list] # [node_ip, gpu_id] pairs + gpu_uuids: list[str | None] + + +@dataclass +class TopologySnapshot(Record): + """Full engine topology; snapshot N is valid until snapshot N+1.""" + + stream: ClassVar[Stream] = Stream.TOPOLOGY + ts: float + engines: list[EngineInfo] + + @classmethod + def from_dict(cls, data: dict) -> TopologySnapshot: + return cls(ts=data["ts"], engines=[EngineInfo(**e) for e in data["engines"]]) + + +class TrajectoryEventKind(StrEnum): + ATTEMPT_START = "attempt_start" + GEN_START = "gen_start" + GEN_END = "gen_end" + TOOL_START = "tool_start" + TOOL_END = "tool_end" + ATTEMPT_END = "attempt_end" + + +@dataclass +class TrajectoryEvent(Record): + """One lifecycle boundary of one sample's rollout (design §18.3). + + ``sample_index`` is the run-global ``Sample.index`` — the same key the + dump join uses, so the read side can resolve a consuming step's batch to + its events. ``turn`` is 1-based for gen/tool segments and -1 for attempt + boundaries; ``weight_version`` is the engine-reported version of the + segment ("" when unknown).""" + + stream: ClassVar[Stream] = Stream.TRAJECTORIES + ts: float + kind: str # a TrajectoryEventKind value + sample_index: int + group_index: int # -1 when the sample carries none + turn: int + weight_version: str + detail: str + + +@dataclass +class GpuSample(Record): + """One NVML sample of one GPU.""" + + stream: ClassVar[Stream] = Stream.GPU_UTIL + ts: float + node: str + gpu: int + util: int + mem_mb: int + power_w: int + + +@dataclass +class EngineSample(Record): + """One scraped sglang engine metric value.""" + + stream: ClassVar[Stream] = Stream.ENGINE_SERIES + ts: float + addr: str + metric: str + labels: dict[str, str] + value: float + + +_RECORD_TYPE_OF_STREAM: dict[Stream, type[Record]] = { + cls.stream: cls for cls in (MetricsRecord, PhaseEvent, TopologySnapshot, GpuSample, EngineSample, TrajectoryEvent) +} + + +@dataclass +class Meta: + run_name: str + start_ts: float + args: dict[str, Any] + schema_version: int = 1 + + FILENAME: ClassVar[str] = "meta.json" + + +def _hour_key(ts: float) -> str: + return time.strftime("%Y%m%d_%H", time.gmtime(ts)) + + +def _hour_bounds(key: str) -> tuple[float, float]: + start = calendar.timegm(time.strptime(key, "%Y%m%d_%H")) + return float(start), float(start + 3600) + + +class _PartitionReader: + """One hourly-partitioned JSONL stream on the read side. + + Files parse lazily at first windowed read into an LRU block cache; a + cached block picks up appended bytes through a per-file byte offset with + the same complete-line rule as ``MetricStore.follow``. A v1 flat + ``{stream}.jsonl`` participates as one "legacy" partition spanning all + time, so every pre-partition dump dir keeps working. + """ + + LEGACY_KEY: ClassVar[str] = "legacy" + + def __init__( + self, + dir: Path, + legacy_path: Path, + *, + parse: Callable[[bytes, Path, int], Any], + concat: Callable[[list], Any], + length: Callable[[Any], int], + line_stamps: Callable[[dict], tuple], + max_blocks: int, + ): + self.dir = dir + self.legacy_path = legacy_path + self._parse = parse + self._concat = concat + self._length = length + self._line_stamps = line_stamps + self._max_blocks = max_blocks + self._blocks: OrderedDict[str, Any] = OrderedDict() + self._offsets: dict[str, int] = {} + + def files(self) -> list[tuple[str, Path]]: + out = [(self.LEGACY_KEY, self.legacy_path)] if self.legacy_path.exists() else [] + if self.dir.is_dir(): + out.extend(sorted((path.stem, path) for path in self.dir.glob("*.jsonl"))) + return out + + def has_data(self) -> bool: + return any(path.stat().st_size > 0 for _, path in self.files()) + + def window(self, t0: float | None, t1: float | None) -> Any: + """Concatenated block over the partitions intersecting [t0, t1].""" + blocks = [] + for key, path in self.files(): + if key != self.LEGACY_KEY: + lo, hi = _hour_bounds(key) + if (t0 is not None and hi <= t0) or (t1 is not None and lo > t1): + continue + blocks.append(self._block(key, path)) + return self._concat([block for block in blocks if self._length(block)]) + + def refresh_cached(self) -> int: + """Tail-refresh every cached block; returns rows appended. Partitions + never read remain lazy — they parse fully at their first read.""" + paths = dict(self.files()) + added = 0 + for key in list(self._blocks): + if key in paths: + before = self._length(self._blocks[key]) + self._block(key, paths[key]) + added += self._length(self._blocks[key]) - before + return added + + def _block(self, key: str, path: Path) -> Any: + offset = self._offsets.get(key, 0) + if key in self._blocks and path.stat().st_size <= offset: + self._blocks.move_to_end(key) + return self._blocks[key] + with open(path, "rb") as f: + f.seek(offset) + chunk = f.read() + end = chunk.rfind(b"\n") + if end >= 0: + new = self._parse(chunk[: end + 1], path, offset) + pieces = [self._blocks[key], new] if key in self._blocks else [new] + self._blocks[key] = self._concat([piece for piece in pieces if self._length(piece)]) + self._offsets[key] = offset + end + 1 + elif key not in self._blocks: + self._blocks[key] = self._concat([]) + self._offsets[key] = offset + self._blocks.move_to_end(key) + while len(self._blocks) > self._max_blocks: + evicted, _ = self._blocks.popitem(last=False) + self._offsets.pop(evicted, None) + return self._blocks[key] + + def edge_stamps(self) -> list[float]: + """Global-range candidates from the first line of the oldest file and + the last complete line of the newest — no full scan.""" + files = [(key, path) for key, path in self.files() if path.stat().st_size > 0] + if not files: + return [] + lines = [] + with open(files[0][1], "rb") as f: + lines.append(f.readline()) + with open(files[-1][1], "rb") as f: + f.seek(0, 2) + f.seek(max(0, f.tell() - 65536)) + tail = f.read() + complete = tail[: tail.rfind(b"\n") + 1] + if complete: + lines.append(complete.splitlines()[-1]) + stamps: list[float] = [] + for line in lines: + try: + stamps.extend(self._line_stamps(json.loads(line))) + except json.JSONDecodeError: + pass # partial line mid-append: no complete record to stamp yet + return stamps + + +class MetricStore: + COLUMNAR_STREAMS: ClassVar[tuple[Stream, ...]] = (Stream.GPU_UTIL, Stream.ENGINE_SERIES) + PARTITIONED_STREAMS: ClassVar[tuple[Stream, ...]] = ( + Stream.GPU_UTIL, + Stream.ENGINE_SERIES, + Stream.PHASES, + Stream.TRAJECTORIES, + ) + MAX_WINDOW_S: ClassVar[float] = 4 * 3600.0 + PARTITION_CACHE_BLOCKS: ClassVar[int] = 24 + CATALOG_FILENAME: ClassVar[str] = "lane_catalog.json" + FRAME_SCHEMAS: ClassVar[dict[Stream, dict[str, Any]]] = { + Stream.GPU_UTIL: dict( + ts=pl.Float64, node=pl.String, gpu=pl.Int64, util=pl.Int64, mem_mb=pl.Int64, power_w=pl.Int64 + ), + # labels dicts are canonicalized to a JSON string column at parse time + Stream.ENGINE_SERIES: dict( + ts=pl.Float64, addr=pl.String, metric=pl.String, labels_json=pl.String, value=pl.Float64 + ), + } + + def __init__(self, dir: Path | str): + self.dir = Path(dir) + self.meta: Meta | None = None + self.records: dict[Stream, list[Record]] = {s: [] for s in Stream if s not in self.PARTITIONED_STREAMS} + self._buffers: dict[Stream, list[Record]] = {s: [] for s in Stream} + self._offsets: dict[Stream, int] = {s: 0 for s in Stream if s not in self.PARTITIONED_STREAMS} + self._readers: dict[Stream, _PartitionReader] = {s: self._make_reader(s) for s in self.PARTITIONED_STREAMS} + self._catalog: dict[str, dict] | None = None # write-side lane catalog accumulator + + def _make_reader(self, stream: Stream) -> _PartitionReader: + if stream in self.COLUMNAR_STREAMS: + empty = pl.DataFrame(schema=self.FRAME_SCHEMAS[stream]) + + def concat_frames(blocks: list, empty: pl.DataFrame = empty) -> pl.DataFrame: + if not blocks: + return empty + merged = blocks[0] if len(blocks) == 1 else pl.concat(blocks, rechunk=False) + # bound chunk fragmentation from long tail sessions + return merged.rechunk() if merged.n_chunks() > 256 else merged + + return _PartitionReader( + self.dir / stream.value, + self.dir / stream.filename, + parse=lambda raw, path, offset, stream=stream: self._parse_columnar(stream, raw, path, offset), + concat=concat_frames, + length=lambda frame: frame.height, + line_stamps=lambda data: (data["ts"],), + max_blocks=self.PARTITION_CACHE_BLOCKS, + ) + assert stream in (Stream.PHASES, Stream.TRAJECTORIES), stream + record_type = _RECORD_TYPE_OF_STREAM[stream] + return _PartitionReader( + self.dir / stream.value, + self.dir / stream.filename, + parse=lambda raw, path, offset, record_type=record_type: self._parse_object_lines( + record_type, raw, path, offset + ), + concat=lambda blocks: [event for block in blocks for event in block], + length=len, + # the record's own timestamps() skips the open-interval t1 sentinel + line_stamps=lambda data, record_type=record_type: record_type.from_dict(data).timestamps(), + max_blocks=self.PARTITION_CACHE_BLOCKS, + ) + + # ------------------------------ write side ------------------------------ + + def append(self, record: Record) -> None: + self._buffers[record.stream].append(record) + + def buffered_count(self, stream: Stream) -> int: + return len(self._buffers[stream]) + + def drop_oldest_buffered(self, stream: Stream, *, keep_ratio: float = 0.9) -> int: + """Drop the oldest buffered (not yet flushed) records of one stream and + return how many were dropped. Lets the collector bound memory when the + disk stays unwritable instead of growing without limit.""" + buffer = self._buffers[stream] + dropped = max(1, int(len(buffer) * (1 - keep_ratio))) + del buffer[:dropped] + return dropped + + def write_meta(self, meta: Meta) -> None: + self.meta = meta + self.dir.mkdir(parents=True, exist_ok=True) + (self.dir / Meta.FILENAME).write_text(json.dumps(asdict(meta), indent=1)) + + def flush(self) -> None: + self.dir.mkdir(parents=True, exist_ok=True) + self._merge_catalog() + for stream, buffer in self._buffers.items(): + if not buffer: + continue + if stream in self.PARTITIONED_STREAMS: + subdir = self.dir / stream.value + subdir.mkdir(exist_ok=True) + groups: dict[str, list[Record]] = {} + for record in buffer: + # phases key by END hour (the completion append is what + # lands on disk, keeping the window lower bound exact); + # OPEN markers have no end yet and key by their start + if stream is Stream.PHASES: + ts = record.t0 if record.open else record.t1 + else: + ts = record.ts + groups.setdefault(_hour_key(ts), []).append(record) + for key, group in sorted(groups.items()): + with open(subdir / f"{key}.jsonl", "a") as f: + f.write("".join(json.dumps(r.to_dict(), separators=(",", ":")) + "\n" for r in group)) + else: + with open(self.dir / stream.filename, "a") as f: + f.write("".join(json.dumps(r.to_dict(), separators=(",", ":")) + "\n" for r in buffer)) + buffer.clear() + + def _merge_catalog(self) -> None: + """Fold buffered records into ``lane_catalog.json``: per-lane train + ranks, engine addrs and roles, so selection resolution and heatmap + rows never need a stream scan. Atomic rewrite, only when changed.""" + if self._catalog is None: + path = self.dir / self.CATALOG_FILENAME + self._catalog = json.loads(path.read_text())["lanes"] if path.exists() else {} + lanes = self._catalog + dirty = False + + def entry(node: str, gpu: int) -> dict: + nonlocal dirty + key = f"{node}:{gpu}" + if key not in lanes: + lanes[key] = dict(ranks=[], engine_addrs=[], roles=[]) + dirty = True + return lanes[key] + + def add(values: list, value) -> None: + nonlocal dirty + if value not in values: + values.append(value) + dirty = True + + for event in self._buffers[Stream.PHASES]: + if event.role != Role.TRAIN: + continue + for gpu in event.gpus: + lane = entry(event.node, gpu) + add(lane["roles"], Role.TRAIN.value) + if event.rank >= 0: + add(lane["ranks"], event.rank) + for snapshot in self._buffers[Stream.TOPOLOGY]: + for engine in snapshot.engines: + for node, gpu in engine.gpus: + lane = entry(node, gpu) + add(lane["roles"], "rollout") + add(lane["engine_addrs"], engine.addr) + for sample in self._buffers[Stream.GPU_UTIL]: + entry(sample.node, sample.gpu) + if dirty: + for lane in lanes.values(): + lane["ranks"].sort() + lane["engine_addrs"].sort() + lane["roles"].sort() + tmp = self.dir / (self.CATALOG_FILENAME + ".tmp") + tmp.write_text(json.dumps(dict(version=1, lanes=lanes))) + os.replace(tmp, self.dir / self.CATALOG_FILENAME) + + # ------------------------------ read side ------------------------------- + + @classmethod + def load(cls, dir: Path | str) -> MetricStore: + store = cls(dir) + meta_path = store.dir / Meta.FILENAME + if meta_path.exists(): + store.meta = Meta(**json.loads(meta_path.read_text())) + store.follow() + return store + + def follow(self) -> int: + """Read records appended since the last load()/follow(). Returns the count. + + Only complete lines (terminated by a newline) are consumed, so a + crash-interrupted partial write at the tail is left for a later call + once its writer completes it. A malformed *complete* line is real + corruption and raises — for hour-partitioned streams that happens at + the first windowed read of the bad partition, since those parse + lazily; here only already-cached partition blocks refresh. + """ + num_new = 0 + for stream in self.records: + path = self.dir / stream.filename + if not path.exists(): + continue + with open(path, "rb") as f: + f.seek(self._offsets[stream]) + chunk = f.read() + end = chunk.rfind(b"\n") + if end < 0: + continue + complete = chunk[: end + 1] + record_type = _RECORD_TYPE_OF_STREAM[stream] + for line in complete.splitlines(): + try: + self.records[stream].append(record_type.from_dict(json.loads(line))) + except (json.JSONDecodeError, TypeError, KeyError) as e: + raise ValueError( + f"corrupt record in {path} near byte {self._offsets[stream]}: {line[:200]!r}" + ) from e + num_new += 1 + self._offsets[stream] += len(complete) + for reader in self._readers.values(): + num_new += reader.refresh_cached() + return num_new + + def _parse_columnar(self, stream: Stream, raw: bytes, path: Path, offset: int) -> pl.DataFrame: + schema = self.FRAME_SCHEMAS[stream] + try: + frame = pl.read_ndjson(io.BytesIO(raw)) + if stream is Stream.ENGINE_SERIES: + # all-empty labels: read_ndjson drops the column entirely + dtype = frame.schema.get("labels") + empty = dtype is None or dtype == pl.Null or (isinstance(dtype, pl.Struct) and not dtype.fields) + # struct schemas unify per chunk (same-writer key order is stable, + # so equal label dicts encode to equal strings within a store) + encoded = pl.lit("{}") if empty else pl.col("labels").struct.json_encode() + frame = frame.with_columns(encoded.alias("labels_json")).drop("labels", strict=False) + if set(frame.columns) != set(schema): + raise ValueError(f"fields {sorted(frame.columns)} != schema {sorted(schema)}") + return frame.select([pl.col(name).cast(dtype) for name, dtype in schema.items()]) + except (pl.exceptions.PolarsError, ValueError) as e: + raise ValueError(f"corrupt record in {path} near byte {offset}: {e}") from e + + def _parse_object_lines(self, record_type: type[Record], raw: bytes, path: Path, offset: int) -> list[Record]: + events = [] + for line in raw.splitlines(): + try: + events.append(record_type.from_dict(json.loads(line))) + except (json.JSONDecodeError, TypeError, KeyError) as e: + raise ValueError(f"corrupt record in {path} near byte {offset}: {line[:200]!r}") from e + return events + + def trajectory_events(self, t0: float | None = None, t1: float | None = None) -> list[TrajectoryEvent]: + """Lifecycle events in [t0, t1], keyed by emission ts (no slack here — + the consuming-step join and its lookback window live in the reader).""" + return self._window_records(self._readers[Stream.TRAJECTORIES].window(t0, t1), t0, t1) + + @staticmethod + def _window_records(events: list[TrajectoryEvent], t0: float | None, t1: float | None) -> list[TrajectoryEvent]: + return [e for e in events if (t0 is None or e.ts >= t0) and (t1 is None or e.ts <= t1)] + + def trajectory_lanes( + self, + *, + t0: float | None = None, + t1: float | None = None, + sample_indices: set[int] | None = None, + ) -> list[dict]: + """Per-sample lifecycle assembly for the batch-anatomy swimlanes: + gen/tool segments paired by (kind, turn), attempt windows, and the + version span observed across segments. A start without an end (still + running, or clipped by the window) yields t1=None and vice versa.""" + grouped: dict[int, list[TrajectoryEvent]] = {} + for event in self.trajectory_events(t0, t1): + if sample_indices is not None and event.sample_index not in sample_indices: + continue + grouped.setdefault(event.sample_index, []).append(event) + + span_kinds = { + TrajectoryEventKind.GEN_START: ("gen", True), + TrajectoryEventKind.GEN_END: ("gen", False), + TrajectoryEventKind.TOOL_START: ("tool", True), + TrajectoryEventKind.TOOL_END: ("tool", False), + } + lanes = [] + for index, events in grouped.items(): + events.sort(key=lambda e: e.ts) + segments: list[dict] = [] + open_spans: dict[tuple[str, int], dict] = {} + attempts: list[dict] = [] + attempt_t0: float | None = None + status = "" + versions = sorted({int(e.weight_version) for e in events if e.weight_version.isdigit()}) + for event in events: + if event.kind == TrajectoryEventKind.ATTEMPT_START: + attempt_t0 = event.ts + elif event.kind == TrajectoryEventKind.ATTEMPT_END: + attempts.append(dict(t0=attempt_t0, t1=event.ts)) + attempt_t0 = None + status = event.detail or status + # an attempt ending closes whatever it left open: the + # single-turn path emits no gen_end (generation ends WITH + # the attempt), and an abort mid-turn stops right here — + # dangling spans must not render to the consume line + for segment in open_spans.values(): + segment["t1"] = event.ts + # coarse gen spans open before the first turn finished, + # so their start event predates any weight_version; the + # attempt_end event carries the version that generated them + if not segment["weight_version"]: + segment["weight_version"] = event.weight_version + open_spans.clear() + else: + base, is_start = span_kinds[TrajectoryEventKind(event.kind)] + if is_start: + segment = dict( + kind=base, t0=event.ts, t1=None, turn=event.turn, weight_version=event.weight_version + ) + open_spans[(base, event.turn)] = segment + segments.append(segment) + else: + segment = open_spans.pop((base, event.turn), None) + if segment is None: + segments.append( + dict( + kind=base, + t0=None, + t1=event.ts, + turn=event.turn, + weight_version=event.weight_version, + ) + ) + else: + segment["t1"] = event.ts + segment["weight_version"] = event.weight_version or segment["weight_version"] + if attempt_t0 is not None: + attempts.append(dict(t0=attempt_t0, t1=None)) # attempt still running + if any(segment["turn"] > 0 for segment in segments if segment["kind"] == "gen"): + # per-turn spans (multi_turn / agentic) supersede the coarse + # attempt-level gen span the generate_and_rm probe emits + segments = [s for s in segments if not (s["kind"] == "gen" and s["turn"] < 0)] + lanes.append( + dict( + sample_index=index, + group_index=events[0].group_index, + first_ts=events[0].ts, + last_ts=events[-1].ts, + segments=segments, + attempts=attempts, + status=status, + versions=versions, + ) + ) + lanes.sort(key=lambda lane: lane["first_ts"]) + return lanes + + def _phase_events(self, t0: float | None, t1: float | None) -> list[PhaseEvent]: + # closed phases partition by END hour (lower bound exact, slack one + # max phase duration FORWARD — design §17); OPEN markers partition by + # their START hour, so the lower bound gets the same slack BACKWARD + lower = None if t0 is None else t0 - self.MAX_WINDOW_S + upper = None if t1 is None else t1 + self.MAX_WINDOW_S + return self._readers[Stream.PHASES].window(lower, upper) + + def has_stream(self, stream: Stream) -> bool: + if stream in self.PARTITIONED_STREAMS: + return self._readers[stream].has_data() + return bool(self.records[stream]) + + def iter_records(self, stream: Stream) -> list[Record]: + """Records as dataclass objects — slow materialization for tests/tools.""" + if stream is Stream.GPU_UTIL: + frame = self._readers[stream].window(None, None) + return [GpuSample(**row) for row in frame.iter_rows(named=True)] + if stream is Stream.ENGINE_SERIES: + return [ + EngineSample( + ts=row["ts"], + addr=row["addr"], + metric=row["metric"], + labels={k: v for k, v in json.loads(row["labels_json"]).items() if v is not None}, + value=row["value"], + ) + for row in self._readers[stream].window(None, None).iter_rows(named=True) + ] + if stream is Stream.PHASES: + return list(self._phase_events(None, None)) + if stream is Stream.TRAJECTORIES: + return self.trajectory_events() + return list(self.records[stream]) + + @staticmethod + def _window(frame: pl.DataFrame, t0: float | None, t1: float | None) -> pl.DataFrame: + if t0 is not None: + frame = frame.filter(pl.col("ts") >= t0) + if t1 is not None: + frame = frame.filter(pl.col("ts") <= t1) + return frame + + # ------------------------------- queries -------------------------------- + + def metric_keys(self) -> list[str]: + keys: set[str] = set() + for record in self.records[Stream.METRICS]: + keys.update(record.metrics.keys()) + return sorted(keys) + + def step_keys(self) -> list[str]: + return sorted({r.step_key for r in self.records[Stream.METRICS] if r.step_key is not None}) + + def metric_series( + self, keys: list[str], *, x_key: str, t0: float | None = None, t1: float | None = None + ) -> dict[str, dict[str, list]]: + """Series for each key, restricted to records logged against ``x_key`` + (e.g. "rollout/step") so values from different step axes never mix.""" + out: dict[str, dict[str, list]] = {k: {"x": [], "y": [], "ts": []} for k in keys} + for record in self.records[Stream.METRICS]: + if record.step_key != x_key: + continue + if (t0 is not None and record.ts < t0) or (t1 is not None and record.ts > t1): + continue + for k in keys: + if k in record.metrics: + series = out[k] + series["x"].append(record.step) + series["y"].append(record.metrics[k]) + series["ts"].append(record.ts) + return out + + def time_range(self) -> tuple[float, float] | None: + stamps = [ts for records in self.records.values() for record in records for ts in record.timestamps()] + for reader in self._readers.values(): + stamps.extend(reader.edge_stamps()) + if not stamps: + return None + return min(stamps), max(stamps) + + # --------------------------- timeline queries --------------------------- + + def lanes(self) -> list[dict]: + """Every (node, gpu) seen in any stream — one timeline lane each.""" + return [dict(node=entry["node"], gpu=entry["gpu"], index=entry["index"]) for entry in self.lane_index()] + + def topology_windows(self) -> list[dict]: + """Engine topology snapshots with validity windows: snapshot N is valid + [ts_N, ts_{N+1}); the last one is open-ended (t1=None).""" + snapshots = sorted(self.records[Stream.TOPOLOGY], key=lambda s: s.ts) + return [ + dict( + t0=snapshot.ts, + t1=snapshots[i + 1].ts if i + 1 < len(snapshots) else None, + engines=[asdict(engine) for engine in snapshot.engines], + ) + for i, snapshot in enumerate(snapshots) + ] + + def phases_by_lane( + self, *, t0: float | None = None, t1: float | None = None, lanes: set[tuple[str, int]] | None = None + ) -> list[dict]: + """Phase intervals resolved onto (node, gpu) lanes. + + Train-side events carry their own GPUs. ``rollout_manager`` events have + none — the manager is a GPU-less driver process — so they are expanded + onto every GPU of every rollout engine, clipped at topology-window + boundaries (an engine restart mid-rollout splits the painted interval). + """ + + def overlaps(a0: float, a1: float) -> bool: + return (t1 is None or a0 < t1) and (t0 is None or a1 > t0) + + # resolve OPEN intervals (Timer.start markers): a closing event with + # the real t1 supersedes its open twin; a still-open one is clipped to + # the newest data timestamp so it renders as a growing in-progress band + events = list(self._phase_events(t0, t1)) + closed = {(e.node, e.rank, e.name, e.t0) for e in events if not e.open} + edge = (self.time_range() or (0.0, 0.0))[1] + events = [ + replace(e, t1=max(edge, e.t0)) if e.open else e + for e in events + if not (e.open and (e.node, e.rank, e.name, e.t0) in closed) + ] + + out = [] + windows = self.topology_windows() + for event in events: + if not overlaps(event.t0, event.t1): + continue + if event.role == Role.ROLLOUT_MANAGER: + for window in windows: + clip0 = max(event.t0, window["t0"]) + clip1 = event.t1 if window["t1"] is None else min(event.t1, window["t1"]) + if clip0 >= clip1: + continue + covered = {(node, gpu) for engine in window["engines"] for node, gpu in engine["gpus"]} + if lanes is not None: + covered &= lanes + for node, gpu in sorted(covered): + out.append( + dict( + name=event.name, + t0=clip0, + t1=clip1, + node=node, + gpu=gpu, + rank=event.rank, + role=event.role, + ) + ) + else: + for gpu in event.gpus: + if lanes is not None and (event.node, gpu) not in lanes: + continue + out.append( + dict( + name=event.name, + t0=event.t0, + t1=event.t1, + node=event.node, + gpu=gpu, + rank=event.rank, + role=event.role, + ) + ) + # synthesize the initialize band: from collector start (meta.start_ts) + # to each lane's first observed event — model loading / engine startup + # has GPU util but no Timer instrumentation. Only when the window + # covers the run start: a later window's first event is not the + # lane's first event ever + if self.meta is not None and (t0 is None or t0 <= self.meta.start_ts): + first_event: dict[tuple[str, int], float] = {} + for event in out: + key = (event["node"], event["gpu"]) + first_event[key] = min(first_event.get(key, float("inf")), event["t0"]) + for (node, gpu), first_t0 in first_event.items(): + if lanes is not None and (node, gpu) not in lanes: + continue + if first_t0 > self.meta.start_ts and overlaps(self.meta.start_ts, first_t0): + out.append( + dict( + name=INITIALIZE_PHASE, + t0=self.meta.start_ts, + t1=first_t0, + node=node, + gpu=gpu, + rank=-1, + role=Role.DERIVED, + ) + ) + out.sort(key=lambda e: (e["node"], e["gpu"], e["t0"])) + return out + + def gpu_series( + self, + *, + t0: float | None = None, + t1: float | None = None, + max_points: int = 2000, + lanes: set[tuple[str, int]] | None = None, + ) -> dict[str, dict]: + """Per-lane NVML series, min/max-downsampled on util (the primary + signal); mem/power take the same indices so all arrays stay aligned.""" + frame = self._window(self._readers[Stream.GPU_UTIL].window(t0, t1), t0, t1) + out = {} + for (node, gpu), part in sorted(frame.partition_by(["node", "gpu"], as_dict=True).items()): + if lanes is not None and (node, gpu) not in lanes: + continue + part = part.sort("ts") + util = part["util"].to_numpy() + indices, _ = minmax_downsample(np.arange(part.height), util, max_points) + out[f"{node}:{gpu}"] = dict( + ts=part["ts"].to_numpy()[indices].tolist(), + util=util[indices].tolist(), + mem_mb=part["mem_mb"].to_numpy()[indices].tolist(), + power_w=part["power_w"].to_numpy()[indices].tolist(), + ) + return out + + def engine_metric_names(self) -> list[str]: + """Distinct scraped engine metrics — the L0 sglang category catalog.""" + return sorted(self._readers[Stream.ENGINE_SERIES].window(None, None).get_column("metric").unique()) + + def engine_series( + self, metric: str, *, t0: float | None = None, t1: float | None = None, max_points: int = 2000 + ) -> list[dict]: + """One series per (engine addr, label set) for the given metric.""" + frame = self._window(self._readers[Stream.ENGINE_SERIES].window(t0, t1), t0, t1).filter( + pl.col("metric") == metric + ) + out = [] + for (addr, labels_json), part in sorted(frame.partition_by(["addr", "labels_json"], as_dict=True).items()): + part = part.sort("ts") + ts, values = stride_downsample(part["ts"].to_numpy(), part["value"].to_numpy(), max_points) + labels = {k: v for k, v in json.loads(labels_json).items() if v is not None} + out.append(dict(addr=addr, labels=labels, ts=ts.tolist(), value=values.tolist())) + return out + + def lane_index(self) -> list[dict]: + """Per-lane metadata for selection resolution: train ranks observed on + the lane, engine addrs that ever covered it, and derived roles. + + Reads ``lane_catalog.json`` (merge-written at flush time) when + present; a legacy dir without one falls back to a full-stream scan. + """ + path = self.dir / self.CATALOG_FILENAME + if path.exists(): + lanes = json.loads(path.read_text())["lanes"] + out = [] + for key, entry in lanes.items(): + node, _, gpu = key.rpartition(":") + out.append( + dict( + node=node, + gpu=int(gpu), + ranks=entry["ranks"], + engine_addrs=entry["engine_addrs"], + roles=entry["roles"], + ) + ) + out.sort(key=lambda entry: (entry["node"], entry["gpu"])) + for i, entry in enumerate(out): + entry["index"] = i + return out + + seen: set[tuple[str, int]] = set( + self._readers[Stream.GPU_UTIL].window(None, None).select("node", "gpu").unique().iter_rows() + ) + events = self._phase_events(None, None) + for event in events: + for gpu in event.gpus: + seen.add((event.node, gpu)) + for snapshot in self.records[Stream.TOPOLOGY]: + for engine in snapshot.engines: + for node, gpu in engine.gpus: + seen.add((node, gpu)) + info: dict[tuple[str, int], dict] = { + (node, gpu): dict(node=node, gpu=gpu, ranks=set(), engine_addrs=set(), roles=set()) + for node, gpu in sorted(seen) + } + for event in events: + if event.role != Role.TRAIN: + continue + for gpu in event.gpus: + entry = info.get((event.node, gpu)) + if entry is not None: + entry["roles"].add(Role.TRAIN.value) + if event.rank >= 0: + entry["ranks"].add(event.rank) + for snapshot in self.records[Stream.TOPOLOGY]: + for engine in snapshot.engines: + for node, gpu in engine.gpus: + entry = info.get((node, gpu)) + if entry is not None: + entry["roles"].add("rollout") + entry["engine_addrs"].add(engine.addr) + return [ + dict( + node=entry["node"], + gpu=entry["gpu"], + index=i, + ranks=sorted(entry["ranks"]), + engine_addrs=sorted(entry["engine_addrs"]), + roles=sorted(entry["roles"]), + ) + for i, entry in enumerate(info.values()) + ] + + def resolve_lanes(self, grammar: str | None) -> set[tuple[str, int]] | None: + """Parse the lane-selection grammar into a lane set (None = all lanes). + + Comma-separated selectors: ``rank:5`` / ``rank:0-7`` (train ranks), + ``g:5`` / ``g:0-31`` (global lane numbers), ``node:``, + ``gpu::``, ``engine:``, + ``role:train`` / ``role:rollout``, or ``all``. Unknown selector syntax + raises; a valid selector matching nothing selects nothing. + """ + if grammar is None or grammar.strip() in ("", "all"): + return None + index = self.lane_index() + selected: set[tuple[str, int]] = set() + for raw_token in grammar.split(","): + token = raw_token.strip() + if not token: + continue + kind, _, value = token.partition(":") + if not value: + raise ValueError(f"bad lane selector {token!r}: expected kind:value") + if kind == "rank": + lo, _, hi = value.partition("-") + ranks = set(range(int(lo), int(hi if hi else lo) + 1)) + selected |= {(e["node"], e["gpu"]) for e in index if ranks & set(e["ranks"])} + elif kind == "g": + # cluster-global lane numbers (the g{index} labels on every view) + lo, _, hi = value.partition("-") + positions = set(range(int(lo), int(hi if hi else lo) + 1)) + selected |= {(e["node"], e["gpu"]) for e in index if e["index"] in positions} + elif kind == "node": + selected |= {(e["node"], e["gpu"]) for e in index if e["node"] == value} + elif kind == "gpu": + node, _, gpu = value.rpartition(":") + if not node: + raise ValueError(f"bad lane selector {token!r}: expected gpu::") + selected.add((node, int(gpu))) + elif kind == "engine": + selected |= {(e["node"], e["gpu"]) for e in index if any(value in a for a in e["engine_addrs"])} + elif kind == "role": + if value not in ("train", "rollout"): + raise ValueError(f"bad lane selector {token!r}: role must be train or rollout") + selected |= {(e["node"], e["gpu"]) for e in index if value in e["roles"]} + else: + raise ValueError(f"unknown lane selector kind {kind!r} in {token!r}") + return selected + + HEATMAP_MAGNITUDE_FIELDS: ClassVar[tuple[str, ...]] = ("util", "mem_mb", "power_w") + + def heatmap( + self, + metric: str, + *, + t0: float | None = None, + t1: float | None = None, + x_buckets: int = 1200, + lanes: set[tuple[str, int]] | None = None, + ) -> dict: + """Rank-carpet matrix: one uint8 cell per (lane, time bucket). + + Magnitude metrics take the bucket MAX (activity must not average away); + ``phase`` paints the covering phase id with non-idle winning over idle. + Returns rows metadata + the raw bytes; the server frames it as binary. + """ + assert x_buckets >= 2, f"{x_buckets=}" + window = self.time_range() + if window is None: + return dict( + rows=[], x_buckets=x_buckets, t0=0.0, t1=0.0, metric=metric, values=b"", scale=None, palette=None + ) + t0 = window[0] if t0 is None else t0 + t1 = window[1] if t1 is None else t1 + assert t1 > t0, f"empty heatmap window [{t0}, {t1})" + if metric == "lifecycle": + return self._lifecycle_heatmap(t0, t1, x_buckets) + rows = [ + dict(node=entry["node"], gpu=entry["gpu"], index=entry["index"], roles=entry["roles"]) + for entry in self.lane_index() + if lanes is None or (entry["node"], entry["gpu"]) in lanes + ] + row_of = {(lane["node"], lane["gpu"]): i for i, lane in enumerate(rows)} + matrix = np.zeros((len(rows), x_buckets), dtype=np.uint8) + span = t1 - t0 + + def bucket(ts: float) -> int: + return min(x_buckets - 1, max(0, int((ts - t0) / span * x_buckets))) + + if metric in self.HEATMAP_MAGNITUDE_FIELDS: + frame = self._window(self._readers[Stream.GPU_UTIL].window(t0, t1), t0, t1) + values = np.zeros((len(rows), x_buckets)) + filled = np.zeros((len(rows), x_buckets), dtype=bool) + for (node, gpu), part in frame.partition_by(["node", "gpu"], as_dict=True).items(): + row = row_of.get((node, gpu)) + if row is None: + continue + buckets = np.clip(((part["ts"].to_numpy() - t0) / span * x_buckets).astype(np.int64), 0, x_buckets - 1) + np.maximum.at(values[row], buckets, part[metric].to_numpy().astype(float)) + filled[row, buckets] = True + peak = float(values[filled].max()) if filled.any() else 1.0 + scale = dict(max=100.0 if metric == "util" else peak) + cell = np.minimum(255, (values[filled] / scale["max"] * 255).astype(np.int64)) + matrix[filled] = cell.astype(np.uint8) + return dict( + rows=rows, + x_buckets=x_buckets, + t0=t0, + t1=t1, + metric=metric, + values=matrix.tobytes(), + scale=scale, + palette=None, + ) + + if metric != "phase": + raise ValueError(f"unknown heatmap metric {metric!r}") + idle = {INITIALIZE_PHASE, "train_wait", "sleep"} + events = self.phases_by_lane(t0=t0, t1=t1, lanes=lanes) + names = sorted({e["name"] for e in events}) + palette = [""] + names # id 0 = no phase observed + phase_id = {name: i + 1 for i, name in enumerate(names)} + # idle painted first so any active phase in the same bucket wins + for event in sorted(events, key=lambda e: e["name"] not in idle): + row = row_of.get((event["node"], event["gpu"])) + if row is None: + continue + b0 = bucket(max(event["t0"], t0)) + b1 = bucket(min(event["t1"], t1)) + matrix[row, b0 : b1 + 1] = phase_id[event["name"]] + return dict( + rows=rows, + x_buckets=x_buckets, + t0=t0, + t1=t1, + metric=metric, + values=matrix.tobytes(), + scale=None, + palette=palette, + ) + + LIFECYCLE_PALETTE: ClassVar[tuple[str, ...]] = ("", "queue", "generating", "tool_wait") + LIFECYCLE_MAX_ROWS: ClassVar[int] = 1024 # y-cap: thousands of samples must not explode the canvas + + def _lifecycle_heatmap(self, t0: float, t1: float, x_buckets: int) -> dict: + """Run-wide trajectory carpet: y = samples ordered by first event, + color = lifecycle state. Queue paints attempt windows first so gen and + tool segments win the shared buckets. Rows cap at LIFECYCLE_MAX_ROWS + (submit order, earliest kept); ``rows_total`` reports the uncapped + count so the frontend can say "showing N/M".""" + all_lanes = self.trajectory_lanes(t0=t0, t1=t1) + lanes = all_lanes[: self.LIFECYCLE_MAX_ROWS] + matrix = np.zeros((len(lanes), x_buckets), dtype=np.uint8) + span = t1 - t0 + state_id = {name: i for i, name in enumerate(self.LIFECYCLE_PALETTE)} + + def bucket(ts: float) -> int: + return min(x_buckets - 1, max(0, int((ts - t0) / span * x_buckets))) + + for row, lane in enumerate(lanes): + for attempt in lane["attempts"]: + matrix[row, bucket(attempt["t0"] or t0) : bucket(attempt["t1"] or t1) + 1] = state_id["queue"] + for segment in lane["segments"]: + cell = state_id["generating" if segment["kind"] == "gen" else "tool_wait"] + matrix[row, bucket(segment["t0"] or t0) : bucket(segment["t1"] or t1) + 1] = cell + rows = [dict(sample_index=lane["sample_index"], group_index=lane["group_index"]) for lane in lanes] + return dict( + rows=rows, + rows_total=len(all_lanes), + x_buckets=x_buckets, + t0=t0, + t1=t1, + metric="lifecycle", + values=matrix.tobytes(), + scale=None, + palette=list(self.LIFECYCLE_PALETTE), + ) + + def outliers( + self, criterion: str, *, t0: float | None = None, t1: float | None = None, top_k: int = 16 + ) -> list[dict]: + """Candidate lanes for the quick-pick buttons; the user confirms the + selection — the machine only proposes.""" + if criterion == "lowest_util": + frame = self._window(self._readers[Stream.GPU_UTIL].window(t0, t1), t0, t1) + grouped = frame.group_by("node", "gpu").agg(pl.col("util").mean().alias("score")) + scored = [(row["score"], row["node"], row["gpu"]) for row in grouped.iter_rows(named=True)] + scored.sort() + elif criterion.startswith("slowest_phase:"): + phase_name = criterion.removeprefix("slowest_phase:") + durations: dict[tuple[str, int], list[float]] = {} + for event in self.phases_by_lane(t0=t0, t1=t1): + if event["name"] != phase_name: + continue + durations.setdefault((event["node"], event["gpu"]), []).append(event["t1"] - event["t0"]) + scored = [(sum(v) / len(v), node, gpu) for (node, gpu), v in durations.items()] + scored.sort(reverse=True) + else: + raise ValueError(f"unknown outlier criterion {criterion!r}") + return [dict(node=node, gpu=gpu, score=score) for score, node, gpu in scored[:top_k]] + + def bubbles(self) -> list[dict]: + """Per-step bubble summary strip: step time and wait ratio from the + perf metrics, with the wall-clock ts as the timeline zoom anchor.""" + out = [] + for record in self.records[Stream.METRICS]: + if record.step_key != "rollout/step": + continue + step_time = record.metrics.get("perf/step_time") + wait_ratio = record.metrics.get("perf/wait_time_ratio") + if step_time is None and wait_ratio is None: + continue + out.append(dict(step=record.step, ts=record.ts, step_time=step_time, wait_ratio=wait_ratio)) + return out + + +# ------------------------------ downsampling -------------------------------- + + +def stride_downsample(xs, ys, max_points: int) -> tuple[np.ndarray, np.ndarray]: + xs, ys = np.asarray(xs), np.asarray(ys) + assert max_points >= 2, f"{max_points=}" + if len(xs) <= max_points: + return xs, ys + idx = np.linspace(0, len(xs) - 1, max_points).astype(np.int64) + return xs[idx], ys[idx] + + +def minmax_downsample(xs, ys, max_points: int) -> tuple[np.ndarray, np.ndarray]: + """Bucketed downsampling keeping each bucket's min and max, so spikes and + dips survive. Output length is at most ``max_points``.""" + xs, ys = np.asarray(xs), np.asarray(ys) + assert max_points >= 2, f"{max_points=}" + if len(xs) <= max_points: + return xs, ys + num_buckets = max_points // 2 + edges = np.linspace(0, len(xs), num_buckets + 1).astype(np.int64) + keep: set[int] = set() + for b in range(num_buckets): + lo, hi = int(edges[b]), int(edges[b + 1]) + if lo >= hi: + continue + segment = ys[lo:hi] + keep.add(lo + int(np.argmin(segment))) + keep.add(lo + int(np.argmax(segment))) + idx = sorted(keep) + return xs[idx], ys[idx] diff --git a/miles/ray/rollout/debug_data.py b/miles/ray/rollout/debug_data.py index 8408521876..3a44e3ebe2 100644 --- a/miles/ray/rollout/debug_data.py +++ b/miles/ray/rollout/debug_data.py @@ -1,3 +1,4 @@ +import json import logging from pathlib import Path @@ -8,6 +9,43 @@ logger = logging.getLogger(__name__) +def trajectory_rows(samples: list[Sample]) -> list[dict]: + """One row per sample that recorded a raw conversation + (``metadata["messages"]``, attached by the session / multi_turn paths).""" + rows = [] + for sample in samples: + messages = sample.metadata.get("messages") if sample.metadata else None + if messages is None: + continue + rows.append( + dict( + sample_index=sample.index, + group_index=sample.group_index, + status=sample.status.value, + reward=sample.reward if isinstance(sample.reward, (int, float)) else None, + prompt=sample.prompt, + messages=messages, + ) + ) + return rows + + +def save_debug_trajectory_data(args, samples: list[Sample], rollout_id, evaluation: bool): + if (path_template := args.save_debug_trajectory_data) is None: + return + rows = trajectory_rows(samples) + if not rows: + return # no conversations: no file (the dashboard keys off its presence) + path = Path(path_template.format(rollout_id=("eval_" if evaluation else "") + str(rollout_id))) + logger.info(f"Save trajectory dump to {path}") + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w") as f: + f.write("".join(json.dumps(row, separators=(",", ":")) + "\n" for row in rows)) + for sample in samples: + if sample.metadata: + sample.metadata.pop("messages", None) # the sidecar is their home; keep the .pt lean + + def load_debug_rollout_data(args, rollout_id: int): data = torch.load( args.load_debug_rollout_data.format(rollout_id=rollout_id), @@ -31,14 +69,9 @@ def save_debug_rollout_data(args, data, rollout_id, evaluation: bool): logger.info(f"Save debug rollout data to {path}") path.parent.mkdir(parents=True, exist_ok=True) - # TODO may improve the format - if evaluation: - dump_data = dict( - samples=[sample.to_dict() for dataset_name, info in data.items() for sample in info["samples"]] - ) - else: - dump_data = dict( - samples=[sample.to_dict() for sample in data], - ) + samples = [sample for info in data.values() for sample in info["samples"]] if evaluation else list(data) + save_debug_trajectory_data(args, samples, rollout_id, evaluation) + # TODO may improve the format + dump_data = dict(samples=[sample.to_dict() for sample in samples]) torch.save(dict(rollout_id=rollout_id, **dump_data), path) diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index c216b440b2..82e6d1bcc6 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -6,6 +6,7 @@ import ray from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS +from miles.dashboard import hooks as dashboard_hooks from miles.ray.rollout.addr_allocator import PortCursors from miles.ray.rollout.debug_data import load_debug_rollout_data, save_debug_rollout_data from miles.ray.rollout.metrics import log_eval_rollout_data, log_rollout_data @@ -28,6 +29,7 @@ from miles.utils.logging_utils import configure_logger from miles.utils.metric_checker import MetricChecker from miles.utils.misc import load_function +from miles.utils.timer import timer from miles.utils.tracking_utils import init_tracking logging.getLogger("httpx").setLevel(logging.WARNING) @@ -75,6 +77,9 @@ def __init__(self, args, pg): init_http_client(args) self.servers = start_rollout_servers(args, pg) start_session_server(args) + # only now does args carry the real router address (init_tracking + # above ran before the router existed) + dashboard_hooks.register_router(args) self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote() self.rollout_id = -1 @@ -107,7 +112,9 @@ async def generate(self, rollout_id): self._health_monitoring_resume() if self.args.ci_test and self.args.use_fault_tolerance and rollout_id >= 2: self._try_ci_fault_injection() - data, metadata, metrics = await self._get_rollout_data(rollout_id=rollout_id) + dashboard_hooks.register_engines(self.servers) + with timer("rollout"): + data, metadata, metrics = await self._get_rollout_data(rollout_id=rollout_id) save_debug_rollout_data(self.args, data, rollout_id=rollout_id, evaluation=False) log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time) data = convert_samples_to_train_data( @@ -125,14 +132,20 @@ async def eval(self, rollout_id): return self._health_monitoring_resume() - if self.use_experimental_refactor: - result = await asyncio.to_thread( - call_rollout_function, self.eval_generate_rollout, RolloutFnEvalInput(rollout_id=rollout_id) - ) - else: - result = await asyncio.to_thread( - call_rollout_fn, self.eval_generate_rollout, self.args, rollout_id, self.data_source, evaluation=True - ) + with timer("eval_rollout"): + if self.use_experimental_refactor: + result = await asyncio.to_thread( + call_rollout_function, self.eval_generate_rollout, RolloutFnEvalInput(rollout_id=rollout_id) + ) + else: + result = await asyncio.to_thread( + call_rollout_fn, + self.eval_generate_rollout, + self.args, + rollout_id, + self.data_source, + evaluation=True, + ) data = result.data save_debug_rollout_data(self.args, data, rollout_id=rollout_id, evaluation=True) metrics = log_eval_rollout_data(rollout_id, self.args, data, result.metrics) diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py index 65bc8d4b6d..165b1e1342 100644 --- a/miles/ray/rollout/train_data_conversion.py +++ b/miles/ray/rollout/train_data_conversion.py @@ -5,6 +5,7 @@ from miles.utils.ray_utils import Box from miles.utils.seqlen_balancing import get_seqlen_balanced_partitions +from miles.utils.timer import Timer from miles.utils.types import Sample @@ -123,6 +124,12 @@ def _post_process_rewards(args, samples: list[Sample] | list[list[Sample]], cust def split_train_data_by_dp(args, data, dp_size): """Split the train data by data parallel size.""" + return [Box(ray.put(shard)) for shard in split_train_data_by_dp_local(args, data, dp_size)] + + +def split_train_data_by_dp_local(args, data, dp_size) -> list[dict]: + """Transport-free core of :func:`split_train_data_by_dp`: returns the plain + per-DP-rank shard dicts (each still carrying its ``partition`` key).""" rollout_data = {} if "prompt" in data: @@ -136,7 +143,7 @@ def split_train_data_by_dp(args, data, dp_size): else: partitions = [range(i, len(total_lengths), dp_size) for i in range(dp_size)] - rollout_data_refs = [] + shards = [] for i in range(dp_size): rollout_data = {} @@ -172,5 +179,18 @@ def split_train_data_by_dp(args, data, dp_size): if key not in data: continue rollout_data[key] = data[key] - rollout_data_refs.append(Box(ray.put(rollout_data))) - return rollout_data_refs + shards.append(rollout_data) + return shards + + +def process_rollout_data_shard(args, rollout_data): + """Train-side completion of the DP split: drop the ``partition`` key and + reorder the batch-global ``total_lengths`` into this shard's row order.""" + partition = rollout_data.pop("partition") + total_lengths = rollout_data["total_lengths"] + + # save the seqlen of the whole rollout batch + Timer().seq_lens = total_lengths + rollout_data["total_lengths"] = [total_lengths[i] for i in partition] + + return rollout_data diff --git a/miles/rollout/generate_hub/multi_turn.py b/miles/rollout/generate_hub/multi_turn.py index 97814ecb3d..bb42edb236 100644 --- a/miles/rollout/generate_hub/multi_turn.py +++ b/miles/rollout/generate_hub/multi_turn.py @@ -3,6 +3,7 @@ """ import argparse +import time from copy import deepcopy from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput @@ -17,6 +18,7 @@ update_sample_with_tool_responses, ) from miles.utils.http_utils import post +from miles.utils.lifecycle import TrajectoryLifecycle from miles.utils.misc import load_function @@ -37,6 +39,10 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: multi_samples = [] + # tool-call markup stays inline in the assistant text (what the model emitted) + record_trajectory = args.save_debug_trajectory_data is not None + trajectory = (list(sample.prompt) if isinstance(sample.prompt, list) else []) if record_trajectory else None + # ----------------------- Initial prompts ------------------------- prompt_tokens_ids = compute_prompt_ids_from_sample(input.state, sample, tools=tool_specs) @@ -56,8 +62,16 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: if args.generate_multi_samples: sample = deepcopy(input.sample) + gen_t0 = time.time() output = await post(url, payload) + sink = None if input.evaluation else TrajectoryLifecycle().sink + if sink is not None: + tokens = output.get("meta_info", {}).get("completion_tokens", "") + sink.gen_span(sample, gen_t0, time.time(), turn=_turn + 1, detail=str(tokens)) await update_sample_from_response(args, sample, payload=payload, output=output, update_loss_mask=True) + if record_trajectory: + trajectory.append({"role": "assistant", "content": output["text"]}) + sample.metadata["messages"] = list(trajectory) # snapshot: multi_samples deepcopies per turn if args.generate_multi_samples: multi_samples.append(deepcopy(sample)) @@ -71,8 +85,14 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: if len(tool_calls) == 0: break + tool_t0 = time.time() tool_messages = await execute_tool_calls(tool_calls, execute_tool_function) + if sink is not None: + sink.tool_span(sample, tool_t0, time.time(), turn=_turn + 1, detail=f"{len(tool_calls)} calls") update_sample_with_tool_responses(sample, tool_messages, tokenizer=tokenizer) + if record_trajectory: + trajectory.extend(tool_messages) + sample.metadata["messages"] = list(trajectory) return GenerateFnOutput(samples=multi_samples if args.generate_multi_samples else sample) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 7d9e9c18ac..f458b647f9 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -13,6 +13,7 @@ ) from miles.rollout.session.types import GetSessionResponse, SessionRecord from miles.utils.http_utils import post +from miles.utils.lifecycle import attach_lifecycle_metadata from miles.utils.types import Sample logger = logging.getLogger(__name__) @@ -151,6 +152,10 @@ def compute_samples_from_openai_records( cursor += matched sample = _compute_sample_from_openai_record(args, input_sample, record, tokenizer, trim_count) + attach_lifecycle_metadata(sample, record, records[i - 1] if i else None, turn=i + 1) + if is_last and args.save_debug_trajectory_data is not None: + # the final turn's request holds the full message history + sample.metadata["messages"] = record.request["messages"] + [record.response["choices"][0]["message"]] samples.append(sample) if accumulated_token_ids is not None: diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 68eb57e726..cfdd25107f 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -70,9 +70,29 @@ def _merge_opd_student_top_logprobs(av, bv): ) return av + [[] for _ in range(obs_len)] + bv + def _pop_lifecycle(metadata): + # per-turn dashboard lifecycle segments (miles.dashboard.lifecycle) + # differ across turns by construction; a merged sample carries them + # as one ordered list instead of failing the equality assert below + if not metadata or "lifecycle" not in metadata: + return metadata, [] + value = metadata["lifecycle"] + rest = {k: v for k, v in metadata.items() if k != "lifecycle"} + return rest, value if isinstance(value, list) else [value] + + def _pop_messages(metadata): + # the later turn's conversation snapshot is a superset; keep it + if not metadata or "messages" not in metadata: + return metadata, None + return {k: v for k, v in metadata.items() if k != "messages"}, metadata["messages"] + def _merge_metadata(): a_metadata, a_top_logprobs = _pop_opd_student_top_logprobs(a.metadata) b_metadata, b_top_logprobs = _pop_opd_student_top_logprobs(b.metadata) + a_metadata, a_lifecycle = _pop_lifecycle(a_metadata) + b_metadata, b_lifecycle = _pop_lifecycle(b_metadata) + a_metadata, a_messages = _pop_messages(a_metadata) + b_metadata, b_messages = _pop_messages(b_metadata) assert a_metadata == b_metadata, f"metadata mismatch: a.metadata={a.metadata}, b.metadata={b.metadata}" merged_metadata = deepcopy(a_metadata) @@ -81,6 +101,14 @@ def _merge_metadata(): if merged_metadata is None: merged_metadata = {} merged_metadata[_OPD_STUDENT_TOP_LOGPROBS_KEY] = merged_top_logprobs + if a_lifecycle or b_lifecycle: + if merged_metadata is None: + merged_metadata = {} + merged_metadata["lifecycle"] = a_lifecycle + b_lifecycle + if (messages := b_messages or a_messages) is not None: + if merged_metadata is None: + merged_metadata = {} + merged_metadata["messages"] = messages return merged_metadata _fill_defaults(a) diff --git a/miles/rollout/session/core.py b/miles/rollout/session/core.py index 30ff0b98a7..e9282958f9 100644 --- a/miles/rollout/session/core.py +++ b/miles/rollout/session/core.py @@ -167,6 +167,7 @@ async def delete_session(self, session_id: str) -> Response: async def chat_completions( self, session_id: str, *, method: str, query: str, headers: dict, body: bytes ) -> Response: + request_timestamp = time.time() """Proxy a chat completion through the backend with TITO token tracking. Three phases around the per-session lock: (1) under the lock, parse the @@ -284,6 +285,7 @@ async def chat_completions( record = SessionRecord( timestamp=time.time(), + request_timestamp=request_timestamp, method=method, path="/v1/chat/completions", status_code=result["status_code"], diff --git a/miles/rollout/session/types.py b/miles/rollout/session/types.py index 6548902093..051f84ae26 100644 --- a/miles/rollout/session/types.py +++ b/miles/rollout/session/types.py @@ -2,7 +2,10 @@ class SessionRecord(BaseModel): - timestamp: float + timestamp: float # stamped after the backend responds (= call END) + # arrival stamp: bounds the agent-side gap between chat calls at the + # server edge (tool time sampling); optional for wire compatibility + request_timestamp: float | None = None method: str path: str request: dict diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index 64662ed522..dec5270a6c 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -21,6 +21,7 @@ from miles.utils.data import Dataset from miles.utils.eval_config import EvalDatasetConfig from miles.utils.http_utils import get, post +from miles.utils.lifecycle import TrajectoryLifecycle from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled from miles.utils.misc import SingletonMeta, load_function from miles.utils.processing_utils import ( @@ -260,11 +261,21 @@ async def generate_and_rm( state = GenerateState(args) + # dashboard lifecycle probe (design §18.3): the semaphore wait IS the + # queue; attempt_end fires once generation is over, before reward + sink = None if evaluation else TrajectoryLifecycle().sink + if sink is not None: + sink.attempt_start(sample) + # generate async with state.semaphore: if state.aborted: sample.status = Sample.Status.ABORTED + if sink is not None: + sink.attempt_end(sample) return sample + if sink is not None: + sink.gen_start(sample) with state.dp_rank_context() as _: # Check sample.generate_function_path for per-sample custom_generate_function_path (e.g., from eval dataset config) @@ -279,6 +290,9 @@ async def generate_and_rm( else: sample = await generate(args, sample, sampling_params) + if sink is not None: + sink.attempt_end(sample) + # for the rm that need the whole group, we will not do the rm here if args.group_rm: return sample diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7d7858b91f..6cdde296d7 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1497,6 +1497,15 @@ def add_debug_arguments(parser): "The file will be saved to `save_debug_rollout_data.format(rollout_id)`." ), ) + parser.add_argument( + "--save-debug-trajectory-data", + type=str, + default=None, + help=( + "Save per-sample role-tagged trajectory text (JSONL) next to the rollout " + "dump. The file will be saved to `save_debug_trajectory_data.format(rollout_id)`." + ), + ) parser.add_argument( "--load-debug-rollout-data", type=str, @@ -1955,6 +1964,9 @@ def add_user_provided_function_arguments(parser): parser = add_mlflow_arguments(parser) parser = add_tensorboard_arguments(parser) parser = add_prometheus_arguments(parser) + from miles.dashboard.args import add_dashboard_arguments + + add_dashboard_arguments(parser) parser = add_router_arguments(parser) parser = add_debug_arguments(parser) parser = add_sglang_arguments(parser) @@ -2114,6 +2126,10 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def miles_validate_args(args): + from miles.dashboard.args import validate_dashboard_args + + validate_dashboard_args(args) + args.eval_datasets = _resolve_eval_datasets(args) if args.recompute_logprobs_via_prefill: @@ -2356,6 +2372,7 @@ def miles_validate_args(args): if args.dump_details is not None: args.save_debug_rollout_data = f"{args.dump_details}/rollout_data/{{rollout_id}}.pt" args.save_debug_train_data = f"{args.dump_details}/train_data/{{rollout_id}}_{{rank}}.pt" + args.save_debug_trajectory_data = f"{args.dump_details}/trajectory/{{rollout_id}}.jsonl" if args.load_debug_rollout_data is not None: logger.info( diff --git a/miles/utils/data.py b/miles/utils/data.py index 0fe7ad483f..1bfd9b8798 100644 --- a/miles/utils/data.py +++ b/miles/utils/data.py @@ -17,7 +17,6 @@ from miles.utils.processing_utils import call_processor from miles.utils.types import MultimodalTypes, Sample -from .timer import Timer __all__ = ["Dataset"] @@ -273,14 +272,8 @@ def get_minimum_num_micro_batch_size(total_lengths, max_tokens_per_gpu): def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size): + from miles.ray.rollout.train_data_conversion import process_rollout_data_shard + assert len(rollout_data_ref) == dp_size rollout_data = ray.get(rollout_data_ref[dp_rank].inner) - - partition = rollout_data.pop("partition") - total_lengths = rollout_data["total_lengths"] - - # save the seqlen of the whole rollout batch - Timer().seq_lens = total_lengths - rollout_data["total_lengths"] = [total_lengths[i] for i in partition] - - return rollout_data + return process_rollout_data_shard(args, rollout_data) diff --git a/miles/utils/lifecycle.py b/miles/utils/lifecycle.py new file mode 100644 index 0000000000..48fe9bb5ab --- /dev/null +++ b/miles/utils/lifecycle.py @@ -0,0 +1,43 @@ +"""Trajectory-lifecycle seam callable from core rollout code. + +A neutral observation point at the same layer as ``Timer.event_sinks``: core +generation paths (``sglang_rollout``, ``generate_hub``) read +``TrajectoryLifecycle().sink`` unconditionally, so this module must never +import anything from the dashboard extra. The sink object itself is constructed and installed by +``miles.dashboard.hooks`` when the dashboard is enabled; while it is ``None`` +every probe site is a guarded no-op. +""" + +from __future__ import annotations + +from miles.utils.misc import SingletonMeta + + +class TrajectoryLifecycle(metaclass=SingletonMeta): + def __init__(self): + # the dashboard TrajectorySink, duck-typed; None = dashboard off + self.sink = None + + +def attach_lifecycle_metadata(sample, record, prev_record, turn: int) -> None: + """Fold one session-server chat record's timing onto the sample it became + (design §18.3 hook 3): ``SessionRecord.timestamp`` is stamped after the + backend returns (= segment END); the start comes from the engine-reported + ``e2e_latency`` when present. Written to ``sample.metadata["lifecycle"]`` + (the ``tito_session_mismatch`` precedent) so it rides the Sample across + the session boundary — records themselves never need to leave it. + """ + meta_info = record.response["choices"][0].get("meta_info", {}) + latency = meta_info.get("e2e_latency") + t1 = record.timestamp + t0 = t1 - latency if latency is not None else None + segment = dict(t0=t0, t1=t1, turn=turn) + if record.request_timestamp is not None: + # server-edge arrival: closes the agent-side gap exactly; the span + # up to gen start is engine/proxy queueing, not agent work + segment["req_ts"] = record.request_timestamp + if prev_record is not None: + # the gap between the previous chat call's end and this one's start is + # agent-side work (tool execution, environment steps) + segment["prev_t1"] = prev_record.timestamp + sample.metadata["lifecycle"] = segment diff --git a/miles/utils/timer.py b/miles/utils/timer.py index f5610842f6..b15de782c6 100644 --- a/miles/utils/timer.py +++ b/miles/utils/timer.py @@ -16,16 +16,28 @@ class Timer(metaclass=SingletonMeta): def __init__(self): self.timers = {} self.start_time = {} + # interval observers called as sink(name, start_ts, end_ts) on every + # end(); a sink additionally exposing .begin(name, start_ts) hears + # starts too, so long-running intervals are observable while open. + # Registered by miles.dashboard when the dashboard is enabled. + self.event_sinks = [] def start(self, name): assert name not in self.start_time, f"Timer {name} already started." self.start_time[name] = time() + for sink in self.event_sinks: + begin = getattr(sink, "begin", None) + if callable(begin): + begin(name, self.start_time[name]) if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: logger.info(f"Timer {name} start") def end(self, name): assert name in self.start_time, f"Timer {name} not started." - elapsed_time = time() - self.start_time[name] + start_time = self.start_time[name] + elapsed_time = time() - start_time + for sink in self.event_sinks: + sink(name, start_time, start_time + elapsed_time) self.add(name, elapsed_time) del self.start_time[name] if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: diff --git a/miles/utils/tracking_utils/base.py b/miles/utils/tracking_utils/base.py index f1f67dc7dc..b0490ed6b7 100644 --- a/miles/utils/tracking_utils/base.py +++ b/miles/utils/tracking_utils/base.py @@ -120,6 +120,26 @@ def finish(self) -> None: return +class MilesDashboardBackend(TrackingBackend): + # Live dashboard collection (see miles/dashboard). Lazy imports keep the + # dashboard package out of processes that never enable it. + + def init(self, args, *, primary: bool = True, **kwargs) -> None: + from miles.dashboard.backend import init_dashboard + + init_dashboard(args, primary=primary, **kwargs) + + def log(self, metrics: dict[str, Any], step: int | None = None, *, step_key: str | None = None, **kwargs) -> None: + from miles.dashboard.backend import dashboard_log + + dashboard_log(metrics, step=step, step_key=step_key) + + def finish(self) -> None: + from miles.dashboard.backend import finish_dashboard + + finish_dashboard() + + # Registry that maps backend name → (class, args-flag attribute) BACKEND_REGISTRY: dict[str, tuple[type[TrackingBackend], str]] = { @@ -127,6 +147,7 @@ def finish(self) -> None: "tensorboard": (TensorboardBackend, "use_tensorboard"), "mlflow": (MlflowBackend, "use_mlflow"), "prometheus": (PrometheusBackend, "use_prometheus"), + "miles_dashboard": (MilesDashboardBackend, "use_miles_dashboard"), } diff --git a/miles/utils/train_dump_utils.py b/miles/utils/train_dump_utils.py index 4b5cbc6a51..d44521083b 100644 --- a/miles/utils/train_dump_utils.py +++ b/miles/utils/train_dump_utils.py @@ -7,8 +7,14 @@ def save_debug_train_data(args, *, rollout_id, rollout_data): + if args.save_debug_train_data is not None: + save_debug_train_data_for_rank( + args, rollout_id=rollout_id, rollout_data=rollout_data, rank=torch.distributed.get_rank() + ) + + +def save_debug_train_data_for_rank(args, *, rollout_id, rollout_data, rank): if (path_template := args.save_debug_train_data) is not None: - rank = torch.distributed.get_rank() path = Path(path_template.format(rollout_id=rollout_id, rank=rank)) logger.info(f"Save debug train data to {path}") path.parent.mkdir(parents=True, exist_ok=True) diff --git a/setup.py b/setup.py index de7281b742..f47ac80965 100644 --- a/setup.py +++ b/setup.py @@ -43,6 +43,7 @@ def get_tag(self): version="0.2.1", packages=find_packages(include=["miles*", "miles_plugins*"]), include_package_data=True, + package_data={"miles.dashboard": ["static/*"]}, install_requires=_fetch_requirements("requirements.txt"), extras_require={ "fsdp": [ @@ -51,6 +52,14 @@ def get_tag(self): "mlflow": [ "mlflow>=2.0", ], + # dashboard server deps; present in the training image via sglang, + # needed explicitly only for standalone offline serving + "dashboard": [ + "fastapi", + "uvicorn", + "prometheus_client", + "polars", + ], }, python_requires=">=3.10", classifiers=[ diff --git a/tests/fast/dashboard/__init__.py b/tests/fast/dashboard/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/fast/dashboard/dummy_dump.py b/tests/fast/dashboard/dummy_dump.py new file mode 100644 index 0000000000..90bd46cccc --- /dev/null +++ b/tests/fast/dashboard/dummy_dump.py @@ -0,0 +1,234 @@ +"""Generate a DUMMY ``--dump-details`` directory through miles' +REAL dumping pipeline. + +Only the *content* is fake (tiny random rollouts); every file on disk is +produced by the actual production code path, so if a PR changes the dumping +logic — schemas, file naming, DP partition / ``sample_indices`` semantics, +reward post-processing, the batch-global ``raw_reward`` layout — dumps built +here change with it and the reader tests break loudly instead of silently +validating a stale hand-written copy. + +Real code invoked per step, in the real call order: + +1. ``save_debug_rollout_data`` (rollout + eval dump files) +2. ``convert_samples_to_train_data`` (column dict incl. ``sample_indices``) +3. ``split_train_data_by_dp_local`` (real seqlen-balanced DP partition) +4. ``process_rollout_data_shard`` (train-side partition pop / reorder) +5. ``save_debug_train_data_for_rank`` (per-rank train dump files) + +The only hand-synthesized part is :func:`_apply_dummy_actor_columns`: the +training actor's GPU forward passes cannot run in a CPU test, so the *values* +of its per-token outputs (``log_probs``, ``ref_log_probs``, ``entropy``, +``ref_entropy``, ``advantages``, ``returns``) are random tensors and the +actor-side tensorization of ``tokens``/``loss_masks``/``rollout_log_probs`` +is mimicked. Column names there mirror real dumps; the end-to-end guard for +them is the realdata test plus the e2e CI run. +""" + +from __future__ import annotations + +import os +import random +import time +from argparse import Namespace +from dataclasses import dataclass +from pathlib import Path + +import torch + +from miles.ray.rollout.debug_data import save_debug_rollout_data +from miles.ray.rollout.train_data_conversion import ( + convert_samples_to_train_data, + process_rollout_data_shard, + split_train_data_by_dp_local, +) +from miles.utils.train_dump_utils import save_debug_train_data_for_rank +from miles.utils.types import Sample + + +@dataclass +class DummyRunTruth: + """Ground truth returned by :func:`dump_dummy_run` for test assertions.""" + + n_samples_per_step: int + # step -> one list of Sample.index per DP shard, in the dumped row order + # (taken from the REAL partition output, not recomputed) + shard_indices: dict[int, list[list[int]]] + eval_ids: list[int] + + +def dump_dummy_run( + dump_dir: Path, + *, + steps: int = 2, + num_prompts: int = 4, + n_samples_per_prompt: int = 2, + dp_size: int = 2, + tp_dup: int = 2, + max_response_len: int = 16, + with_entropy: bool = True, + with_eval: bool = True, + with_tokenizer: bool = True, + remove_sample_indices: tuple[int, ...] = (), + seed: int = 0, +) -> DummyRunTruth: + """``remove_sample_indices`` marks the given within-step positions as + ``remove_sample=True`` in every step; the real conversion then writes an + all-zero loss mask for them.""" + n = num_prompts * n_samples_per_prompt + assert n % dp_size == 0, "equal-size DP partition needs divisible batch" + args = _make_args(dump_dir, num_prompts=num_prompts, n_samples_per_prompt=n_samples_per_prompt) + rng = random.Random(seed) + truth = DummyRunTruth(n_samples_per_step=n, shard_indices={}, eval_ids=[]) + + if with_tokenizer: + _write_tokenizer(dump_dir) + + for rollout_id in range(steps): + samples = [ + _make_sample( + rng, + index=rollout_id * n + i, + group_index=i // n_samples_per_prompt, + max_response_len=max_response_len, + remove_sample=i in remove_sample_indices, + ) + for i in range(n) + ] + for sample in samples[:2]: + sample.metadata["messages"] = _dummy_messages(sample) + save_debug_rollout_data(args, samples, rollout_id=rollout_id, evaluation=False) + if with_eval and rollout_id == 0: + # eval takes the real input shape: {dataset_name: {"samples": [...]}} + eval_data = {"dummy_eval": {"samples": samples[: n // 2]}} + save_debug_rollout_data(args, eval_data, rollout_id=rollout_id, evaluation=True) + truth.eval_ids.append(rollout_id) + + train_data = convert_samples_to_train_data( + args, + samples, + metadata={}, + custom_convert_samples_to_train_data_func=None, + custom_reward_post_process_func=None, + ) + shards = split_train_data_by_dp_local(args, train_data, dp_size) + truth.shard_indices[rollout_id] = [list(shard["sample_indices"]) for shard in shards] + generator = torch.Generator().manual_seed(rng.randint(0, 2**31)) + for shard_index, shard in enumerate(shards): + shard = process_rollout_data_shard(args, shard) + _apply_dummy_actor_columns(shard, generator, with_entropy=with_entropy) + for dup in range(tp_dup): + save_debug_train_data_for_rank( + args, rollout_id=rollout_id, rollout_data=shard, rank=shard_index * tp_dup + dup + ) + + age_files(dump_dir) + return truth + + +def _dummy_messages(sample) -> list[dict]: + return [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": str(sample.prompt)}, + { + "role": "assistant", + "content": "Let me look that up.", + "reasoning_content": "The user asks a question; call the lookup tool first.", + "tool_calls": [{"function": {"name": "lookup", "arguments": '{"query": "answer"}'}}], + }, + {"role": "tool", "name": "lookup", "content": '{"result": 42}'}, + {"role": "assistant", "content": f"The answer for sample {sample.index} is 42."}, + ] + + +def _make_args(dump_dir: Path, *, num_prompts: int, n_samples_per_prompt: int) -> Namespace: + # Exactly the attributes the invoked dump-pipeline functions read; a new + # attribute dependency in main code fails loudly here (drift detection). + return Namespace( + save_debug_rollout_data=f"{dump_dir}/rollout_data/{{rollout_id}}.pt", + save_debug_trajectory_data=f"{dump_dir}/trajectory/{{rollout_id}}.jsonl", + save_debug_train_data=f"{dump_dir}/train_data/{{rollout_id}}_{{rank}}.pt", + balance_data=True, + use_dynamic_global_batch_size=False, + advantage_estimator="grpo", + rewards_normalization=True, + grpo_std_normalization=True, + rollout_batch_size=num_prompts, + n_samples_per_prompt=n_samples_per_prompt, + reward_key=None, + ) + + +def _make_sample( + rng: random.Random, *, index: int, group_index: int, max_response_len: int, remove_sample: bool = False +) -> Sample: + response_length = rng.randint(4, max_response_len) + prompt_length = rng.randint(3, 8) + truncated = response_length == max_response_len + # every third sample is agentic-shaped: multi-turn (mixed weight versions + # under fully async) with a chat prompt carrying tool messages + agentic = index % 3 == 0 + versions = [str(index % 3), str(index % 3 + 1)] if agentic else [str(index % 3)] + prompt: str | list = "What is 1+1?" + if agentic: + prompt = [ + dict(role="user", content="What is 1+1?"), + dict(role="assistant", content="let me check"), + dict(role="tool", content="2"), + ] + return Sample( + group_index=group_index, + index=index, + prompt=prompt, + tokens=[rng.randint(0, _VOCAB_SIZE - 1) for _ in range(prompt_length + response_length)], + response="x" * response_length, + response_length=response_length, + reward=float(rng.random() < 0.5), + rollout_log_probs=[-rng.random() for _ in range(response_length)], + weight_versions=versions, + status=Sample.Status.TRUNCATED if truncated else Sample.Status.COMPLETED, + remove_sample=remove_sample, + ) + + +_VOCAB_SIZE = 100 + + +def _write_tokenizer(dump_dir: Path) -> None: + """Persist a tiny word-level tokenizer (token id i <-> text "t{i}") with the + same ``save_pretrained`` mechanism the real data source uses for + ``{dump_details}/tokenizer/``.""" + from tokenizers import Tokenizer, models + from transformers import PreTrainedTokenizerFast + + vocab = {f"t{i}": i for i in range(_VOCAB_SIZE)} + raw = Tokenizer(models.WordLevel(vocab, unk_token="t0")) + PreTrainedTokenizerFast(tokenizer_object=raw).save_pretrained(dump_dir / "tokenizer") + + +def _apply_dummy_actor_columns(shard: dict, generator: torch.Generator, *, with_entropy: bool) -> None: + """Mimic what the training actor does to a shard before dumping it: + tensorize the converted columns and attach forward-pass outputs. + The values are random; the column names mirror real dumps.""" + response_lengths = shard["response_lengths"] + + def per_token(sign: float = 1.0) -> list[torch.Tensor]: + return [sign * torch.rand(r, generator=generator) for r in response_lengths] + + shard["tokens"] = [torch.tensor(t) for t in shard["tokens"]] + shard["loss_masks"] = [torch.tensor(m, dtype=torch.int32) for m in shard["loss_masks"]] + shard["rollout_log_probs"] = [torch.tensor(p) for p in shard["rollout_log_probs"]] + shard["log_probs"] = per_token(sign=-1.0) + shard["ref_log_probs"] = per_token(sign=-1.0) + shard["advantages"] = per_token() + shard["returns"] = per_token() + if with_entropy: + shard["entropy"] = per_token() + shard["ref_entropy"] = per_token() + + +def age_files(dump_dir: Path, *, seconds: float = 100.0) -> None: + """Push mtimes into the past so files pass the reader's visibility guard.""" + stamp = time.time() - seconds + for path in dump_dir.rglob("*.pt"): + os.utime(path, (stamp, stamp)) diff --git a/tests/fast/dashboard/dummy_telemetry.py b/tests/fast/dashboard/dummy_telemetry.py new file mode 100644 index 0000000000..ee3faec424 --- /dev/null +++ b/tests/fast/dashboard/dummy_telemetry.py @@ -0,0 +1,240 @@ +"""Generate DUMMY live-collection telemetry through miles' REAL store write path. + +Companion of ``dummy_dump.py`` (which fabricates the ``.pt`` dump side): only +the *content* here is fake — every record goes through the actual +``MetricStore`` append/flush pipeline into ``{dump_details}/dashboard/``, so +schema or query changes break the timeline tests instead of drifting +silently. + +This file doubles as the OUTPUT CONTRACT for the live collector (PR-10 in +the implementation plan): the collector must produce files that +``MetricStore.load`` plus the timeline queries interpret exactly like this +fixture. + +The fabricated story deliberately contains the shapes the efficiency view is +designed to surface: + +1. colocate cadence per 100 s step: rollout [T, T+40) then per-rank + ``train_wait`` / ``actor_train`` / ``update_weights`` with cross-rank skew; +2. a rollout long tail: engine A stays busy until T+38 while engine B drains + at T+25 (util drops, ``num_running_reqs`` hits 0 — a bubble); +3. a fault-tolerance engine restart MID-ROLLOUT of step 1: engine B moves + from port 15004 to 15006, so topology windows split and the rollout-lane + expansion must clip intervals at the boundary; +4. one util spike (100%) at a known instant for downsampling assertions. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from miles.dashboard.store import ( + EngineInfo, + EngineSample, + GpuSample, + Meta, + MetricsRecord, + MetricStore, + PhaseEvent, + Role, + TopologySnapshot, + TrajectoryEvent, + TrajectoryEventKind, +) + +BASE_TS = 1_000_000.0 +GPU_NODE = "10.0.0.1" +MANAGER_NODE = "10.0.0.2" # rollout manager runs on a different (GPU-less) node + +ENGINE_A = "http://10.0.0.1:15000" +ENGINE_B_OLD = "http://10.0.0.1:15004" +ENGINE_B_NEW = "http://10.0.0.1:15006" + +STEP_SECONDS = 100 +INIT_SECONDS = 30 # gap between collector start (meta.start_ts) and step 0 +ROLLOUT_SECONDS = 40 +DRAIN_A = 38 # engine A busy almost to the end of the rollout window +DRAIN_B = 25 # engine B drains early -> bubble on its lanes +RESTART_OFFSET = 20 # engine B restarts at step-1 rollout start + 20s (mid-rollout) + +SPIKE = dict(step=0, offset=50, gpu=0, util=100) + + +@dataclass +class DummyTelemetryTruth: + """Ground truth returned by :func:`dump_dummy_telemetry` for test assertions.""" + + steps: int + gpus: int + restart_ts: float + spike_ts: float + + def step_start(self, step: int) -> float: + return BASE_TS + INIT_SECONDS + step * STEP_SECONDS + + def rollout_interval(self, step: int) -> tuple[float, float]: + return self.step_start(step), self.step_start(step) + ROLLOUT_SECONDS + + +SAMPLES_PER_STEP = 8 # dummy_dump's default num_prompts * n_samples_per_prompt + + +def _trajectory_events(store: MetricStore, truth: DummyTelemetryTruth, samples_per_step: int) -> None: + """Lifecycle events for each step's batch, aligned with dummy_dump's + global sample indices (step * n + i) and its agentic pattern (every third + sample: two turns, one tool call, mixed weight versions).""" + for step in range(truth.steps): + t0, t1 = truth.rollout_interval(step) + for i in range(samples_per_step): + index = step * samples_per_step + i + agentic = index % 3 == 0 + # stagger submits so every lifecycle fits inside the rollout window + start = t0 + i * (t1 - t0 - 16.0) / samples_per_step + version = str(index % 3) + + def emit(kind, ts, turn=-1, version=version, detail="", index=index, group=i // 2): + store.append( + TrajectoryEvent( + ts=ts, + kind=kind, + sample_index=index, + group_index=group, + turn=turn, + weight_version=version, + detail=detail, + ) + ) + + emit(TrajectoryEventKind.ATTEMPT_START, start) + emit(TrajectoryEventKind.GEN_START, start + 1.0, turn=1) + emit(TrajectoryEventKind.GEN_END, start + 6.0, turn=1) + if agentic: + emit(TrajectoryEventKind.TOOL_START, start + 6.0, turn=1, detail="1 calls") + emit(TrajectoryEventKind.TOOL_END, start + 9.0, turn=1, detail="1 calls") + bumped = str(index % 3 + 1) + emit(TrajectoryEventKind.GEN_START, start + 9.0, turn=2, version=bumped) + emit(TrajectoryEventKind.GEN_END, start + 14.0, turn=2, version=bumped) + emit(TrajectoryEventKind.ATTEMPT_END, start + (15.0 if agentic else 7.0), detail="completed") + + +def dump_dummy_telemetry( + dump_dir: Path, *, steps: int = 3, gpus: int = 4, samples_per_step: int = SAMPLES_PER_STEP +) -> DummyTelemetryTruth: + truth = DummyTelemetryTruth( + steps=steps, + gpus=gpus, + restart_ts=BASE_TS + INIT_SECONDS + STEP_SECONDS + RESTART_OFFSET, + spike_ts=BASE_TS + INIT_SECONDS + SPIKE["step"] * STEP_SECONDS + SPIKE["offset"], + ) + store = MetricStore(dump_dir / "dashboard") + store.write_meta(Meta(run_name="dummy-telemetry", start_ts=BASE_TS, args={"colocate": True})) + + def engine(addr: str, engine_rank: int, gpu_ids: list[int]) -> EngineInfo: + return EngineInfo( + addr=addr, + worker_type="regular", + engine_rank=engine_rank, + gpus=[[GPU_NODE, g] for g in gpu_ids], + gpu_uuids=[None] * len(gpu_ids), + ) + + store.append(TopologySnapshot(ts=BASE_TS, engines=[engine(ENGINE_A, 0, [0, 1]), engine(ENGINE_B_OLD, 1, [2, 3])])) + store.append( + TopologySnapshot(ts=truth.restart_ts, engines=[engine(ENGINE_A, 0, [0, 1]), engine(ENGINE_B_NEW, 1, [2, 3])]) + ) + + def engine_b_addr(ts: float) -> str: + return ENGINE_B_OLD if ts < truth.restart_ts else ENGINE_B_NEW + + for offset in range(INIT_SECONDS): + for gpu in range(gpus): + store.append(GpuSample(ts=BASE_TS + offset, node=GPU_NODE, gpu=gpu, util=20, mem_mb=30_000, power_w=300)) + + for step in range(steps): + t = truth.step_start(step) + rollout_end = t + ROLLOUT_SECONDS + store.append( + PhaseEvent( + name="rollout", t0=t, t1=rollout_end, node=MANAGER_NODE, gpus=[], rank=-1, role=Role.ROLLOUT_MANAGER + ) + ) + for rank in range(gpus): + wait_end = rollout_end + 1 + 0.2 * rank + train_end = t + 80 + 0.5 * rank + store.append( + PhaseEvent( + name="train_wait", + t0=rollout_end, + t1=wait_end, + node=GPU_NODE, + gpus=[rank], + rank=rank, + role=Role.TRAIN, + ) + ) + store.append( + PhaseEvent( + name="actor_train", + t0=wait_end, + t1=train_end, + node=GPU_NODE, + gpus=[rank], + rank=rank, + role=Role.TRAIN, + ) + ) + store.append( + PhaseEvent( + name="update_weights", + t0=train_end, + t1=train_end + 4 + 0.3 * rank, + node=GPU_NODE, + gpus=[rank], + rank=rank, + role=Role.TRAIN, + ) + ) + store.append( + MetricsRecord( + ts=t + STEP_SECONDS - 5, + step_key="rollout/step", + step=step, + metrics={ + "perf/step_time": 95.0 - step, + "perf/wait_time_ratio": 0.4 - 0.1 * step, + "rollout/rewards_mean": 0.4 + 0.05 * step, + }, + ) + ) + + for offset in range(STEP_SECONDS): + ts = t + offset + for gpu in range(gpus): + drain = DRAIN_A if gpu < 2 else DRAIN_B + if offset < ROLLOUT_SECONDS: + util = 90 if offset < drain else 5 + elif offset < 80: + util = 95 + else: + util = 30 + if step == SPIKE["step"] and gpu == SPIKE["gpu"] and offset == SPIKE["offset"]: + util = SPIKE["util"] + store.append( + GpuSample(ts=ts, node=GPU_NODE, gpu=gpu, util=util, mem_mb=60_000 + 100 * gpu, power_w=600) + ) + + for offset in range(0, STEP_SECONDS, 2): + ts = t + offset + for addr, drain in ((ENGINE_A, DRAIN_A), (engine_b_addr(ts), DRAIN_B)): + running = 16.0 if offset < min(drain, ROLLOUT_SECONDS) else 0.0 + store.append( + EngineSample(ts=ts, addr=addr, metric="sglang_num_running_reqs", labels={}, value=running) + ) + store.append( + EngineSample(ts=ts, addr=addr, metric="sglang_gen_throughput", labels={}, value=running * 50.0) + ) + + _trajectory_events(store, truth, samples_per_step) + store.flush() + return truth diff --git a/tests/fast/dashboard/fixtures/engine_metrics_direct.txt b/tests/fast/dashboard/fixtures/engine_metrics_direct.txt new file mode 100644 index 0000000000..dfe63f3c54 --- /dev/null +++ b/tests/fast/dashboard/fixtures/engine_metrics_direct.txt @@ -0,0 +1,9 @@ +# HELP sglang:num_running_reqs The number of running requests. +# TYPE sglang:num_running_reqs gauge +sglang:num_running_reqs{model_name="qwen3-30b"} 13.0 +# HELP sglang:cache_hit_rate The prefix cache hit rate. +# TYPE sglang:cache_hit_rate gauge +sglang:cache_hit_rate{model_name="qwen3-30b"} 0.71 +# HELP sglang:uptime_seconds Not whitelisted. +# TYPE sglang:uptime_seconds gauge +sglang:uptime_seconds{model_name="qwen3-30b"} 99.0 diff --git a/tests/fast/dashboard/fixtures/engine_metrics_router.txt b/tests/fast/dashboard/fixtures/engine_metrics_router.txt new file mode 100644 index 0000000000..5bd34e2b77 --- /dev/null +++ b/tests/fast/dashboard/fixtures/engine_metrics_router.txt @@ -0,0 +1,17 @@ +# HELP sglang_num_running_reqs The number of running requests. +# TYPE sglang_num_running_reqs gauge +sglang_num_running_reqs{model_name="qwen3-30b",worker_addr="http://10.0.0.1:15000"} 42.0 +sglang_num_running_reqs{model_name="qwen3-30b",worker_addr="http://10.0.0.1:15004"} 7.0 +# HELP sglang_gen_throughput The generation throughput (token/s). +# TYPE sglang_gen_throughput gauge +sglang_gen_throughput{model_name="qwen3-30b",worker_addr="http://10.0.0.1:15000"} 2100.5 +sglang_gen_throughput{model_name="qwen3-30b",worker_addr="http://10.0.0.1:15004"} 350.25 +# HELP sglang_token_usage KV cache usage. +# TYPE sglang_token_usage gauge +sglang_token_usage{model_name="qwen3-30b",engine_type="prefill",worker_addr="http://10.0.0.1:15000"} 0.63 +# HELP sglang_prompt_tokens_total Number of prefill tokens processed. +# TYPE sglang_prompt_tokens_total counter +sglang_prompt_tokens_total{model_name="qwen3-30b",worker_addr="http://10.0.0.1:15000"} 5.4321e+06 +# HELP sglang_uptime_seconds Not in the whitelist either. +# TYPE sglang_uptime_seconds gauge +sglang_uptime_seconds{model_name="qwen3-30b",worker_addr="http://10.0.0.1:15000"} 12345.0 diff --git a/tests/fast/dashboard/test_collector.py b/tests/fast/dashboard/test_collector.py new file mode 100644 index 0000000000..5d740137b7 --- /dev/null +++ b/tests/fast/dashboard/test_collector.py @@ -0,0 +1,268 @@ +import logging +import time +from pathlib import Path + +from tests.fast.dashboard.dummy_telemetry import BASE_TS, dump_dummy_telemetry + +from miles.dashboard.collector import CollectorConfig, DashboardCollector +from miles.dashboard.sglang_scraper import ScrapeMode +from miles.dashboard.store import ( + EngineInfo, + GpuSample, + MetricsRecord, + MetricStore, + PhaseEvent, + Role, + Stream, + TopologySnapshot, +) + +ROUTER_FIXTURE = (Path(__file__).parent / "fixtures" / "engine_metrics_router.txt").read_text() + + +def make_collector(tmp_path, **kwargs) -> DashboardCollector: + config = kwargs.pop("config", None) or CollectorConfig( + dashboard_dir=str(tmp_path / "dashboard"), run_name="collector-test", start_ts=1.0 + ) + return DashboardCollector(config, **kwargs) + + +def test_collector_satisfies_dummy_telemetry_contract(tmp_path): + """Records routed through the collector's RPC surface must be interpreted + by the timeline queries exactly like the fixture's direct writes.""" + dump_dummy_telemetry(tmp_path / "direct") + reference = MetricStore.load(tmp_path / "direct" / "dashboard") + + collector = make_collector( + tmp_path, + config=CollectorConfig( + # start_ts must match the fixture's: the synthesized initialize + # band in phases_by_lane derives from meta.start_ts + dashboard_dir=str(tmp_path / "via_collector" / "dashboard"), + run_name="dummy-telemetry", + start_ts=BASE_TS, + ), + ) + for record in reference.records[Stream.METRICS]: + collector.push_metrics(record) + collector.push_phases(reference.iter_records(Stream.PHASES)) + for snapshot in reference.records[Stream.TOPOLOGY]: + collector.update_topology(snapshot) + by_node: dict[str, list[GpuSample]] = {} + for sample in reference.iter_records(Stream.GPU_UTIL): + by_node.setdefault(sample.node, []).append(sample) + for node, batch in by_node.items(): + collector.push_gpu_samples(node, batch) + for sample in reference.iter_records(Stream.ENGINE_SERIES): + collector._append(sample) # the scraper sink path + collector.flush() + + replayed = MetricStore.load(tmp_path / "via_collector" / "dashboard") + assert replayed.lanes() == reference.lanes() + assert replayed.topology_windows() == reference.topology_windows() + assert replayed.phases_by_lane() == reference.phases_by_lane() + assert replayed.gpu_series() == reference.gpu_series() + assert replayed.engine_series("sglang_num_running_reqs") == reference.engine_series("sglang_num_running_reqs") + assert replayed.bubbles() == reference.bubbles() + assert replayed.meta.run_name == reference.meta.run_name + + +def _engine(addr: str) -> EngineInfo: + return EngineInfo(addr=addr, worker_type="regular", engine_rank=0, gpus=[["n", 0]], gpu_uuids=[None]) + + +def test_topology_dedup_keeps_changes_only(tmp_path): + collector = make_collector(tmp_path) + collector.update_topology(TopologySnapshot(ts=1.0, engines=[_engine("http://a:1")])) + collector.update_topology(TopologySnapshot(ts=2.0, engines=[_engine("http://a:1")])) # same engines + collector.update_topology(TopologySnapshot(ts=3.0, engines=[_engine("http://a:2")])) # restart + collector.flush() + + loaded = MetricStore.load(collector.config.dashboard_dir) + assert [w["t0"] for w in loaded.topology_windows()] == [1.0, 3.0] + + +def test_flush_thread_persists_periodically(tmp_path): + config = CollectorConfig( + dashboard_dir=str(tmp_path / "dashboard"), run_name="r", start_ts=0.0, flush_interval_seconds=0.05 + ) + collector = DashboardCollector(config) + collector.start() + collector.push_metrics(MetricsRecord(ts=1.0, step_key="rollout/step", step=0, metrics={"a": 1})) + time.sleep(0.2) + try: + assert len(MetricStore.load(config.dashboard_dir).records[Stream.METRICS]) == 1 + finally: + collector.shutdown() + + +def test_shutdown_flushes_and_stops_scraper(tmp_path): + collector = make_collector(tmp_path, scraper_http_get=lambda url, timeout: ROUTER_FIXTURE) + collector.set_router("http://router:3000", use_miles_router=False) + assert collector._scraper is not None + collector.push_metrics(MetricsRecord(ts=1.0, step_key="rollout/step", step=0, metrics={"a": 1})) + collector.shutdown() + + assert collector._scraper._stop_event.is_set() + loaded = MetricStore.load(collector.config.dashboard_dir) + assert len(loaded.records[Stream.METRICS]) == 1 + + +def test_scraper_mode_resolution_and_repoint(tmp_path): + collector = make_collector(tmp_path, scraper_http_get=lambda url, timeout: ROUTER_FIXTURE) + collector.set_router("http://router:3000", use_miles_router=False) + assert collector._scraper.mode == ScrapeMode.ROUTER + first = collector._scraper + + collector.set_router("http://router:3000", use_miles_router=False) # no-op + assert collector._scraper is first + + collector.update_topology(TopologySnapshot(ts=1.0, engines=[_engine("http://e:1")])) + collector.set_router("http://router2:3000", use_miles_router=True) # router restarted, miles router now + assert collector._scraper is not first + assert first._stop_event.is_set() + assert collector._scraper.mode == ScrapeMode.DIRECT + # actor-registered engine first, then the externals remembered from the + # first router's scrapes (never actor-known -> kept as scrape targets) + assert collector._scraper.engine_addrs() == [ + "http://e:1", + "http://10.0.0.1:15000", + "http://10.0.0.1:15004", + ] + collector.shutdown() + + +def test_disk_failure_is_loud_and_buffers_stay_bounded(tmp_path, monkeypatch, caplog): + collector = make_collector(tmp_path) + monkeypatch.setattr(DashboardCollector, "MAX_BUFFERED_PER_STREAM", 100) + monkeypatch.setattr(MetricStore, "flush", lambda self: (_ for _ in ()).throw(OSError("no space left"))) + + with caplog.at_level(logging.ERROR): + for step in range(150): + collector.push_metrics(MetricsRecord(ts=float(step), step_key="rollout/step", step=step, metrics={})) + collector.flush() + + assert collector._store.buffered_count(Stream.METRICS) <= 105 # bounded despite dead disk + messages = [r.message for r in caplog.records] + assert any("failed" in m for m in messages) # the OSError is loud + assert any("dropped" in m for m in messages) # so is the data loss + + +def test_prometheus_forwarding_snapshot(tmp_path): + updates: list[dict] = [] + + class FakeRemote: + def remote(self, snapshot): + updates.append(snapshot) + + class FakeHandle: + update = FakeRemote() + + config = CollectorConfig( + dashboard_dir=str(tmp_path / "dashboard"), run_name="r", start_ts=0.0, forward_prometheus=True + ) + collector = DashboardCollector(config, prometheus_handle_factory=lambda: FakeHandle()) + collector.push_gpu_samples( + "10.0.0.1", [GpuSample(ts=1.0, node="10.0.0.1", gpu=0, util=87, mem_mb=1000, power_w=600)] + ) + collector.push_phases( + [PhaseEvent(name="actor_train", t0=1.0, t1=61.0, node="10.0.0.1", gpus=[0], rank=0, role=Role.TRAIN)] + ) + collector.flush() + collector.flush() # latest-value cache must survive the first flush + + assert len(updates) == 2 + snapshot = updates[-1] + assert snapshot["dashboard/gpu_10_0_0_1_0_util"] == 87.0 + assert snapshot["dashboard/phase_actor_train_seconds"] == 60.0 + assert all("." not in k and ":" not in k for k in snapshot) + + +def test_forwarding_disabled_by_default(tmp_path): + called = [] + collector = make_collector(tmp_path, prometheus_handle_factory=lambda: called.append(1)) + collector.flush() + assert called == [] + + +def test_sampler_reconcile_spawns_late_nodes_and_skips_nvml_less(tmp_path, monkeypatch): + from miles.dashboard import collector as collector_mod + + nodes = [("id-a", "10.0.0.1")] + spawned, killed = [], [] + + def fake_spawn(node_id, node_ip, interval): + spawned.append(node_id) + return None if node_id == "id-nvmlless" else f"handle-{node_id}" + + monkeypatch.setattr(collector_mod, "_list_gpu_nodes", lambda: nodes) + monkeypatch.setattr(collector_mod, "_spawn_sampler", fake_spawn) + monkeypatch.setattr(collector_mod, "_kill_sampler", killed.append) + + collector = make_collector(tmp_path) + collector._reconcile_samplers() + assert spawned == ["id-a"] + + # a node joins late (fully async rollout node): next tick picks it up + nodes.append(("id-b", "10.0.0.2")) + collector._reconcile_samplers() + assert spawned == ["id-a", "id-b"] + + # NVML-less node: spawn once, remember None, never retry + nodes.append(("id-nvmlless", "10.0.0.3")) + collector._reconcile_samplers() + collector._reconcile_samplers() + assert spawned.count("id-nvmlless") == 1 + + # node restart: new NodeID replaces the old entry and respawns + nodes[0] = ("id-a2", "10.0.0.1") + collector._reconcile_samplers() + assert "id-a2" in collector._samplers and "id-a" not in collector._samplers + + collector.shutdown() + assert sorted(h for h in killed) == ["handle-id-a2", "handle-id-b"] + + +def test_external_engines_synthesized_from_scrapes(tmp_path): + """Pure sglang engines (no miles actor) never reach register_engines; + the scrape itself is the topology source (disagg report 2026-07-14).""" + from miles.dashboard.store import EngineSample, Stream, TopologySnapshot + + collector = make_collector(tmp_path) + collector.push_metrics(MetricsRecord(ts=1.0, step_key="rollout/step", step=0, metrics={})) + for addr in ("http://10.1.0.5:15000", "http://10.1.0.6:15000"): + collector._append(EngineSample(ts=2.0, addr=addr, metric="sglang_num_running_reqs", labels={}, value=1.0)) + collector._sync_external_topology() # the flush-loop step + collector._sync_external_topology() # steady state: no duplicate snapshot + collector.flush() + + store = MetricStore.load(collector.config.dashboard_dir) + [snapshot] = store.records[Stream.TOPOLOGY] + assert [(e.addr, e.worker_type, e.gpus) for e in snapshot.engines] == [ + ("http://10.1.0.5:15000", "external", []), + ("http://10.1.0.6:15000", "external", []), + ] + + # a real actor registration keeps the synthetic externals merged in + from miles.dashboard.store import EngineInfo + + real = TopologySnapshot( + ts=3.0, + engines=[ + EngineInfo( + addr="http://10.1.0.5:15000", + worker_type="regular", + engine_rank=0, + gpus=[["10.1.0.5", 0]], + gpu_uuids=[None], + ) + ], + ) + collector.update_topology(real) + collector.flush() + snapshots = MetricStore.load(collector.config.dashboard_dir).records[Stream.TOPOLOGY] + latest = snapshots[-1] + assert {e.addr: e.worker_type for e in latest.engines} == { + "http://10.1.0.5:15000": "regular", + "http://10.1.0.6:15000": "external", + } diff --git a/tests/fast/dashboard/test_core_integration.py b/tests/fast/dashboard/test_core_integration.py new file mode 100644 index 0000000000..573fa71894 --- /dev/null +++ b/tests/fast/dashboard/test_core_integration.py @@ -0,0 +1,118 @@ +"""Integration of the dashboard with miles core, on a local Ray instance. + +Exercises the REAL activation path: tracking registry -> MilesDashboardBackend +-> init_dashboard (named collector actor + samplers) -> tracking_utils.log +fan-out -> Timer phase sink -> finish_tracking, then loads the produced +directory with the store and serves it. GPU samples appear only on hosts with +NVML; everything else must work identically on a laptop and a GPU node. +""" + +import time +from argparse import Namespace + +import pytest + +from miles.dashboard.dump_reader import DumpReader +from miles.dashboard.server import make_app +from miles.dashboard.store import MetricStore, Role, Stream +from miles.utils.timer import Timer, timer + + +def _args(tmp_path) -> Namespace: + return Namespace( + use_miles_dashboard=True, + dump_details=str(tmp_path), + wandb_group="wiring-e2e", + use_rollout_entropy=True, + use_miles_router=False, + dashboard_flush_interval=0.2, + dashboard_gpu_sample_interval=0.2, + dashboard_sglang_scrape_interval=2.0, + dashboard_sglang_scrape_mode="auto", + dashboard_sglang_metrics=None, + dashboard_forward_prometheus=False, + ) + + +def test_tracking_backend_end_to_end_local_ray(tmp_path): + ray = pytest.importorskip("ray") + from miles.dashboard import backend, hooks + from miles.utils.tracking_utils import finish_tracking, init_tracking + from miles.utils.tracking_utils import log as tracking_log + + ray.init(num_cpus=2, include_dashboard=False, ignore_reinit_error=True, logging_level="ERROR") + saved_sinks = list(Timer().event_sinks) + args = _args(tmp_path) + try: + # driver: registry picks MilesDashboardBackend from the args flag + init_tracking(args) + assert backend.current_collector() is not None + + # metric fan-out through the real tracking API + tracking_log(args, {"rollout/rewards_mean": 0.5, "rollout/step": 3}, step_key="rollout/step") + + # per-rank phase lane path (driver doubles as a "rank" here) + hooks.register_train_actor(args) + with timer("actor_train"): + time.sleep(0.05) + with timer("update_weights"): + time.sleep(0.01) + + finish_tracking() # flushes sink + synchronous collector shutdown + finally: + Timer().event_sinks[:] = saved_sinks + ray.shutdown() + + store = MetricStore.load(tmp_path / "dashboard") + assert store.meta.run_name == "wiring-e2e" + + [metric] = store.records[Stream.METRICS] + assert metric.step == 3 and metric.metrics["rollout/rewards_mean"] == 0.5 + + phases = {event.name: event for event in store.iter_records(Stream.PHASES)} + assert {"actor_train", "update_weights"} <= set(phases) + actor_train = phases["actor_train"] + assert actor_train.role == Role.TRAIN + assert actor_train.t1 - actor_train.t0 == pytest.approx(0.05, abs=0.05) + + app_client = pytest.importorskip("fastapi.testclient").TestClient( + make_app(store, DumpReader(tmp_path), follow=False) + ) + meta = app_client.get("/api/meta").json() + assert meta["run_name"] == "wiring-e2e" + assert meta["capabilities"]["has_timeline"] is True + assert "rollout/rewards_mean" in meta["metric_keys"] + + +def test_registry_contains_dashboard_backend(): + from miles.utils.tracking_utils.base import BACKEND_REGISTRY, MilesDashboardBackend + + cls, flag = BACKEND_REGISTRY["miles_dashboard"] + assert cls is MilesDashboardBackend + assert flag == "use_miles_dashboard" + + +def test_engine_topology_gpu_range_logic(): + # the pure slice of SGLangEngine.get_topology_info: node-physical range + pytest.importorskip("sglang") + from miles.backends.sglang_utils.sglang_engine import SGLangEngine + + engine = SGLangEngine.__new__(SGLangEngine) + engine.args = Namespace(num_gpus_per_node=8) + engine.base_gpu_id = 4 + engine.num_gpus_per_engine = 2 + engine.worker_type = "regular" + engine.node_rank = 0 + engine.server_host = "10.0.0.9" + engine.server_port = 15000 + + info = engine.get_topology_info() + assert info["url"] == "http://10.0.0.9:15000" + assert info["gpu_ids"] == [4, 5] + assert len(info["gpu_uuids"]) == 2 + assert info["worker_type"] == "regular" + + # multi-node engine: each member covers its whole node (base 0, capped per node) + engine.base_gpu_id = 0 + engine.num_gpus_per_engine = 16 + assert engine.get_topology_info()["gpu_ids"] == list(range(8)) diff --git a/tests/fast/dashboard/test_dashboard_args.py b/tests/fast/dashboard/test_dashboard_args.py new file mode 100644 index 0000000000..c27619fd55 --- /dev/null +++ b/tests/fast/dashboard/test_dashboard_args.py @@ -0,0 +1,73 @@ +import argparse +import logging + +import pytest + +from miles.dashboard.args import add_dashboard_arguments, collector_config_from_args, validate_dashboard_args +from miles.dashboard.sglang_scraper import DEFAULT_METRIC_WHITELIST + + +def parse(argv): + parser = argparse.ArgumentParser() + add_dashboard_arguments(parser) + return parser.parse_args(argv) + + +def full_args(**overrides): + args = parse(["--use-miles-dashboard"]) + args.dump_details = "/tmp/dump" + args.use_rollout_entropy = True + args.wandb_group = "my-run" + args.colocate = True + for key, value in overrides.items(): + setattr(args, key, value) + return args + + +def test_defaults(): + args = parse([]) + assert args.use_miles_dashboard is False + assert args.dashboard_flush_interval == 5.0 + assert args.dashboard_gpu_sample_interval == 1.0 + assert args.dashboard_sglang_scrape_mode == "auto" + assert args.dashboard_sglang_metrics is None + assert args.dashboard_forward_prometheus is False + + +def test_validate_requires_dump_details(): + args = full_args(dump_details=None) + with pytest.raises(AssertionError, match="dump-details"): + validate_dashboard_args(args) + + +def test_validate_warns_without_entropy(caplog): + with caplog.at_level(logging.WARNING): + validate_dashboard_args(full_args(use_rollout_entropy=False)) + assert any("use-rollout-entropy" in r.message for r in caplog.records) + + caplog.clear() + with caplog.at_level(logging.WARNING): + validate_dashboard_args(full_args()) + assert caplog.records == [] + + +def test_validate_disabled_checks_nothing(): + validate_dashboard_args(parse([])) # no dump_details attribute needed + + +def test_collector_config_from_args(): + config = collector_config_from_args(full_args(), start_ts=123.0) + assert config.dashboard_dir == "/tmp/dump/dashboard" + assert config.run_name == "my-run" + assert config.start_ts == 123.0 + assert config.metric_whitelist == DEFAULT_METRIC_WHITELIST + assert config.args_snapshot["colocate"] is True + assert "hf_checkpoint" not in config.args_snapshot # absent attrs are skipped + + +def test_collector_config_whitelist_override_and_run_name_fallback(): + args = full_args(wandb_group=None) + args.dashboard_sglang_metrics = "sglang_gen_throughput,sglang_token_usage" + config = collector_config_from_args(args, start_ts=0.0) + assert config.metric_whitelist == ("sglang_gen_throughput", "sglang_token_usage") + assert config.run_name == "miles-run" diff --git a/tests/fast/dashboard/test_dump_reader_join.py b/tests/fast/dashboard/test_dump_reader_join.py new file mode 100644 index 0000000000..bb10798aa9 --- /dev/null +++ b/tests/fast/dashboard/test_dump_reader_join.py @@ -0,0 +1,190 @@ +import os +import time +from collections import defaultdict + +import pytest +import torch +from tests.fast.dashboard.dummy_dump import dump_dummy_run + +from miles.dashboard.dump_reader import DumpReader, DumpStillWriting, TrainRow +from miles.utils.types import Sample + + +@pytest.fixture +def run(tmp_path): + truth = dump_dummy_run(tmp_path, steps=3, dp_size=2, tp_dup=2, with_eval=True) + return DumpReader(tmp_path), truth + + +def test_rollout_ids_split_and_sorted(run): + reader, truth = run + ids = reader.rollout_ids() + assert ids.train == [0, 1, 2] + assert ids.eval == truth.eval_ids == [0] + + +def test_fresh_rollout_hidden_until_train_companion_exists(run): + reader, _ = run + fresh = reader.rollout_dir / "3.pt" + torch.save(dict(rollout_id=3, samples=[]), fresh) # mtime = now + assert 3 not in reader.rollout_ids().train + + (reader.train_dir / "3_0.pt").touch() # train companion written strictly later + assert 3 in reader.rollout_ids().train + + +def test_fresh_eval_hidden_until_aged(run): + reader, _ = run + fresh = reader.rollout_dir / "eval_2.pt" + torch.save(dict(rollout_id=2, samples=[]), fresh) + assert 2 not in reader.rollout_ids().eval + + stamp = time.time() - 100 + os.utime(fresh, (stamp, stamp)) + assert 2 in reader.rollout_ids().eval + + +def test_load_joined_full_coverage(run): + reader, truth = run + for rollout_id in reader.rollout_ids().train: + joined = reader.load_joined(rollout_id) + assert len(joined.samples) == truth.n_samples_per_step + assert joined.train_coverage == 1.0 + assert set(joined.train_rows) == {s.index for s in joined.samples} + assert all(isinstance(s, Sample) for s in joined.samples) + assert all(isinstance(r, TrainRow) for r in joined.train_rows.values()) + # Sample deserialization went through miles' own from_dict: enums restored. + statuses = {s.status for s in reader.load_joined(0).samples} + assert statuses <= {Sample.Status.COMPLETED, Sample.Status.TRUNCATED} + + +def test_join_row_matches_sample(run): + reader, _ = run + joined = reader.load_joined(1) + sample_of = {s.index: s for s in joined.samples} + for index, row in joined.train_rows.items(): + sample = sample_of[index] + assert row.response_length == sample.response_length + assert row.total_length == len(sample.tokens) + assert torch.equal(row.tokens, torch.tensor(sample.tokens)) + assert torch.allclose(row.rollout_log_probs, torch.tensor(sample.rollout_log_probs)) + assert len(row.log_probs) == len(row.advantages) == len(row.loss_mask) == sample.response_length + + +def test_raw_reward_uses_batch_global_indexing(run): + # Regression test: raw_reward is the one batch-global column + # (split_train_data_by_dp ships it unpartitioned, "splited at train side"); + # with shuffled balanced partitions, indexing it by shard row silently + # misattributes rewards across samples. + reader, _ = run + for rollout_id in reader.rollout_ids().train: + joined = reader.load_joined(rollout_id) + for sample in joined.samples: + assert joined.train_rows[sample.index].raw_reward == sample.reward + + +def test_rewards_went_through_real_group_normalization(run): + # The dummy pipeline runs the real GRPO reward post-processing: per-group + # mean-centering makes shard-local `rewards` sum to ~0 within each group. + reader, _ = run + joined = reader.load_joined(0) + group_of = {s.index: s.group_index for s in joined.samples} + group_rewards = defaultdict(list) + for index, row in joined.train_rows.items(): + group_rewards[group_of[index]].append(row.reward) + for rewards in group_rewards.values(): + assert abs(sum(rewards)) < 1e-4 + + +def test_tp_duplicates_keep_first_rank(run): + reader, truth = run + joined = reader.load_joined(0) + # dp_size=2, tp_dup=2: shard 0 lives on ranks {0,1}, shard 1 on ranks {2,3}; + # dedup must keep the lowest rank of each shard. + assert {row.rank for row in joined.train_rows.values()} == {0, 2} + for shard_idx, indices in enumerate(truth.shard_indices[0]): + for index in indices: + assert joined.train_rows[index].rank == shard_idx * 2 + + +def test_train_index_absent_from_rollout_asserts(run): + reader, _ = run + path = reader.train_dir / "0_0.pt" + pack = torch.load(path, weights_only=False) + pack["rollout_data"]["sample_indices"][0] = 9999 + torch.save(pack, path) + + with pytest.raises(AssertionError, match="9999 absent"): + reader.load_joined(0) + + +def test_inconsistent_tp_duplicate_asserts(run): + reader, _ = run + path = reader.train_dir / "0_1.pt" # TP duplicate of rank 0 + pack = torch.load(path, weights_only=False) + pack["rollout_data"]["response_lengths"][0] += 1 + torch.save(pack, path) + + with pytest.raises(AssertionError, match="disagrees"): + reader.load_joined(0) + + +def test_eval_join_has_no_train_rows(run): + reader, truth = run + joined = reader.load_joined(0, evaluation=True) + assert joined.evaluation + assert joined.train_rows == {} + assert len(joined.samples) == truth.n_samples_per_step // 2 + + +def test_still_writing_vs_corruption(run): + reader, _ = run + path = reader.rollout_dir / "9.pt" + path.write_bytes(b"garbage") # fresh mtime: a torch.save in progress + assert 9 not in reader.rollout_ids().train + with pytest.raises(DumpStillWriting): + reader.load_joined(9) + + stamp = time.time() - 100 + os.utime(path, (stamp, stamp)) # old + unloadable = real corruption + with pytest.raises(Exception) as exc_info: + reader.load_joined(9) + assert not isinstance(exc_info.value, DumpStillWriting) + + +def test_missing_rollout_file_raises(run): + reader, _ = run + with pytest.raises(FileNotFoundError): + reader.load_joined(42) + + +def test_optional_columns_absent(tmp_path): + dump_dummy_run(tmp_path, steps=1, with_entropy=False, with_eval=False) + joined = DumpReader(tmp_path).load_joined(0) + for row in joined.train_rows.values(): + assert row.entropy is None and row.ref_entropy is None + assert row.advantages is not None # unaffected columns still load + + +def test_empty_dump_dir(tmp_path): + ids = DumpReader(tmp_path).rollout_ids() + assert ids.train == [] and ids.eval == [] + + +@pytest.mark.skipif( + "MILES_DASHBOARD_REALDATA_DIR" not in os.environ, + reason="set MILES_DASHBOARD_REALDATA_DIR to a real --dump-details dir", +) +def test_realdata_join(): + reader = DumpReader(os.environ["MILES_DASHBOARD_REALDATA_DIR"]) + ids = reader.rollout_ids() + assert ids.train, "no rollout dumps found" + for rollout_id in ids.train: + joined = reader.load_joined(rollout_id) + assert joined.train_coverage == 1.0 + row = next(iter(joined.train_rows.values())) + assert len(row.log_probs) == row.response_length + # raw_reward correspondence must hold on real data too + sample_of = {s.index: s for s in joined.samples} + for index, train_row in joined.train_rows.items(): + assert train_row.raw_reward == sample_of[index].reward diff --git a/tests/fast/dashboard/test_dump_reader_views.py b/tests/fast/dashboard/test_dump_reader_views.py new file mode 100644 index 0000000000..f37ce03a06 --- /dev/null +++ b/tests/fast/dashboard/test_dump_reader_views.py @@ -0,0 +1,231 @@ +import os +import statistics +import time + +import pytest +from tests.fast.dashboard.dummy_dump import dump_dummy_run + +from miles.dashboard.dump_reader import DumpReader + +REMOVED = (3,) # within-step positions marked remove_sample=True by the fixture + + +@pytest.fixture +def reader(tmp_path): + dump_dummy_run(tmp_path, steps=2, dp_size=2, tp_dup=2, with_eval=True, remove_sample_indices=REMOVED) + return DumpReader(tmp_path) + + +def _df_row(df, sample_index): + matches = df.filter(df["sample_index"] == sample_index) + assert matches.height == 1 + return matches.row(0, named=True) + + +def test_summary_matches_hand_computed(reader): + joined = reader.load_joined(0) + df = reader.summary(0) + assert df.height == len(joined.samples) + + for sample in joined.samples: + entry = _df_row(df, sample.index) + row = joined.train_rows[sample.index] + assert entry["group_index"] == sample.group_index + assert entry["response_length"] == sample.response_length + assert entry["raw_reward"] == sample.reward + assert entry["normalized_reward"] == pytest.approx(row.reward) + assert entry["dumped_rank"] == row.rank + + mask = row.loss_mask > 0 + if mask.any(): + expected = (row.log_probs - row.rollout_log_probs).abs()[mask] + assert entry["mean_abs_lp_diff"] == pytest.approx(float(expected.mean()), rel=1e-5) + assert entry["max_abs_lp_diff"] == pytest.approx(float(expected.max()), rel=1e-5) + assert entry["mean_imp_ratio"] == pytest.approx( + float((row.log_probs - row.rollout_log_probs).exp()[mask].mean()), rel=1e-5 + ) + assert entry["mean_entropy"] == pytest.approx(float(row.entropy[mask].mean()), rel=1e-5) + + +def test_summary_removed_sample_has_no_masked_stats(reader): + df = reader.summary(0) + entry = _df_row(df, REMOVED[0]) # step 0: Sample.index == within-step position + assert entry["remove_sample"] is True + for column in ("mean_entropy", "mean_abs_lp_diff", "mean_imp_ratio", "adv_mean", "return_mean"): + assert entry[column] is None, column + + +def test_summary_parquet_cache_hit_and_invalidation(reader, monkeypatch): + df = reader.summary(0) + assert (reader.cache_dir / "rollout_0.parquet").exists() + + def boom(self, path): + raise RuntimeError("recompute attempted") + + monkeypatch.setattr(DumpReader, "_torch_load", boom) + reader._joined_cache.clear() + assert reader.summary(0).equals(df) # served from parquet, no dump loads + + stamp = time.time() + 5 + os.utime(next(reader.train_dir.glob("0_*.pt")), (stamp, stamp)) + with pytest.raises(RuntimeError, match="recompute attempted"): + reader.summary(0) # source changed -> cache invalidated -> recompute + + +def test_groups(reader): + joined = reader.load_joined(0) + groups_df = reader.groups(0) + rewards_by_group = {} + for sample in joined.samples: + rewards_by_group.setdefault(sample.group_index, []).append(sample.reward) + assert groups_df.height == len(rewards_by_group) + + for entry in groups_df.iter_rows(named=True): + rewards = rewards_by_group[entry["group_index"]] + assert entry["n"] == len(rewards) + assert entry["reward_mean"] == pytest.approx(statistics.mean(rewards)) + assert entry["reward_std"] == pytest.approx(statistics.stdev(rewards)) + assert entry["zero_std"] == (statistics.stdev(rewards) <= 1e-12) + + +def test_step_aggregates(reader): + aggregates = reader.step_aggregates() + assert aggregates["rollout_id"].to_list() == [0, 1] + for entry in aggregates.iter_rows(named=True): + df = reader.summary(entry["rollout_id"]) + assert entry["n_samples"] == df.height + assert entry["reward_mean"] == pytest.approx(df["raw_reward"].mean()) + assert entry["mean_abs_lp_diff"] is not None + + +def test_tokens_full_range(reader): + joined = reader.load_joined(0) + sample = joined.samples[0] + payload = reader.tokens(0, sample.index) + prompt_len = len(sample.tokens) - sample.response_length + + assert payload["total_len"] == len(sample.tokens) + assert payload["prompt_len"] == payload["response_offset"] == prompt_len + assert payload["token_ids"] == sample.tokens + assert len(payload["train_log_probs"]) == sample.response_length + assert len(payload["imp_ratio"]) == sample.response_length + assert payload["rollout_log_probs"] == pytest.approx(sample.rollout_log_probs) + row = joined.train_rows[sample.index] + assert payload["lp_diff"][0] == pytest.approx(float(row.log_probs[0] - row.rollout_log_probs[0])) + + +def test_tokens_window_straddles_response_boundary(reader): + sample = reader.load_joined(0).samples[0] + prompt_len = len(sample.tokens) - sample.response_length + + payload = reader.tokens(0, sample.index, start=prompt_len - 2, end=prompt_len + 3) + assert len(payload["token_ids"]) == 5 + assert payload["response_offset"] == 2 + assert len(payload["train_log_probs"]) == 3 # stats only cover the response overlap + + prompt_only = reader.tokens(0, sample.index, start=0, end=2) + assert prompt_only["train_log_probs"] == [] + assert prompt_only["response_offset"] == 2 + + clamped = reader.tokens(0, sample.index, start=0, end=10**6) + assert clamped["end"] == len(sample.tokens) + + with pytest.raises(ValueError, match="empty token range"): + reader.tokens(0, sample.index, start=5, end=5) + + with pytest.raises(KeyError, match="unknown sample_index"): + reader.tokens(0, 10**6) + + +def test_tokens_decode_via_dumped_tokenizer(reader): + payload = reader.tokens(0, 0) + assert payload["token_text"] == [f"t{tid}" for tid in payload["token_ids"]] + + +def test_tokens_eval_sample_has_rollout_side_only(reader): + payload = reader.tokens(0, 0, evaluation=True) + assert payload["rollout_log_probs"] is not None + for key in ("train_log_probs", "lp_diff", "imp_ratio", "entropy", "advantages", "loss_mask"): + assert payload[key] is None, key + + +def test_tokens_without_entropy_or_tokenizer(tmp_path): + dump_dummy_run(tmp_path, steps=1, with_entropy=False, with_eval=False, with_tokenizer=False) + reader = DumpReader(tmp_path) + payload = reader.tokens(0, 0) + assert payload["entropy"] is None and payload["ref_entropy"] is None + assert payload["train_log_probs"] is not None + assert payload["token_text"] is None + assert _df_row(reader.summary(0), 0)["mean_entropy"] is None + + +def test_joined_lru_eviction(tmp_path): + dump_dummy_run(tmp_path, steps=2, with_eval=False) + reader = DumpReader(tmp_path, tensor_lru=1) + first = reader.joined(0) + assert reader.joined(0) is first # cache hit + reader.joined(1) + assert list(reader._joined_cache) == [(1, False)] # 0 evicted + + +@pytest.mark.skipif( + "MILES_DASHBOARD_REALDATA_DIR" not in os.environ, + reason="set MILES_DASHBOARD_REALDATA_DIR to a real --dump-details dir", +) +def test_realdata_views(tmp_path): + reader = DumpReader(os.environ["MILES_DASHBOARD_REALDATA_DIR"], cache_dir=tmp_path) + df = reader.summary(0) + assert df.height == 256 + assert df["truncated"].cast(int).sum() == 112 # measured on qwen30b-dash step 0 + assert df["mean_abs_lp_diff"].null_count() == 0 + assert df["mean_entropy"].null_count() == 0 + + groups_df = reader.groups(0) + assert groups_df.height == 32 # 256 samples / 8 per prompt + + sample_index = int(df["sample_index"][0]) + payload = reader.tokens(0, sample_index, start=0, end=64) + assert len(payload["token_ids"]) == 64 + assert payload["token_text"] is not None + + aggregates = reader.step_aggregates() + assert aggregates.height == 5 + + +def test_summary_and_tokens_survive_dump_without_log_probs(tmp_path): + """A 744B-scale run may not dump train log_probs at all: everything + derived from them degrades to None instead of KeyError -> HTTP 404 + (disagg report 2026-07-14).""" + import torch + + from tests.fast.dashboard.dummy_dump import dump_dummy_run + + from miles.dashboard.dump_reader import DumpReader + + dump_dummy_run(tmp_path) + for shard in (tmp_path / "train_data").glob("0_*.pt"): + payload = torch.load(shard, map_location="cpu", weights_only=False) + assert payload["rollout_data"].pop("log_probs") is not None # must actually remove it + torch.save(payload, shard) + + reader = DumpReader(tmp_path) + df = reader.summary(0) + assert df["mean_abs_lp_diff"].is_null().all() + assert df["mean_imp_ratio"].is_null().all() + assert df["reward"].null_count() < df.height # the rest of the summary is intact + assert reader.step_aggregates().height >= 1 # the metrics page path +def test_staleness_and_agentic_columns(reader): + import polars as pl + + df = reader.summary(0) + agentic = df.filter(pl.col("sample_index") % 3 == 0) + plain = df.filter(pl.col("sample_index") % 3 != 0) + # dummy dump: every third sample is two-turn with mixed versions + one tool message + assert agentic["mixed_version"].all() and agentic["turns"].min() == 2 + assert agentic["tool_calls"].min() == 1 + assert not plain["mixed_version"].any() and plain["turns"].max() == 1 + assert plain["tool_calls"].is_null().all() # string prompts: nothing to count + assert (agentic["weight_version"].cast(pl.Int64) - agentic["weight_version_min"] == 1).all() + + aggregates = reader.step_aggregates() + assert 0 < aggregates["mixed_version_frac"][0] < 1 diff --git a/tests/fast/dashboard/test_gpu_sampler.py b/tests/fast/dashboard/test_gpu_sampler.py new file mode 100644 index 0000000000..ecc7dd2ead --- /dev/null +++ b/tests/fast/dashboard/test_gpu_sampler.py @@ -0,0 +1,136 @@ +import logging +import time + +import pytest + +from miles.dashboard.gpu_sampler import GpuSampler +from miles.dashboard.store import GpuSample + + +class FakeNvml: + """Just enough of pynvml for the sampler: handle == device index.""" + + def __init__(self, count=2, fail_init=False, failing_devices=()): + self.count = count + self.fail_init = fail_init + self.failing_devices = set(failing_devices) + + def nvmlInit(self): + if self.fail_init: + raise RuntimeError("driver/library version mismatch") + + def nvmlDeviceGetCount(self): + return self.count + + def nvmlDeviceGetHandleByIndex(self, index): + return index + + def nvmlDeviceGetUUID(self, handle): + return f"GPU-fake-{handle}" + + def nvmlDeviceGetUtilizationRates(self, handle): + if handle in self.failing_devices: + raise RuntimeError("GPU is lost") + return type("Util", (), {"gpu": 40 + handle})() + + def nvmlDeviceGetMemoryInfo(self, handle): + return type("Mem", (), {"used": (handle + 1) * 1024 * 1024 * 1024})() # GiB in bytes + + def nvmlDeviceGetPowerUsage(self, handle): + return 600_000 + handle # milliwatts + + +class PushSpy: + def __init__(self): + self.calls: list[tuple[str, list[GpuSample]]] = [] + + def __call__(self, node, batch): + self.calls.append((node, batch)) + + +def test_sample_once_converts_units(): + push = PushSpy() + sampler = GpuSampler(push, node="10.0.0.1", nvml=FakeNvml(count=2)) + assert sampler.available + assert sampler.gpu_uuids() == ["GPU-fake-0", "GPU-fake-1"] + + assert sampler.sample_once(ts=10.0) == 2 + sampler.flush() + [(node, batch)] = push.calls + assert node == "10.0.0.1" + assert batch == [ + GpuSample(ts=10.0, node="10.0.0.1", gpu=0, util=40, mem_mb=1024, power_w=600), + GpuSample(ts=10.0, node="10.0.0.1", gpu=1, util=41, mem_mb=2048, power_w=600), + ] + + +def test_flush_clears_buffer_and_skips_empty(): + push = PushSpy() + sampler = GpuSampler(push, node="n", nvml=FakeNvml(count=1)) + sampler.flush() # empty: no call + assert push.calls == [] + + sampler.sample_once(ts=1.0) + sampler.flush() + sampler.flush() # cleared: no duplicate push + assert len(push.calls) == 1 + + +def test_nvml_init_failure_disables_sampler(caplog): + push = PushSpy() + with caplog.at_level(logging.WARNING): + sampler = GpuSampler(push, node="n", nvml=FakeNvml(fail_init=True)) + assert not sampler.available + assert sampler.start() is False + assert sampler.sample_once(ts=1.0) == 0 + assert push.calls == [] + assert any("NVML unavailable" in r.message for r in caplog.records) + + +def test_failing_device_is_skipped_others_report(caplog): + push = PushSpy() + sampler = GpuSampler(push, node="n", nvml=FakeNvml(count=3, failing_devices={1})) + with caplog.at_level(logging.WARNING): + assert sampler.sample_once(ts=1.0) == 2 + sampler.flush() + [(_, batch)] = push.calls + assert [s.gpu for s in batch] == [0, 2] + assert any("skipping this tick" in r.message for r in caplog.records) + + +def test_thread_lifecycle_flushes_on_stop(): + push = PushSpy() + sampler = GpuSampler(push, node="n", interval=0.01, nvml=FakeNvml(count=1)) + assert sampler.start() is True + time.sleep(0.08) + sampler.stop() + assert push.calls, "stop() must flush buffered samples" + total = sum(len(batch) for _, batch in push.calls) + assert total >= 3 # ~8 ticks at 10ms; generous margin against scheduler jitter + + +def test_interval_must_be_positive(): + with pytest.raises(AssertionError): + GpuSampler(lambda n, b: None, node="n", interval=0, nvml=FakeNvml()) + + +def test_real_nvml_when_gpus_present(): + # Guards the FakeNvml against drifting from the real pynvml API surface; + # runs wherever a GPU + driver exist (devbox/CI-GPU), skips elsewhere. + pynvml = pytest.importorskip("pynvml") + try: + pynvml.nvmlInit() + pynvml.nvmlShutdown() + except Exception: + pytest.skip("no usable NVML device") + + push = PushSpy() + sampler = GpuSampler(push, node="local", nvml=pynvml) + assert sampler.available + assert sampler.sample_once(ts=1.0) >= 1 + sampler.flush() + [(_, batch)] = push.calls + sample = batch[0] + assert 0 <= sample.util <= 100 + assert sample.mem_mb >= 0 and sample.power_w >= 0 + assert sampler.gpu_uuids()[0].startswith("GPU-") diff --git a/tests/fast/dashboard/test_hooks.py b/tests/fast/dashboard/test_hooks.py new file mode 100644 index 0000000000..98bb9f411c --- /dev/null +++ b/tests/fast/dashboard/test_hooks.py @@ -0,0 +1,267 @@ +import logging + +import numpy as np +import pytest + +from miles.dashboard import backend, hooks +from miles.dashboard.hooks import BATCH_MAX_EVENTS, BATCH_MAX_SECONDS, _Identity +from miles.dashboard.store import Role +from miles.utils.timer import Timer + + +class FakeRemoteMethod: + def __init__(self, fail=False): + self.calls = [] + self.fail = fail + + def remote(self, *args, **kwargs): + if self.fail: + raise RuntimeError("collector unreachable") + self.calls.append((args, kwargs)) + + +class FakeHandle: + def __init__(self, fail_push=False): + self.push_phases = FakeRemoteMethod(fail=fail_push) + self.push_metrics = FakeRemoteMethod() + self.update_topology = FakeRemoteMethod() + self.set_router = FakeRemoteMethod() + + +@pytest.fixture(autouse=True) +def clean_state(monkeypatch): + timer = Timer() + saved = list(timer.event_sinks) + timer.event_sinks.clear() + monkeypatch.setattr(hooks, "_phase_sink", None) + monkeypatch.setattr(hooks, "_engines_fingerprint", None) + monkeypatch.setattr(hooks, "_resolve_identity", lambda: _Identity(node="10.0.0.3", gpus=[3], rank=7)) + monkeypatch.setattr(hooks, "_ray_get", lambda refs: refs) + monkeypatch.setattr(backend, "_handle", None) + monkeypatch.setattr(backend, "_is_primary", False) + monkeypatch.setattr(backend, "_resolution_failed", False) + yield + timer.event_sinks[:] = saved + + +# ------------------------------- phase sink --------------------------------- + + +def test_phase_sink_batches_by_count(): + handle = FakeHandle() + hooks.attach_phase_sink(handle, Role.TRAIN) + [sink] = Timer().event_sinks + + for i in range(BATCH_MAX_EVENTS - 1): + sink(f"phase_{i}", float(i), float(i) + 0.5) + assert handle.push_phases.calls == [] + + sink("actor_train", 100.0, 160.0) + [(args, _)] = handle.push_phases.calls + [batch] = args + assert len(batch) == BATCH_MAX_EVENTS + event = batch[-1] + assert (event.name, event.t0, event.t1) == ("actor_train", 100.0, 160.0) + assert (event.node, event.gpus, event.rank, event.role) == ("10.0.0.3", [3], 7, Role.TRAIN) + + +def test_phase_sink_batches_by_time(): + handle = FakeHandle() + hooks.attach_phase_sink(handle, Role.TRAIN) + [sink] = Timer().event_sinks + + sink("a", 1.0, 2.0) + assert handle.push_phases.calls == [] + sink._last_flush -= BATCH_MAX_SECONDS + 1 + sink("b", 2.0, 3.0) + [(args, _)] = handle.push_phases.calls + assert [e.name for e in args[0]] == ["a", "b"] + + +def test_phase_sink_reresolves_until_rank_known(monkeypatch): + handle = FakeHandle() + monkeypatch.setattr(hooks, "_resolve_identity", lambda: _Identity(node="n", gpus=[0], rank=-1)) + hooks.attach_phase_sink(handle, Role.TRAIN) + [sink] = Timer().event_sinks + + sink("early", 1.0, 2.0) # torch.distributed not initialized yet + monkeypatch.setattr(hooks, "_resolve_identity", lambda: _Identity(node="n", gpus=[0], rank=5)) + sink("late", 2.0, 3.0) + hooks.detach_and_flush() + + [(args, _)] = handle.push_phases.calls + assert [event.rank for event in args[0]] == [-1, 5] + + +def test_phase_sink_swallows_push_failures(caplog): + handle = FakeHandle(fail_push=True) + hooks.attach_phase_sink(handle, Role.TRAIN) + [sink] = Timer().event_sinks + with caplog.at_level(logging.WARNING): + for i in range(BATCH_MAX_EVENTS): + sink("p", float(i), float(i) + 1) # must not raise into Timer.end() + assert any("phase sink failed" in r.message for r in caplog.records) + + +def test_attach_is_idempotent_and_detach_flushes(): + handle = FakeHandle() + hooks.attach_phase_sink(handle, Role.TRAIN) + hooks.attach_phase_sink(handle, Role.ROLLOUT_MANAGER) # second attach ignored + assert len(Timer().event_sinks) == 1 + + Timer().event_sinks[0]("tail", 1.0, 2.0) + hooks.detach_and_flush() + assert Timer().event_sinks == [] + [(args, _)] = handle.push_phases.calls + assert [e.name for e in args[0]] == ["tail"] + + +def test_register_train_actor_disabled_is_free(monkeypatch): + monkeypatch.setattr(backend, "resolve_collector", lambda: pytest.fail("must not resolve when disabled")) + hooks.register_train_actor(type("Args", (), {"use_miles_dashboard": False})()) + assert Timer().event_sinks == [] + + +def test_register_train_actor_attaches_train_sink(monkeypatch): + handle = FakeHandle() + monkeypatch.setattr(backend, "resolve_collector", lambda: handle) + hooks.register_train_actor(type("Args", (), {"use_miles_dashboard": True})()) + [sink] = Timer().event_sinks + assert sink.role == Role.TRAIN + + +# ---------------------------- engine registration --------------------------- + + +class FakeEngineHandle: + def __init__(self, info): + self._info = info + self.get_topology_info = self + + def remote(self): + return self._info # hooks._ray_get is patched to the identity function + + +class FakeServerEngine: + def __init__(self, info, alive=True): + self.actor_handle = FakeEngineHandle(info) + self.is_allocated = alive + self.is_alive = alive + + +class FakeGroup: + def __init__(self, engines, nodes_per_engine=1): + self.all_engines = engines + self.nodes_per_engine = nodes_per_engine + + +def _info(url, node, gpus): + return dict(url=url, node_ip=node, gpu_ids=gpus, gpu_uuids=[None] * len(gpus), worker_type="regular", node_rank=0) + + +def _servers(*groups): + server = type("FakeServer", (), {"server_groups": list(groups)})() + return {"default": server} + + +def test_register_engines_groups_multinode_and_dedups(monkeypatch): + handle = FakeHandle() + monkeypatch.setattr(backend, "_handle", handle) + # one multi-node engine (master + worker node) and one single-node engine + master = FakeServerEngine(_info("http://a:1", "node-a", [0, 1])) + worker = FakeServerEngine(_info("http://a-worker:1", "node-b", [0, 1])) + single = FakeServerEngine(_info("http://b:1", "node-a", [2, 3])) + servers = _servers(FakeGroup([master, worker], nodes_per_engine=2), FakeGroup([single])) + + hooks.register_engines(servers) + [(args, _)] = handle.update_topology.calls + [snapshot] = args + assert [e.addr for e in snapshot.engines] == ["http://a:1", "http://b:1"] + multinode = snapshot.engines[0] + assert multinode.gpus == [["node-a", 0], ["node-a", 1], ["node-b", 0], ["node-b", 1]] + + hooks.register_engines(servers) # steady state: no remote traffic + assert len(handle.update_topology.calls) == 1 + + single.actor_handle = FakeEngineHandle(_info("http://b:2", "node-a", [2, 3])) # recovery: new actor + hooks.register_engines(servers) + assert len(handle.update_topology.calls) == 2 + assert handle.update_topology.calls[-1][0][0].engines[1].addr == "http://b:2" + + +def test_register_engines_skips_dead_chunks(monkeypatch): + handle = FakeHandle() + monkeypatch.setattr(backend, "_handle", handle) + alive = FakeServerEngine(_info("http://a:1", "n", [0])) + dead = FakeServerEngine(_info("http://b:1", "n", [1]), alive=False) + hooks.register_engines(_servers(FakeGroup([alive]), FakeGroup([dead]))) + + [(args, _)] = handle.update_topology.calls + assert [e.addr for e in args[0].engines] == ["http://a:1"] + + +def test_register_engines_without_collector_is_noop(): + hooks.register_engines(_servers(FakeGroup([FakeServerEngine(_info("http://a:1", "n", [0]))]))) + assert hooks._engines_fingerprint is None + + +# ------------------------------ dashboard_log ------------------------------- + + +def test_dashboard_log_filters_to_scalars(monkeypatch): + handle = FakeHandle() + monkeypatch.setattr(backend, "_handle", handle) + backend.dashboard_log( + {"a": 1.5, "b": "text", "c": [1, 2], "d": np.float32(2.5), "e": {"nested": 1}}, + step=3, + step_key="rollout/step", + ) + [(args, _)] = handle.push_metrics.calls + [record] = args + assert record.metrics == {"a": 1.5, "b": "text", "d": 2.5} + assert record.step == 3 and record.step_key == "rollout/step" + + +def test_dashboard_log_without_handle_is_noop(): + backend.dashboard_log({"a": 1}) # must not raise + + +# ----------------------------- router registration -------------------------- + + +def _router_args(ip="10.0.0.5", port=3333): + return type("Args", (), {"sglang_router_ip": ip, "sglang_router_port": port, "use_miles_router": False})() + + +def test_register_router_pushes_resolved_addr(monkeypatch): + handle = FakeHandle() + monkeypatch.setattr(backend, "_handle", handle) + hooks.register_router(_router_args()) + [(args, kwargs)] = handle.set_router.calls + assert args == ("http://10.0.0.5:3333",) + assert kwargs == {"use_miles_router": False} + + +def test_register_router_before_router_start_is_a_wiring_bug(monkeypatch): + monkeypatch.setattr(backend, "_handle", FakeHandle()) + with pytest.raises(AssertionError, match="after start_rollout_servers"): + hooks.register_router(_router_args(ip=None)) + + +def test_register_router_without_collector_is_noop(): + hooks.register_router(_router_args()) # must not raise + + +def test_phase_sink_begin_pushes_open_event_immediately(): + from miles.dashboard.store import PhaseEvent + + handle = FakeHandle() + hooks.attach_phase_sink(handle, Role.TRAIN) + [sink] = Timer().event_sinks + + sink.begin("rollout", 100.0) + [(args, _)] = handle.push_phases.calls # no batching for starts + [event] = args[0] + assert event.name == "rollout" and event.t0 == 100.0 + assert event.open and event.t1 == PhaseEvent.OPEN_T1 + assert (event.node, event.rank) == ("10.0.0.3", 7) diff --git a/tests/fast/dashboard/test_lane_selection.py b/tests/fast/dashboard/test_lane_selection.py new file mode 100644 index 0000000000..19958cd1d0 --- /dev/null +++ b/tests/fast/dashboard/test_lane_selection.py @@ -0,0 +1,174 @@ +import json + +import numpy as np +import pytest +from fastapi.testclient import TestClient +from tests.fast.dashboard.dummy_telemetry import DRAIN_A, DRAIN_B, ENGINE_B_NEW, GPU_NODE, dump_dummy_telemetry + +from miles.dashboard.dump_reader import DumpReader +from miles.dashboard.server import make_app +from miles.dashboard.store import GpuSample, Meta, MetricStore + + +@pytest.fixture +def loaded(tmp_path): + truth = dump_dummy_telemetry(tmp_path) + return MetricStore.load(tmp_path / "dashboard"), truth + + +# ------------------------------ lane grammar --------------------------------- + + +def test_resolve_lanes_grammar(loaded): + store, truth = loaded + assert store.resolve_lanes(None) is None + assert store.resolve_lanes("all") is None + # dummy fixture: train rank r runs on gpu r + assert store.resolve_lanes("rank:1") == {(GPU_NODE, 1)} + assert store.resolve_lanes("rank:0-2") == {(GPU_NODE, 0), (GPU_NODE, 1), (GPU_NODE, 2)} + assert store.resolve_lanes(f"gpu:{GPU_NODE}:3") == {(GPU_NODE, 3)} + # global lane numbers: single-node fixture makes them equal the gpu ids + assert store.resolve_lanes("g:2") == {(GPU_NODE, 2)} + assert store.resolve_lanes("g:1-3") == {(GPU_NODE, g) for g in (1, 2, 3)} + assert [e["index"] for e in store.lane_index()] == list(range(truth.gpus)) + assert store.resolve_lanes(f"node:{GPU_NODE}") == {(GPU_NODE, g) for g in range(truth.gpus)} + assert store.resolve_lanes("node:10.9.9.9") == set() # valid selector, no match + # engine B (either addr) covers gpus 2-3 + assert store.resolve_lanes("engine:15004") == {(GPU_NODE, 2), (GPU_NODE, 3)} + assert store.resolve_lanes("rank:0, engine:15000") == {(GPU_NODE, 0), (GPU_NODE, 1)} + # dummy telemetry is colocate-shaped: every lane carries both roles + assert store.resolve_lanes("role:train") == store.resolve_lanes("role:rollout") + + +def test_resolve_lanes_rejects_bad_grammar(loaded): + store, _ = loaded + for bad in ("bananas", "rank", "role:banana", "gpu:3", "g:x"): + with pytest.raises(ValueError): + store.resolve_lanes(bad) + + +def test_queries_respect_lane_filter(loaded): + store, truth = loaded + only = {(GPU_NODE, 2)} + phases = store.phases_by_lane(lanes=only) + assert phases and {(p["node"], p["gpu"]) for p in phases} == only + assert {p["name"] for p in phases} >= {"rollout", "actor_train", "initialize"} + + series = store.gpu_series(lanes=only) + assert set(series) == {f"{GPU_NODE}:2"} + + +# -------------------------------- heatmap ------------------------------------ + + +def _tiny_store(tmp_path): + writer = MetricStore(tmp_path / "dashboard") + writer.write_meta(Meta(run_name="tiny", start_ts=0.0, args={})) + for ts in range(10): + writer.append(GpuSample(ts=float(ts), node="n", gpu=0, util=10, mem_mb=1000, power_w=100)) + writer.append(GpuSample(ts=float(ts), node="n", gpu=1, util=50, mem_mb=4000, power_w=100)) + writer.append(GpuSample(ts=4.0, node="n", gpu=0, util=90, mem_mb=8000, power_w=100)) # spike + writer.flush() + return MetricStore.load(tmp_path / "dashboard") + + +def test_heatmap_util_bucket_max(tmp_path): + store = _tiny_store(tmp_path) + result = store.heatmap("util", t0=0.0, t1=10.0, x_buckets=5) + matrix = np.frombuffer(result["values"], dtype=np.uint8).reshape(2, 5) + assert [r["gpu"] for r in result["rows"]] == [0, 1] + assert result["scale"] == {"max": 100.0} + # lane 0: util 10 everywhere except the spike bucket (t=4 -> bucket 2) + expected_low = int(10 / 100 * 255) + assert matrix[0, 2] == int(90 / 100 * 255) + assert matrix[0, 0] == matrix[0, 4] == expected_low + assert (matrix[1] == int(50 / 100 * 255)).all() + + +def test_heatmap_mem_scales_to_peak(tmp_path): + store = _tiny_store(tmp_path) + result = store.heatmap("mem_mb", t0=0.0, t1=10.0, x_buckets=5) + matrix = np.frombuffer(result["values"], dtype=np.uint8).reshape(2, 5) + assert result["scale"] == {"max": 8000.0} + assert matrix[0, 2] == 255 # the 8000 MB spike + assert matrix[1, 0] == int(4000 / 8000 * 255) + + +def test_heatmap_phase_paints_dominant_and_initialize(loaded): + store, truth = loaded + result = store.heatmap("phase", x_buckets=400) + palette = result["palette"] + matrix = np.frombuffer(result["values"], dtype=np.uint8).reshape(len(result["rows"]), 400) + assert palette[0] == "" + assert {"initialize", "rollout", "actor_train", "train_wait"} <= set(palette) + + # rows carry roles so the carpet can draw the train/rollout divider + assert result["rows"][0]["roles"] == ["rollout", "train"] + + lane0 = matrix[0] + # the run starts with the initialize band, then rollout + assert palette[lane0[0]] == "initialize" + rollout_id = palette.index("rollout") + assert (lane0 == rollout_id).sum() > 0 + # active phases win over idle in shared buckets: actor_train present + assert (lane0 == palette.index("actor_train")).sum() > 0 + + +def test_heatmap_respects_lane_selection(loaded): + store, _ = loaded + result = store.heatmap("util", lanes={(GPU_NODE, 1)}, x_buckets=50) + assert [(r["node"], r["gpu"]) for r in result["rows"]] == [(GPU_NODE, 1)] + assert len(result["values"]) == 50 + + +# -------------------------------- outliers ----------------------------------- + + +def test_outliers_lowest_util(loaded): + store, truth = loaded + ranked = store.outliers("lowest_util", top_k=4) + # engine-B lanes (gpus 2-3) drain earlier each rollout -> lower mean util + assert DRAIN_B < DRAIN_A + assert {entry["gpu"] for entry in ranked[:2]} == {2, 3} + assert ranked[0]["score"] < ranked[-1]["score"] + + +def test_outliers_slowest_phase(loaded): + store, truth = loaded + ranked = store.outliers("slowest_phase:update_weights", top_k=4) + # fixture skew: update_weights duration grows with rank (4 + 0.3*rank) + assert [entry["gpu"] for entry in ranked] == [3, 2, 1, 0] + + with pytest.raises(ValueError): + store.outliers("bananas") + + +# --------------------------------- server ------------------------------------ + + +def test_server_endpoints(tmp_path): + dump_dummy_telemetry(tmp_path) + client = TestClient(make_app(MetricStore.load(tmp_path / "dashboard"), DumpReader(tmp_path))) + + response = client.get("/api/timeline/heatmap", params={"metric": "util", "x_buckets": 100}) + assert response.status_code == 200 + body = response.content + header_len = int.from_bytes(body[:4], "little") + header = json.loads(body[4 : 4 + header_len]) + matrix = np.frombuffer(body[4 + header_len :], dtype=np.uint8) + assert header["metric"] == "util" and header["x_buckets"] == 100 + assert len(matrix) == len(header["rows"]) * 100 + + phases = client.get("/api/timeline/phases", params={"lanes": "rank:0"}).json()["phases"] + assert {(p["node"], p["gpu"]) for p in phases} == {(GPU_NODE, 0)} + + gpu = client.get("/api/timeline/gpu", params={"lanes": f"engine:{ENGINE_B_NEW.rsplit(':', 1)[1]}"}).json() + assert set(gpu["lanes"]) == {f"{GPU_NODE}:2", f"{GPU_NODE}:3"} + + outliers = client.get("/api/timeline/outliers", params={"criterion": "lowest_util", "top_k": 2}).json() + assert len(outliers["outliers"]) == 2 + + assert client.get("/api/timeline/phases", params={"lanes": "bananas"}).status_code == 400 + assert client.get("/api/timeline/heatmap", params={"metric": "nope"}).status_code == 400 + assert client.get("/api/timeline/outliers", params={"criterion": "nope"}).status_code == 400 + assert client.get("/api/timeline/heatmap", params={"x_buckets": 1}).status_code == 400 diff --git a/tests/fast/dashboard/test_partitions.py b/tests/fast/dashboard/test_partitions.py new file mode 100644 index 0000000000..9578024fba --- /dev/null +++ b/tests/fast/dashboard/test_partitions.py @@ -0,0 +1,197 @@ +import json + +import pytest +from fastapi.testclient import TestClient +from tests.fast.dashboard.dummy_telemetry import dump_dummy_telemetry + +from miles.dashboard.dump_reader import DumpReader +from miles.dashboard.server import make_app +from miles.dashboard.store import GpuSample, Meta, MetricStore, PhaseEvent, Role, Stream, _hour_key + +# hour-aligned base so partition boundaries land where the test says they do +BASE = ((1_000_000 // 3600) + 1) * 3600.0 +HOUR = 3600.0 + + +def _three_hour_store(tmp_path): + """Two lanes of gpu samples across three hours, one rollout-role train + event per hour, and one long train event straddling the h0/h1 boundary.""" + writer = MetricStore(tmp_path) + writer.write_meta(Meta(run_name="hours", start_ts=BASE, args={})) + for hour in range(3): + for sec in (10.0, 1800.0, 3590.0): + ts = BASE + hour * HOUR + sec + for gpu in (0, 1): + writer.append(GpuSample(ts=ts, node="n", gpu=gpu, util=10 * (hour + 1), mem_mb=100, power_w=50)) + writer.append( + PhaseEvent( + name="actor_train", + t0=BASE + hour * HOUR + 100, + t1=BASE + hour * HOUR + 200, + node="n", + gpus=[0], + rank=0, + role=Role.TRAIN, + ) + ) + # straddler: starts in hour 0, completes in hour 1 -> lands in the h1 file + writer.append( + PhaseEvent( + name="straddle", + t0=BASE + 3500, + t1=BASE + HOUR + 300, + node="n", + gpus=[1], + rank=1, + role=Role.TRAIN, + ) + ) + writer.flush() + return MetricStore.load(tmp_path) + + +def test_flush_routes_hourly_files(tmp_path): + _three_hour_store(tmp_path) + gpu_files = sorted(p.name for p in (tmp_path / "gpu_util").glob("*.jsonl")) + assert gpu_files == sorted(f"{_hour_key(BASE + h * HOUR)}.jsonl" for h in range(3)) + assert not (tmp_path / Stream.GPU_UTIL.filename).exists() + # the straddler is keyed by its END hour + h1 = tmp_path / "phases" / f"{_hour_key(BASE + HOUR)}.jsonl" + assert any(json.loads(line)["name"] == "straddle" for line in h1.read_text().splitlines()) + + +def test_windowed_read_parses_only_touched_partitions(tmp_path): + store = _three_hour_store(tmp_path) + reader = store._readers[Stream.GPU_UTIL] + parsed = [] + original = reader._parse + reader._parse = lambda raw, path, offset: (parsed.append(path.stem), original(raw, path, offset))[-1] + + series = store.gpu_series(t0=BASE + HOUR + 5, t1=BASE + HOUR + 1900) + assert parsed == [_hour_key(BASE + HOUR)] + assert set(series) == {"n:0", "n:1"} + assert series["n:0"]["util"] == [20, 20] # samples at +10 and +1800 of hour 1 + + # second read inside the same hour: served from cache, no re-parse + store.gpu_series(t0=BASE + HOUR + 5, t1=BASE + HOUR + 2000) + assert parsed == [_hour_key(BASE + HOUR)] + + +def test_cached_partition_tails_appended_rows(tmp_path): + store = _three_hour_store(tmp_path) + window = dict(t0=BASE + 2 * HOUR, t1=BASE + 2 * HOUR + 3599) + assert len(store.gpu_series(**window)["n:0"]["ts"]) == 3 + + writer = MetricStore(tmp_path) + writer.append(GpuSample(ts=BASE + 2 * HOUR + 3595.0, node="n", gpu=0, util=99, mem_mb=100, power_w=50)) + writer.flush() + assert store.follow() == 1 # cached hour-2 block picks up the append + assert len(store.gpu_series(**window)["n:0"]["ts"]) == 4 + + +def test_phases_forward_slack_finds_straddler(tmp_path): + store = _three_hour_store(tmp_path) + # window entirely inside hour 0; the straddler's file is hour 1 + phases = store.phases_by_lane(t0=BASE + 3400, t1=BASE + 3550) + assert "straddle" in {p["name"] for p in phases} + + +def test_initialize_only_when_window_covers_run_start(tmp_path): + store = _three_hour_store(tmp_path) + full = {p["name"] for p in store.phases_by_lane()} + assert "initialize" in full + later = {p["name"] for p in store.phases_by_lane(t0=BASE + 50, t1=BASE + 300)} + assert "initialize" not in later + + +def test_legacy_flat_file_still_reads(tmp_path): + (tmp_path / Meta.FILENAME).write_text(json.dumps(dict(run_name="v1", start_ts=1.0, args={}))) + with open(tmp_path / Stream.GPU_UTIL.filename, "w") as f: + f.write('{"ts": 5.0, "node": "n", "gpu": 0, "util": 42, "mem_mb": 1, "power_w": 1}\n') + store = MetricStore.load(tmp_path) + assert store.has_stream(Stream.GPU_UTIL) + assert store.gpu_series()["n:0"]["util"] == [42] + # no catalog in a v1 dir: lane_index falls back to the stream scan + assert not (tmp_path / MetricStore.CATALOG_FILENAME).exists() + assert [(e["node"], e["gpu"]) for e in store.lane_index()] == [("n", 0)] + + +def test_lane_catalog_matches_stream_scan(tmp_path): + dump_dummy_telemetry(tmp_path) + dashboard = tmp_path / "dashboard" + assert (dashboard / MetricStore.CATALOG_FILENAME).exists() + from_catalog = MetricStore.load(dashboard).lane_index() + (dashboard / MetricStore.CATALOG_FILENAME).unlink() + from_scan = MetricStore.load(dashboard).lane_index() + assert from_catalog == from_scan + + +def test_time_range_from_edge_stamps(tmp_path): + store = _three_hour_store(tmp_path) + t0, t1 = store.time_range() + assert t0 == BASE + 10.0 # first gpu sample + assert t1 == BASE + 2 * HOUR + 3590.0 # last gpu sample + + +def test_open_marker_sentinel_stays_out_of_time_range(tmp_path): + # an OPEN marker (t1=-1.0) partitioned by its start hour can be the first + # line of the oldest phases file; its sentinel must not become the range min + writer = MetricStore(tmp_path) + writer.write_meta(Meta(run_name="open", start_ts=BASE, args={})) + writer.append( + PhaseEvent(name="train", t0=BASE + 5.0, t1=PhaseEvent.OPEN_T1, node="n", gpus=[0], rank=0, role=Role.TRAIN) + ) + writer.append( + PhaseEvent(name="train", t0=BASE + 5.0, t1=BASE + 900.0, node="n", gpus=[0], rank=0, role=Role.TRAIN) + ) + writer.flush() + store = MetricStore.load(tmp_path) + assert store.time_range() == (BASE + 5.0, BASE + 900.0) + + +def test_server_enforces_window_cap(tmp_path): + dump_dummy_telemetry(tmp_path) + client = TestClient(make_app(MetricStore.load(tmp_path / "dashboard"), DumpReader(tmp_path))) + + meta = client.get("/api/meta").json() + assert meta["capabilities"]["max_window_s"] == MetricStore.MAX_WINDOW_S + + over = MetricStore.MAX_WINDOW_S + 1 + for path, params in ( + ("/api/timeline/gpu", {}), + ("/api/timeline/phases", {}), + ("/api/timeline/engine_series", {"metric": "sglang_num_running_reqs"}), + ("/api/timeline/heatmap", {"metric": "util"}), + ): + ok = client.get(path, params=params) + assert ok.status_code == 200, path + bad = client.get(path, params={**params, "t0": 0.0, "t1": over}) + assert bad.status_code == 400, path + assert "max_window_s" in bad.json()["detail"] + + +def test_corrupt_partition_line_surfaces_at_read(tmp_path): + store = _three_hour_store(tmp_path) + partition = tmp_path / "gpu_util" / f"{_hour_key(BASE)}.jsonl" + with open(partition, "a") as f: + f.write("not json\n") + with pytest.raises(ValueError, match="corrupt record"): + store.gpu_series(t0=BASE, t1=BASE + 100) + + +def test_open_marker_partitions_by_start_hour_and_is_found(tmp_path): + writer = MetricStore(tmp_path) + writer.write_meta(Meta(run_name="open", start_ts=BASE, args={})) + writer.append(GpuSample(ts=BASE + 2 * HOUR, node="n", gpu=0, util=1, mem_mb=1, power_w=1)) + writer.append( + PhaseEvent(name="rollout", t0=BASE + 100, t1=PhaseEvent.OPEN_T1, node="n", gpus=[0], rank=0, role=Role.TRAIN) + ) + writer.flush() + # the marker sits in hour 0 (keyed by t0, not by the -1 sentinel) + h0 = tmp_path / "phases" / f"{_hour_key(BASE)}.jsonl" + assert h0.exists() and "rollout" in h0.read_text() + + store = MetricStore.load(tmp_path) + # windowed query two hours later still sees the growing band (backward slack) + phases = [p for p in store.phases_by_lane(t0=BASE + 2 * HOUR - 10, t1=BASE + 2 * HOUR) if p["name"] == "rollout"] + assert phases and phases[0]["t1"] == BASE + 2 * HOUR # clipped to the data edge diff --git a/tests/fast/dashboard/test_scraper.py b/tests/fast/dashboard/test_scraper.py new file mode 100644 index 0000000000..bc83caa2ce --- /dev/null +++ b/tests/fast/dashboard/test_scraper.py @@ -0,0 +1,129 @@ +import logging +from pathlib import Path + +import pytest + +from miles.dashboard.sglang_scraper import DEFAULT_METRIC_WHITELIST, ScrapeMode, SglangScraper +from miles.dashboard.store import EngineSample + +FIXTURES = Path(__file__).parent / "fixtures" +ROUTER_TEXT = (FIXTURES / "engine_metrics_router.txt").read_text() +DIRECT_TEXT = (FIXTURES / "engine_metrics_direct.txt").read_text() + +ENGINE_A = "http://10.0.0.1:15000" +ENGINE_B = "http://10.0.0.1:15004" + + +def router_scraper(sink, http_get, **kwargs): + return SglangScraper(sink, mode=ScrapeMode.ROUTER, router_addr="http://router:3000", http_get=http_get, **kwargs) + + +def test_router_mode_lifts_worker_addr_and_filters(caplog): + records: list[EngineSample] = [] + urls: list[str] = [] + + def fake_get(url, timeout): + urls.append(url) + return ROUTER_TEXT + + scraper = router_scraper(records.append, fake_get) + count = scraper.scrape_once(now=123.0) + + assert urls == ["http://router:3000/engine_metrics"] + assert count == len(records) == 5 # 2 running + 2 throughput + 1 token_usage + assert {r.metric for r in records} == {"sglang_num_running_reqs", "sglang_gen_throughput", "sglang_token_usage"} + assert all(r.ts == 123.0 for r in records) + + running = {r.addr: r.value for r in records if r.metric == "sglang_num_running_reqs"} + assert running == {ENGINE_A: 42.0, ENGINE_B: 7.0} + + # only the engine_type label survives; model_name and worker_addr are dropped + token_usage = next(r for r in records if r.metric == "sglang_token_usage") + assert token_usage.labels == {"engine_type": "prefill"} + assert all("worker_addr" not in r.labels and "model_name" not in r.labels for r in records) + + +def test_direct_mode_normalizes_colon_names_and_attaches_addr(): + records: list[EngineSample] = [] + scraper = SglangScraper( + records.append, + mode=ScrapeMode.DIRECT, + engine_addrs=lambda: [ENGINE_A], + http_get=lambda url, timeout: DIRECT_TEXT if url == f"{ENGINE_A}/metrics" else pytest.fail(url), + ) + count = scraper.scrape_once(now=5.0) + + assert count == 2 # uptime is filtered out + assert {(r.addr, r.metric, r.value) for r in records} == { + (ENGINE_A, "sglang_num_running_reqs", 13.0), + (ENGINE_A, "sglang_cache_hit_rate", 0.71), + } + + +def test_direct_mode_skips_failed_engine(caplog): + records: list[EngineSample] = [] + + def flaky_get(url, timeout): + if ENGINE_B in url: + raise ConnectionError("engine restarting") + return DIRECT_TEXT + + scraper = SglangScraper( + records.append, + mode=ScrapeMode.DIRECT, + engine_addrs=lambda: [ENGINE_A, ENGINE_B], + http_get=flaky_get, + ) + with caplog.at_level(logging.WARNING): + count = scraper.scrape_once(now=1.0) + + assert count == 2 # engine A's samples still land + assert {r.addr for r in records} == {ENGINE_A} + assert len([r for r in caplog.records if "skipping this tick" in r.message]) == 1 + + +def test_warnings_are_rate_limited(caplog): + scraper = SglangScraper( + lambda r: None, + mode=ScrapeMode.DIRECT, + engine_addrs=lambda: [ENGINE_A], + http_get=lambda url, timeout: (_ for _ in ()).throw(ConnectionError("down")), + ) + with caplog.at_level(logging.WARNING): + scraper.scrape_once(now=1.0) + scraper.scrape_once(now=2.0) # within the suppression window + assert len(caplog.records) == 1 + + scraper._warner.reset_window_for_test() # window elapses + with caplog.at_level(logging.WARNING): + scraper.scrape_once(now=3.0) + assert len(caplog.records) == 2 + + +def test_router_sample_without_worker_addr_is_dropped(caplog): + records: list[EngineSample] = [] + text = '# TYPE sglang_num_running_reqs gauge\nsglang_num_running_reqs{model_name="q"} 3.0\n' + scraper = router_scraper(records.append, lambda url, timeout: text) + with caplog.at_level(logging.WARNING): + assert scraper.scrape_once(now=1.0) == 0 + assert records == [] + assert "lacks worker_addr" in caplog.records[0].message + + +def test_whitelist_override(): + records: list[EngineSample] = [] + scraper = router_scraper(records.append, lambda url, timeout: ROUTER_TEXT, whitelist=("sglang_gen_throughput",)) + scraper.scrape_once(now=1.0) + assert {r.metric for r in records} == {"sglang_gen_throughput"} + + +def test_mode_prerequisites(): + with pytest.raises(AssertionError, match="router_addr"): + SglangScraper(lambda r: None, mode=ScrapeMode.ROUTER) + with pytest.raises(AssertionError, match="engine_addrs"): + SglangScraper(lambda r: None, mode=ScrapeMode.DIRECT) + + +def test_default_whitelist_covers_design_set(): + assert "sglang_num_running_reqs" in DEFAULT_METRIC_WHITELIST + assert "sglang_kv_transfer_speed_gb_s" in DEFAULT_METRIC_WHITELIST # PD set present diff --git a/tests/fast/dashboard/test_server.py b/tests/fast/dashboard/test_server.py new file mode 100644 index 0000000000..169afd23f9 --- /dev/null +++ b/tests/fast/dashboard/test_server.py @@ -0,0 +1,177 @@ +import time + +import pytest +import torch +from fastapi.testclient import TestClient +from tests.fast.dashboard.dummy_dump import dump_dummy_run + +from miles.dashboard.dump_reader import DumpReader +from miles.dashboard.server import make_app +from miles.dashboard.store import Meta, MetricsRecord, MetricStore + + +@pytest.fixture +def dump_dir(tmp_path): + dump_dummy_run(tmp_path, steps=2, with_eval=True) + writer = MetricStore(tmp_path / "dashboard") + writer.write_meta(Meta(run_name="dummy-run", start_ts=100.0, args={})) + writer.append(MetricsRecord(ts=101.0, step_key="rollout/step", step=0, metrics={"rollout/rewards_mean": 0.5})) + writer.append(MetricsRecord(ts=102.0, step_key="rollout/step", step=1, metrics={"rollout/rewards_mean": 0.6})) + writer.append(MetricsRecord(ts=103.0, step_key="train/step", step=7, metrics={"train/loss": 1.5})) + writer.flush() + return tmp_path + + +@pytest.fixture +def client(dump_dir): + store = MetricStore.load(dump_dir / "dashboard") + return TestClient(make_app(store, DumpReader(dump_dir))) + + +def test_meta(client): + meta = client.get("/api/meta").json() + assert meta["mode"] == "static" + assert meta["run_name"] == "dummy-run" + assert meta["rollout_ids"] == {"train": [0, 1], "eval": [0]} + assert "rollout/rewards_mean" in meta["metric_keys"] + # telemetry present: the catalog stays wandb-shaped, dump aggregates + # are not advertised (they remain servable on explicit request) + assert "dump/reward_mean" not in meta["metric_keys"] + assert meta["step_keys"] == ["rollout/step", "train/step"] + assert meta["capabilities"]["has_metrics"] is True + assert meta["capabilities"]["has_tokenizer"] is True + assert meta["capabilities"]["has_timeline"] is False + + +def test_metrics_from_store(client): + series = client.get("/api/metrics", params={"keys": "rollout/rewards_mean"}).json() + assert series["rollout/rewards_mean"] == {"x": [0, 1], "y": [0.5, 0.6], "ts": [101.0, 102.0]} + + series = client.get("/api/metrics", params={"keys": "train/loss", "x": "train/step"}).json() + assert series["train/loss"]["x"] == [7] + + series = client.get("/api/metrics", params={"keys": "no/such_key"}).json() + assert series["no/such_key"] == {"x": [], "y": [], "ts": []} + + +def test_metrics_dump_derived(client, dump_dir): + series = client.get("/api/metrics", params={"keys": "dump/reward_mean,rollout/rewards_mean"}).json() + aggregates = DumpReader(dump_dir).step_aggregates() + assert series["dump/reward_mean"]["x"] == [0, 1] + assert series["dump/reward_mean"]["y"] == pytest.approx(aggregates["reward_mean"].to_list()) + assert series["rollout/rewards_mean"]["y"] == [0.5, 0.6] + + assert client.get("/api/metrics", params={"keys": "dump/nope"}).status_code == 400 + assert client.get("/api/metrics", params={"keys": ""}).status_code == 400 + + +def test_summary_and_groups_endpoints(client): + body = client.get("/api/rollout/0/summary").json() + assert body["rollout_id"] == 0 and body["evaluation"] is False + assert len(body["rows"]) == 8 + assert "sample_index" in body["columns"] and "mean_abs_lp_diff" in body["columns"] + + eval_body = client.get("/api/rollout/0/summary", params={"eval": "true"}).json() + assert eval_body["evaluation"] is True + assert len(eval_body["rows"]) == 4 + + groups = client.get("/api/rollout/0/groups").json() + assert len(groups["rows"]) == 4 + assert "zero_std" in groups["columns"] + + assert client.get("/api/rollout/42/summary").status_code == 404 + + +def test_tokens_endpoint(client): + payload = client.get("/api/rollout/0/sample/0/tokens").json() + assert payload["token_text"] is not None + assert len(payload["token_ids"]) == payload["end"] - payload["start"] == payload["total_len"] + + window = client.get("/api/rollout/0/sample/0/tokens", params={"start": 1, "end": 4}).json() + assert len(window["token_ids"]) == 3 + + eval_payload = client.get("/api/rollout/0/sample/0/tokens", params={"eval": "true"}).json() + assert eval_payload["train_log_probs"] is None + assert eval_payload["rollout_log_probs"] is not None + + assert client.get("/api/rollout/0/sample/999999/tokens").status_code == 404 + assert client.get("/api/rollout/0/sample/0/tokens", params={"start": 5, "end": 5}).status_code == 400 + assert client.get("/api/rollout/42/sample/0/tokens").status_code == 404 + + +def test_still_writing_maps_to_503(client, dump_dir): + (dump_dir / "rollout_data" / "9.pt").write_bytes(b"garbage") # fresh mtime + assert client.get("/api/rollout/9/summary").status_code == 503 + + +def test_nan_values_serialize_as_null(dump_dir): + path = dump_dir / "train_data" / "0_0.pt" + pack = torch.load(path, weights_only=False) + pack["rollout_data"]["log_probs"][0][:] = float("nan") + torch.save(pack, path) + stamp = time.time() - 100 + import os + + os.utime(path, (stamp, stamp)) + + client = TestClient(make_app(MetricStore.load(dump_dir / "dashboard"), DumpReader(dump_dir))) + body = client.get("/api/rollout/0/summary").json() # json.loads would fail on bare NaN + poisoned = pack["rollout_data"]["sample_indices"][0] + row = next(r for r in body["rows"] if r["sample_index"] == poisoned) + assert row["mean_abs_lp_diff"] is None + assert row["mean_imp_ratio"] is None + + +def test_static_index(client): + response = client.get("/") + assert response.status_code == 200 + assert "miles dashboard" in response.text + + +def test_dump_keys_advertised_only_without_telemetry(tmp_path): + from tests.fast.dashboard.dummy_dump import dump_dummy_run + + dump_dummy_run(tmp_path) # dumps only: no dashboard/ telemetry at all + client = TestClient(make_app(MetricStore.load(tmp_path / "dashboard"), DumpReader(tmp_path))) + meta = client.get("/api/meta").json() + assert meta["capabilities"]["has_metrics"] is False + assert "dump/reward_mean" in meta["metric_keys"] + + +def test_engine_metric_catalog(tmp_path): + from tests.fast.dashboard.dummy_telemetry import dump_dummy_telemetry + + dump_dummy_telemetry(tmp_path) + from tests.fast.dashboard.dummy_dump import dump_dummy_run + + dump_dummy_run(tmp_path) + client = TestClient(make_app(MetricStore.load(tmp_path / "dashboard"), DumpReader(tmp_path))) + meta = client.get("/api/meta").json() + assert "sglang_num_running_reqs" in meta["engine_metric_keys"] + + +def test_sample_messages_from_trajectory_sidecar(tmp_path): + from tests.fast.dashboard.dummy_dump import dump_dummy_run + + dump_dummy_run(tmp_path) + client = TestClient(make_app(MetricStore.load(tmp_path / "dashboard"), DumpReader(tmp_path))) + + row = client.get("/api/rollout/0/sample/0/messages").json() + assert [m["role"] for m in row["messages"]] == ["system", "user", "assistant", "tool", "assistant"] + assert row["messages"][2]["tool_calls"][0]["function"]["name"] == "lookup" + # sample 2 recorded no conversation; eval dumps carry none at all + assert client.get("/api/rollout/0/sample/2/messages").status_code == 404 + assert client.get("/api/rollout/0/sample/0/messages?eval=1").status_code == 404 + + +def test_make_demo_dir(tmp_path): + from miles.dashboard.serve import make_demo_dir + + make_demo_dir(tmp_path) + reader = DumpReader(tmp_path) + assert reader.rollout_ids().train == [0, 1, 2] + store = MetricStore.load(tmp_path / "dashboard") + client = TestClient(make_app(store, reader)) + meta = client.get("/api/meta").json() + assert meta["capabilities"]["has_timeline"] is True + assert meta["rollout_ids"]["train"] == [0, 1, 2] diff --git a/tests/fast/dashboard/test_store.py b/tests/fast/dashboard/test_store.py new file mode 100644 index 0000000000..cd3ff3639d --- /dev/null +++ b/tests/fast/dashboard/test_store.py @@ -0,0 +1,217 @@ +import numpy as np +import pytest + +from miles.dashboard.store import ( + EngineInfo, + EngineSample, + GpuSample, + Meta, + MetricsRecord, + MetricStore, + PhaseEvent, + Stream, + TopologySnapshot, + TrajectoryEvent, + TrajectoryEventKind, + minmax_downsample, + stride_downsample, +) + + +def _one_of_each() -> list: + return [ + MetricsRecord(ts=10.0, step_key="rollout/step", step=3, metrics={"rollout/rewards_mean": 0.5}), + PhaseEvent(name="actor_train", t0=9.0, t1=11.0, node="10.0.0.2", gpus=[0, 1], rank=5, role="train"), + TopologySnapshot( + ts=8.0, + engines=[ + EngineInfo( + addr="http://10.0.0.2:15000", + worker_type="regular", + engine_rank=0, + gpus=[["10.0.0.2", 0]], + gpu_uuids=[None], + ) + ], + ), + GpuSample(ts=10.5, node="10.0.0.2", gpu=0, util=87, mem_mb=61234, power_w=612), + TrajectoryEvent( + ts=10.7, + kind=TrajectoryEventKind.GEN_START, + sample_index=3, + group_index=0, + turn=1, + weight_version="1", + detail="", + ), + EngineSample(ts=10.6, addr="http://10.0.0.2:15000", metric="sglang_num_running_reqs", labels={}, value=42.0), + ] + + +def test_roundtrip_all_streams(tmp_path): + writer = MetricStore(tmp_path) + writer.write_meta(Meta(run_name="test-run", start_ts=1.0, args={"colocate": True})) + records = _one_of_each() + for record in records: + writer.append(record) + writer.flush() + + reader = MetricStore.load(tmp_path) + assert reader.meta == Meta(run_name="test-run", start_ts=1.0, args={"colocate": True}, schema_version=1) + for record in records: + got = ( + reader.iter_records(record.stream) + if record.stream in MetricStore.PARTITIONED_STREAMS + else reader.records[record.stream] + ) + assert got == [record] + + +def test_flush_clears_buffers_and_appends(tmp_path): + writer = MetricStore(tmp_path) + for record in _one_of_each(): + writer.append(record) + writer.flush() + writer.flush() # empty buffers: no duplicate lines + for record in _one_of_each(): + writer.append(record) + writer.flush() + + reader = MetricStore.load(tmp_path) + for stream in Stream: + got = reader.iter_records(stream) if stream in MetricStore.PARTITIONED_STREAMS else reader.records[stream] + assert len(got) == 2, stream + + +def test_follow_incremental(tmp_path): + writer = MetricStore(tmp_path) + writer.append(MetricsRecord(ts=1.0, step_key="rollout/step", step=0, metrics={"a": 1})) + writer.flush() + + reader = MetricStore.load(tmp_path) + assert len(reader.records[Stream.METRICS]) == 1 + assert reader.follow() == 0 + + assert reader.gpu_series() == {} # prime the (empty) gpu partition cache + + writer.append(MetricsRecord(ts=2.0, step_key="rollout/step", step=1, metrics={"a": 2})) + writer.append(GpuSample(ts=2.5, node="n", gpu=0, util=10, mem_mb=1, power_w=1)) + writer.flush() + # metrics tail eagerly; the gpu partition is new (never cached) so it + # stays lazy and surfaces through the query below + assert reader.follow() == 1 + assert [r.step for r in reader.records[Stream.METRICS]] == [0, 1] + assert list(reader.gpu_series()) == ["n:0"] + assert reader.follow() == 0 + + +def test_partial_trailing_line_left_for_later(tmp_path): + writer = MetricStore(tmp_path) + writer.append(MetricsRecord(ts=1.0, step_key="rollout/step", step=0, metrics={"a": 1})) + writer.flush() + + path = tmp_path / Stream.METRICS.filename + full_line = '{"ts": 2.0, "step_key": "rollout/step", "step": 1, "metrics": {"a": 2}}\n' + with open(path, "a") as f: + f.write(full_line[:20]) # crash mid-write: no trailing newline + + reader = MetricStore.load(tmp_path) + assert len(reader.records[Stream.METRICS]) == 1 + + with open(path, "a") as f: + f.write(full_line[20:]) # writer completes the record + assert reader.follow() == 1 + assert reader.records[Stream.METRICS][1].step == 1 + + +def test_corrupt_complete_line_raises(tmp_path): + writer = MetricStore(tmp_path) + writer.append(MetricsRecord(ts=1.0, step_key="rollout/step", step=0, metrics={"a": 1})) + writer.flush() + with open(tmp_path / Stream.METRICS.filename, "a") as f: + f.write("not json\n") + + with pytest.raises(ValueError, match="corrupt record"): + MetricStore.load(tmp_path) + + +def test_wrong_field_complete_line_raises(tmp_path): + with open(tmp_path / Stream.GPU_UTIL.filename, "w") as f: + f.write('{"ts": 1.0, "unexpected_field": 1}\n') + store = MetricStore.load(tmp_path) # partitions parse lazily: load is clean + with pytest.raises(ValueError, match="corrupt record"): + store.gpu_series() + + +def test_empty_dir(tmp_path): + reader = MetricStore.load(tmp_path / "does_not_exist") + assert reader.meta is None + assert reader.follow() == 0 + assert reader.time_range() is None + assert reader.metric_keys() == [] + assert reader.metric_series(["a"], x_key="rollout/step") == {"a": {"x": [], "y": [], "ts": []}} + + +def test_metric_series_filters_by_step_key_and_time(tmp_path): + writer = MetricStore(tmp_path) + writer.append(MetricsRecord(ts=1.0, step_key="rollout/step", step=0, metrics={"a": 1.0, "b": 10.0})) + writer.append(MetricsRecord(ts=2.0, step_key="train/step", step=100, metrics={"a": -1.0})) + writer.append(MetricsRecord(ts=3.0, step_key="rollout/step", step=1, metrics={"a": 2.0})) + writer.flush() + reader = MetricStore.load(tmp_path) + + series = reader.metric_series(["a", "b"], x_key="rollout/step") + assert series["a"] == {"x": [0, 1], "y": [1.0, 2.0], "ts": [1.0, 3.0]} + assert series["b"] == {"x": [0], "y": [10.0], "ts": [1.0]} + + series = reader.metric_series(["a"], x_key="train/step") + assert series["a"] == {"x": [100], "y": [-1.0], "ts": [2.0]} + + series = reader.metric_series(["a"], x_key="rollout/step", t0=2.5, t1=3.5) + assert series["a"] == {"x": [1], "y": [2.0], "ts": [3.0]} + + +def test_metric_keys_and_step_keys(tmp_path): + writer = MetricStore(tmp_path) + writer.append(MetricsRecord(ts=1.0, step_key="rollout/step", step=0, metrics={"b": 1, "a": 1})) + writer.append(MetricsRecord(ts=2.0, step_key="eval/step", step=0, metrics={"c": 1})) + writer.flush() + reader = MetricStore.load(tmp_path) + assert reader.metric_keys() == ["a", "b", "c"] + assert reader.step_keys() == ["eval/step", "rollout/step"] + + +def test_time_range_spans_all_streams(tmp_path): + writer = MetricStore(tmp_path) + writer.append(PhaseEvent(name="rollout", t0=0.5, t1=7.0, node="n", gpus=[], rank=-1, role="rollout_manager")) + writer.append(GpuSample(ts=9.0, node="n", gpu=0, util=1, mem_mb=1, power_w=1)) + writer.flush() + reader = MetricStore.load(tmp_path) + assert reader.time_range() == (0.5, 9.0) + + +def test_stride_downsample(): + xs, ys = np.arange(1000), np.arange(1000) * 2.0 + dx, dy = stride_downsample(xs, ys, 100) + assert len(dx) == len(dy) == 100 + assert dx[0] == 0 and dx[-1] == 999 # endpoints preserved + + dx, dy = stride_downsample(xs, ys, 2000) # already small enough + assert len(dx) == 1000 + + +def test_minmax_downsample_preserves_extremes(): + rng = np.random.default_rng(0) + ys = rng.normal(size=5000) + ys[1234] = 100.0 # spike + ys[4321] = -100.0 # dip + xs = np.arange(len(ys)) + + dx, dy = minmax_downsample(xs, ys, 200) + assert len(dx) <= 200 + assert dy.max() == 100.0 + assert dy.min() == -100.0 + assert (np.diff(dx) > 0).all() # x stays sorted + + dx, dy = minmax_downsample(xs[:50], ys[:50], 200) + assert len(dx) == 50 # small input unchanged diff --git a/tests/fast/dashboard/test_timeline.py b/tests/fast/dashboard/test_timeline.py new file mode 100644 index 0000000000..99c3628c40 --- /dev/null +++ b/tests/fast/dashboard/test_timeline.py @@ -0,0 +1,184 @@ +import pytest +from fastapi.testclient import TestClient +from tests.fast.dashboard.dummy_telemetry import ( + BASE_TS, + DRAIN_B, + ENGINE_A, + ENGINE_B_NEW, + ENGINE_B_OLD, + GPU_NODE, + ROLLOUT_SECONDS, + dump_dummy_telemetry, +) + +from miles.dashboard.dump_reader import DumpReader +from miles.dashboard.server import make_app +from miles.dashboard.store import MetricStore, Role + + +@pytest.fixture +def loaded(tmp_path): + truth = dump_dummy_telemetry(tmp_path) + return MetricStore.load(tmp_path / "dashboard"), truth + + +def test_lanes(loaded): + store, truth = loaded + assert store.lanes() == [dict(node=GPU_NODE, gpu=g, index=g) for g in range(truth.gpus)] + + +def test_topology_windows(loaded): + store, truth = loaded + windows = store.topology_windows() + assert len(windows) == 2 + assert windows[0]["t0"] == BASE_TS and windows[0]["t1"] == truth.restart_ts + assert windows[1]["t0"] == truth.restart_ts and windows[1]["t1"] is None + assert [e["addr"] for e in windows[0]["engines"]] == [ENGINE_A, ENGINE_B_OLD] + assert [e["addr"] for e in windows[1]["engines"]] == [ENGINE_A, ENGINE_B_NEW] + + +def test_rollout_expansion_covers_engine_lanes(loaded): + store, truth = loaded + phases = store.phases_by_lane() + rollout0 = [p for p in phases if p["role"] == Role.ROLLOUT_MANAGER and p["t0"] == truth.step_start(0)] + # step 0 lies entirely inside the first topology window: one interval per GPU + assert len(rollout0) == truth.gpus + assert {p["gpu"] for p in rollout0} == set(range(truth.gpus)) + assert all(p["node"] == GPU_NODE and p["t1"] == truth.step_start(0) + ROLLOUT_SECONDS for p in rollout0) + + +def test_rollout_expansion_clips_at_topology_restart(loaded): + store, truth = loaded + t0, t1 = truth.rollout_interval(1) # engine B restarts mid-way through this rollout + on_lane_2 = [ + p for p in store.phases_by_lane() if p["role"] == Role.ROLLOUT_MANAGER and p["gpu"] == 2 and t0 <= p["t0"] < t1 + ] + assert [(p["t0"], p["t1"]) for p in on_lane_2] == [(t0, truth.restart_ts), (truth.restart_ts, t1)] + + +def test_train_phases_keep_rank_and_skew(loaded): + store, truth = loaded + updates = [p for p in store.phases_by_lane() if p["name"] == "update_weights" and p["t0"] < truth.step_start(1)] + assert len(updates) == truth.gpus + by_rank = {p["rank"]: p for p in updates} + assert by_rank[truth.gpus - 1]["t1"] > by_rank[0]["t1"] # cross-rank skew survives + assert all(p["gpu"] == p["rank"] for p in updates) + + +def test_phases_time_filter(loaded): + store, truth = loaded + t0, t1 = truth.rollout_interval(0) + phases = store.phases_by_lane(t0=t0, t1=t1 - 1) + assert phases and all(p["t0"] < t1 - 1 and p["t1"] > t0 for p in phases) + assert {p["name"] for p in phases} == {"rollout"} # train events of step 0 start at rollout end + + +def test_gpu_series_downsampling_preserves_spike(loaded): + store, truth = loaded + lanes = store.gpu_series(max_points=40) + lane0 = lanes[f"{GPU_NODE}:0"] + assert len(lane0["ts"]) <= 40 + assert max(lane0["util"]) == 100 # the fixture's known spike survives + assert min(lane0["util"]) == 5 # so does the drain dip + assert len(lane0["ts"]) == len(lane0["util"]) == len(lane0["mem_mb"]) == len(lane0["power_w"]) + assert lane0["ts"] == sorted(lane0["ts"]) + + full = store.gpu_series() + # 1 Hz over initialize + steps, no downsampling needed + assert len(full[f"{GPU_NODE}:0"]["ts"]) == 30 + truth.steps * 100 + + +def test_gpu_series_time_filter(loaded): + store, truth = loaded + t0, t1 = truth.rollout_interval(0) + lanes = store.gpu_series(t0=t0, t1=t1) + for series in lanes.values(): + assert all(t0 <= ts <= t1 for ts in series["ts"]) + + +def test_engine_series_split_on_restart(loaded): + store, truth = loaded + series = store.engine_series("sglang_num_running_reqs") + by_addr = {s["addr"]: s for s in series} + assert set(by_addr) == {ENGINE_A, ENGINE_B_OLD, ENGINE_B_NEW} + assert max(by_addr[ENGINE_B_OLD]["ts"]) < truth.restart_ts <= min(by_addr[ENGINE_B_NEW]["ts"]) + # engine A spans all steps (engines come up after initialize) + assert min(by_addr[ENGINE_A]["ts"]) == truth.step_start(0) + + # long tail: B is drained (0 running) before the rollout window ends + b_new = by_addr[ENGINE_B_NEW] + t0, t1 = truth.rollout_interval(2) + in_window = [v for ts, v in zip(b_new["ts"], b_new["value"], strict=True) if t0 + DRAIN_B <= ts < t1] + assert in_window and all(v == 0 for v in in_window) + + assert store.engine_series("sglang_no_such_metric") == [] + + +def test_initialize_phase_synthesized_per_lane(loaded): + store, truth = loaded + init = [p for p in store.phases_by_lane() if p["name"] == "initialize"] + assert len(init) == truth.gpus # one band per lane + for band in init: + assert band["role"] == Role.DERIVED and band["rank"] == -1 + assert band["t0"] == BASE_TS # collector start (meta.start_ts) + assert band["t1"] == truth.step_start(0) # first observed event + # time-filtering away the init window drops the synthesized bands too + t0, _ = truth.rollout_interval(0) + assert not [p for p in store.phases_by_lane(t0=t0) if p["name"] == "initialize"] + + +def test_bubbles(loaded): + store, truth = loaded + bubbles = store.bubbles() + assert [b["step"] for b in bubbles] == list(range(truth.steps)) + assert bubbles[0]["wait_ratio"] == pytest.approx(0.4) + assert bubbles[0]["step_time"] == pytest.approx(95.0) + + +def test_timeline_endpoints(tmp_path): + truth = dump_dummy_telemetry(tmp_path) + client = TestClient(make_app(MetricStore.load(tmp_path / "dashboard"), DumpReader(tmp_path))) + + meta = client.get("/api/meta").json() + assert meta["capabilities"]["has_timeline"] is True + assert meta["capabilities"]["has_engine_series"] is True + + topology = client.get("/api/timeline/topology").json() + assert len(topology["lanes"]) == truth.gpus + assert len(topology["windows"]) == 2 + + phases = client.get("/api/timeline/phases", params={"t0": truth.step_start(0), "t1": truth.step_start(1)}).json() + assert {p["name"] for p in phases["phases"]} >= {"rollout", "actor_train", "train_wait"} + + gpu = client.get("/api/timeline/gpu", params={"max_points": 50}).json() + assert set(gpu["lanes"]) == {f"{GPU_NODE}:{g}" for g in range(truth.gpus)} + assert client.get("/api/timeline/gpu", params={"max_points": 1}).status_code == 400 + + engines = client.get("/api/timeline/engine_series", params={"metric": "sglang_gen_throughput"}).json() + assert {s["addr"] for s in engines["series"]} == {ENGINE_A, ENGINE_B_OLD, ENGINE_B_NEW} + + bubbles = client.get("/api/timeline/bubbles").json() + assert len(bubbles["bubbles"]) == truth.steps + + +def test_open_interval_clips_to_data_edge_and_closed_twin_wins(tmp_path): + from miles.dashboard.store import GpuSample, Meta, PhaseEvent, Role + + writer = MetricStore(tmp_path) + writer.write_meta(Meta(run_name="open", start_ts=0.0, args={})) + # newest data: a gpu sample at t=500 — the open band must grow to it + writer.append(GpuSample(ts=500.0, node="n", gpu=0, util=10, mem_mb=1, power_w=1)) + writer.append( + PhaseEvent(name="rollout", t0=100.0, t1=PhaseEvent.OPEN_T1, node="n", gpus=[0], rank=0, role=Role.TRAIN) + ) + # a second phase that already closed: its open marker must be superseded + writer.append( + PhaseEvent(name="warmup", t0=10.0, t1=PhaseEvent.OPEN_T1, node="n", gpus=[0], rank=0, role=Role.TRAIN) + ) + writer.append(PhaseEvent(name="warmup", t0=10.0, t1=40.0, node="n", gpus=[0], rank=0, role=Role.TRAIN)) + writer.flush() + + phases = {p["name"]: p for p in MetricStore.load(tmp_path).phases_by_lane() if p["name"] != "initialize"} + assert phases["rollout"]["t1"] == 500.0 # clipped to the data edge, not -1 + assert phases["warmup"]["t1"] == 40.0 # exactly one warmup, the closed one + assert sum(1 for p in MetricStore.load(tmp_path).phases_by_lane() if p["name"] == "warmup") == 1 diff --git a/tests/fast/dashboard/test_trajectory_sink.py b/tests/fast/dashboard/test_trajectory_sink.py new file mode 100644 index 0000000000..59e27670fd --- /dev/null +++ b/tests/fast/dashboard/test_trajectory_sink.py @@ -0,0 +1,163 @@ +import pytest + +from miles.dashboard import hooks +from miles.dashboard.hooks import BATCH_MAX_EVENTS, TrajectorySink +from miles.dashboard.store import Stream, TrajectoryEvent, TrajectoryEventKind +from miles.utils.lifecycle import TrajectoryLifecycle +from miles.utils.types import Sample + + +class FakeRemoteMethod: + def __init__(self, fail=False): + self.calls = [] + self.fail = fail + + def remote(self, *args, **kwargs): + if self.fail: + raise RuntimeError("collector unreachable") + self.calls.append((args, kwargs)) + + +class FakeHandle: + def __init__(self, fail_push=False): + self.push_trajectories = FakeRemoteMethod(fail=fail_push) + + +@pytest.fixture(autouse=True) +def clean_sink(): + TrajectoryLifecycle().sink = None + yield + TrajectoryLifecycle().sink = None + + +def _sample(index=7, group=2, versions=("3", "4")): + return Sample(index=index, group_index=group, weight_versions=list(versions)) + + +def _pushed(handle): + return [event for (args, _) in handle.push_trajectories.calls for event in args[0]] + + +# --------------------------------- emitters ---------------------------------- + + +def test_attempt_and_gen_events_carry_identity_and_version(): + handle = FakeHandle() + sink = TrajectorySink(handle) + sample = _sample() + sink.attempt_start(sample) + sink.gen_start(sample) + sink.attempt_end(sample) + sink.flush() + + kinds = [e.kind for e in _pushed(handle)] + assert kinds == [ + TrajectoryEventKind.ATTEMPT_START, + TrajectoryEventKind.GEN_START, + TrajectoryEventKind.ATTEMPT_END, + ] + for event in _pushed(handle): + assert (event.sample_index, event.group_index) == (7, 2) + assert event.weight_version == "4" + assert _pushed(handle)[-1].detail == Sample.Status.PENDING.value + + +def test_spans_use_explicit_timestamps_and_turns(): + handle = FakeHandle() + sink = TrajectorySink(handle) + sample = _sample() + sink.gen_span(sample, 10.0, 14.5, turn=3, detail="120") + sink.tool_span(sample, 14.5, 16.0, turn=3, detail="2 calls") + sink.flush() + + events = _pushed(handle) + assert [(e.kind, e.ts, e.turn) for e in events] == [ + (TrajectoryEventKind.GEN_START, 10.0, 3), + (TrajectoryEventKind.GEN_END, 14.5, 3), + (TrajectoryEventKind.TOOL_START, 14.5, 3), + (TrajectoryEventKind.TOOL_END, 16.0, 3), + ] + assert events[1].detail == "120" + + +def test_attempt_end_drains_lifecycle_metadata_from_turn_samples(): + handle = FakeHandle() + sink = TrajectorySink(handle) + turns = [] + for i in range(2): + sample = _sample() + sample.metadata["lifecycle"] = dict(t0=100.0 + i, t1=110.0 + i, turn=i + 1) + turns.append(sample) + turns[0].metadata["lifecycle"]["t0"] = None # unknown start: only the end emits + turns[1].metadata["lifecycle"]["prev_t1"] = 110.0 # session gap = agent/tool work + turns[1].metadata["lifecycle"]["req_ts"] = 110.6 # server-edge arrival bounds it + + sink.attempt_end(turns) + sink.flush() + events = _pushed(handle) + assert [(e.kind, e.turn) for e in events] == [ + (TrajectoryEventKind.GEN_END, 1), + (TrajectoryEventKind.TOOL_START, 2), + (TrajectoryEventKind.TOOL_END, 2), + (TrajectoryEventKind.GEN_START, 2), + (TrajectoryEventKind.GEN_END, 2), + (TrajectoryEventKind.ATTEMPT_END, -1), + ] + assert (events[1].ts, events[2].ts) == (110.0, 110.6) # gap ends at arrival, not gen start + assert events[1].detail == "agent gap" + + +def test_batching_and_failure_swallowed(): + handle = FakeHandle() + sink = TrajectorySink(handle) + sample = _sample() + for _ in range(BATCH_MAX_EVENTS - 1): + sink.gen_start(sample) + assert handle.push_trajectories.calls == [] + sink.gen_start(sample) # crosses the batch threshold + assert len(_pushed(handle)) == BATCH_MAX_EVENTS + + failing = TrajectorySink(FakeHandle(fail_push=True)) + for _ in range(BATCH_MAX_EVENTS + 1): + failing.gen_start(sample) # must never raise into the rollout path + failing.flush() + + +def test_attach_detach_via_lifecycle_seam(): + assert TrajectoryLifecycle().sink is None + handle = FakeHandle() + hooks.attach_trajectory_sink(handle) + sink = TrajectoryLifecycle().sink + assert isinstance(sink, TrajectorySink) + hooks.attach_trajectory_sink(FakeHandle()) + assert TrajectoryLifecycle().sink is sink # one sink per process + + sink.gen_start(_sample()) + hooks.detach_and_flush() + assert TrajectoryLifecycle().sink is None + assert len(_pushed(handle)) == 1 # detach flushed the partial batch + + +def test_trajectory_events_roundtrip_store(tmp_path): + from miles.dashboard.store import MetricStore + + writer = MetricStore(tmp_path) + for ts in (10.0, 3700.0): + writer.append( + TrajectoryEvent( + ts=ts, + kind=TrajectoryEventKind.GEN_START, + sample_index=5, + group_index=-1, + turn=1, + weight_version="2", + detail="", + ) + ) + writer.flush() + store = MetricStore.load(tmp_path) + assert store.has_stream(Stream.TRAJECTORIES) + assert len((tmp_path / "trajectories").glob("*.jsonl") and list((tmp_path / "trajectories").glob("*.jsonl"))) == 2 + assert [e.ts for e in store.trajectory_events()] == [10.0, 3700.0] + assert [e.ts for e in store.trajectory_events(t0=3600.0)] == [3700.0] + assert store.iter_records(Stream.TRAJECTORIES)[0].weight_version == "2" diff --git a/tests/fast/dashboard/test_trajectory_views.py b/tests/fast/dashboard/test_trajectory_views.py new file mode 100644 index 0000000000..ec067b92aa --- /dev/null +++ b/tests/fast/dashboard/test_trajectory_views.py @@ -0,0 +1,151 @@ +import numpy as np +import pytest +from fastapi.testclient import TestClient +from tests.fast.dashboard.dummy_dump import dump_dummy_run +from tests.fast.dashboard.dummy_telemetry import SAMPLES_PER_STEP, dump_dummy_telemetry + +from miles.dashboard.dump_reader import DumpReader +from miles.dashboard.server import make_app +from miles.dashboard.store import MetricStore, TrajectoryEvent, TrajectoryEventKind + + +@pytest.fixture +def combined(tmp_path): + """dump + telemetry in one dir — the same composition serve --demo uses.""" + dump_dummy_run(tmp_path) + truth = dump_dummy_telemetry(tmp_path) + return MetricStore.load(tmp_path / "dashboard"), DumpReader(tmp_path), truth + + +# ------------------------------ lane assembly -------------------------------- + + +def _event(ts, kind, index=1, turn=-1, version="2", detail=""): + return TrajectoryEvent( + ts=ts, kind=kind, sample_index=index, group_index=0, turn=turn, weight_version=version, detail=detail + ) + + +def test_trajectory_lanes_pairs_spans_and_attempts(tmp_path): + writer = MetricStore(tmp_path) + kinds = TrajectoryEventKind + for event in ( + _event(10.0, kinds.ATTEMPT_START), + _event(11.0, kinds.GEN_START, turn=1), + _event(15.0, kinds.GEN_END, turn=1), + _event(15.0, kinds.TOOL_START, turn=1), + _event(18.0, kinds.TOOL_END, turn=1), + _event(18.0, kinds.GEN_START, turn=2, version="3"), + # turn-2 gen never ends: aborted mid-generation, then a second attempt + _event(19.0, kinds.ATTEMPT_END, detail="aborted"), + _event(30.0, kinds.ATTEMPT_START), + _event(31.0, kinds.GEN_START, turn=1, version="4"), + _event(35.0, kinds.GEN_END, turn=1, version="4"), + _event(36.0, kinds.ATTEMPT_END, detail="completed"), + ): + writer.append(event) + writer.flush() + + [lane] = MetricStore.load(tmp_path).trajectory_lanes() + assert lane["sample_index"] == 1 + assert lane["status"] == "completed" + assert lane["versions"] == [2, 3, 4] + assert lane["attempts"] == [dict(t0=10.0, t1=19.0), dict(t0=30.0, t1=36.0)] + by_kind = [(s["kind"], s["t0"], s["t1"], s["turn"]) for s in lane["segments"]] + assert by_kind == [ + ("gen", 11.0, 15.0, 1), + ("tool", 15.0, 18.0, 1), + ("gen", 18.0, 19.0, 2), # aborted mid-generation: closed by attempt_end + ("gen", 31.0, 35.0, 1), + ] + + +def test_lifecycle_heatmap_row_cap(combined, monkeypatch): + store, _, truth = combined + monkeypatch.setattr(MetricStore, "LIFECYCLE_MAX_ROWS", 5) + result = store.heatmap("lifecycle", x_buckets=50) + assert len(result["rows"]) == 5 + assert result["rows_total"] == truth.steps * SAMPLES_PER_STEP + # earliest submits kept (rows are submit-ordered) + assert [r["sample_index"] for r in result["rows"]] == [0, 1, 2, 3, 4] + + +def test_lifecycle_heatmap_palette_and_rows(combined): + store, _, truth = combined + result = store.heatmap("lifecycle", x_buckets=200) + assert result["palette"] == ["", "queue", "generating", "tool_wait"] + assert len(result["rows"]) == truth.steps * SAMPLES_PER_STEP + matrix = np.frombuffer(result["values"], dtype=np.uint8).reshape(len(result["rows"]), 200) + assert (matrix == result["palette"].index("generating")).sum() > 0 + assert (matrix == result["palette"].index("tool_wait")).sum() > 0 + # rows are submit-ordered + firsts = [r["sample_index"] for r in result["rows"][: SAMPLES_PER_STEP + 1]] + assert firsts == sorted(firsts) + + +def test_single_turn_gen_closes_at_attempt_end(tmp_path): + # the generate_and_rm probe alone (single-turn path: no gen_end) must + # yield a CLOSED gen span; dangling spans painted every sample's end at + # the consume line (report 2026-07-13) + writer = MetricStore(tmp_path) + kinds = TrajectoryEventKind + for event in ( + _event(10.0, kinds.ATTEMPT_START), + _event(12.0, kinds.GEN_START), # turn=-1: attempt-level probe + _event(50.0, kinds.ATTEMPT_END, detail="completed"), + ): + writer.append(event) + writer.flush() + [lane] = MetricStore.load(tmp_path).trajectory_lanes() + assert [(s["kind"], s["t0"], s["t1"]) for s in lane["segments"]] == [("gen", 12.0, 50.0)] + + +def test_coarse_gen_span_superseded_by_turn_spans(tmp_path): + writer = MetricStore(tmp_path) + kinds = TrajectoryEventKind + for event in ( + _event(10.0, kinds.ATTEMPT_START), + _event(11.0, kinds.GEN_START), # coarse attempt-level span + _event(12.0, kinds.GEN_START, turn=1), + _event(20.0, kinds.GEN_END, turn=1), + _event(30.0, kinds.ATTEMPT_END, detail="completed"), + ): + writer.append(event) + writer.flush() + [lane] = MetricStore.load(tmp_path).trajectory_lanes() + assert [(s["kind"], s["turn"]) for s in lane["segments"]] == [("gen", 1)] + + +# --------------------------------- endpoint ---------------------------------- + + +def test_rollout_trajectories_endpoint(combined): + store, reader, truth = combined + client = TestClient(make_app(store, reader)) + + payload = client.get("/api/rollout/1/trajectories").json() + lanes = payload["lanes"] + assert {lane["sample_index"] for lane in lanes} == set(range(SAMPLES_PER_STEP, 2 * SAMPLES_PER_STEP)) + agentic = next(lane for lane in lanes if lane["sample_index"] % 3 == 0) + plain = next(lane for lane in lanes if lane["sample_index"] % 3 == 1) + assert len(agentic["versions"]) == 2 and [s["kind"] for s in agentic["segments"]] == ["gen", "tool", "gen"] + assert len(plain["versions"]) == 1 and [s["kind"] for s in plain["segments"]] == ["gen"] + assert all(lane["status"] == "completed" for lane in lanes) + # events land within the step's rollout window, before the consume anchor + if payload["consume_ts"] is not None: + assert all(lane["last_ts"] <= payload["consume_ts"] + 1e-6 for lane in lanes) + + +def test_rollout_trajectories_single_sample_filter(combined): + store, reader, _ = combined + client = TestClient(make_app(store, reader)) + payload = client.get("/api/rollout/1/trajectories", params={"sample_index": SAMPLES_PER_STEP + 1}).json() + [lane] = payload["lanes"] + assert lane["sample_index"] == SAMPLES_PER_STEP + 1 + + +def test_rollout_trajectories_empty_for_probe_less_runs(tmp_path): + dump_dummy_run(tmp_path) # dump only: no telemetry, no events + client = TestClient(make_app(MetricStore.load(tmp_path / "dashboard"), DumpReader(tmp_path))) + payload = client.get("/api/rollout/0/trajectories").json() + assert payload["lanes"] == [] diff --git a/tests/fast/rollout/generate_hub/test_multi_turn.py b/tests/fast/rollout/generate_hub/test_multi_turn.py index 40cfe36e33..aa67353865 100644 --- a/tests/fast/rollout/generate_hub/test_multi_turn.py +++ b/tests/fast/rollout/generate_hub/test_multi_turn.py @@ -127,9 +127,9 @@ def verify_samples(actual: Sample | list[Sample], expected: list[ExpectedSampleI prefix_cache_info=Sample.PrefixCacheInfo(), ) # Session server populates diagnostic metadata (token IDs, - # trim config, mismatch analysis) that varies with mock setup. - # Strip these before comparing sample structure. - for key in ("tito_session_mismatch", "accumulated_token_ids", "max_trim_tokens"): + # trim config, mismatch analysis, dashboard lifecycle timing) that + # varies with mock setup. Strip these before comparing structure. + for key in ("tito_session_mismatch", "accumulated_token_ids", "max_trim_tokens", "lifecycle"): actual_partial.metadata.pop(key, None) assert actual_partial == expected_item.partial_sample @@ -272,6 +272,26 @@ def test_two_turns_with_tool_call(self, variant, generation_env): class TestExitConditions: + @pytest.mark.parametrize( + "generation_env", + [{"args_kwargs": {"extra_argv": ["--save-debug-trajectory-data", "/unused/{rollout_id}.jsonl"]}}], + indirect=True, + ) + def test_conversation_metadata_records_turns(self, variant, generation_env): + if variant != "multi_turn_single_sample": + pytest.skip("conversation recording asserted once, on the engine multi-turn path") + generation_env.mock_server.process_fn = TwoTurnStub.process_fn + + S = TwoTurnStub + result = _run_generate(variant, generation_env, make_sample(prompt=S.PROMPT)) + + messages = result.sample.metadata["messages"] + # prompt message + turn-1 assistant + its two tool results + turn-2 assistant + assert [m["role"] for m in messages] == ["user", "assistant", "tool", "tool", "assistant"] + assert messages[1]["content"] == S.FIRST_RESPONSE + assert {m["name"] for m in messages[2:4]} == {"get_year", "get_temperature"} + assert messages[4]["content"] == S.SECOND_RESPONSE + def test_partial_rollout_not_supported(self, variant, generation_env): if is_agentic_variant(variant): pytest.skip("agentic_tool_call does not check partial_rollout flag") diff --git a/tests/fast/rollout/generate_utils/test_lifecycle_metadata.py b/tests/fast/rollout/generate_utils/test_lifecycle_metadata.py new file mode 100644 index 0000000000..e3756461a5 --- /dev/null +++ b/tests/fast/rollout/generate_utils/test_lifecycle_metadata.py @@ -0,0 +1,39 @@ +"""Pin for the dashboard lifecycle call site in +``compute_samples_from_openai_records`` (design doc §18.3): sample assembly +must keep chat-record timing on the Sample. If a session refactor moves the +assembly and drops the one-line call, this fails loudly.""" + +from types import SimpleNamespace + +from miles.utils.lifecycle import TrajectoryLifecycle, attach_lifecycle_metadata +from miles.utils.types import Sample + + +def _record(ts, latency=None, arrived=None): + meta_info = {} if latency is None else {"e2e_latency": latency} + return SimpleNamespace(timestamp=ts, request_timestamp=arrived, response={"choices": [{"meta_info": meta_info}]}) + + +def test_segment_shape_with_and_without_latency(): + first, second = Sample(index=1), Sample(index=1) + r1, r2 = _record(100.0, latency=8.5), _record(130.0) + + attach_lifecycle_metadata(first, r1, None, turn=1) + assert first.metadata["lifecycle"] == dict(t0=91.5, t1=100.0, turn=1) + + attach_lifecycle_metadata(second, r2, r1, turn=2) + assert second.metadata["lifecycle"] == dict(t0=None, t1=130.0, turn=2, prev_t1=100.0) + + +def test_request_timestamp_bounds_the_agent_gap(): + sample = Sample(index=1) + prev = _record(100.0) + record = _record(130.0, latency=10.0, arrived=104.0) + attach_lifecycle_metadata(sample, record, prev, turn=2) + # req_ts closes the agent gap at the server edge; [104, 120] is queueing + assert sample.metadata["lifecycle"] == dict(t0=120.0, t1=130.0, turn=2, req_ts=104.0, prev_t1=100.0) + + +def test_sink_defaults_to_none(): + # without the dashboard installed every core probe is a guarded no-op + assert TrajectoryLifecycle().sink is None diff --git a/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py b/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py index b09126dea9..21fafb8198 100644 --- a/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py +++ b/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py @@ -19,7 +19,8 @@ # ── helpers ────────────────────────────────────────────────────────── -_ARGS = SimpleNamespace() +_ARGS = SimpleNamespace(save_debug_trajectory_data=None) +_ARGS_RECORDING = SimpleNamespace(save_debug_trajectory_data="/unused/{rollout_id}.jsonl") def _mock_tokenizer(): @@ -155,6 +156,32 @@ def test_multiple_records_produce_multiple_samples(self): assert samples[0].tokens == [1, 2, 10] assert samples[1].tokens == [1, 2, 10, 20, 30] + def test_last_sample_carries_raw_conversation(self): + tok = _mock_tokenizer() + records = [ + _make_record(prompt_token_ids=[1, 2], output_token_ids=[10]), + _make_record(prompt_token_ids=[1, 2, 10, 20], output_token_ids=[30]), + ] + + samples = compute_samples_from_openai_records(_ARGS_RECORDING, _make_input_sample(), records, tok) + + assert "messages" not in (samples[0].metadata or {}) + assert samples[1].metadata["messages"] == records[1].request["messages"] + [ + records[1].response["choices"][0]["message"] + ] + + def test_merge_keeps_last_conversation_snapshot(self): + tok = _mock_tokenizer() + records = [ + _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10, 11]), + _make_record(prompt_token_ids=[1, 2, 3, 10, 11, 20, 21], output_token_ids=[30, 31]), + ] + + samples = compute_samples_from_openai_records(_ARGS_RECORDING, _make_input_sample(), records, tok) + merged = merge_samples(samples, tok) + + assert merged.metadata["messages"] == samples[-1].metadata["messages"] + def test_finish_reason_length_gives_truncated(self): tok = _mock_tokenizer() record = _make_record( diff --git a/tests/fast/rollout/test_trajectory_dump.py b/tests/fast/rollout/test_trajectory_dump.py new file mode 100644 index 0000000000..c27a1f8176 --- /dev/null +++ b/tests/fast/rollout/test_trajectory_dump.py @@ -0,0 +1,79 @@ +import argparse +import json + +import pytest +import torch + +from miles.ray.rollout.debug_data import save_debug_rollout_data, trajectory_rows +from miles.utils.types import Sample + + +def make_sample(**kwargs): + defaults = dict(index=7, group_index=2, prompt="rendered prompt", status=Sample.Status.COMPLETED) + return Sample(**{**defaults, **kwargs}) + + +MESSAGES = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + {"role": "tool", "name": "get_year", "content": "2026"}, + {"role": "assistant", "content": "done"}, +] + + +def test_rows_carry_the_raw_conversation(): + sample = make_sample(reward=1.0, metadata={"messages": MESSAGES}) + (row,) = trajectory_rows([sample]) + assert row["sample_index"] == 7 and row["group_index"] == 2 + assert row["reward"] == 1.0 and row["prompt"] == "rendered prompt" + assert row["messages"] == MESSAGES + + +def test_rows_skip_samples_without_a_conversation(): + with_messages = make_sample(index=1, metadata={"messages": MESSAGES}) + without = make_sample(index=2) + rows = trajectory_rows([without, with_messages]) + assert [row["sample_index"] for row in rows] == [1] + + +def test_non_numeric_reward_becomes_null(): + sample = make_sample(reward={"score": 1}, metadata={"messages": MESSAGES}) + (row,) = trajectory_rows([sample]) + assert row["reward"] is None + + +@pytest.fixture +def dump_args(tmp_path): + return argparse.Namespace( + save_debug_rollout_data=str(tmp_path / "rollout_data" / "{rollout_id}.pt"), + save_debug_trajectory_data=str(tmp_path / "trajectory" / "{rollout_id}.jsonl"), + ) + + +def test_save_writes_sidecar_and_strips_messages_from_pt(dump_args, tmp_path): + samples = [make_sample(index=i, metadata={"messages": MESSAGES}) for i in range(2)] + save_debug_rollout_data(dump_args, samples, 5, evaluation=False) + lines = (tmp_path / "trajectory" / "5.jsonl").read_text().splitlines() + assert [json.loads(line)["sample_index"] for line in lines] == [0, 1] + saved = torch.load(tmp_path / "rollout_data" / "5.pt", weights_only=False) + assert all("messages" not in s["metadata"] for s in saved["samples"]) + + +def test_save_eval_flattens_datasets(dump_args, tmp_path): + data = {"math": {"samples": [make_sample(metadata={"messages": MESSAGES})]}} + save_debug_rollout_data(dump_args, data, 1, evaluation=True) + (line,) = (tmp_path / "trajectory" / "eval_1.jsonl").read_text().splitlines() + assert json.loads(line)["messages"] == MESSAGES + + +def test_no_conversations_means_no_file(dump_args, tmp_path): + save_debug_rollout_data(dump_args, [make_sample()], 9, evaluation=False) + assert (tmp_path / "rollout_data" / "9.pt").exists() + assert not (tmp_path / "trajectory").exists() + + +def test_messages_stay_in_pt_when_sidecar_disabled(dump_args, tmp_path): + dump_args.save_debug_trajectory_data = None + save_debug_rollout_data(dump_args, [make_sample(metadata={"messages": MESSAGES})], 8, evaluation=False) + saved = torch.load(tmp_path / "rollout_data" / "8.pt", weights_only=False) + assert saved["samples"][0]["metadata"]["messages"] == MESSAGES diff --git a/tests/fast/test_timer_event_sinks.py b/tests/fast/test_timer_event_sinks.py new file mode 100644 index 0000000000..5250187944 --- /dev/null +++ b/tests/fast/test_timer_event_sinks.py @@ -0,0 +1,77 @@ +import time + +import pytest + +from miles.utils.timer import Timer, inverse_timer, timer + + +@pytest.fixture +def fresh_timer(): + instance = Timer() + saved = (instance.timers, instance.start_time, instance.event_sinks) + instance.timers, instance.start_time, instance.event_sinks = {}, {}, [] + yield instance + instance.timers, instance.start_time, instance.event_sinks = saved + + +def test_sink_receives_interval(fresh_timer): + events = [] + fresh_timer.event_sinks.append(lambda name, t0, t1: events.append((name, t0, t1))) + + before = time.time() + with timer("phase_a"): + time.sleep(0.01) + after = time.time() + + [(name, t0, t1)] = events + assert name == "phase_a" + assert before <= t0 <= t1 <= after + assert t1 - t0 == pytest.approx(fresh_timer.timers["phase_a"]) + + +def test_no_sinks_keeps_behavior_identical(fresh_timer): + with timer("phase_b"): + pass + assert list(fresh_timer.timers) == ["phase_b"] + assert fresh_timer.start_time == {} + + +def test_multiple_sinks_and_nested_timers(fresh_timer): + events = [] + fresh_timer.event_sinks.append(lambda *e: events.append(("s1", e[0]))) + fresh_timer.event_sinks.append(lambda *e: events.append(("s2", e[0]))) + + fresh_timer.start("outer") + with inverse_timer("outer"): + with timer("inner"): + pass + fresh_timer.end("outer") + + tags = [tag for tag, _ in events] + assert tags == ["s1", "s2"] * 3 # every end() fans out to every sink, in order + timer_names = [name for _, name in events] + # inverse_timer ends "outer" first, "inner" ends inside, "outer" ends last + assert timer_names == ["outer", "outer", "inner", "inner", "outer", "outer"] + + +def test_sink_with_begin_hears_starts(): + timer = Timer() + begins, ends = [], [] + + class Sink: + def begin(self, name, t0): + begins.append((name, t0)) + + def __call__(self, name, t0, t1): + ends.append(name) + + timer.event_sinks.append(Sink()) + try: + timer.start("long_phase") + assert [name for name, _ in begins] == ["long_phase"] # visible while open + assert ends == [] + timer.end("long_phase") + assert ends == ["long_phase"] + finally: + timer.event_sinks.clear() + timer.reset()