Make consumers optional: a URL enables its mode, empty disables it

OLLAMA_URL and COMFY_URL no longer have defaults; each consumer
(listener, client, startup probe, lock participation) is enabled by
setting its URL and disabled by leaving it empty. At least one must be
set. Ollama-only mode is a pure pass-through; ComfyUI-only mode skips
the unload and warm-reload steps. /metrics is now served on both
listeners. This is the extension pattern for future consumers such as
local game detection.
This commit is contained in:
mram
2026-09-20 22:29:16 +02:00
parent 642cc36a39
commit 30a1f55aae
8 changed files with 237 additions and 61 deletions
+9 -3
View File
@@ -29,6 +29,12 @@ Open WebUI / n8n ────► :8188 ───┘
- Everything else (including websockets and all streaming) passes through - Everything else (including websockets and all streaming) passes through
transparently and unbuffered. 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
Configuration comes from environment variables and/or an `.env`-style 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_OLLAMA` | `:11434` | Ollama-facing listener |
| `LISTEN_COMFY` | `:8188` | ComfyUI-facing listener | | `LISTEN_COMFY` | `:8188` | ComfyUI-facing listener |
| `OLLAMA_URL` | `http://127.0.0.1:11435` | Ollama upstream | | `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer |
| `COMFY_URL` | `http://127.0.0.1:8189` | ComfyUI upstream | | `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer |
| `UNLOAD_TIMEOUT` | `60s` | Wait for Ollama to unload before an image job | | `UNLOAD_TIMEOUT` | `60s` | Wait for Ollama to unload before an image job |
| `JOB_TIMEOUT` | `15m` | Wait for a ComfyUI job to finish | | `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) | | `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 ## Observability
- `GET /healthz` (both listeners): `{"state":"idle|llm|image","llm_inflight":N,"image_pending":B}` - `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_llm_inflight`, `gpu_turnstile_image_pending`,
`gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds` `gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
(histogram, `kind="llm|image"`), `gpu_turnstile_unload_seconds`. (histogram, `kind="llm|image"`), `gpu_turnstile_unload_seconds`.
+23 -6
View File
@@ -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 (ComfyUI `/ws`) and streaming bodies (Ollama NDJSON / SSE) must pass through
unbuffered (`FlushInterval = -1`). 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 ### Lock semantics
Two-mode lock with image priority (writer-preferring RW lock, where "readers" 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`) ### Image job flow (`POST /prompt`)
1. `AcquireImage()`. 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 {"model":M,"keep_alive":0}`; if that returns non-2xx (embedding-only
models), `POST /api/embed {"model":M,"input":"x","keep_alive":0}`. Poll models), `POST /api/embed {"model":M,"input":"x","keep_alive":0}`. Poll
`/api/ps` every `UNLOAD_POLL_INTERVAL` (default 500 ms) until empty or `/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_OLLAMA` | `:11434` | Ollama-facing listener |
| `LISTEN_COMFY` | `:8188` | ComfyUI-facing listener | | `LISTEN_COMFY` | `:8188` | ComfyUI-facing listener |
| `OLLAMA_URL` | `http://127.0.0.1:11435` | upstream | | `OLLAMA_URL` | _(empty = disabled)_ | Ollama upstream; set to enable the Ollama consumer |
| `COMFY_URL` | `http://127.0.0.1:8189` | upstream | | `COMFY_URL` | _(empty = disabled)_ | ComfyUI upstream; set to enable the ComfyUI consumer |
| `UNLOAD_TIMEOUT` | `60s` | wait for Ollama to unload | | `UNLOAD_TIMEOUT` | `60s` | wait for Ollama to unload |
| `JOB_TIMEOUT` | `15m` | wait for ComfyUI job | | `JOB_TIMEOUT` | `15m` | wait for ComfyUI job |
| `LLM_WAIT_TIMEOUT` | `10m` | max time an LLM request waits for the lock before 503 (wait mode) | | `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_REPO` | `https://git.rambossek.at/PUBLIC/gpu-turnstile` | repository to check for releases |
| `UPDATE_ASSET` | `gpu-turnstile.exe` | release asset to download | | `UPDATE_ASSET` | `gpu-turnstile.exe` | release asset to download |
Startup fails fast on unparsable values. Both upstreams are probed once at Startup fails fast on unparsable values and when neither consumer URL is
start (`/api/version`, `/system_stats`); failure is logged, not fatal. set. Enabled upstreams are probed once at start (`/api/version`,
`/system_stats`); failure is logged, not fatal.
## Native Windows deployment ## Native Windows deployment
@@ -184,7 +201,7 @@ Docker.
- `GET /healthz` on both listeners: 200 with JSON - `GET /healthz` on both listeners: 200 with JSON
`{"state":"idle|llm|image","llm_inflight":N,"image_pending":B}`. `{"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: dependency needed:
`gpu_turnstile_state{state="…"} 1`, `gpu_turnstile_llm_inflight`, `gpu_turnstile_state{state="…"} 1`, `gpu_turnstile_llm_inflight`,
`gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds` `gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
+48 -20
View File
@@ -180,6 +180,14 @@ func serviceCommand(configPath string, args []string) int {
return 0 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 { 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 // 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. // 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, "version", version,
"listen_ollama", cfg.ListenOllama, "listen_ollama", cfg.ListenOllama,
"listen_comfy", cfg.ListenComfy, "listen_comfy", cfg.ListenComfy,
"ollama_url", cfg.OllamaURL, "ollama_url", orDisabled(cfg.OllamaURL),
"comfy_url", cfg.ComfyURL, "comfy_url", orDisabled(cfg.ComfyURL),
"unload_timeout", cfg.UnloadTimeout, "unload_timeout", cfg.UnloadTimeout,
"job_timeout", cfg.JobTimeout, "job_timeout", cfg.JobTimeout,
"llm_wait_timeout", cfg.LLMWaitTimeout, "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) lk := lock.New(log)
ollamaClient, err := ollama.New(cfg.OllamaURL, log) // Each consumer is enabled by setting its URL; a disabled consumer gets
if err != nil { // no client, no listener and no probe.
return err 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) var comfyClient *comfy.Client
if err != nil { if cfg.ComfyURL != "" {
return err if comfyClient, err = comfy.New(cfg.ComfyURL, log); err != nil {
return err
}
} }
srv, err := proxy.New(proxy.Config{ 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 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) probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout)
if err := ollamaClient.Probe(probeCtx); err != nil { if ollamaClient != nil {
log.Warn("ollama probe failed", "url", cfg.OllamaURL, "err", err) if err := ollamaClient.Probe(probeCtx); err != nil {
log.Warn("ollama probe failed", "url", cfg.OllamaURL, "err", err)
}
} }
if err := comfyClient.Probe(probeCtx); err != nil { if comfyClient != nil {
log.Warn("comfy probe failed", "url", cfg.ComfyURL, "err", err) if err := comfyClient.Probe(probeCtx); err != nil {
log.Warn("comfy probe failed", "url", cfg.ComfyURL, "err", err)
}
} }
probeCancel() probeCancel()
ollamaSrv := &http.Server{Addr: cfg.ListenOllama, Handler: srv.OllamaHandler()} var servers []*http.Server
comfySrv := &http.Server{Addr: cfg.ListenComfy, Handler: srv.ComfyHandler()} 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) errCh := make(chan error, len(servers))
go func() { errCh <- ollamaSrv.ListenAndServe() }() for _, s := range servers {
go func() { errCh <- comfySrv.ListenAndServe() }() go func(s *http.Server) { errCh <- s.ListenAndServe() }(s)
}
if cfg.AutoUpdate { if cfg.AutoUpdate {
go updateLoop(ctx, cfg, log, lk, isService) 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) shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), cfg.ShutdownTimeout)
defer shutdownCancel() defer shutdownCancel()
ollamaSrv.Shutdown(shutdownCtx) for _, s := range servers {
comfySrv.Shutdown(shutdownCtx) s.Shutdown(shutdownCtx)
}
return nil return nil
} }
+2
View File
@@ -9,6 +9,8 @@ services:
image: git.rambossek.at/public/gpu-turnstile:v0.1.3 image: git.rambossek.at/public/gpu-turnstile:v0.1.3
restart: unless-stopped restart: unless-stopped
environment: 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: # Services on the Docker host itself:
OLLAMA_URL: http://host.docker.internal:11435 OLLAMA_URL: http://host.docker.internal:11435
COMFY_URL: http://host.docker.internal:8189 COMFY_URL: http://host.docker.internal:8189
+5 -3
View File
@@ -51,13 +51,12 @@ type Config struct {
} }
// Defaults returns the configuration used when neither the environment nor // 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 { func Defaults() Config {
return Config{ return Config{
ListenOllama: ":11434", ListenOllama: ":11434",
ListenComfy: ":8188", ListenComfy: ":8188",
OllamaURL: "http://127.0.0.1:11435",
ComfyURL: "http://127.0.0.1:8189",
UnloadTimeout: time.Minute, UnloadTimeout: time.Minute,
JobTimeout: 15 * time.Minute, JobTimeout: 15 * time.Minute,
LLMWaitTimeout: 10 * time.Minute, LLMWaitTimeout: 10 * time.Minute,
@@ -219,5 +218,8 @@ func Load(getenv func(string) string) (Config, error) {
default: default:
return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"") 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 return cfg, nil
} }
+19 -1
View File
@@ -8,13 +8,21 @@ import (
) )
func TestDefaults(t *testing.T) { 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 { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if cfg.ListenOllama != ":11434" || cfg.ListenComfy != ":8188" { if cfg.ListenOllama != ":11434" || cfg.ListenComfy != ":8188" {
t.Fatalf("listen addrs = %s %s", cfg.ListenOllama, cfg.ListenComfy) 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 { if cfg.UnloadTimeout != time.Minute || cfg.JobTimeout != 15*time.Minute {
t.Fatalf("timeouts = %v %v", cfg.UnloadTimeout, cfg.JobTimeout) 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) { func TestParseEnvFile(t *testing.T) {
input := `# comment input := `# comment
OLLAMA_URL=http://host:11435 OLLAMA_URL=http://host:11435
@@ -91,6 +106,9 @@ func TestLoadErrors(t *testing.T) {
if k == tc.key { if k == tc.key {
return tc.value return tc.value
} }
if k == "OLLAMA_URL" {
return "http://127.0.0.1:11435"
}
return "" return ""
}) })
if err == nil { if err == nil {
+52 -28
View File
@@ -102,24 +102,37 @@ type Server struct {
comfyProxy *httputil.ReverseProxy 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) { func New(cfg Config) (*Server, error) {
ollamaURL, err := url.Parse(cfg.OllamaURL) if cfg.OllamaURL == "" && cfg.ComfyURL == "" {
if err != nil || ollamaURL.Scheme == "" || ollamaURL.Host == "" { return nil, fmt.Errorf("at least one of OllamaURL or ComfyURL is required")
return nil, fmt.Errorf("invalid OLLAMA_URL %q", cfg.OllamaURL)
} }
comfyURL, err := url.Parse(cfg.ComfyURL) var ollamaURL, comfyURL *url.URL
if err != nil || comfyURL.Scheme == "" || comfyURL.Host == "" { if cfg.OllamaURL != "" {
return nil, fmt.Errorf("invalid COMFY_URL %q", cfg.ComfyURL) 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 log := cfg.Log
if log == nil { if log == nil {
log = slog.Default() log = slog.Default()
} }
if cfg.UnloadPollInterval > 0 { if cfg.UnloadPollInterval > 0 && cfg.Ollama != nil {
cfg.Ollama.PollInterval = cfg.UnloadPollInterval cfg.Ollama.PollInterval = cfg.UnloadPollInterval
} }
if cfg.HistoryPollInterval > 0 { if cfg.HistoryPollInterval > 0 && cfg.Comfy != nil {
cfg.Comfy.PollInterval = cfg.HistoryPollInterval cfg.Comfy.PollInterval = cfg.HistoryPollInterval
} }
freeTimeout := cfg.FreeTimeout freeTimeout := cfg.FreeTimeout
@@ -160,7 +173,7 @@ func New(cfg Config) (*Server, error) {
max: backoffMax, max: backoffMax,
log: log, log: log,
} }
return &Server{ s := &Server{
cfg: cfg, cfg: cfg,
log: log, log: log,
logWriter: cfg.LogWriter, logWriter: cfg.LogWriter,
@@ -172,9 +185,14 @@ func New(cfg Config) (*Server, error) {
busyMode: busyMode, busyMode: busyMode,
busyStatus: busyStatus, busyStatus: busyStatus,
busyRetryAfter: busyRetryAfter, busyRetryAfter: busyRetryAfter,
ollamaProxy: newReverseProxy(ollamaURL, retry, log.With("upstream", "ollama")), }
comfyProxy: newReverseProxy(comfyURL, retry, log.With("upstream", "comfy")), if ollamaURL != nil {
}, 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 // 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. // ComfyHandler serves the ComfyUI-facing listener.
func (s *Server) ComfyHandler() http.Handler { func (s *Server) ComfyHandler() http.Handler {
return s.logRequests("comfy", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 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) s.writeHealthz(w)
return return
case "/metrics":
s.writeMetrics(w)
return
} }
if r.Method == http.MethodPost && r.URL.Path == "/prompt" { if r.Method == http.MethodPost && r.URL.Path == "/prompt" {
s.handlePrompt(w, r) 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()) s.cfg.Metrics.ObserveLockWait("image", time.Since(start).Seconds())
log.Info("image lock acquired") log.Info("image lock acquired")
uctx, ucancel := context.WithTimeout(r.Context(), s.cfg.UnloadTimeout) if s.cfg.Ollama != nil {
elapsed, uerr := s.cfg.Ollama.UnloadAll(uctx) uctx, ucancel := context.WithTimeout(r.Context(), s.cfg.UnloadTimeout)
ucancel() elapsed, uerr := s.cfg.Ollama.UnloadAll(uctx)
s.cfg.Metrics.ObserveUnload(elapsed.Seconds()) ucancel()
switch { s.cfg.Metrics.ObserveUnload(elapsed.Seconds())
case r.Context().Err() != nil: switch {
s.cfg.Lock.ReleaseImage() case r.Context().Err() != nil:
return s.cfg.Lock.ReleaseImage()
case uerr != nil: return
// Degrade, don't fail the user's request on a misbehaving neighbour. case uerr != nil:
log.Warn("ollama unload incomplete; continuing", "err", uerr) // Degrade, don't fail the user's request on a misbehaving neighbour.
default: log.Warn("ollama unload incomplete; continuing", "err", uerr)
log.Info("ollama models unloaded", "seconds", elapsed.Seconds()) default:
log.Info("ollama models unloaded", "seconds", elapsed.Seconds())
}
} }
cw := &captureWriter{ResponseWriter: w, status: http.StatusOK, limit: s.captureLimit} cw := &captureWriter{ResponseWriter: w, status: http.StatusOK, limit: s.captureLimit}
@@ -585,7 +609,7 @@ func (s *Server) finishImageJob(promptID string) {
s.cfg.Lock.ReleaseImage() s.cfg.Lock.ReleaseImage()
log.Info("image lock released") 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 { if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle {
wctx, wcancel := context.WithTimeout(context.Background(), s.warmTimeout) wctx, wcancel := context.WithTimeout(context.Background(), s.warmTimeout)
if err := s.cfg.Ollama.Warm(wctx, s.cfg.WarmModel); err != nil { if err := s.cfg.Ollama.Warm(wctx, s.cfg.WarmModel); err != nil {
+79
View File
@@ -451,3 +451,82 @@ func TestLLMBusyWaitTimeoutRetryAfter(t *testing.T) {
t.Fatalf("Retry-After = %q", got) 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")
}
}