BUGFIX: modifier: respect GPU volume-mount device requests

The gated modifiers used to add support for GDS, Mofed, and CUDA Forward Comatibility only check the NVIDIA_VISIBLE_DEVICES envvar to determine whether GPUs are requested and modifications should be made. This means that use cases where volume mounts are used to request devices (e.g. when using the GPU Device Plugin) are not supported.

This patch takes visibleDevicesFromEnvVar private, making VisibleDevices the only exported method to query valid devices.
And edits the gated modifiers to use this func, ensuring device requests
via mounts are also taken into acount.

Signed-off-by: Carlos Eduardo Arango Gutierrez <[email protected]>

Author: ArangoGutierrez

internal/config/image/cuda_image.go

@@ -212,7 +212,7 @@ 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.visibleDevicesFromEnvVar() {
 		if !parser.IsQualifiedName(device) {
 			return false
 		}
@@ -258,7 +258,7 @@ func (i CUDA) VisibleDevices() []string {
 	}
 
 	// Get the Fallback to reading from the environment variable if privileges are correct
-	envVarDeviceRequests := i.VisibleDevicesFromEnvVar()
+	envVarDeviceRequests := i.visibleDevicesFromEnvVar()
 	if len(envVarDeviceRequests) == 0 {
 		return nil
 	}
@@ -278,11 +278,11 @@ func (i CUDA) VisibleDevices() []string {
 	return nil
 }
 
-// VisibleDevicesFromEnvVar returns the set of visible devices requested through environment variables.
+// 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
 // NVIDIA_VISIBLE_DEVICES environment variable is used.
-func (i CUDA) VisibleDevicesFromEnvVar() []string {
+func (i CUDA) visibleDevicesFromEnvVar() []string {
 	envVars := i.visibleEnvVars()
 	return i.DevicesFromEnvvars(envVars...).List()
 }

internal/config/image/cuda_image_test.go

@@ -429,7 +429,7 @@ func TestGetDevicesFromEnvvar(t *testing.T) {
 			)
 
 			require.NoError(t, err)
-			devices := image.VisibleDevicesFromEnvVar()
+			devices := image.visibleDevicesFromEnvVar()
 			require.EqualValues(t, tc.expectedDevices, devices)
 		})
 	}
@@ -508,13 +508,15 @@ func TestGetVisibleDevicesFromMounts(t *testing.T) {
 
 func TestVisibleDevices(t *testing.T) {
 	var tests = []struct {
-		description        string
-		mountDevices       []specs.Mount
-		envvarDevices      string
-		privileged         bool
-		acceptUnprivileged bool
-		acceptMounts       bool
-		expectedDevices    []string
+		description                   string
+		mountDevices                  []specs.Mount
+		envvarDevices                 string
+		privileged                    bool
+		acceptUnprivileged            bool
+		acceptMounts                  bool
+		preferredVisibleDeviceEnvVars []string
+		env                           map[string]string
+		expectedDevices               []string
 	}{
 		{
 			description: "Mount devices, unprivileged, no accept unprivileged",
@@ -597,20 +599,92 @@ func TestVisibleDevices(t *testing.T) {
 			acceptMounts:       false,
 			expectedDevices:    nil,
 		},
+		// New test cases for visibleEnvVars functionality
+		{
+			description:                   "preferred env var set and present in env, privileged",
+			mountDevices:                  nil,
+			envvarDevices:                 "",
+			privileged:                    true,
+			acceptUnprivileged:            false,
+			acceptMounts:                  true,
+			preferredVisibleDeviceEnvVars: []string{"DOCKER_RESOURCE_GPUS"},
+			env: map[string]string{
+				"DOCKER_RESOURCE_GPUS": "GPU-12345",
+			},
+			expectedDevices: []string{"GPU-12345"},
+		},
+		{
+			description:                   "preferred env var set and present in env, unprivileged but accepted",
+			mountDevices:                  nil,
+			envvarDevices:                 "",
+			privileged:                    false,
+			acceptUnprivileged:            true,
+			acceptMounts:                  true,
+			preferredVisibleDeviceEnvVars: []string{"DOCKER_RESOURCE_GPUS"},
+			env: map[string]string{
+				"DOCKER_RESOURCE_GPUS": "GPU-12345",
+			},
+			expectedDevices: []string{"GPU-12345"},
+		},
+		{
+			description:                   "preferred env var set and present in env, unprivileged and not accepted",
+			mountDevices:                  nil,
+			envvarDevices:                 "",
+			privileged:                    false,
+			acceptUnprivileged:            false,
+			acceptMounts:                  true,
+			preferredVisibleDeviceEnvVars: []string{"DOCKER_RESOURCE_GPUS"},
+			env: map[string]string{
+				"DOCKER_RESOURCE_GPUS": "GPU-12345",
+			},
+			expectedDevices: nil,
+		},
+		{
+			description:                   "multiple preferred env vars, both present, privileged",
+			mountDevices:                  nil,
+			envvarDevices:                 "",
+			privileged:                    true,
+			acceptUnprivileged:            false,
+			acceptMounts:                  true,
+			preferredVisibleDeviceEnvVars: []string{"DOCKER_RESOURCE_GPUS", "DOCKER_RESOURCE_GPUS_ADDITIONAL"},
+			env: map[string]string{
+				"DOCKER_RESOURCE_GPUS":            "GPU-12345",
+				"DOCKER_RESOURCE_GPUS_ADDITIONAL": "GPU-67890",
+			},
+			expectedDevices: []string{"GPU-12345", "GPU-67890"},
+		},
+		{
+			description:                   "preferred env var not present, fallback to NVIDIA_VISIBLE_DEVICES, privileged",
+			mountDevices:                  nil,
+			envvarDevices:                 "GPU-12345",
+			privileged:                    true,
+			acceptUnprivileged:            false,
+			acceptMounts:                  true,
+			preferredVisibleDeviceEnvVars: []string{"DOCKER_RESOURCE_GPUS"},
+			env: map[string]string{
+				EnvVarNvidiaVisibleDevices: "GPU-12345",
+			},
+			expectedDevices: []string{"GPU-12345"},
+		},
 	}
 	for _, tc := range tests {
 		t.Run(tc.description, func(t *testing.T) {
-			// Wrap the call to getDevices() in a closure.
+			// Create env map with both NVIDIA_VISIBLE_DEVICES and any additional env vars
+			env := make(map[string]string)
+			if tc.envvarDevices != "" {
+				env[EnvVarNvidiaVisibleDevices] = tc.envvarDevices
+			}
+			for k, v := range tc.env {
+				env[k] = v
+			}
+
 			image, err := New(
-				WithEnvMap(
-					map[string]string{
-						EnvVarNvidiaVisibleDevices: tc.envvarDevices,
-					},
-				),
+				WithEnvMap(env),
 				WithMounts(tc.mountDevices),
 				WithPrivileged(tc.privileged),
 				WithAcceptDeviceListAsVolumeMounts(tc.acceptMounts),
 				WithAcceptEnvvarUnprivileged(tc.acceptUnprivileged),
+				WithPreferredVisibleDevicesEnvVars(tc.preferredVisibleDeviceEnvVars...),
 			)
 			require.NoError(t, err)
 			require.Equal(t, tc.expectedDevices, image.VisibleDevices())

internal/modifier/cdi.go

@@ -86,37 +86,38 @@ func getDevicesFromSpec(logger logger.Interface, ociSpec oci.Spec, cfg *config.C
 	if err != nil {
 		return nil, err
 	}
-	if cfg.AcceptDeviceListAsVolumeMounts {
-		mountDevices := container.CDIDevicesFromMounts()
-		if len(mountDevices) > 0 {
-			return mountDevices, nil
-		}
-	}
 
 	var devices []string
 	seen := make(map[string]bool)
-	for _, name := range container.VisibleDevicesFromEnvVar() {
+	for _, name := range container.VisibleDevices() {
+		// Skip empty device names
+		if strings.TrimSpace(name) == "" {
+			continue
+		}
+
 		if !parser.IsQualifiedName(name) {
 			name = fmt.Sprintf("%s=%s", cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.DefaultKind, name)
 		}
+
+		// Skip devices that result in empty device IDs after qualification
+		if strings.HasSuffix(name, "=") {
+			logger.Debugf("Ignoring device with empty ID: %q", name)
+			continue
+		}
+
 		if seen[name] {
 			logger.Debugf("Ignoring duplicate device %q", name)
 			continue
 		}
 		devices = append(devices, name)
+		seen[name] = true
 	}
 
 	if len(devices) == 0 {
 		return nil, nil
 	}
 
-	if cfg.AcceptEnvvarUnprivileged || image.IsPrivileged((*image.OCISpec)(rawSpec)) {
-		return devices, nil
-	}
-
-	logger.Warningf("Ignoring devices specified in NVIDIA_VISIBLE_DEVICES: %v", devices)
-
-	return nil, nil
+	return devices, nil
 }
 
 // getAnnotationDevices returns a list of devices specified in the annotations.