12630: Use CUDADeviceCount instead of EnableCUDA

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

diff --git a/lib/crunchrun/docker.go b/lib/crunchrun/docker.go
index 573df7faf0b66ac0fa3ac97195a14b8ed5cde0cc..c64ff83da38ada46a19fc15153f11a42270e28bf 100644 (file)
--- a/lib/crunchrun/docker.go
+++ b/lib/crunchrun/docker.go
@@ -106,10 +106,10 @@ func (e *dockerExecutor) Create(spec containerSpec) error {
                        KernelMemory: spec.RAM, // kernel portion
                },
        }
-       if spec.EnableCUDA {
+       if spec.CUDADeviceCount != 0 {
                hostCfg.Resources.DeviceRequests = append(hostCfg.Resources.DeviceRequests, dockercontainer.DeviceRequest{
                        Driver:       "nvidia",
-                       Count:        -1,
+                       Count:        spec.CUDADeviceCount,
                        Capabilities: [][]string{[]string{"gpu", "nvidia", "compute"}},
                })
        }