Add CUDA_MPS_ACTIVE_THREAD_PERCENTAGE information in MPS situation.

Signed-off-by: 630445639 <630445639@qq.com>

Author: Sakuralbj

pkg/gpu/nvidia/allocate.go

@@ -21,18 +21,21 @@ func init() {
 	kubeInit()
 }
 
-func buildErrResponse(reqs *pluginapi.AllocateRequest, podReqGPU uint) *pluginapi.AllocateResponse {
+func (m *NvidiaDevicePlugin) buildErrResponse(reqs *pluginapi.AllocateRequest, podReqGPU uint) *pluginapi.AllocateResponse {
 	responses := pluginapi.AllocateResponse{}
 	for _, req := range reqs.ContainerRequests {
 		response := pluginapi.ContainerAllocateResponse{
 			Envs: map[string]string{
-				envNVGPU:               fmt.Sprintf("no-gpu-has-%dMiB-to-run", podReqGPU),
+				envNVGPU:               fmt.Sprintf("no-gpu-has-%dGiB-to-run", podReqGPU),
 				EnvResourceIndex:       fmt.Sprintf("-1"),
 				EnvResourceByPod:       fmt.Sprintf("%d", podReqGPU),
 				EnvResourceByContainer: fmt.Sprintf("%d", uint(len(req.DevicesIDs))),
 				EnvResourceByDev:       fmt.Sprintf("%d", getGPUMemory()),
 			},
 		}
+		if m.mps {
+			response.Envs[EnvPercentage] = fmt.Sprintf("%d", 100*uint(len(req.DevicesIDs))/getGPUMemory()) //
+		}
 		responses.ContainerResponses = append(responses.ContainerResponses, &response)
 	}
 	return &responses
@@ -62,7 +65,7 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
 	pods, err := getCandidatePods()
 	if err != nil {
 		log.Infof("invalid allocation requst: Failed to find candidate pods due to %v", err)
-		return buildErrResponse(reqs, podReqGPU), nil
+		return m.buildErrResponse(reqs, podReqGPU), nil
 	}
 
 	if log.V(4) {
@@ -106,7 +109,7 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
 		}
 
 		if id < 0 {
-			return buildErrResponse(reqs, podReqGPU), nil
+			return m.buildErrResponse(reqs, podReqGPU), nil
 		}
 
 		// 1. Create container requests
@@ -121,6 +124,9 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
 					EnvResourceByDev:       fmt.Sprintf("%d", getGPUMemory()),
 				},
 			}
+			if m.mps {
+				response.Envs[EnvPercentage] = fmt.Sprintf("%d", 100*reqGPU/getGPUMemory())
+			}
 			responses.ContainerResponses = append(responses.ContainerResponses, &response)
 		}
 
@@ -134,25 +140,25 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
 				pod, err := clientset.CoreV1().Pods(assumePod.Namespace).Get(assumePod.Name, metav1.GetOptions{})
 				if err != nil {
 					log.Warningf("Failed due to %v", err)
-					return buildErrResponse(reqs, podReqGPU), nil
+					return m.buildErrResponse(reqs, podReqGPU), nil
 				}
 				newPod = updatePodAnnotations(pod)
 				_, err = clientset.CoreV1().Pods(newPod.Namespace).Update(newPod)
 				if err != nil {
 					log.Warningf("Failed due to %v", err)
-					return buildErrResponse(reqs, podReqGPU), nil
+					return m.buildErrResponse(reqs, podReqGPU), nil
 				}
 			} else {
 				log.Warningf("Failed due to %v", err)
-				return buildErrResponse(reqs, podReqGPU), nil
+				return m.buildErrResponse(reqs, podReqGPU), 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
+		return m.buildErrResponse(reqs, podReqGPU), nil
 	}
 
 	log.Infof("new allocated GPUs info %v", &responses)

pkg/gpu/nvidia/const.go

@@ -28,6 +28,7 @@ const (
 	EnvResourceByDev       = "ALIYUN_COM_GPU_MEM_DEV"
 	EnvAssignedFlag        = "ALIYUN_COM_GPU_MEM_ASSIGNED"
 	EnvResourceAssumeTime  = "ALIYUN_COM_GPU_MEM_ASSUME_TIME"
+	EnvPercentage          = "CUDA_MPS_ACTIVE_THREAD_PERCENTAGE"
 	EnvResourceAssignTime  = "ALIYUN_COM_GPU_MEM_ASSIGN_TIME"
 
 	GiBPrefix = MemoryUnit("GiB")