llmgw/internal/router/tokenmetrics.go

package router

import (
	"slices"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
)

// Token-level timing: how long each output token took, the gaps between
// streamed tokens, and output tokens a second — per model and per tenant.
// What a model's SLO is written in, and what admission reads back.

type tokenMetrics struct {
	tpot       *prometheus.HistogramVec
	gap        *prometheus.HistogramVec
	rate       *prometheus.HistogramVec
	tenantTPOT *prometheus.HistogramVec
	tenantRate *prometheus.HistogramVec
	slo        *prometheus.GaugeVec
}

var tokenBuckets = []float64{0.005, 0.01, 0.02, 0.035, 0.05, 0.075, 0.1, 0.15, 0.25, 0.5, 1, 2}

func newTokenMetrics() *tokenMetrics {
	rateBuckets := []float64{1, 2, 5, 10, 15, 20, 30, 50, 75, 100, 150, 250}
	return &tokenMetrics{
		tpot: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "llmgw_tpot_seconds",
			Help:    "Time per output token after the first, per streamed request: what a model's slo.tpot_p99 is written in.",
			Buckets: tokenBuckets,
		}, []string{"model", "pool"}),
		gap: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "llmgw_inter_token_seconds",
			Help:    "Time between streamed events, about one token each: the stalls a mean per-token time hides.",
			Buckets: tokenBuckets,
		}, []string{"model"}),
		rate: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "llmgw_output_tokens_per_second",
			Help:    "Output tokens a second per request: after the first token when streamed, over the whole request otherwise.",
			Buckets: rateBuckets,
		}, []string{"model", "pool"}),
		tenantTPOT: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "llmgw_tenant_tpot_seconds",
			Help:    "Time per output token by tenant (the level router.tenant_label names), for streamed requests.",
			Buckets: tokenBuckets,
		}, []string{"tenant"}),
		slo: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Name: "llmgw_model_slo_seconds",
			Help: "A model's SLO from the catalog, by kind: ttft_p99 or tpot_p99. For alerting against the histograms.",
		}, []string{"model", "kind"}),
		tenantRate: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "llmgw_tenant_output_tokens_per_second",
			Help:    "Output tokens a second per request, by tenant.",
			Buckets: rateBuckets,
		}, []string{"tenant"}),
	}
}

func (t *tokenMetrics) collectors() []prometheus.Collector {
	return []prometheus.Collector{t.tpot, t.gap, t.rate, t.tenantTPOT, t.tenantRate, t.slo}
}

// tokenTiming is one finished request's per-token figures; ok is false when
// there are none to give.
func tokenTiming(streamed bool, o proxy.Outcome, took time.Duration) (tpot time.Duration, rate float64, ok bool) {
	if o.Usage == nil || o.Usage.CompletionTokens <= 0 {
		return 0, 0, false
	}
	n := o.Usage.CompletionTokens
	if streamed {
		if n < 2 || o.Decode <= 0 {
			return 0, 0, false
		}
		tpot = o.Decode / time.Duration(n-1)
		return tpot, float64(n-1) / o.Decode.Seconds(), true
	}
	if took <= 0 {
		return 0, 0, false
	}
	return 0, float64(n) / took.Seconds(), true
}

func (r *Router) observeTokens(req *parsedRequest, res dispatchResult, took time.Duration) {
	tpot, rate, ok := tokenTiming(req.Stream, res.outcome, took)
	if !ok {
		return
	}
	model, pool := r.metricModel(req.Model), res.pool
	if pool == "" {
		pool = "none"
	}
	tenant := r.m.tenant.tenantOf(req)
	if tpot > 0 {
		r.m.token.tpot.WithLabelValues(model, pool).Observe(tpot.Seconds())
		if tenant != "" {
			r.m.token.tenantTPOT.WithLabelValues(tenant).Observe(tpot.Seconds())
		}
	}
	r.m.token.rate.WithLabelValues(model, pool).Observe(rate)
	if tenant != "" {
		r.m.token.tenantRate.WithLabelValues(tenant).Observe(rate)
	}
	if req.Stream {
		r.slo.add(r.canonicalModel(req.Model), res.outcome.TTFT, tpot)
	}
}

func (r *Router) canonicalModel(name string) string {
	c, _ := r.resolve(name)
	return c
}

// sloTracker keeps each model's recent first-token and per-token times, for
// admission to compare against its SLO.
type sloTracker struct {
	mu sync.Mutex
	by map[string]*sloSamples
}

const (
	sloSamplesKept = 256
	sloWindow      = 5 * time.Minute
	// sloMinSamples is how many recent requests a p99 needs before it is
	// believed.
	sloMinSamples = 20
)

type sloSample struct {
	at         time.Time
	ttft, tpot time.Duration
}

type sloSamples struct {
	ring []sloSample
	next int
}

func (t *sloTracker) add(model string, ttft, tpot time.Duration) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.by == nil {
		t.by = map[string]*sloSamples{}
	}
	s := t.by[model]
	if s == nil {
		s = &sloSamples{}
		t.by[model] = s
	}
	x := sloSample{at: time.Now(), ttft: ttft, tpot: tpot}
	if len(s.ring) < sloSamplesKept {
		s.ring = append(s.ring, x)
		return
	}
	s.ring[s.next] = x
	s.next = (s.next + 1) % sloSamplesKept
}

// p99 is a model's recent 99th percentiles, and false with too few.
func (t *sloTracker) p99(model string, now time.Time) (ttft, tpot time.Duration, ok bool) {
	t.mu.Lock()
	s := t.by[model]
	var a, b []time.Duration
	if s != nil {
		for _, x := range s.ring {
			if now.Sub(x.at) > sloWindow {
				continue
			}
			if x.ttft > 0 {
				a = append(a, x.ttft)
			}
			if x.tpot > 0 {
				b = append(b, x.tpot)
			}
		}
	}
	t.mu.Unlock()
	pick := func(v []time.Duration) time.Duration {
		if len(v) < sloMinSamples {
			return 0
		}
		slices.Sort(v)
		return v[(len(v)*99)/100]
	}
	ttft, tpot = pick(a), pick(b)
	return ttft, tpot, ttft > 0 || tpot > 0
}

// sloPressure is how far a model is over its SLO: 1 at it, above 1 past
// it, 0 when it has none or too little is known.
func (r *Router) sloPressure(model string) float64 {
	_, entry := r.resolve(model)
	slo := entry.SLO
	if slo.TTFTp99 <= 0 && slo.TPOTp99 <= 0 {
		return 0
	}
	ttft, tpot, ok := r.slo.p99(entry.Name, time.Now())
	if !ok {
		return 0
	}
	p := 0.0
	if slo.TTFTp99 > 0 && ttft > 0 {
		p = max(p, float64(ttft)/float64(slo.TTFTp99))
	}
	if slo.TPOTp99 > 0 && tpot > 0 {
		p = max(p, float64(tpot)/float64(slo.TPOTp99))
	}
	return p
}

// setSLOs publishes the catalog's SLOs, all at once, so a model removed or
// given none stops being reported.
func (t *tokenMetrics) setSLOs(models []catalog.Model) {
	t.slo.Reset()
	for _, m := range models {
		if m.SLO.TTFTp99 > 0 {
			t.slo.WithLabelValues(m.Name, "ttft_p99").Set(m.SLO.TTFTp99.Seconds())
		}
		if m.SLO.TPOTp99 > 0 {
			t.slo.WithLabelValues(m.Name, "tpot_p99").Set(m.SLO.TPOTp99.Seconds())
		}
	}
}