llmgw/internal/router/kvroute_test.go

package router

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus/testutil"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
)

// kvFakeVLLM is an engine with vLLM's /tokenize beside the completion: it
// "tokenizes" a chat as the bytes of each message's role and content, which
// is enough for the router and a test to agree on token ids.
type kvFakeVLLM struct {
	srv       *httptest.Server
	served    atomic.Int64
	tokenized atomic.Int64
	failTok   atomic.Bool
	lastTok   atomic.Value
}

func newKVFakeVLLM(t *testing.T) *kvFakeVLLM {
	t.Helper()
	f := &kvFakeVLLM{}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprintf(w, "vllm:num_requests_running 0\nvllm:num_requests_waiting 0\nvllm:kv_cache_usage_perc 0.1\n")
	})
	mux.HandleFunc("POST /tokenize", func(w http.ResponseWriter, r *http.Request) {
		f.tokenized.Add(1)
		b, _ := io.ReadAll(r.Body)
		f.lastTok.Store(b)
		if f.failTok.Load() {
			http.Error(w, "overloaded", http.StatusServiceUnavailable)
			return
		}
		var req struct {
			Messages []struct{ Role, Content string } `json:"messages"`
		}
		if err := json.Unmarshal(b, &req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		toks := fakeTokens(req.Messages...)
		_ = json.NewEncoder(w).Encode(map[string]any{"count": len(toks), "max_model_len": 4096, "tokens": toks})
	})
	mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, _ *http.Request) {
		f.served.Add(1)
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":4,"completion_tokens":2}}`)
	})
	f.srv = httptest.NewServer(mux)
	t.Cleanup(f.srv.Close)
	return f
}

func fakeTokens(msgs ...struct{ Role, Content string }) []int {
	var out []int
	for _, m := range msgs {
		for _, c := range []byte(m.Role + ":" + m.Content + "\n") {
			out = append(out, int(c))
		}
	}
	return out
}

func chatOf(prompt string) []struct{ Role, Content string } {
	return []struct{ Role, Content string }{{"user", prompt}}
}

// kvRouter is a router over two fake vLLM replicas, a and b, of a model
// whose profile publishes KV cache events.
func kvRouter(t *testing.T, tier1 bool) (*Router, *kvFakeVLLM, *kvFakeVLLM) {
	t.Helper()
	a, b := newKVFakeVLLM(t), newKVFakeVLLM(t)
	cfg := baseConfig(
		config.Endpoint{ID: "a", URL: a.srv.URL, Models: []string{"m"}},
		config.Endpoint{ID: "b", URL: b.srv.URL, Models: []string{"m"}},
		config.Endpoint{ID: "o-1", URL: a.srv.URL, Models: []string{"other"}},
	)
	m := catModel("m", nil)
	m.Profiles[0].KVEvents = true
	m.RoutingProfile = "chat"
	other := catModel("other", nil)
	other.RoutingProfile = "chat"
	cfg.Models = []catalog.Model{m, other}
	cfg.Router.Prefix.Enabled = tier1
	rt := newTestRouter(t, cfg)
	// Both scraped before anything is sent: a replica not yet heard from is
	// passed over while another has been, which is admission's rule and not
	// what these tests are about.
	deadline := time.Now().Add(3 * time.Second)
	for {
		fresh := 0
		for _, sn := range rt.reg.Snapshots("m") {
			if !sn.Stale {
				fresh++
			}
		}
		if fresh == 2 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("the replicas were never scraped")
		}
		time.Sleep(10 * time.Millisecond)
	}
	return rt, a, b
}

// longPrompt is long enough for several blocks of sixteen.
var longPrompt = strings.Repeat("You are a careful assistant who answers briefly. ", 6)

func sendChat(t *testing.T, rt *Router, prompt string) {
	t.Helper()
	rec := post(t, rt.Handler(), map[string]any{
		"model": "m", "messages": []map[string]any{{"role": "user", "content": prompt}},
	}, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
}

func kvDecisions(rt *Router, d kvindex.Decision) float64 {
	return testutil.ToFloat64(rt.m.kvDecisions.WithLabelValues(string(d)))
}

func TestRequestsGoToTheReplicaHoldingTheLongestBlockPrefix(t *testing.T) {
	rt, a, b := kvRouter(t, false)
	hs := kvevents.BlockHashes(fakeTokens(chatOf(longPrompt)...), 16, 0)
	if len(hs) < 6 {
		t.Fatalf("the prompt should make several blocks, made %d", len(hs))
	}
	// a holds the first two blocks — the opening a chat template shares —
	// and b the whole prompt.
	rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: hs[:2]})
	rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16, Stored: hs})

	for range 10 {
		sendChat(t, rt, longPrompt)
	}
	if b.served.Load() != 10 || a.served.Load() != 0 {
		t.Fatalf("served a=%d b=%d; every request should go to the replica holding its blocks",
			a.served.Load(), b.served.Load())
	}
	if n := kvDecisions(rt, kvindex.Hit); n != 10 {
		t.Fatalf("kv_hit counted %v times, want 10", n)
	}

	// b evicts everything past the first block; now a holds the longer
	// prefix, and the index follows the engines rather than history.
	rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16, Removed: hs[1:]})
	rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: hs[2:]})
	for range 5 {
		sendChat(t, rt, longPrompt)
	}
	if a.served.Load() != 5 {
		t.Fatalf("served a=%d after b evicted; want 5", a.served.Load())
	}
}

func TestTooShortAMatchLeavesTheChoiceToLoad(t *testing.T) {
	rt, _, _ := kvRouter(t, false)
	hs := kvevents.BlockHashes(fakeTokens(chatOf(longPrompt)...), 16, 0)
	rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: hs[:1]}) // below min_blocks
	rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16})
	sendChat(t, rt, longPrompt)
	if kvDecisions(rt, kvindex.MissMinLoad) != 1 {
		t.Fatal("one shared block is coincidence, and should be counted as a miss")
	}
}

func TestWithNoEventsTheTreeDecidesAndNothingIsTokenized(t *testing.T) {
	rt, a, b := kvRouter(t, true)
	before := testutil.ToFloat64(rt.m.prefixDecisions.WithLabelValues("cache_miss_min_load"))
	sendChat(t, rt, longPrompt)
	if a.tokenized.Load()+b.tokenized.Load() != 0 {
		t.Fatal("with nothing in the index a tokenize call is a round trip for nothing")
	}
	if kvDecisions(rt, kvindex.FallbackNoIndex) != 1 {
		t.Fatal("the fallback should be counted by its reason")
	}
	if testutil.ToFloat64(rt.m.prefixDecisions.WithLabelValues("cache_miss_min_load")) != before+1 {
		t.Fatal("tier 1 should have decided, and said so in its own metric")
	}
}

func TestAFailedTokenizeFallsBackAndTheRequestStillSucceeds(t *testing.T) {
	rt, a, b := kvRouter(t, true)
	a.failTok.Store(true)
	b.failTok.Store(true)
	rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: []uint64{1}})
	rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16, Stored: []uint64{2}})
	sendChat(t, rt, longPrompt)
	if kvDecisions(rt, kvindex.FallbackTokenizeFailed) != 1 {
		t.Fatal("a failed tokenize should fall back to tier 1, counted as such")
	}
}

func TestTheTokenizeCallCarriesWhatShapesTheTokens(t *testing.T) {
	rt, a, b := kvRouter(t, false)
	rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16})
	rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16})
	rec := post(t, rt.Handler(), map[string]any{
		"model": "m", "messages": []map[string]any{{"role": "user", "content": "hi"}},
		"tools":       []any{map[string]any{"type": "function", "function": map[string]any{"name": "f"}}},
		"temperature": 0.2, "max_tokens": 5,
	}, nil)
	if rec.Code != http.StatusOK {
		t.Fatal(rec.Body.String())
	}
	var last []byte
	for _, f := range []*kvFakeVLLM{a, b} {
		if v, ok := f.lastTok.Load().([]byte); ok {
			last = v
		}
	}
	var got map[string]any
	if err := json.Unmarshal(last, &got); err != nil {
		t.Fatalf("tokenize body %q: %v", last, err)
	}
	if got["tools"] == nil || got["messages"] == nil || got["model"] != "m" {
		t.Fatalf("the chat template sees the tools, so must the tokenize call: %s", last)
	}
	if _, ok := got["temperature"]; ok {
		t.Fatalf("sampling settings do not change the tokens and should not be sent: %s", last)
	}
}

func TestAnImageMakesARequestUnindexable(t *testing.T) {
	rt, a, b := kvRouter(t, false)
	rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16})
	rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16})
	rec := post(t, rt.Handler(), map[string]any{
		"model": "m", "messages": []map[string]any{{"role": "user", "content": []any{
			map[string]any{"type": "text", "text": "what is this"},
			map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,AAAA"}},
		}}},
	}, nil)
	if rec.Code != http.StatusOK {
		t.Fatal(rec.Body.String())
	}
	if a.tokenized.Load()+b.tokenized.Load() != 0 || kvDecisions(rt, kvindex.FallbackUnsupported) != 1 {
		t.Fatal("an image's blocks hash with its identity, which the index cannot know")
	}
}

func TestAModelWithoutKVEventsIsNeverTokenized(t *testing.T) {
	rt, a, _ := kvRouter(t, false)
	rt.ApplyKV(kvevents.Update{Replica: "o-1", BlockSize: 16})
	rec := post(t, rt.Handler(), map[string]any{
		"model": "other", "messages": []map[string]any{{"role": "user", "content": longPrompt}},
	}, nil)
	if rec.Code != http.StatusOK {
		t.Fatal(rec.Body.String())
	}
	if a.tokenized.Load() != 0 {
		t.Fatal("the precise index is opt-in per model")
	}
}

func TestAReplicaThatLeavesIsForgottenByTheIndex(t *testing.T) {
	rt, _, _ := kvRouter(t, false)
	rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: []uint64{1, 2}})
	rt.SetEndpoints(t.Context(), nil)
	if len(rt.kv.Stats()) != 0 {
		t.Fatal("a gone replica's cache went with it")
	}
}

// BenchmarkTokenize is the round trip the precise index adds before placing
// a request, against a fake engine on loopback: the HTTP cost alone, without
// a real tokenizer behind it.
func BenchmarkTokenize(b *testing.B) {
	f := &kvFakeVLLM{}
	mux := http.NewServeMux()
	toks := make([]int, 2048)
	mux.HandleFunc("POST /tokenize", func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
		_ = json.NewEncoder(w).Encode(map[string]any{"tokens": toks})
	})
	f.srv = httptest.NewServer(mux)
	defer f.srv.Close()
	body, _ := kvindex.TokenizeBody("/v1/chat/completions",
		[]byte(`{"model":"m","messages":[{"role":"user","content":"`+strings.Repeat("x", 8000)+`"}]}`), "m")
	client := newKVClient()
	b.ResetTimer()
	for b.Loop() {
		got, err := kvindex.Tokenize(b.Context(), client, f.srv.URL, "", body)
		if err != nil || len(got) != 2048 {
			b.Fatal(err)
		}
		_ = kvevents.BlockHashes(got, 16, 0)
	}
}