Add volume information

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

Author: Sakuralbj

cmd/nvidia/main.go

@@ -11,12 +11,13 @@ 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'")
+	mpspipe     = flag.String("mps-pipe", "/tmp/nvidia-mps", " pipes and UNIX domain sockets")
 )
 
 func main() {
 	flag.Parse()
 	log.V(1).Infoln("Start gpushare device plugin")
-	ngm := nvidia.NewSharedGPUManager(*mps, *healthCheck, translatememoryUnits(*memoryUnit))
+	ngm := nvidia.NewSharedGPUManager(*mps, *healthCheck, *mpspipe, translatememoryUnits(*memoryUnit))
 	err := ngm.Run()
 	if err != nil {
 		log.Fatalf("Failed due to %v", err)

device-plugin-ds.yaml

@@ -25,6 +25,8 @@ spec:
           - gpushare-device-plugin-v2
           - -logtostderr
           - --v=5
+        # - --mps-pipe=/root/nvidia-mps  // mps-client and mps-server communicate through the directory.You can modify it.
+        # - --mps=true     //if you want to use mps
           - --memory-unit=GiB
         resources:
           limits:

pkg/gpu/nvidia/allocate.go

@@ -34,7 +34,8 @@ func (m *NvidiaDevicePlugin) buildErrResponse(reqs *pluginapi.AllocateRequest, p
 			},
 		}
 		if m.mps {
-			response.Envs[EnvPercentage] = fmt.Sprintf("%d", 100*uint(len(req.DevicesIDs))/getGPUMemory()) //
+			response.Envs[EnvMPSActiveThreadPercentage] = fmt.Sprintf("%d", 100*uint(len(req.DevicesIDs))/getGPUMemory())
+			response.Envs[EnvMPSPipeDirectory] = fmt.Sprintf(m.mpspipe)
 		}
 		responses.ContainerResponses = append(responses.ContainerResponses, &response)
 	}
@@ -125,7 +126,13 @@ func (m *NvidiaDevicePlugin) Allocate(ctx context.Context,
 				},
 			}
 			if m.mps {
-				response.Envs[EnvPercentage] = fmt.Sprintf("%d", 100*reqGPU/getGPUMemory())
+				response.Envs[EnvMPSActiveThreadPercentage] = fmt.Sprintf("%d", 100*reqGPU/getGPUMemory())
+				response.Envs[EnvMPSPipeDirectory] = fmt.Sprintf(m.mpspipe)
+				mount := pluginapi.Mount{
+					ContainerPath: m.mpspipe,
+					HostPath:      m.mpspipe,
+				}
+				response.Mounts = append(response.Mounts, &mount)
 			}
 			responses.ContainerResponses = append(responses.ContainerResponses, &response)
 		}

pkg/gpu/nvidia/const.go

@@ -21,15 +21,16 @@ 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"
-	EnvPercentage          = "CUDA_MPS_ACTIVE_THREAD_PERCENTAGE"
-	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"
+	EnvMPSPipeDirectory          = "CUDA_MPS_PIPE_DIRECTORY"
+	EnvMPSActiveThreadPercentage = "CUDA_MPS_ACTIVE_THREAD_PERCENTAGE"
+	EnvResourceAssignTime        = "ALIYUN_COM_GPU_MEM_ASSIGN_TIME"
 
 	GiBPrefix = MemoryUnit("GiB")
 	MiBPrefix = MemoryUnit("MiB")

pkg/gpu/nvidia/gpumanager.go

@@ -14,13 +14,15 @@ import (
 type sharedGPUManager struct {
 	enableMPS   bool
 	healthCheck bool
+	mpspipe     string
 }
 
-func NewSharedGPUManager(enableMPS, healthCheck bool, bp MemoryUnit) *sharedGPUManager {
+func NewSharedGPUManager(enableMPS, healthCheck bool, mpspipe string, bp MemoryUnit) *sharedGPUManager {
 	metric = bp
 	return &sharedGPUManager{
 		enableMPS:   enableMPS,
 		healthCheck: healthCheck,
+		mpspipe:     mpspipe,
 	}
 }
 
@@ -61,7 +63,7 @@ L:
 				devicePlugin.Stop()
 			}
 
-			devicePlugin = NewNvidiaDevicePlugin(ngm.enableMPS, ngm.healthCheck)
+			devicePlugin = NewNvidiaDevicePlugin(ngm.enableMPS, ngm.healthCheck, ngm.mpspipe)
 			if err := devicePlugin.Serve(); err != nil {
 				log.Warningf("Failed to start device plugin due to %v", err)
 			} else {

pkg/gpu/nvidia/server.go

@@ -22,6 +22,7 @@ type NvidiaDevicePlugin struct {
 	devIndxMap   map[uint]string
 	socket       string
 	mps          bool
+	mpspipe      string
 	healthCheck  bool
 
 	stop   chan struct{}
@@ -32,7 +33,7 @@ type NvidiaDevicePlugin struct {
 }
 
 // NewNvidiaDevicePlugin returns an initialized NvidiaDevicePlugin
-func NewNvidiaDevicePlugin(mps, healthCheck bool) *NvidiaDevicePlugin {
+func NewNvidiaDevicePlugin(mps, healthCheck bool, mpspipe string) *NvidiaDevicePlugin {
 	devs, devNameMap := getDevices()
 	devList := []string{}
 
@@ -54,10 +55,10 @@ func NewNvidiaDevicePlugin(mps, healthCheck bool) *NvidiaDevicePlugin {
 		devNameMap:   devNameMap,
 		socket:       serverSock,
 		mps:          mps,
+		mpspipe:      mpspipe,
 		healthCheck:  healthCheck,
-
-		stop:   make(chan struct{}),
-		health: make(chan *pluginapi.Device),
+		stop:         make(chan struct{}),
+		health:       make(chan *pluginapi.Device),
 	}
 }