gpu-turnstile/internal/config/config.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

	// 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",

		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
	}
	// 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
}