12630: Remove distinction between cubin/PTX hardware capabilities

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

diff --git a/lib/crunchrun/docker.go b/lib/crunchrun/docker.go
index 861f8c8c1913f07bab8d7ea722dfa3c643678059..c64ff83da38ada46a19fc15153f11a42270e28bf 100644 (file)
--- a/lib/crunchrun/docker.go
+++ b/lib/crunchrun/docker.go
@@ -11,6 +11,7 @@ import (
        "strings"
        "time"
 
+       "git.arvados.org/arvados.git/sdk/go/arvados"
        dockertypes "github.com/docker/docker/api/types"
        dockercontainer "github.com/docker/docker/api/types/container"
        dockerclient "github.com/docker/docker/client"
@@ -45,13 +46,17 @@ func newDockerExecutor(containerUUID string, logf func(string, ...interface{}),
        }, err
 }
 
-func (e *dockerExecutor) ImageLoaded(imageID string) bool {
+func (e *dockerExecutor) Runtime() string { return "docker" }
+
+func (e *dockerExecutor) LoadImage(imageID string, imageTarballPath string, container arvados.Container, arvMountPoint string,
+       containerClient *arvados.Client) error {
        _, _, err := e.dockerclient.ImageInspectWithRaw(context.TODO(), imageID)
-       return err == nil
-}
+       if err == nil {
+               // already loaded
+               return nil
+       }
 
-func (e *dockerExecutor) LoadImage(filename string) error {
-       f, err := os.Open(filename)
+       f, err := os.Open(imageTarballPath)
        if err != nil {
                return err
        }
@@ -101,6 +106,13 @@ func (e *dockerExecutor) Create(spec containerSpec) error {
                        KernelMemory: spec.RAM, // kernel portion
                },
        }
+       if spec.CUDADeviceCount != 0 {
+               hostCfg.Resources.DeviceRequests = append(hostCfg.Resources.DeviceRequests, dockercontainer.DeviceRequest{
+                       Driver:       "nvidia",
+                       Count:        spec.CUDADeviceCount,
+                       Capabilities: [][]string{[]string{"gpu", "nvidia", "compute"}},
+               })
+       }
        for path, mount := range spec.BindMounts {
                bind := mount.HostPath + ":" + path
                if mount.ReadOnly {