8 Commits
Author SHA1 Message Date
mram 5f6a22c2cf Pin compose example to v0.1.2
ci / test (push) Successful in 47s
ci / docker (push) Successful in 1m5s
2026-09-20 20:07:38 +02:00
mram 42e1386811 Extend retry backoff; rework logging (LOGLEVEL, request lines, colors)
Backoff now covers all dial-phase errors (refused, timeout, DNS), TLS
handshake failures, and 5xx responses with replayable bodies; streamed
POSTs are never replayed to avoid duplicate work. First retry of an
episode logs at WARN, subsequent attempts at INFO.

LOGLEVEL (LOG_LEVEL kept as alias) now defaults to warn: startup logs
version plus every setting; INFO adds one line per incoming request and
per response with status/duration, ANSI-colored in text mode (bypasses
slog's escaping so colors render in docker compose logs); NO_COLOR or
LOG_FORMAT=json disables colors.
2026-09-20 20:00:57 +02:00
mram 20d5439b5f Retry upstream connection-refused with exponential backoff
A refused dial (service down/restarting) is retried with a wait that
doubles from BACKOFF_INITIAL (1s) up to BACKOFF_MAX (60s) until the
upstream answers or the client disconnects. Handles the Windows WSA
errno (10061) as well as POSIX ECONNREFUSED. compose.yaml.example now
uses host.docker.internal like the working local deployment.
2026-09-20 19:45:48 +02:00
mram 3aaa5d80a9 CI: lowercase image repository name; compose example too
ci / test (push) Successful in 49s
ci / docker (push) Successful in 1m7s
Docker registry names must be lowercase, so PUBLIC/gpu-turnstile was
rejected by buildx. The workflow now lowercases gitea.repository.
2026-09-20 19:25:37 +02:00
mram cef845c2e3 Add compose.yaml.example; ignore local compose files 2026-09-20 19:24:26 +02:00
mram cf9be55a6c Make poll intervals and operational timeouts configurable
ci / test (push) Successful in 47s
ci / docker (push) Failing after 30s
New env vars: UNLOAD_POLL_INTERVAL, HISTORY_POLL_INTERVAL, PROBE_TIMEOUT,
FREE_TIMEOUT, WARM_TIMEOUT, SHUTDOWN_TIMEOUT, PROMPT_CAPTURE_LIMIT.
Defaults unchanged; invalid values fail fast at startup.
2026-09-20 19:21:25 +02:00
mram 0f950f3134 CI: use gitea.actor as registry login username
ci / test (push) Successful in 48s
ci / docker (push) Skipped
Gitea authenticates registry pushes by token alone; a single
REGISTRY_TOKEN secret suffices.
2026-09-20 18:55:46 +02:00
mram 73fb8ac3ad CI: log in to registry with REGISTRY_USERNAME/REGISTRY_TOKEN secrets
ci / test (push) Successful in 47s
ci / docker (push) Skipped
The automatic GITEA_TOKEN has no write:package scope, so docker login to
git.rambossek.at failed with unauthorized.
2026-09-20 18:24:54 +02:00
8 changed files with 679 additions and 53 deletions
+9 -3
View File
@@ -37,13 +37,19 @@ jobs:
exit 1
fi
- name: Compute lowercase image name
id: meta
run: |
REPO=git.rambossek.at/$(echo "${{ gitea.repository }}" | tr '[:upper:]' '[:lower:]')
echo "image=$REPO" >> "$GITHUB_OUTPUT"
- uses: docker/setup-buildx-action@v3
- uses: docker/login-action@v3
with:
registry: git.rambossek.at
username: ${{ gitea.actor }}
password: ${{ secrets.GITEA_TOKEN }}
password: ${{ secrets.REGISTRY_TOKEN }}
- uses: docker/build-push-action@v6
with:
@@ -52,5 +58,5 @@ jobs:
build-args: |
VERSION=${{ gitea.ref_name }}
tags: |
git.rambossek.at/${{ gitea.repository }}:${{ gitea.ref_name }}
git.rambossek.at/${{ gitea.repository }}:latest
${{ steps.meta.outputs.image }}:${{ gitea.ref_name }}
${{ steps.meta.outputs.image }}:latest
+2
View File
@@ -1,2 +1,4 @@
/gpu-turnstile
/gpu-turnstile.exe
/compose.yaml
/compose.yml
+20 -1
View File
@@ -41,8 +41,17 @@ startup.
| `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 |
| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` works as an alias |
| `LOG_FORMAT` | `text` | `json` for structured JSON logs |
| `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
| `PROBE_TIMEOUT` | `5s` | Startup probe of both upstreams |
| `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job |
| `WARM_TIMEOUT` | `2m` | Warm-model reload after an image job |
| `SHUTDOWN_TIMEOUT` | `10s` | Graceful shutdown on SIGINT/SIGTERM |
| `BACKOFF_INITIAL` | `1s` | First retry wait when an upstream refuses a connection |
| `BACKOFF_MAX` | `60s` | Cap for the exponential retry backoff |
| `PROMPT_CAPTURE_LIMIT` | `65536` | Bytes of the `/prompt` response buffered to find `prompt_id` (pass-through is unaffected) |
## Observability
@@ -51,6 +60,12 @@ startup.
`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`.
- Logs: startup logs the version and every setting (visible even at the
default `warn` level). With `LOGLEVEL=info` or `debug`, every request
logs a `-->` incoming line and a `<--` response line with status and
duration — ANSI-colored (cyan incoming; green/yellow/red by status
class) in text mode, which renders in `docker compose logs` on Windows
Terminal. Set `NO_COLOR` to disable colors.
## Build and run
@@ -73,6 +88,10 @@ runs `go vet` and `go test -race`, and pushing a semantic-version tag
`git.rambossek.at/<owner>/gpu-turnstile:vX.Y.Z` (and updates `:latest`).
No images are built from branches.
The registry login needs one repository secret (Settings → Actions →
Secrets): `REGISTRY_TOKEN` — an access token with `write:package` scope.
The automatic `GITEA_TOKEN` cannot push packages.
## Development
```sh
+36 -7
View File
@@ -81,12 +81,14 @@ ComfyUI listener (`:8188` → `COMFY_URL`):
2. Unload Ollama: `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 500 ms until empty or `UNLOAD_TIMEOUT`. On timeout: log and
`/api/ps` every `UNLOAD_POLL_INTERVAL` (default 500 ms) until empty or
`UNLOAD_TIMEOUT`. On timeout: log and
continue (degrade, don't fail the user's request).
3. Forward the original request body to ComfyUI `/prompt`, return status,
headers and body to the caller unchanged, flush.
4. If the response is 200 and contains `prompt_id`: in a goroutine, poll
`GET /history/<prompt_id>` every 1 s until the entry has
`GET /history/<prompt_id>` every `HISTORY_POLL_INTERVAL` (default 1 s)
until the entry has
`status.completed == true`, `status.status_str == "error"`, or
`JOB_TIMEOUT`. Then `POST /free {"unload_models":true,"free_memory":true}`.
Then release the image lock.
@@ -110,7 +112,16 @@ load time. Off by default.
| `JOB_TIMEOUT` | `15m` | wait for ComfyUI job |
| `LLM_WAIT_TIMEOUT` | `10m` | max time an LLM request waits for the lock before 503 |
| `WARM_MODEL` | `` | optional model to reload after an image job |
| `LOG_LEVEL` | `info` | `debug` logs every lock transition |
| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` is accepted as an alias |
| `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
| `PROBE_TIMEOUT` | `5s` | startup probe of both upstreams |
| `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job |
| `WARM_TIMEOUT` | `2m` | warm-model reload after an image job |
| `SHUTDOWN_TIMEOUT` | `10s` | graceful shutdown on SIGINT/SIGTERM |
| `BACKOFF_INITIAL` | `1s` | first retry wait when an upstream refuses a connection |
| `BACKOFF_MAX` | `60s` | cap for the exponential retry backoff |
| `PROMPT_CAPTURE_LIMIT` | `65536` | bytes of the `/prompt` response buffered to find `prompt_id` (pass-through is unaffected) |
Startup fails fast on unparsable values. Both upstreams are probed once at
start (`/api/version`, `/system_stats`); failure is logged, not fatal.
@@ -125,7 +136,13 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal.
`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`.
state transition and per image job phase with `prompt_id`. Startup logs
the version and every setting (visible even at the default `warn`
level). With `LOGLEVEL=info` or `debug`, every request logs a `-->`
incoming line and a `<--` response line with status and duration —
ANSI-colored (cyan incoming; green/yellow/red by status class) in text
mode, which renders in `docker compose logs` on Windows Terminal. Set
`NO_COLOR` to disable colors.
## Edge cases to handle
@@ -137,6 +154,14 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal.
`JOB_TIMEOUT` releases the lock; log at warn.
- Ollama unreachable during unload: continue with the image job; the whole
point is not to block users on a misbehaving neighbour.
- Upstream unreachable while proxying (connection refused, dial timeout,
DNS failure, TLS handshake error): retry with exponential backoff —
`BACKOFF_INITIAL`, doubling per attempt, capped at `BACKOFF_MAX` — until
the upstream answers or the client disconnects. These are safe to retry:
the request never reached the upstream application. 5xx responses are
retried the same way, but only when the request body can be replayed
(GETs, or bodies with `GetBody`); streamed POSTs are never replayed to
avoid duplicate work such as a double-enqueued ComfyUI prompt.
- `POST /prompt` with a body that ComfyUI rejects (400): lock released
immediately, body passed back.
- Websocket `/ws` connections are long-lived and never take the lock.
@@ -185,8 +210,12 @@ are new.
available in the runner image
2. on a version tag only (`vX.Y.Z`, enforced): build the image with buildx
and push it to the Gitea registry
`git.rambossek.at/<owner>/gpu-turnstile` tagged `:<tag>` and `:latest`,
using the workflow token (`${{ secrets.GITEA_TOKEN }}` / `gitea.actor`)
`git.rambossek.at/<owner>/gpu-turnstile` tagged `:<tag>` and `:latest`
(the repository path is lowercased in the workflow; Docker registry
names must be lowercase).
Login uses the repo secret `REGISTRY_TOKEN` (an access token with
`write:package` scope) because the automatic `GITEA_TOKEN` cannot push
packages; the username is just `gitea.actor`.
- Release: a git tag `vX.Y.Z` produces the versioned image; the Open WebUI
compose pins that tag. No images are built from branches.
@@ -194,7 +223,7 @@ are new.
```yaml
gpu-turnstile:
image: git.rambossek.at/<owner>/gpu-turnstile:v0.1.0
image: git.rambossek.at/<owner>/gpu-turnstile:v0.1.0 # owner lowercased, e.g. "public"
environment:
OLLAMA_URL: http://<workstation-ip>:11435
COMFY_URL: http://<workstation-ip>:8189
+75 -7
View File
@@ -10,6 +10,7 @@ import (
"net/http"
"os"
"os/signal"
"strconv"
"strings"
"syscall"
"time"
@@ -32,6 +33,17 @@ type config struct {
unloadTimeout time.Duration
jobTimeout time.Duration
llmWaitTimeout time.Duration
unloadPollInterval time.Duration
historyPollInterval time.Duration
probeTimeout time.Duration
freeTimeout time.Duration
warmTimeout time.Duration
shutdownTimeout time.Duration
backoffInitial time.Duration
backoffMax time.Duration
promptCaptureLimit int64
warmModel string
logLevel slog.Level
logJSON bool
@@ -59,7 +71,18 @@ func loadConfig(getenv func(string) string) (config, error) {
unloadTimeout: time.Minute,
jobTimeout: 15 * time.Minute,
llmWaitTimeout: 10 * time.Minute,
logLevel: slog.LevelInfo,
unloadPollInterval: 500 * time.Millisecond,
historyPollInterval: time.Second,
probeTimeout: 5 * time.Second,
freeTimeout: 30 * time.Second,
warmTimeout: 2 * time.Minute,
shutdownTimeout: 10 * time.Second,
backoffInitial: time.Second,
backoffMax: time.Minute,
promptCaptureLimit: 64 * 1024,
logLevel: slog.LevelWarn,
}
for _, e := range []struct {
name string
@@ -82,15 +105,35 @@ func loadConfig(getenv func(string) string) (config, error) {
{"UNLOAD_TIMEOUT", &cfg.unloadTimeout},
{"JOB_TIMEOUT", &cfg.jobTimeout},
{"LLM_WAIT_TIMEOUT", &cfg.llmWaitTimeout},
{"UNLOAD_POLL_INTERVAL", &cfg.unloadPollInterval},
{"HISTORY_POLL_INTERVAL", &cfg.historyPollInterval},
{"PROBE_TIMEOUT", &cfg.probeTimeout},
{"FREE_TIMEOUT", &cfg.freeTimeout},
{"WARM_TIMEOUT", &cfg.warmTimeout},
{"SHUTDOWN_TIMEOUT", &cfg.shutdownTimeout},
{"BACKOFF_INITIAL", &cfg.backoffInitial},
{"BACKOFF_MAX", &cfg.backoffMax},
} {
if err := envDuration(getenv, e.name, e.dst); err != nil {
return cfg, err
}
}
if v := getenv("LOG_LEVEL"); v != "" {
if v := getenv("PROMPT_CAPTURE_LIMIT"); v != "" {
n, err := strconv.ParseInt(v, 10, 64)
if err != nil || n < 0 {
return cfg, fmt.Errorf("PROMPT_CAPTURE_LIMIT: must be a non-negative integer (bytes)")
}
cfg.promptCaptureLimit = n
}
// LOGLEVEL is the canonical spelling; LOG_LEVEL is kept as an alias.
logLevelValue := getenv("LOGLEVEL")
if logLevelValue == "" {
logLevelValue = getenv("LOG_LEVEL")
}
if logLevelValue != "" {
var level slog.Level
if err := level.UnmarshalText([]byte(v)); err != nil {
return cfg, fmt.Errorf("LOG_LEVEL: %w", err)
if err := level.UnmarshalText([]byte(logLevelValue)); err != nil {
return cfg, fmt.Errorf("LOGLEVEL: %w", err)
}
cfg.logLevel = level
}
@@ -119,12 +162,29 @@ func main() {
log := slog.New(handler)
slog.SetDefault(log)
log.Info("starting gpu-turnstile",
// 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.
log.Log(context.Background(), slog.LevelWarn, "starting gpu-turnstile",
"version", version,
"listen_ollama", cfg.listenOllama,
"listen_comfy", cfg.listenComfy,
"ollama_url", cfg.ollamaURL,
"comfy_url", cfg.comfyURL,
"unload_timeout", cfg.unloadTimeout,
"job_timeout", cfg.jobTimeout,
"llm_wait_timeout", cfg.llmWaitTimeout,
"unload_poll_interval", cfg.unloadPollInterval,
"history_poll_interval", cfg.historyPollInterval,
"probe_timeout", cfg.probeTimeout,
"free_timeout", cfg.freeTimeout,
"warm_timeout", cfg.warmTimeout,
"shutdown_timeout", cfg.shutdownTimeout,
"backoff_initial", cfg.backoffInitial,
"backoff_max", cfg.backoffMax,
"prompt_capture_limit", cfg.promptCaptureLimit,
"warm_model", cfg.warmModel,
"log_level", cfg.logLevel,
"log_format", map[bool]string{true: "json", false: "text"}[cfg.logJSON],
)
ollamaClient, err := ollama.New(cfg.ollamaURL, log)
@@ -146,9 +206,17 @@ func main() {
Comfy: comfyClient,
Metrics: metrics.New(),
Log: log,
LogColor: !cfg.logJSON && os.Getenv("NO_COLOR") == "",
LLMWaitTimeout: cfg.llmWaitTimeout,
UnloadTimeout: cfg.unloadTimeout,
JobTimeout: cfg.jobTimeout,
UnloadPollInterval: cfg.unloadPollInterval,
HistoryPollInterval: cfg.historyPollInterval,
FreeTimeout: cfg.freeTimeout,
WarmTimeout: cfg.warmTimeout,
BackoffInitial: cfg.backoffInitial,
BackoffMax: cfg.backoffMax,
PromptCaptureLimit: cfg.promptCaptureLimit,
WarmModel: cfg.warmModel,
})
if err != nil {
@@ -157,7 +225,7 @@ func main() {
}
// Probe both upstreams once; failure is logged, not fatal.
probeCtx, probeCancel := context.WithTimeout(context.Background(), 5*time.Second)
probeCtx, probeCancel := context.WithTimeout(context.Background(), cfg.probeTimeout)
if err := ollamaClient.Probe(probeCtx); err != nil {
log.Warn("ollama probe failed", "url", cfg.ollamaURL, "err", err)
}
@@ -186,7 +254,7 @@ func main() {
log.Info("shutting down")
}
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), cfg.shutdownTimeout)
defer shutdownCancel()
ollamaSrv.Shutdown(shutdownCtx)
comfySrv.Shutdown(shutdownCtx)
+31
View File
@@ -0,0 +1,31 @@
# Example deployment for gpu-turnstile. Copy to compose.yaml and adjust.
#
# gpu-turnstile listens on the ports the services normally use; the actual
# Ollama and ComfyUI instances run one port higher (11435 / 8189) and must
# bind 0.0.0.0 so the container can reach them (OLLAMA_HOST=0.0.0.0:11435,
# ComfyUI --listen 0.0.0.0 --port 8189).
services:
gpu-turnstile:
image: git.rambossek.at/public/gpu-turnstile:v0.1.2
restart: unless-stopped
environment:
# Services on the Docker host itself:
OLLAMA_URL: http://host.docker.internal:11435
COMFY_URL: http://host.docker.internal:8189
# Services on another machine: use its LAN IP instead, e.g.
# OLLAMA_URL: http://192.168.1.10:11435
# COMFY_URL: http://192.168.1.10:8189
# UNLOAD_TIMEOUT: 60s
# JOB_TIMEOUT: 15m
# LLM_WAIT_TIMEOUT: 10m
# WARM_MODEL: qwen3:14b
# LOG_LEVEL: info
# LOG_FORMAT: json
ports:
- "11434:11434" # LiteLLM api_base -> http://gpu-turnstile:11434
- "8188:8188" # Open WebUI COMFYUI_BASE_URL -> http://gpu-turnstile:8188
networks: [internal]
networks:
internal:
external: true
+280 -16
View File
@@ -4,15 +4,22 @@
package proxy
import (
"bufio"
"bytes"
"context"
"crypto/tls"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net"
"net/http"
"net/http/httputil"
"net/url"
"os"
"strconv"
"strings"
"time"
"gpu-turnstile/internal/comfy"
@@ -21,10 +28,10 @@ import (
"gpu-turnstile/internal/ollama"
)
// captureLimit bounds how much of a /prompt response body is buffered while
// looking for prompt_id. The body still passes through to the client
// unchanged regardless of size.
const captureLimit = 64 * 1024
// defaultCaptureLimit bounds how much of a /prompt response body is
// buffered while looking for prompt_id. The body still passes through to
// the client unchanged regardless of size.
const defaultCaptureLimit = 64 * 1024
// Config wires a Server.
type Config struct {
@@ -37,9 +44,32 @@ type Config struct {
Metrics *metrics.Metrics
Log *slog.Logger
// LogColor enables ANSI colors in per-request log lines. Ignored when
// the log level is above INFO (request lines are not emitted at all).
LogColor bool
LLMWaitTimeout time.Duration
UnloadTimeout time.Duration
JobTimeout time.Duration
// BackoffInitial and BackoffMax control the exponential retry backoff
// when an upstream refuses a connection: the wait doubles from
// BackoffInitial up to BackoffMax between attempts. Zero selects the
// defaults (1s / 60s).
BackoffInitial time.Duration
BackoffMax time.Duration
// UnloadPollInterval and HistoryPollInterval override the clients'
// /api/ps and /history poll intervals when > 0.
UnloadPollInterval time.Duration
HistoryPollInterval time.Duration
// FreeTimeout and WarmTimeout bound the /free call and the warm-model
// reload; zero selects the defaults.
FreeTimeout time.Duration
WarmTimeout time.Duration
// PromptCaptureLimit overrides defaultCaptureLimit when > 0.
PromptCaptureLimit int64
WarmModel string
}
@@ -47,6 +77,11 @@ type Config struct {
type Server struct {
cfg Config
log *slog.Logger
freeTimeout time.Duration
warmTimeout time.Duration
captureLimit int64
backoffInitial time.Duration
backoffMax time.Duration
ollamaProxy *httputil.ReverseProxy
comfyProxy *httputil.ReverseProxy
@@ -66,16 +101,160 @@ func New(cfg Config) (*Server, error) {
if log == nil {
log = slog.Default()
}
if cfg.UnloadPollInterval > 0 {
cfg.Ollama.PollInterval = cfg.UnloadPollInterval
}
if cfg.HistoryPollInterval > 0 {
cfg.Comfy.PollInterval = cfg.HistoryPollInterval
}
freeTimeout := cfg.FreeTimeout
if freeTimeout <= 0 {
freeTimeout = 30 * time.Second
}
warmTimeout := cfg.WarmTimeout
if warmTimeout <= 0 {
warmTimeout = 2 * time.Minute
}
captureLimit := cfg.PromptCaptureLimit
if captureLimit <= 0 {
captureLimit = defaultCaptureLimit
}
backoffInitial := cfg.BackoffInitial
if backoffInitial <= 0 {
backoffInitial = time.Second
}
backoffMax := cfg.BackoffMax
if backoffMax <= 0 {
backoffMax = time.Minute
}
retry := &retryTransport{
base: http.DefaultTransport,
initial: backoffInitial,
max: backoffMax,
log: log,
}
return &Server{
cfg: cfg,
log: log,
ollamaProxy: newReverseProxy(ollamaURL, log.With("upstream", "ollama")),
comfyProxy: newReverseProxy(comfyURL, log.With("upstream", "comfy")),
freeTimeout: freeTimeout,
warmTimeout: warmTimeout,
captureLimit: captureLimit,
backoffInitial: backoffInitial,
backoffMax: backoffMax,
ollamaProxy: newReverseProxy(ollamaURL, retry, log.With("upstream", "ollama")),
comfyProxy: newReverseProxy(comfyURL, retry, log.With("upstream", "comfy")),
}, nil
}
func newReverseProxy(target *url.URL, log *slog.Logger) *httputil.ReverseProxy {
// retryTransport retries requests whose failure means the upstream never
// saw them — any dial-phase error (connection refused, dial timeout, DNS
// failure), TLS handshake errors — plus 5xx responses when the request
// body can be replayed (GETs and requests with GetBody set). The wait
// doubles from initial up to max between attempts. The loop runs until
// the request succeeds, fails in a non-retryable way, or the client's
// context is cancelled.
type retryTransport struct {
base http.RoundTripper
initial time.Duration
max time.Duration
log *slog.Logger
}
// shouldRetry reports whether a RoundTrip error means the request never
// reached the upstream application and is therefore safe to send again.
func shouldRetry(err error) bool {
// Dial-phase failures: refused, timeout, unreachable, DNS (wrapped).
var opErr *net.OpError
if errors.As(err, &opErr) && opErr.Op == "dial" {
return true
}
var dnsErr *net.DNSError
if errors.As(err, &dnsErr) {
return true
}
// TLS handshake failures: the HTTP request was never written.
var recordErr tls.RecordHeaderError
if errors.As(err, &recordErr) {
return true
}
var certErr *tls.CertificateVerificationError
if errors.As(err, &certErr) {
return true
}
var alertErr tls.AlertError
if errors.As(err, &alertErr) {
return true
}
return false
}
// replayable reports whether the request body can be sent again. Bodies
// streamed from the client (GetBody == nil) cannot, so 5xx responses to
// POSTs are not retried: the upstream may have partially processed them,
// and re-sending could duplicate work (e.g. a second ComfyUI prompt).
func replayable(req *http.Request) bool {
return req.Body == nil || req.Body == http.NoBody || req.GetBody != nil
}
func (t *retryTransport) logAt(level slog.Level, msg string, args ...any) {
if t.log != nil && t.log.Enabled(context.Background(), level) {
t.log.Log(context.Background(), level, msg, args...)
}
}
func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
wait := t.initial
attempt := 0
for {
resp, err := t.base.RoundTrip(req)
switch {
case err != nil && !shouldRetry(err):
return nil, err
case err == nil && (resp.StatusCode < 500 || !replayable(req)):
return resp, nil
}
// Retryable failure: a transport error, or a 5xx response.
var reason string
if err != nil {
reason = err.Error()
} else {
reason = resp.Status
io.Copy(io.Discard, resp.Body)
resp.Body.Close()
if req.GetBody != nil {
if body, berr := req.GetBody(); berr == nil {
req.Body = body
}
}
}
attempt++
// One WARN per outage episode; subsequent attempts at INFO.
level := slog.LevelInfo
if attempt == 1 {
level = slog.LevelWarn
}
t.logAt(level, "upstream unavailable; retrying with backoff",
"path", req.URL.Path, "reason", reason, "attempt", attempt, "retry_in", wait)
select {
case <-req.Context().Done():
if err != nil {
return nil, err
}
return nil, req.Context().Err()
case <-time.After(wait):
}
wait *= 2
if wait > t.max {
wait = t.max
}
}
}
func newReverseProxy(target *url.URL, transport http.RoundTripper, log *slog.Logger) *httputil.ReverseProxy {
return &httputil.ReverseProxy{
Transport: transport,
Rewrite: func(pr *httputil.ProxyRequest) {
pr.SetURL(target)
pr.SetXForwarded()
@@ -106,6 +285,90 @@ func (s *Server) writeMetrics(w http.ResponseWriter) {
s.cfg.Metrics.Render(w, string(state), n, pending)
}
// ANSI colors for per-request log lines.
const (
ansiReset = "\x1b[0m"
ansiCyan = "\x1b[36m"
ansiGreen = "\x1b[32m"
ansiYellow = "\x1b[33m"
ansiRed = "\x1b[31m"
)
// statusRecorder remembers the response status while passing everything
// through, including streaming flushes and websocket hijacks.
type statusRecorder struct {
http.ResponseWriter
status int
}
func (r *statusRecorder) WriteHeader(code int) {
r.status = code
r.ResponseWriter.WriteHeader(code)
}
func (r *statusRecorder) Flush() {
if f, ok := r.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
func (r *statusRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
h, ok := r.ResponseWriter.(http.Hijacker)
if !ok {
return nil, nil, errors.New("response writer does not support hijacking")
}
return h.Hijack()
}
func (r *statusRecorder) Unwrap() http.ResponseWriter { return r.ResponseWriter }
// reqLine emits one request log line. With color enabled slog cannot be
// used: its text handler escapes the ANSI sequences, so the line is
// written to stderr directly in the same key=value shape. Without color
// it is a plain slog INFO line.
func (s *Server) reqLine(log *slog.Logger, code, line string, attrs ...any) {
if !s.cfg.LogColor {
log.Info(line, attrs...)
return
}
var sb strings.Builder
sb.WriteString("time=" + time.Now().Format("2006-01-02T15:04:05.000Z07:00") + " level=INFO ")
sb.WriteString(code + line + ansiReset)
for i := 0; i+1 < len(attrs); i += 2 {
fmt.Fprintf(&sb, " %v=%v", attrs[i], attrs[i+1])
}
sb.WriteByte('\n')
os.Stderr.WriteString(sb.String())
}
// logRequests logs one line per incoming request and one per completed
// response at INFO level, colored when enabled: cyan "-->" for incoming,
// green/yellow/red "<--" for responses by status class. At log levels
// above INFO it is a pass-through.
func (s *Server) logRequests(listener string, next http.Handler) http.Handler {
log := s.log.With("listener", listener)
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if !log.Enabled(r.Context(), slog.LevelInfo) {
next.ServeHTTP(w, r)
return
}
start := time.Now()
s.reqLine(log, ansiCyan, "--> "+r.Method+" "+r.URL.RequestURI(),
"listener", listener, "remote", r.RemoteAddr)
rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
next.ServeHTTP(rec, r)
code := ansiGreen
switch {
case rec.status >= 500:
code = ansiRed
case rec.status >= 400:
code = ansiYellow
}
s.reqLine(log, code, "<-- "+strconv.Itoa(rec.status)+" "+r.Method+" "+r.URL.RequestURI(),
"listener", listener, "ms", time.Since(start).Milliseconds())
})
}
// llmPaths are the Ollama endpoints that load models into VRAM and therefore
// take the LLM lock. Everything else passes through unlocked.
var llmPaths = map[string]bool{
@@ -124,7 +387,7 @@ func isLLMRequest(r *http.Request) bool {
// OllamaHandler serves the Ollama-facing listener.
func (s *Server) OllamaHandler() http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
return s.logRequests("ollama", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/healthz":
s.writeHealthz(w)
@@ -151,12 +414,12 @@ func (s *Server) OllamaHandler() http.Handler {
}
defer s.cfg.Lock.ReleaseLLM()
s.ollamaProxy.ServeHTTP(w, r)
})
}))
}
// ComfyHandler serves the ComfyUI-facing listener.
func (s *Server) ComfyHandler() http.Handler {
return 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" {
s.writeHealthz(w)
return
@@ -166,15 +429,16 @@ func (s *Server) ComfyHandler() http.Handler {
return
}
s.comfyProxy.ServeHTTP(w, r)
})
}))
}
// captureWriter passes the response through unchanged while recording the
// status code and the first captureLimit bytes of the body.
// status code and the first limit bytes of the body.
type captureWriter struct {
http.ResponseWriter
status int
buf bytes.Buffer
limit int64
}
func (w *captureWriter) WriteHeader(code int) {
@@ -183,7 +447,7 @@ func (w *captureWriter) WriteHeader(code int) {
}
func (w *captureWriter) Write(p []byte) (int, error) {
if w.buf.Len() < captureLimit {
if int64(w.buf.Len()) < w.limit {
w.buf.Write(p)
}
return w.ResponseWriter.Write(p)
@@ -228,7 +492,7 @@ func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
log.Info("ollama models unloaded", "seconds", elapsed.Seconds())
}
cw := &captureWriter{ResponseWriter: w, status: http.StatusOK}
cw := &captureWriter{ResponseWriter: w, status: http.StatusOK, limit: s.captureLimit}
s.comfyProxy.ServeHTTP(cw, r)
var accepted struct {
@@ -262,7 +526,7 @@ func (s *Server) finishImageJob(promptID string) {
log.Info("image job completed")
}
freeCtx, freeCancel := context.WithTimeout(context.Background(), 30*time.Second)
freeCtx, freeCancel := context.WithTimeout(context.Background(), s.freeTimeout)
if err := s.cfg.Comfy.Free(freeCtx); err != nil {
log.Warn("failed to free ComfyUI models", "err", err)
}
@@ -273,7 +537,7 @@ func (s *Server) finishImageJob(promptID string) {
if s.cfg.WarmModel != "" {
if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle {
wctx, wcancel := context.WithTimeout(context.Background(), 2*time.Minute)
wctx, wcancel := context.WithTimeout(context.Background(), s.warmTimeout)
if err := s.cfg.Ollama.Warm(wctx, s.cfg.WarmModel); err != nil {
log.Warn("warm model reload failed", "model", s.cfg.WarmModel, "err", err)
} else {
+207
View File
@@ -0,0 +1,207 @@
package proxy
import (
"context"
"crypto/tls"
"io"
"net"
"net/http"
"strings"
"syscall"
"testing"
"time"
)
// stubTransport fails with ECONNREFUSED for the first fails requests, then
// returns a 200 response.
type stubTransport struct {
fails int
calls int
}
func (s *stubTransport) RoundTrip(req *http.Request) (*http.Response, error) {
s.calls++
if s.calls <= s.fails {
return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
}
return &http.Response{
StatusCode: 200,
Body: io.NopCloser(strings.NewReader("ok")),
Header: make(http.Header),
}, nil
}
func TestRetryTransportBackoff(t *testing.T) {
st := &stubTransport{fails: 3}
rt := &retryTransport{
base: st,
initial: 10 * time.Millisecond,
max: 25 * time.Millisecond,
}
start := time.Now()
req, _ := http.NewRequest(http.MethodGet, "http://upstream/api/version", nil)
resp, err := rt.RoundTrip(req)
if err != nil {
t.Fatal(err)
}
resp.Body.Close()
if st.calls != 4 {
t.Fatalf("calls = %d, want 4", st.calls)
}
// Waits: 10ms + 20ms + 25ms (capped) = 55ms minimum.
elapsed := time.Since(start)
if elapsed < 50*time.Millisecond {
t.Fatalf("elapsed = %v, want >= ~55ms of backoff", elapsed)
}
if elapsed > 5*time.Second {
t.Fatalf("elapsed = %v, suspiciously long", elapsed)
}
}
func TestRetryTransportNonRefusedErrorNotRetried(t *testing.T) {
rt := &retryTransport{
base: &stubTransport{fails: 0},
initial: time.Millisecond,
max: time.Millisecond,
}
req, _ := http.NewRequest(http.MethodGet, "http://upstream/", nil)
resp, err := rt.RoundTrip(req)
if err != nil || resp.StatusCode != 200 {
t.Fatalf("resp=%v err=%v", resp, err)
}
}
type alwaysRefused struct{ calls int }
func (a *alwaysRefused) RoundTrip(*http.Request) (*http.Response, error) {
a.calls++
return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
}
func TestRetryTransportContextCancel(t *testing.T) {
ar := &alwaysRefused{}
rt := &retryTransport{base: ar, initial: time.Second, max: time.Second}
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://upstream/", nil)
start := time.Now()
_, err := rt.RoundTrip(req)
if err == nil {
t.Fatal("expected error after context cancellation")
}
if time.Since(start) > 2*time.Second {
t.Fatal("retry loop did not stop on context cancellation")
}
}
func TestRetryViaReverseProxy(t *testing.T) {
// Point the proxy at a port nothing listens on; the request should be
// retried (not instantly 502) until the client context ends.
srv, err := New(Config{
OllamaURL: "http://127.0.0.1:1",
ComfyURL: "http://127.0.0.1:1",
Lock: nil,
Metrics: nil,
BackoffInitial: 10 * time.Millisecond,
BackoffMax: 20 * time.Millisecond,
})
if err != nil {
t.Fatal(err)
}
_ = srv // construction must not panic with minimal config
rt := &retryTransport{base: http.DefaultTransport, initial: 10 * time.Millisecond, max: 20 * time.Millisecond}
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
defer cancel()
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://127.0.0.1:1/", nil)
_, err = rt.RoundTrip(req)
if err == nil || !shouldRetry(err) {
t.Fatalf("err = %v, want connection refused", err)
}
}
// flakyStatus returns 500 for the first fails requests, then 200.
type flakyStatus struct {
fails int
calls int
}
func (s *flakyStatus) RoundTrip(req *http.Request) (*http.Response, error) {
s.calls++
code := 200
if s.calls <= s.fails {
code = 500
}
return &http.Response{
StatusCode: code,
Status: http.StatusText(code),
Body: io.NopCloser(strings.NewReader("")),
Header: make(http.Header),
}, nil
}
func TestRetryTransport5xxGet(t *testing.T) {
fs := &flakyStatus{fails: 2}
rt := &retryTransport{base: fs, initial: time.Millisecond, max: 2 * time.Millisecond}
req, _ := http.NewRequest(http.MethodGet, "http://upstream/api/version", nil)
resp, err := rt.RoundTrip(req)
if err != nil {
t.Fatal(err)
}
resp.Body.Close()
if resp.StatusCode != 200 {
t.Fatalf("status = %d, want 200", resp.StatusCode)
}
if fs.calls != 3 {
t.Fatalf("calls = %d, want 3", fs.calls)
}
}
func TestRetryTransport5xxPostNotRetried(t *testing.T) {
fs := &flakyStatus{fails: 10}
rt := &retryTransport{base: fs, initial: time.Millisecond, max: time.Millisecond}
// A streamed body (no GetBody) must not be replayed after a 500.
req, _ := http.NewRequest(http.MethodPost, "http://upstream/prompt", io.NopCloser(strings.NewReader("{}")))
resp, err := rt.RoundTrip(req)
if err != nil {
t.Fatal(err)
}
resp.Body.Close()
if resp.StatusCode != 500 {
t.Fatalf("status = %d, want 500", resp.StatusCode)
}
if fs.calls != 1 {
t.Fatalf("calls = %d, want 1 (no retry for streamed POST)", fs.calls)
}
}
func TestShouldRetryClassification(t *testing.T) {
cases := []struct {
name string
err error
want bool
}{
{"dial refused", &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}, true},
{"dial timeout", &net.OpError{Op: "dial", Net: "tcp", Err: timeoutErr{}}, true},
{"dns", &net.DNSError{Err: "no such host", IsNotFound: true}, true},
{"tls record", tls.RecordHeaderError{Msg: "bad"}, true},
{"tls alert", tls.AlertError(42), true},
{"read error mid-request", &net.OpError{Op: "read", Net: "tcp", Err: syscall.ECONNRESET}, false},
{"plain error", io.EOF, false},
}
for _, c := range cases {
if got := shouldRetry(c.err); got != c.want {
t.Errorf("%s: shouldRetry = %v, want %v", c.name, got, c.want)
}
}
}
type timeoutErr struct{}
func (timeoutErr) Error() string { return "i/o timeout" }
func (timeoutErr) Timeout() bool { return true }
func (timeoutErr) Temporary() bool { return true }