llmgw/internal/bench/sampler.go

package bench

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape"
)

// sampler reads an engine's metrics during a step.
//
// It is what turns a latency measurement into something admission can use: the
// queue depth and cache pressure observed while the SLO was still being met are
// the values target_waiting and kv_ceiling should have, rather than the round
// numbers somebody typed.
//
// Optional, because a run pointed at the gateway has no single engine to read.
type sampler struct {
	url    string
	apiKey string
	engine string
	client *http.Client
	cancel context.CancelFunc

	mu      sync.Mutex
	wg      sync.WaitGroup
	waiting float64
	running float64
	kv      float64
	n       int
	// lastErr says why a step has no numbers: averages over no samples are
	// zeros, and zeros read as an idle engine.
	lastErr error
}

func newSampler(metricsURL, engine, key string) *sampler {
	return &sampler{
		url:    metricsURL,
		engine: engine,
		apiKey: key,
		client: &http.Client{Timeout: 2 * time.Second},
	}
}

func (s *sampler) start(ctx context.Context) {
	if s.url == "" {
		return
	}
	ctx, s.cancel = context.WithCancel(ctx)
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		t := time.NewTicker(time.Second)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				s.sample(ctx)
			}
		}
	}()
}

// MetricsRefusedError is an engine refusing to show its metrics without its
// key: llama.cpp guards /metrics with it, and llmgw's TensorRT-LLM launcher
// guards everything but /health.
type MetricsRefusedError struct {
	URL     string
	Status  string
	KeySent bool
}

func (e *MetricsRefusedError) Error() string {
	if e.KeySent {
		return fmt.Sprintf("%s refused the key sent for its metrics (%s)", e.URL, e.Status)
	}
	return fmt.Sprintf("%s guards its metrics with the engine's key, and none was given (%s)", e.URL, e.Status)
}

// read takes one sample.
func (s *sampler) read(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.url, nil)
	if err != nil {
		return err
	}
	if s.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+s.apiKey)
	}
	resp, err := s.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	switch {
	case resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden:
		return &MetricsRefusedError{URL: s.url, Status: resp.Status, KeySent: s.apiKey != ""}
	case resp.StatusCode != http.StatusOK:
		return fmt.Errorf("%s: %s", s.url, resp.Status)
	}

	// The same mapping the router reads with, so a bench result cannot disagree
	// with what the scheduler sees. The engine has to be named: reading a
	// llama.cpp endpoint with vLLM's names finds nothing, and the run would
	// report a queue depth of zero and suggest thresholds built on it.
	stats, err := scrape.Parse(resp.Body, scrape.MappingFor(s.engine))
	if err != nil {
		var missing *scrape.MissingError
		if errors.As(err, &missing) {
			return fmt.Errorf("%s does not serve %v: is %q the engine there?", s.url, missing.Missing, s.engine)
		}
		return err
	}
	s.mu.Lock()
	s.waiting += float64(stats.Waiting)
	s.running += float64(stats.Running)
	s.kv += stats.KVCacheUsage
	s.n++
	s.mu.Unlock()
	return nil
}

func (s *sampler) sample(ctx context.Context) {
	err := s.read(ctx)
	if err == nil || ctx.Err() != nil {
		return
	}
	s.mu.Lock()
	s.lastErr = err
	s.mu.Unlock()
}

func (s *sampler) stop() {
	if s.cancel != nil {
		s.cancel()
	}
	s.wg.Wait()
}

func (s *sampler) averages() (waiting, running, kv float64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.n == 0 {
		return 0, 0, 0
	}
	n := float64(s.n)
	return s.waiting / n, s.running / n, s.kv / n
}

// unread says why a step that tried to read the metrics has none of them, or
// "" when it has some or never tried.
func (s *sampler) unread() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.n > 0 || s.lastErr == nil {
		return ""
	}
	return s.lastErr.Error()
}