mram
bacb26772a
LLM_BUSY_MODE=reject answers blocked LLM requests immediately with LLM_BUSY_STATUS (default 503, 429 works) and Retry-After, so routers like LiteLLM can cool down and retry instead of holding a hung connection. The default wait mode now also sends Retry-After when LLM_WAIT_TIMEOUT expires. Document the service account (LocalSystem default, NT SERVICE virtual-account hardening) and the Program Files / ProgramData install layout.
347 lines
10 KiB
Go
347 lines
10 KiB
Go
// 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
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"path/filepath"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
|
|
"gpu-turnstile/internal/comfy"
|
|
"gpu-turnstile/internal/config"
|
|
"gpu-turnstile/internal/lock"
|
|
"gpu-turnstile/internal/metrics"
|
|
"gpu-turnstile/internal/ollama"
|
|
"gpu-turnstile/internal/proxy"
|
|
"gpu-turnstile/internal/service"
|
|
"gpu-turnstile/internal/update"
|
|
)
|
|
|
|
// version is injected at build time via -ldflags "-X main.version=...".
|
|
var version = "dev"
|
|
|
|
// exitCodeUpdate tells the service recovery configuration to restart the
|
|
// process: a signed update has been staged and the GPU lock is idle.
|
|
const exitCodeUpdate = 3
|
|
|
|
func main() {
|
|
configPath, args := splitConfigFlag(os.Args[1:])
|
|
if len(args) > 0 && args[0] == "service" {
|
|
os.Exit(serviceCommand(configPath, args[1:]))
|
|
}
|
|
if len(args) > 0 {
|
|
fmt.Fprintf(os.Stderr, "usage: gpu-turnstile [-config path] | gpu-turnstile service install|remove [-config path]\n")
|
|
os.Exit(2)
|
|
}
|
|
|
|
if exePath, err := os.Executable(); err == nil {
|
|
update.CleanupOld(exePath)
|
|
}
|
|
|
|
cfg, err := loadMergedConfig(configPath)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "gpu-turnstile: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
log, logOut, logCloser := newLogger(cfg)
|
|
defer logCloser.Close()
|
|
|
|
if service.IsService() {
|
|
if err := service.Run(func(ctx context.Context) error { return run(ctx, cfg, log, logOut, true) }); err != nil {
|
|
log.Error("service failed", "err", err)
|
|
os.Exit(1)
|
|
}
|
|
return
|
|
}
|
|
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
|
defer stop()
|
|
if err := run(ctx, cfg, log, logOut, false); err != nil {
|
|
log.Error("listener failed", "err", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
// splitConfigFlag extracts -config <path> (or -config=<path>) from args.
|
|
func splitConfigFlag(args []string) (string, []string) {
|
|
var configPath string
|
|
rest := args[:0]
|
|
for i := 0; i < len(args); i++ {
|
|
switch {
|
|
case args[i] == "-config" && i+1 < len(args):
|
|
configPath = args[i+1]
|
|
i++
|
|
case strings.HasPrefix(args[i], "-config="):
|
|
configPath = strings.TrimPrefix(args[i], "-config=")
|
|
default:
|
|
rest = append(rest, args[i])
|
|
}
|
|
}
|
|
return configPath, rest
|
|
}
|
|
|
|
// defaultConfigPath returns gpu-turnstile.env next to the executable.
|
|
func defaultConfigPath() string {
|
|
exe, err := os.Executable()
|
|
if err != nil {
|
|
return "gpu-turnstile.env"
|
|
}
|
|
return filepath.Join(filepath.Dir(exe), "gpu-turnstile.env")
|
|
}
|
|
|
|
// resolveConfigPath applies the precedence: -config flag, then
|
|
// GPU_TURNSTILE_CONFIG, then the default next to the executable.
|
|
func resolveConfigPath(flagValue string) string {
|
|
if flagValue != "" {
|
|
return flagValue
|
|
}
|
|
if v := os.Getenv("GPU_TURNSTILE_CONFIG"); v != "" {
|
|
return v
|
|
}
|
|
return defaultConfigPath()
|
|
}
|
|
|
|
// loadMergedConfig reads the config file (if present) and overlays process
|
|
// environment variables on top. A missing file is fine; an unreadable or
|
|
// malformed file is fatal.
|
|
func loadMergedConfig(flagValue string) (config.Config, error) {
|
|
path := resolveConfigPath(flagValue)
|
|
values := map[string]string{}
|
|
if f, err := os.Open(path); err == nil {
|
|
defer f.Close()
|
|
parsed, err := config.ParseEnvFile(f)
|
|
if err != nil {
|
|
return config.Config{}, fmt.Errorf("%s: %w", path, err)
|
|
}
|
|
values = parsed
|
|
} else if !errors.Is(err, os.ErrNotExist) {
|
|
return config.Config{}, fmt.Errorf("read config file: %w", err)
|
|
}
|
|
getenv := func(key string) string {
|
|
if v := os.Getenv(key); v != "" {
|
|
return v
|
|
}
|
|
return values[key]
|
|
}
|
|
return config.Load(getenv)
|
|
}
|
|
|
|
// newLogger builds the slog logger and returns the output writer (stderr or
|
|
// the opened LOG_FILE) plus a closer for it.
|
|
func newLogger(cfg config.Config) (*slog.Logger, io.Writer, io.Closer) {
|
|
out := io.Writer(os.Stderr)
|
|
closer := io.NopCloser(nil)
|
|
if cfg.LogFile != "" {
|
|
if f, err := os.OpenFile(cfg.LogFile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644); err == nil {
|
|
out = f
|
|
closer = f
|
|
} else {
|
|
fmt.Fprintf(os.Stderr, "gpu-turnstile: cannot open LOG_FILE %s: %v (logging to stderr)\n", cfg.LogFile, err)
|
|
}
|
|
}
|
|
opts := &slog.HandlerOptions{Level: cfg.LogLevel}
|
|
var handler slog.Handler = slog.NewTextHandler(out, opts)
|
|
if cfg.LogJSON {
|
|
handler = slog.NewJSONHandler(out, opts)
|
|
}
|
|
log := slog.New(handler)
|
|
slog.SetDefault(log)
|
|
return log, out, closer
|
|
}
|
|
|
|
func serviceCommand(configPath string, args []string) int {
|
|
if len(args) != 1 || (args[0] != "install" && args[0] != "remove") {
|
|
fmt.Fprintf(os.Stderr, "usage: gpu-turnstile service install|remove [-config path]\n")
|
|
return 2
|
|
}
|
|
var err error
|
|
if args[0] == "install" {
|
|
path := resolveConfigPath(configPath)
|
|
if abs, absErr := filepath.Abs(path); absErr == nil {
|
|
path = abs
|
|
}
|
|
err = service.Install(path)
|
|
} else {
|
|
err = service.Remove()
|
|
}
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "gpu-turnstile service %s: %v\n", args[0], err)
|
|
return 1
|
|
}
|
|
fmt.Printf("service %s: %sd\n", service.Name, args[0])
|
|
return 0
|
|
}
|
|
|
|
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.
|
|
log.Log(ctx, 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,
|
|
"llm_busy_mode", cfg.LLMBusyMode,
|
|
"llm_busy_status", cfg.LLMBusyStatus,
|
|
"busy_retry_after", cfg.BusyRetryAfter,
|
|
"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,
|
|
"auto_update", cfg.AutoUpdate,
|
|
"update_interval", cfg.UpdateInterval,
|
|
"update_repo", cfg.UpdateRepo,
|
|
"update_asset", cfg.UpdateAsset,
|
|
"log_level", cfg.LogLevel,
|
|
"log_format", map[bool]string{true: "json", false: "text"}[cfg.LogJSON],
|
|
"log_file", cfg.LogFile,
|
|
)
|
|
|
|
lk := lock.New(log)
|
|
ollamaClient, err := ollama.New(cfg.OllamaURL, log)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
comfyClient, err := comfy.New(cfg.ComfyURL, log)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
srv, err := proxy.New(proxy.Config{
|
|
OllamaURL: cfg.OllamaURL,
|
|
ComfyURL: cfg.ComfyURL,
|
|
Lock: lk,
|
|
Ollama: ollamaClient,
|
|
Comfy: comfyClient,
|
|
Metrics: metrics.New(),
|
|
Log: log,
|
|
LogColor: !cfg.LogJSON && cfg.LogFile == "" && os.Getenv("NO_COLOR") == "",
|
|
LogWriter: logOut,
|
|
LLMWaitTimeout: cfg.LLMWaitTimeout,
|
|
UnloadTimeout: cfg.UnloadTimeout,
|
|
JobTimeout: cfg.JobTimeout,
|
|
UnloadPollInterval: cfg.UnloadPollInterval,
|
|
HistoryPollInterval: cfg.HistoryPollInterval,
|
|
FreeTimeout: cfg.FreeTimeout,
|
|
WarmTimeout: cfg.WarmTimeout,
|
|
LLMBusyMode: cfg.LLMBusyMode,
|
|
LLMBusyStatus: cfg.LLMBusyStatus,
|
|
BusyRetryAfter: cfg.BusyRetryAfter,
|
|
BackoffInitial: cfg.BackoffInitial,
|
|
BackoffMax: cfg.BackoffMax,
|
|
PromptCaptureLimit: cfg.PromptCaptureLimit,
|
|
WarmModel: cfg.WarmModel,
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Probe both 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 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()}
|
|
|
|
errCh := make(chan error, 2)
|
|
go func() { errCh <- ollamaSrv.ListenAndServe() }()
|
|
go func() { errCh <- comfySrv.ListenAndServe() }()
|
|
|
|
if cfg.AutoUpdate {
|
|
go updateLoop(ctx, cfg, log, lk, isService)
|
|
}
|
|
|
|
select {
|
|
case err := <-errCh:
|
|
if err != nil && !errors.Is(err, http.ErrServerClosed) {
|
|
return err
|
|
}
|
|
case <-ctx.Done():
|
|
log.Info("shutting down")
|
|
}
|
|
|
|
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), cfg.ShutdownTimeout)
|
|
defer shutdownCancel()
|
|
ollamaSrv.Shutdown(shutdownCtx)
|
|
comfySrv.Shutdown(shutdownCtx)
|
|
return nil
|
|
}
|
|
|
|
// updateLoop checks for signed updates on startup and every UPDATE_INTERVAL.
|
|
// In service mode a staged update is applied by exiting with exitCodeUpdate
|
|
// once the GPU lock is idle; the service recovery configuration restarts the
|
|
// process with the new binary. Interactively it only logs.
|
|
func updateLoop(ctx context.Context, cfg config.Config, log *slog.Logger, lk *lock.Lock, isService bool) {
|
|
exePath, err := os.Executable()
|
|
if err != nil {
|
|
log.Warn("auto-update disabled: cannot locate executable", "err", err)
|
|
return
|
|
}
|
|
u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Log: log}
|
|
for {
|
|
staged, err := u.Check(ctx, exePath)
|
|
if err != nil && ctx.Err() == nil {
|
|
log.Warn("auto-update check failed", "err", err)
|
|
}
|
|
if staged {
|
|
if !isService {
|
|
log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply")
|
|
return
|
|
}
|
|
log.Warn("auto-update: staged; restarting once the GPU is idle")
|
|
if waitForIdle(ctx, lk, 24*time.Hour) {
|
|
log.Warn("auto-update: restarting to apply update")
|
|
os.Exit(exitCodeUpdate)
|
|
}
|
|
return
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-time.After(cfg.UpdateInterval):
|
|
}
|
|
}
|
|
}
|
|
|
|
// waitForIdle polls the lock until no LLM or image work is active or
|
|
// pending, max at most. Returns false on timeout or cancellation.
|
|
func waitForIdle(ctx context.Context, lk *lock.Lock, max time.Duration) bool {
|
|
deadline := time.Now().Add(max)
|
|
for {
|
|
if state, _, pending := lk.Snapshot(); state == lock.StateIdle && !pending {
|
|
return true
|
|
}
|
|
if time.Now().After(deadline) {
|
|
return false
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return false
|
|
case <-time.After(5 * time.Second):
|
|
}
|
|
}
|
|
}
|