Optimize monitor (Project-HAMi#683)

zhegemingzimeibanquan · Dec 5, 2024 · 54b8f74 · 54b8f74
1 parent 2066410
commit 54b8f74
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git a/cmd/vGPUmonitor/metrics.go b/cmd/vGPUmonitor/metrics.go
@@ -18,8 +18,10 @@ package main
 
 import (
 	"fmt"
+	"github.com/Project-HAMi/HAMi/pkg/util"
 	"log"
 	"net/http"
+	"os"
 	"strings"
 	"time"
 
@@ -229,8 +231,8 @@ func (cc ClusterManagerCollector) Collect(ch chan<- prometheus.Metric) {
 
 		}
 	}
-
-	pods, err := cc.ClusterManager.PodLister.List(labels.Everything())
+	nodeName := os.Getenv(util.NodeNameEnvName)
+	pods, err := cc.ClusterManager.PodLister.List(labels.SelectorFromSet(labels.Set{util.AssignedNodeAnnotations: nodeName}))
 	if err != nil {
 		klog.Error("failed to list pods with err=", err.Error())
 	}

diff --git a/pkg/monitor/nvidia/cudevshr.go b/pkg/monitor/nvidia/cudevshr.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/Project-HAMi/HAMi/pkg/util"
 	"os"
 	"path/filepath"
 	"strings"
@@ -122,7 +123,13 @@ func (l *ContainerLister) Clientset() *kubernetes.Clientset {
 }
 
 func (l *ContainerLister) Update() error {
-	pods, err := l.clientset.CoreV1().Pods("").List(context.Background(), metav1.ListOptions{})
+	nodename := os.Getenv(util.NodeNameEnvName)
+	if nodename == "" {
+		return fmt.Errorf("env %s not set", util.NodeNameEnvName)
+	}
+	pods, err := l.clientset.CoreV1().Pods("").List(context.Background(), metav1.ListOptions{
+		FieldSelector: fmt.Sprintf("spec.nodeName=%s", nodename),
+	})
 	if err != nil {
 		return err
 	}

diff --git a/pkg/util/util.go b/pkg/util/util.go
@@ -363,6 +363,7 @@ func PatchNodeAnnotations(node *corev1.Node, annotations map[string]string) erro
 func PatchPodAnnotations(pod *corev1.Pod, annotations map[string]string) error {
 	type patchMetadata struct {
 		Annotations map[string]string `json:"annotations,omitempty"`
+		Labels      map[string]string `json:"labels,omitempty"`
 	}
 	type patchPod struct {
 		Metadata patchMetadata `json:"metadata"`
@@ -371,6 +372,11 @@ func PatchPodAnnotations(pod *corev1.Pod, annotations map[string]string) error {
 
 	p := patchPod{}
 	p.Metadata.Annotations = annotations
+	label := make(map[string]string)
+	if v, ok := annotations[AssignedNodeAnnotations]; ok && v != "" {
+		label[AssignedNodeAnnotations] = v
+		p.Metadata.Labels = label
+	}
 
 	bytes, err := json.Marshal(p)
 	if err != nil {