diff --git a/README.md b/README.md index b2cb20f..c107af9 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,12 @@ Open WebUI / n8n ────► :8188 ───┘ - Everything else (including websockets and all streaming) passes through transparently and unbuffered. +Each consumer is enabled by setting its URL (`OLLAMA_URL`, `COMFY_URL`) and +disabled by leaving it empty — at least one is required. With only Ollama +the proxy is a pass-through (no image jobs can arrive); with only ComfyUI +the Ollama unload/warm steps are skipped. Future consumers (e.g. local game +detection) plug into the same lock the same way. + ## Configuration Configuration comes from environment variables and/or an `.env`-style @@ -41,8 +47,8 @@ override file values. Invalid values fail at startup. |---|---|---| | `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 | +| `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer | +| `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer | | `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 (wait mode) | @@ -70,7 +76,7 @@ override file values. Invalid values fail at startup. ## Observability - `GET /healthz` (both listeners): `{"state":"idle|llm|image","llm_inflight":N,"image_pending":B}` -- `GET /metrics` (Ollama listener): Prometheus text format — `gpu_turnstile_state`, +- `GET /metrics` (both listeners): 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`. diff --git a/SPEC.md b/SPEC.md index 08193ab..8d7c80c 100644 --- a/SPEC.md +++ b/SPEC.md @@ -42,6 +42,21 @@ listener is an `httputil.ReverseProxy` to its upstream. Websocket upgrades (ComfyUI `/ws`) and streaming bodies (Ollama NDJSON / SSE) must pass through unbuffered (`FlushInterval = -1`). +### Modes of operation + +Each GPU consumer is enabled by setting its URL and disabled by leaving it +empty — no separate flags. At least one URL must be set; a disabled +consumer gets no listener, no startup probe, and no lock participation: + +- **Both set** (default deployment): full arbitration as described below. +- **Only `OLLAMA_URL`**: pure pass-through for Ollama; the LLM lock never + blocks since no image jobs can arrive. +- **Only `COMFY_URL`**: image jobs are tracked and ComfyUI's VRAM is freed + afterwards, but the Ollama unload and warm-reload steps are skipped. +- Future consumers (e.g. detecting a local game holding VRAM) plug into the + same lock the same way: enabled by their config knob, excluded when + absent. + ### Lock semantics Two-mode lock with image priority (writer-preferring RW lock, where "readers" @@ -83,7 +98,8 @@ ComfyUI listener (`:8188` → `COMFY_URL`): ### Image job flow (`POST /prompt`) 1. `AcquireImage()`. -2. Unload Ollama: `GET /api/ps`; for each model `POST /api/generate +2. Unload Ollama (skipped when `OLLAMA_URL` is unset): `GET /api/ps`; for + each model `POST /api/generate {"model":M,"keep_alive":0}`; if that returns non-2xx (embedding-only models), `POST /api/embed {"model":M,"input":"x","keep_alive":0}`. Poll `/api/ps` every `UNLOAD_POLL_INTERVAL` (default 500 ms) until empty or @@ -117,8 +133,8 @@ override file values. A missing file is fine; a malformed one is fatal. |---|---|---| | `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 | +| `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer | +| `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer | | `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 (wait mode) | @@ -143,8 +159,9 @@ override file values. A missing file is fine; a malformed one is fatal. | `UPDATE_REPO` | `https://git.rambossek.at/PUBLIC/gpu-turnstile` | repository to check for releases | | `UPDATE_ASSET` | `gpu-turnstile.exe` | release asset to download | -Startup fails fast on unparsable values. Both upstreams are probed once at -start (`/api/version`, `/system_stats`); failure is logged, not fatal. +Startup fails fast on unparsable values and when neither consumer URL is +set. Enabled upstreams are probed once at start (`/api/version`, +`/system_stats`); failure is logged, not fatal. ## Native Windows deployment @@ -184,7 +201,7 @@ Docker. - `GET /healthz` on both listeners: 200 with JSON `{"state":"idle|llm|image","llm_inflight":N,"image_pending":B}`. -- `GET /metrics` on the Ollama listener: Prometheus text format, no external +- `GET /metrics` on both listeners: Prometheus text format, no external dependency needed: `gpu_turnstile_state{state="…"} 1`, `gpu_turnstile_llm_inflight`, `gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds` diff --git a/cmd/gpu-turnstile/main.go b/cmd/gpu-turnstile/main.go index dcd6f4f..babed5f 100644 --- a/cmd/gpu-turnstile/main.go +++ b/cmd/gpu-turnstile/main.go @@ -180,6 +180,14 @@ func serviceCommand(configPath string, args []string) int { return 0 } +// orDisabled renders an empty URL as "disabled" for the startup dump. +func orDisabled(url string) string { + if url == "" { + return "disabled" + } + return url +} + func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Writer, isService bool) error { // The startup line carries the version and every setting and is emitted // at WARN so it is visible even with the default (quiet) log level. @@ -187,8 +195,8 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri "version", version, "listen_ollama", cfg.ListenOllama, "listen_comfy", cfg.ListenComfy, - "ollama_url", cfg.OllamaURL, - "comfy_url", cfg.ComfyURL, + "ollama_url", orDisabled(cfg.OllamaURL), + "comfy_url", orDisabled(cfg.ComfyURL), "unload_timeout", cfg.UnloadTimeout, "job_timeout", cfg.JobTimeout, "llm_wait_timeout", cfg.LLMWaitTimeout, @@ -215,13 +223,20 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri ) lk := lock.New(log) - ollamaClient, err := ollama.New(cfg.OllamaURL, log) - if err != nil { - return err + // Each consumer is enabled by setting its URL; a disabled consumer gets + // no client, no listener and no probe. + var ollamaClient *ollama.Client + var err error + if cfg.OllamaURL != "" { + if ollamaClient, err = ollama.New(cfg.OllamaURL, log); err != nil { + return err + } } - comfyClient, err := comfy.New(cfg.ComfyURL, log) - if err != nil { - return err + var comfyClient *comfy.Client + if cfg.ComfyURL != "" { + if comfyClient, err = comfy.New(cfg.ComfyURL, log); err != nil { + return err + } } srv, err := proxy.New(proxy.Config{ @@ -253,22 +268,34 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri return err } - // Probe both upstreams once; failure is logged, not fatal. + // Probe the enabled upstreams once; failure is logged, not fatal. probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout) - if err := ollamaClient.Probe(probeCtx); err != nil { - log.Warn("ollama probe failed", "url", cfg.OllamaURL, "err", err) + if ollamaClient != nil { + if err := ollamaClient.Probe(probeCtx); err != nil { + log.Warn("ollama probe failed", "url", cfg.OllamaURL, "err", err) + } } - if err := comfyClient.Probe(probeCtx); err != nil { - log.Warn("comfy probe failed", "url", cfg.ComfyURL, "err", err) + if comfyClient != nil { + if err := comfyClient.Probe(probeCtx); err != nil { + log.Warn("comfy probe failed", "url", cfg.ComfyURL, "err", err) + } } probeCancel() - ollamaSrv := &http.Server{Addr: cfg.ListenOllama, Handler: srv.OllamaHandler()} - comfySrv := &http.Server{Addr: cfg.ListenComfy, Handler: srv.ComfyHandler()} + var servers []*http.Server + if ollamaClient != nil { + servers = append(servers, &http.Server{Addr: cfg.ListenOllama, Handler: srv.OllamaHandler()}) + log.Warn("listening", "consumer", "ollama", "addr", cfg.ListenOllama) + } + if comfyClient != nil { + servers = append(servers, &http.Server{Addr: cfg.ListenComfy, Handler: srv.ComfyHandler()}) + log.Warn("listening", "consumer", "comfy", "addr", cfg.ListenComfy) + } - errCh := make(chan error, 2) - go func() { errCh <- ollamaSrv.ListenAndServe() }() - go func() { errCh <- comfySrv.ListenAndServe() }() + errCh := make(chan error, len(servers)) + for _, s := range servers { + go func(s *http.Server) { errCh <- s.ListenAndServe() }(s) + } if cfg.AutoUpdate { go updateLoop(ctx, cfg, log, lk, isService) @@ -285,8 +312,9 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), cfg.ShutdownTimeout) defer shutdownCancel() - ollamaSrv.Shutdown(shutdownCtx) - comfySrv.Shutdown(shutdownCtx) + for _, s := range servers { + s.Shutdown(shutdownCtx) + } return nil } diff --git a/compose.yaml.example b/compose.yaml.example index ee0cc2a..6a3ff7f 100644 --- a/compose.yaml.example +++ b/compose.yaml.example @@ -9,6 +9,8 @@ services: image: git.rambossek.at/public/gpu-turnstile:v0.1.3 restart: unless-stopped environment: + # Each consumer is enabled by setting its URL; leave one unset to + # disable that side (no listener, no probe, no lock participation). # Services on the Docker host itself: OLLAMA_URL: http://host.docker.internal:11435 COMFY_URL: http://host.docker.internal:8189 diff --git a/internal/config/config.go b/internal/config/config.go index 5785dd1..d5efa85 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -51,13 +51,12 @@ type Config struct { } // Defaults returns the configuration used when neither the environment nor -// a config file sets a value. +// a config file sets a value. The upstream URLs default to empty: a +// consumer is enabled by setting its URL, disabled by leaving it empty. func Defaults() Config { return Config{ 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, @@ -219,5 +218,8 @@ func Load(getenv func(string) string) (Config, error) { default: return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"") } + if cfg.OllamaURL == "" && cfg.ComfyURL == "" { + return cfg, fmt.Errorf("at least one of OLLAMA_URL or COMFY_URL must be set (each URL enables its consumer)") + } return cfg, nil } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index fa51d5a..17411cd 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -8,13 +8,21 @@ import ( ) func TestDefaults(t *testing.T) { - cfg, err := Load(func(string) string { return "" }) + cfg, err := Load(func(k string) string { + if k == "OLLAMA_URL" { + return "http://127.0.0.1:11435" + } + return "" + }) if err != nil { t.Fatal(err) } if cfg.ListenOllama != ":11434" || cfg.ListenComfy != ":8188" { t.Fatalf("listen addrs = %s %s", cfg.ListenOllama, cfg.ListenComfy) } + if cfg.ComfyURL != "" { + t.Fatalf("ComfyURL default = %q, want empty (disabled)", cfg.ComfyURL) + } if cfg.UnloadTimeout != time.Minute || cfg.JobTimeout != 15*time.Minute { t.Fatalf("timeouts = %v %v", cfg.UnloadTimeout, cfg.JobTimeout) } @@ -26,6 +34,13 @@ func TestDefaults(t *testing.T) { } } +func TestLoadRequiresConsumer(t *testing.T) { + _, err := Load(func(string) string { return "" }) + if err == nil || !strings.Contains(err.Error(), "OLLAMA_URL") { + t.Fatalf("err = %v, want missing-consumer error", err) + } +} + func TestParseEnvFile(t *testing.T) { input := `# comment OLLAMA_URL=http://host:11435 @@ -91,6 +106,9 @@ func TestLoadErrors(t *testing.T) { if k == tc.key { return tc.value } + if k == "OLLAMA_URL" { + return "http://127.0.0.1:11435" + } return "" }) if err == nil { diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index c4217c9..8471846 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -102,24 +102,37 @@ type Server struct { comfyProxy *httputil.ReverseProxy } -// New builds a Server, validating the upstream URLs. +// New builds a Server, validating the upstream URLs. At least one of +// OllamaURL / ComfyURL must be set; an empty URL disables that consumer — +// its handler is then never served, its client may be nil, and the image +// job flow skips the Ollama unload/warm steps. func New(cfg Config) (*Server, error) { - ollamaURL, err := url.Parse(cfg.OllamaURL) - if err != nil || ollamaURL.Scheme == "" || ollamaURL.Host == "" { - return nil, fmt.Errorf("invalid OLLAMA_URL %q", cfg.OllamaURL) + if cfg.OllamaURL == "" && cfg.ComfyURL == "" { + return nil, fmt.Errorf("at least one of OllamaURL or ComfyURL is required") } - comfyURL, err := url.Parse(cfg.ComfyURL) - if err != nil || comfyURL.Scheme == "" || comfyURL.Host == "" { - return nil, fmt.Errorf("invalid COMFY_URL %q", cfg.ComfyURL) + var ollamaURL, comfyURL *url.URL + if cfg.OllamaURL != "" { + u, err := url.Parse(cfg.OllamaURL) + if err != nil || u.Scheme == "" || u.Host == "" { + return nil, fmt.Errorf("invalid OLLAMA_URL %q", cfg.OllamaURL) + } + ollamaURL = u + } + if cfg.ComfyURL != "" { + u, err := url.Parse(cfg.ComfyURL) + if err != nil || u.Scheme == "" || u.Host == "" { + return nil, fmt.Errorf("invalid COMFY_URL %q", cfg.ComfyURL) + } + comfyURL = u } log := cfg.Log if log == nil { log = slog.Default() } - if cfg.UnloadPollInterval > 0 { + if cfg.UnloadPollInterval > 0 && cfg.Ollama != nil { cfg.Ollama.PollInterval = cfg.UnloadPollInterval } - if cfg.HistoryPollInterval > 0 { + if cfg.HistoryPollInterval > 0 && cfg.Comfy != nil { cfg.Comfy.PollInterval = cfg.HistoryPollInterval } freeTimeout := cfg.FreeTimeout @@ -160,7 +173,7 @@ func New(cfg Config) (*Server, error) { max: backoffMax, log: log, } - return &Server{ + s := &Server{ cfg: cfg, log: log, logWriter: cfg.LogWriter, @@ -172,9 +185,14 @@ func New(cfg Config) (*Server, error) { busyMode: busyMode, busyStatus: busyStatus, busyRetryAfter: busyRetryAfter, - ollamaProxy: newReverseProxy(ollamaURL, retry, log.With("upstream", "ollama")), - comfyProxy: newReverseProxy(comfyURL, retry, log.With("upstream", "comfy")), - }, nil + } + if ollamaURL != nil { + s.ollamaProxy = newReverseProxy(ollamaURL, retry, log.With("upstream", "ollama")) + } + if comfyURL != nil { + s.comfyProxy = newReverseProxy(comfyURL, retry, log.With("upstream", "comfy")) + } + return s, nil } // retryTransport retries requests whose failure means the upstream never @@ -470,9 +488,13 @@ func (s *Server) OllamaHandler() http.Handler { // ComfyHandler serves the ComfyUI-facing listener. func (s *Server) ComfyHandler() http.Handler { return s.logRequests("comfy", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/healthz" { + switch r.URL.Path { + case "/healthz": s.writeHealthz(w) return + case "/metrics": + s.writeMetrics(w) + return } if r.Method == http.MethodPost && r.URL.Path == "/prompt" { s.handlePrompt(w, r) @@ -527,19 +549,21 @@ func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) { s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds()) log.Info("image lock acquired") - uctx, ucancel := context.WithTimeout(r.Context(), s.cfg.UnloadTimeout) - elapsed, uerr := s.cfg.Ollama.UnloadAll(uctx) - ucancel() - s.cfg.Metrics.ObserveUnload(elapsed.Seconds()) - switch { - case r.Context().Err() != nil: - s.cfg.Lock.ReleaseImage() - return - case uerr != nil: - // Degrade, don't fail the user's request on a misbehaving neighbour. - log.Warn("ollama unload incomplete; continuing", "err", uerr) - default: - log.Info("ollama models unloaded", "seconds", elapsed.Seconds()) + if s.cfg.Ollama != nil { + uctx, ucancel := context.WithTimeout(r.Context(), s.cfg.UnloadTimeout) + elapsed, uerr := s.cfg.Ollama.UnloadAll(uctx) + ucancel() + s.cfg.Metrics.ObserveUnload(elapsed.Seconds()) + switch { + case r.Context().Err() != nil: + s.cfg.Lock.ReleaseImage() + return + case uerr != nil: + // Degrade, don't fail the user's request on a misbehaving neighbour. + log.Warn("ollama unload incomplete; continuing", "err", uerr) + default: + log.Info("ollama models unloaded", "seconds", elapsed.Seconds()) + } } cw := &captureWriter{ResponseWriter: w, status: http.StatusOK, limit: s.captureLimit} @@ -585,7 +609,7 @@ func (s *Server) finishImageJob(promptID string) { s.cfg.Lock.ReleaseImage() log.Info("image lock released") - if s.cfg.WarmModel != "" { + if s.cfg.Ollama != nil && s.cfg.WarmModel != "" { if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle { wctx, wcancel := context.WithTimeout(context.Background(), s.warmTimeout) if err := s.cfg.Ollama.Warm(wctx, s.cfg.WarmModel); err != nil { diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index 2b68702..92b3cb2 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -451,3 +451,82 @@ func TestLLMBusyWaitTimeoutRetryAfter(t *testing.T) { t.Fatalf("Retry-After = %q", got) } } + +func TestComfyOnlyModeSkipsOllama(t *testing.T) { + f := newFakes(t) + + comfyClient, err := comfy.New(f.comfy.URL, nil) + if err != nil { + t.Fatal(err) + } + comfyClient.PollInterval = 5 * time.Millisecond + srv, err := New(Config{ + ComfyURL: f.comfy.URL, + Lock: lock.New(nil), + Comfy: comfyClient, + Metrics: metrics.New(), + JobTimeout: 2 * time.Second, + }) + if err != nil { + t.Fatal(err) + } + front := httptest.NewServer(srv.ComfyHandler()) + defer front.Close() + + resp, err := http.Post(front.URL+"/prompt", "application/json", strings.NewReader(`{}`)) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != 200 { + t.Fatalf("prompt status = %d", resp.StatusCode) + } + f.completeJob() + select { + case <-f.freeCh: + case <-time.After(3 * time.Second): + t.Fatal("/free never called") + } + // With Ollama disabled the unload steps must not happen. + if i := f.rec.index("unload"); i >= 0 { + f.rec.mu.Lock() + t.Fatalf("unload called with ollama disabled; events: %v", f.rec.events) + } +} + +func TestOllamaOnlyMode(t *testing.T) { + f := newFakes(t) + + ollamaClient, err := ollama.New(f.ollama.URL, nil) + if err != nil { + t.Fatal(err) + } + srv, err := New(Config{ + OllamaURL: f.ollama.URL, + Lock: lock.New(nil), + Ollama: ollamaClient, + Metrics: metrics.New(), + LLMWaitTimeout: time.Second, + }) + if err != nil { + t.Fatal(err) + } + front := httptest.NewServer(srv.OllamaHandler()) + defer front.Close() + + resp, err := http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{}`)) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != 200 { + t.Fatalf("chat status = %d", resp.StatusCode) + } +} + +func TestNewRequiresConsumer(t *testing.T) { + _, err := New(Config{Lock: lock.New(nil), Metrics: metrics.New()}) + if err == nil { + t.Fatal("New with no upstream URLs should fail") + } +}