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. Each field's env tag names the
// environment variable / config-file key that sets it; user-facing output
// (e.g. the reload diff) uses those names, never the Go field names.
type Config struct {
	ListenOllama   string        `env:"LISTEN_OLLAMA"`
	ListenComfy    string        `env:"LISTEN_COMFY"`
	OllamaURL      string        `env:"OLLAMA_URL"`
	ComfyURL       string        `env:"COMFY_URL"`
	UnloadTimeout  time.Duration `env:"UNLOAD_TIMEOUT"`
	JobTimeout     time.Duration `env:"JOB_TIMEOUT"`
	LLMWaitTimeout time.Duration `env:"LLM_WAIT_TIMEOUT"`

	UnloadPollInterval  time.Duration `env:"UNLOAD_POLL_INTERVAL"`
	HistoryPollInterval time.Duration `env:"HISTORY_POLL_INTERVAL"`
	ProbeTimeout        time.Duration `env:"PROBE_TIMEOUT"`
	HealthInterval      time.Duration `env:"HEALTH_INTERVAL"`
	FreeTimeout         time.Duration `env:"FREE_TIMEOUT"`
	WarmTimeout         time.Duration `env:"WARM_TIMEOUT"`
	ShutdownTimeout     time.Duration `env:"SHUTDOWN_TIMEOUT"`
	BackoffInitial      time.Duration `env:"BACKOFF_INITIAL"`
	BackoffMax          time.Duration `env:"BACKOFF_MAX"`
	PromptCaptureLimit  int64         `env:"PROMPT_CAPTURE_LIMIT"`

	AutoUpdate     bool          `env:"AUTO_UPDATE"`
	UpdateInterval time.Duration `env:"UPDATE_INTERVAL"`
	UpdateRepo     string        `env:"UPDATE_REPO"`
	UpdateAsset    string        `env:"UPDATE_ASSET"`

	// 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 `env:"APP_VER"`

	// 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 `env:"LLM_BUSY_MODE"`
	LLMBusyStatus  int    `env:"LLM_BUSY_STATUS"`
	BusyRetryAfter int    `env:"BUSY_RETRY_AFTER"`

	// ComfyCmd spawns and supervises a ComfyUI server on demand. When
	// ComfyCmd is empty but ComfyDir is set, management is enabled with the
	// standard venv layout under ComfyDir (.venv + main.py or
	// ComfyUI/main.py; --port from the COMFY_URL port) — ComfyCmd is the
	// override for other layouts and doubles as the working directory when
	// set explicitly. The managed server is stopped after ComfyIdleTimeout
	// without requests, freeing its VRAM; ComfyStartTimeout bounds how long
	// a request waits for it to come up.
	ComfyCmd          string        `env:"COMFY_CMD"`
	ComfyDir          string        `env:"COMFY_DIR"`
	ComfyIdleTimeout  time.Duration `env:"COMFY_IDLE_TIMEOUT"`
	ComfyStartTimeout time.Duration `env:"COMFY_START_TIMEOUT"`

	// GameProcs (GAME_PROCS) is a watch list of process names; while any of
	// them runs, the GPU is treated as held by a foreign process. The
	// nvidia-smi path (GPUForeignVRAMMB, GPU_FOREIGN_VRAM_MB) does the same
	// when a process not in GPUIgnoreProcs (GPU_IGNORE_PROCS) holds more than
	// that many MiB of VRAM, and the PDH path (GPUForeignUtilPct,
	// GPU_FOREIGN_UTIL_PCT) when such a process uses more than that many
	// percent of the GPU 3D engine (Windows only). GamePollInterval
	// (GAME_POLL_INTERVAL) is how often all checks run.
	GameProcs         []string      `env:"GAME_PROCS"`
	GPUForeignVRAMMB  int           `env:"GPU_FOREIGN_VRAM_MB"`
	GPUForeignUtilPct int           `env:"GPU_FOREIGN_UTIL_PCT"`
	GPUIgnoreProcs    []string      `env:"GPU_IGNORE_PROCS"`
	GamePollInterval  time.Duration `env:"GAME_POLL_INTERVAL"`

	WarmModel string     `env:"WARM_MODEL"`
	LogLevel  slog.Level `env:"LOGLEVEL"`
	LogJSON   bool       `env:"LOG_FORMAT"`
	LogFile   string     `env:"LOG_FILE"`
}

// 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,
		HealthInterval:      30 * 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,

		ComfyIdleTimeout:  5 * time.Minute,
		ComfyStartTimeout: 2 * time.Minute,

		// ComfyUI runs under python; excluding it (and Ollama) by name keeps
		// our own consumers from tripping the foreign-VRAM check. dwm is the
		// desktop compositor — it always shows some 3D-engine usage.
		GPUIgnoreProcs:   []string{"ollama", "ollama app", "ollama_llama_server", "python", "pythonw", "dwm"},
		GamePollInterval: 15 * time.Second,

		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()
}

// splitList parses a comma-separated setting into trimmed, non-empty items.
func splitList(v string) []string {
	var out []string
	for _, item := range strings.Split(v, ",") {
		if item = strings.TrimSpace(item); item != "" {
			out = append(out, item)
		}
	}
	return out
}

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},
		{"COMFY_CMD", &cfg.ComfyCmd},
		{"COMFY_DIR", &cfg.ComfyDir},
		{"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},
		{"HEALTH_INTERVAL", &cfg.HealthInterval},
		{"FREE_TIMEOUT", &cfg.FreeTimeout},
		{"WARM_TIMEOUT", &cfg.WarmTimeout},
		{"SHUTDOWN_TIMEOUT", &cfg.ShutdownTimeout},
		{"BACKOFF_INITIAL", &cfg.BackoffInitial},
		{"BACKOFF_MAX", &cfg.BackoffMax},
		{"COMFY_IDLE_TIMEOUT", &cfg.ComfyIdleTimeout},
		{"COMFY_START_TIMEOUT", &cfg.ComfyStartTimeout},
		{"UPDATE_INTERVAL", &cfg.UpdateInterval},
		{"GAME_POLL_INTERVAL", &cfg.GamePollInterval},
	} {
		if err := envDuration(getenv, e.name, e.dst); err != nil {
			return cfg, err
		}
	}
	if v := getenv("GAME_PROCS"); v != "" {
		cfg.GameProcs = splitList(v)
	}
	if v := getenv("GPU_IGNORE_PROCS"); v != "" {
		cfg.GPUIgnoreProcs = splitList(v)
	}
	if v := getenv("GPU_FOREIGN_VRAM_MB"); v != "" {
		n, err := strconv.Atoi(v)
		if err != nil || n < 0 {
			return cfg, fmt.Errorf("GPU_FOREIGN_VRAM_MB: must be a non-negative integer (MiB, 0 = disabled)")
		}
		cfg.GPUForeignVRAMMB = n
	}
	if v := getenv("GPU_FOREIGN_UTIL_PCT"); v != "" {
		n, err := strconv.Atoi(v)
		if err != nil || n < 0 || n > 100 {
			return cfg, fmt.Errorf("GPU_FOREIGN_UTIL_PCT: must be an integer in 0-100 (percent of the GPU 3D engine, 0 = disabled)")
		}
		cfg.GPUForeignUtilPct = n
	}
	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.ComfyCmd != "" && cfg.ComfyURL == "" {
		return cfg, fmt.Errorf("COMFY_CMD requires COMFY_URL to be set (the proxy needs somewhere to forward)")
	}
	if cfg.ComfyCmd == "" && cfg.ComfyDir != "" && cfg.ComfyURL == "" {
		return cfg, fmt.Errorf("COMFY_DIR without COMFY_CMD requires COMFY_URL to be set (it enables the managed ComfyUI)")
	}
	if cfg.OllamaURL == "" && cfg.ComfyURL == "" {
		return cfg, ErrNoConsumer
	}
	return cfg, nil
}