neuron-tai/packages/node/meshnet_node/dashboard.py

"""Live node status dashboard — rich TUI with plain-text fallback."""

from __future__ import annotations

import os
import sys
import time
from collections import deque
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    pass


def is_interactive_tty() -> bool:
    """Return True when stdout is a real terminal (not CI / redirected / WSL2 dumb)."""
    if not sys.stdout.isatty():
        return False
    term = os.environ.get("TERM", "")
    if term in ("dumb", ""):
        return False
    return True


def _format_uptime(seconds: float) -> str:
    s = int(seconds)
    h, rem = divmod(s, 3600)
    m, sec = divmod(rem, 60)
    return f"{h:02d}:{m:02d}:{sec:02d}"


def _gpu_stats() -> list[dict]:
    """Return per-GPU utilization and VRAM stats, or empty list on CPU."""
    try:
        import torch  # type: ignore[import]

        if not torch.cuda.is_available():
            return []
        stats = []
        for i in range(torch.cuda.device_count()):
            props = torch.cuda.get_device_properties(i)
            used = torch.cuda.memory_allocated(i)
            total = props.total_memory
            # Utilization requires pynvml; skip gracefully if not available
            util = _nvml_gpu_util(i)
            stats.append(
                {
                    "index": i,
                    "name": props.name,
                    "used_gb": used / 1e9,
                    "total_gb": total / 1e9,
                    "util_pct": util,
                }
            )
        return stats
    except ImportError:
        return []


def _nvml_gpu_util(index: int) -> int | None:
    """Return GPU utilization % via pynvml, or None if unavailable."""
    try:
        import pynvml  # type: ignore[import]

        pynvml.nvmlInit()
        handle = pynvml.nvmlDeviceGetHandleByIndex(index)
        rates = pynvml.nvmlDeviceGetUtilizationRates(handle)
        return rates.gpu
    except Exception:
        return None


class _EMA:
    """Exponential moving average for tokens/sec."""

    def __init__(self, alpha: float = 0.1):
        self._alpha = alpha
        self._value: float | None = None

    def update(self, sample: float) -> float:
        if self._value is None:
            self._value = sample
        else:
            self._value = self._alpha * sample + (1 - self._alpha) * self._value
        return self._value

    @property
    def value(self) -> float:
        return self._value or 0.0


def _make_bar(pct: float, width: int = 10) -> str:
    filled = round(pct / 100 * width)
    return "█" * filled + "░" * (width - filled)


def run_dashboard(node, config: dict, start_time: float) -> None:
    """Start the live dashboard. Blocks until Ctrl-C. Returns cleanly."""
    if not is_interactive_tty():
        _run_plain_loop(node, config, start_time)
        return

    try:
        from rich.live import Live  # type: ignore[import]

        _run_rich_dashboard(node, config, start_time)
    except ImportError:
        _run_plain_loop(node, config, start_time)


def _build_rich_renderable(
    node, config: dict, start_time: float, tps_ema: _EMA, prev_req: list[int]
):
    from rich.table import Table  # type: ignore[import]
    from rich.panel import Panel  # type: ignore[import]
    from rich.columns import Columns  # type: ignore[import]
    from rich.text import Text  # type: ignore[import]

    uptime = time.monotonic() - start_time
    req_count = getattr(node, "chat_completion_count", 0)

    # Tokens/sec EMA (approximate: 20 tokens per request heuristic when no real counter)
    delta_req = req_count - prev_req[0]
    prev_req[0] = req_count
    if delta_req > 0:
        approx_tokens = delta_req * 20
        tps_ema.update(approx_tokens / 2.0)  # 2s interval

    gpu_stats = _gpu_stats()

    model_name = config.get("model_name") or config.get("model_hf_repo", "unknown").split("/")[-1]
    shard = ""
    if config.get("shard_start") is not None:
        shard = f"  shard {config['shard_start']}–{config['shard_end']}"

    # Header line
    header = Text(
        f"meshnet-node  {model_name} [{config.get('quantization', 'bf16')}]{shard}"
        f"  up {_format_uptime(uptime)}",
        style="bold white",
    )

    # GPU table
    gpu_table = Table(show_header=False, box=None, padding=(0, 1))
    gpu_table.add_column("label", style="dim", no_wrap=True)
    gpu_table.add_column("bar", no_wrap=True)
    gpu_table.add_column("vram", no_wrap=True, style="cyan")

    if gpu_stats:
        for g in gpu_stats:
            util = g["util_pct"]
            util_str = f"{_make_bar(util)}  {util:3d}%" if util is not None else "            n/a"
            vram_str = f"VRAM {g['used_gb']:.1f}/{g['total_gb']:.1f} GB"
            gpu_table.add_row(f"GPU {g['index']}  {g['name'][:20]}", util_str, vram_str)
    else:
        gpu_table.add_row("CPU mode", "", "no GPU detected")

    # Stats panel
    tps = tps_ema.value
    bar_len = min(8, max(0, int(tps / 10)))
    tps_bar = "▁▂▃▄▅▆▇█"[:bar_len].ljust(8)

    stats_lines = [
        f"Tokens/sec   {tps_bar}  {tps:.1f} t/s (EMA)",
        f"Requests     {req_count:,} served",
        f"Peers        0 connected  (gossip: US-017)",
        f"TAI earned   0.00 TAI  (payments: US-006)",
        f"Uptime       {_format_uptime(uptime)}",
        "",
        "[q] quit   [c] compact view",
    ]

    from rich.console import Group  # type: ignore[import]

    return Panel(
        Group(header, gpu_table, Text("\n".join(stats_lines))),
        title="[bold green]meshnet-node[/bold green]",
        border_style="green",
    )


def _run_rich_dashboard(node, config: dict, start_time: float) -> None:
    from rich.live import Live  # type: ignore[import]

    tps_ema = _EMA()
    prev_req = [0]

    try:
        with Live(
            _build_rich_renderable(node, config, start_time, tps_ema, prev_req),
            refresh_per_second=0.5,
            screen=False,
        ) as live:
            while True:
                time.sleep(2)
                live.update(
                    _build_rich_renderable(node, config, start_time, tps_ema, prev_req)
                )
    except KeyboardInterrupt:
        pass


def _run_plain_loop(node, config: dict, start_time: float) -> None:
    model_name = config.get("model_name") or config.get("model_hf_repo", "unknown").split("/")[-1]
    try:
        while True:
            uptime = time.monotonic() - start_time
            req = getattr(node, "chat_completion_count", 0)
            gpu_stats = _gpu_stats()
            vram_str = ""
            if gpu_stats:
                g = gpu_stats[0]
                vram_str = f"  VRAM{g['used_gb']:.1f}GB"
            print(
                f"[{model_name}  req{req}{vram_str}  up{_format_uptime(uptime)}]",
                flush=True,
            )
            time.sleep(2)
    except KeyboardInterrupt:
        pass