modelbeast/server/sysinfo.py

"""System snapshot for the dashboard — CPU/RAM/disk (psutil), Apple GPU (macmon,
no sudo), queue/lane occupancy, recent jobs. Everything cached so refresh-spam is
cheap; the slow bits (du of data/ and HF cache) refresh at most every 5 min."""
import json
import os
import shutil
import subprocess
import time
from pathlib import Path

import psutil

from . import db, store

_cache: dict = {}


def _cached(key: str, ttl: float, fn):
    now = time.time()
    hit = _cache.get(key)
    if hit and now - hit[0] < ttl:
        return hit[1]
    val = fn()
    _cache[key] = (now, val)
    return val


def _macmon_gpu():
    def fetch():
        try:
            out = subprocess.run(["macmon", "pipe", "-s", "1"], capture_output=True,
                                 text=True, timeout=6).stdout.strip().splitlines()
            sample = json.loads(out[-1])
            gpu = sample.get("gpu_usage") or [0, 0.0]
            return {"util_percent": round(gpu[1] * 100, 1),
                    "power_w": round(sample.get("gpu_power", 0.0), 1),
                    "temp_c": round(sample.get("temp", {}).get("gpu_temp_avg", 0.0), 1),
                    "source": "macmon"}
        except FileNotFoundError:
            return {"source": "unavailable", "hint": "brew install macmon for GPU stats"}
        except Exception as e:
            return {"source": "unavailable", "hint": str(e)[:80]}
    return _cached("gpu", 5.0, fetch)


def _du_gb(path: Path):
    try:
        out = subprocess.run(["du", "-sk", str(path)], capture_output=True,
                             text=True, timeout=60).stdout.split()
        return round(int(out[0]) / (1024 * 1024), 1)
    except Exception:
        return None


def _lanes(con, runner):
    # gpu lane capacity = number of gpu nodes in the pool (this Mac + remote workers)
    gpu_limit = max(1, len(getattr(runner, "gpu_nodes", []) or []))
    lanes = {"gpu": {"limit": gpu_limit, "running": 0, "queued": 0},
             "cpu": {"limit": 3, "running": 0, "queued": 0},
             "net": {"limit": 6, "running": 0, "queued": 0}}
    rows = con.execute("SELECT operator, status FROM jobs WHERE status IN ('queued','running')").fetchall()
    for r in rows:
        lane = runner.lane_of(r["operator"])
        if lane not in lanes:
            lane = "cpu"
        lanes[lane]["running" if r["status"] == "running" else "queued"] += 1
    return lanes


def gpu_pool(runner) -> list[dict]:
    """Per-node status for the dashboard: name, remote?, currently busy."""
    return [{"name": n.get("name", n.get("ssh", "?")),
             "remote": not n.get("local", False),
             "busy": n.get("busy", False)} for n in getattr(runner, "gpu_nodes", [])]


def _jobs_24h(con):
    cutoff = time.time() - 86400
    rows = con.execute(
        "SELECT operator, status FROM jobs WHERE finished_at >= ?", (cutoff,)).fetchall()
    out = {"done": 0, "error": 0, "cancelled": 0, "by_operator": {}}
    for r in rows:
        out[r["status"]] = out.get(r["status"], 0) + 1
        out["by_operator"][r["operator"]] = out["by_operator"].get(r["operator"], 0) + 1
    return out


def snapshot(con, runner) -> dict:
    def stats():
        per_core = psutil.cpu_percent(interval=0.15, percpu=True)
        vm = psutil.virtual_memory()
        du = shutil.disk_usage(str(db.DATA))
        try:
            load1 = os.getloadavg()[0]
        except OSError:
            load1 = 0.0
        return {
            "cpu": {"percent": round(sum(per_core) / len(per_core), 1),
                    "cores": len(per_core), "per_core": [round(c, 0) for c in per_core],
                    "load1": round(load1, 2)},
            "ram": {"used_gb": round((vm.total - vm.available) / 1e9, 1),
                    "total_gb": round(vm.total / 1e9, 1), "percent": vm.percent},
            "disk": {"free_gb": round(du.free / 1e9, 1), "total_gb": round(du.total / 1e9, 1)},
            "uptime_s": round(time.time() - psutil.boot_time()),
        }
    base = _cached("stats", 5.0, stats)
    hf = Path.home() / ".cache" / "huggingface"
    return {
        **base,
        "gpu": _macmon_gpu(),
        "lanes": _lanes(con, runner),
        "nodes": gpu_pool(runner),
        "jobs_24h": _jobs_24h(con),
        "data_dir_gb": _cached("data_du", 300.0, lambda: _du_gb(db.DATA)),
        "hf_cache_gb": _cached("hf_du", 300.0, lambda: _du_gb(hf) if hf.exists() else None),
        "operators": len(runner.operators),
    }


def recent_jobs(con, runner, limit: int = 20, user_id: str | None = None) -> list[dict]:
    cutoff = time.time() - 86400
    if user_id is None:  # owner: all users
        rows = con.execute(
            "SELECT * FROM jobs WHERE status IN ('done','error','cancelled') "
            "AND finished_at >= ? ORDER BY finished_at DESC LIMIT ?", (cutoff, limit)).fetchall()
    else:  # guest: only their own
        rows = con.execute(
            "SELECT * FROM jobs WHERE status IN ('done','error','cancelled') "
            "AND finished_at >= ? AND user_id = ? ORDER BY finished_at DESC LIMIT ?",
            (cutoff, user_id, limit)).fetchall()
    umap = {u_id: name for u_id, name in
            con.execute("SELECT id, username FROM users").fetchall()}
    out = []
    for r in rows:
        j = db.row_to_dict(r)
        outs = con.execute(
            "SELECT id, kind, name FROM assets WHERE parent_job = ?", (j["id"],)).fetchall()
        dur = (j["finished_at"] - j["started_at"]) if j.get("started_at") and j.get("finished_at") else None
        out.append({
            "id": j["id"], "operator": j["operator"], "status": j["status"],
            "username": umap.get(j.get("user_id")) or ("—" if j.get("user_id") else "owner"),
            "duration_s": round(dur, 1) if dur else None,
            "outputs": [{"id": o["id"], "kind": o["kind"], "name": o["name"]} for o in outs],
        })
    return out