feat: GPU process list with per-process VRAM/utilization and kill
Add a GPU process table to the Performance tab showing which processes use the selected GPU, with per-process VRAM, GPU utilization, CPU, host memory, and command. Includes a kill endpoint (SIGTERM/SIGKILL, with optional parent kill) behind a confirm dialog. - nvcurve/hal/processes.py: NVML + /proc process enumeration - server.py: GET /api/processes, POST /api/processes/kill (refuses to signal init/kthreadd) - frontend: ProcessList component, API client, types - tests: test_processes.py - docs: API reference for the new endpoints
This commit is contained in:
ARIA committed
1 parent
b75b9d43e9
commit
c11ea73ead
10 files changed
+1008
-23
No files matched your search
@@ -0,0 +1,346 @@
|
||||
"""Tests for the GPU process list + kill endpoints.
|
||||
|
||||
Standalone (no pytest required):
|
||||
|
||||
python tests/test_processes.py
|
||||
|
||||
Also works under pytest if available. Covers the /proc reader (including
|
||||
comm values with spaces and parentheses), the kill endpoint's input
|
||||
validation, a real SIGTERM round-trip, and the GET /api/processes payload.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
|
||||
from nvcurve import server # noqa: E402
|
||||
from nvcurve.hal import processes # noqa: E402
|
||||
|
||||
PASS = 0
|
||||
FAIL = 0
|
||||
|
||||
|
||||
def check(name: str, cond: bool) -> None:
|
||||
global PASS, FAIL
|
||||
if cond:
|
||||
PASS += 1
|
||||
print(f" PASS {name}")
|
||||
else:
|
||||
FAIL += 1
|
||||
print(f" FAIL {name}")
|
||||
|
||||
|
||||
def test_proc_info_self() -> None:
|
||||
"""_proc_info must return sane data for our own process."""
|
||||
info = processes._proc_info(os.getpid())
|
||||
check("self info not None", info is not None)
|
||||
if info is None:
|
||||
return
|
||||
check(
|
||||
"user is non-empty string",
|
||||
bool(isinstance(info["user"], str) and info["user"]),
|
||||
)
|
||||
check(
|
||||
"cpu_pct is None or non-negative number",
|
||||
info["cpu_pct"] is None or (isinstance(info["cpu_pct"], float) and info["cpu_pct"] >= 0),
|
||||
)
|
||||
check("mem_bytes > 0", info["mem_bytes"] is not None and info["mem_bytes"] > 0)
|
||||
check("command non-empty", isinstance(info["command"], str) and len(info["command"]) > 0)
|
||||
|
||||
|
||||
def test_proc_info_missing() -> None:
|
||||
"""A pid that does not exist must yield None (not an exception)."""
|
||||
check("missing pid -> None", processes._proc_info(99999999) is None)
|
||||
|
||||
|
||||
def test_proc_info_comm_with_parens() -> None:
|
||||
"""comm containing spaces and parentheses must not break /proc/stat parsing.
|
||||
|
||||
The kernel comm is set via prctl(PR_SET_NAME); `exec -a` only changes
|
||||
argv[0] and would leave comm as the executable name, so the parens path
|
||||
would never be exercised.
|
||||
"""
|
||||
proc = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
"import ctypes, time\n"
|
||||
"ctypes.CDLL(None).prctl(15, b'proc (test) x', 0, 0, 0)\n"
|
||||
"time.sleep(30)",
|
||||
]
|
||||
)
|
||||
try:
|
||||
# Wait until /proc/<pid>/stat exists and comm actually carries the
|
||||
# parenthesised name (prctl may take a moment after exec).
|
||||
comm = ""
|
||||
for _ in range(100):
|
||||
try:
|
||||
with open(f"/proc/{proc.pid}/stat") as f:
|
||||
raw = f.read()
|
||||
comm = raw[raw.index("(") + 1 : raw.rfind(")")]
|
||||
except OSError:
|
||||
pass
|
||||
if "(" in comm:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
check("fixture comm contains parens", "(" in comm)
|
||||
info = processes._proc_info(proc.pid)
|
||||
check("weird-comm info not None", info is not None)
|
||||
if info is not None:
|
||||
check("weird-comm command parsed", "prctl" in info["command"])
|
||||
finally:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
|
||||
|
||||
def test_kill_endpoint_validation() -> None:
|
||||
"""The kill endpoint must reject bad input and unknown pids."""
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
client = TestClient(server.app)
|
||||
r = client.post(
|
||||
"/api/processes/kill", json={"pid": os.getpid(), "signal": "BOGUS"}
|
||||
)
|
||||
check("invalid signal -> 400", r.status_code == 400)
|
||||
r = client.post("/api/processes/kill", json={"pid": -5, "signal": "TERM"})
|
||||
check("negative pid -> 400", r.status_code == 400)
|
||||
r = client.post("/api/processes/kill", json={"pid": 0, "signal": "TERM"})
|
||||
check("zero pid -> 400", r.status_code == 400)
|
||||
r = client.post("/api/processes/kill", json={"pid": 99999999, "signal": "TERM"})
|
||||
check("unknown pid -> 404", r.status_code == 404)
|
||||
|
||||
|
||||
def test_kill_endpoint_real() -> None:
|
||||
"""A real SIGTERM round-trip: the process must actually exit."""
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
client = TestClient(server.app)
|
||||
proc = subprocess.Popen(["sleep", "30"])
|
||||
try:
|
||||
r = client.post(
|
||||
"/api/processes/kill", json={"pid": proc.pid, "signal": "TERM"}
|
||||
)
|
||||
check("kill -> 200", r.status_code == 200)
|
||||
check("response ok", r.json().get("ok") is True)
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
check("process exited", True)
|
||||
except subprocess.TimeoutExpired:
|
||||
check("process exited", False)
|
||||
finally:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
|
||||
|
||||
def test_kill_endpoint_refuses_init_kthreadd() -> None:
|
||||
"""The endpoint must refuse to signal init (pid 1) or kthreadd (pid 2).
|
||||
|
||||
SIGTERM to init is a system-shutdown vector; kernel threads don't handle
|
||||
signals anyway. This fires before any NVML/state access, so no GPU needed.
|
||||
"""
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
client = TestClient(server.app)
|
||||
for pid in (1, 2):
|
||||
r = client.post(
|
||||
"/api/processes/kill", json={"pid": pid, "signal": "KILL"}
|
||||
)
|
||||
check(f"pid {pid} refused -> 400", r.status_code == 400)
|
||||
|
||||
|
||||
def test_get_parent_pid() -> None:
|
||||
"""get_parent_pid must return the kernel ppid and degrade gracefully."""
|
||||
check("self ppid matches", processes.get_parent_pid(os.getpid()) == os.getppid())
|
||||
check("missing pid -> None", processes.get_parent_pid(99999999) is None)
|
||||
check("pid 1 has no parent", processes.get_parent_pid(1) == 0)
|
||||
|
||||
|
||||
def test_kill_parent_endpoint() -> None:
|
||||
"""parent=true must signal the parent, leaving the child alive."""
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
client = TestClient(server.app)
|
||||
|
||||
# Validation cases first (no fixture needed).
|
||||
r = client.post(
|
||||
"/api/processes/kill", json={"pid": 99999999, "signal": "KILL", "parent": True}
|
||||
)
|
||||
check("parent of missing pid -> 404", r.status_code == 404)
|
||||
r = client.post(
|
||||
"/api/processes/kill", json={"pid": 1, "signal": "KILL", "parent": True}
|
||||
)
|
||||
check("pid 1 has no parent -> 400", r.status_code == 400)
|
||||
|
||||
# bash forks a background sleep (child), then sleeps itself (parent).
|
||||
parent = subprocess.Popen(["bash", "-c", "sleep 30 & sleep 30"])
|
||||
child: int | None = None
|
||||
for _ in range(100):
|
||||
try:
|
||||
with open(f"/proc/{parent.pid}/task/{parent.pid}/children") as f:
|
||||
kids = f.read().split()
|
||||
if kids:
|
||||
child = int(kids[0])
|
||||
break
|
||||
except OSError:
|
||||
pass
|
||||
time.sleep(0.05)
|
||||
try:
|
||||
check("fixture has a child", child is not None)
|
||||
if child is None:
|
||||
return
|
||||
r = client.post(
|
||||
"/api/processes/kill",
|
||||
json={"pid": child, "signal": "KILL", "parent": True},
|
||||
)
|
||||
check("kill parent -> 200", r.status_code == 200)
|
||||
check("response targets the parent", r.json().get("pid") == parent.pid)
|
||||
check("response reports parent flag", r.json().get("parent") is True)
|
||||
try:
|
||||
parent.wait(timeout=5)
|
||||
check("parent died", True)
|
||||
except subprocess.TimeoutExpired:
|
||||
check("parent died", False)
|
||||
time.sleep(0.2)
|
||||
try:
|
||||
os.kill(child, 0)
|
||||
check("child survived (orphaned)", True)
|
||||
except ProcessLookupError:
|
||||
check("child survived (orphaned)", False)
|
||||
finally:
|
||||
for victim in (parent,):
|
||||
with contextlib.suppress(ProcessLookupError, subprocess.TimeoutExpired):
|
||||
victim.kill()
|
||||
victim.wait(timeout=5)
|
||||
if child is not None:
|
||||
with contextlib.suppress(ProcessLookupError):
|
||||
os.kill(child, 9)
|
||||
|
||||
|
||||
def test_processes_endpoint() -> None:
|
||||
"""GET /api/processes must return the HAL payload for a known GPU."""
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
fake = {
|
||||
"processes": [
|
||||
{
|
||||
"pid": 1,
|
||||
"user": "root",
|
||||
"dev": 3,
|
||||
"type": "C",
|
||||
"gpu_util_pct": 42.0,
|
||||
"mem_util_pct": 1.0,
|
||||
"vram_bytes": 1000,
|
||||
"vram_pct": 0.1,
|
||||
"cpu_pct": 1.5,
|
||||
"mem_host_bytes": 2048,
|
||||
"command": "/usr/bin/fake",
|
||||
}
|
||||
],
|
||||
"mem_total_bytes": 1000000,
|
||||
}
|
||||
orig = server.list_gpu_processes
|
||||
server._state["gpus"][3] = {"gpu": object(), "gpu_name": "fake"}
|
||||
server.list_gpu_processes = lambda idx: fake
|
||||
try:
|
||||
client = TestClient(server.app)
|
||||
r = client.get("/api/processes", params={"gpu_index": 3})
|
||||
check("processes -> 200", r.status_code == 200)
|
||||
check("payload matches HAL", r.json() == fake)
|
||||
r = client.get("/api/processes", params={"gpu_index": 99})
|
||||
check("unknown gpu -> 404", r.status_code == 404)
|
||||
finally:
|
||||
server.list_gpu_processes = orig
|
||||
server._state["gpus"].pop(3, None)
|
||||
|
||||
|
||||
def test_list_gpu_processes_shape() -> None:
|
||||
"""With NVML available, list_gpu_processes must return a well-shaped payload."""
|
||||
from nvcurve.hal import monitoring
|
||||
|
||||
if not monitoring._NVML_AVAILABLE:
|
||||
print(" SKIP NVML not available")
|
||||
return
|
||||
if not monitoring.init_nvml():
|
||||
print(" SKIP NVML init failed (no driver?)")
|
||||
return
|
||||
try:
|
||||
data = processes.list_gpu_processes(0)
|
||||
check("has processes list", isinstance(data.get("processes"), list))
|
||||
check("has mem_total_bytes", "mem_total_bytes" in data)
|
||||
keys = {
|
||||
"pid",
|
||||
"user",
|
||||
"dev",
|
||||
"type",
|
||||
"gpu_util_pct",
|
||||
"mem_util_pct",
|
||||
"vram_bytes",
|
||||
"vram_pct",
|
||||
"cpu_pct",
|
||||
"mem_host_bytes",
|
||||
"command",
|
||||
}
|
||||
for p in data["processes"]:
|
||||
if not keys.issubset(p.keys()):
|
||||
check(f"process {p.get('pid')} has all keys", False)
|
||||
break
|
||||
else:
|
||||
check("all processes have all keys", True)
|
||||
# Sorted by VRAM descending.
|
||||
vram = [p["vram_bytes"] or 0 for p in data["processes"]]
|
||||
check("sorted by VRAM desc", vram == sorted(vram, reverse=True))
|
||||
finally:
|
||||
monitoring.shutdown_nvml()
|
||||
|
||||
|
||||
def test_list_gpu_processes_no_nvml() -> None:
|
||||
"""With NVML unavailable, list_gpu_processes returns an empty payload.
|
||||
|
||||
This is the branch the shape test skips (no driver): the UI uses
|
||||
mem_total_bytes is None to render 'data unavailable' instead of
|
||||
'no processes'.
|
||||
"""
|
||||
from nvcurve.hal import monitoring
|
||||
|
||||
orig_avail = monitoring._NVML_AVAILABLE
|
||||
orig_init = monitoring._nvml_initialized
|
||||
monitoring._NVML_AVAILABLE = False
|
||||
monitoring._nvml_initialized = False
|
||||
try:
|
||||
data = processes.list_gpu_processes(0)
|
||||
check("empty processes list", data["processes"] == [])
|
||||
check("mem_total_bytes is None", data["mem_total_bytes"] is None)
|
||||
finally:
|
||||
monitoring._NVML_AVAILABLE = orig_avail
|
||||
monitoring._nvml_initialized = orig_init
|
||||
|
||||
|
||||
def main() -> None:
|
||||
tests = [
|
||||
test_proc_info_self,
|
||||
test_proc_info_missing,
|
||||
test_proc_info_comm_with_parens,
|
||||
test_kill_endpoint_validation,
|
||||
test_kill_endpoint_real,
|
||||
test_kill_endpoint_refuses_init_kthreadd,
|
||||
test_get_parent_pid,
|
||||
test_kill_parent_endpoint,
|
||||
test_processes_endpoint,
|
||||
test_list_gpu_processes_shape,
|
||||
test_list_gpu_processes_no_nvml,
|
||||
]
|
||||
for t in tests:
|
||||
print(f"== {t.__name__} ==")
|
||||
t()
|
||||
print()
|
||||
print(f"{PASS} passed, {FAIL} failed")
|
||||
if FAIL:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user