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?, live occupancy + free slots.

    Occupancy is read from the runner's `inflight` counters. It used to read a
    `busy` key, but nothing ever wrote to that after load_gpu_nodes() set it False
    — so every node reported idle even while running jobs. Don't reintroduce it.
    """
    pool = []
    for n in getattr(runner, "gpu_nodes", []):
        inflight = n.get("inflight") or {}
        gpu, cpu = int(inflight.get("gpu", 0)), int(inflight.get("cpu", 0))
        gpu_cap, cpu_cap = runner._node_capacity(n, "gpu"), runner._node_capacity(n, "cpu")
        ops = n.get("operators")
        pool.append({
            "name": n.get("name", n.get("ssh", "?")),
            "remote": not n.get("local", False),
            "busy": (gpu + cpu) > 0,
            "gpu": {"running": gpu, "slots": gpu_cap, "free": max(0, gpu_cap - gpu)},
            "cpu": {"running": cpu, "slots": cpu_cap, "free": max(0, cpu_cap - cpu)},
            "operators": len(ops) if ops else "all",
        })
    return pool


def nodes_detail(con, runner) -> dict:
    """Per-node capability + live work — the real per-machine view.

    /api/system's gpu_pool only carries name/remote/busy. This joins the runner's
    live inflight counters with the job table so the UI can show who is running
    what, and why a heavy job waits while a light one flies.

    `running` is attributed via jobs.node, written at dispatch (2026-08-05).
    Before that column existed this had to infer placement, and every job was
    credited to the primary while busy remotes rendered as idle.
    """
    nodes = getattr(runner, "gpu_nodes", [])
    local_name = next((n.get("name") for n in nodes if n.get("local")), "local")

    rows = con.execute(
        "SELECT id, operator, status, node, created_at, started_at FROM jobs "
        "WHERE status IN ('running','queued') ORDER BY created_at"
    ).fetchall()

    running_by_node: dict[str, list] = {}
    queued: list[dict] = []
    for r in rows:
        rec = {"id": r["id"], "operator": r["operator"],
               "created_at": r["created_at"], "started_at": r["started_at"]}
        if r["status"] == "running":
            # legacy rows predate jobs.node; credit them to the primary as before
            running_by_node.setdefault(r["node"] or local_name, []).append(rec)
        else:
            queued.append(rec)

    out, by_operator = [], {}
    for n in nodes:
        name = n.get("name", n.get("ssh", "?"))
        ops = n.get("operators")
        explicit = bool(ops)
        op_ids = list(ops) if explicit else sorted(runner.operators or {})
        for op in op_ids:
            by_operator.setdefault(op, []).append(name)

        gpu_cap = runner._node_capacity(n, "gpu")
        inflight = n.get("inflight") or {}
        # a queued job this node could pick up if it had a free slot
        eligible = [q for q in queued
                    if (not explicit or q["operator"] in op_ids)
                    and runner.lane_of(q["operator"]) == "gpu"]

        out.append({
            "name": name,
            "remote": not n.get("local", False),
            "ssh": n.get("ssh"),
            "gpu_capacity": gpu_cap,
            "cpu_capacity": runner._node_capacity(n, "cpu"),
            "gpu_inflight": int(inflight.get("gpu", 0)),
            "cpu_inflight": int(inflight.get("cpu", 0)),
            "operators_explicit": explicit,
            "operators": op_ids,
            "operator_count": len(op_ids),
            "running": running_by_node.get(name, []),
            "eligible_queued": eligible,
            "notes": n.get("_notes") or [],
        })

    placed = {j["id"] for v in running_by_node.values() for j in v}
    gpu_ops = sorted(op for op in (runner.operators or {})
                     if runner.lane_of(op) == "gpu")
    return {"nodes": out, "by_operator": by_operator,
            "queued_unassigned": [q for q in queued if q["id"] not in placed],
            "gpu_ops": gpu_ops}


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