Repository: vllm-project/aibrix
Pull request #1096: feature: use worker pool management for periodic metrics update
Author: login: googs1025
Created: 2025-05-18T04:16:49Z | Merged: 2025-05-29T23:57:18Z

## Pull Request Description
[Please provide a clear and concise description of your changes here]

1. Introduce a worker pool in Store initialization to handle pod metrics updates
2. Adjust `updatePodMetrics` method to utilize the job channel for distributing work among workers
3. use context timeouts to prevent hanging operations.

## Related Issues
Resolves: part of https://github.com/vllm-project/aibrix/issues/970

**Important: Before submitting, please complete the description above and review the checklist below.**

---

<details>
<summary><strong>Contribution Guidelines (Expand for Details)</strong></summary>

<p>We appreciate your contribution to aibrix! To ensure a smooth review process and maintain high code quality, please adhere to the following guidelines:</p>

<h3>Pull Request Title Format</h3>
<p>Your PR title should start with one of these prefixes to indicate the nature of the change:</p>
<ul>
    <li><code>[Bug]</code>: Corrections to existing functionality</li>
    <li><code>[CI]</code>: Changes to build process or CI pipeline</li>
    <li><code>[Docs]</code>: Updates or additions to documentation</li>
    <li><code>[API]</code>: Modifications to aibrix's API or interface</li>
    <li><code>[CLI]</code>: Changes or additions to the Command Line Interface</li>
    <li><code>[Misc]</code>: For changes not covered above (use sparingly)</li>
</ul>
<p><em>Note: For changes spanning multiple categories, use multiple prefixes in order of importance.</em></p>

<h3>Submission Checklist</h3>
<ul>
    <li>[ ] PR title includes appropriate prefix(es)</li>
    <li>[ ] Changes are clearly explained in the PR description</li>
    <li>[ ] New and existing tests pass successfully</li>
    <li>[ ] Code adheres to project style and best practices</li>
    <li>[ ] Documentation updated to reflect changes (if applicable)</li>
    <li>[ ] Thorough testing completed, no regressions introduced</li>
</ul>

<p>By submitting this PR, you confirm that you've read these guidelines and your changes align with the project's contribution standards.</p>

</details>


Diff:
diff --git a/pkg/cache/cache_init.go b/pkg/cache/cache_init.go
index e599be3ef..34b5c63df 100644
--- a/pkg/cache/cache_init.go
+++ b/pkg/cache/cache_init.go
@@ -59,6 +59,13 @@ type Store struct {
 	// buffer for sync map operations
 	bufferPod   *Pod
 	bufferModel *Model
+
+	// podMetricsWorkerCount
+	// Number of concurrent workers used to update Pod metrics
+	podMetricsWorkerCount int
+
+	// podMetricsJobs Channel for sending Pod metrics update jobs to workers
+	podMetricsJobs chan *Pod
 }
 
 // Get retrieves the cache instance
@@ -83,12 +90,22 @@ func Get() (Cache, error) {
 //
 //	Store: Initialized cache store instance
 func New(redisClient *redis.Client, prometheusApi prometheusv1.API) *Store {
-	return &Store{
-		initialized:   true,
-		redisClient:   redisClient,
-		prometheusApi: prometheusApi,
-		requestTrace:  &utils.SyncMap[string, *RequestTrace]{},
+
+	store = &Store{
+		initialized:           true,
+		redisClient:           redisClient,
+		prometheusApi:         prometheusApi,
+		requestTrace:          &utils.SyncMap[string, *RequestTrace]{},
+		podMetricsWorkerCount: defaultPodMetricsWorkerCount,
+		podMetricsJobs:        make(chan *Pod, 100), // Initialize the job channel with a buffer size of 100
 	}
+
+	// Start podMetrics worker pool
+	for w := 0; w < store.podMetricsWorkerCount; w++ {
+		go store.worker(store.podMetricsJobs)
+	}
+
+	return store
 }
 
 func NewTestCacheWithPods(pods []*v1.Pod, model string) *Store {
diff --git a/pkg/cache/cache_metrics.go b/pkg/cache/cache_metrics.go
index c3d43e15e..d3330f6a8 100644
--- a/pkg/cache/cache_metrics.go
+++ b/pkg/cache/cache_metrics.go
@@ -110,37 +110,48 @@ func (c *Store) getPodModelMetricName(modelName string, metricName string) strin
 	return fmt.Sprintf("%s/%s", modelName, metricName)
 }
 
+const defaultPodMetricsWorkerCount = 10
+
 func (c *Store) updatePodMetrics() {
 	c.metaPods.Range(func(key string, metaPod *Pod) bool {
 		if !utils.FilterReadyPod(metaPod.Pod) {
 			// Skip unready pod
 			return true
 		}
+		c.podMetricsJobs <- metaPod // Send the job to the worker pool
+		return true
+	})
+}
 
+func (c *Store) worker(jobs <-chan *Pod) {
+	for pod := range jobs {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		// We should use the primary container port. In the future, we can decide whether to use sidecar container's port
-		url := fmt.Sprintf("http://%s:%d/metrics", metaPod.Status.PodIP, podPort)
-		allMetrics, err := metrics.ParseMetricsURL(url)
+		url := fmt.Sprintf("http://%s:%d/metrics", pod.Status.PodIP, podPort)
+		allMetrics, err := metrics.ParseMetricsURLWithContext(ctx, url)
 		if err != nil {
 			klog.V(4).Infof("Error parsing metric families: %v\n", err)
+			cancel()
+			continue
 		}
 
 		// parse counterGaugeMetricsNames
-		c.updateSimpleMetricFromRawMetrics(metaPod, allMetrics)
+		c.updateSimpleMetricFromRawMetrics(pod, allMetrics)
 
 		// parse histogramMetrics
-		c.updateHistogramMetricFromRawMetrics(metaPod, allMetrics)
+		c.updateHistogramMetricFromRawMetrics(pod, allMetrics)
 
 		// parse QueryLabel metrics
-		c.updateQueryLabelMetricFromRawMetrics(metaPod, allMetrics)
+		c.updateQueryLabelMetricFromRawMetrics(pod, allMetrics)
 
 		if c.prometheusApi == nil {
 			klog.V(4).InfoS("Prometheus api is not initialized, PROMETHEUS_ENDPOINT is not configured, skip fetching prometheus metrics")
-			return true
+		} else {
+			// parse prometheus metrics
+			c.updateMetricFromPromQL(ctx, pod)
 		}
-		// parse prometheus metrics
-		c.updateMetricFromPromQL(metaPod)
-		return true
-	})
+		cancel()
+	}
 }
 
 func (c *Store) updateSimpleMetricFromRawMetrics(pod *Pod, allMetrics map[string]*dto.MetricFamily) {
@@ -247,7 +258,7 @@ func (c *Store) updateQueryLabelMetricFromRawMetrics(pod *Pod, allMetrics map[st
 	}
 }
 
-func (c *Store) updateMetricFromPromQL(pod *Pod) {
+func (c *Store) updateMetricFromPromQL(ctx context.Context, pod *Pod) {
 	podName := pod.Name
 
 	for _, metricName := range prometheusMetricNames {
@@ -261,7 +272,7 @@ func (c *Store) updateMetricFromPromQL(pod *Pod) {
 		}
 		scope := metric.MetricScope
 		if scope == metrics.PodMetricScope {
-			err := c.queryUpdatePromQLMetrics(metric, queryLabels, pod, "", metricName)
+			err := c.queryUpdatePromQLMetrics(ctx, metric, queryLabels, pod, "", metricName)
 			if err != nil {
 				klog.V(4).Infof("Failed to query and update PromQL metrics: %v", err)
 				continue
@@ -270,7 +281,7 @@ func (c *Store) updateMetricFromPromQL(pod *Pod) {
 			if pod.Models.Len() > 0 {
 				for _, modelName := range pod.Models.Array() {
 					queryLabels["model_name"] = modelName
-					err := c.queryUpdatePromQLMetrics(metric, queryLabels, pod, modelName, metricName)
+					err := c.queryUpdatePromQLMetrics(ctx, metric, queryLabels, pod, modelName, metricName)
 					if err != nil {
 						klog.V(4).Infof("Failed to query and update PromQL metrics: %v", err)
 						continue
@@ -285,11 +296,11 @@ func (c *Store) updateMetricFromPromQL(pod *Pod) {
 	}
 }
 
-func (c *Store) queryUpdatePromQLMetrics(metric metrics.Metric, queryLabels map[string]string, pod *Pod, modelName string, metricName string) error {
+func (c *Store) queryUpdatePromQLMetrics(ctx context.Context, metric metrics.Metric, queryLabels map[string]string, pod *Pod, modelName string, metricName string) error {
 	scope := metric.MetricScope
 	query := metrics.BuildQuery(metric.PromQL, queryLabels)
 	// Querying metrics
-	result, warnings, err := c.prometheusApi.Query(context.Background(), query, time.Now())
+	result, warnings, err := c.prometheusApi.Query(ctx, query, time.Now())
 	if err != nil {
 		// Skip this model fetching if an error is thrown
 		return fmt.Errorf("error executing query: %v", err)
diff --git a/pkg/metrics/utils.go b/pkg/metrics/utils.go
index ea00649d1..47f47d343 100644
--- a/pkg/metrics/utils.go
+++ b/pkg/metrics/utils.go
@@ -17,6 +17,7 @@ limitations under the License.
 package metrics
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"regexp"
@@ -201,10 +202,15 @@ func GetHistogramValue(metric *dto.Metric) (*HistogramMetricValue, error) {
 	return histogram, nil
 }
 
-func ParseMetricsURL(url string) (map[string]*dto.MetricFamily, error) {
-	resp, err := http.Get(url)
+func ParseMetricsURLWithContext(ctx context.Context, url string) (map[string]*dto.MetricFamily, error) {
+	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
 	if err != nil {
-		return make(map[string]*dto.MetricFamily), fmt.Errorf("Failed to fetch metrics from %s: %v", url, err)
+		return nil, fmt.Errorf("Failed to create request for %s: %v\n", url, err)
+	}
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return make(map[string]*dto.MetricFamily), fmt.Errorf("Failed to fetch metrics from %s: %v\n", url, err)
 	}
 	defer func() {
 		if err := resp.Body.Close(); err != nil {
@@ -212,6 +218,10 @@ func ParseMetricsURL(url string) (map[string]*dto.MetricFamily, error) {
 		}
 	}()
 
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("Bad status code while fetching metrics from %s: %d\n", url, resp.StatusCode)
+	}
+
 	var parser expfmt.TextParser
 	allMetrics, err := parser.TextToMetricFamilies(resp.Body)
 	if err != nil {
diff --git a/pkg/metrics/utils_test.go b/pkg/metrics/utils_test.go
index 6bada534a..cb2e0966b 100644
--- a/pkg/metrics/utils_test.go
+++ b/pkg/metrics/utils_test.go
@@ -17,7 +17,13 @@ limitations under the License.
 package metrics
 
 import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -89,3 +95,88 @@ func TestBuildQueryWithModelLabel(t *testing.T) {
 		assert.Contains(t, query, `model_name="Qwen/Qwen2.5-1.5B-Instruct"`)
 	})
 }
+
+func TestParseMetricsURLWithContext(t *testing.T) {
+	tests := []struct {
+		name           string
+		metricsContent string
+		statusCode     int
+		timeout        time.Duration
+		wantErr        bool
+		errContains    string
+	}{
+		{
+			name: "Success with valid metrics",
+			metricsContent: `
+				# HELP http_requests_total The total number of HTTP requests.
+				# TYPE http_requests_total counter
+				http_requests_total{method="post",code="200"} 1027 1395066363000
+			`,
+			statusCode: http.StatusOK,
+			timeout:    5 * time.Second,
+			wantErr:    false,
+		},
+		{
+			name:           "HTTP non-200 status code",
+			metricsContent: "",
+			statusCode:     http.StatusServiceUnavailable,
+			timeout:        5 * time.Second,
+			wantErr:        true,
+			errContains:    "Bad status code while fetching metrics",
+		},
+		{
+			name: "Context timeout",
+			metricsContent: `
+				http_requests_total{method="post",code="200"} 1027
+			`,
+			statusCode:  http.StatusOK,
+			timeout:     1 * time.Nanosecond,
+			wantErr:     true,
+			errContains: "context deadline exceeded",
+		},
+		{
+			name:           "Invalid metric format",
+			metricsContent: "invalid_metric_format_here",
+			statusCode:     http.StatusOK,
+			timeout:        5 * time.Second,
+			wantErr:        true,
+			errContains:    "Error parsing metric families",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Setup mock server
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(tt.statusCode)
+				_, err := io.WriteString(w, tt.metricsContent)
+				if err != nil {
+					return
+				}
+			}))
+			defer server.Close()
+
+			ctx, cancel := context.WithTimeout(context.Background(), tt.timeout)
+			defer cancel()
+			metrics, err := ParseMetricsURLWithContext(ctx, server.URL)
+
+			if tt.wantErr {
+				if err == nil {
+					t.Fatal("expected error but got none")
+				}
+				if tt.errContains == "context deadline exceeded" {
+					if !strings.Contains(err.Error(), tt.errContains) {
+						t.Fatalf("expected error to contain %q, got %q", tt.errContains, err.Error())
+					}
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+			if len(metrics) == 0 {
+				t.Errorf("expected non-empty metrics map")
+			}
+		})
+	}
+}


Reviews:
user: login: googs1025
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-19T02:23:08Z
user: login: googs1025
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-27T02:22:12Z
user: login: googs1025
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-27T02:22:53Z
user: login: varungup90
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-29T23:57:05Z

Review comments:
diff_hunk: @@ -285,11 +296,11 @@ func (c *Store) updateMetricFromPromQL(pod *Pod) {
 	}
 }
 
-func (c *Store) queryUpdatePromQLMetrics(metric metrics.Metric, queryLabels map[string]string, pod *Pod, modelName string, metricName string) error {
+func (c *Store) queryUpdatePromQLMetrics(ctx context.Context, metric metrics.Metric, queryLabels map[string]string, pod *Pod, modelName string, metricName string) error {
path: pkg/cache/cache_metrics.go
user: login: googs1025
body: pass `ctx` to handle timeout case 
created_at: 2025-05-19T02:23:08Z
author_association: COLLABORATOR
diff_hunk: @@ -110,37 +110,48 @@ func (c *Store) getPodModelMetricName(modelName string, metricName string) strin
 	return fmt.Sprintf("%s/%s", modelName, metricName)
 }
 
+const defaultPodMetricsWorkerCount = 10
+
 func (c *Store) updatePodMetrics() {
 	c.metaPods.Range(func(key string, metaPod *Pod) bool {
 		if !utils.FilterReadyPod(metaPod.Pod) {
 			// Skip unready pod
 			return true
 		}
+		c.podMetricsJobs <- metaPod // Send the job to the worker pool
+		return true
+	})
+}
 
+func (c *Store) worker(jobs <-chan *Pod) {
+	for pod := range jobs {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
path: pkg/cache/cache_metrics.go
user: login: googs1025
body: Using a timeout of 5 seconds, I'm not sure if this is appropriate? 🤔 
created_at: 2025-05-27T02:22:12Z
author_association: COLLABORATOR
diff_hunk: @@ -110,37 +110,48 @@ func (c *Store) getPodModelMetricName(modelName string, metricName string) strin
 	return fmt.Sprintf("%s/%s", modelName, metricName)
 }
 
+const defaultPodMetricsWorkerCount = 10
+
 func (c *Store) updatePodMetrics() {
 	c.metaPods.Range(func(key string, metaPod *Pod) bool {
 		if !utils.FilterReadyPod(metaPod.Pod) {
 			// Skip unready pod
 			return true
 		}
+		c.podMetricsJobs <- metaPod // Send the job to the worker pool
+		return true
+	})
+}
 
+func (c *Store) worker(jobs <-chan *Pod) {
+	for pod := range jobs {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		// We should use the primary container port. In the future, we can decide whether to use sidecar container's port
-		url := fmt.Sprintf("http://%s:%d/metrics", metaPod.Status.PodIP, podPort)
-		allMetrics, err := metrics.ParseMetricsURL(url)
+		url := fmt.Sprintf("http://%s:%d/metrics", pod.Status.PodIP, podPort)
+		allMetrics, err := metrics.ParseMetricsURLWithContext(ctx, url)
 		if err != nil {
 			klog.V(4).Infof("Error parsing metric families: %v\n", err)
+			cancel()
+			continue
path: pkg/cache/cache_metrics.go
user: login: googs1025
body: if parsing error, we don't update cache metrics
created_at: 2025-05-27T02:22:53Z
author_association: COLLABORATOR

PR comments:
user: login: googs1025
created_at: 2025-05-19T02:28:11Z
body: In large-scale scenarios, we want multiple goroutines to share the processing of querying metrics. We use the method of `multiple goroutines` + `chan`. 🤔 
author_association: COLLABORATOR
user: login: googs1025
created_at: 2025-05-19T02:29:00Z
body: friendly ping @Jeffwan @varungup90  😄 
author_association: COLLABORATOR
user: login: varungup90
created_at: 2025-05-19T16:41:14Z
body: Can you add more unit tests. Overall LGTM.

@zhangjyr PTAL
author_association: COLLABORATOR
user: login: Jeffwan
created_at: 2025-05-20T21:35:06Z
body: /hold
we can merge this one after release. 
author_association: COLLABORATOR
user: login: googs1025
created_at: 2025-05-25T07:27:26Z
body: > Can you add more unit tests.

I'll deal with it in the next few days. 😄 


author_association: COLLABORATOR
user: login: googs1025
created_at: 2025-05-27T02:58:18Z
body: > Can you add more unit tests. Overall LGTM.

Looking at the current structure of the Store, it is not easy to write unit tests. Maybe I hope to fix this in other PRs.
I hope to optimize `Store` structure in other PRs to make it better for writing unit tests 


author_association: COLLABORATOR