ARIA
c11ea73ead
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
347 lines
11 KiB
Python
347 lines
11 KiB
Python
"""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()
|