diff --git a/internal/config/image/cuda_image.go b/internal/config/image/cuda_image.go index 55990f61..4fbeb68e 100644 --- a/internal/config/image/cuda_image.go +++ b/internal/config/image/cuda_image.go @@ -56,8 +56,12 @@ type CUDA struct { // NewCUDAImageFromSpec creates a CUDA image from the input OCI runtime spec. // The process environment is read (if present) to construc the CUDA Image. func NewCUDAImageFromSpec(spec *specs.Spec, opts ...Option) (CUDA, error) { + if spec == nil { + return New(opts...) + } + var env []string - if spec != nil && spec.Process != nil { + if spec.Process != nil { env = spec.Process.Env } @@ -219,19 +223,12 @@ func parseMajorMinorVersion(version string) (string, error) { // OnlyFullyQualifiedCDIDevices returns true if all devices requested in the image are requested as CDI devices/ func (i CUDA) OnlyFullyQualifiedCDIDevices() bool { var hasCDIdevice bool - for _, device := range i.VisibleDevicesFromEnvVar() { + for _, device := range i.VisibleDevices() { if !parser.IsQualifiedName(device) { return false } hasCDIdevice = true } - - for _, device := range i.DevicesFromMounts() { - if !strings.HasPrefix(device, "cdi/") { - return false - } - hasCDIdevice = true - } return hasCDIdevice } @@ -309,20 +306,27 @@ func (i CUDA) VisibleDevicesFromEnvVar() []string { // visibleDevicesFromMounts returns the set of visible devices requested as mounts. func (i CUDA) visibleDevicesFromMounts() []string { var devices []string - for _, device := range i.DevicesFromMounts() { + for _, device := range i.requestsFromMounts() { switch { - case strings.HasPrefix(device, volumeMountDevicePrefixCDI): - continue case strings.HasPrefix(device, volumeMountDevicePrefixImex): continue + case strings.HasPrefix(device, volumeMountDevicePrefixCDI): + name, err := cdiDeviceMountRequest(device).qualifiedName() + if err != nil { + i.logger.Warningf("Ignoring invalid mount request for CDI device %v: %v", device, err) + continue + } + devices = append(devices, name) + default: + devices = append(devices, device) } - devices = append(devices, device) + } return devices } -// DevicesFromMounts returns a list of device specified as mounts. -func (i CUDA) DevicesFromMounts() []string { +// requestsFromMounts returns a list of device specified as mounts. +func (i CUDA) requestsFromMounts() []string { root := filepath.Clean(DeviceListAsVolumeMountsRoot) seen := make(map[string]bool) var devices []string @@ -354,23 +358,30 @@ func (i CUDA) DevicesFromMounts() []string { return devices } -// CDIDevicesFromMounts returns a list of CDI devices specified as mounts on the image. -func (i CUDA) CDIDevicesFromMounts() []string { - var devices []string - for _, mountDevice := range i.DevicesFromMounts() { - if !strings.HasPrefix(mountDevice, volumeMountDevicePrefixCDI) { - continue - } - parts := strings.SplitN(strings.TrimPrefix(mountDevice, volumeMountDevicePrefixCDI), "/", 3) - if len(parts) != 3 { - continue - } - vendor := parts[0] - class := parts[1] - device := parts[2] - devices = append(devices, fmt.Sprintf("%s/%s=%s", vendor, class, device)) +// a cdiDeviceMountRequest represents a CDI device requests as a mount. +// Here the host path /dev/null is mounted to a particular path in the container. +// The container path has the form: +// /var/run/nvidia-container-devices/cdi/// +// or +// /var/run/nvidia-container-devices/cdi//= +type cdiDeviceMountRequest string + +// qualifiedName returns the fully-qualified name of the CDI device. +func (m cdiDeviceMountRequest) qualifiedName() (string, error) { + if !strings.HasPrefix(string(m), volumeMountDevicePrefixCDI) { + return "", fmt.Errorf("invalid mount CDI device request: %s", m) } - return devices + + requestedDevice := strings.TrimPrefix(string(m), volumeMountDevicePrefixCDI) + if parser.IsQualifiedName(requestedDevice) { + return requestedDevice, nil + } + + parts := strings.SplitN(requestedDevice, "/", 3) + if len(parts) != 3 { + return "", fmt.Errorf("invalid mount CDI device request: %s", m) + } + return fmt.Sprintf("%s/%s=%s", parts[0], parts[1], parts[2]), nil } // ImexChannelsFromEnvVar returns the list of IMEX channels requested for the image. @@ -385,7 +396,7 @@ func (i CUDA) ImexChannelsFromEnvVar() []string { // ImexChannelsFromMounts returns the list of IMEX channels requested for the image. func (i CUDA) ImexChannelsFromMounts() []string { var channels []string - for _, mountDevice := range i.DevicesFromMounts() { + for _, mountDevice := range i.requestsFromMounts() { if !strings.HasPrefix(mountDevice, volumeMountDevicePrefixImex) { continue } diff --git a/internal/config/image/cuda_image_test.go b/internal/config/image/cuda_image_test.go index f2dbbee9..d60b6e0c 100644 --- a/internal/config/image/cuda_image_test.go +++ b/internal/config/image/cuda_image_test.go @@ -487,9 +487,9 @@ func TestGetVisibleDevicesFromMounts(t *testing.T) { expectedDevices: []string{"GPU0-MIG0/0/1", "GPU1-MIG0/0/1"}, }, { - description: "cdi devices are ignored", - mounts: makeTestMounts("GPU0", "cdi/nvidia.com/gpu=all", "GPU1"), - expectedDevices: []string{"GPU0", "GPU1"}, + description: "cdi devices are included", + mounts: makeTestMounts("GPU0", "nvidia.com/gpu=all", "GPU1"), + expectedDevices: []string{"GPU0", "nvidia.com/gpu=all", "GPU1"}, }, { description: "imex devices are ignored", diff --git a/internal/info/auto_test.go b/internal/info/auto_test.go index ad147526..c2ab93d7 100644 --- a/internal/info/auto_test.go +++ b/internal/info/auto_test.go @@ -184,7 +184,7 @@ func TestResolveAutoMode(t *testing.T) { expectedMode: "legacy", }, { - description: "cdi mount and non-CDI envvar resolves to legacy", + description: "cdi mount and non-CDI envvar resolves to cdi", mode: "auto", envmap: map[string]string{ "NVIDIA_VISIBLE_DEVICES": "0", @@ -197,6 +197,22 @@ func TestResolveAutoMode(t *testing.T) { "tegra": false, "nvgpu": false, }, + expectedMode: "cdi", + }, + { + description: "non-cdi mount and CDI envvar resolves to legacy", + mode: "auto", + envmap: map[string]string{ + "NVIDIA_VISIBLE_DEVICES": "nvidia.com/gpu=0", + }, + mounts: []string{ + "/var/run/nvidia-container-devices/0", + }, + info: map[string]bool{ + "nvml": true, + "tegra": false, + "nvgpu": false, + }, expectedMode: "legacy", }, } @@ -232,6 +248,8 @@ func TestResolveAutoMode(t *testing.T) { image, _ := image.New( image.WithEnvMap(tc.envmap), image.WithMounts(mounts), + image.WithAcceptDeviceListAsVolumeMounts(true), + image.WithAcceptEnvvarUnprivileged(true), ) mode := resolveMode(logger, tc.mode, image, properties) require.EqualValues(t, tc.expectedMode, mode) diff --git a/internal/modifier/cdi.go b/internal/modifier/cdi.go index 6139a3a8..c011b4e6 100644 --- a/internal/modifier/cdi.go +++ b/internal/modifier/cdi.go @@ -34,10 +34,12 @@ import ( // CDI specifications available on the system. The NVIDIA_VISIBLE_DEVICES environment variable is // used to select the devices to include. func NewCDIModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) { - devices, err := getDevicesFromImage(logger, cfg, image) - if err != nil { - return nil, fmt.Errorf("failed to get required devices from OCI specification: %v", err) - } + deviceRequestor := newCDIDeviceRequestor( + logger, + image, + cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.DefaultKind, + ) + devices := deviceRequestor.DeviceRequests() if len(devices) == 0 { logger.Debugf("No devices requested; no modification required.") return nil, nil @@ -64,46 +66,47 @@ func NewCDIModifier(logger logger.Interface, cfg *config.Config, image image.CUD ) } -func getDevicesFromImage(logger logger.Interface, cfg *config.Config, container image.CUDA) ([]string, error) { - annotationDevices, err := getAnnotationDevices(container) +type deviceRequestor interface { + DeviceRequests() []string +} + +type cdiDeviceRequestor struct { + image image.CUDA + logger logger.Interface + defaultKind string +} + +func newCDIDeviceRequestor(logger logger.Interface, image image.CUDA, defaultKind string) deviceRequestor { + c := &cdiDeviceRequestor{ + logger: logger, + image: image, + defaultKind: defaultKind, + } + return withUniqueDevices(c) +} + +func (c *cdiDeviceRequestor) DeviceRequests() []string { + if c == nil { + return nil + } + annotationDevices, err := getAnnotationDevices(c.image) if err != nil { - return nil, fmt.Errorf("failed to parse container annotations: %v", err) + c.logger.Warningf("failed to get device requests from container annotations: %v; ignoring.", err) + annotationDevices = nil } if len(annotationDevices) > 0 { - return annotationDevices, nil - } - - if cfg.AcceptDeviceListAsVolumeMounts { - mountDevices := container.CDIDevicesFromMounts() - if len(mountDevices) > 0 { - return mountDevices, nil - } + return annotationDevices } var devices []string - seen := make(map[string]bool) - for _, name := range container.VisibleDevicesFromEnvVar() { + for _, name := range c.image.VisibleDevices() { if !parser.IsQualifiedName(name) { - name = fmt.Sprintf("%s=%s", cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.DefaultKind, name) - } - if seen[name] { - logger.Debugf("Ignoring duplicate device %q", name) - continue + name = fmt.Sprintf("%s=%s", c.defaultKind, name) } devices = append(devices, name) } - if len(devices) == 0 { - return nil, nil - } - - if cfg.AcceptEnvvarUnprivileged || container.IsPrivileged() { - return devices, nil - } - - logger.Warningf("Ignoring devices specified in NVIDIA_VISIBLE_DEVICES: %v", devices) - - return nil, nil + return devices } // getAnnotationDevices returns a list of devices specified in the annotations. @@ -111,16 +114,11 @@ func getDevicesFromImage(logger logger.Interface, cfg *config.Config, container // fully-qualified CDI devices names. If any device name is not fully-quality an error is returned. // The list of returned devices is deduplicated. func getAnnotationDevices(image image.CUDA) ([]string, error) { - seen := make(map[string]bool) var annotationDevices []string for _, device := range image.CDIDeviceRequestsFromAnnotations() { if !parser.IsQualifiedName(device) { return nil, fmt.Errorf("invalid device name %q in annotations", device) } - if seen[device] { - continue - } - seen[device] = true annotationDevices = append(annotationDevices, device) } return annotationDevices, nil @@ -147,7 +145,7 @@ func newAutomaticCDISpecModifier(logger logger.Interface, cfg *config.Config, de if err != nil { return nil, fmt.Errorf("failed to generate CDI spec: %w", err) } - cdiModifier, err := cdi.New( + cdiDeviceRequestor, err := cdi.New( cdi.WithLogger(logger), cdi.WithSpec(spec.Raw()), ) @@ -155,7 +153,7 @@ func newAutomaticCDISpecModifier(logger logger.Interface, cfg *config.Config, de return nil, fmt.Errorf("failed to construct CDI modifier: %w", err) } - return cdiModifier, nil + return cdiDeviceRequestor, nil } func generateAutomaticCDISpec(logger logger.Interface, cfg *config.Config, devices []string) (spec.Interface, error) { @@ -193,3 +191,42 @@ func generateAutomaticCDISpec(logger logger.Interface, cfg *config.Config, devic spec.WithClass("gpu"), ) } + +func deviceRequestorFromImage(image image.CUDA) deviceRequestor { + return &fromImage{image} +} + +type fromImage struct { + image.CUDA +} + +func (f *fromImage) DeviceRequests() []string { + if f == nil { + return nil + } + return f.CUDA.VisibleDevices() +} + +type deduplicatedDeviceRequestor struct { + deviceRequestor +} + +func withUniqueDevices(deviceRequestor deviceRequestor) deviceRequestor { + return &deduplicatedDeviceRequestor{deviceRequestor: deviceRequestor} +} + +func (d *deduplicatedDeviceRequestor) DeviceRequests() []string { + if d == nil { + return nil + } + seen := make(map[string]bool) + var devices []string + for _, device := range d.deviceRequestor.DeviceRequests() { + if seen[device] { + continue + } + seen[device] = true + devices = append(devices, device) + } + return devices +} diff --git a/internal/modifier/cdi_test.go b/internal/modifier/cdi_test.go index 401a1e83..cf7811a5 100644 --- a/internal/modifier/cdi_test.go +++ b/internal/modifier/cdi_test.go @@ -20,6 +20,8 @@ import ( "fmt" "testing" + "github.com/opencontainers/runtime-spec/specs-go" + testlog "github.com/sirupsen/logrus/hooks/test" "github.com/stretchr/testify/require" "github.com/NVIDIA/nvidia-container-toolkit/internal/config/image" @@ -67,7 +69,7 @@ func TestGetAnnotationDevices(t *testing.T) { "prefix/foo": "example.com/device=bar", "another-prefix/bar": "example.com/device=bar", }, - expectedDevices: []string{"example.com/device=bar"}, + expectedDevices: []string{"example.com/device=bar", "example.com/device=bar"}, }, { description: "invalid devices", @@ -98,3 +100,67 @@ func TestGetAnnotationDevices(t *testing.T) { }) } } + +func TestDeviceRequests(t *testing.T) { + logger, _ := testlog.NewNullLogger() + + testCases := []struct { + description string + input cdiDeviceRequestor + spec *specs.Spec + expectedDevices []string + }{ + { + description: "empty spec yields no devices", + }, + { + description: "cdi devices from mounts", + input: cdiDeviceRequestor{ + defaultKind: "nvidia.com/gpu", + }, + spec: &specs.Spec{ + Mounts: []specs.Mount{ + { + Destination: "/var/run/nvidia-container-devices/cdi/nvidia.com/gpu/0", + Source: "/dev/null", + }, + { + Destination: "/var/run/nvidia-container-devices/cdi/nvidia.com/gpu/1", + Source: "/dev/null", + }, + }, + }, + expectedDevices: []string{"nvidia.com/gpu=0", "nvidia.com/gpu=1"}, + }, + { + description: "cdi devices from envvar", + input: cdiDeviceRequestor{ + defaultKind: "nvidia.com/gpu", + }, + spec: &specs.Spec{ + Process: &specs.Process{ + Env: []string{"NVIDIA_VISIBLE_DEVICES=0,example.com/class=device"}, + }, + }, + expectedDevices: []string{"nvidia.com/gpu=0", "example.com/class=device"}, + }, + } + + for _, tc := range testCases { + tc.input.logger = logger + + image, err := image.NewCUDAImageFromSpec( + tc.spec, + image.WithAcceptDeviceListAsVolumeMounts(true), + image.WithAcceptEnvvarUnprivileged(true), + ) + require.NoError(t, err) + tc.input.image = image + + t.Run(tc.description, func(t *testing.T) { + devices := tc.input.DeviceRequests() + require.NoError(t, err) + require.EqualValues(t, tc.expectedDevices, devices) + }) + } +}