llmgw/internal/bench/sampler_test.go

package bench

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"
)

// llamaMetrics is the part of llama-server's /metrics the router reads.
const llamaMetrics = `# TYPE llamacpp:requests_processing gauge
llamacpp:requests_processing 1
# TYPE llamacpp:requests_deferred gauge
llamacpp:requests_deferred 2
`

// guardedEngine is a llama.cpp-like engine: it streams completions to
// clientKey and serves /metrics only to engineKey, as llama-server does with
// --api-key. refuseAfter, when set, makes /metrics refuse every read after
// that many.
type guardedEngine struct {
	clientKey, engineKey string
	refuseAfter          int64
	chats, reads         atomic.Int64
}

func (g *guardedEngine) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case "/metrics":
		n := g.reads.Add(1)
		if r.Header.Get("Authorization") != "Bearer "+g.engineKey || (g.refuseAfter > 0 && n > g.refuseAfter) {
			http.Error(w, `{"error":"Invalid API Key"}`, http.StatusUnauthorized)
			return
		}
		_, _ = w.Write([]byte(llamaMetrics))
	case "/v1/chat/completions":
		if r.Header.Get("Authorization") != "Bearer "+g.clientKey {
			http.Error(w, "no", http.StatusUnauthorized)
			return
		}
		g.chats.Add(1)
		w.Header().Set("Content-Type", "text/event-stream")
		for i := range 3 {
			fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"t%d\"}}]}\n\n", i)
		}
		fmt.Fprint(w, "data: [DONE]\n\n")
		time.Sleep(20 * time.Millisecond)
	default:
		http.NotFound(w, r)
	}
}

func guardedRun(t *testing.T, g *guardedEngine, metricsKey string) (Result, error) {
	t.Helper()
	srv := httptest.NewServer(g)
	t.Cleanup(srv.Close)
	return Run(context.Background(), Config{
		Target: srv.URL, APIKey: g.clientKey, Model: "m",
		MetricsURL: srv.URL + "/metrics", Engine: "llamacpp", MetricsKey: metricsKey,
		Concurrency: []int{1}, Duration: 1500 * time.Millisecond,
		SLO: SLO{TTFTp99: time.Minute},
	}, nil)
}

func TestAnEngineThatGuardsItsMetricsIsReadWithItsOwnKey(t *testing.T) {
	// The requests go with the client's key and the metrics with the
	// engine's: llama.cpp's /metrics takes only its own.
	g := &guardedEngine{clientKey: "client", engineKey: "engine"}
	res, err := guardedRun(t, g, "engine")
	if err != nil {
		t.Fatal(err)
	}
	s := res.Steps[0]
	if s.MetricsUnread != "" || s.Waiting != 2 || s.Running != 1 {
		t.Errorf("metrics not read: waiting=%v running=%v unread=%q", s.Waiting, s.Running, s.MetricsUnread)
	}
	if res.SuggestedTargetWaiting != 2 {
		t.Errorf("target_waiting suggested %d", res.SuggestedTargetWaiting)
	}
}

func TestMetricsThatRefuseTheRunStopItBeforeAnyLoad(t *testing.T) {
	// Without the engine's key every sample is refused, and the run used to
	// report a queue depth of zero as though the engine had been idle.
	g := &guardedEngine{clientKey: "client", engineKey: "engine"}
	_, err := guardedRun(t, g, "")
	var refused *MetricsRefusedError
	if !errors.As(err, &refused) || refused.KeySent {
		t.Fatalf("expected the metrics refused for want of a key, got %v", err)
	}
	if n := g.chats.Load(); n != 0 {
		t.Errorf("%d requests were sent before the metrics were found unreadable", n)
	}

	_, err = guardedRun(t, &guardedEngine{clientKey: "client", engineKey: "engine"}, "client")
	if !errors.As(err, &refused) || !refused.KeySent {
		t.Fatalf("expected the key sent reported refused, got %v", err)
	}
}

func TestAStepWhoseMetricsWereNotReadSaysSo(t *testing.T) {
	// Readable before the ramp and not during it — an engine restarted with
	// another key, say. The step's zeros are not an idle engine.
	g := &guardedEngine{clientKey: "client", engineKey: "engine", refuseAfter: 1}
	res, err := guardedRun(t, g, "engine")
	if err != nil {
		t.Fatal(err)
	}
	if s := res.Steps[0]; s.MetricsUnread == "" {
		t.Errorf("a step with no metrics read reported none missing: %+v", s)
	}
	if res.SuggestedTargetWaiting != 0 || res.SuggestedKVCeiling != 0 {
		t.Errorf("thresholds suggested from metrics never read: %+v", res)
	}
}

func TestMetricsInAnotherEnginesNamesStopTheRun(t *testing.T) {
	g := &guardedEngine{clientKey: "client", engineKey: "engine"}
	srv := httptest.NewServer(g)
	defer srv.Close()
	_, err := Run(context.Background(), Config{
		Target: srv.URL, APIKey: "client", Model: "m",
		MetricsURL: srv.URL + "/metrics", Engine: "vllm", MetricsKey: "engine",
		Concurrency: []int{1}, Duration: time.Second,
	}, nil)
	if err == nil || g.chats.Load() != 0 {
		t.Fatalf("a run reading llama.cpp's metrics with vLLM's names started: %v", err)
	}
}