skip to search pod when count of gpu devices is only one (#46)

Author: happy2048

cmd/nvidia/main.go

@@ -3,19 +3,20 @@ 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"
-	"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'")
-	queryFromKubelet = flag.Bool("query-kubelet", true, "Query pending pods from kubelet instead of kube-apiserver")
+	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")

pkg/gpu/nvidia/allocate.go

@@ -148,6 +148,34 @@ 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)