feat(metrics): add request prompt, generation, max_tokens and success metrics

Signed-off-by: googs1025 <[email protected]>

Author: googs1025

manifests/config_with_fake.yaml

@@ -7,10 +7,18 @@ time-to-first-token: 2000
 inter-token-latency: 1000
 kv-cache-transfer-latency: 100
 seed: 100100100
-fake-metrics: 
+fake-metrics:
   running-requests: 16
   waiting-requests: 3 
   kv-cache-usage: 0.3
+  request-success-total:
+    stop: 20
+    length: 0
+    tool_calls: 0
+    remote_decode: 0
+  request-prompt-tokens: [ 10, 20, 30, 15 ]
+  request-generation-tokens: [ 50, 60, 40 ]
+  request-params-max-tokens: [ 128, 256, 512 ]
   loras:
   - '{"running":"lora1,lora2","waiting":"lora3","timestamp":1257894567}'
   - '{"running":"lora1,lora3","waiting":"","timestamp":1257894569}'

pkg/common/config.go

@@ -186,6 +186,13 @@ type Metrics struct {
 	WaitingRequests int64 `yaml:"waiting-requests" json:"waiting-requests"`
 	// KVCacheUsagePercentage  is the fraction of KV-cache blocks currently in use (from 0 to 1)
 	KVCacheUsagePercentage float32 `yaml:"kv-cache-usage" json:"kv-cache-usage"`
+	// RequestPromptTokens RequestGenerationTokens RequestParamsMaxTokens Histogram fake-observation arrays for init.
+	// Each value will be passed to Observe() once at start-up.
+	RequestPromptTokens     []float64 `yaml:"request-prompt-tokens" json:"request-prompt-tokens"`         // prompt-length samples
+	RequestGenerationTokens []float64 `yaml:"request-generation-tokens" json:"request-generation-tokens"` // generation-length samples
+	RequestParamsMaxTokens  []float64 `yaml:"request-params-max-tokens" json:"request-params-max-tokens"` // max_tokens parameter samples
+	// RequestSuccessTotal is the number of successful requests, key: finish-reason (stop, length, etc.).
+	RequestSuccessTotal map[string]int64 `yaml:"request-success-total" json:"request-success-total"`
 }
 
 type LorasMetrics struct {

pkg/common/config_test.go

@@ -203,6 +203,15 @@ var _ = Describe("Simulator configuration", func() {
 			"{\"running\":\"lora1,lora2\",\"waiting\":\"lora3\",\"timestamp\":1257894567}",
 			"{\"running\":\"lora1,lora3\",\"waiting\":\"\",\"timestamp\":1257894569}",
 		},
+		RequestPromptTokens:     []float64{10, 20, 30, 15},
+		RequestGenerationTokens: []float64{50, 60, 40},
+		RequestParamsMaxTokens:  []float64{128, 256, 512},
+		RequestSuccessTotal: map[string]int64{
+			StopFinishReason:         20,
+			LengthFinishReason:       0,
+			ToolsFinishReason:        0,
+			RemoteDecodeFinishReason: 0,
+		},
 	}
 	test = testCase{
 		name:           "config with fake metrics file",

pkg/llm-d-inference-sim/metrics.go

@@ -20,6 +20,7 @@ package llmdinferencesim
 
 import (
 	"context"
+	"math"
 	"strconv"
 	"strings"
 	"sync"
@@ -94,6 +95,61 @@ func (s *VllmSimulator) createAndRegisterPrometheus() error {
 		return err
 	}
 
+	s.requestPromptTokens = prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Subsystem: "",
+			Name:      "vllm:request_prompt_tokens",
+			Help:      "Number of prefill tokens processed.",
+			Buckets:   build125Buckets(s.config.MaxModelLen),
+		},
+		[]string{vllmapi.PromLabelModelName},
+	)
+	if err := s.registry.Register(s.requestPromptTokens); err != nil {
+		s.logger.Error(err, "Prometheus request_prompt_tokens histogram register failed")
+		return err
+	}
+
+	s.requestGenerationTokens = prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Subsystem: "",
+			Name:      "vllm:request_generation_tokens",
+			Help:      "Number of generation tokens processed.",
+			Buckets:   build125Buckets(s.config.MaxModelLen),
+		},
+		[]string{vllmapi.PromLabelModelName},
+	)
+	if err := s.registry.Register(s.requestGenerationTokens); err != nil {
+		s.logger.Error(err, "Prometheus request_generation_tokens histogram register failed")
+		return err
+	}
+
+	s.requestParamsMaxTokens = prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Subsystem: "",
+			Name:      "vllm:request_params_max_tokens",
+			Help:      "Histogram of the max_tokens request parameter.",
+			Buckets:   build125Buckets(s.config.MaxModelLen),
+		},
+		[]string{vllmapi.PromLabelModelName},
+	)
+	if err := s.registry.Register(s.requestParamsMaxTokens); err != nil {
+		s.logger.Error(err, "Prometheus request_params_max_tokens histogram register failed")
+		return err
+	}
+
+	s.requestSuccessTotal = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Subsystem: "",
+			Name:      "vllm:request_success_total",
+			Help:      "Count of successfully processed requests.",
+		},
+		[]string{vllmapi.PromLabelModelName, vllmapi.PromLabelFinishReason},
+	)
+	if err := s.registry.Register(s.requestSuccessTotal); err != nil {
+		s.logger.Error(err, "Prometheus request_success_total counter register failed")
+		return err
+	}
+
 	s.setInitialPrometheusMetrics()
 
 	return nil
@@ -103,12 +159,25 @@ func (s *VllmSimulator) createAndRegisterPrometheus() error {
 // the fake metrics if set
 func (s *VllmSimulator) setInitialPrometheusMetrics() {
 	var nRunningReqs, nWaitingReqs, kvCacheUsage float64
+	modelName := s.getDisplayedModelName(s.config.Model)
 	if s.config.FakeMetrics != nil {
 		nRunningReqs = float64(s.config.FakeMetrics.RunningRequests)
 		nWaitingReqs = float64(s.config.FakeMetrics.WaitingRequests)
 		kvCacheUsage = float64(s.config.FakeMetrics.KVCacheUsagePercentage)
+		for _, requestPromptToken := range s.config.FakeMetrics.RequestPromptTokens {
+			s.requestPromptTokens.WithLabelValues(modelName).Observe(requestPromptToken)
+		}
+		for _, requestGenerationToken := range s.config.FakeMetrics.RequestGenerationTokens {
+			s.requestGenerationTokens.WithLabelValues(modelName).Observe(requestGenerationToken)
+		}
+		for _, requestParamsMaxToken := range s.config.FakeMetrics.RequestParamsMaxTokens {
+			s.requestParamsMaxTokens.WithLabelValues(modelName).Observe(requestParamsMaxToken)
+		}
+		for reason, requestSuccessTotal := range s.config.FakeMetrics.RequestSuccessTotal {
+			s.requestSuccessTotal.WithLabelValues(modelName, reason).Add(float64(requestSuccessTotal))
+		}
+
 	}
-	modelName := s.getDisplayedModelName(s.config.Model)
 	s.runningRequests.WithLabelValues(modelName).Set(nRunningReqs)
 	s.waitingRequests.WithLabelValues(modelName).Set(nWaitingReqs)
 	s.kvCacheUsagePercentage.WithLabelValues(modelName).Set(kvCacheUsage)
@@ -198,6 +267,7 @@ func (s *VllmSimulator) startMetricsUpdaters(ctx context.Context) {
 	go s.runningRequestsUpdater(ctx)
 	go s.lorasUpdater(ctx)
 	go s.kvCacheUsageUpdater(ctx)
+	go s.recordRequestUpdater(ctx)
 }
 
 // waitingRequestsUpdater updates the waiting requests metric by listening on the relevant channel
@@ -282,3 +352,75 @@ func (s *VllmSimulator) decrementLoraRefCount(lora string, theMap *sync.Map) {
 		s.logger.Error(nil, "Zero model reference", "model", lora)
 	}
 }
+
+// recordRequestUpdater listens on requestSuccessChan and drives the Prometheus metric
+// for successfully completed requests.
+func (s *VllmSimulator) recordRequestUpdater(ctx context.Context) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case event := <-s.requestSuccessChan:
+			s.recordRequestMetricsOnSuccess(
+				event.promptTokens,
+				event.generationTokens,
+				event.maxTokens,
+				event.finishReason,
+			)
+		}
+	}
+}
+
+// requestSuccessEvent represents the data associated with a successfully completed request,
+// which is sent through the requestSuccessChan for asynchronous metrics recording.
+type requestSuccessEvent struct {
+	// promptTokens is the number of input (prompt) tokens in the request
+	promptTokens int
+	// generationTokens is the number of generated (output) tokens in the response
+	generationTokens int
+	// maxTokens is the maximum number of tokens allowed for generation (if specified in the request)
+	maxTokens *int64
+	// finishReason indicates why the generation stopped (e.g., "stop", "length", "tool_calls")
+	finishReason string
+}
+
+// recordRequestMetricsOnSuccess records metrics for a successfully completed request
+func (s *VllmSimulator) recordRequestMetricsOnSuccess(promptTokens,
+	generationTokens int, maxTokens *int64, finishReason string) {
+	modelName := s.getDisplayedModelName(s.config.Model)
+	s.requestPromptTokens.WithLabelValues(modelName).Observe(float64(promptTokens))
+	s.requestGenerationTokens.WithLabelValues(modelName).Observe(float64(generationTokens))
+	if maxTokens != nil {
+		s.requestParamsMaxTokens.WithLabelValues(modelName).Observe(float64(*maxTokens))
+	}
+	s.requestSuccessTotal.WithLabelValues(modelName, finishReason).Inc()
+}
+
+// build125Buckets generates histogram buckets in powers of 10 scaled by [1,2,5].
+// This matches vLLM's build_1_2_5_buckets() in metrics.py.
+//
+// Reference: https://github.com/vllm-project/vllm/blob/main/vllm/engine/metrics.py#L175
+func build125Buckets(maxValue int) []float64 {
+	if maxValue <= 0 {
+		return []float64{}
+	}
+	var buckets []float64
+	exponent := 0
+	mantissa := []int{1, 2, 5}
+
+	for {
+		complete := true
+		for _, m := range mantissa {
+			value := m * int(math.Pow10(exponent))
+			if value <= maxValue {
+				buckets = append(buckets, float64(value))
+				complete = false
+			}
+		}
+		if complete {
+			break
+		}
+		exponent++
+	}
+	return buckets
+}

pkg/llm-d-inference-sim/metrics_test.go

@@ -22,11 +22,13 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"reflect"
 	"regexp"
 	"sort"
 	"strconv"
 	"strings"
 	"sync"
+	"testing"
 	"time"
 
 	"github.com/llm-d/llm-d-inference-sim/pkg/common"
@@ -114,6 +116,49 @@ var _ = Describe("Simulator metrics", Ordered, func() {
 		wg.Wait()
 	})
 
+	It("Should record correct prompt and generation token counts", func() {
+		modelName := "testmodel"
+		prompt := strings.Repeat("hello ", 10)
+		maxTokens := 25
+
+		ctx := context.TODO()
+		args := []string{"cmd", "--model", modelName, "--mode", common.ModeRandom,
+			"--time-to-first-token", "100", "--max-num-seqs", "4"}
+
+		client, err := startServerWithArgs(ctx, common.ModeRandom, args, nil)
+		Expect(err).NotTo(HaveOccurred())
+
+		openaiclient := openai.NewClient(
+			option.WithBaseURL(baseURL),
+			option.WithHTTPClient(client))
+
+		params := openai.ChatCompletionNewParams{
+			Messages: []openai.ChatCompletionMessageParamUnion{
+				openai.UserMessage(prompt),
+			},
+			Model:       modelName,
+			MaxTokens:   openai.Int(int64(maxTokens)),
+			Temperature: openai.Float(0.0),
+		}
+
+		_, err = openaiclient.Chat.Completions.New(ctx, params)
+		Expect(err).NotTo(HaveOccurred())
+
+		time.Sleep(500 * time.Millisecond)
+
+		metricsResp, err := client.Get(metricsUrl)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(metricsResp.StatusCode).To(Equal(http.StatusOK))
+
+		data, err := io.ReadAll(metricsResp.Body)
+		Expect(err).NotTo(HaveOccurred())
+		metrics := string(data)
+		Expect(metrics).To(ContainSubstring(`vllm:request_prompt_tokens_bucket{model_name="testmodel",le="50"} 1`))
+		Expect(metrics).To(ContainSubstring(`vllm:request_params_max_tokens_bucket{model_name="testmodel",le="50"} 1`))
+		Expect(metrics).To(ContainSubstring(`vllm:request_generation_tokens_count{model_name="testmodel"} 1`))
+		Expect(metrics).To(ContainSubstring(`vllm:request_success_total{finish_reason="stop",model_name="testmodel"} 1`))
+	})
+
 	It("Should send correct lora metrics", func() {
 		ctx := context.TODO()
 		args := []string{"cmd", "--model", model, "--mode", common.ModeRandom,
@@ -572,3 +617,77 @@ func splitString(str string) []string {
 	}
 	return strings.Split(str, ",")
 }
+
+// TestBuild125Buckets tests the build125Buckets function with various inputs.
+func TestBuild125Buckets(t *testing.T) {
+	tests := []struct {
+		name     string
+		maxValue int
+		want     []float64
+	}{
+		{
+			name:     "max_value zero",
+			maxValue: 0,
+			want:     []float64{}, // no bucket <= 0
+		},
+		{
+			name:     "max_value one",
+			maxValue: 1,
+			want:     []float64{1},
+		},
+		{
+			name:     "max_value five",
+			maxValue: 5,
+			want:     []float64{1, 2, 5},
+		},
+		{
+			name:     "max_value ten",
+			maxValue: 10,
+			want:     []float64{1, 2, 5, 10},
+		},
+		{
+			name:     "max_value 100",
+			maxValue: 100,
+			want:     []float64{1, 2, 5, 10, 20, 50, 100},
+		},
+		{
+			name:     "max_value 999",
+			maxValue: 999,
+			want:     []float64{1, 2, 5, 10, 20, 50, 100, 200, 500},
+		},
+		{
+			name:     "max_value 1024",
+			maxValue: 1024,
+			want:     []float64{1, 2, 5, 10, 20, 50, 100, 200, 500, 1000},
+		},
+		{
+			name:     "max_value 4096",
+			maxValue: 4096,
+			want:     []float64{1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000},
+		},
+		{
+			name:     "max_value 32768",
+			maxValue: 32768,
+			want:     []float64{1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000},
+		},
+		{
+			name:     "max_value just below power of 10",
+			maxValue: 99,
+			want:     []float64{1, 2, 5, 10, 20, 50},
+		},
+		{
+			name:     "max_value negative",
+			maxValue: -1,
+			want:     []float64{}, // no positive bucket <= -1
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := build125Buckets(tt.maxValue)
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("build125Buckets(%d) = %v, want %v", tt.maxValue, got, tt.want)
+			}
+		})
+	}
+}