Request GPU without deprecated --runtime=nvidia (#913)

(cherry picked from commit 730e8b0)

Author: Egor-S

runner/internal/shim/docker.go

@@ -82,11 +82,10 @@ func pullImage(ctx context.Context, client docker.APIClient, imageName string, r
 }
 
 func createContainer(ctx context.Context, client docker.APIClient, params DockerParameters) (string, error) {
-	runtime, err := getRuntime(ctx, client)
+	gpuRequest, err := requestGpuIfAvailable(ctx, client)
 	if err != nil {
 		return "", gerrors.Wrap(err)
 	}
-
 	mounts, err := params.DockerMounts()
 	if err != nil {
 		return "", gerrors.Wrap(err)
@@ -103,8 +102,10 @@ func createContainer(ctx context.Context, client docker.APIClient, params Docker
 		PortBindings:    bindPorts(params.DockerPorts()...),
 		PublishAllPorts: true,
 		Sysctls:         map[string]string{},
-		Runtime:         runtime,
-		Mounts:          mounts,
+		Resources: container.Resources{
+			DeviceRequests: gpuRequest,
+		},
+		Mounts: mounts,
 	}
 	resp, err := client.ContainerCreate(ctx, containerConfig, hostConfig, nil, nil, "")
 	if err != nil {
@@ -178,17 +179,21 @@ func getNetworkMode() container.NetworkMode {
 	return "default"
 }
 
-func getRuntime(ctx context.Context, client docker.APIClient) (string, error) {
+func requestGpuIfAvailable(ctx context.Context, client docker.APIClient) ([]container.DeviceRequest, error) {
 	info, err := client.Info(ctx)
 	if err != nil {
-		return "", gerrors.Wrap(err)
+		return nil, gerrors.Wrap(err)
 	}
-	for name := range info.Runtimes {
-		if name == consts.NVIDIA_RUNTIME {
-			return name, nil
+
+	for runtime := range info.Runtimes {
+		if runtime == consts.NVIDIA_RUNTIME {
+			return []container.DeviceRequest{
+				{Capabilities: [][]string{{"gpu"}}, Count: -1}, // --gpus=all
+			}, nil
 		}
 	}
-	return info.DefaultRuntime, nil
+
+	return nil, nil
 }
 
 /* DockerParameters interface implementation for CLIArgs */