diff --git a/cmd/mavgpud/gpu.go b/cmd/mavgpud/gpu.go new file mode 100644 index 0000000..89120c6 --- /dev/null +++ b/cmd/mavgpud/gpu.go @@ -0,0 +1,117 @@ +package main + +import ( + "os" + "path/filepath" + "strconv" + "strings" +) + +// The card is an AMD 7900 GRE with 16GB, driven by amdgpu and ROCm. Everything +// here reads sysfs and forks nothing: rocm-smi is not even installed on the +// workstation, and a poll that costs a subprocess every second is a poll that +// gets tuned down until it is useless. + +// gpuProc — one process holding the compute engine. +type gpuProc struct { + PID int + Comm string + VRAM int64 // bytes, as the kernel accounts them to this process +} + +// probe reads the two sysfs trees the supervisor decides from. +// +// kfdRoot is /sys/class/kfd/kfd/proc, one directory per ROCm process. The +// directory appears when the process initialises HIP, which is well before it +// allocates anything large. That is the whole reason this works: the job that +// is about to want the card announces itself while it is still starting up, +// so we see the contender rather than only the winner of an allocation race. +// +// drmDev is /sys/class/drm/cardN/device, which reports total and used VRAM for +// the card as a whole. +type probe struct { + kfdRoot string + drmDev string +} + +// foreign lists every ROCm process that is not ours. selfPID is the supervisor's +// llama-server child, or 0 when it is not running. +// +// An unreadable kfd tree returns no processes and no error. That is deliberate +// and it is the safe direction only because startVRAM also has to agree before +// anything launches: a supervisor that cannot see the KFD never sees free VRAM +// either, because the CPT run holding the card shows up in the drm totals. +func (p probe) foreign(selfPID int) []gpuProc { + entries, err := os.ReadDir(p.kfdRoot) + if err != nil { + return nil + } + var out []gpuProc + for _, e := range entries { + pid, err := strconv.Atoi(e.Name()) + if err != nil || pid == selfPID { + continue + } + out = append(out, gpuProc{ + PID: pid, + Comm: readComm(pid), + VRAM: p.procVRAM(e.Name()), + }) + } + return out +} + +// procVRAM sums the per-node vram_* files under one process directory. The +// suffix is the KFD topology node id (vram_35881 on this card), so it is +// globbed rather than named, and a machine with two cards sums both. +func (p probe) procVRAM(pid string) int64 { + matches, err := filepath.Glob(filepath.Join(p.kfdRoot, pid, "vram_*")) + if err != nil { + return 0 + } + var total int64 + for _, m := range matches { + total += readInt(m) + } + return total +} + +// freeVRAM reports the bytes the card has left. Used only to decide whether to +// start: a shortfall here means llama-server would refuse to load anyway. It is +// never used to decide to stop, because by the time free VRAM has dropped the +// other job has already failed its allocation, which is exactly the outcome +// yielding exists to prevent. +func (p probe) freeVRAM() int64 { + total := readInt(filepath.Join(p.drmDev, "mem_info_vram_total")) + used := readInt(filepath.Join(p.drmDev, "mem_info_vram_used")) + if total <= 0 { + return 0 + } + if free := total - used; free > 0 { + return free + } + return 0 +} + +func readInt(path string) int64 { + b, err := os.ReadFile(path) + if err != nil { + return 0 + } + n, err := strconv.ParseInt(strings.TrimSpace(string(b)), 10, 64) + if err != nil { + return 0 + } + return n +} + +// readComm names the contender for the log. The log is the instrument for the +// open question in Vikunja #488: whether a process can want this card without +// ever registering on the KFD, which a Vulkan or video-decode job would. +func readComm(pid int) string { + b, err := os.ReadFile(filepath.Join("/proc", strconv.Itoa(pid), "comm")) + if err != nil { + return "?" + } + return strings.TrimSpace(string(b)) +} diff --git a/cmd/mavgpud/gpu_test.go b/cmd/mavgpud/gpu_test.go new file mode 100644 index 0000000..073874b --- /dev/null +++ b/cmd/mavgpud/gpu_test.go @@ -0,0 +1,103 @@ +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_ 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 +} diff --git a/cmd/mavgpud/main.go b/cmd/mavgpud/main.go new file mode 100644 index 0000000..04c6123 --- /dev/null +++ b/cmd/mavgpud/main.go @@ -0,0 +1,247 @@ +// mavgpud — the workstation's GPU supervisor. +// +// It runs on the workstation (an AMD 7900 GRE, 16GB), not on homesrv, and it is +// deployed separately from the Maven daemons. Maven does not participate in any +// of this and never asks for a start: it reads /health through internal/llm.Pair +// and either gets the big model or falls back to the resident 1.7B. +// +// The rule, from Vikunja #488: keep llama-server loaded whenever the card is +// free, unload it when it has been idle too long or when another process needs +// the card. Not on demand, because a 7-14B takes tens of seconds to load and a +// world question would be answered by a gap every time the card had been quiet. +// Not always on, because that holds 16GB against the owner's own jobs. +package main + +import ( + "context" + "encoding/json" + "flag" + "log" + "net/http" + "net/http/httputil" + "net/url" + "os" + "os/signal" + "sync/atomic" + "syscall" + "time" +) + +type config struct { + Listen string `json:"listen"` // what Maven talks to + LlamaAddr string `json:"llama_addr"` // where llama-server binds + LlamaBin string `json:"llama_bin"` + // LlamaArgs must include the flags that bind LlamaAddr. They are passed + // through untouched so the model, context size and layer count stay the + // owner's business and not this daemon's schema. + LlamaArgs []string `json:"llama_args"` + + KFDRoot string `json:"kfd_root"` + DRMDevice string `json:"drm_device"` + + Poll duration `json:"poll"` + IdleTimeout duration `json:"idle_timeout"` + StopGrace duration `json:"stop_grace"` + MinFreeVRAM int64 `json:"min_free_vram_bytes"` + // EvictAfter and StartAfter are counted in polls, not seconds. Both exist + // to damp flapping: a one-tick blip from a short-lived rocm process must + // not evict the model, and a card that has just been released must not be + // grabbed before the previous job has finished unmapping. + EvictAfter int `json:"evict_after_polls"` + StartAfter int `json:"start_after_polls"` +} + +func defaults() config { + return config{ + Listen: ":8080", + LlamaAddr: "127.0.0.1:8081", + KFDRoot: "/sys/class/kfd/kfd/proc", + DRMDevice: "/sys/class/drm/card1/device", + Poll: duration(time.Second), + IdleTimeout: duration(15 * time.Minute), + StopGrace: duration(20 * time.Second), + MinFreeVRAM: 15 << 30, + EvictAfter: 2, + StartAfter: 5, + } +} + +// duration lets the config file say "15m" instead of counting nanoseconds. +type duration time.Duration + +func (d *duration) UnmarshalJSON(b []byte) error { + var s string + if err := json.Unmarshal(b, &s); err != nil { + return err + } + v, err := time.ParseDuration(s) + if err != nil { + return err + } + *d = duration(v) + return nil +} + +func main() { + path := flag.String("config", "/etc/mavgpud.json", "config file") + flag.Parse() + + cfg := defaults() + b, err := os.ReadFile(*path) + if err != nil { + log.Fatalf("mavgpud: read config: %v", err) + } + if err := json.Unmarshal(b, &cfg); err != nil { + log.Fatalf("mavgpud: parse config: %v", err) + } + if cfg.LlamaBin == "" { + log.Fatal("mavgpud: llama_bin is required") + } + + base := "http://" + cfg.LlamaAddr + run := newRunner(cfg.LlamaBin, cfg.LlamaArgs, base+"/health") + sup := &supervisor{ + cfg: cfg, + probe: probe{kfdRoot: cfg.KFDRoot, drmDev: cfg.DRMDevice}, + run: run, + } + sup.touch() + + ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer cancel() + + target, err := url.Parse(base) + if err != nil { + log.Fatalf("mavgpud: llama_addr: %v", err) + } + srv := &http.Server{Addr: cfg.Listen, Handler: sup.handler(target)} + go func() { + log.Printf("mavgpud: listening on %s, model %s", cfg.Listen, cfg.LlamaBin) + if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed { + log.Fatalf("mavgpud: listen: %v", err) + } + }() + + sup.loop(ctx) + + // The card must come back before we do. A supervisor that exits leaving + // llama-server holding 14GB is worse than one that never ran. + shut, done := context.WithTimeout(context.Background(), 5*time.Second) + defer done() + _ = srv.Shutdown(shut) + run.stop(time.Duration(cfg.StopGrace)) +} + +type supervisor struct { + cfg config + probe probe + run *runner + + lastReq atomic.Int64 // unix nanos of the last request Maven sent + + foreignStreak int + clearStreak int +} + +func (s *supervisor) touch() { s.lastReq.Store(time.Now().UnixNano()) } + +func (s *supervisor) idle() time.Duration { + return time.Since(time.Unix(0, s.lastReq.Load())) +} + +// handler serves the two things the workstation exposes. +// +// /health is answered locally and always, with no GPU cost and no round trip, +// because it is the only thing Maven reads and Maven reads it on a timer +// forever. Everything else is llama-server's API, reverse-proxied. Proxying +// rather than pointing Maven straight at llama-server is what makes the idle +// window measurable: the supervisor cannot otherwise know when the model was +// last used. +func (s *supervisor) handler(target *url.URL) http.Handler { + proxy := httputil.NewSingleHostReverseProxy(target) + mux := http.NewServeMux() + mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { + if !s.run.isReady() { + http.Error(w, "model not loaded", http.StatusServiceUnavailable) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"status":"ok"}`)) + }) + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + if !s.run.isReady() { + http.Error(w, "model not loaded", http.StatusServiceUnavailable) + return + } + s.touch() + proxy.ServeHTTP(w, r) + }) + return mux +} + +func (s *supervisor) loop(ctx context.Context) { + t := time.NewTicker(time.Duration(s.cfg.Poll)) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + s.tick(ctx) + } + } +} + +// tick is the whole decision. Yielding is checked before starting, and presence +// on the KFD is what triggers it — not a VRAM threshold. A ROCm process +// registers under /sys/class/kfd/kfd/proc when it initialises HIP, before it +// allocates, so we see a contender during its startup rather than after it has +// already failed to get the memory it wanted. +func (s *supervisor) tick(ctx context.Context) { + others := s.probe.foreign(s.run.pid()) + if len(others) > 0 { + s.foreignStreak++ + s.clearStreak = 0 + } else { + s.foreignStreak = 0 + s.clearStreak++ + } + + if s.run.running() { + s.run.refreshReady(ctx) + switch { + case s.foreignStreak >= s.cfg.EvictAfter: + log.Printf("mavgpud: yielding the card to %s", describe(others)) + s.run.stop(time.Duration(s.cfg.StopGrace)) + case s.idle() > time.Duration(s.cfg.IdleTimeout): + log.Printf("mavgpud: idle for %s, unloading", s.idle().Round(time.Second)) + s.run.stop(time.Duration(s.cfg.StopGrace)) + } + return + } + + if s.clearStreak < s.cfg.StartAfter { + return + } + if free := s.probe.freeVRAM(); free < s.cfg.MinFreeVRAM { + return + } + s.touch() // the idle clock starts at load, not at the last request before it + if err := s.run.start(); err != nil { + log.Printf("mavgpud: start llama-server: %v", err) + } +} + +// describe names the contenders in the log. This log is the instrument for the +// open question in #488: whether polling the KFD misses a job that wants the +// card without registering there. +func describe(procs []gpuProc) string { + out := "" + for i, p := range procs { + if i > 0 { + out += ", " + } + out += p.Comm + } + return out +} diff --git a/cmd/mavgpud/runner.go b/cmd/mavgpud/runner.go new file mode 100644 index 0000000..39ce341 --- /dev/null +++ b/cmd/mavgpud/runner.go @@ -0,0 +1,132 @@ +package main + +import ( + "context" + "log" + "net/http" + "os/exec" + "sync" + "syscall" + "time" +) + +// runner owns one llama-server process. Owning it is the point of the daemon: +// the workstation cannot keep a 7-14B resident, because that holds 16GB against +// the owner's CPT runs, Correx and the manga-recap pipeline. So the thing that +// stays up is this, which costs no VRAM, and the model comes and goes under it. +type runner struct { + bin string + args []string + // ready is llama-server's own /health, which answers "is a model loaded". + // Loading a 7-14B takes tens of seconds, so started is not ready. + readyURL string + + mu sync.Mutex + cmd *exec.Cmd + ready bool + http *http.Client +} + +func newRunner(bin string, args []string, readyURL string) *runner { + return &runner{ + bin: bin, args: args, readyURL: readyURL, + http: &http.Client{Timeout: 2 * time.Second}, + } +} + +// pid is the child's, or 0. The GPU probe needs it to tell our own model apart +// from a contender. +func (r *runner) pid() int { + r.mu.Lock() + defer r.mu.Unlock() + if r.cmd == nil || r.cmd.Process == nil { + return 0 + } + return r.cmd.Process.Pid +} + +func (r *runner) running() bool { return r.pid() != 0 } + +// isReady reports the cached readiness. The supervisor loop refreshes it; the +// health handler only reads, so answering /health never costs a round trip. +func (r *runner) isReady() bool { + r.mu.Lock() + defer r.mu.Unlock() + return r.ready +} + +// start launches llama-server. It returns as soon as the process exists, not +// when the model is loaded. +func (r *runner) start() error { + r.mu.Lock() + defer r.mu.Unlock() + if r.cmd != nil { + return nil + } + cmd := exec.Command(r.bin, r.args...) + // Own process group, so stop kills anything llama-server spawned rather + // than leaving it holding VRAM after we have declared the card yielded. + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + if err := cmd.Start(); err != nil { + return err + } + r.cmd, r.ready = cmd, false + log.Printf("mavgpud: started llama-server pid=%d", cmd.Process.Pid) + go func() { + err := cmd.Wait() + r.mu.Lock() + r.cmd, r.ready = nil, false + r.mu.Unlock() + log.Printf("mavgpud: llama-server exited: %v", err) + }() + return nil +} + +// stop ends llama-server and waits for the VRAM to come back. SIGTERM first so +// it unmaps cleanly, SIGKILL after the grace window. Returning before the +// process is gone would let the supervisor report a free card while 14GB is +// still mapped, which is the one lie that would make yielding useless. +func (r *runner) stop(grace time.Duration) { + r.mu.Lock() + cmd := r.cmd + r.ready = false + r.mu.Unlock() + if cmd == nil || cmd.Process == nil { + return + } + pgid := -cmd.Process.Pid + _ = syscall.Kill(pgid, syscall.SIGTERM) + deadline := time.Now().Add(grace) + for time.Now().Before(deadline) { + if !r.running() { + return + } + time.Sleep(100 * time.Millisecond) + } + log.Printf("mavgpud: llama-server did not exit in %s, killing", grace) + _ = syscall.Kill(pgid, syscall.SIGKILL) +} + +// refreshReady asks llama-server whether the model is loaded. Called once per +// supervisor tick, never per request. +func (r *runner) refreshReady(ctx context.Context) { + if !r.running() { + return + } + ok := false + req, err := http.NewRequestWithContext(ctx, http.MethodGet, r.readyURL, nil) + if err == nil { + resp, err := r.http.Do(req) + if err == nil { + ok = resp.StatusCode == http.StatusOK + resp.Body.Close() + } + } + r.mu.Lock() + was := r.ready + r.ready = ok + r.mu.Unlock() + if ok && !was { + log.Printf("mavgpud: model ready") + } +}