Game detection: foreign GPU holders take an external lock hold (GAME_PROCS, GPU_FOREIGN_VRAM_MB)
This commit is contained in:
@@ -0,0 +1,159 @@
|
||||
// Package game detects processes outside gpu-turnstile's control that hold
|
||||
// the GPU — typically a game — so the proxy can block new GPU work and free
|
||||
// VRAM while they run. Two detection paths: an explicit process watch list
|
||||
// (GAME_PROCS) and a foreign-VRAM threshold via nvidia-smi
|
||||
// (GPU_FOREIGN_VRAM_MB) that catches anything not on the ignore list.
|
||||
package game
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Process is one running OS process.
|
||||
type Process struct {
|
||||
PID int
|
||||
Name string
|
||||
}
|
||||
|
||||
// computeApp is one process holding GPU memory, as reported by nvidia-smi.
|
||||
type computeApp struct {
|
||||
PID int
|
||||
UsedMB int
|
||||
}
|
||||
|
||||
// Detector checks whether a foreign process holds the GPU. The zero value
|
||||
// (no watch list, no threshold) never detects anything; main only starts the
|
||||
// poll loop when at least one path is configured.
|
||||
type Detector struct {
|
||||
procs map[string]bool // normalized names from GAME_PROCS
|
||||
vramMB int // foreign VRAM threshold; 0 = disabled
|
||||
ignore map[string]bool // normalized names never counted as foreign
|
||||
log *slog.Logger
|
||||
noNvidia bool // nvidia-smi was not found; VRAM path disabled for good
|
||||
}
|
||||
|
||||
// New builds a Detector from the configured watch list, VRAM threshold in
|
||||
// MiB (0 disables the nvidia-smi path) and ignore list. Names are matched
|
||||
// case-insensitively, with or without a trailing ".exe".
|
||||
func New(procs []string, vramMB int, ignore []string, log *slog.Logger) *Detector {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Detector{
|
||||
procs: nameSet(procs),
|
||||
vramMB: vramMB,
|
||||
ignore: nameSet(ignore),
|
||||
log: log,
|
||||
}
|
||||
}
|
||||
|
||||
// normName lowercases a process name and strips a trailing ".exe" so the
|
||||
// watch and ignore lists match on Windows and Linux spellings alike.
|
||||
func normName(s string) string {
|
||||
return strings.TrimSuffix(strings.ToLower(strings.TrimSpace(s)), ".exe")
|
||||
}
|
||||
|
||||
func nameSet(names []string) map[string]bool {
|
||||
set := make(map[string]bool, len(names))
|
||||
for _, n := range names {
|
||||
if n = normName(n); n != "" {
|
||||
set[n] = true
|
||||
}
|
||||
}
|
||||
return set
|
||||
}
|
||||
|
||||
// Check looks once for foreign GPU holders and returns a human-readable
|
||||
// description of each (empty when the GPU is free for gpu-turnstile's
|
||||
// consumers). A failing nvidia-smi call is returned as an error only when
|
||||
// the process list found nothing; a missing nvidia-smi binary disables the
|
||||
// VRAM path permanently (logged once).
|
||||
func (d *Detector) Check(ctx context.Context) ([]string, error) {
|
||||
ps, psErr := processes()
|
||||
if d.vramMB <= 0 || d.noNvidia {
|
||||
return d.detect(ps, nil), psErr
|
||||
}
|
||||
apps, err := queryComputeApps(ctx)
|
||||
if errors.Is(err, exec.ErrNotFound) {
|
||||
d.noNvidia = true
|
||||
d.log.Warn("GPU_FOREIGN_VRAM_MB is set but nvidia-smi was not found; VRAM detection disabled")
|
||||
return d.detect(ps, nil), nil
|
||||
}
|
||||
if err != nil {
|
||||
return d.detect(ps, nil), err
|
||||
}
|
||||
return d.detect(ps, apps), nil
|
||||
}
|
||||
|
||||
// detect is the pure core of Check: given the process table and (optionally)
|
||||
// the nvidia-smi compute-apps list, it returns the foreign holders.
|
||||
func (d *Detector) detect(ps []Process, apps []computeApp) []string {
|
||||
var holders []string
|
||||
for _, p := range ps {
|
||||
if d.procs[normName(p.Name)] {
|
||||
holders = append(holders, fmt.Sprintf("%s (pid %d)", p.Name, p.PID))
|
||||
}
|
||||
}
|
||||
if d.vramMB > 0 && apps != nil {
|
||||
names := make(map[int]string, len(ps))
|
||||
for _, p := range ps {
|
||||
names[p.PID] = p.Name
|
||||
}
|
||||
for _, a := range apps {
|
||||
name := names[a.PID]
|
||||
if d.ignore[normName(name)] || a.UsedMB < d.vramMB {
|
||||
continue
|
||||
}
|
||||
if name == "" {
|
||||
name = "unknown process"
|
||||
}
|
||||
holders = append(holders, fmt.Sprintf("%s (pid %d) using %d MiB VRAM", name, a.PID, a.UsedMB))
|
||||
}
|
||||
}
|
||||
return holders
|
||||
}
|
||||
|
||||
// queryComputeApps runs nvidia-smi and parses the per-process VRAM list.
|
||||
// Note: under Windows' WDDM driver, nvidia-smi only sees compute
|
||||
// allocations, so graphics-only games may not appear there — GAME_PROCS is
|
||||
// the reliable path on Windows; on Linux both work.
|
||||
func queryComputeApps(ctx context.Context) ([]computeApp, error) {
|
||||
out, err := exec.CommandContext(ctx, "nvidia-smi",
|
||||
"--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits").Output()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return parseComputeApps(string(out))
|
||||
}
|
||||
|
||||
// parseComputeApps parses "pid, used_memory" CSV lines (no header, MiB
|
||||
// units). Unsupported rows ("N/A" on WDDM) are skipped.
|
||||
func parseComputeApps(out string) ([]computeApp, error) {
|
||||
var apps []computeApp
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
pidStr, memStr, ok := strings.Cut(line, ",")
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("nvidia-smi: unexpected line %q", line)
|
||||
}
|
||||
pid, err := strconv.Atoi(strings.TrimSpace(pidStr))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("nvidia-smi: unexpected pid in %q", line)
|
||||
}
|
||||
mem, err := strconv.Atoi(strings.TrimSpace(memStr))
|
||||
if err != nil {
|
||||
continue // "N/A" and friends: unsupported under WDDM
|
||||
}
|
||||
apps = append(apps, computeApp{PID: pid, UsedMB: mem})
|
||||
}
|
||||
return apps, nil
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package game
|
||||
|
||||
import (
|
||||
"runtime"
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNormName(t *testing.T) {
|
||||
for in, want := range map[string]string{
|
||||
"Cyberpunk2077.exe": "cyberpunk2077",
|
||||
"ollama": "ollama",
|
||||
"OLLAMA APP.EXE": "ollama app",
|
||||
" python.exe ": "python",
|
||||
"hl2": "hl2",
|
||||
} {
|
||||
if got := normName(in); got != want {
|
||||
t.Errorf("normName(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseComputeApps(t *testing.T) {
|
||||
apps, err := parseComputeApps("1234, 512\n 42 , 8192 \n\n")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := []computeApp{{PID: 1234, UsedMB: 512}, {PID: 42, UsedMB: 8192}}
|
||||
if !slices.Equal(apps, want) {
|
||||
t.Errorf("got %+v, want %+v", apps, want)
|
||||
}
|
||||
|
||||
// WDDM reports "N/A" for memory; those rows are skipped, not fatal.
|
||||
apps, err = parseComputeApps("1234, N/A\n42, 1024\n")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !slices.Equal(apps, []computeApp{{PID: 42, UsedMB: 1024}}) {
|
||||
t.Errorf("got %+v", apps)
|
||||
}
|
||||
|
||||
if _, err := parseComputeApps("garbage\n"); err == nil {
|
||||
t.Error("expected an error for a malformed line")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetect(t *testing.T) {
|
||||
d := New([]string{"Cyberpunk2077.exe", "hl2"}, 1024,
|
||||
[]string{"ollama", "python", "pythonw"}, nil)
|
||||
ps := []Process{
|
||||
{PID: 10, Name: "ollama.exe"},
|
||||
{PID: 20, Name: "python.exe"},
|
||||
{PID: 30, Name: "cyberpunk2077.exe"},
|
||||
}
|
||||
apps := []computeApp{
|
||||
{PID: 10, UsedMB: 8192}, // ignored: ollama
|
||||
{PID: 20, UsedMB: 4096}, // ignored: python (ComfyUI)
|
||||
{PID: 40, UsedMB: 2048}, // foreign, above threshold
|
||||
{PID: 50, UsedMB: 100}, // foreign but below threshold
|
||||
}
|
||||
holders := d.detect(ps, apps)
|
||||
if len(holders) != 2 {
|
||||
t.Fatalf("got %v, want 2 holders", holders)
|
||||
}
|
||||
if holders[0] != "cyberpunk2077.exe (pid 30)" {
|
||||
t.Errorf("holders[0] = %q", holders[0])
|
||||
}
|
||||
if holders[1] != "unknown process (pid 40) using 2048 MiB VRAM" {
|
||||
t.Errorf("holders[1] = %q", holders[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectNothingConfigured(t *testing.T) {
|
||||
d := New(nil, 0, nil, nil)
|
||||
if got := d.detect([]Process{{PID: 1, Name: "game.exe"}}, nil); len(got) != 0 {
|
||||
t.Errorf("got %v, want none", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessesLive(t *testing.T) {
|
||||
if runtime.GOOS != "windows" && runtime.GOOS != "linux" {
|
||||
t.Skip("no process listing on this platform")
|
||||
}
|
||||
ps, err := processes()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(ps) == 0 {
|
||||
t.Fatal("no processes listed")
|
||||
}
|
||||
for _, p := range ps {
|
||||
if p.Name == "" {
|
||||
t.Errorf("pid %d has an empty name", p.PID)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
//go:build linux
|
||||
|
||||
package game
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// processes lists the running processes from /proc/<pid>/comm.
|
||||
func processes() ([]Process, error) {
|
||||
entries, err := os.ReadDir("/proc")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var ps []Process
|
||||
for _, e := range entries {
|
||||
pid, err := strconv.Atoi(e.Name())
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
comm, err := os.ReadFile("/proc/" + e.Name() + "/comm")
|
||||
if err != nil {
|
||||
continue // process vanished mid-walk
|
||||
}
|
||||
ps = append(ps, Process{PID: pid, Name: strings.TrimSpace(string(comm))})
|
||||
}
|
||||
return ps, nil
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//go:build !windows && !linux
|
||||
|
||||
package game
|
||||
|
||||
// processes is unsupported on this platform; the process watch list never
|
||||
// matches and the nvidia-smi path reports PIDs without names.
|
||||
func processes() ([]Process, error) { return nil, nil }
|
||||
@@ -0,0 +1,35 @@
|
||||
//go:build windows
|
||||
|
||||
package game
|
||||
|
||||
import (
|
||||
"unsafe"
|
||||
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
// processes lists the running processes via the Toolhelp32 snapshot API.
|
||||
func processes() ([]Process, error) {
|
||||
h, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer windows.CloseHandle(h) //nolint:errcheck // best effort
|
||||
|
||||
var entry windows.ProcessEntry32
|
||||
entry.Size = uint32(unsafe.Sizeof(entry))
|
||||
if err := windows.Process32First(h, &entry); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var ps []Process
|
||||
for {
|
||||
ps = append(ps, Process{
|
||||
PID: int(entry.ProcessID),
|
||||
Name: windows.UTF16ToString(entry.ExeFile[:]),
|
||||
})
|
||||
if err := windows.Process32Next(h, &entry); err != nil {
|
||||
break // ERROR_NO_MORE_FILES ends the walk
|
||||
}
|
||||
}
|
||||
return ps, nil
|
||||
}
|
||||
Reference in New Issue
Block a user