llmgw/internal/router/scrape/parse.go

// Package scrape reads engines' /metrics output and turns it into something
// the router can make decisions from (plan §7.6).
package scrape

import (
	"fmt"
	"io"
	"strconv"

	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
	"github.com/prometheus/common/model"

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

// Mapping says which metric names the contract fields are read from.
//
// Being configurable is essential: engine versions rename their metrics
// (plan §14). Once the agent reports an engine version the control plane can
// pick the right mapping; in phase 1 it comes from configuration.
type Mapping struct {
	Running      string
	Waiting      string
	KVCacheUsage string

	CacheConfig      string
	PromptTokens     string
	GenerationTokens string
	SuccessTotal     string

	// Path is where the engine serves these metrics, when it is not the
	// poller's configured path: trtllm-serve's /metrics is JSON, and its
	// Prometheus text is under /prometheus/metrics.
	Path string
}

// MappingFor returns the metric names for an engine.
//
// Per engine rather than per cluster: a cluster running llama.cpp on the CPU
// nodes and vLLM on the GPU nodes has two vocabularies for the same three
// numbers, and one global mapping would leave half the fleet unmeasurable.
// An unknown engine falls back to vLLM's names, which is also what an endpoint
// that names no engine gets.
func MappingFor(engine string) Mapping {
	switch engine {
	case "llamacpp":
		return LlamaCppMapping()
	case "trtllm":
		return TRTLLMMapping()
	case "sglang":
		return SGLangMapping()
	default:
		return VLLMMapping()
	}
}

// LlamaCppMapping returns llama.cpp's metric names, which llama-server only
// exposes when started with --metrics.
//
// "Deferred" is llama.cpp's word for a request that arrived with every slot
// busy, which is the same thing vLLM calls waiting.
//
// KV cache usage is deliberately absent: current llama.cpp does not publish it
// (b10991 serves fifteen families and none of them is a cache ratio), and
// naming a metric that is not there would fail the contract on every scrape and
// leave the router with no numbers at all rather than two out of three. The
// loss is smaller here than it would be for vLLM — llama.cpp serves a fixed
// number of slots, so its queue depth already says when it is full — but it is
// a loss: saturation falls back to queue depth alone.
//
// The rest are metrics llama.cpp does not have either. It counts decodes rather
// than requests, so there is no success counter and no cache configuration to
// read block sizes from.
func LlamaCppMapping() Mapping {
	return Mapping{
		Running:          "llamacpp:requests_processing",
		Waiting:          "llamacpp:requests_deferred",
		PromptTokens:     "llamacpp:prompt_tokens_total",
		GenerationTokens: "llamacpp:tokens_predicted_total",
	}
}

// TRTLLMMapping returns TensorRT-LLM's metric names, as trtllm-serve
// publishes them under /prometheus/metrics.
//
// They are there only when the server runs with return_perf_metrics, and
// running and waiting move only with enable_iter_perf_stats; llmgw's launcher
// turns both on. The gauges are set from each iteration's statistics, so
// they say what the engine was doing when it last iterated: an engine that
// has gone idle keeps showing its last batch until the next request moves it.
// Running and waiting first appeared in the 1.3 release candidates; an older
// server serves none of the three and is reported as breaking the contract,
// which is what it does.
//
// KV cache utilization is used blocks over the pool, a fraction like vLLM's.
// The cache's configuration is published under other label names than the
// block counts read from vLLM, so it is left unread.
func TRTLLMMapping() Mapping {
	return Mapping{
		Running:          "trtllm_num_requests_running",
		Waiting:          "trtllm_num_requests_waiting",
		KVCacheUsage:     "trtllm_kv_cache_utilization",
		PromptTokens:     "trtllm_prompt_tokens_total",
		GenerationTokens: "trtllm_generation_tokens_total",
		SuccessTotal:     "trtllm_request_success_total",
		Path:             "/prometheus/metrics",
	}
}

// SGLangMapping returns SGLang's metric names, as sglang.launch_server
// publishes them under /metrics (checked against 0.5.21's collector).
//
// They are there only when the server runs with --enable-metrics, which
// llmgw's launch command passes; without it /metrics is not served at all.
// SGLang exempts /metrics from its API key, so the key the poller sends is
// ignored rather than needed.
//
// "Queue" is SGLang's word for vLLM's waiting. The gauges are set by the
// scheduler as it reports its batches, and at least every thirty seconds
// while it is idle, so a server that has only just started may serve none of
// them yet and is reported incomplete until it does.
//
// token_usage is the fraction of the KV pool in use, as vLLM's
// kv_cache_usage_perc is (for a model with sliding-window layers, the fuller
// of its two pools). cache_hit_rate is a different number — how much of the
// prompts the prefix cache answered — and is not what saturation means.
//
// num_requests_total counts the requests that finished, the nearest thing it
// has to vLLM's success counter; like the token counters it is split by
// is_streaming and summed. There is no cache configuration to read block
// sizes from.
func SGLangMapping() Mapping {
	return Mapping{
		Running:          "sglang:num_running_reqs",
		Waiting:          "sglang:num_queue_reqs",
		KVCacheUsage:     "sglang:token_usage",
		PromptTokens:     "sglang:prompt_tokens_total",
		GenerationTokens: "sglang:generation_tokens_total",
		SuccessTotal:     "sglang:num_requests_total",
	}
}

// VLLMMapping returns vLLM's current metric names.
func VLLMMapping() Mapping {
	return Mapping{
		Running:          "vllm:num_requests_running",
		Waiting:          "vllm:num_requests_waiting",
		KVCacheUsage:     "vllm:kv_cache_usage_perc",
		CacheConfig:      "vllm:cache_config_info",
		PromptTokens:     "vllm:prompt_tokens_total",
		GenerationTokens: "vllm:generation_tokens_total",
		SuccessTotal:     "vllm:request_success_total",
	}
}

// normalize fills in vLLM's names for an empty mapping.
//
// Only for an empty one. A mapping that names its own Running metric is a
// complete statement about some other engine, and the fields it leaves blank
// are metrics that engine does not expose — filling those with vLLM's names
// would have the router hunting for metrics that cannot be there.
func (m *Mapping) normalize() {
	if m.Running != "" {
		return
	}
	d := VLLMMapping()
	if m.Running == "" {
		m.Running = d.Running
	}
	if m.Waiting == "" {
		m.Waiting = d.Waiting
	}
	if m.KVCacheUsage == "" {
		m.KVCacheUsage = d.KVCacheUsage
	}
	if m.CacheConfig == "" {
		m.CacheConfig = d.CacheConfig
	}
	if m.PromptTokens == "" {
		m.PromptTokens = d.PromptTokens
	}
	if m.GenerationTokens == "" {
		m.GenerationTokens = d.GenerationTokens
	}
	if m.SuccessTotal == "" {
		m.SuccessTotal = d.SuccessTotal
	}
}

// MissingError reports that mandatory contract metrics are absent.
type MissingError struct {
	Missing []string
}

func (e *MissingError) Error() string {
	return fmt.Sprintf("engine metric contract incomplete: %v", e.Missing)
}

// Parse converts Prometheus text-format metrics into Stats.
//
// A metric the mapping names and the engine does not serve is a contract
// violation and returns *MissingError: routing to that endpoint would be blind.
// A metric the mapping leaves unnamed is one the engine is known not to
// publish, and stays zero — the router loses that signal for those endpoints
// rather than losing the endpoints.
func Parse(r io.Reader, m Mapping) (endpoint.Stats, error) {
	m.normalize()

	// TextParser's zero value is invalid; the scheme must be explicit.
	// The legacy scheme accepts vLLM's names containing ':'.
	p := expfmt.NewTextParser(model.LegacyValidation)
	fams, err := p.TextToMetricFamilies(r)
	if err != nil {
		return endpoint.Stats{}, fmt.Errorf("failed to parse metrics: %w", err)
	}

	var s endpoint.Stats
	var missing []string

	// Running and waiting have no such escape: an engine that reports neither
	// how much it is doing nor how much is queued cannot be scheduled against
	// at all.
	if v, ok := firstValue(fams, m.Running); ok {
		s.Running = int(v)
	} else {
		missing = append(missing, m.Running)
	}
	if v, ok := firstValue(fams, m.Waiting); ok {
		s.Waiting = int(v)
	} else {
		missing = append(missing, m.Waiting)
	}
	if m.KVCacheUsage != "" {
		if v, ok := firstValue(fams, m.KVCacheUsage); ok {
			s.KVCacheUsage = v
		} else {
			missing = append(missing, m.KVCacheUsage)
		}
	}
	if len(missing) > 0 {
		return endpoint.Stats{}, &MissingError{Missing: missing}
	}

	// Optional fields: silently stay zero when absent.
	s.PromptTokens, _ = sumValues(fams, m.PromptTokens)
	s.GenerationTokens, _ = sumValues(fams, m.GenerationTokens)
	// request_success_total is split by finished_reason; sum across labels.
	s.SuccessTotal, _ = sumValues(fams, m.SuccessTotal)

	if f, ok := fams[m.CacheConfig]; ok && len(f.Metric) > 0 {
		for _, lp := range f.Metric[0].Label {
			switch lp.GetName() {
			case "block_size":
				s.BlockSize, _ = strconv.Atoi(lp.GetValue())
			case "num_gpu_blocks":
				s.NumGPUBlocks, _ = strconv.Atoi(lp.GetValue())
			}
		}
	}

	return s, nil
}

// metricValue extracts a numeric value regardless of the metric type.
func metricValue(mt *dto.Metric) (float64, bool) {
	switch {
	case mt.Gauge != nil:
		return mt.Gauge.GetValue(), true
	case mt.Counter != nil:
		return mt.Counter.GetValue(), true
	case mt.Untyped != nil:
		return mt.Untyped.GetValue(), true
	default:
		return 0, false
	}
}

func firstValue(fams map[string]*dto.MetricFamily, name string) (float64, bool) {
	f, ok := fams[name]
	if !ok || len(f.Metric) == 0 {
		return 0, false
	}
	return metricValue(f.Metric[0])
}

func sumValues(fams map[string]*dto.MetricFamily, name string) (float64, bool) {
	f, ok := fams[name]
	if !ok {
		return 0, false
	}
	var sum float64
	var any bool
	for _, mt := range f.Metric {
		if v, ok := metricValue(mt); ok {
			sum += v
			any = true
		}
	}
	return sum, any
}