gpu-turnstile/cmd/gpulock/main.go

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