diff --git a/deploy/examples/server.yaml b/deploy/examples/server.yaml index 19c4167..6998421 100644 --- a/deploy/examples/server.yaml +++ b/deploy/examples/server.yaml @@ -110,6 +110,39 @@ router: failure_threshold: 3 default_profile: chat # chat | batch | long-context | debug + + # Approximate prefix routing (plan §7.5 tier 1): send a request to the replica + # that has probably already prefilled its prefix. + # + # It earns its keep when the working set does not fit in one replica's cache. + # With a handful of prompts and a large cache, scattering is free — every + # replica ends up holding everything — and this only adds constraint. Measured + # with sim-engines whose caches were too small for the working set: cache hits + # went from 22.5% to 81.7%. + # + # Watch llmgw_prefix_decisions_total to tune it. A high + # load_balance_override count means the guard is doing its job and affinity is + # not getting a chance; a high cache_miss_min_load means match_ratio or + # min_match is too strict for your traffic. + prefix: + enabled: false + # The shortest match worth acting on, in bytes. Below this the overlap is + # coincidence — two English prompts share "The". + min_match: 64 + # The share of the prompt that must match. A short shared header on a very + # long prompt saves almost no prefill. + match_ratio: 0.3 + # The imbalance guard, and both must trip before affinity is abandoned. + # Absolute alone fires on an idle cluster where one replica has two requests + # and another none; relative alone fires on a busy one where 200 against 180 + # is a rounding error. Without the guard the cache hit rate rises while p99 + # collapses. + balance_abs: 8 + balance_rel: 1.5 + # An engine evicts on its own schedule, so an old record may describe a + # cache that is gone. + ttl: 10m + max_nodes: 100000 session_header: x-session-id # affinity key; falls back to the `user` field # Ask engines for usage even when the client did not, so token accounting # works regardless of client behaviour. The extra chunk is dropped on the way diff --git a/internal/config/config.go b/internal/config/config.go index e7ccb17..d26f9d0 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -149,6 +149,11 @@ type Router struct { // DefaultProfile names the routing profile used when a model picks none. DefaultProfile string `koanf:"default_profile"` + // Prefix turns on the approximate prefix tree (plan §7.5 tier 1). Off, the + // router uses session affinity alone, which is tier 0 and already most of + // the win for multi-turn traffic. + Prefix Prefix `koanf:"prefix"` + // SessionHeader carries the session id for affinity (plan §7.5 tier 0). SessionHeader string `koanf:"session_header"` // InjectUsage asks engines for usage even when the client did not, so token @@ -181,6 +186,24 @@ type Admission struct { Interval time.Duration `koanf:"interval"` } +// Prefix configures approximate prefix routing. +type Prefix struct { + Enabled bool `koanf:"enabled"` + // MaxNodes bounds the tree, MinMatch is the shortest match worth acting on + // in bytes, and TTL is how long an entry is believed — an engine evicts on + // its own schedule, so an old record may describe a cache that is gone. + MaxNodes int `koanf:"max_nodes"` + MinMatch int `koanf:"min_match"` + TTL time.Duration `koanf:"ttl"` + + // MatchRatio is the share of the prompt that must match to be worth acting + // on. BalanceAbs and BalanceRel are the imbalance guard, and both must trip + // before affinity is abandoned. + MatchRatio float64 `koanf:"match_ratio"` + BalanceAbs float64 `koanf:"balance_abs"` + BalanceRel float64 `koanf:"balance_rel"` +} + // Band is one priority class. type Band struct { Name string `koanf:"name"` diff --git a/internal/router/metrics.go b/internal/router/metrics.go index 8558514..9565e8c 100644 --- a/internal/router/metrics.go +++ b/internal/router/metrics.go @@ -6,6 +6,7 @@ import ( "github.com/prometheus/client_golang/prometheus" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix" "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" "gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler" ) @@ -26,10 +27,17 @@ type metrics struct { admissionWait *prometheus.HistogramVec admissionRejected *prometheus.CounterVec + + prefixDecisions *prometheus.CounterVec } func newMetrics() *metrics { return &metrics{ + prefixDecisions: prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "llmgw_prefix_decisions_total", + Help: "Prefix routing decisions by branch: cache_hit, cache_miss_min_load, " + + "load_balance_override, no_tree, no_healthy_worker.", + }, []string{"decision"}), requests: prometheus.NewCounterVec(prometheus.CounterOpts{ Name: "llmgw_requests_total", Help: "Requests handled, by model, pool and response status.", @@ -86,10 +94,18 @@ func (m *metrics) collectors() []prometheus.Collector { return []prometheus.Collector{ m.requests, m.duration, m.ttft, m.tokens, m.scorerValue, m.filtered, m.considered, - m.admissionWait, m.admissionRejected, + m.admissionWait, m.admissionRejected, m.prefixDecisions, } } +// observePrefix records which branch a prefix decision took. +// +// Without it, "affinity is not helping" and "affinity never applied" look +// identical from the outside, and the thresholds can only be tuned by guessing. +func (m *metrics) observePrefix(d prefix.Decision) { + m.prefixDecisions.WithLabelValues(string(d)).Inc() +} + func (m *metrics) observeAdmission(band, model string, wait time.Duration) { if band == "" { band = "none" diff --git a/internal/router/prefix/scorer.go b/internal/router/prefix/scorer.go new file mode 100644 index 0000000..f60e739 --- /dev/null +++ b/internal/router/prefix/scorer.go @@ -0,0 +1,126 @@ +package prefix + +import ( + "math" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// Decision is which branch a prefix lookup took. +// +// Recorded as a metric because tuning prefix routing without it is guesswork: +// a cache hit rate says nothing about how often affinity was overridden, and +// "affinity is not helping" and "affinity never applied" look identical from +// the outside. +type Decision string + +// The branches. +const ( + Hit Decision = "cache_hit" + MissMinLoad Decision = "cache_miss_min_load" + LoadOverride Decision = "load_balance_override" + NoTree Decision = "no_tree" + NoCandidates Decision = "no_healthy_worker" +) + +// Thresholds decide when the prefix is worth following. +type Thresholds struct { + // MatchRatio is the share of the prompt that must match. Below it, the + // overlap is not enough to save a meaningful prefill. + MatchRatio float64 + // BalanceAbs and BalanceRel are the imbalance guard, and both must trip + // before it does. Without this the cache hit rate rises while p99 latency + // collapses: affinity keeps sending a hot prefix to one replica long after + // that replica has stopped being able to serve it quickly. + // + // Absolute alone would fire on an idle cluster where one replica has two + // requests and another none. Relative alone would fire on a busy one where + // 200 against 180 is a rounding error. Together they mean "meaningfully + // more loaded, and by enough to matter". + BalanceAbs float64 + BalanceRel float64 +} + +// DefaultThresholds returns values that work without measurement. +func DefaultThresholds() Thresholds { + return Thresholds{MatchRatio: 0.3, BalanceAbs: 8, BalanceRel: 1.5} +} + +func (t *Thresholds) normalize() { + d := DefaultThresholds() + if t.MatchRatio <= 0 { + t.MatchRatio = d.MatchRatio + } + if t.BalanceAbs <= 0 { + t.BalanceAbs = d.BalanceAbs + } + if t.BalanceRel <= 1 { + t.BalanceRel = d.BalanceRel + } +} + +// Choose picks the replica a prompt should go to, and says why. +// +// It returns an index into eps, or -1 when the caller should decide on load +// alone — which is the common case and not a failure. +func Choose(t *Tree, prompt string, eps []endpoint.Snapshot, th Thresholds) (int, Decision) { + th.normalize() + if len(eps) == 0 { + return -1, NoCandidates + } + if t == nil || prompt == "" { + return -1, NoTree + } + + // The guard runs before the lookup, not after. Checking afterwards would + // mean paying for a match only to discard it, and — worse — would invite + // the shape where a cache hit quietly wins because somebody moved the + // check. + if imbalanced(eps, th) { + return -1, LoadOverride + } + + replica, matched := t.Match(prompt) + if replica == "" || float64(matched)/float64(len(prompt)) < th.MatchRatio { + return -1, MissMinLoad + } + for i, e := range eps { + if e.ID == replica { + return i, Hit + } + } + // The match names a replica that is no longer a candidate: gone, draining, + // or its breaker is open. That is a miss, not an error. + return -1, MissMinLoad +} + +// imbalanced reports whether load is lopsided enough to ignore affinity. +func imbalanced(eps []endpoint.Snapshot, th Thresholds) bool { + if len(eps) < 2 { + return false + } + maxL, minL := math.Inf(-1), math.Inf(1) + for _, e := range eps { + l := load(e) + maxL = math.Max(maxL, l) + minL = math.Min(minL, l) + } + if maxL-minL <= th.BalanceAbs { + return false + } + // A ratio against zero is infinite, and an idle replica is exactly where + // the next request should go, so treat it as tripping the guard. + if minL <= 0 { + return true + } + return maxL/minL > th.BalanceRel +} + +// load is what the guard compares. +// +// Queue depth plus what the router has dispatched but not seen finish: the +// metric alone lags by a scrape interval, which is long enough for affinity to +// pile a burst onto one replica before anything notices. +func load(e endpoint.Snapshot) float64 { + return float64(e.Stats.Waiting + e.Stats.Running + e.Inflight) +} diff --git a/internal/router/prefix/scorer_test.go b/internal/router/prefix/scorer_test.go new file mode 100644 index 0000000..dfa71b1 --- /dev/null +++ b/internal/router/prefix/scorer_test.go @@ -0,0 +1,151 @@ +package prefix + +import ( + "strings" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +func ep(id string, waiting, running, inflight int) endpoint.Snapshot { + return endpoint.Snapshot{ + Info: endpoint.Info{ID: id}, + Stats: endpoint.Stats{Waiting: waiting, Running: running}, + Inflight: inflight, + } +} + +const prompt = "You are a helpful assistant answering questions about distributed systems. " + +func TestAMatchIsFollowedWhenLoadIsEven(t *testing.T) { + tr := New(Config{MinMatch: 16}) + tr.Observe(prompt, "b") + + eps := []endpoint.Snapshot{ep("a", 1, 1, 0), ep("b", 1, 1, 0)} + i, d := Choose(tr, prompt+"What is a quorum?", eps, DefaultThresholds()) + if d != Hit || i != 1 { + t.Errorf("index %d, decision %s; want a hit on b", i, d) + } +} + +func TestAffinityIsAbandonedWhenLoadIsLopsided(t *testing.T) { + // The guard the plan insists on. Without it the cache hit rate rises while + // p99 collapses: affinity keeps feeding a hot prefix to one replica long + // after it stopped being able to serve it quickly. + tr := New(Config{MinMatch: 16}) + tr.Observe(prompt, "b") + + eps := []endpoint.Snapshot{ep("a", 0, 1, 0), ep("b", 40, 8, 4)} + i, d := Choose(tr, prompt+"What is a quorum?", eps, DefaultThresholds()) + if d != LoadOverride { + t.Errorf("decision %s; the guard should have fired", d) + } + if i != -1 { + t.Errorf("index %d; the caller decides on load alone", i) + } +} + +func TestBothHalvesOfTheGuardMustTrip(t *testing.T) { + tr := New(Config{MinMatch: 16}) + tr.Observe(prompt, "b") + full := prompt + "What is a quorum?" + th := Thresholds{MatchRatio: 0.1, BalanceAbs: 8, BalanceRel: 1.5} + + cases := []struct { + name string + eps []endpoint.Snapshot + want Decision + }{ + { + // Relative is huge, absolute is tiny: an idle cluster where one + // replica has two requests and another none. + "a big ratio over a small difference", + []endpoint.Snapshot{ep("a", 0, 0, 0), ep("b", 2, 0, 0)}, + Hit, + }, + { + // Absolute is large, relative is nothing: a busy cluster where 200 + // against 180 is a rounding error. + "a big difference at a small ratio", + []endpoint.Snapshot{ep("a", 180, 0, 0), ep("b", 200, 0, 0)}, + Hit, + }, + { + "both, which is what the guard is for", + []endpoint.Snapshot{ep("a", 2, 0, 0), ep("b", 60, 0, 0)}, + LoadOverride, + }, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if _, d := Choose(tr, full, c.eps, th); d != c.want { + t.Errorf("decision %s, want %s", d, c.want) + } + }) + } +} + +func TestAnIdleReplicaAlwaysWinsOverAffinity(t *testing.T) { + // A ratio against zero is infinite, and an idle replica is exactly where + // the next request should go. + tr := New(Config{MinMatch: 16}) + tr.Observe(prompt, "b") + + eps := []endpoint.Snapshot{ep("a", 0, 0, 0), ep("b", 20, 0, 0)} + if _, d := Choose(tr, prompt+"x", eps, DefaultThresholds()); d != LoadOverride { + t.Errorf("decision %s; an idle replica should win", d) + } +} + +func TestASmallMatchRatioIsAMiss(t *testing.T) { + // A short shared header on a very long prompt saves almost no prefill, and + // pretending otherwise trades load balance for nothing. + tr := New(Config{MinMatch: 8}) + tr.Observe("SYSTEM: be brief. ", "b") + + long := "SYSTEM: be brief. " + strings.Repeat("unique content here. ", 200) + eps := []endpoint.Snapshot{ep("a", 1, 0, 0), ep("b", 1, 0, 0)} + if _, d := Choose(tr, long, eps, Thresholds{MatchRatio: 0.5}); d != MissMinLoad { + t.Errorf("decision %s, want a miss", d) + } +} + +func TestAMatchOnAReplicaThatIsGoneIsAMiss(t *testing.T) { + // The tree outlives a replica leaving the candidate set — drained, breaker + // open, removed. That is a miss, not an error. + tr := New(Config{MinMatch: 16}) + tr.Observe(prompt, "vanished") + + eps := []endpoint.Snapshot{ep("a", 1, 0, 0), ep("b", 1, 0, 0)} + if i, d := Choose(tr, prompt+"x", eps, DefaultThresholds()); d != MissMinLoad || i != -1 { + t.Errorf("index %d decision %s", i, d) + } +} + +func TestTheDegenerateCasesAreNamed(t *testing.T) { + // Every branch is a distinct decision, because "affinity is not helping" + // and "affinity never applied" look identical from the outside otherwise. + tr := New(Config{MinMatch: 4}) + if _, d := Choose(tr, "p", nil, DefaultThresholds()); d != NoCandidates { + t.Errorf("no endpoints: %s", d) + } + eps := []endpoint.Snapshot{ep("a", 0, 0, 0)} + if _, d := Choose(nil, "p", eps, DefaultThresholds()); d != NoTree { + t.Errorf("no tree: %s", d) + } + if _, d := Choose(tr, "", eps, DefaultThresholds()); d != NoTree { + t.Errorf("no prompt: %s", d) + } +} + +func TestOneCandidateNeverTripsTheGuard(t *testing.T) { + // There is nothing to balance against, so overriding affinity would throw + // away a cache hit for no reason. + tr := New(Config{MinMatch: 16}) + tr.Observe(prompt, "only") + + eps := []endpoint.Snapshot{ep("only", 500, 500, 500)} + if _, d := Choose(tr, prompt+"x", eps, DefaultThresholds()); d != Hit { + t.Errorf("decision %s, want a hit", d) + } +} diff --git a/internal/router/prefix/tree.go b/internal/router/prefix/tree.go new file mode 100644 index 0000000..0e165f9 --- /dev/null +++ b/internal/router/prefix/tree.go @@ -0,0 +1,317 @@ +// Package prefix estimates which replica already holds the prefix of a request +// (plan §7.5 tier 1). +// +// It is an estimate and says so. The router does not know an engine's real +// cache; it knows what it sent where, which is a good enough proxy because +// engines cache what they are sent. Nothing here is ever authoritative — a +// wrong guess costs a prefill, not a wrong answer. +// +// Tokens are deliberately absent. Matching on bytes needs no tokenizer, and +// needing the engine's exact tokenizer is what makes tier 2 expensive enough to +// be a separate project. +package prefix + +import ( + "sync" + "time" +) + +// Config tunes the tree. +type Config struct { + // MaxNodes bounds the tree. Past it the least recently used branches go, or + // a long-running router accumulates every prompt it has ever seen. + MaxNodes int + // MinMatch is the shortest match worth acting on, in bytes. Below it the + // overlap is coincidence — two English prompts share "The" — and treating + // coincidence as affinity is how prefix routing drowns one replica. + MinMatch int + // TTL is how long an entry is believed. An engine evicts on its own + // schedule, so an old record describes a cache that may be long gone. + TTL time.Duration +} + +// DefaultConfig returns values that work without measurement. +func DefaultConfig() Config { + return Config{MaxNodes: 100_000, MinMatch: 64, TTL: 10 * time.Minute} +} + +func (c *Config) normalize() { + d := DefaultConfig() + if c.MaxNodes <= 0 { + c.MaxNodes = d.MaxNodes + } + if c.MinMatch <= 0 { + c.MinMatch = d.MinMatch + } + if c.TTL <= 0 { + c.TTL = d.TTL + } +} + +// seen is how often and how recently a replica was sent this prefix. +// +// The count is what makes the estimate stable. Picking the most *recent* +// replica sounds right and is not: with a randomizing picker every replica ends +// up in the node eventually, the freshest one flips with every request, and the +// tree spends its time chasing whatever the picker last did. Measured: prefix +// routing made cache hits worse than no prefix routing at all, 28.7% against +// 32.8%. +type seen struct { + count int + last time.Time +} + +// Tree is a radix tree of request prefixes, per replica. +type Tree struct { + cfg Config + now func() time.Time + + mu sync.Mutex + root *node + nodes int +} + +type node struct { + // label is this node's slice of the prefix, shared by everything below it. + label string + children map[byte]*node + // replicas is which endpoints were sent a request passing through this + // node — not only ending at it. Recorded along the whole path, because a + // prompt that diverges partway through a label still shares everything + // before the divergence, and shared leading bytes are exactly what a prefix + // cache holds. + // + // Per replica rather than a set, so one replica's entry can expire without + // taking the others with it. + replicas map[string]*seen + // touched drives eviction. + touched time.Time +} + +// New builds an empty tree. +func New(cfg Config) *Tree { + cfg.normalize() + return &Tree{cfg: cfg, now: time.Now, root: &node{children: map[byte]*node{}}, nodes: 1} +} + +// Match reports the replica with the longest usable prefix in common, and how +// many bytes matched. +// +// Empty when nothing matched far enough to be worth acting on — which is a +// normal answer, not a failure, and the caller must be able to proceed on load +// alone when it comes back. +func (t *Tree) Match(prompt string) (replica string, matched int) { + if prompt == "" { + return "", 0 + } + now := t.now() + + t.mu.Lock() + defer t.mu.Unlock() + + cur := t.root + consumed := 0 + best, bestLen := "", 0 + + for consumed < len(prompt) { + child, ok := cur.children[prompt[consumed]] + if !ok { + break + } + shared := commonPrefix(child.label, prompt[consumed:]) + if shared == 0 { + break + } + consumed += shared + cur = child + child.touched = now + + // A partial label counts. Everything recorded at or below this node + // began with the bytes just matched, so a replica here holds that much + // of this prompt however the rest diverges — and a cache hit is about + // shared leading bytes, not about where two prompts stop agreeing. + // + // Requiring the whole label to match made this tree inert: the first + // request for a prompt becomes one long label, the second diverges in + // its last few bytes, and every lookup after that returned nothing. + // Measured as 144 misses and zero hits before this was fixed. + if r, ok := cur.owner(now, t.cfg.TTL); ok && consumed > bestLen { + best, bestLen = r, consumed + } + if shared < len(child.label) { + break + } + } + + if bestLen < t.cfg.MinMatch { + return "", 0 + } + return best, bestLen +} + +// Observe records that a replica was sent this prompt. +func (t *Tree) Observe(prompt, replica string) { + if prompt == "" || replica == "" { + return + } + now := t.now() + + t.mu.Lock() + defer t.mu.Unlock() + + cur := t.root + rest := prompt + for len(rest) > 0 { + child, ok := cur.children[rest[0]] + if !ok { + leaf := &node{label: rest, children: map[byte]*node{}, touched: now} + cur.children[rest[0]] = leaf + t.nodes++ + cur = leaf + break + } + + shared := commonPrefix(child.label, rest) + if shared < len(child.label) { + // The existing label diverges partway, so it splits into a shared + // head and what used to hang below it. + tail := &node{ + label: child.label[shared:], children: child.children, + replicas: child.replicas, touched: child.touched, + } + child.label = child.label[:shared] + child.children = map[byte]*node{tail.label[0]: tail} + child.replicas = nil + t.nodes++ + } + child.touched = now + cur = child + // Marked on the way through, not only at the end: a later prompt that + // diverges here still shares everything above, and that is what makes + // a partial match meaningful. + cur.record(replica, now) + rest = rest[shared:] + } + + cur.record(replica, now) + + if t.nodes > t.cfg.MaxNodes { + t.evictLocked(now) + } +} + +// Forget drops a replica from the tree, for one that has gone away. +// +// Leaving it would have the router steering traffic towards a cache that no +// longer exists, and a replica id is reused by nothing, so the entries would +// never be corrected — only expire. +func (t *Tree) Forget(replica string) { + t.mu.Lock() + defer t.mu.Unlock() + forget(t.root, replica) +} + +func forget(n *node, replica string) { + delete(n.replicas, replica) + for _, c := range n.children { + forget(c, replica) + } +} + +// Size reports how many nodes the tree holds, for metrics and tests. +func (t *Tree) Size() int { + t.mu.Lock() + defer t.mu.Unlock() + return t.nodes +} + +// record marks that a replica was sent something passing through this node. +func (n *node) record(replica string, now time.Time) { + if n.replicas == nil { + n.replicas = map[string]*seen{} + } + if e, ok := n.replicas[replica]; ok { + e.count++ + e.last = now + } else { + n.replicas[replica] = &seen{count: 1, last: now} + } + n.touched = now +} + +// owner returns the replica most likely to hold this prefix. +// +// Most often sent, not most recently: a count is stable under a randomizing +// picker, and it still moves when traffic genuinely does. Ties go to the more +// recent one, which is the only thing recency is good for here. +func (n *node) owner(now time.Time, ttl time.Duration) (string, bool) { + var best string + var bestSeen *seen + for r, s := range n.replicas { + if now.Sub(s.last) > ttl { + continue + } + switch { + case bestSeen == nil, + s.count > bestSeen.count, + s.count == bestSeen.count && s.last.After(bestSeen.last): + best, bestSeen = r, s + } + } + return best, best != "" +} + +// evictLocked drops the coldest half of the tree's branches. +// +// Half rather than just enough: trimming to exactly the limit means evicting +// again on the very next insert, and the walk is the expensive part. +func (t *Tree) evictLocked(now time.Time) { + cutoff := t.coldestHalf(now) + t.nodes = 1 + prune(t.root, cutoff) +} + +// coldestHalf finds the touch time that splits the tree in two. +func (t *Tree) coldestHalf(now time.Time) time.Time { + var times []time.Time + collect(t.root, ×) + if len(times) == 0 { + return now + } + // A partial sort would do, but the tree is evicted rarely and this is + // simpler to be sure of. + for i := 1; i < len(times); i++ { + for j := i; j > 0 && times[j].Before(times[j-1]); j-- { + times[j], times[j-1] = times[j-1], times[j] + } + } + return times[len(times)/2] +} + +func collect(n *node, out *[]time.Time) { + for _, c := range n.children { + *out = append(*out, c.touched) + collect(c, out) + } +} + +// prune removes branches untouched since the cutoff, returning what remains. +func prune(n *node, cutoff time.Time) int { + kept := 0 + for b, c := range n.children { + if c.touched.Before(cutoff) { + delete(n.children, b) + continue + } + kept += 1 + prune(c, cutoff) + } + return kept +} + +func commonPrefix(a, b string) int { + n := min(len(a), len(b)) + i := 0 + for i < n && a[i] == b[i] { + i++ + } + return i +} diff --git a/internal/router/prefix/tree_test.go b/internal/router/prefix/tree_test.go new file mode 100644 index 0000000..1e71198 --- /dev/null +++ b/internal/router/prefix/tree_test.go @@ -0,0 +1,165 @@ +package prefix + +import ( + "fmt" + "strings" + "sync" + "testing" + "time" +) + +func tree(t *testing.T, cfg Config) *Tree { + t.Helper() + return New(cfg) +} + +func TestTheLongestMatchWins(t *testing.T) { + // Two replicas have seen overlapping prefixes; the one that has seen more + // of this prompt is the one that has more of it cached. + tr := tree(t, Config{MinMatch: 4}) + tr.Observe("You are a helpful assistant. ", "a") + tr.Observe("You are a helpful assistant. Answer in French. ", "b") + + replica, n := tr.Match("You are a helpful assistant. Answer in French. What is 2+2?") + if replica != "b" { + t.Errorf("matched %q with %d bytes; the longer prefix belongs to b", replica, n) + } + if n < len("You are a helpful assistant. Answer in French. ") { + t.Errorf("matched only %d bytes", n) + } +} + +func TestAShortOverlapIsNotAffinity(t *testing.T) { + // Two English prompts share "The". Treating coincidence as affinity is how + // prefix routing drowns one replica. + tr := tree(t, Config{MinMatch: 64}) + tr.Observe("The quick brown fox", "a") + + if replica, n := tr.Match("The slow green turtle"); replica != "" { + t.Errorf("a %d-byte coincidence was treated as a match to %q", n, replica) + } +} + +func TestSplittingAnExistingBranch(t *testing.T) { + // The case a radix tree exists for: a new prompt diverges partway through + // a label that was one node. + tr := tree(t, Config{MinMatch: 4}) + tr.Observe("abcdefgh", "a") + tr.Observe("abcdXYZW", "b") + + if r, _ := tr.Match("abcdefgh and more"); r != "a" { + t.Errorf("the original branch matched %q", r) + } + if r, _ := tr.Match("abcdXYZW and more"); r != "b" { + t.Errorf("the new branch matched %q", r) + } + // The shared head belongs to both, because both were sent something + // beginning with it — and a prefix cache holds leading bytes, not whole + // prompts. An earlier version treated this as no match, which made the + // whole tree inert in practice. + if r, n := tr.Match("abcd"); r == "" || n != 4 { + t.Errorf("the shared head matched %q with %d bytes; both replicas hold it", r, n) + } +} + +func TestAPartialLabelIsStillAMatch(t *testing.T) { + // The correction that made this feature work at all. + // + // The first request for a prompt becomes one long label. The second differs + // in its last few bytes — a new question after the same system prompt, + // which is the single most common shape there is. Requiring the whole label + // to match meant every lookup after the first returned nothing: measured as + // 144 misses and zero hits. + // + // What matters is the shared head, because that is what the engine cached. + tr := tree(t, Config{MinMatch: 2}) + tr.Observe("abcdefghij", "a") + + r, n := tr.Match("abcdeZZZZ") + if r != "a" { + t.Errorf("a partial match returned %q", r) + } + if n != 5 { + t.Errorf("matched %d bytes, want the 5 they share", n) + } +} + +func TestEntriesExpire(t *testing.T) { + // An engine evicts on its own schedule, so an old record describes a cache + // that may be long gone. Believing it forever sends traffic towards + // something that is not there. + now := time.Now() + tr := New(Config{MinMatch: 4, TTL: time.Minute}) + tr.now = func() time.Time { return now } + + tr.Observe("shared prefix here", "a") + if r, _ := tr.Match("shared prefix here and more"); r != "a" { + t.Fatalf("fresh entry matched %q", r) + } + + now = now.Add(2 * time.Minute) + if r, _ := tr.Match("shared prefix here and more"); r != "" { + t.Errorf("a stale entry still matched %q", r) + } +} + +func TestForgettingAReplica(t *testing.T) { + // A replica that has gone away must stop attracting traffic. Nothing reuses + // a replica id, so the entries would otherwise only ever expire. + tr := tree(t, Config{MinMatch: 4}) + tr.Observe("a long enough shared prefix", "a") + tr.Observe("a long enough shared prefix", "b") + + tr.Forget("a") + r, _ := tr.Match("a long enough shared prefix and more") + if r == "a" { + t.Error("a forgotten replica still matches") + } + if r != "b" { + t.Errorf("forgetting one replica lost the other: %q", r) + } +} + +func TestTheTreeIsBounded(t *testing.T) { + // A router that runs for months would otherwise accumulate every prompt it + // has ever seen. + tr := tree(t, Config{MinMatch: 4, MaxNodes: 200}) + for i := range 2000 { + tr.Observe(fmt.Sprintf("prompt number %d with some trailing text", i), "a") + } + if got := tr.Size(); got > 400 { + t.Errorf("the tree grew to %d nodes against a limit of 200", got) + } + // And it still works afterwards. + tr.Observe("a fresh prompt after all that eviction", "b") + if r, _ := tr.Match("a fresh prompt after all that eviction, continued"); r != "b" { + t.Errorf("the tree stopped working after eviction: %q", r) + } +} + +func TestConcurrentUse(t *testing.T) { + tr := tree(t, Config{MinMatch: 4, MaxNodes: 500}) + var wg sync.WaitGroup + for i := range 32 { + wg.Add(1) + go func() { + defer wg.Done() + p := strings.Repeat("x", 40) + fmt.Sprint(i%4) + for range 50 { + tr.Observe(p, fmt.Sprintf("r%d", i%4)) + tr.Match(p) + tr.Forget("r0") + } + }() + } + wg.Wait() +} + +func TestEmptyInputsAreHarmless(t *testing.T) { + tr := tree(t, Config{}) + tr.Observe("", "a") + tr.Observe("something", "") + if r, n := tr.Match(""); r != "" || n != 0 { + t.Errorf("an empty prompt matched %q/%d", r, n) + } +} diff --git a/internal/router/router.go b/internal/router/router.go index 52eff47..2a09dbf 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -23,6 +23,7 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/router/admission" "gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker" "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix" "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" "gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler" "gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape" @@ -59,6 +60,9 @@ type Router struct { acct *accounting.Accountant recorder *accounting.Recorder + // prefix is nil unless approximate prefix routing is turned on. + prefix *scheduler.PrefixScorer + // clusterView is nil until the control plane hands it over, and stays nil // in a router-only deployment. clusterView web.Source @@ -102,6 +106,27 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { KVCeiling: cfg.Router.Scoring.KVCeiling, SessionLoadFactor: cfg.Router.Scoring.SessionLoadFactor, } + // The metrics come first because the prefix scorer reports into them, and + // the decision metric is the whole difference between tuning this and + // guessing at it. + metrics := newMetrics() + var prefixScorer *scheduler.PrefixScorer + if cfg.Router.Prefix.Enabled { + prefixScorer = scheduler.NewPrefixScorer( + prefix.New(prefix.Config{ + MaxNodes: cfg.Router.Prefix.MaxNodes, + MinMatch: cfg.Router.Prefix.MinMatch, + TTL: cfg.Router.Prefix.TTL, + }), + prefix.Thresholds{ + MatchRatio: cfg.Router.Prefix.MatchRatio, + BalanceAbs: cfg.Router.Prefix.BalanceAbs, + BalanceRel: cfg.Router.Prefix.BalanceRel, + }, + metrics.observePrefix) + sc.Prefix = prefixScorer + } + profiles := map[string]*scheduler.Profile{} for _, name := range []string{"chat", "batch", "long-context", "debug"} { profiles[name] = scheduler.Builtin(name, sc) @@ -125,7 +150,8 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { reg: reg, profiles: profiles, models: models, - m: newMetrics(), + prefix: prefixScorer, + m: metrics, px: proxy.New(proxy.Config{ DialTimeout: cfg.Router.Timeouts.Dial, ResponseHeaderTimeout: cfg.Router.Timeouts.ResponseHeader, @@ -297,6 +323,23 @@ func (r *Router) Registry() *endpoint.Registry { return r.reg } // — so a replica that was already serving does not start from cold just // because another one moved. func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) { + if r.prefix != nil { + // A replica that has gone takes its cache with it. Leaving its entries + // in the tree would steer traffic towards something that is not there, + // and since nothing reuses a replica id the entries would never be + // corrected — only expire, minutes later. + before := map[string]struct{}{} + for _, s := range r.reg.Snapshots("") { + before[s.ID] = struct{}{} + } + for _, i := range infos { + delete(before, i.ID) + } + for id := range before { + r.prefix.Forget(id) + } + } + r.reg.Replace(infos) r.poller.Sync(ctx) } @@ -425,6 +468,14 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool, outcome: out, err: perr} r.m.observeScores(dec) + // Recorded after the decision, not during it: the scorer proposes and + // the picker disposes, so recording the proposal would fill the tree + // with replicas that were never sent anything — the one thing that + // would make the estimate worse than no estimate at all. + if r.prefix != nil { + r.prefix.Observe(req.PromptPrefix, dec.Endpoint.ID) + } + if perr == nil || out.Wrote || !out.Retryable { return res } diff --git a/internal/router/scheduler/prefix.go b/internal/router/scheduler/prefix.go new file mode 100644 index 0000000..3c64ef1 --- /dev/null +++ b/internal/router/scheduler/prefix.go @@ -0,0 +1,65 @@ +package scheduler + +import ( + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix" +) + +// PrefixScorer steers a request towards the replica that has probably already +// prefilled its prefix (plan §7.5 tier 1). +// +// A scorer rather than a filter, deliberately. A filter would make affinity a +// rule, and a rule that sends every request carrying a popular system prompt to +// one replica is how prefix routing produces a beautiful cache hit rate and an +// unusable p99. As a scorer it is one voice among several, and the queue and +// inflight scorers can outvote it — on top of the imbalance guard inside the +// lookup itself, which stops it speaking at all when load has gone lopsided. +type PrefixScorer struct { + tree *prefix.Tree + thresholds prefix.Thresholds + // observe reports which branch the decision took, for the metric without + // which tuning this is guesswork. + observe func(prefix.Decision) +} + +// NewPrefixScorer builds the scorer. A nil tree makes it inert, which is what +// a model with prefix routing turned off should get. +func NewPrefixScorer(t *prefix.Tree, th prefix.Thresholds, observe func(prefix.Decision)) *PrefixScorer { + if observe == nil { + observe = func(prefix.Decision) {} + } + return &PrefixScorer{tree: t, thresholds: th, observe: observe} +} + +// Name identifies the scorer in weights and metrics. +func (*PrefixScorer) Name() string { return "prefix" } + +// Score gives the matched replica 1 and everything else 0. +func (p *PrefixScorer) Score(c *Ctx, eps []endpoint.Snapshot) []float64 { + out := make([]float64, len(eps)) + idx, decision := prefix.Choose(p.tree, c.Req.PromptPrefix, eps, p.thresholds) + p.observe(decision) + if idx >= 0 { + out[idx] = 1 + } + return out +} + +// Observe records where a request was actually sent. +// +// After the decision rather than during it: the scorer proposes, the picker +// disposes, and recording the proposal would fill the tree with replicas that +// were never sent anything — which is the one thing that would make the +// estimate worse than no estimate at all. +func (p *PrefixScorer) Observe(promptPrefix, replicaID string) { + if p.tree != nil { + p.tree.Observe(promptPrefix, replicaID) + } +} + +// Forget drops a replica that has gone away. +func (p *PrefixScorer) Forget(replicaID string) { + if p.tree != nil { + p.tree.Forget(replicaID) + } +} diff --git a/internal/router/scheduler/profiles.go b/internal/router/scheduler/profiles.go index eabfcfd..bcedf4b 100644 --- a/internal/router/scheduler/profiles.go +++ b/internal/router/scheduler/profiles.go @@ -10,6 +10,11 @@ type ProfileConfig struct { KVCeiling float64 // SessionLoadFactor is the load ceiling multiplier for session affinity. SessionLoadFactor float64 + + // Prefix, when set, adds the approximate prefix tree to the profiles that + // benefit from it. Nil leaves them exactly as they were — session affinity + // alone, which is tier 0 and already most of the win. + Prefix *PrefixScorer } // DefaultProfileConfig returns sensible starting values. @@ -46,17 +51,30 @@ func Chat(c ProfileConfig) *Profile { return &Profile{ Name: "chat", Filters: []Filter{BreakerOpen(), PoolPriority()}, - Scorers: []WeightedScorer{ + Scorers: c.withPrefix([]WeightedScorer{ {Session(c.SessionLoadFactor), 3.0}, {Queue(c.TargetWaiting), 1.0}, {KV(c.KVCeiling), 1.0}, {Inflight(c.TargetInflight), 1.0}, {Weight(), 0.5}, - }, + }, 2.0), Picker: P2C(), } } +// withPrefix appends the prefix scorer when one is configured. +// +// Weighted below session affinity on purpose. A session id is a statement from +// the client that these requests belong together; a prefix match is the +// router's guess that they might. When the two disagree the client is more +// likely to be right. +func (c ProfileConfig) withPrefix(scorers []WeightedScorer, weight float64) []WeightedScorer { + if c.Prefix == nil { + return scorers + } + return append(scorers, WeightedScorer{c.Prefix, weight}) +} + // Batch is the profile for bulk work: throughput matters, TTFT does not, so // finding the shortest queue outranks stickiness. func Batch(c ProfileConfig) *Profile { @@ -82,12 +100,14 @@ func LongContext(c ProfileConfig) *Profile { return &Profile{ Name: "long-context", Filters: []Filter{BreakerOpen(), PoolPriority()}, - Scorers: []WeightedScorer{ + // Prefix matching earns more here than anywhere else: a long prompt is + // mostly prefill, and prefill is what a cache hit skips. + Scorers: c.withPrefix([]WeightedScorer{ {KV(c.KVCeiling), 3.0}, {Session(c.SessionLoadFactor), 2.0}, {Queue(c.TargetWaiting), 1.0}, {Inflight(c.TargetInflight), 1.0}, - }, + }, 3.0), Picker: Softmax(0.2), } } diff --git a/llmgw-plan.md b/llmgw-plan.md index f80ff12..18e7575 100644 --- a/llmgw-plan.md +++ b/llmgw-plan.md @@ -1424,7 +1424,11 @@ The launch path is now built too: no identity to consult, which let any client promote itself. - [ ] Budget fallback - [ ] Routing profiles from configuration plus the `x-llmgw-profile` header -- [ ] Approximate radix tree (§7.5 tier 1), with its thresholds and decision metrics +- [x] Approximate radix tree (§7.5 tier 1), with its thresholds and the + `llmgw_prefix_decisions_total` breakdown. Measured against sim-engines + whose caches were too small for the working set: 22.5% → 81.7% cache + hits. It does nothing when the working set fits in one replica's cache, + which is worth knowing rather than assuming. - [ ] Fallback models and pool spillover - [x] `ctl bench`: a load ramp that reports capacity as the last concurrency level meeting the SLO, plus the queue depth and KV usage observed there, diff --git a/test/e2e/harness_test.go b/test/e2e/harness_test.go index 09a6d21..407e3f7 100644 --- a/test/e2e/harness_test.go +++ b/test/e2e/harness_test.go @@ -68,12 +68,27 @@ func withLoadedEngines() func(*config.Server) { return nil // marker; see start() } +// tightPrefixCache shrinks the engines' prefix cache so the working set does +// not fit in any one of them. +// +// That condition is the whole reason tier 1 exists. With a cache large enough +// for every prompt, scattering is free — every replica ends up holding +// everything — and prefix routing can only add constraint. It is worth knowing +// the feature does nothing there rather than assuming it helps everywhere. +var tightPrefixCache = func(c *sim.Config) { c.PrefixCacheBlocks = 96 } + // start brings up n sim-engines behind a router. // // A nil tune function is the "small engines" marker, which keeps the call sites // readable without a second variadic parameter. func start(t *testing.T, n int, tune ...func(*config.Server)) *harness { t.Helper() + return startTuned(t, n, nil, tune...) +} + +// startTuned is start with the sim-engines themselves adjusted. +func startTuned(t *testing.T, n int, engineOpts []engineOpt, tune ...func(*config.Server)) *harness { + t.Helper() small := false kept := tune[:0] @@ -99,6 +114,9 @@ func start(t *testing.T, n int, tune ...func(*config.Server)) *harness { sc.DecodeTokensPerSec = 800 sc.BatchContention = 0.005 } + for _, opt := range engineOpts { + opt(&sc) + } e := sim.New(sc) srv := httptest.NewServer(sim.NewServer(e, slog.New(slog.DiscardHandler))) h.engines = append(h.engines, engine{eng: e, srv: srv}) diff --git a/test/e2e/prefix_test.go b/test/e2e/prefix_test.go new file mode 100644 index 0000000..91f27b9 --- /dev/null +++ b/test/e2e/prefix_test.go @@ -0,0 +1,211 @@ +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 +}