llmgw/internal/router/scrape/llamacpp_test.go

package scrape

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

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

// llamaMetrics is what llama-server --metrics actually serves, trimmed to the
// families that matter.
const llamaMetrics = `# HELP llamacpp:prompt_tokens_total Number of prompt tokens processed.
# TYPE llamacpp:prompt_tokens_total counter
llamacpp:prompt_tokens_total 1200
# HELP llamacpp:tokens_predicted_total Number of generation tokens processed.
# TYPE llamacpp:tokens_predicted_total counter
llamacpp:tokens_predicted_total 3400
# HELP llamacpp:requests_processing Number of requests processing.
# TYPE llamacpp:requests_processing gauge
llamacpp:requests_processing 3
# HELP llamacpp:requests_deferred Number of requests deferred.
# TYPE llamacpp:requests_deferred gauge
llamacpp:requests_deferred 7
`

func TestLlamaCppMetricsSatisfyTheContract(t *testing.T) {
	s, err := Parse(strings.NewReader(llamaMetrics), LlamaCppMapping())
	if err != nil {
		t.Fatal(err)
	}
	if s.Running != 3 {
		t.Errorf("running=%d", s.Running)
	}
	// "Deferred" is llama.cpp's word for a request that arrived with every slot
	// busy, which is what the scheduler needs to read as queue depth.
	if s.Waiting != 7 {
		t.Errorf("waiting=%d", s.Waiting)
	}
	// Current llama.cpp publishes no cache ratio at all, so this stays zero
	// rather than failing the whole scrape.
	if s.KVCacheUsage != 0 {
		t.Errorf("kv=%v", s.KVCacheUsage)
	}
	if s.PromptTokens != 1200 || s.GenerationTokens != 3400 {
		t.Errorf("tokens=%v/%v", s.PromptTokens, s.GenerationTokens)
	}
}

func TestLlamaCppMappingDoesNotInheritVLLMNames(t *testing.T) {
	// llama.cpp has no request-success counter. Filling that field with vLLM's
	// name would send the router hunting for a metric that cannot be there, and
	// would hide the fact that this engine simply does not report it.
	m := LlamaCppMapping()
	m.normalize()
	if m.SuccessTotal != "" || m.CacheConfig != "" || m.KVCacheUsage != "" {
		t.Errorf("vLLM names leaked into the llama.cpp mapping: %+v", m)
	}
	if m.Running != "llamacpp:requests_processing" {
		t.Errorf("running=%q", m.Running)
	}
}

func TestEmptyMappingStillDefaultsToVLLM(t *testing.T) {
	var m Mapping
	m.normalize()
	if m.Running != "vllm:num_requests_running" || m.SuccessTotal == "" {
		t.Errorf("an unconfigured mapping should be vLLM's: %+v", m)
	}
}

func TestVLLMNamesAreNotReadAsLlamaCpp(t *testing.T) {
	// The failure this guards against is silent: reading a vLLM endpoint with
	// llama.cpp's vocabulary finds nothing, and an endpoint with no numbers
	// must be reported as unmeasurable rather than as idle.
	var missing *MissingError
	_, err := Parse(strings.NewReader(llamaMetrics), VLLMMapping())
	if !errors.As(err, &missing) {
		t.Fatalf("expected the contract to be reported incomplete, got %v", err)
	}
	if len(missing.Missing) != 3 {
		t.Errorf("every metric vLLM's mapping expects should be named: %v", missing.Missing)
	}
}

func TestMappingForPicksByEngine(t *testing.T) {
	if MappingFor("llamacpp").Running != "llamacpp:requests_processing" {
		t.Error("llamacpp did not get its own mapping")
	}
	// An endpoint that names no engine, or one this build has never heard of,
	// is read with vLLM's names rather than left unreadable.
	for _, engine := range []string{"", "vllm", "something-new"} {
		if MappingFor(engine).Running != "vllm:num_requests_running" {
			t.Errorf("engine %q did not fall back to vLLM", engine)
		}
	}
}

func TestAMetricTheMappingNamesIsStillMandatory(t *testing.T) {
	// Leaving a name out says "this engine does not publish it". Naming one the
	// engine does not serve is a contract violation, and must stay one: the
	// alternative is an endpoint quietly scheduled against numbers that are
	// zero because nothing ever set them.
	m := LlamaCppMapping()
	m.KVCacheUsage = "llamacpp:kv_cache_usage_ratio"

	var missing *MissingError
	if _, err := Parse(strings.NewReader(llamaMetrics), m); !errors.As(err, &missing) {
		t.Fatalf("a named but absent metric was tolerated: %v", err)
	}
}

func TestRunningAndWaitingCannotBeOmitted(t *testing.T) {
	// There is no engine worth routing to that reports neither how much it is
	// doing nor how much is queued.
	_, err := Parse(strings.NewReader(llamaMetrics), Mapping{
		Running: "llamacpp:requests_processing", Waiting: "nope:waiting",
	})
	if err == nil {
		t.Error("an endpoint with no queue signal was accepted")
	}
}

// logCapture collects records so a test can assert what an operator would see.
type logCapture struct {
	mu      sync.Mutex
	records []slog.Record
}

func (c *logCapture) Enabled(context.Context, slog.Level) bool { return true }
func (c *logCapture) WithAttrs([]slog.Attr) slog.Handler       { return c }
func (c *logCapture) WithGroup(string) slog.Handler            { return c }
func (c *logCapture) Handle(_ context.Context, r slog.Record) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.records = append(c.records, r)
	return nil
}

func (c *logCapture) count(level slog.Level, substr string) int {
	c.mu.Lock()
	defer c.mu.Unlock()
	n := 0
	for _, r := range c.records {
		if r.Level == level && strings.Contains(r.Message, substr) {
			n++
		}
	}
	return n
}

func TestABrokenContractIsReportedOnceAndLoudly(t *testing.T) {
	// The failure this guards against is a cluster that looks healthy while
	// none of its routing decisions are informed by anything. Debug-level
	// logging was not enough: a contract violation fails every scrape forever,
	// and at the default level nothing said so.
	served := llamaMetrics
	var mu sync.Mutex
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		mu.Lock()
		defer mu.Unlock()
		_, _ = io.WriteString(w, served)
	}))
	defer srv.Close()

	cap := &logCapture{}
	reg := endpoint.NewRegistry(endpoint.Config{})
	in := endpoint.Info{ID: "e1", URL: srv.URL, Engine: "vllm", Models: []string{"m"}}
	reg.Replace([]endpoint.Info{in})

	// A generous timeout: a scrape that timed out would be reported as
	// unreachable, and the violation reported again once it answered — two
	// warnings where the test counts one, on a slow runner only.
	p := NewPoller(reg, Config{Interval: 20 * time.Millisecond, Timeout: 2 * time.Second},
		slog.New(cap))
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	p.Sync(ctx)
	defer p.Stop()

	// vLLM's names against llama.cpp's output: nothing matches.
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) && cap.count(slog.LevelWarn, "does not serve the metrics") == 0 {
		time.Sleep(10 * time.Millisecond)
	}
	if n := cap.count(slog.LevelWarn, "does not serve the metrics"); n != 1 {
		t.Fatalf("the violation was reported %d times; expected exactly one", n)
	}

	// It must not turn into a per-scrape stream either.
	time.Sleep(200 * time.Millisecond)
	if n := cap.count(slog.LevelWarn, "does not serve the metrics"); n != 1 {
		t.Errorf("the warning repeated: %d times", n)
	}

	// And recovery has to be visible, or an operator cannot tell a fixed
	// cluster from a broken one.
	mu.Lock()
	served = `# TYPE vllm:num_requests_running gauge
vllm:num_requests_running 1
# TYPE vllm:num_requests_waiting gauge
vllm:num_requests_waiting 2
# TYPE vllm:kv_cache_usage_perc gauge
vllm:kv_cache_usage_perc 0.5
`
	mu.Unlock()

	deadline = time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) && cap.count(slog.LevelInfo, "readable again") == 0 {
		time.Sleep(10 * time.Millisecond)
	}
	if cap.count(slog.LevelInfo, "readable again") == 0 {
		t.Error("recovery was never reported")
	}
}