e52c616592
The fixtures are the live numbers sampled from the box on 02-08-2026, where the CPT run held 12.8GB of 16 as proc/478104/vram_35881. The cases that matter are the ones where a mistake is silent: our own llama-server counting as a contender, an unreadable card reading as free, and /health hanging or proxying into a closed port instead of answering 503.
104 lines
3.1 KiB
Go
104 lines
3.1 KiB
Go
package main
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"testing"
|
|
)
|
|
|
|
// fakeKFD builds the sysfs shape the workstation actually has: one directory
|
|
// per ROCm process, each holding a vram_<node> file. Sampled from the live box
|
|
// on 02-08-2026, where the CPT run appeared as proc/478104/vram_35881.
|
|
func fakeKFD(t *testing.T, vramByPID map[int]int64) string {
|
|
t.Helper()
|
|
root := t.TempDir()
|
|
for pid, vram := range vramByPID {
|
|
dir := filepath.Join(root, strconv.Itoa(pid))
|
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
f := filepath.Join(dir, "vram_35881")
|
|
if err := os.WriteFile(f, []byte(strconv.FormatInt(vram, 10)+"\n"), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
return root
|
|
}
|
|
|
|
func TestForeignExcludesOurChild(t *testing.T) {
|
|
root := fakeKFD(t, map[int]int64{478104: 12791693312, 999: 4096})
|
|
p := probe{kfdRoot: root}
|
|
|
|
all := p.foreign(0)
|
|
if len(all) != 2 {
|
|
t.Fatalf("with no child running, both processes are foreign, got %d", len(all))
|
|
}
|
|
|
|
ours := p.foreign(999)
|
|
if len(ours) != 1 || ours[0].PID != 478104 {
|
|
t.Fatalf("our own llama-server must not count as a contender, got %+v", ours)
|
|
}
|
|
if ours[0].VRAM != 12791693312 {
|
|
t.Errorf("per-process VRAM = %d, want the value from vram_35881", ours[0].VRAM)
|
|
}
|
|
}
|
|
|
|
// An empty KFD tree is the state that permits a start, so it must read as empty
|
|
// rather than as an error the caller has to interpret.
|
|
func TestForeignEmptyAndMissing(t *testing.T) {
|
|
if got := (probe{kfdRoot: t.TempDir()}).foreign(0); len(got) != 0 {
|
|
t.Errorf("empty kfd tree: got %d processes, want 0", len(got))
|
|
}
|
|
if got := (probe{kfdRoot: "/nonexistent"}).foreign(0); got != nil {
|
|
t.Errorf("missing kfd tree: got %+v, want nil", got)
|
|
}
|
|
}
|
|
|
|
func TestFreeVRAM(t *testing.T) {
|
|
dev := t.TempDir()
|
|
write := func(name, v string) {
|
|
if err := os.WriteFile(filepath.Join(dev, name), []byte(v), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
// The live numbers from the workstation while the CPT run held the card.
|
|
write("mem_info_vram_total", "17163091968\n")
|
|
write("mem_info_vram_used", "13396389888\n")
|
|
p := probe{drmDev: dev}
|
|
if got, want := p.freeVRAM(), int64(3766702080); got != want {
|
|
t.Errorf("freeVRAM = %d, want %d", got, want)
|
|
}
|
|
if got := (probe{drmDev: "/nonexistent"}).freeVRAM(); got != 0 {
|
|
t.Errorf("unreadable card reports %d free, want 0 so nothing starts", got)
|
|
}
|
|
}
|
|
|
|
// With no model loaded the supervisor must still answer, and it must answer 503
|
|
// rather than hanging or proxying into a closed port. Maven reads this endpoint
|
|
// on a timer forever, including while the workstation is busy.
|
|
func TestHealthAndProxyRefuseWhenNotReady(t *testing.T) {
|
|
s := &supervisor{run: newRunner("/bin/true", nil, "")}
|
|
h := s.handler(mustURL(t, "http://127.0.0.1:1"))
|
|
|
|
for _, path := range []string{"/health", "/v1/chat/completions"} {
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, path, nil))
|
|
if w.Code != http.StatusServiceUnavailable {
|
|
t.Errorf("%s with no model: got %d, want 503", path, w.Code)
|
|
}
|
|
}
|
|
}
|
|
|
|
func mustURL(t *testing.T, s string) *url.URL {
|
|
t.Helper()
|
|
u, err := url.Parse(s)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return u
|
|
}
|