llmgw/internal/router/scrape/trtllm_test.go

package scrape

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

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

// trtllmMetrics is what trtllm-serve's /prometheus/metrics serves, in the
// shape prometheus_client's multiprocess collector writes it: every family
// labelled with the model and the backend, gauges with the pid of the
// process that set them, counters named with their _total. Trimmed to the
// families that matter and one histogram, which the parser has to step over.
const trtllmMetrics = `# HELP trtllm_request_success_total Count of successfully processed requests.
# TYPE trtllm_request_success_total counter
trtllm_request_success_total{engine_type="pytorch",finished_reason="stop",model_name="llama-3.1-70b"} 40.0
trtllm_request_success_total{engine_type="pytorch",finished_reason="length",model_name="llama-3.1-70b"} 2.0
# HELP trtllm_e2e_request_latency_seconds Histogram of end to end request latency in seconds.
# TYPE trtllm_e2e_request_latency_seconds histogram
trtllm_e2e_request_latency_seconds_bucket{engine_type="pytorch",le="0.5",model_name="llama-3.1-70b"} 10.0
trtllm_e2e_request_latency_seconds_bucket{engine_type="pytorch",le="+Inf",model_name="llama-3.1-70b"} 42.0
trtllm_e2e_request_latency_seconds_count{engine_type="pytorch",model_name="llama-3.1-70b"} 42.0
trtllm_e2e_request_latency_seconds_sum{engine_type="pytorch",model_name="llama-3.1-70b"} 61.5
# HELP trtllm_prompt_tokens_total Number of prefill tokens processed.
# TYPE trtllm_prompt_tokens_total counter
trtllm_prompt_tokens_total{engine_type="pytorch",model_name="llama-3.1-70b"} 18234.0
# HELP trtllm_generation_tokens_total Number of generation tokens processed.
# TYPE trtllm_generation_tokens_total counter
trtllm_generation_tokens_total{engine_type="pytorch",model_name="llama-3.1-70b"} 9120.0
# HELP trtllm_kv_cache_utilization KV cache utilization
# TYPE trtllm_kv_cache_utilization gauge
trtllm_kv_cache_utilization{engine_type="pytorch",model_name="llama-3.1-70b",pid="4121"} 0.375
# HELP trtllm_kv_cache_config_info KV cache configuration info
# TYPE trtllm_kv_cache_config_info gauge
trtllm_kv_cache_config_info{enable_block_reuse="True",engine_type="pytorch",model_name="llama-3.1-70b",page_size="32",pid="4121"} 1.0
# HELP trtllm_num_requests_running Number of active requests
# TYPE trtllm_num_requests_running gauge
trtllm_num_requests_running{engine_type="pytorch",model_name="llama-3.1-70b",pid="4121"} 5.0
# HELP trtllm_num_requests_waiting Number of queued requests
# TYPE trtllm_num_requests_waiting gauge
trtllm_num_requests_waiting{engine_type="pytorch",model_name="llama-3.1-70b",pid="4121"} 2.0
`

func TestTRTLLMMetricsSatisfyTheContract(t *testing.T) {
	s, err := Parse(strings.NewReader(trtllmMetrics), MappingFor("trtllm"))
	if err != nil {
		t.Fatal(err)
	}
	if s.Running != 5 || s.Waiting != 2 {
		t.Errorf("running=%d waiting=%d", s.Running, s.Waiting)
	}
	// A fraction of the pool, as vLLM's is: the scorer compares the two.
	if s.KVCacheUsage != 0.375 {
		t.Errorf("kv=%v", s.KVCacheUsage)
	}
	if s.PromptTokens != 18234 || s.GenerationTokens != 9120 {
		t.Errorf("tokens=%v/%v", s.PromptTokens, s.GenerationTokens)
	}
	// Split by finished_reason, summed.
	if s.SuccessTotal != 42 {
		t.Errorf("success=%v", s.SuccessTotal)
	}
	// Its cache configuration carries no block counts to read.
	if s.BlockSize != 0 || s.NumGPUBlocks != 0 {
		t.Errorf("block size %d and blocks %d read from labels that do not hold them", s.BlockSize, s.NumGPUBlocks)
	}
}

func TestATRTLLMWithoutItsQueueGaugesIsReportedNotIdle(t *testing.T) {
	// TensorRT-LLM 1.2 publishes the cache's utilization and none of the queue
	// gauges. Read as zeros, it would look like the emptiest endpoint in the
	// cluster and take every request.
	only := `# TYPE trtllm_kv_cache_utilization gauge
trtllm_kv_cache_utilization{engine_type="pytorch",model_name="m"} 0.1
`
	var missing *MissingError
	_, err := Parse(strings.NewReader(only), TRTLLMMapping())
	if !errors.As(err, &missing) || len(missing.Missing) != 2 {
		t.Fatalf("expected running and waiting reported missing, got %v", err)
	}
}

func TestTRTLLMIsScrapedWhereItPublishesPrometheusText(t *testing.T) {
	// trtllm-serve's /metrics answers JSON — iteration statistics, drained by
	// each read — and the launcher puts every path but /health behind the key.
	var gotKey atomic.Value
	gotKey.Store("")
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/prometheus/metrics":
			gotKey.Store(r.Header.Get("Authorization"))
			if r.Header.Get("Authorization") != "Bearer k1" {
				http.Error(w, "a valid API key is required", http.StatusUnauthorized)
				return
			}
			_, _ = w.Write([]byte(trtllmMetrics))
		case "/metrics":
			_, _ = w.Write([]byte(`[{"iter": 154, "numActiveRequests": 1}]`))
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()

	reg := endpoint.NewRegistry(endpoint.Config{})
	reg.Replace([]endpoint.Info{{ID: "t1", URL: srv.URL, Engine: "trtllm", APIKey: "k1", Models: []string{"m"}}})
	p := NewPoller(reg, Config{Interval: 20 * time.Millisecond, Timeout: 2 * time.Second}, nil)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	p.Sync(ctx)
	defer p.Stop()

	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		for _, snap := range reg.Snapshots("m") {
			if snap.ID == "t1" && snap.Stats.Running == 5 {
				if snap.Stats.Waiting != 2 || snap.Stats.KVCacheUsage != 0.375 {
					t.Errorf("stats read wrong: %+v", snap.Stats)
				}
				return
			}
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("the endpoint was never measured (key sent: %q)", gotKey.Load())
}