Return annotation devices from VisibleDevices

This change includes annotation devices in CUDA.VisibleDevices
with the highest priority. This allows for the CDI device
request extraction to be consistent across all request mechanisms.

Note that this does change behaviour in the following ways:
1. Annotations are considered when resolving the runtime mode.
2. Incorrectly formed device names in annotations are no longer treated as an error.

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

Author: elezar

internal/config/image/cuda_image.go

@@ -238,6 +238,12 @@ func (i CUDA) OnlyFullyQualifiedCDIDevices() bool {
 // In cases where environment variable requests required privileged containers,
 // such devices requests are ignored.
 func (i CUDA) VisibleDevices() []string {
+	// If annotation device requests are present, these are preferred.
+	annotationDeviceRequests := i.cdiDeviceRequestsFromAnnotations()
+	if len(annotationDeviceRequests) > 0 {
+		return annotationDeviceRequests
+	}
+
 	// If enabled, try and get the device list from volume mounts first
 	if i.acceptDeviceListAsVolumeMounts {
 		volumeMountDeviceRequests := i.visibleDevicesFromMounts()
@@ -264,29 +270,28 @@ func (i CUDA) VisibleDevices() []string {
 	return nil
 }
 
-// CDIDeviceRequestsFromAnnotations returns a list of devices specified in the annotations.
+// cdiDeviceRequestsFromAnnotations 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.
 // The format of the requested devices is not checked and the list is not
 // deduplicated.
-func (i CUDA) CDIDeviceRequestsFromAnnotations() []string {
+func (i CUDA) cdiDeviceRequestsFromAnnotations() []string {
 	if len(i.annotationsPrefixes) == 0 || len(i.annotations) == 0 {
 		return nil
 	}
 
-	var deviceKeys []string
-	for key := range i.annotations {
+	var devices []string
+	for key, value := range i.annotations {
 		for _, prefix := range i.annotationsPrefixes {
 			if strings.HasPrefix(key, prefix) {
-				deviceKeys = append(deviceKeys, key)
+				devices = append(devices, strings.Split(value, ",")...)
+				// There is no need to check additional prefixes since we
+				// typically deduplicate devices in any case.
+				break
 			}
 		}
 	}
-
-	var devices []string
-	for _, key := range deviceKeys {
-		devices = append(devices, strings.Split(i.annotations[key], ",")...)
-	}
 	return devices
 }
 

internal/config/image/cuda_image_test.go

@@ -710,7 +710,7 @@ func TestCDIDeviceRequestsFromAnnotations(t *testing.T) {
 			)
 			require.NoError(t, err)
 
-			devices := image.CDIDeviceRequestsFromAnnotations()
+			devices := image.cdiDeviceRequestsFromAnnotations()
 			require.ElementsMatch(t, tc.expectedDevices, devices)
 		})
 	}

internal/modifier/cdi.go

@@ -89,15 +89,6 @@ func (c *cdiDeviceRequestor) DeviceRequests() []string {
 	if c == nil {
 		return nil
 	}
-	annotationDevices, err := getAnnotationDevices(c.image)
-	if err != nil {
-		c.logger.Warningf("failed to get device requests from container annotations: %v; ignoring.", err)
-		annotationDevices = nil
-	}
-	if len(annotationDevices) > 0 {
-		return annotationDevices
-	}
-
 	var devices []string
 	for _, name := range c.image.VisibleDevices() {
 		if !parser.IsQualifiedName(name) {
@@ -109,21 +100,6 @@ func (c *cdiDeviceRequestor) DeviceRequests() []string {
 	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) {
-	var annotationDevices []string
-	for _, device := range image.CDIDeviceRequestsFromAnnotations() {
-		if !parser.IsQualifiedName(device) {
-			return nil, fmt.Errorf("invalid device name %q in annotations", device)
-		}
-		annotationDevices = append(annotationDevices, device)
-	}
-	return annotationDevices, nil
-}
-
 // filterAutomaticDevices searches for "automatic" device names in the input slice.
 // "Automatic" devices are a well-defined list of CDI device names which, when requested,
 // trigger the generation of a CDI spec at runtime. This removes the need to generate a
@@ -192,21 +168,6 @@ func generateAutomaticCDISpec(logger logger.Interface, cfg *config.Config, devic
 	)
 }
 
-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
 }

internal/modifier/cdi_test.go

@@ -17,7 +17,6 @@
 package modifier
 
 import (
-	"fmt"
 	"testing"
 
 	"github.com/opencontainers/runtime-spec/specs-go"
@@ -27,122 +26,116 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
 )
 
-func TestGetAnnotationDevices(t *testing.T) {
+func TestDeviceRequests(t *testing.T) {
+	logger, _ := testlog.NewNullLogger()
+
 	testCases := []struct {
 		description     string
+		input           cdiDeviceRequestor
+		spec            *specs.Spec
 		prefixes        []string
-		annotations     map[string]string
 		expectedDevices []string
-		expectedError   error
 	}{
 		{
-			description: "no annotations",
+			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"},
 		},
 		{
 			description: "no matching annotations",
 			prefixes:    []string{"not-prefix/"},
-			annotations: map[string]string{
-				"prefix/foo": "example.com/device=bar",
+			spec: &specs.Spec{
+				Annotations: map[string]string{
+					"prefix/foo": "example.com/device=bar",
+				},
 			},
 		},
 		{
 			description: "single matching annotation",
 			prefixes:    []string{"prefix/"},
-			annotations: map[string]string{
-				"prefix/foo": "example.com/device=bar",
+			spec: &specs.Spec{
+				Annotations: map[string]string{
+					"prefix/foo": "example.com/device=bar",
+				},
 			},
 			expectedDevices: []string{"example.com/device=bar"},
 		},
 		{
 			description: "multiple matching annotations",
 			prefixes:    []string{"prefix/", "another-prefix/"},
-			annotations: map[string]string{
-				"prefix/foo":         "example.com/device=bar",
-				"another-prefix/bar": "example.com/device=baz",
+			spec: &specs.Spec{
+				Annotations: map[string]string{
+					"prefix/foo":         "example.com/device=bar",
+					"another-prefix/bar": "example.com/device=baz",
+				},
 			},
 			expectedDevices: []string{"example.com/device=bar", "example.com/device=baz"},
 		},
 		{
 			description: "multiple matching annotations with duplicate devices",
 			prefixes:    []string{"prefix/", "another-prefix/"},
-			annotations: map[string]string{
-				"prefix/foo":         "example.com/device=bar",
-				"another-prefix/bar": "example.com/device=bar",
+			spec: &specs.Spec{
+				Annotations: map[string]string{
+					"prefix/foo":         "example.com/device=bar",
+					"another-prefix/bar": "example.com/device=bar",
+				},
 			},
 			expectedDevices: []string{"example.com/device=bar", "example.com/device=bar"},
 		},
 		{
-			description: "invalid devices",
-			prefixes:    []string{"prefix/"},
-			annotations: map[string]string{
-				"prefix/foo": "example.com/device",
-			},
-			expectedError: fmt.Errorf("invalid device %q", "example.com/device"),
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.description, func(t *testing.T) {
-			image, err := image.New(
-				image.WithAnnotations(tc.annotations),
-				image.WithAnnotationsPrefixes(tc.prefixes),
-			)
-			require.NoError(t, err)
-
-			devices, err := getAnnotationDevices(image)
-			if tc.expectedError != nil {
-				require.Error(t, err)
-				return
-			}
-
-			require.NoError(t, err)
-			require.ElementsMatch(t, tc.expectedDevices, devices)
-		})
-	}
-}
-
-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",
+			description: "devices in annotations are expanded",
 			input: cdiDeviceRequestor{
 				defaultKind: "nvidia.com/gpu",
 			},
+			prefixes: []string{"prefix/"},
 			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",
-					},
+				Annotations: map[string]string{
+					"prefix/foo": "device",
 				},
 			},
-			expectedDevices: []string{"nvidia.com/gpu=0", "nvidia.com/gpu=1"},
+			expectedDevices: []string{"nvidia.com/gpu=device"},
 		},
 		{
-			description: "cdi devices from envvar",
+			description: "invalid devices in annotations are treated as strings",
 			input: cdiDeviceRequestor{
 				defaultKind: "nvidia.com/gpu",
 			},
+			prefixes: []string{"prefix/"},
 			spec: &specs.Spec{
-				Process: &specs.Process{
-					Env: []string{"NVIDIA_VISIBLE_DEVICES=0,example.com/class=device"},
+				Annotations: map[string]string{
+					"prefix/foo": "example.com/device",
 				},
 			},
-			expectedDevices: []string{"nvidia.com/gpu=0", "example.com/class=device"},
+			expectedDevices: []string{"nvidia.com/gpu=example.com/device"},
 		},
 	}
 
@@ -153,6 +146,7 @@ func TestDeviceRequests(t *testing.T) {
 			tc.spec,
 			image.WithAcceptDeviceListAsVolumeMounts(true),
 			image.WithAcceptEnvvarUnprivileged(true),
+			image.WithAnnotationsPrefixes(tc.prefixes),
 		)
 		require.NoError(t, err)
 		tc.input.image = image