12630: first try at adding GPU support, no tests yet

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

diff --git a/lib/crunchrun/docker.go b/lib/crunchrun/docker.go
index 07f79bbcc2d11f0239a6231288a94d84a89f87fb..573df7faf0b66ac0fa3ac97195a14b8ed5cde0cc 100644 (file)
--- a/lib/crunchrun/docker.go
+++ b/lib/crunchrun/docker.go
@@ -106,6 +106,13 @@ func (e *dockerExecutor) Create(spec containerSpec) error {
                        KernelMemory: spec.RAM, // kernel portion
                },
        }
+       if spec.EnableCUDA {
+               hostCfg.Resources.DeviceRequests = append(hostCfg.Resources.DeviceRequests, dockercontainer.DeviceRequest{
+                       Driver:       "nvidia",
+                       Count:        -1,
+                       Capabilities: [][]string{[]string{"gpu", "nvidia", "compute"}},
+               })
+       }
        for path, mount := range spec.BindMounts {
                bind := mount.HostPath + ":" + path
                if mount.ReadOnly {