12630: Remove distinction between cubin/PTX hardware capabilities

[arvados.git] / lib / crunchrun / crunchrun.go

diff --git a/lib/crunchrun/crunchrun.go b/lib/crunchrun/crunchrun.go
index 589a046a34c0edb94ed97c48b7776cd63c1091cd..7e68dcd3314685bf4132739bb17ff0ad2e9ba7b8 100644 (file)
--- 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,