249 lines
7.0 KiB
Go
249 lines
7.0 KiB
Go
// Package config loads gpu-turnstile's configuration from environment
|
|
// variables and an optional .env-style config file.
|
|
package config
|
|
|
|
import (
|
|
"bufio"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// Config holds every gpu-turnstile setting.
|
|
type Config struct {
|
|
ListenOllama string
|
|
ListenComfy string
|
|
OllamaURL string
|
|
ComfyURL string
|
|
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
|
|
|
|
AutoUpdate bool
|
|
UpdateInterval time.Duration
|
|
UpdateRepo string
|
|
UpdateAsset string
|
|
|
|
// AppVersion is the version the user wants to run: "dev" disables
|
|
// updates, "stable" tracks the latest release, anything else is an
|
|
// exact vX.Y.Z release to pin. From APP_VER; defaults to "stable".
|
|
AppVersion string
|
|
|
|
// LLMBusyMode is "wait" (hold requests until the lock is free or
|
|
// LLMWaitTimeout expires) or "reject" (immediately answer with
|
|
// LLMBusyStatus + Retry-After when an image job is active or pending).
|
|
LLMBusyMode string
|
|
LLMBusyStatus int
|
|
BusyRetryAfter int
|
|
|
|
WarmModel string
|
|
LogLevel slog.Level
|
|
LogJSON bool
|
|
LogFile string
|
|
}
|
|
|
|
// Defaults returns the configuration used when neither the environment nor
|
|
// 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",
|
|
UnloadTimeout: time.Minute,
|
|
JobTimeout: 15 * time.Minute,
|
|
LLMWaitTimeout: 10 * time.Minute,
|
|
|
|
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,
|
|
|
|
AutoUpdate: true,
|
|
UpdateInterval: 6 * time.Hour,
|
|
UpdateRepo: "https://git.rambossek.at/PUBLIC/gpu-turnstile",
|
|
UpdateAsset: "gpu-turnstile.exe",
|
|
AppVersion: "stable",
|
|
|
|
LLMBusyMode: "wait",
|
|
LLMBusyStatus: 503,
|
|
BusyRetryAfter: 30,
|
|
|
|
LogLevel: slog.LevelWarn,
|
|
}
|
|
}
|
|
|
|
// ErrNoConsumer is returned by Load when neither OLLAMA_URL nor COMFY_URL
|
|
// is set. Commands that never talk to an upstream (--force-update) may
|
|
// ignore it and proceed with the rest of the configuration.
|
|
var ErrNoConsumer = errors.New("at least one of OLLAMA_URL or COMFY_URL must be set (each URL enables its consumer)")
|
|
|
|
// ParseEnvFile parses a .env-style file: KEY=VALUE lines, blank lines and
|
|
// #-comments are ignored, no quoting. A line without '=' is an error.
|
|
func ParseEnvFile(r io.Reader) (map[string]string, error) {
|
|
values := make(map[string]string)
|
|
scanner := bufio.NewScanner(r)
|
|
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
|
|
lineNo := 0
|
|
for scanner.Scan() {
|
|
lineNo++
|
|
line := strings.TrimSpace(scanner.Text())
|
|
if line == "" || strings.HasPrefix(line, "#") {
|
|
continue
|
|
}
|
|
key, value, ok := strings.Cut(line, "=")
|
|
if !ok {
|
|
return nil, fmt.Errorf("line %d: expected KEY=VALUE", lineNo)
|
|
}
|
|
key = strings.TrimSpace(key)
|
|
if key == "" {
|
|
return nil, fmt.Errorf("line %d: empty key", lineNo)
|
|
}
|
|
values[key] = strings.TrimSpace(value)
|
|
}
|
|
return values, scanner.Err()
|
|
}
|
|
|
|
func envDuration(getenv func(string) string, name string, dst *time.Duration) error {
|
|
v := getenv(name)
|
|
if v == "" {
|
|
return nil
|
|
}
|
|
d, err := time.ParseDuration(v)
|
|
if err != nil {
|
|
return fmt.Errorf("%s: %w", name, err)
|
|
}
|
|
*dst = d
|
|
return nil
|
|
}
|
|
|
|
// Load overlays values from getenv onto the Defaults. Unknown keys are
|
|
// ignored. Invalid values are fatal.
|
|
func Load(getenv func(string) string) (Config, error) {
|
|
cfg := Defaults()
|
|
for _, e := range []struct {
|
|
name string
|
|
dst *string
|
|
}{
|
|
{"LISTEN_OLLAMA", &cfg.ListenOllama},
|
|
{"LISTEN_COMFY", &cfg.ListenComfy},
|
|
{"OLLAMA_URL", &cfg.OllamaURL},
|
|
{"COMFY_URL", &cfg.ComfyURL},
|
|
{"WARM_MODEL", &cfg.WarmModel},
|
|
{"UPDATE_REPO", &cfg.UpdateRepo},
|
|
{"UPDATE_ASSET", &cfg.UpdateAsset},
|
|
{"LOG_FILE", &cfg.LogFile},
|
|
} {
|
|
if v := getenv(e.name); v != "" {
|
|
*e.dst = v
|
|
}
|
|
}
|
|
for _, e := range []struct {
|
|
name string
|
|
dst *time.Duration
|
|
}{
|
|
{"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},
|
|
{"UPDATE_INTERVAL", &cfg.UpdateInterval},
|
|
} {
|
|
if err := envDuration(getenv, e.name, e.dst); err != nil {
|
|
return cfg, err
|
|
}
|
|
}
|
|
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
|
|
}
|
|
if v := getenv("AUTO_UPDATE"); v != "" {
|
|
b, err := strconv.ParseBool(v)
|
|
if err != nil {
|
|
return cfg, fmt.Errorf("AUTO_UPDATE: must be a boolean (true/false)")
|
|
}
|
|
cfg.AutoUpdate = b
|
|
}
|
|
if v := getenv("LLM_BUSY_MODE"); v != "" {
|
|
if v != "wait" && v != "reject" {
|
|
return cfg, fmt.Errorf("LLM_BUSY_MODE: must be \"wait\" or \"reject\"")
|
|
}
|
|
cfg.LLMBusyMode = v
|
|
}
|
|
if v := getenv("LLM_BUSY_STATUS"); v != "" {
|
|
n, err := strconv.Atoi(v)
|
|
if err != nil || n < 400 || n > 599 {
|
|
return cfg, fmt.Errorf("LLM_BUSY_STATUS: must be an HTTP status in 400-599")
|
|
}
|
|
cfg.LLMBusyStatus = n
|
|
}
|
|
if v := getenv("BUSY_RETRY_AFTER"); v != "" {
|
|
n, err := strconv.Atoi(v)
|
|
if err != nil || n <= 0 {
|
|
return cfg, fmt.Errorf("BUSY_RETRY_AFTER: must be a positive integer (seconds)")
|
|
}
|
|
cfg.BusyRetryAfter = n
|
|
}
|
|
if v := getenv("APP_VER"); v != "" {
|
|
switch {
|
|
case v == "dev" || v == "stable":
|
|
cfg.AppVersion = v
|
|
default:
|
|
if _, ok := parseVersion(v); !ok {
|
|
return cfg, fmt.Errorf("APP_VER: must be \"dev\", \"stable\" or a vX.Y.Z version")
|
|
}
|
|
cfg.AppVersion = "v" + strings.TrimPrefix(v, "v")
|
|
}
|
|
}
|
|
// 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(logLevelValue)); err != nil {
|
|
return cfg, fmt.Errorf("LOGLEVEL: %w", err)
|
|
}
|
|
cfg.LogLevel = level
|
|
}
|
|
switch strings.ToLower(getenv("LOG_FORMAT")) {
|
|
case "", "text":
|
|
case "json":
|
|
cfg.LogJSON = true
|
|
default:
|
|
return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"")
|
|
}
|
|
if cfg.OllamaURL == "" && cfg.ComfyURL == "" {
|
|
return cfg, ErrNoConsumer
|
|
}
|
|
return cfg, nil
|
|
}
|