llmgw/internal/router/kvroute.go

package router

import (
	"context"
	"errors"
	"net/http"
	"time"

	"github.com/prometheus/client_golang/prometheus"

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

// ApplyKV folds in a change to a replica's KV cache, as its node's agent
// forwarded it. The control plane calls it for every KVEvents message.
func (r *Router) ApplyKV(u kvevents.Update) { r.kv.Apply(u) }

// kvSettings are router.kv_index with the defaults filled in.
func (r *Router) kvSettings() (timeout time.Duration, maxBlocks int) {
	timeout, maxBlocks = r.cfg.KVIndex.TokenizeTimeout, r.cfg.KVIndex.MaxBlocks
	if timeout <= 0 {
		timeout = 250 * time.Millisecond
	}
	if maxBlocks <= 0 {
		maxBlocks = 4096
	}
	return timeout, maxBlocks
}

// kvQuery asks the precise index about a request for a model that uses it,
// before scheduling: which of the candidates hold how many of the prompt's
// leading blocks. Whatever goes wrong, the answer is a fallback to tier 1, and
// the request is never refused or held up past the tokenize timeout for it.
func (r *Router) kvQuery(ctx context.Context, req *parsedRequest, model string, eps []endpoint.Snapshot) *scheduler.KVQuery {
	ids := make([]string, len(eps))
	for i, e := range eps {
		ids[i] = e.ID
	}
	// Asked first because it is free: with no events from any candidate
	// there is nothing to look up, and nothing worth a round trip.
	sizes := r.kv.BlockSizes(ids)
	if len(sizes) == 0 {
		return &scheduler.KVQuery{Fallback: kvindex.FallbackNoIndex}
	}
	// An adapter's blocks are hashed with its name, a form's are not text.
	if _, ok := r.adapter(req.Model); ok || req.ContentType != "" {
		return &scheduler.KVQuery{Fallback: kvindex.FallbackUnsupported}
	}

	// Tokenized by the least busy candidate the index knows: any of them
	// runs the same template and tokenizer, and the busiest is the one
	// least able to spare the time.
	var by *endpoint.Snapshot
	for i := range eps {
		e := &eps[i]
		if _, _, known := r.kv.Leading(e.ID, nil); !known {
			continue
		}
		if by == nil || e.Inflight+e.Stats.Waiting < by.Inflight+by.Stats.Waiting {
			by = e
		}
	}
	served := model
	if by.UpstreamModel != "" {
		served = by.UpstreamModel
	}
	body, err := kvindex.TokenizeBody(req.Path, req.UpstreamBody, served)
	if err != nil {
		return &scheduler.KVQuery{Fallback: kvindex.FallbackUnsupported}
	}

	timeout, maxBlocks := r.kvSettings()
	tctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	start := time.Now()
	tokens, err := kvindex.Tokenize(tctx, r.kvClient, by.URL, by.APIKey, body)
	outcome := "ok"
	if err != nil {
		outcome = "error"
		if errors.Is(err, context.DeadlineExceeded) {
			outcome = "timeout"
		}
	}
	r.m.kvTokenize.WithLabelValues(outcome).Observe(time.Since(start).Seconds())
	if err != nil {
		r.log.Debug("tokenizing for the kv index failed; routing by tier 1",
			"model", model, "endpoint", by.ID, "err", err)
		return &scheduler.KVQuery{Fallback: kvindex.FallbackTokenizeFailed}
	}

	hashes := make(map[int][]uint64, len(sizes))
	for _, s := range sizes {
		hashes[s] = kvevents.BlockHashes(tokens, s, maxBlocks)
	}
	q := &scheduler.KVQuery{Matched: map[string]int{}, Total: map[string]int{}}
	for _, id := range ids {
		if m, t, known := r.kv.Leading(id, hashes); known {
			q.Matched[id], q.Total[id] = m, t
		}
	}
	return q
}

// newKVClient is the client the tokenize calls use: engines are on the
// management network, and a call that has not connected in a second is one
// the timeout will end anyway.
func newKVClient() *http.Client {
	return &http.Client{Transport: &http.Transport{
		MaxIdleConnsPerHost: 16,
		IdleConnTimeout:     90 * time.Second,
	}}
}

// kvCollector reports what the index holds of each replica.
type kvCollector struct {
	idx     *kvindex.Index
	blocks  *prometheus.Desc
	problem *prometheus.Desc
}

func newKVCollector(idx *kvindex.Index) *kvCollector {
	return &kvCollector{
		idx: idx,
		blocks: prometheus.NewDesc("llmgw_kv_index_blocks",
			"KV cache blocks each replica reports holding, as the router's precise index has them.",
			[]string{"replica"}, nil),
		problem: prometheus.NewDesc("llmgw_kv_index_problem",
			"1 for a replica whose KV cache events cannot be used (its engine hashes another way); "+
				"the agent's log says why.",
			[]string{"replica"}, nil),
	}
}

func (c *kvCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- c.blocks
	ch <- c.problem
}

func (c *kvCollector) Collect(ch chan<- prometheus.Metric) {
	for _, s := range c.idx.Stats() {
		ch <- prometheus.MustNewConstMetric(c.blocks, prometheus.GaugeValue, float64(s.Blocks), s.Replica)
		p := 0.0
		if s.Problem != "" {
			p = 1
		}
		ch <- prometheus.MustNewConstMetric(c.problem, prometheus.GaugeValue, p, s.Replica)
	}
}