[WIP] miles dashboard

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(

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,

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
+```

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,
+    )

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