gpulock: GPU arbitration proxy for Ollama + ComfyUI

Implements SPEC.md: two listeners, one writer-preferring two-mode lock, Ollama unload before image jobs, ComfyUI history polling + VRAM free, optional model warm-up, healthz/metrics endpoints, streaming-safe reverse proxies, Dockerfile and Gitea Actions CI.
2026-09-20 18:05:20 +02:00
commit 065c294a96
14 changed files with 2027 additions and 0 deletions
@@ -0,0 +1,193 @@
+// gpulock 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"
+	"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
+	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:   ":11435",
+		listenComfy:    ":8189",
+		ollamaURL:      "http://127.0.0.1:11434",
+		comfyURL:       "http://127.0.0.1:8188",
+		unloadTimeout:  time.Minute,
+		jobTimeout:     15 * time.Minute,
+		llmWaitTimeout: 10 * time.Minute,
+		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},
+	} {
+		if err := envDuration(getenv, e.name, e.dst); err != nil {
+			return cfg, err
+		}
+	}
+	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, "gpulock: %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 gpulock",
+		"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,
+		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(), 5*time.Second)
+	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(), 10*time.Second)
+	defer shutdownCancel()
+	ollamaSrv.Shutdown(shutdownCtx)
+	comfySrv.Shutdown(shutdownCtx)
+}