diff --git a/devices/gpu/nvidia/device.go b/devices/gpu/nvidia/device.go index 110287681..b288181c0 100644 --- a/devices/gpu/nvidia/device.go +++ b/devices/gpu/nvidia/device.go @@ -31,7 +31,7 @@ const ( const ( // Nvidia-container-runtime environment variable names - nvidiaVisibleDevices = "NVIDIA_VISIBLE_DEVICES" + NvidiaVisibleDevices = "NVIDIA_VISIBLE_DEVICES" ) var ( @@ -181,7 +181,7 @@ func (d *NvidiaDevice) Reserve(deviceIDs []string) (*device.ContainerReservation return &device.ContainerReservation{ Envs: map[string]string{ - nvidiaVisibleDevices: strings.Join(deviceIDs, ","), + NvidiaVisibleDevices: strings.Join(deviceIDs, ","), }, }, nil } diff --git a/devices/gpu/nvidia/device_test.go b/devices/gpu/nvidia/device_test.go index e6daec967..717491f2b 100644 --- a/devices/gpu/nvidia/device_test.go +++ b/devices/gpu/nvidia/device_test.go @@ -73,7 +73,7 @@ func TestReserve(t *testing.T) { Name: "All RequestedIDs are managed by Device", ExpectedReservation: &device.ContainerReservation{ Envs: map[string]string{ - nvidiaVisibleDevices: "UUID1,UUID2,UUID3", + NvidiaVisibleDevices: "UUID1,UUID2,UUID3", }, }, ExpectedError: nil, diff --git a/drivers/docker/config.go b/drivers/docker/config.go index 2c1e5567a..0398137e0 100644 --- a/drivers/docker/config.go +++ b/drivers/docker/config.go @@ -101,6 +101,12 @@ func PluginLoader(opts map[string]string) (map[string]interface{}, error) { if v, err := strconv.ParseBool(opts["docker.privileged.enabled"]); err == nil { conf["allow_privileged"] = v } + + // nvidia_runtime + if v, ok := opts["docker.nvidia_runtime"]; ok { + conf["nvidia_runtime"] = v + } + return conf, nil } @@ -153,6 +159,7 @@ var ( // } // allow_privileged = false // allow_caps = ["CHOWN", "NET_RAW" ... ] + // nvidia_runtime = "nvidia" // } // } configSpec = hclspec.NewObject(map[string]*hclspec.Spec{ @@ -204,6 +211,10 @@ var ( hclspec.NewAttr("allow_caps", "list(string)", false), hclspec.NewLiteral(`["CHOWN","DAC_OVERRIDE","FSETID","FOWNER","MKNOD","NET_RAW","SETGID","SETUID","SETFCAP","SETPCAP","NET_BIND_SERVICE","SYS_CHROOT","KILL","AUDIT_WRITE"]`), ), + "nvidia_runtime": hclspec.NewDefault( + hclspec.NewAttr("nvidia_runtime", "string", false), + hclspec.NewLiteral(`"nvidia"`), + ), }) // taskConfigSpec is the hcl specification for the driver config section of @@ -470,6 +481,7 @@ type DriverConfig struct { Volumes VolumeConfig `codec:"volumes"` AllowPrivileged bool `codec:"allow_privileged"` AllowCaps []string `codec:"allow_caps"` + GPURuntimeName string `codec:"nvidia_runtime"` } type AuthConfig struct { diff --git a/drivers/docker/driver.go b/drivers/docker/driver.go index 7c2cd8b1f..d20de68ad 100644 --- a/drivers/docker/driver.go +++ b/drivers/docker/driver.go @@ -17,6 +17,7 @@ import ( hclog "github.com/hashicorp/go-hclog" multierror "github.com/hashicorp/go-multierror" "github.com/hashicorp/nomad/client/taskenv" + "github.com/hashicorp/nomad/devices/gpu/nvidia" "github.com/hashicorp/nomad/drivers/docker/docklog" "github.com/hashicorp/nomad/drivers/shared/eventer" nstructs "github.com/hashicorp/nomad/nomad/structs" @@ -84,6 +85,9 @@ type Driver struct { // logger will log to the Nomad agent logger hclog.Logger + + // gpuRuntime indicates nvidia-docker runtime availability + gpuRuntime bool } // NewDockerDriver returns a docker implementation of a driver plugin @@ -625,6 +629,13 @@ func (d *Driver) createContainerConfig(task *drivers.TaskConfig, driverConfig *T PidsLimit: driverConfig.PidsLimit, } + if _, ok := task.DeviceEnv[nvidia.NvidiaVisibleDevices]; ok { + if !d.gpuRuntime { + return c, fmt.Errorf("requested docker-runtime %q was not found", d.config.GPURuntimeName) + } + hostConfig.Runtime = d.config.GPURuntimeName + } + // Calculate CPU Quota // cfs_quota_us is the time per core, so we must // multiply the time by the number of cores available diff --git a/drivers/docker/driver_test.go b/drivers/docker/driver_test.go index 9b3124626..e2d681a82 100644 --- a/drivers/docker/driver_test.go +++ b/drivers/docker/driver_test.go @@ -23,6 +23,7 @@ import ( "github.com/hashicorp/nomad/client/allocdir" "github.com/hashicorp/nomad/client/taskenv" "github.com/hashicorp/nomad/client/testutil" + "github.com/hashicorp/nomad/devices/gpu/nvidia" "github.com/hashicorp/nomad/helper/testlog" "github.com/hashicorp/nomad/helper/uuid" "github.com/hashicorp/nomad/nomad/structs" @@ -95,8 +96,9 @@ func dockerTask(t *testing.T) (*drivers.TaskConfig, *TaskConfig, []int) { Args: busyboxLongRunningCmd[1:], } task := &drivers.TaskConfig{ - ID: uuid.Generate(), - Name: "redis-demo", + ID: uuid.Generate(), + Name: "redis-demo", + DeviceEnv: make(map[string]string), Resources: &drivers.Resources{ NomadResources: &structs.AllocatedTaskResources{ Memory: structs.AllocatedMemoryResources{ @@ -1136,6 +1138,82 @@ func TestDockerDriver_CreateContainerConfig(t *testing.T) { require.Equal(t, "org/repo:0.1", c.Config.Image) require.EqualValues(t, opt, c.HostConfig.StorageOpt) } + +func TestDockerDriver_CreateContainerConfigWithRuntimes(t *testing.T) { + if !tu.IsTravis() { + t.Parallel() + } + if !testutil.DockerIsConnected(t) { + t.Skip("Docker not connected") + } + if runtime.GOOS != "linux" { + t.Skip("nvidia plugin supports only linux") + } + testCases := []struct { + description string + gpuRuntimeSet bool + expectToReturnError bool + expectedRuntime string + nvidiaDevicesProvided bool + }{ + { + description: "gpu devices are provided, docker driver was able to detect nvidia-runtime 1", + gpuRuntimeSet: true, + expectToReturnError: false, + expectedRuntime: "nvidia", + nvidiaDevicesProvided: true, + }, + { + description: "gpu devices are provided, docker driver was able to detect nvidia-runtime 2", + gpuRuntimeSet: true, + expectToReturnError: false, + expectedRuntime: "nvidia-runtime-modified-name", + nvidiaDevicesProvided: true, + }, + { + description: "no gpu devices provided - no runtime should be set", + gpuRuntimeSet: true, + expectToReturnError: false, + expectedRuntime: "nvidia", + nvidiaDevicesProvided: false, + }, + { + description: "no gpuRuntime supported by docker driver", + gpuRuntimeSet: false, + expectToReturnError: true, + expectedRuntime: "nvidia", + nvidiaDevicesProvided: true, + }, + } + for _, testCase := range testCases { + t.Run(testCase.description, func(t *testing.T) { + task, cfg, _ := dockerTask(t) + + dh := dockerDriverHarness(t, nil) + driver := dh.Impl().(*Driver) + + driver.gpuRuntime = testCase.gpuRuntimeSet + driver.config.GPURuntimeName = testCase.expectedRuntime + if testCase.nvidiaDevicesProvided { + task.DeviceEnv[nvidia.NvidiaVisibleDevices] = "GPU_UUID_1" + } + + c, err := driver.createContainerConfig(task, cfg, "org/repo:0.1") + if testCase.expectToReturnError { + require.NotNil(t, err) + } else { + require.NoError(t, err) + if testCase.nvidiaDevicesProvided { + require.Equal(t, testCase.expectedRuntime, c.HostConfig.Runtime) + } else { + // no nvidia devices provided -> no point to use nvidia runtime + require.Equal(t, "", c.HostConfig.Runtime) + } + } + }) + } +} + func TestDockerDriver_Capabilities(t *testing.T) { if !tu.IsTravis() { t.Parallel() diff --git a/drivers/docker/fingerprint.go b/drivers/docker/fingerprint.go index d20550153..acdad87fc 100644 --- a/drivers/docker/fingerprint.go +++ b/drivers/docker/fingerprint.go @@ -2,6 +2,8 @@ package docker import ( "context" + "sort" + "strings" "time" "github.com/hashicorp/nomad/plugins/drivers" @@ -88,5 +90,23 @@ func (d *Driver) buildFingerprint() *drivers.Fingerprint { } } + if dockerInfo, err := client.Info(); err != nil { + d.logger.Warn("failed to get Docker system info", "error", err) + } else { + runtimeNames := make([]string, 0, len(dockerInfo.Runtimes)) + for name := range dockerInfo.Runtimes { + if d.config.GPURuntimeName == name { + // Nvidia runtime is detected by Docker. + // It makes possible to run GPU workloads using Docker driver on this host. + d.gpuRuntime = true + } + runtimeNames = append(runtimeNames, name) + } + sort.Strings(runtimeNames) + + fp.Attributes["runtimes"] = pstructs.NewStringAttribute( + strings.Join(runtimeNames, ",")) + } + return fp } diff --git a/website/source/docs/drivers/docker.html.md b/website/source/docs/drivers/docker.html.md index 0a69f6801..a9c6f7a0c 100644 --- a/website/source/docs/drivers/docker.html.md +++ b/website/source/docs/drivers/docker.html.md @@ -688,6 +688,8 @@ options](/docs/configuration/client.html#options): disable Nomad from removing a container when the task exits. Under a name conflict, Nomad may still remove the dead container. +* `docker.nvidia_runtime`: Defaults to `nvidia`. This option allows operators to select the runtime that should be used in order to expose Nvidia GPUs to the container. + Note: When testing or using the `-dev` flag you can use `DOCKER_HOST`, `DOCKER_TLS_VERIFY`, and `DOCKER_CERT_PATH` to customize Nomad's behavior. If `docker.endpoint` is set Nomad will **only** read client configuration from the