mram
065c294a96
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.
88 lines
3.2 KiB
Markdown
88 lines
3.2 KiB
Markdown
# 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
|
|
```
|