282 lines
7.0 KiB
Go
282 lines
7.0 KiB
Go
// 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. An
|
|
// external hold (SetExternal) blocks new grants of both kinds while a
|
|
// foreign process — e.g. a game — holds the GPU; in-flight work drains.
|
|
package lock
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// State is the current GPU occupancy state.
|
|
type State string
|
|
|
|
const (
|
|
StateIdle State = "idle"
|
|
StateLLM State = "llm"
|
|
StateImage State = "image"
|
|
StateExternal State = "external"
|
|
)
|
|
|
|
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
|
|
llmWaiting int // LLM requests blocked waiting for the GPU
|
|
imageActive bool // an image job holds the GPU
|
|
imageQ []imageWaiter
|
|
nextID uint64
|
|
external string // non-empty: a foreign process (e.g. a game) holds the GPU
|
|
detail string // what the current holder is doing (best effort)
|
|
since time.Time // when the current state began
|
|
|
|
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, since: time.Now()}
|
|
}
|
|
|
|
// 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...)
|
|
}
|
|
}
|
|
|
|
// SetExternal records that a process outside gpu-turnstile's control (a
|
|
// game, another ML job) holds the GPU: new LLM and image grants block until
|
|
// ClearExternal. In-flight work is not preempted. holder describes the
|
|
// process for logs and busy responses.
|
|
func (l *Lock) SetExternal(holder string) {
|
|
l.mu.Lock()
|
|
l.external = holder
|
|
l.since = time.Now()
|
|
l.broadcast()
|
|
l.mu.Unlock()
|
|
l.logTransition("lock transition", "state", StateExternal, "holder", holder)
|
|
}
|
|
|
|
// ClearExternal lifts the external hold; waiting LLM and image requests
|
|
// proceed.
|
|
func (l *Lock) ClearExternal() {
|
|
l.mu.Lock()
|
|
l.external = ""
|
|
l.since = time.Now()
|
|
l.broadcast()
|
|
l.mu.Unlock()
|
|
l.logTransition("lock transition", "state", StateIdle)
|
|
}
|
|
|
|
// External returns the current external holder, or "" when none.
|
|
func (l *Lock) External() string {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
return l.external
|
|
}
|
|
|
|
// 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()
|
|
waiting := false
|
|
for l.imageActive || len(l.imageQ) > 0 || l.external != "" {
|
|
if !waiting {
|
|
l.llmWaiting++
|
|
waiting = true
|
|
}
|
|
ch := l.change
|
|
l.mu.Unlock()
|
|
select {
|
|
case <-ctx.Done():
|
|
l.mu.Lock()
|
|
l.llmWaiting--
|
|
l.mu.Unlock()
|
|
return ctx.Err()
|
|
case <-ch:
|
|
}
|
|
l.mu.Lock()
|
|
}
|
|
if waiting {
|
|
l.llmWaiting--
|
|
}
|
|
l.n++
|
|
if l.n == 1 {
|
|
l.since = time.Now()
|
|
}
|
|
n := l.n
|
|
l.mu.Unlock()
|
|
l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
|
|
return nil
|
|
}
|
|
|
|
// TryAcquireLLM acquires one in-flight LLM slot without waiting and
|
|
// reports whether it succeeded. It fails when an image job is active or
|
|
// pending or an external hold is set.
|
|
func (l *Lock) TryAcquireLLM() bool {
|
|
l.mu.Lock()
|
|
if l.imageActive || len(l.imageQ) > 0 || l.external != "" {
|
|
l.mu.Unlock()
|
|
return false
|
|
}
|
|
l.n++
|
|
n := l.n
|
|
l.mu.Unlock()
|
|
l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
|
|
return true
|
|
}
|
|
|
|
// ReleaseLLM marks one LLM request as finished.
|
|
func (l *Lock) ReleaseLLM() {
|
|
l.mu.Lock()
|
|
l.n--
|
|
n := l.n
|
|
if l.n == 0 {
|
|
l.since = time.Now()
|
|
l.detail = ""
|
|
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.external == "" {
|
|
l.imageQ = l.imageQ[1:]
|
|
l.imageActive = true
|
|
l.since = time.Now()
|
|
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.since = time.Now()
|
|
l.detail = ""
|
|
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
|
|
case l.external != "":
|
|
state = StateExternal
|
|
default:
|
|
state = StateIdle
|
|
}
|
|
return state, l.n, l.imageActive || len(l.imageQ) > 0
|
|
}
|
|
|
|
// SetDetail records what the current holder is doing (e.g. the request
|
|
// path), for status displays. Best effort: overwritten by each new holder,
|
|
// cleared when the GPU goes idle.
|
|
func (l *Lock) SetDetail(detail string) {
|
|
l.mu.Lock()
|
|
l.detail = detail
|
|
l.mu.Unlock()
|
|
}
|
|
|
|
// Status is a point-in-time view of the lock for monitoring.
|
|
type Status struct {
|
|
State State
|
|
Detail string
|
|
LLMInflight int
|
|
LLMWaiting int
|
|
ImageActive bool
|
|
ImageQueue int
|
|
External string
|
|
Since time.Time
|
|
}
|
|
|
|
// Status reports the full lock state, including waiters and how long the
|
|
// current state has held.
|
|
func (l *Lock) Status() Status {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
s := Status{
|
|
Detail: l.detail,
|
|
LLMInflight: l.n,
|
|
LLMWaiting: l.llmWaiting,
|
|
ImageActive: l.imageActive,
|
|
ImageQueue: len(l.imageQ),
|
|
External: l.external,
|
|
Since: l.since,
|
|
}
|
|
switch {
|
|
case l.imageActive:
|
|
s.State = StateImage
|
|
case l.n > 0:
|
|
s.State = StateLLM
|
|
case l.external != "":
|
|
s.State = StateExternal
|
|
default:
|
|
s.State = StateIdle
|
|
}
|
|
return s
|
|
}
|