detect node label to disable cgpu (#32)

Author: happy2048

pkg/gpu/nvidia/allocate.go

@@ -121,6 +121,9 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
 					EnvResourceByDev:       fmt.Sprintf("%d", getGPUMemory()),
 				},
 			}
+            if m.disableCGPUIsolation {
+                response.Envs["CGPU_DISABLE"] = "true"
+            }
 			responses.ContainerResponses = append(responses.ContainerResponses, &response)
 		}
 

pkg/gpu/nvidia/const.go

@@ -21,14 +21,15 @@ const (
 	containerLogPathLabelKey    = "io.kubernetes.container.logpath"
 	sandboxIDLabelKey           = "io.kubernetes.sandbox.id"
 
-	envNVGPU               = "NVIDIA_VISIBLE_DEVICES"
-	EnvResourceIndex       = "ALIYUN_COM_GPU_MEM_IDX"
-	EnvResourceByPod       = "ALIYUN_COM_GPU_MEM_POD"
-	EnvResourceByContainer = "ALIYUN_COM_GPU_MEM_CONTAINER"
-	EnvResourceByDev       = "ALIYUN_COM_GPU_MEM_DEV"
-	EnvAssignedFlag        = "ALIYUN_COM_GPU_MEM_ASSIGNED"
-	EnvResourceAssumeTime  = "ALIYUN_COM_GPU_MEM_ASSUME_TIME"
-	EnvResourceAssignTime  = "ALIYUN_COM_GPU_MEM_ASSIGN_TIME"
+	envNVGPU                   = "NVIDIA_VISIBLE_DEVICES"
+	EnvResourceIndex           = "ALIYUN_COM_GPU_MEM_IDX"
+	EnvResourceByPod           = "ALIYUN_COM_GPU_MEM_POD"
+	EnvResourceByContainer     = "ALIYUN_COM_GPU_MEM_CONTAINER"
+	EnvResourceByDev           = "ALIYUN_COM_GPU_MEM_DEV"
+	EnvAssignedFlag            = "ALIYUN_COM_GPU_MEM_ASSIGNED"
+	EnvResourceAssumeTime      = "ALIYUN_COM_GPU_MEM_ASSUME_TIME"
+	EnvResourceAssignTime      = "ALIYUN_COM_GPU_MEM_ASSIGN_TIME"
+	EnvNodeLabelForDisableCGPU = "cgpu.disable.isolation"
 
 	GiBPrefix = MemoryUnit("GiB")
 	MiBPrefix = MemoryUnit("MiB")

pkg/gpu/nvidia/podmanager.go

@@ -57,6 +57,21 @@ 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
+}
+
 func patchGPUCount(gpuCount int) error {
 	node, err := clientset.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
 	if err != nil {

pkg/gpu/nvidia/server.go

@@ -16,16 +16,16 @@ import (
 
 // NvidiaDevicePlugin implements the Kubernetes device plugin API
 type NvidiaDevicePlugin struct {
-	devs         []*pluginapi.Device
-	realDevNames []string
-	devNameMap   map[string]uint
-	devIndxMap   map[uint]string
-	socket       string
-	mps          bool
-	healthCheck  bool
-
-	stop   chan struct{}
-	health chan *pluginapi.Device
+	devs                 []*pluginapi.Device
+	realDevNames         []string
+	devNameMap           map[string]uint
+	devIndxMap           map[uint]string
+	socket               string
+	mps                  bool
+	healthCheck          bool
+	disableCGPUIsolation bool
+	stop                 chan struct{}
+	health               chan *pluginapi.Device
 
 	server *grpc.Server
 	sync.RWMutex
@@ -47,17 +47,20 @@ func NewNvidiaDevicePlugin(mps, healthCheck bool) (*NvidiaDevicePlugin, error) {
 	if err != nil {
 		return nil, err
 	}
-
+	disableCGPUIsolation, err := disableCGPUIsolationOrNot()
+	if err != nil {
+		return nil, err
+	}
 	return &NvidiaDevicePlugin{
-		devs:         devs,
-		realDevNames: devList,
-		devNameMap:   devNameMap,
-		socket:       serverSock,
-		mps:          mps,
-		healthCheck:  healthCheck,
-
-		stop:   make(chan struct{}),
-		health: make(chan *pluginapi.Device),
+		devs:                 devs,
+		realDevNames:         devList,
+		devNameMap:           devNameMap,
+		socket:               serverSock,
+		mps:                  mps,
+		healthCheck:          healthCheck,
+		disableCGPUIsolation: disableCGPUIsolation,
+		stop:                 make(chan struct{}),
+		health:               make(chan *pluginapi.Device),
 	}, nil
 }