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,138 @@
+// Package comfy is a minimal client for the ComfyUI endpoints gpulock needs
+// after a prompt has been accepted: polling /history and freeing VRAM.
+package comfy
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/url"
+	"time"
+)
+
+// Client talks to a ComfyUI upstream.
+type Client struct {
+	base string
+	hc   *http.Client
+	log  *slog.Logger
+
+	// PollInterval is how often /history/<id> is polled. Defaults to 1s;
+	// 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 ComfyUI URL %q", baseURL)
+	}
+	return &Client{
+		base:         u.String(),
+		hc:           &http.Client{Timeout: 30 * time.Second},
+		log:          log,
+		PollInterval: time.Second,
+	}, nil
+}
+
+// Probe checks that ComfyUI answers on /system_stats.
+func (c *Client) Probe(ctx context.Context) error {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/system_stats", 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("/system_stats returned %s", resp.Status)
+	}
+	return nil
+}
+
+type historyEntry struct {
+	Status struct {
+		Completed bool   `json:"completed"`
+		StatusStr string `json:"status_str"`
+	} `json:"status"`
+}
+
+// jobDone reports whether the history entry for promptID shows a finished
+// job (completed or errored).
+func (c *Client) jobDone(ctx context.Context, promptID string) (bool, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/history/"+promptID, nil)
+	if err != nil {
+		return false, err
+	}
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return false, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode/100 != 2 {
+		io.Copy(io.Discard, resp.Body)
+		return false, fmt.Errorf("/history/%s returned %s", promptID, resp.Status)
+	}
+	var history map[string]historyEntry
+	if err := json.NewDecoder(resp.Body).Decode(&history); err != nil {
+		return false, err
+	}
+	entry, ok := history[promptID]
+	if !ok {
+		return false, nil
+	}
+	return entry.Status.Completed || entry.Status.StatusStr == "error", nil
+}
+
+// WaitJob polls until the job for promptID is finished. The context should
+// carry the job deadline; transient polling errors are logged and retried,
+// so a returned error is always from the context.
+func (c *Client) WaitJob(ctx context.Context, promptID string) error {
+	for {
+		done, err := c.jobDone(ctx, promptID)
+		if err != nil {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			c.log.Warn("failed to poll job history", "prompt_id", promptID, "err", err)
+		}
+		if done {
+			return nil
+		}
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(c.PollInterval):
+		}
+	}
+}
+
+// Free asks ComfyUI to unload its models and free VRAM.
+func (c *Client) Free(ctx context.Context) error {
+	buf, err := json.Marshal(map[string]bool{"unload_models": true, "free_memory": true})
+	if err != nil {
+		return err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.base+"/free", bytes.NewReader(buf))
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	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("/free returned %s", resp.Status)
+	}
+	return nil
+}