llmgw/test/e2e/prefix_test.go

package e2e

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
)

// tenants is how many distinct system prompts the load uses.
//
// Enough that the working set does not fit in one replica's cache, which is the
// condition tier 1 exists for. With a handful of prompts and a large cache,
// scattering is free — every replica ends up holding everything — and prefix
// routing can only add constraint. Measured that way it made cache hits worse,
// which is a true statement about a scenario the feature is not for.
const tenants = 24

func systemPrompt(i int) string {
	return fmt.Sprintf("SYSTEM PROMPT %d. %s ", i,
		strings.Repeat("You are a careful assistant with detailed instructions. ", 12))
}

// askAll sends one request per tenant, repeated, and reports how many prompt
// tokens the engines served from cache.
func askAll(t *testing.T, h *harness, rounds int) (cached, prompt int64) {
	t.Helper()
	var mu sync.Mutex
	var wg sync.WaitGroup

	for round := range rounds {
		for tenant := range tenants {
			wg.Add(1)
			go func() {
				defer wg.Done()
				body, _ := json.Marshal(map[string]any{
					"model":      testModel,
					"max_tokens": 4,
					"messages": []map[string]string{
						{"role": "system", "content": systemPrompt(tenant)},
						{"role": "user", "content": fmt.Sprintf("question %d", round)},
					},
				})
				req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
					h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
				if err != nil {
					return
				}
				req.Header.Set("Content-Type", "application/json")
				resp, err := http.DefaultClient.Do(req)
				if err != nil {
					return
				}
				defer resp.Body.Close()
				raw, _ := io.ReadAll(resp.Body)

				var out struct {
					Usage struct {
						PromptTokens        int `json:"prompt_tokens"`
						PromptTokensDetails struct {
							CachedTokens int `json:"cached_tokens"`
						} `json:"prompt_tokens_details"`
					} `json:"usage"`
				}
				if json.Unmarshal(raw, &out) != nil {
					return
				}
				mu.Lock()
				cached += int64(out.Usage.PromptTokensDetails.CachedTokens)
				prompt += int64(out.Usage.PromptTokens)
				mu.Unlock()
			}()
		}
		wg.Wait()
	}
	return cached, prompt
}

// TestPrefixRoutingRaisesCacheHits is the measurement the tier exists for.
//
// Many tenants with distinct system prompts against four replicas whose caches
// are too small to hold all of them. Without prefix awareness a tenant's
// requests scatter, each replica re-prefills the same prompt, and they evict
// each other's; with it a tenant converges on the replica that already has it.
func TestPrefixRoutingRaisesCacheHits(t *testing.T) {
	tight := []engineOpt{tightPrefixCache}

	withoutPrefix := startTuned(t, 4, tight, func(c *config.Server) {
		c.Router.DefaultProfile = "long-context"
	})
	baseCached, basePrompt := askAll(t, withoutPrefix, 6)

	withPrefix := startTuned(t, 4, tight, func(c *config.Server) {
		c.Router.DefaultProfile = "long-context"
		c.Router.Prefix = config.Prefix{
			Enabled: true, MinMatch: 32, MatchRatio: 0.2,
			BalanceAbs: 8, BalanceRel: 1.5,
		}
	})
	prefixCached, prefixPrompt := askAll(t, withPrefix, 6)

	if basePrompt == 0 || prefixPrompt == 0 {
		t.Fatal("no requests completed")
	}
	baseRate := float64(baseCached) / float64(basePrompt)
	prefixRate := float64(prefixCached) / float64(prefixPrompt)
	t.Logf("cache hit rate: %.1f%% without prefix routing, %.1f%% with",
		baseRate*100, prefixRate*100)
	t.Logf("prefix decisions: %v", prefixDecisions(t, withPrefix))

	if prefixRate <= baseRate {
		t.Errorf("prefix routing did not improve cache hits: %.3f vs %.3f", prefixRate, baseRate)
	}
}

// TestPrefixRoutingDoesNotDrownOneReplica is the other half, and the reason the
// guard exists. A cache hit rate bought by sending everything to one replica is
// not a win.
func TestPrefixRoutingDoesNotDrownOneReplica(t *testing.T) {
	h := start(t, 4, func(c *config.Server) {
		c.Router.DefaultProfile = "long-context"
		c.Router.Prefix = config.Prefix{
			Enabled: true, MinMatch: 32, MatchRatio: 0.2,
			BalanceAbs: 8, BalanceRel: 1.5,
		}
	})

	// One tenant, hammered: the shape that makes affinity dangerous.
	var wg sync.WaitGroup
	for range 60 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			body, _ := json.Marshal(map[string]any{
				"model":      testModel,
				"max_tokens": 8,
				"messages": []map[string]string{
					{"role": "system", "content": systemPrompt(0)},
					{"role": "user", "content": "go"},
				},
			})
			req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost,
				h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
			req.Header.Set("Content-Type", "application/json")
			resp, err := http.DefaultClient.Do(req)
			if err != nil {
				return
			}
			_, _ = io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
		}()
	}
	wg.Wait()

	// Every replica should have done some of it. PrefixQueries counts requests
	// the engine looked up, which is one per request it was given.
	var served int
	for _, e := range h.engines {
		if e.eng.Stats().PrefixQueries > 0 {
			served++
		}
	}
	if served < 2 {
		t.Errorf("only %d of %d replicas served anything; affinity drowned the rest",
			served, len(h.engines))
	}
}

// prefixDecisions reads the router's own count of which branch each prefix
// lookup took. This is the metric that exists so tuning is not guesswork.
func prefixDecisions(t *testing.T, h *harness) map[string]int {
	t.Helper()
	mux, err := h.rt.AdminHandler()
	if err != nil {
		t.Fatal(err)
	}
	srv := httptest.NewServer(mux)
	defer srv.Close()

	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, srv.URL+"/metrics", nil)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)

	out := map[string]int{}
	for _, line := range strings.Split(string(body), "\n") {
		if !strings.HasPrefix(line, "llmgw_prefix_decisions_total{") {
			continue
		}
		name := line[strings.Index(line, `"`)+1 : strings.LastIndex(line, `"`)]
		fields := strings.Fields(line)
		n, _ := strconv.Atoi(strings.TrimSuffix(fields[len(fields)-1], ".0"))
		if f, err := strconv.ParseFloat(fields[len(fields)-1], 64); err == nil {
			n = int(f)
		}
		out[name] = n
	}
	return out
}