Merge pull request #407 from sebpie/gpu

Gpu

worker/dockerengine.go

@@ -182,6 +182,8 @@ func (dockerengine *DockerEngine) StartTask(task *ScheduledTaskInstance, brokerU
 		commands = append(commands, setOutputCmd)
 	}
 
+	hostConfig := docker.HostConfig{}
+
 	// check if it is required to set up the portmapping for its endpoint services
 	servicePorts := make([]string, 0)
 
@@ -190,6 +192,31 @@ func (dockerengine *DockerEngine) StartTask(task *ScheduledTaskInstance, brokerU
 		if parameter.Name == "service_port" {
 			servicePorts = strings.Split(parameter.Value, ";")
 		}
+
+		var gpu_count = 0
+		var err error
+		if parameter.Name == "gpus" {
+			DEBUG.Println("Enable NVIDIA gpus")
+			hostConfig.Runtime = "nvidia"
+
+			if parameter.Value == "all" {
+				gpu_count = -1
+			} else {
+				count, err = strconv.Atoi(parameter.Value)
+				DEBUG.Println("Error converting parameter to int: ", err)
+			}
+
+			if count != 0 {
+				// Configure GPU resources
+				// DEBUG.Println("Request GPU devices")
+				hostConfig.DeviceRequests = append(hostConfig.DeviceRequests, docker.DeviceRequest{
+					Driver:       "nvidia",
+					Count:        gpu_count, // Allocate all available GPUs
+					Capabilities: [][]string{{"compute", "utility"}},
+				})
+			}
+			// DEBUG.Printf("hostConfig: %v", hostConfig)
+		}
 	}
 
 	// prepare the configuration for a docker container, host mode for the container network
@@ -202,8 +229,6 @@ func (dockerengine *DockerEngine) StartTask(task *ScheduledTaskInstance, brokerU
 
 	config := docker.Config{Image: dockerImage, Env: evs}
 
-	hostConfig := docker.HostConfig{}
-
 	//if runtime.GOOS == "darwin" {   already use the bridge model
 	internalPort := docker.Port(freePort + "/tcp")
 	portBindings := map[docker.Port][]docker.PortBinding{