Add totalRunningRequestsMetric to EPP pod metrics for upcomming latency prediciton features (#1899)

BenjaminBraunDev · web-flow · commit 5debae8c380b · 2025-11-25T09:06:37.000-08:00
diff --git a/cmd/epp/runner/runner.go b/cmd/epp/runner/runner.go
@@ -126,6 +126,7 @@ var (
 		"then a self-signed certificate is used.")
 	// metric flags
 	totalQueuedRequestsMetric    = flag.String("total-queued-requests-metric", runserver.DefaultTotalQueuedRequestsMetric, "Prometheus metric for the number of queued requests.")
+	totalRunningRequestsMetric   = flag.String("total-running-requests-metric", runserver.DefaultTotalRunningRequestsMetric, "Prometheus metric for the number of running requests.")
 	kvCacheUsagePercentageMetric = flag.String("kv-cache-usage-percentage-metric", runserver.DefaultKvCacheUsagePercentageMetric, "Prometheus metric for the fraction of KV-cache blocks currently in use (from 0 to 1).")
 	// LoRA metrics
 	loraInfoMetric = flag.String("lora-info-metric", runserver.DefaultLoraInfoMetric, "Prometheus metric for the LoRA info metrics (must be in vLLM label format).")
@@ -553,6 +554,7 @@ func (r *Runner) setupMetricsCollection(setupLog logr.Logger, useExperimentalDat
 func setupMetricsV1(setupLog logr.Logger) (datalayer.EndpointFactory, error) {
 	mapping, err := backendmetrics.NewMetricMapping(
 		*totalQueuedRequestsMetric,
+		*totalRunningRequestsMetric,
 		*kvCacheUsagePercentageMetric,
 		*loraInfoMetric,
 		*cacheInfoMetric,
@@ -601,6 +603,7 @@ func setupDatalayer(logger logr.Logger) (datalayer.EndpointFactory, error) {
 		*modelServerMetricsHttpsInsecureSkipVerify,
 		nil)
 	extractor, err := dlmetrics.NewExtractor(*totalQueuedRequestsMetric,
+		*totalRunningRequestsMetric,
 		*kvCacheUsagePercentageMetric,
 		*loraInfoMetric, *cacheInfoMetric)
 
diff --git a/pkg/epp/backend/metrics/metrics.go b/pkg/epp/backend/metrics/metrics.go
@@ -97,6 +97,15 @@ func (p *PodMetricsClientImpl) promToPodMetrics(
 		}
 	}
 
+	if p.MetricMapping.TotalRunningRequests != nil {
+		running, err := p.getMetric(metricFamilies, *p.MetricMapping.TotalRunningRequests)
+		if err == nil {
+			updated.RunningQueueSize = int(running.GetGauge().GetValue())
+		} else {
+			errs = multierr.Append(errs, err)
+		}
+	}
+
 	if p.MetricMapping.KVCacheUtilization != nil {
 		usage, err := p.getMetric(metricFamilies, *p.MetricMapping.KVCacheUtilization)
 		if err == nil {
diff --git a/pkg/epp/backend/metrics/metrics_spec.go b/pkg/epp/backend/metrics/metrics_spec.go
@@ -29,10 +29,11 @@ type MetricSpec struct {
 
 // MetricMapping holds named MetricSpecs.
 type MetricMapping struct {
-	TotalQueuedRequests *MetricSpec
-	KVCacheUtilization  *MetricSpec
-	LoraRequestInfo     *MetricSpec
-	CacheConfigInfo     *MetricSpec
+	TotalQueuedRequests  *MetricSpec
+	TotalRunningRequests *MetricSpec
+	KVCacheUtilization   *MetricSpec
+	LoraRequestInfo      *MetricSpec
+	CacheConfigInfo      *MetricSpec
 }
 
 // stringToMetricSpec converts a string to a MetricSpec.
@@ -94,11 +95,15 @@ func stringToMetricSpec(specStr string) (*MetricSpec, error) {
 }
 
 // NewMetricMapping creates a MetricMapping from string values.
-func NewMetricMapping(queuedStr, kvUsageStr, loraReqInfoStr, cacheInfoMetric string) (*MetricMapping, error) {
+func NewMetricMapping(queuedStr, runningStr, kvUsageStr, loraReqInfoStr, cacheInfoMetric string) (*MetricMapping, error) {
 	queuedSpec, err := stringToMetricSpec(queuedStr)
 	if err != nil {
 		return nil, fmt.Errorf("error parsing WaitingRequests: %w", err)
 	}
+	runningSpec, err := stringToMetricSpec(runningStr)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing RunningRequests: %w", err)
+	}
 	kvUsageSpec, err := stringToMetricSpec(kvUsageStr)
 	if err != nil {
 		return nil, fmt.Errorf("error parsing KVCacheUsage: %w", err)
@@ -114,10 +119,11 @@ func NewMetricMapping(queuedStr, kvUsageStr, loraReqInfoStr, cacheInfoMetric str
 	}
 
 	mapping := &MetricMapping{
-		TotalQueuedRequests: queuedSpec,
-		KVCacheUtilization:  kvUsageSpec,
-		LoraRequestInfo:     loraReqInfoSpec,
-		CacheConfigInfo:     cacheInfoSpec,
+		TotalQueuedRequests:  queuedSpec,
+		TotalRunningRequests: runningSpec,
+		KVCacheUtilization:   kvUsageSpec,
+		LoraRequestInfo:      loraReqInfoSpec,
+		CacheConfigInfo:      cacheInfoSpec,
 	}
 
 	return mapping, nil
diff --git a/pkg/epp/datalayer/metrics/datasource_test.go b/pkg/epp/datalayer/metrics/datasource_test.go
@@ -29,7 +29,7 @@ import (
 
 func TestDatasource(t *testing.T) {
 	source := NewDataSource("https", "/metrics", true, nil)
-	extractor, err := NewExtractor(defaultTotalQueuedRequestsMetric, "", "", "")
+	extractor, err := NewExtractor(defaultTotalQueuedRequestsMetric, "", "", "", "")
 	assert.Nil(t, err, "failed to create extractor")
 
 	name := source.Name()
diff --git a/pkg/epp/datalayer/metrics/extractor.go b/pkg/epp/datalayer/metrics/extractor.go
@@ -66,8 +66,8 @@ func Produces() map[string]any {
 // configured with the given metrics' specifications.
 // These are mandatory metrics per the MSP specification, and are used
 // as the basis for the built-in scheduling plugins.
-func NewExtractor(queueSpec, kvusageSpec, loraSpec, cacheInfoSpec string) (*Extractor, error) {
-	mapping, err := NewMapping(queueSpec, kvusageSpec, loraSpec, cacheInfoSpec)
+func NewExtractor(queueSpec, runningSpec, kvusageSpec, loraSpec, cacheInfoSpec string) (*Extractor, error) {
+	mapping, err := NewMapping(queueSpec, runningSpec, kvusageSpec, loraSpec, cacheInfoSpec)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create extractor metrics Mapping - %w", err)
 	}
@@ -109,6 +109,15 @@ func (ext *Extractor) Extract(ctx context.Context, data any, ep datalayer.Endpoi
 		}
 	}
 
+	if spec := ext.mapping.TotalRunningRequests; spec != nil { // extract running requests
+		if metric, err := spec.getLatestMetric(families); err != nil {
+			errs = append(errs, err)
+		} else {
+			clone.RunningQueueSize = int(extractValue(metric))
+			updated = true
+		}
+	}
+
 	if spec := ext.mapping.KVCacheUtilization; spec != nil { // extract KV cache usage
 		if metric, err := spec.getLatestMetric(families); err != nil {
 			errs = append(errs, err)
diff --git a/pkg/epp/datalayer/metrics/extractor_test.go b/pkg/epp/datalayer/metrics/extractor_test.go
@@ -31,6 +31,7 @@ import (
 const (
 	// use hardcoded values - importing causes cycle
 	defaultTotalQueuedRequestsMetric    = "vllm:num_requests_waiting"
+	defaultTotalRunningRequestsMetric   = "vllm:num_requests_running"
 	defaultKvCacheUsagePercentageMetric = "vllm:gpu_cache_usage_perc"
 	defaultLoraInfoMetric               = "vllm:lora_requests_info"
 	defaultCacheInfoMetric              = "vllm:cache_config_info"
@@ -39,11 +40,11 @@ const (
 func TestExtractorExtract(t *testing.T) {
 	ctx := context.Background()
 
-	if _, err := NewExtractor("vllm: dummy", "", "", ""); err == nil {
+	if _, err := NewExtractor("vllm: dummy", "", "", "", ""); err == nil {
 		t.Error("expected to fail to create extractor with invalid specification")
 	}
 
-	extractor, err := NewExtractor(defaultTotalQueuedRequestsMetric,
+	extractor, err := NewExtractor(defaultTotalQueuedRequestsMetric, defaultTotalRunningRequestsMetric,
 		defaultKvCacheUsagePercentageMetric, defaultLoraInfoMetric, defaultCacheInfoMetric)
 	if err != nil {
 		t.Fatalf("failed to create extractor: %v", err)
@@ -106,6 +107,14 @@ func TestExtractorExtract(t *testing.T) {
 						},
 					},
 				},
+				defaultTotalRunningRequestsMetric: &dto.MetricFamily{
+					Type: dto.MetricType_GAUGE.Enum(),
+					Metric: []*dto.Metric{
+						{
+							Gauge: &dto.Gauge{Value: ptr.To(1.0)},
+						},
+					},
+				},
 				defaultKvCacheUsagePercentageMetric: &dto.MetricFamily{
 					Type: dto.MetricType_GAUGE.Enum(),
 					Metric: []*dto.Metric{
diff --git a/pkg/epp/datalayer/metrics/mapping.go b/pkg/epp/datalayer/metrics/mapping.go
@@ -23,20 +23,25 @@ import (
 // Mapping holds specifications for the well-known metrics defined
 // in the Model Server Protocol.
 type Mapping struct {
-	TotalQueuedRequests *Spec
-	KVCacheUtilization  *Spec
-	LoraRequestInfo     *LoRASpec
-	CacheInfo           *Spec
+	TotalQueuedRequests  *Spec
+	TotalRunningRequests *Spec
+	KVCacheUtilization   *Spec
+	LoraRequestInfo      *LoRASpec
+	CacheInfo            *Spec
 }
 
 // NewMapping creates a metrics.Mapping from the input specification strings.
-func NewMapping(queue, kvusage, lora, cacheInfo string) (*Mapping, error) {
+func NewMapping(queue, running, kvusage, lora, cacheInfo string) (*Mapping, error) {
 	var errs []error
 
 	queueSpec, err := parseStringToSpec(queue)
 	if err != nil {
 		errs = append(errs, err)
 	}
+	runningSpec, err := parseStringToSpec(running)
+	if err != nil {
+		errs = append(errs, err)
+	}
 	kvusageSpec, err := parseStringToSpec(kvusage)
 	if err != nil {
 		errs = append(errs, err)
@@ -55,9 +60,10 @@ func NewMapping(queue, kvusage, lora, cacheInfo string) (*Mapping, error) {
 		return nil, errors.Join(errs...)
 	}
 	return &Mapping{
-		TotalQueuedRequests: queueSpec,
-		KVCacheUtilization:  kvusageSpec,
-		LoraRequestInfo:     loraSpec,
-		CacheInfo:           cacheInfoSpec,
+		TotalQueuedRequests:  queueSpec,
+		TotalRunningRequests: runningSpec,
+		KVCacheUtilization:   kvusageSpec,
+		LoraRequestInfo:      loraSpec,
+		CacheInfo:            cacheInfoSpec,
 	}, nil
 }
diff --git a/pkg/epp/server/runserver.go b/pkg/epp/server/runserver.go
@@ -78,6 +78,7 @@ const (
 	DefaultHealthChecking                   = false                         // default for --health-checking
 	DefaultEnablePprof                      = true                          // default for --enable-pprof
 	DefaultTotalQueuedRequestsMetric        = "vllm:num_requests_waiting"   // default for --total-queued-requests-metric
+	DefaultTotalRunningRequestsMetric       = "vllm:num_requests_running"   // default for --total-running-requests-metric
 	DefaultKvCacheUsagePercentageMetric     = "vllm:gpu_cache_usage_perc"   // default for --kv-cache-usage-percentage-metric
 	DefaultLoraInfoMetric                   = "vllm:lora_requests_info"     // default for --lora-info-metric
 	DefaultCacheInfoMetric                  = "vllm:cache_config_info"      // default for --cache-info-metric
