feat: merage new

.gitignore

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

cmd/nvidia/main.go

@@ -2,21 +2,62 @@ package main
 
 import (
  "flag"
+ "fmt"
+ "io/ioutil"
+ "time"
 
  "github.com/AliyunContainerService/gpushare-device-plugin/pkg/gpu/nvidia"
+ "github.com/AliyunContainerService/gpushare-device-plugin/pkg/kubelet/client"
  log "github.com/golang/glog"
+ "k8s.io/client-go/rest"
 )
 
 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", false, "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
@@ -108,13 +108,13 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
  if id < 0 {
  return buildErrResponse(reqs, podReqGPU), nil
  }
-
+ log.Infof("gpu index %v,uuid: %v", id, candidateDevID)
  // 1. Create container requests
  for _, req := range reqs.ContainerRequests {
  reqGPU := uint(len(req.DevicesIDs))
  response := pluginapi.ContainerAllocateResponse{
  Envs: map[string]string{
- envNVGPU: candidateDevID,
+ envNVGPU: fmt.Sprintf("%v", id),
  EnvResourceIndex: fmt.Sprintf("%d", id),
  EnvResourceByPod: fmt.Sprintf("%d", podReqGPU),
  EnvResourceByContainer: fmt.Sprintf("%d", reqGPU),
@@ -148,15 +148,47 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
  }
  }
 
+ } else if len(m.devNameMap) == 1 {
+ var devName string
+ var devIndex uint
+ for d, index := range m.devNameMap {
+ devName = d
+ devIndex = index
+ break
+ }
+ log.Infof("this node has only one gpu device,skip to search pod and directly specify the device %v(%v) for container", devIndex, devName)
+ for _, req := range reqs.ContainerRequests {
+ reqGPU := uint(len(req.DevicesIDs))
+ response := pluginapi.ContainerAllocateResponse{
+ Envs: map[string]string{
+ envNVGPU: devName,
+ EnvResourceIndex: fmt.Sprintf("%d", devIndex),
+ EnvResourceByPod: fmt.Sprintf("%d", podReqGPU),
+ EnvResourceByContainer: fmt.Sprintf("%d", reqGPU),
+ EnvResourceByDev: fmt.Sprintf("%d", getGPUMemory()),
+ },
+ }
+ if m.disableCGPUIsolation {
+ response.Envs["CGPU_DISABLE"] = "true"
+ }
+ responses.ContainerResponses = append(responses.ContainerResponses, &response)
+ }
+ log.Infof("get allocated GPUs info %v", responses)
+ return &responses, nil
+
  } else {
  log.Warningf("invalid allocation requst: request GPU memory %d can't be satisfied.",
  podReqGPU)
  // return &responses, fmt.Errorf("invalid allocation requst: request GPU memory %d can't be satisfied", reqGPU)
  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,7 +2,9 @@ package nvidia
 
 import (
  "fmt"
+ "github.com/AliyunContainerService/gpushare-device-plugin/pkg/kubelet/client"
  "syscall"
+ "os"
  "time"
 
  "github.com/NVIDIA/gpu-monitoring-tools/bindings/go/nvml"
@@ -12,15 +14,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,11 +67,13 @@ 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)
+ os.Exit(1)
  } else if err = devicePlugin.Serve(); err != nil {
  log.Warningf("Failed to start device plugin due to %v", err)
+ os.Exit(2)
  } else {
  restart = false
  }