Unify GetDevices logic at internal/config/image

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

Author: ArangoGutierrez

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

@@ -81,7 +81,9 @@ type HookState struct {
 	BundlePath string `json:"bundlePath"`
 }
 
-func loadSpec(path string) (spec *Spec) {
+func loadSpec(path string) *specs.Spec {
+	var spec Spec
+
 	f, err := os.Open(path)
 	if err != nil {
 		log.Panicln("could not open OCI spec:", err)
@@ -100,85 +102,57 @@ func loadSpec(path string) (spec *Spec) {
 	if spec.Root == nil {
 		log.Panicln("Root is empty in OCI spec")
 	}
-	return
-}
 
-func isPrivileged(s *Spec) bool {
-	if s.Process.Capabilities == nil {
-		return false
+	process := specs.Process{
+		Env: spec.Process.Env,
 	}
-
 	var caps []string
 	// If v1.0.0-rc1 <= OCI version < v1.0.0-rc5 parse s.Process.Capabilities as:
 	// github.com/opencontainers/runtime-spec/blob/v1.0.0-rc1/specs-go/config.go#L30-L54
-	rc1cmp := semver.Compare("v"+*s.Version, "v1.0.0-rc1")
-	rc5cmp := semver.Compare("v"+*s.Version, "v1.0.0-rc5")
+	rc1cmp := semver.Compare("v"+*spec.Version, "v1.0.0-rc1")
+	rc5cmp := semver.Compare("v"+*spec.Version, "v1.0.0-rc5")
 	if (rc1cmp == 1 || rc1cmp == 0) && (rc5cmp == -1) {
-		err := json.Unmarshal(*s.Process.Capabilities, &caps)
+		err := json.Unmarshal(*spec.Process.Capabilities, &caps)
 		if err != nil {
 			log.Panicln("could not decode Process.Capabilities in OCI spec:", err)
 		}
 		for _, c := range caps {
 			if c == capSysAdmin {
-				return true
+				process.Capabilities = &specs.LinuxCapabilities{
+					Bounding: caps,
+				}
+				break
 			}
 		}
-		return false
-	}
-
-	// Otherwise, parse s.Process.Capabilities as:
-	// github.com/opencontainers/runtime-spec/blob/v1.0.0/specs-go/config.go#L30-L54
-	process := specs.Process{
-		Env: s.Process.Env,
-	}
-
-	err := json.Unmarshal(*s.Process.Capabilities, &process.Capabilities)
-	if err != nil {
-		log.Panicln("could not decode Process.Capabilities in OCI spec:", err)
-	}
-
-	fullSpec := specs.Spec{
-		Version: *s.Version,
-		Process: &process,
-	}
-
-	return image.IsPrivileged(&fullSpec)
-}
-
-func getDevicesFromEnvvar(containerImage image.CUDA, swarmResourceEnvvars []string) []string {
-	// We check if the image has at least one of the Swarm resource envvars defined and use this
-	// if specified.
-	for _, envvar := range swarmResourceEnvvars {
-		if containerImage.HasEnvvar(envvar) {
-			return containerImage.DevicesFromEnvvars(swarmResourceEnvvars...).List()
+	} else {
+		// Otherwise, parse s.Process.Capabilities as:
+		// github.com/opencontainers/runtime-spec/blob/v1.0.0/specs-go/config.go#L30-L54
+		err := json.Unmarshal(*spec.Process.Capabilities, &process.Capabilities)
+		if err != nil {
+			log.Panicln("could not decode Process.Capabilities in OCI spec:", err)
 		}
 	}
 
-	return containerImage.VisibleDevicesFromEnvVar()
-}
+	root := specs.Root{
+		Path: spec.Root.Path,
+	}
 
-func (hookConfig *hookConfig) getDevices(image image.CUDA, privileged bool) []string {
-	// If enabled, try and get the device list from volume mounts first
-	if hookConfig.AcceptDeviceListAsVolumeMounts {
-		devices := image.VisibleDevicesFromMounts()
-		if len(devices) > 0 {
-			return devices
+	mounts := make([]specs.Mount, len(spec.Mounts))
+	for i, m := range spec.Mounts {
+		mounts[i] = specs.Mount{
+			Source:      m.Source,
+			Destination: m.Destination,
+			Type:        m.Type,
+			Options:     m.Options,
 		}
 	}
 
-	// Fallback to reading from the environment variable if privileges are correct
-	devices := getDevicesFromEnvvar(image, hookConfig.getSwarmResourceEnvvars())
-	if len(devices) == 0 {
-		return nil
-	}
-	if privileged || hookConfig.AcceptEnvvarUnprivileged {
-		return devices
+	return &specs.Spec{
+		Version: *spec.Version,
+		Process: &process,
+		Root:    &root,
+		Mounts:  mounts,
 	}
-
-	configName := hookConfig.getConfigOption("AcceptEnvvarUnprivileged")
-	log.Printf("Ignoring devices specified in NVIDIA_VISIBLE_DEVICES (privileged=%v, %v=%v) ", privileged, configName, hookConfig.AcceptEnvvarUnprivileged)
-
-	return nil
 }
 
 func getMigConfigDevices(i image.CUDA) *string {
@@ -225,7 +199,6 @@ func (hookConfig *hookConfig) getDriverCapabilities(cudaImage image.CUDA, legacy
 	// We use the default driver capabilities by default. This is filtered to only include the
 	// supported capabilities
 	supportedDriverCapabilities := image.NewDriverCapabilities(hookConfig.SupportedDriverCapabilities)
-
 	capabilities := supportedDriverCapabilities.Intersection(image.DefaultDriverCapabilities)
 
 	capsEnvSpecified := cudaImage.HasEnvvar(image.EnvVarNvidiaDriverCapabilities)
@@ -251,7 +224,7 @@ func (hookConfig *hookConfig) getDriverCapabilities(cudaImage image.CUDA, legacy
 func (hookConfig *hookConfig) getNvidiaConfig(image image.CUDA, privileged bool) *nvidiaConfig {
 	legacyImage := image.IsLegacy()
 
-	devices := hookConfig.getDevices(image, privileged)
+	devices := image.GetDevices(hookConfig.AcceptDeviceListAsVolumeMounts, hookConfig.AcceptEnvvarUnprivileged)
 	if len(devices) == 0 {
 		// empty devices means this is not a GPU container.
 		return nil
@@ -305,21 +278,26 @@ func (hookConfig *hookConfig) getContainerConfig() (config containerConfig) {
 	}
 
 	s := loadSpec(path.Join(b, "config.json"))
-
-	image, err := image.New(
+	opts := []image.Option{
 		image.WithEnv(s.Process.Env),
 		image.WithMounts(s.Mounts),
+		image.WithSpec(s),
 		image.WithDisableRequire(hookConfig.DisableRequire),
-	)
+	}
+
+	if len(hookConfig.getSwarmResourceEnvvars()) > 0 {
+		opts = append(opts, image.WithSwarmResource(hookConfig.getSwarmResourceEnvvars()...))
+	}
+
+	i, err := image.New(opts...)
 	if err != nil {
 		log.Panicln(err)
 	}
 
-	privileged := isPrivileged(s)
 	return containerConfig{
 		Pid:    h.Pid,
 		Rootfs: s.Root.Path,
-		Image:  image,
-		Nvidia: hookConfig.getNvidiaConfig(image, privileged),
+		Image:  i,
+		Nvidia: hookConfig.getNvidiaConfig(i, i.IsPrivileged()),
 	}
 }

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

@@ -477,9 +477,18 @@ func TestGetNvidiaConfig(t *testing.T) {
 	}
 	for _, tc := range tests {
 		t.Run(tc.description, func(t *testing.T) {
-			image, _ := image.New(
+			opts := []image.Option{
 				image.WithEnvMap(tc.env),
-			)
+				image.WithPrivileged(tc.privileged),
+			}
+
+			if tc.hookConfig != nil {
+				if tc.hookConfig.SwarmResource != "" {
+					opts = append(opts, image.WithSwarmResource(tc.hookConfig.SwarmResource))
+				}
+			}
+			image, _ := image.New(opts...)
+
 			// Wrap the call to getNvidiaConfig() in a closure.
 			var cfg *nvidiaConfig
 			getConfig := func() {
@@ -622,12 +631,13 @@ func TestDeviceListSourcePriority(t *testing.T) {
 						},
 					),
 					image.WithMounts(tc.mountDevices),
+					image.WithPrivileged(tc.privileged),
 				)
 				defaultConfig, _ := config.GetDefault()
 				cfg := &hookConfig{defaultConfig}
 				cfg.AcceptEnvvarUnprivileged = tc.acceptUnprivileged
 				cfg.AcceptDeviceListAsVolumeMounts = tc.acceptMounts
-				devices = cfg.getDevices(image, tc.privileged)
+				devices = image.GetDevices(tc.acceptMounts, tc.acceptUnprivileged)
 			}
 
 			// For all other tests, just grab the devices and check the results
@@ -843,10 +853,18 @@ func TestGetDevicesFromEnvvar(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.description, func(t *testing.T) {
-			image, _ := image.New(
+			opts := []image.Option{
 				image.WithEnvMap(tc.env),
-			)
-			devices := getDevicesFromEnvvar(image, tc.swarmResourceEnvvars)
+				image.WithPrivileged(true),
+			}
+
+			if len(tc.swarmResourceEnvvars) > 0 {
+				opts = append(opts, image.WithSwarmResource(tc.swarmResourceEnvvars...))
+			}
+
+			image, _ := image.New(opts...)
+
+			devices := image.GetDevices(false, false)
 			require.EqualValues(t, tc.expectedDevices, devices)
 		})
 	}