Repository: vllm-project/aibrix
File: pkg/cache/cache.go (Go)

Code before:

func (c *Cache) updatePodMetrics() {
	c.mu.Lock()
	defer c.mu.Unlock()

	for _, pod := range c.pods {
		if pod.Status.PodIP == "" {
			continue
		}
		podName := pod.Name
		for _, metricName := range metricNames {
			metricVal, err := getMetricsForPod(pod, metricName, 8000)
			if err != nil {
				continue
			}
			if len(c.podMetrics[podName]) == 0 {
				c.podMetrics[podName] = map[string]float64{}
			}
			c.podMetrics[pod.Name][metricName] = metricVal
		}
	}
}

func getMetricsForPod(pod *v1.Pod, metricName string, metricsPort int) (float64, error) {
	// We should use the primary container port. In future, we can decide whether to use sidecar container's port
	url := fmt.Sprintf("http://%s:%d/metrics", pod.Status.PodIP, metricsPort)
	//// TODO a temp for debugging
	//url := fmt.Sprintf("http://%s:%d/metrics", "127.0.0.1", metricsPort)

	// scrape metrics
	resp, err := http.Get(url)
	if err != nil {
		return 0, fmt.Errorf("failed to fetch metrics from pod %s %s %d: %v", pod.Name, pod.Status.PodIP, metricsPort, err)
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			// Handle the error here. For example, log it or take appropriate corrective action.
			klog.InfoS("Error closing response body:", err)
		}
	}()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return 0, fmt.Errorf("failed to read response from pod %s %s %d: %v", pod.Name, pod.Status.PodIP, metricsPort, err)
	}

	metricValue, err := parseMetricFromBody(body, metricName)
	if err != nil {
		return 0, fmt.Errorf("failed to parse metrics from pod %s %s %d: %v", pod.Name, pod.Status.PodIP, metricsPort, err)
	}

	klog.V(5).InfoS("Successfully parsed metrics", "metric", metricName, "PodIP", pod.Status.PodIP, "Port", metricsPort, "metricValue", metricValue)

Review comment (suggestion):
there's no need to fetch the `/metrics` endpoint for every metrics. Let's just fetch once and follow up metric reading can just parse it. 

Code after:
func (c *Cache) updatePodMetrics() {
	c.mu.Lock()
	defer c.mu.Unlock()

	for _, pod := range c.pods {
		if pod.Status.PodIP == "" {
			continue
		}
		podName := pod.Name
		if len(c.podMetrics[podName]) == 0 {
			c.podMetrics[podName] = map[string]float64{}
		}

		// We should use the primary container port. In future, we can decide whether to use sidecar container's port
		url := fmt.Sprintf("http://%s:%d/metrics", pod.Status.PodIP, podPort)
		resp, err := http.Get(url)
		if err != nil {
			klog.Errorf("failed to fetch metrics from pod %s %s %d: %v", pod.Name, pod.Status.PodIP, podPort, err)
			continue
		}
		defer func() {
			if err := resp.Body.Close(); err != nil {
				klog.Errorf("Error closing response body: %v", err)
			}
		}()

		body, err := io.ReadAll(resp.Body)
		if err != nil {
			klog.Errorf("failed to read response from pod %s %s %d: %v", pod.Name, pod.Status.PodIP, podPort, err)
			continue
		}

		for _, metricName := range metricNames {
			metricValue, err := parseMetricFromBody(body, metricName)
			if err != nil {
				klog.Errorf("failed to parse metrics from pod %s %s %d: %v", pod.Name, pod.Status.PodIP, podPort, err)
				continue
			}

			c.podMetrics[pod.Name][metricName] = metricValue
			klog.V(5).InfoS("Successfully parsed metrics", "metric", metricName, "PodIP", pod.Status.PodIP, "Port", podPort, "metricValue", metricValue)
		}
	}
}

func parseMetricFromBody(body []byte, metricName string) (float64, error) {
	lines := strings.Split(string(body), "\n")
	for _, line := range lines {
		if !strings.HasPrefix(line, "#") && strings.Contains(line, metricName) {
			// format is `http_requests_total 1234.56`
			parts := strings.Fields(line)