gpu_data.py

"""Data collection for AMD GPUs.

Standalone module (no StreamController imports) so it can be tested outside
the app. Prefers streaming JSON from `amdgpu_top -J`, falls back to reading
sysfs (/sys/class/drm/card*/device) if amdgpu_top is not available.

All samples land in shared, thread-safe history deques (one per GPU/metric),
so every key on the deck reads from the same collector and new keys instantly
show past history.
"""

import glob
import json
import os
import subprocess
import threading
from collections import deque

HISTORY_LENGTH = 600  # samples (~10 min at 1 Hz)
SAMPLE_PERIOD_MS = 1000


def is_in_flatpak() -> bool:
    return os.path.isfile("/.flatpak-info")


def _value(node):
    """amdgpu_top wraps most values as {"unit": ..., "value": ...}."""
    if isinstance(node, dict):
        node = node.get("value")
    return node


def _sensor(dev: dict, name: str):
    return _value((dev.get("Sensors") or {}).get(name))


def _activity(dev: dict, name: str):
    return _value((dev.get("gpu_activity") or {}).get(name))


def _vram(dev: dict, name: str):
    return _value((dev.get("VRAM") or {}).get(name))


def _power(dev: dict):
    for name in ("Average Power", "GFX Power", "Input Power"):
        v = _sensor(dev, name)
        if v is not None:
            return v
    return None


class Metric:
    def __init__(self, key: str, title: str, short: str, fmt,
                 from_json, from_sysfs, scale_key: str = None,
                 default_max: float = None, color=(255, 255, 255)):
        self.key = key
        self.title = title          # shown in the config dropdown
        self.short = short          # shown as the key's top label
        self.fmt = fmt              # value -> display string
        self.from_json = from_json  # amdgpu_top device dict -> value
        self.from_sysfs = from_sysfs  # _SysfsCard -> value
        self.scale_key = scale_key  # static-info key holding the fixed max
        self.default_max = default_max  # None -> auto-scale to the data
        self.color = color          # default line color (r, g, b)

    def format(self, value) -> str:
        if value is None:
            return "--"
        return self.fmt(value)


def _fmt_pct(v): return f"{v:.0f}%"
def _fmt_temp(v): return f"{v:.0f}°C"
def _fmt_watt(v): return f"{v:.0f}W"
def _fmt_mhz(v): return f"{v:.0f}MHz"
def _fmt_rpm(v): return f"{v:.0f}rpm"


def _fmt_mib(v):
    return f"{v / 1024:.1f}G" if v >= 1024 else f"{v:.0f}M"


METRICS = {m.key: m for m in [
    Metric("activity_gfx", "GPU Activity (GFX)", "GPU", _fmt_pct,
           lambda dev: _activity(dev, "GFX"),
           lambda card: card.read_value("gpu_busy_percent"),
           default_max=100, color=(227, 54, 54)),
    Metric("activity_mem", "Memory Activity", "MEM", _fmt_pct,
           lambda dev: _activity(dev, "Memory"),
           lambda card: card.read_value("mem_busy_percent"),
           default_max=100, color=(227, 139, 54)),
    Metric("activity_media", "Media Engine", "MEDIA", _fmt_pct,
           lambda dev: _activity(dev, "MediaEngine"),
           lambda card: None,
           default_max=100, color=(227, 200, 54)),
    Metric("temp_edge", "Edge Temperature", "EDGE", _fmt_temp,
           lambda dev: _sensor(dev, "Edge Temperature"),
           lambda card: card.temp("edge"),
           scale_key="temp_edge_crit", default_max=105, color=(255, 112, 67)),
    Metric("temp_junction", "Junction Temperature", "TEMP", _fmt_temp,
           lambda dev: _sensor(dev, "Junction Temperature"),
           lambda card: card.temp("junction"),
           scale_key="temp_junction_crit", default_max=105, color=(255, 87, 34)),
    Metric("temp_mem", "Memory Temperature", "VRAM °C", _fmt_temp,
           lambda dev: _sensor(dev, "Memory Temperature"),
           lambda card: card.temp("mem"),
           scale_key="temp_mem_crit", default_max=105, color=(255, 160, 0)),
    Metric("power", "Power Draw", "POWER", _fmt_watt,
           _power,
           lambda card: card.power(),
           scale_key="power_cap", default_max=None, color=(255, 214, 64)),
    Metric("vram_used", "VRAM Usage", "VRAM", _fmt_mib,
           lambda dev: _vram(dev, "Total VRAM Usage"),
           lambda card: card.mem_mib("mem_info_vram_used"),
           scale_key="vram_total", default_max=None, color=(171, 99, 227)),
    Metric("gtt_used", "GTT Usage", "GTT", _fmt_mib,
           lambda dev: _vram(dev, "Total GTT Usage"),
           lambda card: card.mem_mib("mem_info_gtt_used"),
           scale_key="gtt_total", default_max=None, color=(121, 134, 203)),
    Metric("sclk", "GPU Clock (SCLK)", "SCLK", _fmt_mhz,
           lambda dev: _sensor(dev, "GFX_SCLK"),
           lambda card: card.freq("sclk"),
           scale_key="sclk_max", default_max=None, color=(66, 165, 245)),
    Metric("mclk", "Memory Clock (MCLK)", "MCLK", _fmt_mhz,
           lambda dev: _sensor(dev, "GFX_MCLK"),
           lambda card: card.freq("mclk"),
           default_max=None, color=(38, 198, 218)),
    Metric("fan", "Fan Speed", "FAN", _fmt_rpm,
           lambda dev: _sensor(dev, "Fan"),
           lambda card: card.fan(),
           scale_key="fan_max", default_max=None, color=(102, 187, 106)),
]}

METRIC_ORDER = list(METRICS.keys())


class _SysfsCard:
    """Reader for one AMD GPU under /sys/class/drm/cardN/device."""

    def __init__(self, path: str):
        self.path = path
        hwmons = sorted(glob.glob(os.path.join(path, "hwmon", "hwmon*")))
        self.hwmon = hwmons[0] if hwmons else None

        self.temp_files = {}
        self.temp_crits = {}
        self.freq_files = {}
        if self.hwmon:
            for label_file in glob.glob(os.path.join(self.hwmon, "temp*_label")):
                label = self._read_str(label_file)
                if label:
                    self.temp_files[label] = label_file.replace("_label", "_input")
                    self.temp_crits[label] = label_file.replace("_label", "_crit")
            for label_file in glob.glob(os.path.join(self.hwmon, "freq*_label")):
                label = self._read_str(label_file)
                if label:
                    self.freq_files[label] = label_file.replace("_label", "_input")

    @staticmethod
    def _read_str(path):
        try:
            with open(path) as f:
                return f.read().strip()
        except OSError:
            return None

    def _read_num(self, path):
        raw = self._read_str(path)
        try:
            return float(raw)
        except (TypeError, ValueError):
            return None

    def read_value(self, name):
        return self._read_num(os.path.join(self.path, name))

    def temp(self, label):
        f = self.temp_files.get(label)
        v = self._read_num(f) if f else None
        return v / 1000 if v is not None else None

    def temp_crit(self, label):
        f = self.temp_crits.get(label)
        v = self._read_num(f) if f else None
        return v / 1000 if v is not None else None

    def freq(self, label):
        f = self.freq_files.get(label)
        v = self._read_num(f) if f else None
        return v / 1e6 if v is not None else None  # Hz -> MHz

    def power(self):
        if not self.hwmon:
            return None
        for name in ("power1_average", "power1_input"):
            v = self._read_num(os.path.join(self.hwmon, name))
            if v is not None:
                return v / 1e6  # µW -> W
        return None

    def power_cap(self):
        if not self.hwmon:
            return None
        v = self._read_num(os.path.join(self.hwmon, "power1_cap"))
        return v / 1e6 if v is not None else None

    def fan(self):
        if not self.hwmon:
            return None
        return self._read_num(os.path.join(self.hwmon, "fan1_input"))

    def fan_max(self):
        if not self.hwmon:
            return None
        return self._read_num(os.path.join(self.hwmon, "fan1_max"))

    def mem_mib(self, name):
        v = self.read_value(name)
        return v / (1024 * 1024) if v is not None else None

    def name(self):
        return "AMD GPU"

    @staticmethod
    def find_cards():
        cards = []
        for dev in sorted(glob.glob("/sys/class/drm/card[0-9]*/device")):
            vendor = _SysfsCard._read_str(os.path.join(dev, "vendor"))
            if vendor == "0x1002":
                cards.append(_SysfsCard(dev))
        return cards


class GPUDataProvider:
    """Samples GPU metrics once per second on a background thread."""

    def __init__(self):
        self._lock = threading.Lock()
        self._stop = threading.Event()
        self._thread = None
        self._proc = None
        self._sysfs_cards = None
        self._hist: dict = {}     # (gpu_index, metric_key) -> deque
        self._last: dict = {}     # (gpu_index, metric_key) -> latest value
        self._static: dict = {}   # gpu_index -> static info dict
        self.source = "starting"  # "amdgpu_top" | "sysfs" | "none"

    # -- lifecycle -----------------------------------------------------------

    def ensure_started(self):
        with self._lock:
            if self._thread is not None:
                return
            self._thread = threading.Thread(
                target=self._run, name="AMDGPUTop-Collector", daemon=True)
            self._thread.start()

    def stop(self, *args):
        self._stop.set()
        proc = self._proc
        if proc is not None:
            try:
                proc.kill()
            except OSError:
                pass

    # -- public accessors ----------------------------------------------------

    def device_count(self) -> int:
        with self._lock:
            return len(self._static)

    def device_name(self, gpu: int) -> str:
        with self._lock:
            return self._static.get(gpu, {}).get("name", "AMD GPU")

    def get_window(self, gpu: int, metric_key: str, n: int) -> list:
        """Last n samples, left-padded with None so the list is always n long."""
        with self._lock:
            dq = self._hist.get((gpu, metric_key))
            data = list(dq)[-n:] if dq else []
        return [None] * (n - len(data)) + data

    def get_latest(self, gpu: int, metric_key: str):
        with self._lock:
            return self._last.get((gpu, metric_key))

    def get_scale_max(self, gpu: int, metric: Metric):
        """Fixed y-axis max for the metric, or None if it should auto-scale."""
        with self._lock:
            static = self._static.get(gpu, {})
        if metric.scale_key:
            v = static.get(metric.scale_key)
            if v:
                return v
        return metric.default_max

    # -- collection ----------------------------------------------------------

    def _run(self):
        empty_streaks = 0
        while not self._stop.is_set():
            produced = self._stream_amdgpu_top()
            if self._stop.is_set():
                return
            empty_streaks = 0 if produced else empty_streaks + 1
            if empty_streaks >= 2:
                break  # amdgpu_top unusable -> sysfs
            self._stop.wait(2)
        if not self._stop.is_set():
            self._sysfs_loop()

    def _stream_amdgpu_top(self) -> bool:
        cmd = ["amdgpu_top", "-J", "-s", str(SAMPLE_PERIOD_MS), "-n", "0", "--no-pc"]
        if is_in_flatpak():
            cmd = ["flatpak-spawn", "--host"] + cmd
        try:
            proc = subprocess.Popen(
                cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
                text=True)
        except OSError:
            return False
        self._proc = proc
        produced = False
        try:
            for line in proc.stdout:
                if self._stop.is_set():
                    break
                try:
                    data = json.loads(line)
                except ValueError:
                    continue
                self._ingest_json(data)
                produced = True
                self.source = "amdgpu_top"
        finally:
            self._proc = None
            try:
                proc.kill()
                proc.wait(timeout=2)  # reap so restarts don't pile up zombies
            except (OSError, subprocess.TimeoutExpired):
                pass
        return produced

    def _ingest_json(self, data: dict):
        devices = data.get("devices") or []
        with self._lock:
            for i, dev in enumerate(devices):
                if i not in self._static:
                    static = self._static_from_json(dev)
                    self._enrich_from_sysfs(static, i)
                    self._static[i] = static
                for key, metric in METRICS.items():
                    try:
                        val = metric.from_json(dev)
                    except (KeyError, TypeError, ValueError):
                        val = None
                    self._append(i, key, val)

    @staticmethod
    def _static_from_json(dev: dict) -> dict:
        info = dev.get("Info") or {}
        return {
            "name": info.get("DeviceName") or "AMD GPU",
            "vram_total": _vram(dev, "Total VRAM"),
            "gtt_total": _vram(dev, "Total GTT"),
            "fan_max": _sensor(dev, "Fan Max"),
            "sclk_max": (info.get("GPU Clock") or {}).get("max"),
            "temp_edge_crit": _sensor(dev, "Edge Critical Temperature"),
            "temp_junction_crit": _sensor(dev, "Junction Critical Temperature"),
            "temp_mem_crit": _sensor(dev, "Memory Critical Temperature"),
            "power_cap": _sensor(dev, "Power Cap"),
        }

    def _enrich_from_sysfs(self, static: dict, index: int):
        """Fill gaps in amdgpu_top's static info (e.g. it reports no power cap)."""
        if self._sysfs_cards is None:
            self._sysfs_cards = _SysfsCard.find_cards()
        if index >= len(self._sysfs_cards):
            return
        card = self._sysfs_cards[index]
        if not static.get("power_cap"):
            static["power_cap"] = card.power_cap()
        if not static.get("fan_max"):
            static["fan_max"] = card.fan_max()

    def _sysfs_loop(self):
        cards = _SysfsCard.find_cards()
        if not cards:
            self.source = "none"
            return
        self.source = "sysfs"
        with self._lock:
            for i, card in enumerate(cards):
                self._static.setdefault(i, {
                    "name": card.name(),
                    "vram_total": card.mem_mib("mem_info_vram_total"),
                    "gtt_total": card.mem_mib("mem_info_gtt_total"),
                    "fan_max": card.fan_max(),
                    "sclk_max": None,
                    "temp_edge_crit": card.temp_crit("edge"),
                    "temp_junction_crit": card.temp_crit("junction"),
                    "temp_mem_crit": card.temp_crit("mem"),
                    "power_cap": card.power_cap(),
                })
        while not self._stop.is_set():
            samples = []  # read files outside the lock, it's contended by renderers
            for i, card in enumerate(cards):
                for key, metric in METRICS.items():
                    try:
                        val = metric.from_sysfs(card)
                    except OSError:
                        val = None
                    samples.append((i, key, val))
            with self._lock:
                for i, key, val in samples:
                    self._append(i, key, val)
            self._stop.wait(SAMPLE_PERIOD_MS / 1000)

    def _append(self, gpu: int, key: str, val):
        dq = self._hist.get((gpu, key))
        if dq is None:
            dq = self._hist[(gpu, key)] = deque(maxlen=HISTORY_LENGTH)
        dq.append(val)
        self._last[(gpu, key)] = val