Add logic to extract annotation device requests to image type

This change updates the image.CUDA type to also extract CDI
device requests. These are only relevant IF CDI prefixes are
specifically set.

Signed-off-by: Evan Lezar <elezar@nvidia.com>

internal/config/image/builder.go

@@ -50,7 +50,6 @@ func New(opt ...Option) (CUDA, error) {
 	if b.logger == nil {
 		b.logger = logger.New()
 	}
-
 	if b.env == nil {
 		b.env = make(map[string]string)
 	}
@@ -81,6 +80,20 @@ func WithAcceptEnvvarUnprivileged(acceptEnvvarUnprivileged bool) Option {
 	}
 }
 
+func WithAnnotations(annotations map[string]string) Option {
+	return func(b *builder) error {
+		b.annotations = annotations
+		return nil
+	}
+}
+
+func WithAnnotationsPrefixes(annotationsPrefixes []string) Option {
+	return func(b *builder) error {
+		b.annotationsPrefixes = annotationsPrefixes
+		return nil
+	}
+}
+
 // WithDisableRequire sets the disable require option.
 func WithDisableRequire(disableRequire bool) Option {
 	return func(b *builder) error {

internal/config/image/cuda_image.go

@@ -42,10 +42,12 @@ const (
 type CUDA struct {
 	logger logger.Interface
 
+	annotations  map[string]string
 	env          map[string]string
 	isPrivileged bool
 	mounts       []specs.Mount
 
+	annotationsPrefixes            []string
 	acceptDeviceListAsVolumeMounts bool
 	acceptEnvvarUnprivileged       bool
 	preferredVisibleDeviceEnvVars  []string
@@ -60,6 +62,7 @@ func NewCUDAImageFromSpec(spec *specs.Spec, opts ...Option) (CUDA, error) {
 	}
 
 	specOpts := []Option{
+		WithAnnotations(spec.Annotations),
 		WithEnv(env),
 		WithMounts(spec.Mounts),
 		WithPrivileged(IsPrivileged((*OCISpec)(spec))),
@@ -264,6 +267,32 @@ func (i CUDA) VisibleDevices() []string {
 	return nil
 }
 
+// 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 {
+	if len(i.annotationsPrefixes) == 0 || len(i.annotations) == 0 {
+		return nil
+	}
+
+	var deviceKeys []string
+	for key := range i.annotations {
+		for _, prefix := range i.annotationsPrefixes {
+			if strings.HasPrefix(key, prefix) {
+				deviceKeys = append(deviceKeys, key)
+			}
+		}
+	}
+
+	var devices []string
+	for _, key := range deviceKeys {
+		devices = append(devices, strings.Split(i.annotations[key], ",")...)
+	}
+	return devices
+}
+
 // VisibleDevicesFromEnvVar returns the set of visible devices requested through environment variables.
 // If any of the preferredVisibleDeviceEnvVars are present in the image, they
 // are used to determine the visible devices. If this is not the case, the

internal/config/image/cuda_image_test.go

@@ -649,6 +649,73 @@ func TestImexChannelsFromEnvVar(t *testing.T) {
 	}
 }
 
+func TestCDIDeviceRequestsFromAnnotations(t *testing.T) {
+	testCases := []struct {
+		description     string
+		prefixes        []string
+		annotations     map[string]string
+		expectedDevices []string
+	}{
+		{
+			description: "no annotations",
+		},
+		{
+			description: "no matching annotations",
+			prefixes:    []string{"not-prefix/"},
+			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",
+			},
+			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",
+			},
+			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",
+			},
+			expectedDevices: []string{"example.com/device=bar", "example.com/device=bar"},
+		},
+		{
+			description: "invalid devices are returned as is",
+			prefixes:    []string{"prefix/"},
+			annotations: map[string]string{
+				"prefix/foo": "example.com/device",
+			},
+			expectedDevices: []string{"example.com/device"},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			image, err := New(
+				WithAnnotationsPrefixes(tc.prefixes),
+				WithAnnotations(tc.annotations),
+			)
+			require.NoError(t, err)
+
+			devices := image.CDIDeviceRequestsFromAnnotations()
+			require.ElementsMatch(t, tc.expectedDevices, devices)
+		})
+	}
+}
+
 func makeTestMounts(paths ...string) []specs.Mount {
 	var mounts []specs.Mount
 	for _, path := range paths {

internal/modifier/cdi.go

@@ -18,7 +18,6 @@ package modifier
 
 import (
 	"fmt"
-	"strings"
 
 	"tags.cncf.io/container-device-interface/pkg/parser"
 
@@ -71,7 +70,18 @@ func getDevicesFromSpec(logger logger.Interface, ociSpec oci.Spec, cfg *config.C
 		return nil, fmt.Errorf("failed to load OCI spec: %v", err)
 	}
 
-	annotationDevices, err := getAnnotationDevices(cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.AnnotationPrefixes, rawSpec.Annotations)
+	container, err := image.NewCUDAImageFromSpec(
+		rawSpec,
+		image.WithLogger(logger),
+		image.WithAcceptDeviceListAsVolumeMounts(cfg.AcceptDeviceListAsVolumeMounts),
+		image.WithAcceptEnvvarUnprivileged(cfg.AcceptEnvvarUnprivileged),
+		image.WithAnnotationsPrefixes(cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.AnnotationPrefixes),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	annotationDevices, err := getAnnotationDevices(container)
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse container annotations: %v", err)
 	}
@@ -79,13 +89,6 @@ func getDevicesFromSpec(logger logger.Interface, ociSpec oci.Spec, cfg *config.C
 		return annotationDevices, nil
 	}
 
-	container, err := image.NewCUDAImageFromSpec(
-		rawSpec,
-		image.WithLogger(logger),
-	)
-	if err != nil {
-		return nil, err
-	}
 	if cfg.AcceptDeviceListAsVolumeMounts {
 		mountDevices := container.CDIDevicesFromMounts()
 		if len(mountDevices) > 0 {
@@ -123,31 +126,19 @@ func getDevicesFromSpec(logger logger.Interface, ociSpec oci.Spec, cfg *config.C
 // 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(prefixes []string, annotations map[string]string) ([]string, error) {
-	devicesByKey := make(map[string][]string)
-	for key, value := range annotations {
-		for _, prefix := range prefixes {
-			if strings.HasPrefix(key, prefix) {
-				devicesByKey[key] = strings.Split(value, ",")
-			}
-		}
-	}
-
+func getAnnotationDevices(image image.CUDA) ([]string, error) {
 	seen := make(map[string]bool)
 	var annotationDevices []string
-	for key, devices := range devicesByKey {
-		for _, device := range devices {
-			if !parser.IsQualifiedName(device) {
-				return nil, fmt.Errorf("invalid device name %q in annotation %q", device, key)
-			}
-			if seen[device] {
-				continue
-			}
-			annotationDevices = append(annotationDevices, device)
-			seen[device] = true
+	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
 }
 

internal/modifier/cdi_test.go

@@ -21,6 +21,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
 )
 
 func TestGetAnnotationDevices(t *testing.T) {
@@ -79,7 +81,13 @@ func TestGetAnnotationDevices(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			devices, err := getAnnotationDevices(tc.prefixes, tc.annotations)
+			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