gpu-turnstile/cmd/gpu-turnstile/main.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"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"gpu-turnstile/internal/comfy"
	"gpu-turnstile/internal/lock"
	"gpu-turnstile/internal/metrics"
	"gpu-turnstile/internal/ollama"
	"gpu-turnstile/internal/proxy"
)

// version is injected at build time via -ldflags "-X main.version=...".
var version = "dev"

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
	promptCaptureLimit  int64

	warmModel string
	logLevel  slog.Level
	logJSON   bool
}

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
}

func loadConfig(getenv func(string) string) (config, error) {
	cfg := config{
		listenOllama:   ":11434",
		listenComfy:    ":8188",
		ollamaURL:      "http://127.0.0.1:11435",
		comfyURL:       "http://127.0.0.1:8189",
		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,
		promptCaptureLimit:  64 * 1024,

		logLevel: slog.LevelInfo,
	}
	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},
	} {
		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},
	} {
		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("LOG_LEVEL"); v != "" {
		var level slog.Level
		if err := level.UnmarshalText([]byte(v)); err != nil {
			return cfg, fmt.Errorf("LOG_LEVEL: %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\"")
	}
	return cfg, nil
}

func main() {
	cfg, err := loadConfig(os.Getenv)
	if err != nil {
		fmt.Fprintf(os.Stderr, "gpu-turnstile: %v\n", err)
		os.Exit(1)
	}

	opts := &slog.HandlerOptions{Level: cfg.logLevel}
	var handler slog.Handler = slog.NewTextHandler(os.Stderr, opts)
	if cfg.logJSON {
		handler = slog.NewJSONHandler(os.Stderr, opts)
	}
	log := slog.New(handler)
	slog.SetDefault(log)

	log.Info("starting gpu-turnstile",
		"version", version,
		"listen_ollama", cfg.listenOllama,
		"listen_comfy", cfg.listenComfy,
		"ollama_url", cfg.ollamaURL,
		"comfy_url", cfg.comfyURL,
	)

	ollamaClient, err := ollama.New(cfg.ollamaURL, log)
	if err != nil {
		log.Error("invalid configuration", "err", err)
		os.Exit(1)
	}
	comfyClient, err := comfy.New(cfg.comfyURL, log)
	if err != nil {
		log.Error("invalid configuration", "err", err)
		os.Exit(1)
	}

	srv, err := proxy.New(proxy.Config{
		OllamaURL:           cfg.ollamaURL,
		ComfyURL:            cfg.comfyURL,
		Lock:                lock.New(log),
		Ollama:              ollamaClient,
		Comfy:               comfyClient,
		Metrics:             metrics.New(),
		Log:                 log,
		LLMWaitTimeout:      cfg.llmWaitTimeout,
		UnloadTimeout:       cfg.unloadTimeout,
		JobTimeout:          cfg.jobTimeout,
		UnloadPollInterval:  cfg.unloadPollInterval,
		HistoryPollInterval: cfg.historyPollInterval,
		FreeTimeout:         cfg.freeTimeout,
		WarmTimeout:         cfg.warmTimeout,
		PromptCaptureLimit:  cfg.promptCaptureLimit,
		WarmModel:           cfg.warmModel,
	})
	if err != nil {
		log.Error("invalid configuration", "err", err)
		os.Exit(1)
	}

	// Probe both upstreams once; failure is logged, not fatal.
	probeCtx, probeCancel := context.WithTimeout(context.Background(), 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() }()

	sigCtx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	select {
	case err := <-errCh:
		if err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Error("listener failed", "err", err)
			os.Exit(1)
		}
	case <-sigCtx.Done():
		log.Info("shutting down")
	}

	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), cfg.shutdownTimeout)
	defer shutdownCancel()
	ollamaSrv.Shutdown(shutdownCtx)
	comfySrv.Shutdown(shutdownCtx)
}