X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/3134405ebc155a8a51738b7c6d0d4be348c65087..2f344e8b8dde661e74307ed7e561a758809382e1:/lib/crunchrun/crunchrun.go

diff --git a/lib/crunchrun/crunchrun.go b/lib/crunchrun/crunchrun.go
index 33558b5d9b..7e68dcd331 100644
--- a/lib/crunchrun/crunchrun.go
+++ b/lib/crunchrun/crunchrun.go
@@ -986,20 +986,33 @@ func (runner *ContainerRunner) CreateContainer(imageID string, bindmounts map[st
 	runner.executorStdin = stdin
 	runner.executorStdout = stdout
 	runner.executorStderr = stderr
+
+	cudaDeviceCount := 0
+	if runner.Container.RuntimeConstraints.CUDADriverVersion != "" ||
+		runner.Container.RuntimeConstraints.CUDAHardwareCapability != "" ||
+		runner.Container.RuntimeConstraints.CUDADeviceCount != 0 {
+		// if any of these are set, enable CUDA GPU support
+		cudaDeviceCount = runner.Container.RuntimeConstraints.CUDADeviceCount
+		if cudaDeviceCount == 0 {
+			cudaDeviceCount = 1
+		}
+	}
+
 	return runner.executor.Create(containerSpec{
-		Image:         imageID,
-		VCPUs:         runner.Container.RuntimeConstraints.VCPUs,
-		RAM:           ram,
-		WorkingDir:    workdir,
-		Env:           env,
-		BindMounts:    bindmounts,
-		Command:       runner.Container.Command,
-		EnableNetwork: enableNetwork,
-		NetworkMode:   runner.networkMode,
-		CgroupParent:  runner.setCgroupParent,
-		Stdin:         stdin,
-		Stdout:        stdout,
-		Stderr:        stderr,
+		Image:           imageID,
+		VCPUs:           runner.Container.RuntimeConstraints.VCPUs,
+		RAM:             ram,
+		WorkingDir:      workdir,
+		Env:             env,
+		BindMounts:      bindmounts,
+		Command:         runner.Container.Command,
+		EnableNetwork:   enableNetwork,
+		CUDADeviceCount: cudaDeviceCount,
+		NetworkMode:     runner.networkMode,
+		CgroupParent:    runner.setCgroupParent,
+		Stdin:           stdin,
+		Stdout:          stdout,
+		Stderr:          stderr,
 	})
 }
 
@@ -1699,11 +1712,11 @@ func (command) RunCommand(prog string, args []string, stdin io.Reader, stdout, s
 		ignoreDetachFlag = true
 	}
 
-	if err := flags.Parse(args); err == flag.ErrHelp {
-		return 0
-	} else if err != nil {
-		log.Print(err)
-		return 1
+	if ok, code := cmd.ParseFlags(flags, prog, args, "container-uuid", stderr); !ok {
+		return code
+	} else if !*list && flags.NArg() != 1 {
+		fmt.Fprintf(stderr, "missing required argument: container-uuid (try -help)\n")
+		return 2
 	}
 
 	containerUUID := flags.Arg(0)