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

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

Author: googs1025

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

@@ -94,6 +94,61 @@ func (s *VllmSimulator) createAndRegisterPrometheus() error {
 		return err
 	}
 
+	s.requestPromptTokens = prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Subsystem: "",
+			Name:      "vllm:request_prompt_tokens",
+			Help:      "Number of input prompt tokens in the request.",
+			Buckets:   []float64{1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000},
+		},
+		[]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 generated tokens so far in the request.",
+			Buckets:   []float64{1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000},
+		},
+		[]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:      "The 'max_tokens' parameter from the request.",
+			Buckets:   []float64{1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000},
+		},
+		[]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:      "Total number of successful inference 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
@@ -282,3 +337,15 @@ func (s *VllmSimulator) decrementLoraRefCount(lora string, theMap *sync.Map) {
 		s.logger.Error(nil, "Zero model reference", "model", lora)
 	}
 }
+
+// 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()
+}

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

@@ -108,6 +108,14 @@ type VllmSimulator struct {
 	waitingRequests *prometheus.GaugeVec
 	// kvCacheUsagePercentage is prometheus gauge
 	kvCacheUsagePercentage *prometheus.GaugeVec
+	// requestPromptTokens is prometheus histogram for number of input (prompt) tokens in request
+	requestPromptTokens *prometheus.HistogramVec
+	// requestGenerationTokens is prometheus histogram for number of generated tokens in request
+	requestGenerationTokens *prometheus.HistogramVec
+	// requestParamsMaxTokens is prometheus histogram for 'max_tokens' parameter in request
+	requestParamsMaxTokens *prometheus.HistogramVec
+	// requestSuccessTotal is prometheus counter for total number of successful requests
+	requestSuccessTotal *prometheus.CounterVec
 	// channel for requeasts to be passed to workers
 	reqChan chan *openaiserverapi.CompletionReqCtx
 	// schema validator for tools parameters
@@ -597,9 +605,12 @@ func (s *VllmSimulator) reqProcessingWorker(ctx context.Context, id int) {
 						// in case this is prefill pod processing, return special finish reason
 						finishReason = common.RemoteDecodeFinishReason
 					}
-
 					s.sendResponse(reqCtx, responseTokens, toolCalls, displayModel, finishReason, &usageData)
 				}
+
+				maxTokens, _ := common.GetMaxTokens(reqCtx.CompletionReq.GetMaxCompletionTokens(), nil)
+				s.recordRequestMetricsOnSuccess(usageData.PromptTokens, usageData.CompletionTokens,
+					maxTokens, finishReason)
 			}
 			reqCtx.Wg.Done()
 		}

pkg/vllm-api/vllm-models.go

@@ -25,6 +25,7 @@ const (
 	PromLabelRunningLoraAdapters = "running_lora_adapters"
 	PromLabelMaxLora             = "max_lora"
 	PromLabelModelName           = "model_name"
+	PromLabelFinishReason        = "finish_reason"
 
 	VllmLoraRequestInfo    = "vllm:lora_requests_info"
 	VllmNumRequestsRunning = "vllm:num_requests_running"