Merge pull request #33 from liu-cong/fix

Fix mutierr appending; add a unit test.

Author: k8s-ci-robot

pkg/ext-proc/backend/fake.go

@@ -1,5 +1,7 @@
 package backend
 
+import "context"
+
 type FakePodLister struct {
 	Err  error
 	Pods PodSet
@@ -10,7 +12,7 @@ type FakePodMetricsClient struct {
 	Res map[Pod]*PodMetrics
 }
 
-func (f *FakePodMetricsClient) FetchMetrics(pod Pod, existing *PodMetrics) (*PodMetrics, error) {
+func (f *FakePodMetricsClient) FetchMetrics(ctx context.Context, pod Pod, existing *PodMetrics) (*PodMetrics, error) {
 	if err, ok := f.Err[pod]; ok {
 		return nil, err
 	}

pkg/ext-proc/backend/provider.go

@@ -1,6 +1,7 @@
 package backend
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"time"
@@ -9,6 +10,10 @@ import (
 	klog "k8s.io/klog/v2"
 )
 
+const (
+	fetchMetricsTimeout = 5 * time.Second
+)
+
 func NewProvider(pmc PodMetricsClient, pl PodLister) *Provider {
 	p := &Provider{
 		podMetrics: sync.Map{},
@@ -27,7 +32,7 @@ type Provider struct {
 }
 
 type PodMetricsClient interface {
-	FetchMetrics(pod Pod, existing *PodMetrics) (*PodMetrics, error)
+	FetchMetrics(ctx context.Context, pod Pod, existing *PodMetrics) (*PodMetrics, error)
 }
 
 type PodLister interface {
@@ -60,7 +65,8 @@ func (p *Provider) Init(refreshPodsInterval, refreshMetricsInterval time.Duratio
 	if err := p.refreshPodsOnce(); err != nil {
 		return fmt.Errorf("failed to init pods: %v", err)
 	}
-	if err := p.refreshMetricsOnce(); err != nil {
+	err := p.refreshMetricsOnce()
+	if err != nil {
 		return fmt.Errorf("failed to init metrics: %v", err)
 	}
 
@@ -132,35 +138,48 @@ func (p *Provider) refreshPodsOnce() error {
 }
 
 func (p *Provider) refreshMetricsOnce() error {
+	ctx, cancel := context.WithTimeout(context.Background(), fetchMetricsTimeout)
+	defer cancel()
 	start := time.Now()
 	defer func() {
 		d := time.Since(start)
 		// TODO: add a metric instead of logging
 		klog.V(4).Infof("Refreshed metrics in %v", d)
 	}()
 	var wg sync.WaitGroup
-	var errs error
+	errCh := make(chan error)
 	processOnePod := func(key, value any) bool {
 		klog.V(4).Infof("Processing pod %v and metric %v", key, value)
 		pod := key.(Pod)
 		existing := value.(*PodMetrics)
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			updated, err := p.pmc.FetchMetrics(pod, existing)
+			updated, err := p.pmc.FetchMetrics(ctx, pod, existing)
 			if err != nil {
-				multierr.Append(errs, fmt.Errorf("failed to parse metrics from %s: %v", pod, err))
+				errCh <- fmt.Errorf("failed to parse metrics from %s: %v", pod, err)
 				return
 			}
-			klog.V(4).Infof("Updated metrics for pod %s: %v", pod, updated.Metrics)
-			if err != nil {
-				multierr.Append(errs, fmt.Errorf("failed to get all pod metrics updated from prometheus: %v", err))
-			}
 			p.UpdatePodMetrics(pod, updated)
+			klog.V(4).Infof("Updated metrics for pod %s: %v", pod, updated.Metrics)
 		}()
 		return true
 	}
 	p.podMetrics.Range(processOnePod)
-	wg.Wait()
+
+	// Wait for metric collection for all pods to complete and close the error channel in a
+	// goroutine so this is unblocking, allowing the code to proceed to the error collection code
+	// below.
+	// Note we couldn't use a buffered error channel with a size because the size of the podMetrics
+	// sync.Map is unknown beforehand.
+	go func() {
+		wg.Wait()
+		close(errCh)
+	}()
+
+	var errs error
+	for err := range errCh {
+		errs = multierr.Append(errs, err)
+	}
 	return errs
 }

pkg/ext-proc/backend/provider_test.go

@@ -0,0 +1,111 @@
+package backend
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+)
+
+var (
+	pod1 = &PodMetrics{
+		Pod: Pod{Name: "pod1"},
+		Metrics: Metrics{
+			WaitingQueueSize:    0,
+			KVCacheUsagePercent: 0.2,
+			MaxActiveModels:     2,
+			ActiveModels: map[string]int{
+				"foo": 1,
+				"bar": 1,
+			},
+		},
+	}
+	pod2 = &PodMetrics{
+		Pod: Pod{Name: "pod2"},
+		Metrics: Metrics{
+			WaitingQueueSize:    1,
+			KVCacheUsagePercent: 0.2,
+			MaxActiveModels:     2,
+			ActiveModels: map[string]int{
+				"foo1": 1,
+				"bar1": 1,
+			},
+		},
+	}
+)
+
+func TestProvider(t *testing.T) {
+	tests := []struct {
+		name    string
+		pmc     PodMetricsClient
+		pl      PodLister
+		initErr bool
+		want    []*PodMetrics
+	}{
+		{
+			name: "Init success",
+			pl: &FakePodLister{
+				Pods: map[Pod]bool{
+					pod1.Pod: true,
+					pod2.Pod: true,
+				},
+			},
+			pmc: &FakePodMetricsClient{
+				Res: map[Pod]*PodMetrics{
+					pod1.Pod: pod1,
+					pod2.Pod: pod2,
+				},
+			},
+			want: []*PodMetrics{pod1, pod2},
+		},
+		{
+			name: "Fetch metrics error",
+			pl: &FakePodLister{
+				Pods: map[Pod]bool{
+					pod1.Pod: true,
+					pod2.Pod: true,
+				},
+			},
+			pmc: &FakePodMetricsClient{
+				Err: map[Pod]error{
+					pod2.Pod: errors.New("injected error"),
+				},
+				Res: map[Pod]*PodMetrics{
+					pod1.Pod: pod1,
+				},
+			},
+			initErr: true,
+			want: []*PodMetrics{
+				pod1,
+				// Failed to fetch pod2 metrics so it remains the default values.
+				&PodMetrics{
+					Pod: Pod{Name: "pod2"},
+					Metrics: Metrics{
+						WaitingQueueSize:    0,
+						KVCacheUsagePercent: 0,
+						MaxActiveModels:     0,
+						ActiveModels:        map[string]int{},
+					},
+				}},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			p := NewProvider(test.pmc, test.pl)
+			err := p.Init(time.Millisecond, time.Millisecond)
+			if test.initErr != (err != nil) {
+				t.Fatalf("Unexpected error, got: %v, want: %v", err, test.initErr)
+			}
+			metrics := p.AllPodMetrics()
+			lessFunc := func(a, b *PodMetrics) bool {
+				return a.String() < b.String()
+			}
+			if diff := cmp.Diff(test.want, metrics, cmpopts.SortSlices(lessFunc)); diff != "" {
+				t.Errorf("Unexpected output (-want +got): %v", diff)
+			}
+		})
+	}
+}

pkg/ext-proc/backend/vllm/metrics.go

@@ -2,6 +2,7 @@
 package vllm
 
 import (
+	"context"
 	"ext-proc/backend"
 	"fmt"
 	"net/http"
@@ -15,8 +16,7 @@ import (
 )
 
 const (
-	ActiveLoRAAdaptersMetricName        = "vllm:info_active_adapters_info"
-	LoRAAdapterPendingRequestMetricName = "vllm:active_lora_adapters"
+	ActiveLoRAAdaptersMetricName = "vllm:info_active_adapters_info"
 	// TODO: Replace these with the num_tokens_running/waiting below once we add those to the fork.
 	RunningQueueSizeMetricName = "vllm:num_requests_running"
 	WaitingQueueSizeMetricName = "vllm:num_requests_waiting"
@@ -32,11 +32,15 @@ type PodMetricsClientImpl struct {
 }
 
 // FetchMetrics fetches metrics from a given pod.
-func (p *PodMetricsClientImpl) FetchMetrics(pod backend.Pod, existing *backend.PodMetrics) (*backend.PodMetrics, error) {
+func (p *PodMetricsClientImpl) FetchMetrics(ctx context.Context, pod backend.Pod, existing *backend.PodMetrics) (*backend.PodMetrics, error) {
 	// Currently the metrics endpoint is hard-coded, which works with vLLM.
 	// TODO(https://github.com/kubernetes-sigs/llm-instance-gateway/issues/16): Consume this from LLMServerPool config.
 	url := fmt.Sprintf("http://%s/metrics", pod.Address)
-	resp, err := http.Get(url)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create request: %v", err)
+	}
+	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		klog.Errorf("failed to fetch metrics from %s: %v", pod, err)
 		return nil, fmt.Errorf("failed to fetch metrics from %s: %w", pod, err)
@@ -63,23 +67,23 @@ func promToPodMetrics(metricFamilies map[string]*dto.MetricFamily, existing *bac
 	var errs error
 	updated := existing.Clone()
 	runningQueueSize, _, err := getLatestMetric(metricFamilies, RunningQueueSizeMetricName)
-	multierr.Append(errs, err)
+	errs = multierr.Append(errs, err)
 	if err == nil {
 		updated.RunningQueueSize = int(runningQueueSize.GetGauge().GetValue())
 	}
 	waitingQueueSize, _, err := getLatestMetric(metricFamilies, WaitingQueueSizeMetricName)
-	multierr.Append(errs, err)
+	errs = multierr.Append(errs, err)
 	if err == nil {
 		updated.WaitingQueueSize = int(waitingQueueSize.GetGauge().GetValue())
 	}
 	cachePercent, _, err := getLatestMetric(metricFamilies, KVCacheUsagePercentMetricName)
-	multierr.Append(errs, err)
+	errs = multierr.Append(errs, err)
 	if err == nil {
 		updated.KVCacheUsagePercent = cachePercent.GetGauge().GetValue()
 	}
 	/* TODO: uncomment once this is available in vllm.
 	kvCap, _, err := getGaugeLatestValue(metricFamilies, KvCacheMaxTokenCapacityMetricName)
-	multierr.Append(errs, err)
+	errs = multierr.Append(errs, err)
 	if err != nil {
 		updated.KvCacheMaxTokenCapacity = int(kvCap)
 	}
@@ -107,7 +111,7 @@ func promToPodMetrics(metricFamilies map[string]*dto.MetricFamily, existing *bac
 		}
 	} else {
 		klog.Warningf("metric family %q not found", ActiveLoRAAdaptersMetricName)
-		multierr.Append(errs, fmt.Errorf("metric family %q not found", ActiveLoRAAdaptersMetricName))
+		errs = multierr.Append(errs, fmt.Errorf("metric family %q not found", ActiveLoRAAdaptersMetricName))
 	}
 
 	return updated, errs