X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/7be9cb0ae8aeb5a49d5450aa38ff9f652761c2d3..f04d5211ed026a4e0cbdca77dad447700eb88772:/lib/crunchrun/crunchrun.go

diff --git a/lib/crunchrun/crunchrun.go b/lib/crunchrun/crunchrun.go
index 589a046a34..7e68dcd331 100644
--- a/lib/crunchrun/crunchrun.go
+++ b/lib/crunchrun/crunchrun.go
@@ -986,6 +986,18 @@ 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,
@@ -995,7 +1007,7 @@ func (runner *ContainerRunner) CreateContainer(imageID string, bindmounts map[st
 		BindMounts:      bindmounts,
 		Command:         runner.Container.Command,
 		EnableNetwork:   enableNetwork,
-		CUDADeviceCount: runner.Container.RuntimeConstraints.CUDADeviceCount,
+		CUDADeviceCount: cudaDeviceCount,
 		NetworkMode:     runner.networkMode,
 		CgroupParent:    runner.setCgroupParent,
 		Stdin:           stdin,