server.py

#!/usr/bin/env python3
import json
import os
import shutil
import signal
import subprocess
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import psutil

HOST = "0.0.0.0"
PORT = 9090
ROOT = os.path.dirname(os.path.abspath(__file__))
INDEX = os.path.join(ROOT, "index.html")

NVIDIA_SMI = shutil.which("nvidia-smi")

GPU_QUERY = (
    "index,name,utilization.gpu,utilization.memory,"
    "memory.total,memory.used,memory.free,"
    "temperature.gpu,fan.speed,power.draw,power.limit,"
    "clocks.current.graphics,clocks.current.memory"
)


def read_gpus():
    if not NVIDIA_SMI:
        return []
    try:
        out = subprocess.check_output(
            [NVIDIA_SMI, f"--query-gpu={GPU_QUERY}", "--format=csv,noheader,nounits"],
            stderr=subprocess.DEVNULL,
            timeout=2,
        ).decode("utf-8", errors="ignore")
    except Exception:
        return []

    gpus = []
    for line in out.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 13:
            continue

        def num(v):
            try:
                if v in ("[N/A]", "[Not Supported]", "N/A", ""):
                    return None
                return float(v)
            except ValueError:
                return None

        gpus.append({
            "index": int(parts[0]),
            "name": parts[1],
            "util": num(parts[2]),
            "mem_util": num(parts[3]),
            "mem_total": num(parts[4]),
            "mem_used": num(parts[5]),
            "mem_free": num(parts[6]),
            "temp": num(parts[7]),
            "fan": num(parts[8]),
            "power": num(parts[9]),
            "power_limit": num(parts[10]),
            "clock_gpu": num(parts[11]),
            "clock_mem": num(parts[12]),
        })
    return gpus


_last_net = psutil.net_io_counters()
_last_net_t = time.time()
_net_lock = threading.Lock()

_proc_cache = {}
_proc_lock = threading.Lock()


def _refresh_proc_cache():
    current = set(psutil.pids())
    with _proc_lock:
        for pid in list(_proc_cache):
            if pid not in current:
                _proc_cache.pop(pid, None)
        for pid in current:
            if pid not in _proc_cache:
                try:
                    p = psutil.Process(pid)
                    p.cpu_percent(None)
                    _proc_cache[pid] = p
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass


def read_top_processes(limit=25):
    _refresh_proc_cache()
    procs = []
    with _proc_lock:
        items = list(_proc_cache.items())
    for pid, p in items:
        try:
            with p.oneshot():
                cpu = p.cpu_percent(None)
                mem = p.memory_info().rss
                name = p.name()
                user = p.username()
                status = p.status()
        except (psutil.NoSuchProcess, psutil.AccessDenied, ProcessLookupError):
            continue
        procs.append({
            "pid": pid,
            "name": name,
            "user": user,
            "cpu": cpu,
            "mem": mem,
            "status": status,
        })
    procs.sort(key=lambda x: (x["cpu"], x["mem"]), reverse=True)
    return procs[:limit]


def read_gpu_processes():
    if not NVIDIA_SMI:
        return []
    try:
        out = subprocess.check_output(
            [NVIDIA_SMI,
             "--query-compute-apps=pid,process_name,used_memory",
             "--format=csv,noheader,nounits"],
            stderr=subprocess.DEVNULL,
            timeout=2,
        ).decode("utf-8", errors="ignore")
    except Exception:
        return []

    procs = []
    for line in out.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 3:
            continue
        try:
            pid = int(parts[0])
        except ValueError:
            continue
        try:
            mem_mb = float(parts[2])
        except ValueError:
            mem_mb = 0.0
        user = None
        cmdline = None
        try:
            pp = psutil.Process(pid)
            user = pp.username()
            cmdline = " ".join(pp.cmdline()) or None
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
        procs.append({
            "pid": pid,
            "name": parts[1],
            "user": user,
            "cmdline": cmdline,
            "vram_bytes": mem_mb * 1024 * 1024,
        })
    procs.sort(key=lambda x: x["vram_bytes"], reverse=True)
    return procs


def kill_pid(pid, sig="TERM"):
    signame = "SIG" + str(sig).upper().lstrip("SIG")
    signum = getattr(signal, signame, None)
    if signum is None:
        return False, f"unknown signal {sig}"
    if pid in (0, 1, os.getpid()):
        return False, "refused (init/self)"
    try:
        os.kill(pid, signum)
        return True, signame
    except ProcessLookupError:
        return False, "no such process"
    except PermissionError:
        return False, "permission denied"
    except Exception as e:
        return False, str(e)


def read_net_rate():
    global _last_net, _last_net_t
    with _net_lock:
        now = time.time()
        cur = psutil.net_io_counters()
        dt = max(now - _last_net_t, 1e-3)
        rx = (cur.bytes_recv - _last_net.bytes_recv) / dt
        tx = (cur.bytes_sent - _last_net.bytes_sent) / dt
        _last_net = cur
        _last_net_t = now
        return rx, tx


def collect_stats():
    per_cpu = psutil.cpu_percent(interval=None, percpu=True)
    cpu_total = sum(per_cpu) / len(per_cpu) if per_cpu else 0.0
    freq = psutil.cpu_freq()
    vm = psutil.virtual_memory()
    sm = psutil.swap_memory()
    la1, la5, la15 = os.getloadavg()
    rx, tx = read_net_rate()

    temps = {}
    try:
        raw = psutil.sensors_temperatures(fahrenheit=False) or {}
        for name, entries in raw.items():
            vals = [e.current for e in entries if e.current is not None]
            if vals:
                temps[name] = max(vals)
    except Exception:
        pass

    return {
        "ts": time.time(),
        "cpu": {
            "total": cpu_total,
            "per_core": per_cpu,
            "count": psutil.cpu_count(logical=True),
            "physical": psutil.cpu_count(logical=False),
            "freq_mhz": freq.current if freq else None,
            "load": [la1, la5, la15],
        },
        "mem": {
            "total": vm.total,
            "used": vm.used,
            "available": vm.available,
            "percent": vm.percent,
            "swap_total": sm.total,
            "swap_used": sm.used,
            "swap_percent": sm.percent,
        },
        "net": {"rx_bps": rx, "tx_bps": tx},
        "temps": temps,
        "gpus": read_gpus(),
        "uptime": time.time() - psutil.boot_time(),
    }


# Prime CPU counter so first call returns sensible values
psutil.cpu_percent(interval=None, percpu=True)


class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        pass  # silence default access log

    def _send_html(self, path):
        try:
            with open(path, "rb") as f:
                data = f.read()
        except OSError:
            self.send_error(404)
            return
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):
        if self.path in ("/", "/index.html"):
            self._send_html(INDEX)
            return

        if self.path == "/stats":
            payload = json.dumps(collect_stats()).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.send_header("Cache-Control", "no-store")
            self.send_header("Access-Control-Allow-Origin", "*")
            self.end_headers()
            self.wfile.write(payload)
            return

        if self.path == "/processes":
            payload = json.dumps({
                "ts": time.time(),
                "gpu": read_gpu_processes(),
                "top": read_top_processes(25),
                "runner": {"uid": os.geteuid(), "user": psutil.Process().username()},
            }).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.send_header("Cache-Control", "no-store")
            self.send_header("Access-Control-Allow-Origin", "*")
            self.end_headers()
            self.wfile.write(payload)
            return

        if self.path == "/stream":
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("Cache-Control", "no-cache")
            self.send_header("Connection", "keep-alive")
            self.send_header("X-Accel-Buffering", "no")
            self.send_header("Access-Control-Allow-Origin", "*")
            self.end_headers()
            try:
                while True:
                    payload = json.dumps(collect_stats())
                    self.wfile.write(b"data: " + payload.encode("utf-8") + b"\n\n")
                    self.wfile.flush()
                    time.sleep(1.0)
            except (BrokenPipeError, ConnectionResetError):
                return
            return

        self.send_error(404)

    def do_POST(self):
        if self.path == "/kill":
            try:
                length = int(self.headers.get("Content-Length", "0") or 0)
                body = self.rfile.read(length) if length else b"{}"
                data = json.loads(body.decode("utf-8") or "{}")
                pid = int(data["pid"])
                sig = data.get("signal", "TERM")
            except Exception as e:
                self._send_json({"ok": False, "error": f"bad request: {e}"}, 400)
                return
            ok, info = kill_pid(pid, sig)
            self._send_json(
                {"ok": ok, "pid": pid, "signal": info if ok else None,
                 "error": None if ok else info},
                200 if ok else 403 if info == "permission denied" else 404 if info == "no such process" else 500,
            )
            return
        self.send_error(404)

    def _send_json(self, obj, code=200):
        payload = json.dumps(obj).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(payload)))
        self.send_header("Cache-Control", "no-store")
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(payload)


def main():
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    print(f"sysmon listening on http://{HOST}:{PORT}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        server.shutdown()


if __name__ == "__main__":
    main()