Evan's comments

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

Author: elezar

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

@@ -258,17 +258,15 @@ func (hookConfig *hookConfig) getContainerConfig() (config containerConfig) {
 
 	privileged := isPrivileged(s)
 
-	opts := []image.Option{
+	i, err := image.New(
 		image.WithEnv(s.Process.Env),
 		image.WithMounts(s.Mounts),
 		image.WithPrivileged(privileged),
 		image.WithDisableRequire(hookConfig.DisableRequire),
 		image.WithAcceptDeviceListAsVolumeMounts(hookConfig.AcceptDeviceListAsVolumeMounts),
 		image.WithAcceptEnvvarUnprivileged(hookConfig.AcceptEnvvarUnprivileged),
 		image.WithAdditionalVisibleDevicesEnvVars(hookConfig.getSwarmResourceEnvvars()...),
-	}
-
-	i, err := image.New(opts...)
+	)
 	if err != nil {
 		log.Panicln(err)
 	}

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

@@ -477,18 +477,11 @@ func TestGetNvidiaConfig(t *testing.T) {
 	}
 	for _, tc := range tests {
 		t.Run(tc.description, func(t *testing.T) {
-			opts := []image.Option{
+			image, _ := image.New(
 				image.WithEnvMap(tc.env),
 				image.WithPrivileged(tc.privileged),
-				image.WithAcceptEnvvarUnprivileged(true),
-			}
-
-			if tc.hookConfig != nil {
-				if tc.hookConfig.SwarmResource != "" {
-					opts = append(opts, image.WithAdditionalVisibleDevicesEnvVars(tc.hookConfig.SwarmResource))
-				}
-			}
-			image, _ := image.New(opts...)
+				image.WithAdditionalVisibleDevicesEnvVars(tc.hookConfig.getSwarmResourceEnvvars()...),
+			)
 
 			// Wrap the call to getNvidiaConfig() in a closure.
 			var cfg *nvidiaConfig
@@ -852,15 +845,13 @@ func TestGetDevicesFromEnvvar(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.description, func(t *testing.T) {
-			opts := []image.Option{
+			image, _ := image.New(
 				image.WithEnvMap(tc.env),
 				image.WithPrivileged(true),
 				image.WithAcceptDeviceListAsVolumeMounts(false),
 				image.WithAcceptEnvvarUnprivileged(false),
 				image.WithAdditionalVisibleDevicesEnvVars(tc.swarmResourceEnvvars...),
-			}
-
-			image, _ := image.New(opts...)
+			)
 
 			devices := image.VisibleDevices()
 			require.EqualValues(t, tc.expectedDevices, devices)

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

@@ -88,7 +88,7 @@ func (c hookConfig) getConfigOption(fieldName string) string {
 
 // getSwarmResourceEnvvars returns the swarm resource envvars for the config.
 func (c *hookConfig) getSwarmResourceEnvvars() []string {
-	if c.SwarmResource == "" {
+	if c == nil || c.SwarmResource == "" {
 		return nil
 	}
 

internal/config/image/builder.go

@@ -21,6 +21,8 @@ import (
 	"strings"
 
 	"github.com/opencontainers/runtime-spec/specs-go"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 )
 
 type builder struct {
@@ -34,12 +36,21 @@ type Option func(*builder) error
 
 // New creates a new CUDA image from the input options.
 func New(opt ...Option) (CUDA, error) {
-	b := &builder{}
+	b := &builder{
+		CUDA: CUDA{
+			acceptEnvvarUnprivileged: true,
+		},
+	}
 	for _, o := range opt {
 		if err := o(b); err != nil {
 			return CUDA{}, err
 		}
 	}
+
+	if b.logger == nil {
+		b.logger = logger.New()
+	}
+
 	if b.env == nil {
 		b.env = make(map[string]string)
 	}
@@ -103,6 +114,14 @@ func WithEnvMap(env map[string]string) Option {
 	}
 }
 
+// WithLogger sets the logger to use when creating the CUDA image.
+func WithLogger(logger logger.Interface) Option {
+	return func(b *builder) error {
+		b.logger = logger
+		return nil
+	}
+}
+
 // WithMounts sets the mounts associated with the CUDA image.
 func WithMounts(mounts []specs.Mount) Option {
 	return func(b *builder) error {
@@ -126,16 +145,7 @@ func WithPrivileged(isPrivileged bool) Option {
 //   - Concatenate all of this
 func WithAdditionalVisibleDevicesEnvVars(visibleDevicesEnvVars ...string) Option {
 	return func(b *builder) error {
-		var result []string
-		for _, v := range visibleDevicesEnvVars {
-			for _, c := range strings.Split(v, ",") {
-				trimmed := strings.TrimSpace(c)
-				if trimmed != "" {
-					result = append(result, trimmed)
-				}
-			}
-		}
-		b.additionalVisibleDevicesEnvVars = result
+		b.additionalVisibleDevicesEnvVars = visibleDevicesEnvVars
 		return nil
 	}
 }

internal/config/image/builder_test.go

@@ -25,6 +25,8 @@ import (
 	"github.com/opencontainers/runtime-spec/specs-go"
 	"golang.org/x/mod/semver"
 	"tags.cncf.io/container-device-interface/pkg/parser"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 )
 
 const (
@@ -38,30 +40,33 @@ const (
 // a map of environment variable to values that can be used to perform lookups
 // such as requirements.
 type CUDA struct {
-	additionalVisibleDevicesEnvVars []string
-
-	env map[string]string
+	logger logger.Interface
 
+	env    map[string]string
 	mounts []specs.Mount
 
-	acceptDeviceListAsVolumeMounts bool
-	acceptEnvvarUnprivileged       bool
-	isPrivileged                   bool
+	isPrivileged bool
+
+	additionalVisibleDevicesEnvVars []string
+	acceptDeviceListAsVolumeMounts  bool
+	acceptEnvvarUnprivileged        bool
 }
 
 // NewCUDAImageFromSpec creates a CUDA image from the input OCI runtime spec.
 // The process environment is read (if present) to construc the CUDA Image.
-func NewCUDAImageFromSpec(spec *specs.Spec) (CUDA, error) {
+func NewCUDAImageFromSpec(spec *specs.Spec, opts ...Option) (CUDA, error) {
 	var env []string
 	if spec != nil && spec.Process != nil {
 		env = spec.Process.Env
 	}
 
-	return New(
+	specOpts := []Option{
 		WithEnv(env),
 		WithMounts(spec.Mounts),
 		WithPrivileged(IsPrivileged((*OCISpec)(spec))),
-	)
+	}
+
+	return New(append(opts, specOpts...)...)
 }
 
 // newCUDAImageFromEnv creates a CUDA image from the input environment. The environment
@@ -163,7 +168,7 @@ func (i CUDA) DevicesFromEnvvars(envVars ...string) VisibleDevices {
 
 // GetDriverCapabilities returns the requested driver capabilities.
 func (i CUDA) GetDriverCapabilities() DriverCapabilities {
-	env := i.Getenv(EnvVarNvidiaDriverCapabilities)
+	env := i.env[EnvVarNvidiaDriverCapabilities]
 
 	capabilities := make(DriverCapabilities)
 	for _, c := range strings.Split(env, ",") {
@@ -174,7 +179,7 @@ func (i CUDA) GetDriverCapabilities() DriverCapabilities {
 }
 
 func (i CUDA) legacyVersion() (string, error) {
-	cudaVersion := i.Getenv(EnvVarCudaVersion)
+	cudaVersion := i.env[EnvVarCudaVersion]
 	majorMinor, err := parseMajorMinorVersion(cudaVersion)
 	if err != nil {
 		return "", fmt.Errorf("invalid CUDA version %v: %v", cudaVersion, err)
@@ -224,9 +229,42 @@ func (i CUDA) OnlyFullyQualifiedCDIDevices() bool {
 	return hasCDIdevice
 }
 
-// VisibleDevicesFromEnvVar returns the set of visible devices requested through
-// the NVIDIA_VISIBLE_DEVICES environment variable or any variables specified
-// in visibleDevicesEnvVars.
+// VisibleDevices returns a list of devices requested in the container image.
+// If volume mount requests are enabled these are returned if requested,
+// otherwise device requests through environment variables are considered.
+// In cases where environment variable requests required privileged containers,
+// such devices requests are ignored.
+func (i CUDA) VisibleDevices() []string {
+	// If enabled, try and get the device list from volume mounts first
+	if i.acceptDeviceListAsVolumeMounts {
+		volumeMountDeviceRequests := i.visibleDevicesFromMounts()
+		if len(volumeMountDeviceRequests) > 0 {
+			return volumeMountDeviceRequests
+		}
+	}
+
+	// Get the Fallback to reading from the environment variable if privileges are correct
+	envVarDeviceRequests := i.VisibleDevicesFromEnvVar()
+	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.
+	i.logger.Warningf("Ignoring devices specified in NVIDIA_VISIBLE_DEVICES in unprivileged container")
+
+	return nil
+}
+
+// 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 caes, the
+// NVIDIA_VISIBLE_DEVICES environment variable is used.
 func (i CUDA) VisibleDevicesFromEnvVar() []string {
 	for _, envVar := range i.additionalVisibleDevicesEnvVars {
 		if i.HasEnvvar(envVar) {
@@ -284,28 +322,6 @@ func (i CUDA) DevicesFromMounts() []string {
 	return devices
 }
 
-func (i CUDA) VisibleDevices() []string {
-	// If enabled, try and get the device list from volume mounts first
-	if i.acceptDeviceListAsVolumeMounts {
-		devices := i.visibleDevicesFromMounts()
-		if len(devices) > 0 {
-			return devices
-		}
-	}
-
-	// Fallback to reading from the environment variable if privileges are correct
-	devices := i.VisibleDevicesFromEnvVar()
-	if len(devices) == 0 {
-		return nil
-	}
-
-	if i.isPrivileged || i.acceptEnvvarUnprivileged {
-		return devices
-	}
-
-	return nil
-}
-
 // CDIDevicesFromMounts returns a list of CDI devices specified as mounts on the image.
 func (i CUDA) CDIDevicesFromMounts() []string {
 	var devices []string

internal/config/image/cuda_image.go

@@ -16,7 +16,9 @@
 
 package image
 
-import "github.com/opencontainers/runtime-spec/specs-go"
+import (
+	"github.com/opencontainers/runtime-spec/specs-go"
+)
 
 const (
 	capSysAdmin = "CAP_SYS_ADMIN"

internal/config/image/privileged.go

@@ -79,7 +79,10 @@ func getDevicesFromSpec(logger logger.Interface, ociSpec oci.Spec, cfg *config.C
 		return annotationDevices, nil
 	}
 
-	container, err := image.NewCUDAImageFromSpec(rawSpec)
+	container, err := image.NewCUDAImageFromSpec(
+		rawSpec,
+		image.WithLogger(logger),
+	)
 	if err != nil {
 		return nil, err
 	}

internal/modifier/cdi.go

@@ -92,7 +92,6 @@ func TestGraphicsModifier(t *testing.T) {
 		t.Run(tc.description, func(t *testing.T) {
 			image, _ := image.New(
 				image.WithEnvMap(tc.envmap),
-				image.WithAcceptEnvvarUnprivileged(true),
 			)
 			required, _ := requiresGraphicsModifier(image)
 			require.EqualValues(t, tc.expectedRequired, required)

internal/modifier/graphics_test.go

@@ -70,7 +70,10 @@ func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Sp
 		return nil, fmt.Errorf("failed to load OCI spec: %v", err)
 	}
 
-	image, err := image.NewCUDAImageFromSpec(rawSpec)
+	image, err := image.NewCUDAImageFromSpec(
+		rawSpec,
+		image.WithLogger(logger),
+	)
 	if err != nil {
 		return nil, err
 	}