Remove redundant CDI annotations

cmd/nvidia-ctk/cdi/generate/generate.go

@@ -447,15 +447,22 @@ type deviceSpecs []specs.Device
 func (d deviceSpecs) splitOnAnnotation(key string) map[string][]specs.Device {
 	splitSpecs := make(map[string][]specs.Device)
 
+	var specsToRemoveAnnotations []*specs.Device
 	for _, deviceSpec := range d {
-		if len(deviceSpec.Annotations) == 0 {
-			continue
-		}
 		value, ok := deviceSpec.Annotations[key]
 		if !ok {
 			continue
 		}
 		splitSpecs[key+"="+value] = append(splitSpecs[key+"="+value], deviceSpec)
+		specsToRemoveAnnotations = append(specsToRemoveAnnotations, &deviceSpec)
+	}
+
+	// We also remove the annotations that were used to split the devices:
+	for _, deviceSpec := range specsToRemoveAnnotations {
+		if _, ok := deviceSpec.Annotations[key]; !ok {
+			continue
+		}
+		delete(deviceSpec.Annotations, key)
 	}
 
 	return splitSpecs

cmd/nvidia-ctk/cdi/generate/generate_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/NVIDIA/go-nvml/pkg/nvml/mock/dgxa100"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
+	"tags.cncf.io/container-device-interface/specs-go"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/test"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi"
@@ -491,3 +492,162 @@ containerEdits:
 		})
 	}
 }
+
+func TestSplitOnAnnotation(t *testing.T) {
+	testCases := []struct {
+		description            string
+		input                  deviceSpecs
+		annotation             string
+		expectedInputPostSplit deviceSpecs
+		expectedOutput         map[string][]specs.Device
+	}{
+		{
+			description: "non-matching annotation",
+			input: deviceSpecs{
+				specs.Device{
+					Name: "foo",
+					Annotations: map[string]string{
+						"key": "value",
+					},
+				},
+			},
+			annotation: "not-key",
+			expectedInputPostSplit: deviceSpecs{
+				specs.Device{
+					Name: "foo",
+					Annotations: map[string]string{
+						"key": "value",
+					},
+				},
+			},
+			expectedOutput: map[string][]specs.Device{},
+		},
+		{
+			description: "single annotation present",
+			input: deviceSpecs{
+				specs.Device{
+					Name: "foo",
+					Annotations: map[string]string{
+						"key": "value",
+					},
+				},
+			},
+			annotation: "key",
+			expectedInputPostSplit: deviceSpecs{
+				specs.Device{
+					Name:        "foo",
+					Annotations: map[string]string{},
+				},
+			},
+			expectedOutput: map[string][]specs.Device{
+				"key=value": {
+					{
+						Name:        "foo",
+						Annotations: map[string]string{},
+					},
+				},
+			},
+		},
+		{
+			description: "non-matching annotations are not removed",
+			input: deviceSpecs{
+				specs.Device{
+					Name: "foo",
+					Annotations: map[string]string{
+						"key":     "value",
+						"another": "foo",
+					},
+				},
+			},
+			annotation: "key",
+			expectedInputPostSplit: deviceSpecs{
+				specs.Device{
+					Name:        "foo",
+					Annotations: map[string]string{"another": "foo"},
+				},
+			},
+			expectedOutput: map[string][]specs.Device{
+				"key=value": {
+					{
+						Name:        "foo",
+						Annotations: map[string]string{"another": "foo"},
+					},
+				},
+			},
+		},
+		{
+			description: "duplicated annotations different names",
+			input: func() deviceSpecs {
+				annotations := map[string]string{
+					"key": "value",
+				}
+				return deviceSpecs{
+					{Name: "0", Annotations: annotations},
+					{Name: "first", Annotations: annotations},
+				}
+			}(),
+			annotation: "key",
+			expectedInputPostSplit: deviceSpecs{
+				specs.Device{Name: "0", Annotations: map[string]string{}},
+				specs.Device{Name: "first", Annotations: map[string]string{}},
+			},
+			expectedOutput: map[string][]specs.Device{
+				"key=value": {
+					{Name: "0", Annotations: map[string]string{}},
+					{Name: "first", Annotations: map[string]string{}},
+				},
+			},
+		},
+		{
+			description: "annotation with different values",
+			input: deviceSpecs{
+				specs.Device{
+					Name: "foo",
+					Annotations: map[string]string{
+						"key": "value",
+					},
+				},
+				specs.Device{
+					Name: "bar",
+					Annotations: map[string]string{
+						"key": "another",
+					},
+				},
+			},
+			annotation: "key",
+			expectedInputPostSplit: deviceSpecs{
+				specs.Device{
+					Name:        "foo",
+					Annotations: map[string]string{},
+				},
+				specs.Device{
+					Name:        "bar",
+					Annotations: map[string]string{},
+				},
+			},
+			expectedOutput: map[string][]specs.Device{
+				"key=value": {
+					{
+						Name:        "foo",
+						Annotations: map[string]string{},
+					},
+				},
+				"key=another": {
+					{
+						Name:        "bar",
+						Annotations: map[string]string{},
+					},
+				},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			result := tc.input.splitOnAnnotation(tc.annotation)
+
+			require.EqualValues(t, tc.expectedOutput, result)
+			require.EqualValues(t, tc.expectedInputPostSplit, tc.input)
+		})
+	}
+}