fix #39: get pods from kubelet client rather than list cluster-scope pods from apiserver (#41)

* fix: add .gitignore

* fix: get pods from kubelet client rather than list cluster-scope pods from apiserver

* feat: add query-kubelet flag for control the mothod of query pending podLists

Co-authored-by: menglong.qi <[email protected]>

Author: qmloong

.gitignore

@@ -0,0 +1 @@
+.idea

cmd/nvidia/main.go

@@ -2,21 +2,61 @@ package main
 
 import (
 	"flag"
-
+	"fmt"
 	"github.com/AliyunContainerService/gpushare-device-plugin/pkg/gpu/nvidia"
+	"github.com/AliyunContainerService/gpushare-device-plugin/pkg/kubelet/client"
 	log "github.com/golang/glog"
+	"io/ioutil"
+	"k8s.io/client-go/rest"
+	"time"
 )
 
 var (
-	mps         = flag.Bool("mps", false, "Enable or Disable MPS")
-	healthCheck = flag.Bool("health-check", false, "Enable or disable Health check")
-	memoryUnit  = flag.String("memory-unit", "GiB", "Set memoryUnit of the GPU Memroy, support 'GiB' and 'MiB'")
+	mps              = flag.Bool("mps", false, "Enable or Disable MPS")
+	healthCheck      = flag.Bool("health-check", false, "Enable or disable Health check")
+	memoryUnit       = flag.String("memory-unit", "GiB", "Set memoryUnit of the GPU Memroy, support 'GiB' and 'MiB'")
+	queryFromKubelet = flag.Bool("query-kubelet", true, "Query pending pods from kubelet instead of kube-apiserver")
+	kubeletAddress   = flag.String("kubelet-address", "0.0.0.0", "Kubelet IP Address")
+	kubeletPort      = flag.Uint("kubelet-port", 10250, "Kubelet listened Port")
+	clientCert       = flag.String("client-cert", "", "Kubelet TLS client certificate")
+	clientKey        = flag.String("client-key", "", "Kubelet TLS client key")
+	token            = flag.String("token", "", "Kubelet client bearer token")
+	timeout          = flag.Int("timeout", 10, "Kubelet client http timeout duration")
 )
 
+func buildKubeletClient() *client.KubeletClient {
+	if *clientCert == "" && *clientKey == "" && *token == "" {
+		tokenByte, err := ioutil.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/token")
+		if err != nil {
+			panic(fmt.Errorf("in cluster mode, find token failed, error: %v", err))
+		}
+		tokenStr := string(tokenByte)
+		token = &tokenStr
+	}
+	kubeletClient, err := client.NewKubeletClient(&client.KubeletClientConfig{
+		Address: *kubeletAddress,
+		Port:    *kubeletPort,
+		TLSClientConfig: rest.TLSClientConfig{
+			Insecure:   true,
+			ServerName: "gpushare-device-plugin",
+			CertFile:   *clientCert,
+			KeyFile:    *clientKey,
+		},
+		BearerToken: *token,
+		HTTPTimeout: time.Duration(*timeout) * time.Second,
+	})
+	if err != nil {
+		panic(err)
+	}
+	return kubeletClient
+}
+
 func main() {
 	flag.Parse()
 	log.V(1).Infoln("Start gpushare device plugin")
-	ngm := nvidia.NewSharedGPUManager(*mps, *healthCheck, translatememoryUnits(*memoryUnit))
+
+	kubeletClient := buildKubeletClient()
+	ngm := nvidia.NewSharedGPUManager(*mps, *healthCheck, *queryFromKubelet, translatememoryUnits(*memoryUnit), kubeletClient)
 	err := ngm.Run()
 	if err != nil {
 		log.Fatalf("Failed due to %v", err)

cmd/podgetter/main.go

@@ -0,0 +1,57 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"github.com/AliyunContainerService/gpushare-device-plugin/pkg/kubelet/client"
+	"io/ioutil"
+	"k8s.io/client-go/rest"
+	"time"
+)
+
+var (
+	clientCert string
+	clientKey  string
+	token      string
+	timeout    int
+)
+
+func main() {
+	flag.StringVar(&clientCert, "client-cert", "", "")
+	flag.StringVar(&clientKey, "client-key", "", "")
+	flag.StringVar(&token, "token", "", "")
+	flag.IntVar(&timeout, "timeout", 10, "")
+
+	flag.Parse()
+
+	if clientCert == "" && clientKey == "" && token == "" {
+		tokenByte, err := ioutil.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/token")
+		if err != nil {
+			panic(fmt.Errorf("in cluster mode, find token failed, error: %v", err))
+		}
+		token = string(tokenByte)
+	}
+
+	c, err := client.NewKubeletClient(&client.KubeletClientConfig{
+		Address: "127.0.0.1",
+		Port:    10250,
+		TLSClientConfig: rest.TLSClientConfig{
+			Insecure:   true,
+			ServerName: "kubelet",
+			CertFile:   clientCert,
+			KeyFile:    clientKey,
+		},
+		BearerToken: token,
+		HTTPTimeout: time.Duration(timeout) * time.Second,
+	})
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+	podsList, err := c.GetNodeRunningPods()
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+	fmt.Println(podsList)
+}

device-plugin-rbac.yaml

@@ -9,6 +9,7 @@ rules:
   - ""
   resources:
   - nodes
+  - nodes/proxy
   verbs:
   - get
   - list

pkg/gpu/nvidia/allocate.go

@@ -59,7 +59,7 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
 	m.Lock()
 	defer m.Unlock()
 	log.Infoln("checking...")
-	pods, err := getCandidatePods()
+	pods, err := getCandidatePods(m.queryKubelet, m.kubeletClient)
 	if err != nil {
 		log.Infof("invalid allocation requst: Failed to find candidate pods due to %v", err)
 		return buildErrResponse(reqs, podReqGPU), nil
@@ -155,8 +155,12 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
 		return buildErrResponse(reqs, podReqGPU), nil
 	}
 
-	log.Infof("new allocated GPUs info %v", &responses)
-	log.Infoln("----Allocating GPU for gpu mem is ended----")
+	podName := ""
+	if assumePod != nil {
+		podName = assumePod.Name
+	}
+	log.Infof("pod %v, new allocated GPUs info %v", podName, &responses)
+	log.Infof("----Allocating GPU for gpu mem for %v is ended----", podName)
 	// // Add this to make sure the container is created at least
 	// currentTime := time.Now()
 

pkg/gpu/nvidia/gpumanager.go

@@ -2,6 +2,7 @@ package nvidia
 
 import (
 	"fmt"
+	"github.com/AliyunContainerService/gpushare-device-plugin/pkg/kubelet/client"
 	"syscall"
 	"time"
 
@@ -12,15 +13,19 @@ import (
 )
 
 type sharedGPUManager struct {
-	enableMPS   bool
-	healthCheck bool
+	enableMPS     bool
+	healthCheck   bool
+	queryKubelet  bool
+	kubeletClient *client.KubeletClient
 }
 
-func NewSharedGPUManager(enableMPS, healthCheck bool, bp MemoryUnit) *sharedGPUManager {
+func NewSharedGPUManager(enableMPS, healthCheck, queryKubelet bool, bp MemoryUnit, client *client.KubeletClient) *sharedGPUManager {
 	metric = bp
 	return &sharedGPUManager{
-		enableMPS:   enableMPS,
-		healthCheck: healthCheck,
+		enableMPS:     enableMPS,
+		healthCheck:   healthCheck,
+		queryKubelet:  queryKubelet,
+		kubeletClient: client,
 	}
 }
 
@@ -61,7 +66,7 @@ L:
 				devicePlugin.Stop()
 			}
 
-			devicePlugin, err = NewNvidiaDevicePlugin(ngm.enableMPS, ngm.healthCheck)
+			devicePlugin, err = NewNvidiaDevicePlugin(ngm.enableMPS, ngm.healthCheck, ngm.queryKubelet, ngm.kubeletClient)
 			if err != nil {
 				log.Warningf("Failed to get device plugin due to %v", err)
 			} else if err = devicePlugin.Serve(); err != nil {

pkg/gpu/nvidia/podmanager.go

@@ -1,30 +1,29 @@
 package nvidia
 
 import (
+	"encoding/json"
 	"fmt"
-	"os"
-	"sort"
-	"time"
-
+	"github.com/AliyunContainerService/gpushare-device-plugin/pkg/kubelet/client"
 	log "github.com/golang/glog"
-	"k8s.io/apimachinery/pkg/labels"
-
 	"k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/fields"
-
+	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
 	nodeutil "k8s.io/kubernetes/pkg/util/node"
+	"os"
+	"sort"
+	"time"
 )
 
 var (
 	clientset *kubernetes.Clientset
 	nodeName  string
-	retries   = 5
+	retries   = 8
 )
 
 func kubeInit() {
@@ -58,18 +57,18 @@ func kubeInit() {
 }
 
 func disableCGPUIsolationOrNot() (bool, error) {
-    disable := false
-    node, err := clientset.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
-    if err != nil {
-        return disable, err
-    }
-    labels := node.ObjectMeta.Labels
-    value, ok := labels[EnvNodeLabelForDisableCGPU]
-    if ok && value == "true" {
-        log.Infof("enable gpusharing mode and disable cgpu mode")
-        disable = true
-    }
-    return disable, nil
+	disable := false
+	node, err := clientset.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
+	if err != nil {
+		return disable, err
+	}
+	labels := node.ObjectMeta.Labels
+	value, ok := labels[EnvNodeLabelForDisableCGPU]
+	if ok && value == "true" {
+		log.Infof("enable gpusharing mode and disable cgpu mode")
+		disable = true
+	}
+	return disable, nil
 }
 
 func patchGPUCount(gpuCount int) error {
@@ -99,31 +98,90 @@ func patchGPUCount(gpuCount int) error {
 	return err
 }
 
-func getPendingPodsInNode() ([]v1.Pod, error) {
-	// pods, err := m.lister.List(labels.Everything())
-	// if err != nil {
-	// 	return nil, err
-	// }
-	pods := []v1.Pod{}
+func getPodList(kubeletClient *client.KubeletClient) (*v1.PodList, error) {
+	podList, err := kubeletClient.GetNodeRunningPods()
+	if err != nil {
+		return nil, err
+	}
 
-	podIDMap := map[types.UID]bool{}
+	list, _ := json.Marshal(podList)
+	log.V(8).Infof("get pods list %v", string(list))
+
+	resultPodList := &v1.PodList{}
+	for _, metaPod := range podList.Items {
+		if metaPod.Status.Phase != v1.PodPending {
+			continue
+		}
+		resultPodList.Items = append(resultPodList.Items, metaPod)
+	}
+
+	if len(resultPodList.Items) == 0 {
+		return nil, fmt.Errorf("not found pending pod")
+	}
+
+	return resultPodList, nil
+}
 
+func getPodListsByQueryKubelet(kubeletClient *client.KubeletClient) (*v1.PodList, error) {
+	podList, err := getPodList(kubeletClient)
+	for i := 0; i < retries && err != nil; i++ {
+		podList, err = getPodList(kubeletClient)
+		log.Warningf("failed to get pending pod list, retry")
+		time.Sleep(100 * time.Millisecond)
+	}
+	if err != nil {
+		log.Warningf("not found from kubelet /pods api, start to list apiserver")
+		podList, err = getPodListsByListAPIServer()
+		if err != nil {
+			return nil, err
+		}
+	}
+	return podList, nil
+}
+
+func getPodListsByListAPIServer() (*v1.PodList, error) {
 	selector := fields.SelectorFromSet(fields.Set{"spec.nodeName": nodeName, "status.phase": "Pending"})
 	podList, err := clientset.CoreV1().Pods(v1.NamespaceAll).List(metav1.ListOptions{
 		FieldSelector: selector.String(),
 		LabelSelector: labels.Everything().String(),
 	})
-	for i := 0; i < retries && err != nil; i++ {
+	for i := 0; i < 3 && err != nil; i++ {
 		podList, err = clientset.CoreV1().Pods(v1.NamespaceAll).List(metav1.ListOptions{
 			FieldSelector: selector.String(),
 			LabelSelector: labels.Everything().String(),
 		})
-		time.Sleep(100 * time.Millisecond)
+		time.Sleep(1 * time.Second)
 	}
 	if err != nil {
 		return nil, fmt.Errorf("failed to get Pods assigned to node %v", nodeName)
 	}
 
+	return podList, nil
+}
+
+func getPendingPodsInNode(queryKubelet bool, kubeletClient *client.KubeletClient) ([]v1.Pod, error) {
+	// pods, err := m.lister.List(labels.Everything())
+	// if err != nil {
+	// 	return nil, err
+	// }
+	pods := []v1.Pod{}
+
+	podIDMap := map[types.UID]bool{}
+
+	var podList *v1.PodList
+	var err error
+	if queryKubelet {
+		podList, err = getPodListsByQueryKubelet(kubeletClient)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		podList, err = getPodListsByListAPIServer()
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	log.V(5).Infof("all pod list %v", podList.Items)
 
 	// if log.V(5) {
@@ -154,9 +212,9 @@ func getPendingPodsInNode() ([]v1.Pod, error) {
 }
 
 // pick up the gpushare pod with assigned status is false, and
-func getCandidatePods() ([]*v1.Pod, error) {
+func getCandidatePods(queryKubelet bool, client *client.KubeletClient) ([]*v1.Pod, error) {
 	candidatePods := []*v1.Pod{}
-	allPods, err := getPendingPodsInNode()
+	allPods, err := getPendingPodsInNode(queryKubelet, client)
 	if err != nil {
 		return candidatePods, err
 	}