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
}

// 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",
	}
}

func (m *Mapping) normalize() {
	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.
//
// If any of the mandatory three (running, waiting, kv usage) is absent it
// returns *MissingError: routing to that endpoint would be blind.
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

	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 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
}