PUBLIC/gpu-turnstile
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 6 Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
internal/proxy/proxy.go
+52 -28
View File
@@ -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 {