gpulock: GPU arbitration proxy for Ollama + ComfyUI

Implements SPEC.md: two listeners, one writer-preferring two-mode lock, Ollama unload before image jobs, ComfyUI history polling + VRAM free, optional model warm-up, healthz/metrics endpoints, streaming-safe reverse proxies, Dockerfile and Gitea Actions CI.
2026-09-20 18:05:20 +02:00
commit 065c294a96
14 changed files with 2027 additions and 0 deletions
@@ -0,0 +1,157 @@
+// Package lock implements the two-mode GPU arbitration lock: any number of
+// concurrent LLM requests ("readers") or exactly one image job ("writer"),
+// with image jobs taking priority over newly arriving LLM requests.
+package lock
+
+import (
+	"context"
+	"log/slog"
+	"sync"
+)
+
+// State is the current GPU occupancy state.
+type State string
+
+const (
+	StateIdle  State = "idle"
+	StateLLM   State = "llm"
+	StateImage State = "image"
+)
+
+type imageWaiter struct{ id uint64 }
+
+// Lock is a writer-preferring two-mode lock. The zero value is not usable;
+// construct with New.
+type Lock struct {
+	mu     sync.Mutex
+	change chan struct{} // closed and replaced on every state change
+
+	n           int  // LLM requests in flight
+	imageActive bool // an image job holds the GPU
+	imageQ      []imageWaiter
+	nextID      uint64
+
+	log *slog.Logger
+}
+
+// New returns a ready-to-use Lock. log may be nil; if set, every state
+// transition is logged at debug level.
+func New(log *slog.Logger) *Lock {
+	return &Lock{change: make(chan struct{}), log: log}
+}
+
+// broadcast wakes all waiters. Call with mu held.
+func (l *Lock) broadcast() {
+	close(l.change)
+	l.change = make(chan struct{})
+}
+
+func (l *Lock) logTransition(msg string, args ...any) {
+	if l.log != nil {
+		l.log.Debug(msg, args...)
+	}
+}
+
+// AcquireLLM blocks until no image job is active or pending, then registers
+// one in-flight LLM request. Returns ctx.Err() if the context is cancelled
+// while waiting; no state is changed in that case.
+func (l *Lock) AcquireLLM(ctx context.Context) error {
+	l.mu.Lock()
+	for l.imageActive || len(l.imageQ) > 0 {
+		ch := l.change
+		l.mu.Unlock()
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ch:
+		}
+		l.mu.Lock()
+	}
+	l.n++
+	n := l.n
+	l.mu.Unlock()
+	l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
+	return nil
+}
+
+// ReleaseLLM marks one LLM request as finished.
+func (l *Lock) ReleaseLLM() {
+	l.mu.Lock()
+	l.n--
+	n := l.n
+	if l.n == 0 {
+		l.broadcast()
+	}
+	l.mu.Unlock()
+	if n == 0 {
+		l.logTransition("lock transition", "state", StateIdle, "llm_inflight", 0)
+	}
+}
+
+// AcquireImage queues the caller FIFO behind other image jobs, blocks new
+// LLM requests immediately, and waits for in-flight LLM requests to drain
+// before granting exclusive GPU access. Returns ctx.Err() if the context is
+// cancelled while waiting; the caller is removed from the queue.
+func (l *Lock) AcquireImage(ctx context.Context) error {
+	l.mu.Lock()
+	l.nextID++
+	w := imageWaiter{id: l.nextID}
+	l.imageQ = append(l.imageQ, w)
+	pending := len(l.imageQ)
+	l.broadcast() // new LLM requests must block from now on
+	l.mu.Unlock()
+	l.logTransition("image job pending", "image_queue", pending)
+
+	for {
+		l.mu.Lock()
+		if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive {
+			l.imageQ = l.imageQ[1:]
+			l.imageActive = true
+			l.mu.Unlock()
+			l.logTransition("lock transition", "state", StateImage)
+			return nil
+		}
+		ch := l.change
+		l.mu.Unlock()
+		select {
+		case <-ctx.Done():
+			l.mu.Lock()
+			for i, q := range l.imageQ {
+				if q.id == w.id {
+					l.imageQ = append(l.imageQ[:i], l.imageQ[i+1:]...)
+					break
+				}
+			}
+			l.broadcast()
+			l.mu.Unlock()
+			l.logTransition("image job wait cancelled", "image_queue", len(l.imageQ))
+			return ctx.Err()
+		case <-ch:
+		}
+	}
+}
+
+// ReleaseImage frees the GPU after an image job.
+func (l *Lock) ReleaseImage() {
+	l.mu.Lock()
+	l.imageActive = false
+	l.broadcast()
+	l.mu.Unlock()
+	l.logTransition("lock transition", "state", StateIdle)
+}
+
+// Snapshot reports the current state, the number of in-flight LLM requests,
+// and whether an image job is active or waiting.
+func (l *Lock) Snapshot() (state State, llmInflight int, imagePending bool) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	switch {
+	case l.imageActive:
+		state = StateImage
+	case l.n > 0:
+		state = StateLLM
+	default:
+		state = StateIdle
+	}
+	return state, l.n, l.imageActive || len(l.imageQ) > 0
+}