llmgw/internal/sim/metrics.go

package sim

import (
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

// ttftBuckets and tpotBuckets use second boundaries close to vLLM's own.
var (
	ttftBuckets = []float64{0.001, 0.005, 0.01, 0.02, 0.04, 0.06, 0.08, 0.1, 0.25, 0.5, 0.75, 1, 2.5, 5, 7.5, 10, 20, 40, 80}
	tpotBuckets = []float64{0.01, 0.025, 0.05, 0.075, 0.1, 0.15, 0.2, 0.3, 0.4, 0.5, 0.75, 1, 2.5}
)

// histogram is a minimal Prometheus-compatible cumulative histogram.
type histogram struct {
	bounds []float64
	counts []uint64
	sum    float64
	count  uint64
}

func newHistogram(bounds []float64) *histogram {
	return &histogram{bounds: bounds, counts: make([]uint64, len(bounds))}
}

func (h *histogram) observe(v float64) {
	h.sum += v
	h.count++
	for i, b := range h.bounds {
		if v <= b {
			h.counts[i]++
		}
	}
}

// write renders the histogram in Prometheus text format (cumulative le buckets).
func (h *histogram) write(w io.Writer, name, labels string) {
	var cum uint64
	for i, b := range h.bounds {
		cum = h.counts[i]
		fmt.Fprintf(w, "%s_bucket{%s,le=\"%g\"} %d\n", name, labels, b, cum)
	}
	fmt.Fprintf(w, "%s_bucket{%s,le=\"+Inf\"} %d\n", name, labels, h.count)
	fmt.Fprintf(w, "%s_sum{%s} %g\n", name, labels, h.sum)
	fmt.Fprintf(w, "%s_count{%s} %d\n", name, labels, h.count)
}

type metrics struct {
	received  atomic.Int64
	finished  atomic.Int64
	failed    atomic.Int64
	cancelled atomic.Int64

	promptTokens atomic.Int64
	genTokens    atomic.Int64

	mu   sync.Mutex
	ttft *histogram
	tpot *histogram
}

func newMetrics() *metrics {
	return &metrics{
		ttft: newHistogram(ttftBuckets),
		tpot: newHistogram(tpotBuckets),
	}
}

func (m *metrics) observeTTFT(d time.Duration) {
	m.mu.Lock()
	m.ttft.observe(d.Seconds())
	m.mu.Unlock()
}

func (m *metrics) observeTPOT(d time.Duration) {
	m.mu.Lock()
	m.tpot.observe(d.Seconds())
	m.mu.Unlock()
}

// WriteMetrics renders the engine's state in Prometheus text format using
// vLLM's metric names.
//
// Those names are the engine contract from plan §7.6; the router parses
// against them. Matching real vLLM names exactly is the point: sim-engine
// exists to stand in for a real engine as far as the router can tell.
func (e *Engine) WriteMetrics(w io.Writer) {
	s := e.Stats()
	lbl := fmt.Sprintf("model_name=%q", e.cfg.Model)

	fmt.Fprintf(w, "# HELP vllm:num_requests_running Number of requests currently running on GPU.\n")
	fmt.Fprintf(w, "# TYPE vllm:num_requests_running gauge\n")
	fmt.Fprintf(w, "vllm:num_requests_running{%s} %d\n", lbl, s.Running)

	fmt.Fprintf(w, "# HELP vllm:num_requests_waiting Number of requests waiting to be processed.\n")
	fmt.Fprintf(w, "# TYPE vllm:num_requests_waiting gauge\n")
	fmt.Fprintf(w, "vllm:num_requests_waiting{%s} %d\n", lbl, s.Waiting)

	fmt.Fprintf(w, "# HELP vllm:kv_cache_usage_perc KV-cache usage. 1 means 100 percent usage.\n")
	fmt.Fprintf(w, "# TYPE vllm:kv_cache_usage_perc gauge\n")
	fmt.Fprintf(w, "vllm:kv_cache_usage_perc{%s} %g\n", lbl, s.KVCacheUsage)

	fmt.Fprintf(w, "# HELP vllm:cache_config_info Information of the LLMEngine CacheConfig.\n")
	fmt.Fprintf(w, "# TYPE vllm:cache_config_info gauge\n")
	fmt.Fprintf(w, "vllm:cache_config_info{block_size=\"%d\",num_gpu_blocks=\"%d\"} 1\n", s.BlockSize, s.NumGPUBlocks)

	fmt.Fprintf(w, "# HELP vllm:prefix_cache_queries_total Prefix cache queries, in terms of number of queried blocks.\n")
	fmt.Fprintf(w, "# TYPE vllm:prefix_cache_queries_total counter\n")
	fmt.Fprintf(w, "vllm:prefix_cache_queries_total{%s} %d\n", lbl, s.PrefixQueries)

	fmt.Fprintf(w, "# HELP vllm:prefix_cache_hits_total Prefix cache hits, in terms of number of cached blocks.\n")
	fmt.Fprintf(w, "# TYPE vllm:prefix_cache_hits_total counter\n")
	fmt.Fprintf(w, "vllm:prefix_cache_hits_total{%s} %d\n", lbl, s.PrefixHits)

	fmt.Fprintf(w, "# HELP vllm:prompt_tokens_total Number of prefill tokens processed.\n")
	fmt.Fprintf(w, "# TYPE vllm:prompt_tokens_total counter\n")
	fmt.Fprintf(w, "vllm:prompt_tokens_total{%s} %d\n", lbl, e.m.promptTokens.Load())

	fmt.Fprintf(w, "# HELP vllm:generation_tokens_total Number of generation tokens processed.\n")
	fmt.Fprintf(w, "# TYPE vllm:generation_tokens_total counter\n")
	fmt.Fprintf(w, "vllm:generation_tokens_total{%s} %d\n", lbl, e.m.genTokens.Load())

	fmt.Fprintf(w, "# HELP vllm:request_success_total Count of successfully processed requests.\n")
	fmt.Fprintf(w, "# TYPE vllm:request_success_total counter\n")
	fmt.Fprintf(w, "vllm:request_success_total{%s,finished_reason=\"length\"} %d\n", lbl, e.m.finished.Load())

	// Simulation-specific (absent in real vLLM; useful for chaos tests).
	fmt.Fprintf(w, "# TYPE sim:requests_cancelled_total counter\n")
	fmt.Fprintf(w, "sim:requests_cancelled_total{%s} %d\n", lbl, e.m.cancelled.Load())
	fmt.Fprintf(w, "# TYPE sim:requests_failed_total counter\n")
	fmt.Fprintf(w, "sim:requests_failed_total{%s} %d\n", lbl, e.m.failed.Load())

	e.m.mu.Lock()
	defer e.m.mu.Unlock()
	fmt.Fprintf(w, "# TYPE vllm:time_to_first_token_seconds histogram\n")
	e.m.ttft.write(w, "vllm:time_to_first_token_seconds", lbl)
	fmt.Fprintf(w, "# TYPE vllm:time_per_output_token_seconds histogram\n")
	e.m.tpot.write(w, "vllm:time_per_output_token_seconds", lbl)
}