Files

T

mram 065c294a96 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

3.2 KiB

Raw Blame History

gpulock

GPU arbitration proxy for Ollama + ComfyUI. One consumer GPU is shared by an LLM server (Ollama) and an image generator (ComfyUI); gpulock sits in front of both and guarantees the GPU is always in exactly one of three states: idle, llm (N ≥ 1 Ollama requests in flight), or image (exactly one ComfyUI job, Ollama models unloaded). See SPEC.md for the full design.

LiteLLM / Open WebUI ──► :11435 ─┐                     ┌─► Ollama  :11434
                                 ├── gpulock (1 lock) ──┤
Open WebUI / n8n ────► :8189  ───┘                     └─► ComfyUI :8188

LLM endpoints (/api/generate, /api/chat, /api/embed, /v1/*) take the LLM lock: concurrent requests allowed, but blocked while an image job is active or waiting (image priority).
POST /prompt on the ComfyUI listener takes the image lock: new LLM requests block, in-flight LLMs drain, Ollama models are unloaded, the prompt is forwarded, and the lock is held until the job finishes and ComfyUI frees its VRAM.
Everything else (including websockets and all streaming) passes through transparently and unbuffered.

Configuration

All configuration is via environment variables; invalid values fail at startup.

Var	Default	Meaning
`LISTEN_OLLAMA`	`:11435`	Ollama-facing listener
`LISTEN_COMFY`	`:8189`	ComfyUI-facing listener
`OLLAMA_URL`	`http://127.0.0.1:11434`	Ollama upstream
`COMFY_URL`	`http://127.0.0.1:8188`	ComfyUI upstream
`UNLOAD_TIMEOUT`	`60s`	Wait for Ollama to unload before an image job
`JOB_TIMEOUT`	`15m`	Wait for a ComfyUI job to finish
`LLM_WAIT_TIMEOUT`	`10m`	Max lock wait for an LLM request before 503
`WARM_MODEL`	(empty)	Model to reload after an image job (off by default)
`LOG_LEVEL`	`info`	`debug` logs every lock transition
`LOG_FORMAT`	`text`	`json` for structured JSON logs

Observability

GET /healthz (both listeners): {"state":"idle|llm|image","llm_inflight":N,"image_pending":B}
GET /metrics (Ollama listener): Prometheus text format — gpulock_state, gpulock_llm_inflight, gpulock_image_pending, gpulock_image_jobs_total, gpulock_lock_wait_seconds (histogram, kind="llm|image"), gpulock_unload_seconds.

Build and run

go build ./cmd/gpulock
./gpulock

docker build -t gpulock .
docker run --rm -p 11435:11435 -p 8189:8189 \
  -e OLLAMA_URL=http://<workstation-ip>:11434 \
  -e COMFY_URL=http://<workstation-ip>:8188 \
  gpulock

Releases are built by Gitea Actions (.gitea/workflows/ci.yml): pushes run go vet and go test -race and publish git.rambossek.at/<owner>/gpulock:sha-<short>; main additionally gets :latest, and a git tag vX.Y.Z produces the versioned image.

Development

go vet ./...
go test -race ./...

Stdlib only, Go 1.23+. Layout:

cmd/gpulock/main.go        wiring, config, listeners
internal/lock/             two-mode lock (LLM readers / image writer, FIFO)
internal/ollama/           ps / unload / warm client
internal/comfy/            history poll / free client
internal/proxy/            handlers for both listeners
internal/metrics/          Prometheus exposition, no dependencies