Refactor IMEX channel requests from image

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

Author: elezar

cmd/nvidia-container-runtime-hook/container_config.go

@@ -148,30 +148,6 @@ func getMigDevices(image image.CUDA, envvar string) *string {
 	return &devices
 }
 
-func (hookConfig *hookConfig) getImexChannels(image image.CUDA, privileged bool) []string {
-	if hookConfig.Features.IgnoreImexChannelRequests.IsEnabled() {
-		return nil
-	}
-
-	// If enabled, try and get the device list from volume mounts first
-	if hookConfig.AcceptDeviceListAsVolumeMounts {
-		devices := image.ImexChannelsFromMounts()
-		if len(devices) > 0 {
-			return devices
-		}
-	}
-	devices := image.ImexChannelsFromEnvVar()
-	if len(devices) == 0 {
-		return nil
-	}
-
-	if privileged || hookConfig.AcceptEnvvarUnprivileged {
-		return devices
-	}
-
-	return nil
-}
-
 func (hookConfig *hookConfig) getDriverCapabilities(cudaImage image.CUDA, legacyImage bool) image.DriverCapabilities {
 	// We use the default driver capabilities by default. This is filtered to only include the
 	// supported capabilities
@@ -223,8 +199,6 @@ func (hookConfig *hookConfig) getNvidiaConfig(image image.CUDA, privileged bool)
 		log.Panicln("cannot set MIG_MONITOR_DEVICES in non privileged container")
 	}
 
-	imexChannels := hookConfig.getImexChannels(image, privileged)
-
 	driverCapabilities := hookConfig.getDriverCapabilities(image, legacyImage).String()
 
 	requirements, err := image.GetRequirements()
@@ -236,7 +210,7 @@ func (hookConfig *hookConfig) getNvidiaConfig(image image.CUDA, privileged bool)
 		Devices:            devices,
 		MigConfigDevices:   migConfigDevices,
 		MigMonitorDevices:  migMonitorDevices,
-		ImexChannels:       imexChannels,
+		ImexChannels:       image.ImexChannelRequests(),
 		DriverCapabilities: driverCapabilities,
 		Requirements:       requirements,
 	}
@@ -273,6 +247,7 @@ func (hookConfig *hookConfig) getContainerConfig() (config *containerConfig) {
 		image.WithAcceptDeviceListAsVolumeMounts(hookConfig.AcceptDeviceListAsVolumeMounts),
 		image.WithAcceptEnvvarUnprivileged(hookConfig.AcceptEnvvarUnprivileged),
 		image.WithPreferredVisibleDevicesEnvVars(hookConfig.getSwarmResourceEnvvars()...),
+		image.WithIgnoreImexChannelRequests(hookConfig.Features.IgnoreImexChannelRequests.IsEnabled()),
 	)
 	if err != nil {
 		log.Panicln(err)

internal/config/image/builder.go

@@ -127,6 +127,15 @@ func WithEnvMap(env map[string]string) Option {
 	}
 }
 
+// WithIgnoreImexChannelRequests sets whether per-container IMEX channel
+// requests are supported.
+func WithIgnoreImexChannelRequests(ignoreImexChannelRequests bool) Option {
+	return func(b *builder) error {
+		b.ignoreImexChannelRequests = ignoreImexChannelRequests
+		return nil
+	}
+}
+
 // WithLogger sets the logger to use when creating the CUDA image.
 func WithLogger(logger logger.Interface) Option {
 	return func(b *builder) error {

internal/config/image/cuda_image.go

@@ -51,6 +51,7 @@ type CUDA struct {
 	annotationsPrefixes            []string
 	acceptDeviceListAsVolumeMounts bool
 	acceptEnvvarUnprivileged       bool
+	ignoreImexChannelRequests      bool
 	preferredVisibleDeviceEnvVars  []string
 }
 
@@ -412,17 +413,51 @@ func (m cdiDeviceMountRequest) qualifiedName() (string, error) {
 	return fmt.Sprintf("%s/%s=%s", parts[0], parts[1], parts[2]), nil
 }
 
-// ImexChannelsFromEnvVar returns the list of IMEX channels requested for the image.
-func (i CUDA) ImexChannelsFromEnvVar() []string {
+func (i CUDA) ImexChannelRequests() []string {
+	if i.ignoreImexChannelRequests {
+		return nil
+	}
+
+	// If enabled, try and get the device list from volume mounts first
+	if i.acceptDeviceListAsVolumeMounts {
+		volumeMountDeviceRequests := i.imexChannelsFromMounts()
+		if len(volumeMountDeviceRequests) > 0 {
+			return volumeMountDeviceRequests
+		}
+	}
+
+	// Get the Fallback to reading from the environment variable if privileges are correct
+	envVarDeviceRequests := i.imexChannelsFromEnvVar()
+	if len(envVarDeviceRequests) == 0 {
+		return nil
+	}
+
+	// If the container is privileged, or environment variable requests are
+	// allowed for unprivileged containers, these devices are returned.
+	if i.isPrivileged || i.acceptEnvvarUnprivileged {
+		return envVarDeviceRequests
+	}
+
+	// We log a warning if we are ignoring the environment variable requests.
+	envVars := []string{EnvVarNvidiaImexChannels}
+	if len(envVars) > 0 {
+		i.logger.Warningf("Ignoring request by environment variable(s) in unprivileged container: %v", envVars)
+	}
+
+	return nil
+}
+
+// imexChannelsFromEnvVar returns the list of IMEX channels requested for the image.
+func (i CUDA) imexChannelsFromEnvVar() []string {
 	imexChannels := i.devicesFromEnvvars(EnvVarNvidiaImexChannels)
 	if len(imexChannels) == 1 && imexChannels[0] == "all" {
 		return nil
 	}
 	return imexChannels
 }
 
-// ImexChannelsFromMounts returns the list of IMEX channels requested for the image.
-func (i CUDA) ImexChannelsFromMounts() []string {
+// imexChannelsFromMounts returns the list of IMEX channels requested for the image.
+func (i CUDA) imexChannelsFromMounts() []string {
 	var channels []string
 	for _, mountDevice := range i.requestsFromMounts() {
 		if !strings.HasPrefix(mountDevice, volumeMountDevicePrefixImex) {

internal/config/image/cuda_image_test.go

@@ -716,7 +716,7 @@ func TestImexChannelsFromEnvVar(t *testing.T) {
 				i, err := newCUDAImageFromEnv(append(baseEnvvars, tc.env...))
 				require.NoError(t, err)
 
-				channels := i.ImexChannelsFromEnvVar()
+				channels := i.imexChannelsFromEnvVar()
 				require.EqualValues(t, tc.expected, channels)
 			})
 		}