diff --git a/.gitignore b/.gitignore index 2d9900c..a2e7d88 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,2 @@ -/gpulock -/gpulock.exe +/gpu-turnstile +/gpu-turnstile.exe diff --git a/Dockerfile b/Dockerfile index 015fb3f..a489c90 100644 --- a/Dockerfile +++ b/Dockerfile @@ -4,10 +4,10 @@ COPY go.mod ./ COPY cmd ./cmd COPY internal ./internal ARG VERSION=dev -RUN CGO_ENABLED=0 go build -ldflags="-s -w -X main.version=${VERSION}" -o /gpulock ./cmd/gpulock +RUN CGO_ENABLED=0 go build -ldflags="-s -w -X main.version=${VERSION}" -o /gpu-turnstile ./cmd/gpu-turnstile FROM gcr.io/distroless/static:nonroot -COPY --from=build /gpulock /gpulock -EXPOSE 8189 11435 +COPY --from=build /gpu-turnstile /gpu-turnstile +EXPOSE 8188 11434 USER nonroot -ENTRYPOINT ["/gpulock"] +ENTRYPOINT ["/gpu-turnstile"] diff --git a/README.md b/README.md index 46de94c..30ee6a0 100644 --- a/README.md +++ b/README.md @@ -1,15 +1,19 @@ -# gpulock +# gpu-turnstile 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. +LLM server (Ollama) and an image generator (ComfyUI); gpu-turnstile 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. + +gpu-turnstile listens on the ports the services normally use; the actual +services run one port higher (Ollama on 11435, ComfyUI on 8189). ``` -LiteLLM / Open WebUI ──► :11435 ─┐ ┌─► Ollama :11434 - ├── gpulock (1 lock) ──┤ -Open WebUI / n8n ────► :8189 ───┘ └─► ComfyUI :8188 +LiteLLM / Open WebUI ──► :11434 ─┐ ┌─► Ollama :11435 + ├── gpu-turnstile (1 lock) ──┤ +Open WebUI / n8n ────► :8188 ───┘ └─► ComfyUI :8189 ``` - LLM endpoints (`/api/generate`, `/api/chat`, `/api/embed`, `/v1/*`) take the @@ -29,10 +33,10 @@ 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 | +| `LISTEN_OLLAMA` | `:11434` | Ollama-facing listener | +| `LISTEN_COMFY` | `:8188` | ComfyUI-facing listener | +| `OLLAMA_URL` | `http://127.0.0.1:11435` | Ollama upstream | +| `COMFY_URL` | `http://127.0.0.1:8189` | 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 | @@ -43,29 +47,29 @@ startup. ## 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`. +- `GET /metrics` (Ollama listener): Prometheus text format — `gpu_turnstile_state`, + `gpu_turnstile_llm_inflight`, `gpu_turnstile_image_pending`, + `gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds` + (histogram, `kind="llm|image"`), `gpu_turnstile_unload_seconds`. ## Build and run ```sh -go build ./cmd/gpulock -./gpulock +go build ./cmd/gpu-turnstile +./gpu-turnstile ``` ```sh -docker build -t gpulock . -docker run --rm -p 11435:11435 -p 8189:8189 \ - -e OLLAMA_URL=http://:11434 \ - -e COMFY_URL=http://:8188 \ - gpulock +docker build -t gpu-turnstile . +docker run --rm -p 11434:11434 -p 8188:8188 \ + -e OLLAMA_URL=http://:11435 \ + -e COMFY_URL=http://:8189 \ + gpu-turnstile ``` Releases are built by Gitea Actions (`.gitea/workflows/ci.yml`): pushes run `go vet` and `go test -race` and publish -`git.rambossek.at//gpulock:sha-`; `main` additionally gets +`git.rambossek.at//gpu-turnstile:sha-`; `main` additionally gets `:latest`, and a git tag `vX.Y.Z` produces the versioned image. ## Development @@ -78,7 +82,7 @@ go test -race ./... Stdlib only, Go 1.23+. Layout: ``` -cmd/gpulock/main.go wiring, config, listeners +cmd/gpu-turnstile/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 diff --git a/SPEC.md b/SPEC.md index 81311c1..b5096d1 100644 --- a/SPEC.md +++ b/SPEC.md @@ -1,4 +1,4 @@ -# gpulock — GPU arbitration proxy for Ollama + ComfyUI +# gpu-turnstile — GPU arbitration proxy for Ollama + ComfyUI ## Problem @@ -16,8 +16,9 @@ at any moment the GPU is in exactly one of three states: - `llm` — N ≥ 1 Ollama inference requests in flight (concurrency allowed) - `image` — exactly one ComfyUI job in flight, Ollama models unloaded -Clients (LiteLLM, Open WebUI, n8n) point at gpulock instead of at the services. -gpulock is transparent for everything that does not touch the GPU. +Clients (LiteLLM, Open WebUI, n8n) point at gpu-turnstile instead of at the +services. gpu-turnstile is transparent for everything that does not touch the +GPU. ## Non-goals @@ -30,15 +31,16 @@ gpulock is transparent for everything that does not touch the GPU. ## Architecture ``` -LiteLLM / Open WebUI ──► :11435 ─┐ ┌─► Ollama :11434 - ├── gpulock (1 lock) ──┤ -Open WebUI / n8n ────► :8189 ───┘ └─► ComfyUI :8188 +LiteLLM / Open WebUI ──► :11434 ─┐ ┌─► Ollama :11435 + ├── gpu-turnstile (1 lock) ──┤ +Open WebUI / n8n ────► :8188 ───┘ └─► ComfyUI :8189 ``` -Two listeners, one process, one lock. Each listener is an -`httputil.ReverseProxy` to its upstream. Websocket upgrades (ComfyUI `/ws`) -and streaming bodies (Ollama NDJSON / SSE) must pass through unbuffered -(`FlushInterval = -1`). +gpu-turnstile listens on the ports the services normally use; the actual +services run one port higher. Two listeners, one process, one lock. Each +listener is an `httputil.ReverseProxy` to its upstream. Websocket upgrades +(ComfyUI `/ws`) and streaming bodies (Ollama NDJSON / SSE) must pass through +unbuffered (`FlushInterval = -1`). ### Lock semantics @@ -58,7 +60,7 @@ are LLM requests and the single "writer" is an image job): ### Endpoint classification -Ollama listener (`:11435` → `OLLAMA_URL`): +Ollama listener (`:11434` → `OLLAMA_URL`): | Path | Handling | |---|---| @@ -66,7 +68,7 @@ Ollama listener (`:11435` → `OLLAMA_URL`): | `POST /v1/chat/completions`, `/v1/completions`, `/v1/embeddings` | LLM lock | | everything else (`/api/tags`, `/api/ps`, `/api/show`, `/api/version`, `/v1/models`, `/api/pull`, …) | pass-through, no lock | -ComfyUI listener (`:8189` → `COMFY_URL`): +ComfyUI listener (`:8188` → `COMFY_URL`): | Path | Handling | |---|---| @@ -100,10 +102,10 @@ load time. Off by default. | Var | Default | Meaning | |---|---|---| -| `LISTEN_OLLAMA` | `:11435` | Ollama-facing listener | -| `LISTEN_COMFY` | `:8189` | ComfyUI-facing listener | -| `OLLAMA_URL` | `http://127.0.0.1:11434` | upstream | -| `COMFY_URL` | `http://127.0.0.1:8188` | upstream | +| `LISTEN_OLLAMA` | `:11434` | Ollama-facing listener | +| `LISTEN_COMFY` | `:8188` | ComfyUI-facing listener | +| `OLLAMA_URL` | `http://127.0.0.1:11435` | upstream | +| `COMFY_URL` | `http://127.0.0.1:8189` | upstream | | `UNLOAD_TIMEOUT` | `60s` | wait for Ollama to unload | | `JOB_TIMEOUT` | `15m` | wait for ComfyUI job | | `LLM_WAIT_TIMEOUT` | `10m` | max time an LLM request waits for the lock before 503 | @@ -119,9 +121,9 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal. `{"state":"idle|llm|image","llm_inflight":N,"image_pending":B}`. - `GET /metrics` on the Ollama listener: Prometheus text format, no external dependency needed: - `gpulock_state{state="…"} 1`, `gpulock_llm_inflight`, - `gpulock_image_jobs_total`, `gpulock_lock_wait_seconds` (histogram, label - `kind="llm|image"`), `gpulock_unload_seconds`. + `gpu_turnstile_state{state="…"} 1`, `gpu_turnstile_llm_inflight`, + `gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds` + (histogram, label `kind="llm|image"`), `gpu_turnstile_unload_seconds`. - Structured logs (`log/slog`, JSON when `LOG_FORMAT=json`), one line per state transition and per image job phase with `prompt_id`. @@ -143,8 +145,8 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal. ## Repository layout ``` -gpulock/ - cmd/gpulock/main.go # wiring, config, listeners +gpu-turnstile/ + cmd/gpu-turnstile/main.go # wiring, config, listeners internal/lock/lock.go # two-mode lock + tests internal/ollama/client.go # ps / unload / warm internal/comfy/client.go # history poll / free @@ -176,30 +178,31 @@ are new. - Go 1.23+, stdlib only. `CGO_ENABLED=0`, `-ldflags="-s -w"`, version from `git describe` injected via `-X main.version=`. - Dockerfile: multi-stage, final image `gcr.io/distroless/static` (or - `scratch`), non-root user, `EXPOSE 8189 11435`, `ENTRYPOINT ["/gpulock"]`. + `scratch`), non-root user, `EXPOSE 8188 11434`, + `ENTRYPOINT ["/gpu-turnstile"]`. - `.gitea/workflows/ci.yml` (Gitea Actions): 1. on push and tag: `go vet`, `go test -race ./...`, `golangci-lint` if available in the runner image 2. build image with buildx, tags `:sha-` and `:latest` on main, `:` on tags - 3. push to the Gitea registry `git.rambossek.at//gpulock` using the - workflow token (`${{ secrets.GITEA_TOKEN }}` / `gitea.actor`) + 3. push to the Gitea registry `git.rambossek.at//gpu-turnstile` + using the workflow token (`${{ secrets.GITEA_TOKEN }}` / `gitea.actor`) - Release: a git tag `vX.Y.Z` produces the versioned image; the Open WebUI compose pins that tag. ## Deployment (target) ```yaml - gpulock: - image: git.rambossek.at//gpulock:v0.1.0 + gpu-turnstile: + image: git.rambossek.at//gpu-turnstile:v0.1.0 environment: - OLLAMA_URL: http://:11434 - COMFY_URL: http://:8188 + OLLAMA_URL: http://:11435 + COMFY_URL: http://:8189 networks: [internal] ``` -LiteLLM `api_base` → `http://gpulock:11435`; Open WebUI -`COMFYUI_BASE_URL` → `http://gpulock:8189`. Nothing else talks to the +LiteLLM `api_base` → `http://gpu-turnstile:11434`; Open WebUI +`COMFYUI_BASE_URL` → `http://gpu-turnstile:8188`. Nothing else talks to the workstation directly. ## Open questions @@ -207,5 +210,5 @@ workstation directly. - Should embedding requests (`/api/embed`, `/v1/embeddings`) count as LLM traffic for the lock? They do in this spec (they hold VRAM); reconsider if RAG indexing starves image jobs for too long. -- Whether to add a `POST /gpulock/release` admin endpoint to force-reset the - lock without restarting. Cheap to add; decide once it's been stuck once. +- Whether to add a `POST /gpu-turnstile/release` admin endpoint to force-reset + the lock without restarting. Cheap to add; decide once it's been stuck once. diff --git a/cmd/gpulock/main.go b/cmd/gpu-turnstile/main.go similarity index 93% rename from cmd/gpulock/main.go rename to cmd/gpu-turnstile/main.go index 2757a08..c848bfd 100644 --- a/cmd/gpulock/main.go +++ b/cmd/gpu-turnstile/main.go @@ -1,4 +1,4 @@ -// gpulock is a GPU arbitration proxy that sits in front of Ollama and +// gpu-turnstile is a GPU arbitration proxy that sits in front of Ollama and // ComfyUI and guarantees only one of them uses the GPU at a time. package main @@ -52,10 +52,10 @@ func envDuration(getenv func(string) string, name string, dst *time.Duration) er func loadConfig(getenv func(string) string) (config, error) { cfg := config{ - listenOllama: ":11435", - listenComfy: ":8189", - ollamaURL: "http://127.0.0.1:11434", - comfyURL: "http://127.0.0.1:8188", + listenOllama: ":11434", + listenComfy: ":8188", + ollamaURL: "http://127.0.0.1:11435", + comfyURL: "http://127.0.0.1:8189", unloadTimeout: time.Minute, jobTimeout: 15 * time.Minute, llmWaitTimeout: 10 * time.Minute, @@ -107,7 +107,7 @@ func loadConfig(getenv func(string) string) (config, error) { func main() { cfg, err := loadConfig(os.Getenv) if err != nil { - fmt.Fprintf(os.Stderr, "gpulock: %v\n", err) + fmt.Fprintf(os.Stderr, "gpu-turnstile: %v\n", err) os.Exit(1) } @@ -119,7 +119,7 @@ func main() { log := slog.New(handler) slog.SetDefault(log) - log.Info("starting gpulock", + log.Info("starting gpu-turnstile", "version", version, "listen_ollama", cfg.listenOllama, "listen_comfy", cfg.listenComfy, diff --git a/internal/comfy/client.go b/internal/comfy/client.go index 1ccacb2..270fb71 100644 --- a/internal/comfy/client.go +++ b/internal/comfy/client.go @@ -1,5 +1,5 @@ -// Package comfy is a minimal client for the ComfyUI endpoints gpulock needs -// after a prompt has been accepted: polling /history and freeing VRAM. +// Package comfy is a minimal client for the ComfyUI endpoints gpu-turnstile +// needs after a prompt has been accepted: polling /history and freeing VRAM. package comfy import ( diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index d024c88..8305beb 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -1,4 +1,4 @@ -// Package metrics provides the Prometheus text exposition for gpulock +// Package metrics provides the Prometheus text exposition for gpu-turnstile // without any external dependencies. package metrics @@ -51,7 +51,7 @@ func trimLabel(label string) string { return label } -// Metrics holds all gpulock metric values. +// Metrics holds all gpu-turnstile metric values. type Metrics struct { llmWait *histogram imageWait *histogram @@ -86,36 +86,36 @@ func (m *Metrics) IncImageJobs() { m.imageJobs.Add(1) } // Render writes the Prometheus text exposition for the given lock snapshot. func (m *Metrics) Render(w io.Writer, state string, llmInflight int, imagePending bool) { - fmt.Fprint(w, `# HELP gpulock_state Current GPU state (1 for the active state). -# TYPE gpulock_state gauge + fmt.Fprint(w, `# HELP gpu_turnstile_state Current GPU state (1 for the active state). +# TYPE gpu_turnstile_state gauge `) for _, s := range []string{"idle", "llm", "image"} { v := 0 if s == state { v = 1 } - fmt.Fprintf(w, "gpulock_state{state=%q} %d\n", s, v) + fmt.Fprintf(w, "gpu_turnstile_state{state=%q} %d\n", s, v) } pending := 0 if imagePending { pending = 1 } - fmt.Fprintf(w, `# HELP gpulock_llm_inflight LLM requests currently in flight. -# TYPE gpulock_llm_inflight gauge -gpulock_llm_inflight %d -# HELP gpulock_image_pending Whether an image job is active or waiting. -# TYPE gpulock_image_pending gauge -gpulock_image_pending %d -# HELP gpulock_image_jobs_total Image jobs accepted by ComfyUI. -# TYPE gpulock_image_jobs_total counter -gpulock_image_jobs_total %d -# HELP gpulock_lock_wait_seconds Time spent waiting to acquire the GPU lock. -# TYPE gpulock_lock_wait_seconds histogram + fmt.Fprintf(w, `# HELP gpu_turnstile_llm_inflight LLM requests currently in flight. +# TYPE gpu_turnstile_llm_inflight gauge +gpu_turnstile_llm_inflight %d +# HELP gpu_turnstile_image_pending Whether an image job is active or waiting. +# TYPE gpu_turnstile_image_pending gauge +gpu_turnstile_image_pending %d +# HELP gpu_turnstile_image_jobs_total Image jobs accepted by ComfyUI. +# TYPE gpu_turnstile_image_jobs_total counter +gpu_turnstile_image_jobs_total %d +# HELP gpu_turnstile_lock_wait_seconds Time spent waiting to acquire the GPU lock. +# TYPE gpu_turnstile_lock_wait_seconds histogram `, llmInflight, pending, m.imageJobs.Load()) - m.llmWait.write(w, "gpulock_lock_wait_seconds", `kind="llm",`) - m.imageWait.write(w, "gpulock_lock_wait_seconds", `kind="image",`) - fmt.Fprint(w, `# HELP gpulock_unload_seconds Time spent unloading Ollama models before an image job. -# TYPE gpulock_unload_seconds histogram + m.llmWait.write(w, "gpu_turnstile_lock_wait_seconds", `kind="llm",`) + m.imageWait.write(w, "gpu_turnstile_lock_wait_seconds", `kind="image",`) + fmt.Fprint(w, `# HELP gpu_turnstile_unload_seconds Time spent unloading Ollama models before an image job. +# TYPE gpu_turnstile_unload_seconds histogram `) - m.unload.write(w, "gpulock_unload_seconds", "") + m.unload.write(w, "gpu_turnstile_unload_seconds", "") } diff --git a/internal/ollama/client.go b/internal/ollama/client.go index 2a7fbdb..8c57e72 100644 --- a/internal/ollama/client.go +++ b/internal/ollama/client.go @@ -1,5 +1,6 @@ // Package ollama is a minimal client for the Ollama management endpoints -// gpulock needs: listing loaded models, unloading them, and warming a model. +// gpu-turnstile needs: listing loaded models, unloading them, and warming a +// model. package ollama import ( diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 53e2155..9c13177 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -1,4 +1,4 @@ -// Package proxy contains the HTTP handlers for both gpulock listeners: +// Package proxy contains the HTTP handlers for both gpu-turnstile listeners: // reverse proxies to Ollama and ComfyUI with GPU lock arbitration in front // of the endpoints that load models. package proxy @@ -43,7 +43,7 @@ type Config struct { WarmModel string } -// Server serves both gpulock listeners. +// Server serves both gpu-turnstile listeners. type Server struct { cfg Config log *slog.Logger diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index 0d26f96..c04a0f3 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -40,13 +40,13 @@ func (r *recorder) index(e string) int { return -1 } -// fakes wires up fake Ollama and ComfyUI upstreams plus the gpulock server. +// fakes wires up fake Ollama and ComfyUI upstreams plus the gpu-turnstile server. type fakes struct { rec *recorder ollama *httptest.Server comfy *httptest.Server - server *httptest.Server // Ollama-facing gpulock listener - comfySrv *httptest.Server // ComfyUI-facing gpulock listener + server *httptest.Server // Ollama-facing gpu-turnstile listener + comfySrv *httptest.Server // ComfyUI-facing gpu-turnstile listener freeCh chan struct{} chatCh chan struct{} historyMu sync.Mutex @@ -279,7 +279,7 @@ func TestStreamingNotBuffered(t *testing.T) { } defer resp.Body.Close() - // If gpulock buffered the stream, this read would never complete while + // If gpu-turnstile buffered the stream, this read would never complete while // the gate is closed. line, err := bufio.NewReader(resp.Body).ReadString('\n') if err != nil { @@ -321,11 +321,11 @@ func TestHealthzAndMetrics(t *testing.T) { resp.Body.Close() text := string(body) for _, want := range []string{ - `gpulock_state{state="idle"} 1`, - "gpulock_llm_inflight 0", - "gpulock_image_jobs_total 0", - `gpulock_lock_wait_seconds_bucket{kind="llm",le="+Inf"} 0`, - "gpulock_unload_seconds_count{} 0", + `gpu_turnstile_state{state="idle"} 1`, + "gpu_turnstile_llm_inflight 0", + "gpu_turnstile_image_jobs_total 0", + `gpu_turnstile_lock_wait_seconds_bucket{kind="llm",le="+Inf"} 0`, + "gpu_turnstile_unload_seconds_count{} 0", } { if !strings.Contains(text, want) { t.Fatalf("metrics missing %q", want)