gpu-turnstile/internal/ollama/client.go

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