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,180 @@
+// Package ollama is a minimal client for the Ollama management endpoints
+// gpulock needs: listing loaded models, unloading them, and warming a model.
+package ollama
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/url"
+	"time"
+)
+
+// Client talks to an Ollama upstream.
+type Client struct {
+	base string
+	hc   *http.Client
+	log  *slog.Logger
+
+	// PollInterval is how often /api/ps is re-checked while waiting for
+	// models to unload. Defaults to 500ms; tests can shorten it.
+	PollInterval time.Duration
+}
+
+// New validates baseURL and returns a Client.
+func New(baseURL string, log *slog.Logger) (*Client, error) {
+	u, err := url.Parse(baseURL)
+	if err != nil || u.Scheme == "" || u.Host == "" {
+		return nil, fmt.Errorf("invalid Ollama URL %q", baseURL)
+	}
+	return &Client{
+		base:         u.String(),
+		hc:           &http.Client{Timeout: 30 * time.Second},
+		log:          log,
+		PollInterval: 500 * time.Millisecond,
+	}, nil
+}
+
+// Probe checks that Ollama answers on /api/version.
+func (c *Client) Probe(ctx context.Context) error {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/api/version", nil)
+	if err != nil {
+		return err
+	}
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	io.Copy(io.Discard, resp.Body)
+	if resp.StatusCode/100 != 2 {
+		return fmt.Errorf("/api/version returned %s", resp.Status)
+	}
+	return nil
+}
+
+type psResponse struct {
+	Models []struct {
+		Name  string `json:"name"`
+		Model string `json:"model"`
+	} `json:"models"`
+}
+
+// LoadedModels returns the names of models currently held in memory.
+func (c *Client) LoadedModels(ctx context.Context) ([]string, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/api/ps", nil)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode/100 != 2 {
+		io.Copy(io.Discard, resp.Body)
+		return nil, fmt.Errorf("/api/ps returned %s", resp.Status)
+	}
+	var ps psResponse
+	if err := json.NewDecoder(resp.Body).Decode(&ps); err != nil {
+		return nil, err
+	}
+	models := make([]string, 0, len(ps.Models))
+	for _, m := range ps.Models {
+		if m.Name != "" {
+			models = append(models, m.Name)
+		} else {
+			models = append(models, m.Model)
+		}
+	}
+	return models, nil
+}
+
+func (c *Client) post(ctx context.Context, path string, body any) (int, error) {
+	buf, err := json.Marshal(body)
+	if err != nil {
+		return 0, err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.base+path, bytes.NewReader(buf))
+	if err != nil {
+		return 0, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return 0, err
+	}
+	defer resp.Body.Close()
+	io.Copy(io.Discard, resp.Body)
+	return resp.StatusCode, nil
+}
+
+// unloadModel asks Ollama to evict one model. Embedding-only models reject
+// /api/generate, so fall back to /api/embed on non-2xx.
+func (c *Client) unloadModel(ctx context.Context, model string) error {
+	status, err := c.post(ctx, "/api/generate", map[string]any{"model": model, "keep_alive": 0})
+	if err == nil && status/100 == 2 {
+		return nil
+	}
+	status, err2 := c.post(ctx, "/api/embed", map[string]any{"model": model, "input": "x", "keep_alive": 0})
+	if err2 != nil {
+		return fmt.Errorf("generate: %v; embed: %v", err, err2)
+	}
+	if status/100 != 2 {
+		return fmt.Errorf("embed unload returned status %d", status)
+	}
+	return nil
+}
+
+// UnloadAll evicts every loaded model and waits until /api/ps is empty. The
+// context should carry the unload deadline; when it expires UnloadAll returns
+// ctx.Err() so the caller can degrade instead of failing. It returns the
+// elapsed wall time.
+func (c *Client) UnloadAll(ctx context.Context) (time.Duration, error) {
+	start := time.Now()
+	models, err := c.LoadedModels(ctx)
+	if err != nil {
+		return time.Since(start), fmt.Errorf("list models: %w", err)
+	}
+	for _, m := range models {
+		if err := c.unloadModel(ctx, m); err != nil {
+			c.log.Warn("failed to unload model", "model", m, "err", err)
+		}
+	}
+	for {
+		models, err := c.LoadedModels(ctx)
+		if err == nil && len(models) == 0 {
+			return time.Since(start), nil
+		}
+		if err != nil {
+			c.log.Warn("failed to poll /api/ps", "err", err)
+		}
+		select {
+		case <-ctx.Done():
+			return time.Since(start), ctx.Err()
+		case <-time.After(c.PollInterval):
+		}
+	}
+}
+
+// Warm reloads a model and pins it in memory (keep_alive=-1) so the next
+// chat request does not pay the load time.
+func (c *Client) Warm(ctx context.Context, model string) error {
+	status, err := c.post(ctx, "/api/generate", map[string]any{
+		"model":      model,
+		"prompt":     "",
+		"stream":     false,
+		"keep_alive": -1,
+	})
+	if err != nil {
+		return err
+	}
+	if status/100 != 2 {
+		return fmt.Errorf("warm generate returned status %d", status)
+	}
+	return nil
+}