Make CDI device requests consistent with other methods

Following the refactoring of device request extraction, we can
now make CDI device requests consistent with other methods.

This change moves to using image.VisibleDevices instead of
separate calls to CDIDevicesFromMounts and VisibleDevicesFromEnvVar.

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

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/<vendor>/<class>/<device>
+// or
+// /var/run/nvidia-container-devices/cdi/<vendor>/<class>=<device>
+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
 		}

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",

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)

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,63 +66,59 @@ 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.
 // Keys starting with the specified prefixes are considered and expected to contain a comma-separated list of
 // 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,15 +145,15 @@ 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()),
 	)
 	if err != nil {
 		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
+}