gpu-turnstile/internal/comfy/client.go

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