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,87 @@
+# 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](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
+
+```sh
+go build ./cmd/gpulock
+./gpulock
+```
+
+```sh
+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
+
+```sh
+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
+```