llmgw/internal/router/metrics.go

package router

import (
	"slices"
	"strconv"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
)

// metrics holds the router's own Prometheus series (plan §8).
//
// The scorer series matter more than they look: routing profile weights cannot
// be tuned without seeing what each scorer actually produced in production.
type metrics struct {
	requests *prometheus.CounterVec
	duration *prometheus.HistogramVec
	ttft     *prometheus.HistogramVec
	tokens   *prometheus.CounterVec

	scorerValue *prometheus.GaugeVec
	filtered    *prometheus.CounterVec
	considered  prometheus.Histogram

	admissionWait     *prometheus.HistogramVec
	admissionRejected *prometheus.CounterVec

	prefixDecisions *prometheus.CounterVec
	// kvDecisions is the precise index's counterpart, and kvTokenize how
	// long the tokenize call it needs took, by outcome.
	kvDecisions *prometheus.CounterVec
	kvTokenize  *prometheus.HistogramVec

	// adapterAffinity is where requests for adapters that load on demand
	// found them; adapterLoads and adapterEvictions what that cost.
	adapterAffinity  *prometheus.CounterVec
	adapterLoads     *prometheus.CounterVec
	adapterEvictions *prometheus.CounterVec

	// cacheRequests counts what the answer cache did, cacheBytes and
	// cacheEntries what it holds.
	cacheRequests *prometheus.CounterVec
	cacheBytes    prometheus.Gauge
	cacheEntries  prometheus.Gauge

	tenant *tenantMetrics
	guard  *guardMetrics
	token  *tokenMetrics
}

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"}),
		adapterAffinity: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "llmgw_adapter_affinity_total",
			Help: "Requests for LoRA adapters that load on demand, by where the chosen replica had it: " +
				"hit (loaded), joined (being loaded for another request), loaded (loaded for this one).",
		}, []string{"model", "result"}),
		adapterLoads: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "llmgw_adapter_loads_total",
			Help: "LoRA adapters the router had a replica load, by outcome: ok or failed.",
		}, []string{"model", "outcome"}),
		adapterEvictions: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "llmgw_adapter_evictions_total",
			Help: "LoRA adapters the router had a replica unload to make room, by outcome: ok or failed.",
		}, []string{"model", "outcome"}),
		kvDecisions: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "llmgw_kv_index_decisions_total",
			Help: "Precise KV index routing decisions, for models with kv_events, by branch: kv_hit, " +
				"kv_miss_min_load, load_balance_override, and fallback_no_index, fallback_unsupported, " +
				"fallback_tokenize_failed, where tier 1 decided instead.",
		}, []string{"decision"}),
		kvTokenize: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "llmgw_kv_index_tokenize_seconds",
			Help:    "Time the precise KV index waited for a replica to tokenize a request, by outcome: ok, timeout, error.",
			Buckets: []float64{0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5},
		}, []string{"outcome"}),
		requests: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "llmgw_requests_total",
			Help: "Requests handled, by model, pool and response status.",
		}, []string{"model", "pool", "status"}),

		duration: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "llmgw_request_duration_seconds",
			Help:    "End-to-end request duration.",
			Buckets: []float64{0.01, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, 120, 300},
		}, []string{"model", "pool"}),

		ttft: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "llmgw_ttft_seconds",
			Help:    "Time to first token, as observed by the router.",
			Buckets: []float64{0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10, 30, 60},
		}, []string{"model", "pool"}),

		tokens: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "llmgw_tokens_total",
			Help: "Tokens accounted for, by direction.",
		}, []string{"model", "direction"}),

		scorerValue: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Name: "llmgw_scorer_value",
			Help: "Last score each scorer produced for the endpoint that was chosen.",
		}, []string{"scorer", "replica"}),

		filtered: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "llmgw_candidates_filtered_total",
			Help: "Candidate endpoints removed, by filter.",
		}, []string{"filter"}),

		considered: prometheus.NewHistogram(prometheus.HistogramOpts{
			Name:    "llmgw_candidates_considered",
			Help:    "Candidates remaining after filtering, per scheduling decision.",
			Buckets: []float64{1, 2, 3, 5, 8, 16, 32, 64},
		}),

		admissionWait: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "llmgw_admission_wait_seconds",
			Help:    "Time spent waiting in the admission queue.",
			Buckets: []float64{0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 2.5, 5, 10, 30},
		}, []string{"band", "model"}),

		admissionRejected: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "llmgw_admission_rejected_total",
			Help: "Requests refused by flow control, by reason.",
		}, []string{"band", "reason"}),
		cacheRequests: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "llmgw_cache_requests_total",
			Help: "Requests of projects that reuse answers, by what the answer cache did: " +
				"hit (answered from it), miss (asked the model, and kept the answer if it could), " +
				"bypass (not one whose answer is reused).",
		}, []string{"model", "result"}),
		cacheBytes: prometheus.NewGauge(prometheus.GaugeOpts{
			Name: "llmgw_cache_bytes",
			Help: "Memory the answers kept for reuse take in this process, roughly.",
		}),
		cacheEntries: prometheus.NewGauge(prometheus.GaugeOpts{
			Name: "llmgw_cache_entries",
			Help: "Answers kept for reuse in this process.",
		}),
		tenant: newTenantMetrics(),
		guard:  newGuardMetrics(),
		token:  newTokenMetrics(),
	}
}

// collectors returns everything to register, so registration lives in one place.
func (m *metrics) collectors() []prometheus.Collector {
	return append([]prometheus.Collector{
		m.requests, m.duration, m.ttft, m.tokens,
		m.scorerValue, m.filtered, m.considered,
		m.admissionWait, m.admissionRejected, m.prefixDecisions,
		m.adapterAffinity, m.adapterLoads, m.adapterEvictions,
		m.kvDecisions, m.kvTokenize,
		m.cacheRequests, m.cacheBytes, m.cacheEntries,
	}, slices.Concat(m.tenant.collectors(), m.guard.collectors(), m.token.collectors())...)
}

// 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()
}

// observeAdapter counts a load or an eviction the router made.
func (m *metrics) observeAdapter(model, what, outcome string) {
	switch what {
	case "load":
		m.adapterLoads.WithLabelValues(model, outcome).Inc()
	case "evict":
		m.adapterEvictions.WithLabelValues(model, outcome).Inc()
	}
}

// observeKV records which branch a precise index decision took.
func (m *metrics) observeKV(d kvindex.Decision) {
	m.kvDecisions.WithLabelValues(string(d)).Inc()
}

func (m *metrics) observeAdmission(band, model string, wait time.Duration) {
	if band == "" {
		band = "none"
	}
	m.admissionWait.WithLabelValues(band, model).Observe(wait.Seconds())
}

func (m *metrics) observeAdmissionRejected(band, reason string) {
	if band == "" {
		band = "unknown"
	}
	m.admissionRejected.WithLabelValues(band, reason).Inc()
}

func (m *metrics) observeRequest(model, pool string, status int, d time.Duration, u *proxy.Usage, ttft time.Duration) {
	if model == "" {
		model = "unknown"
	}
	if pool == "" {
		pool = "none"
	}
	m.requests.WithLabelValues(model, pool, strconv.Itoa(status)).Inc()
	m.duration.WithLabelValues(model, pool).Observe(d.Seconds())

	if ttft > 0 {
		m.ttft.WithLabelValues(model, pool).Observe(ttft.Seconds())
	}
	if u != nil {
		m.tokens.WithLabelValues(model, "prompt").Add(float64(u.PromptTokens))
		m.tokens.WithLabelValues(model, "completion").Add(float64(u.CompletionTokens))
		if u.CachedTokens > 0 {
			m.tokens.WithLabelValues(model, "cached").Add(float64(u.CachedTokens))
		}
	}
}

func (m *metrics) observeScores(d scheduler.Decision) {
	for name, v := range d.Scores {
		m.scorerValue.WithLabelValues(name, d.Endpoint.ID).Set(v)
	}
	for filter, n := range d.Dropped {
		m.filtered.WithLabelValues(filter).Add(float64(n))
	}
	m.considered.Observe(float64(d.Considered))
}

// metricModel is the model label a request is counted under: the catalog's
// name for it, or "unknown". Taken as sent, any key holder could mint a new
// series per request — a refused request counted its made-up model name —
// and grow Prometheus without bound.
func (r *Router) metricModel(name string) string {
	if name == "" {
		return ""
	}
	if m, ok := r.modelMap()[name]; ok {
		return m.Name
	}
	return "unknown"
}