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.
This commit is contained in:
mram
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,228 @@
|
||||
package lock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func waitFor(t *testing.T, ch <-chan struct{}, what string) {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-ch:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatalf("timed out waiting for %s", what)
|
||||
}
|
||||
}
|
||||
|
||||
func assertBlocked(t *testing.T, ch <-chan struct{}, what string) {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-ch:
|
||||
t.Fatalf("%s should still be blocked", what)
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
|
||||
func TestLLMConcurrency(t *testing.T) {
|
||||
lk := New(nil)
|
||||
ctx := context.Background()
|
||||
for i := 0; i < 3; i++ {
|
||||
if err := lk.AcquireLLM(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
state, n, pending := lk.Snapshot()
|
||||
if state != StateLLM || n != 3 || pending {
|
||||
t.Fatalf("got state=%s n=%d pending=%v", state, n, pending)
|
||||
}
|
||||
lk.ReleaseLLM()
|
||||
lk.ReleaseLLM()
|
||||
if state, _, _ := lk.Snapshot(); state != StateLLM {
|
||||
t.Fatalf("state=%s, want llm", state)
|
||||
}
|
||||
lk.ReleaseLLM()
|
||||
if state, n, _ := lk.Snapshot(); state != StateIdle || n != 0 {
|
||||
t.Fatalf("got state=%s n=%d, want idle", state, n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageWaitsForLLMDrainAndBlocksNewLLM(t *testing.T) {
|
||||
lk := New(nil)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := lk.AcquireLLM(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := lk.AcquireLLM(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
imageAcquired := make(chan struct{})
|
||||
go func() {
|
||||
if err := lk.AcquireImage(ctx); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
close(imageAcquired)
|
||||
}()
|
||||
assertBlocked(t, imageAcquired, "image acquire while LLMs in flight")
|
||||
|
||||
// While an image job is pending, new LLM requests must block even
|
||||
// though LLM concurrency is otherwise allowed.
|
||||
llmAcquired := make(chan struct{})
|
||||
go func() {
|
||||
if err := lk.AcquireLLM(ctx); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
close(llmAcquired)
|
||||
}()
|
||||
assertBlocked(t, llmAcquired, "LLM acquire while image pending")
|
||||
|
||||
lk.ReleaseLLM()
|
||||
assertBlocked(t, imageAcquired, "image acquire with one LLM still in flight")
|
||||
|
||||
lk.ReleaseLLM()
|
||||
waitFor(t, imageAcquired, "image acquire after drain")
|
||||
assertBlocked(t, llmAcquired, "LLM acquire while image active")
|
||||
|
||||
state, _, pending := lk.Snapshot()
|
||||
if state != StateImage || !pending {
|
||||
t.Fatalf("got state=%s pending=%v", state, pending)
|
||||
}
|
||||
|
||||
lk.ReleaseImage()
|
||||
waitFor(t, llmAcquired, "LLM acquire after image release")
|
||||
lk.ReleaseLLM()
|
||||
}
|
||||
|
||||
func TestImageFIFO(t *testing.T) {
|
||||
lk := New(nil)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := lk.AcquireImage(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
order := make(chan int, 2)
|
||||
for i, id := range []int{1, 2} {
|
||||
go func() {
|
||||
if err := lk.AcquireImage(ctx); err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
order <- id
|
||||
}()
|
||||
// Ensure this waiter is queued before the next one starts.
|
||||
for deadline := time.Now().Add(2 * time.Second); ; {
|
||||
lk.mu.Lock()
|
||||
qlen := len(lk.imageQ)
|
||||
lk.mu.Unlock()
|
||||
if qlen == i+1 {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("image waiter %d never queued", id)
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
lk.ReleaseImage()
|
||||
if got := <-order; got != 1 {
|
||||
t.Fatalf("first image job = %d, want 1", got)
|
||||
}
|
||||
select {
|
||||
case got := <-order:
|
||||
t.Fatalf("second image job %d acquired while first active", got)
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
lk.ReleaseImage()
|
||||
if got := <-order; got != 2 {
|
||||
t.Fatalf("second image job = %d, want 2", got)
|
||||
}
|
||||
lk.ReleaseImage()
|
||||
}
|
||||
|
||||
func TestContextCancelRemovesLLMWaiter(t *testing.T) {
|
||||
lk := New(nil)
|
||||
if err := lk.AcquireImage(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- lk.AcquireLLM(ctx) }()
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
cancel()
|
||||
if err := <-done; err == nil {
|
||||
t.Fatal("expected context error")
|
||||
}
|
||||
|
||||
lk.ReleaseImage()
|
||||
if state, n, _ := lk.Snapshot(); state != StateIdle || n != 0 {
|
||||
t.Fatalf("got state=%s n=%d, want idle", state, n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestContextCancelRemovesImageWaiter(t *testing.T) {
|
||||
lk := New(nil)
|
||||
ctx := context.Background()
|
||||
if err := lk.AcquireLLM(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
wctx, cancel := context.WithCancel(ctx)
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- lk.AcquireImage(wctx) }()
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
cancel()
|
||||
if err := <-done; err == nil {
|
||||
t.Fatal("expected context error")
|
||||
}
|
||||
|
||||
lk.ReleaseLLM()
|
||||
// The cancelled waiter must be gone: a new LLM request acquires
|
||||
// immediately instead of blocking behind it.
|
||||
if err := lk.AcquireLLM(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, pending := lk.Snapshot(); pending {
|
||||
t.Fatal("image_pending=true after waiter was cancelled")
|
||||
}
|
||||
lk.ReleaseLLM()
|
||||
}
|
||||
|
||||
// TestRace hammers the lock from many goroutines; run with -race.
|
||||
func TestRace(t *testing.T) {
|
||||
lk := New(nil)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 8; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
for j := 0; j < 25; j++ {
|
||||
if i%4 == 3 {
|
||||
if err := lk.AcquireImage(context.Background()); err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
lk.Snapshot()
|
||||
lk.ReleaseImage()
|
||||
} else {
|
||||
if err := lk.AcquireLLM(context.Background()); err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
lk.Snapshot()
|
||||
lk.ReleaseLLM()
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
state, n, pending := lk.Snapshot()
|
||||
if state != StateIdle || n != 0 || pending {
|
||||
t.Fatalf("leaked lock state: state=%s n=%d pending=%v", state, n, pending)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user