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
	HealthInterval      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

	// 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
	ComfyDir          string
	ComfyIdleTimeout  time.Duration
	ComfyStartTimeout time.Duration

	// 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. GamePollInterval (GAME_POLL_INTERVAL) is how
	// often both checks run.
	GameProcs        []string
	GPUForeignVRAMMB int
	GPUIgnoreProcs   []string
	GamePollInterval time.Duration

	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,
		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.
		GPUIgnoreProcs:   []string{"ollama", "ollama app", "ollama_llama_server", "python", "pythonw"},
		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("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
}