[no-relnote] Use image.CUDA to extract visible devices

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

Author: elezar

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

@@ -6,7 +6,6 @@ import (
 	"log"
 	"os"
 	"path"
-	"path/filepath"
 
 	"github.com/opencontainers/runtime-spec/specs-go"
 	"golang.org/x/mod/semver"
@@ -15,11 +14,11 @@ import (
 )
 
 const (
-	envCUDAVersion          = "CUDA_VERSION"
-	envNVRequirePrefix      = "NVIDIA_REQUIRE_"
-	envNVRequireCUDA        = envNVRequirePrefix + "CUDA"
-	envNVDisableRequire     = "NVIDIA_DISABLE_REQUIRE"
-	envNVVisibleDevices     = "NVIDIA_VISIBLE_DEVICES"
+	envCUDAVersion      = "CUDA_VERSION"
+	envNVRequirePrefix  = "NVIDIA_REQUIRE_"
+	envNVRequireCUDA    = envNVRequirePrefix + "CUDA"
+	envNVDisableRequire = "NVIDIA_DISABLE_REQUIRE"
+
 	envNVMigConfigDevices   = "NVIDIA_MIG_CONFIG_DEVICES"
 	envNVMigMonitorDevices  = "NVIDIA_MIG_MONITOR_DEVICES"
 	envNVImexChannels       = "NVIDIA_IMEX_CHANNELS"
@@ -30,10 +29,6 @@ const (
 	capSysAdmin = "CAP_SYS_ADMIN"
 )
 
-const (
-	deviceListAsVolumeMountsRoot = "/var/run/nvidia-container-devices"
-)
-
 type nvidiaConfig struct {
 	Devices            []string
 	MigConfigDevices   string
@@ -76,23 +71,14 @@ type LinuxCapabilities struct {
 	Ambient     []string `json:"ambient,omitempty" platform:"linux"`
 }
 
-// Mount from OCI runtime spec
-// https://github.com/opencontainers/runtime-spec/blob/v1.0.0/specs-go/config.go#L103
-type Mount struct {
-	Destination string   `json:"destination"`
-	Type        string   `json:"type,omitempty" platform:"linux,solaris"`
-	Source      string   `json:"source,omitempty"`
-	Options     []string `json:"options,omitempty"`
-}
-
 // Spec from OCI runtime spec
 // We use pointers to structs, similarly to the latest version of runtime-spec:
 // https://github.com/opencontainers/runtime-spec/blob/v1.0.0/specs-go/config.go#L5-L28
 type Spec struct {
-	Version *string  `json:"ociVersion"`
-	Process *Process `json:"process,omitempty"`
-	Root    *Root    `json:"root,omitempty"`
-	Mounts  []Mount  `json:"mounts,omitempty"`
+	Version *string       `json:"ociVersion"`
+	Process *Process      `json:"process,omitempty"`
+	Root    *Root         `json:"root,omitempty"`
+	Mounts  []specs.Mount `json:"mounts,omitempty"`
 }
 
 // HookState holds state information about the hook
@@ -171,58 +157,22 @@ func isPrivileged(s *Spec) bool {
 	return image.IsPrivileged(&fullSpec)
 }
 
-func getDevicesFromEnvvar(image image.CUDA, swarmResourceEnvvars []string) []string {
+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 image.HasEnvvar(envvar) {
-			return image.DevicesFromEnvvars(swarmResourceEnvvars...).List()
+		if containerImage.HasEnvvar(envvar) {
+			return containerImage.DevicesFromEnvvars(swarmResourceEnvvars...).List()
 		}
 	}
 
-	return image.DevicesFromEnvvars(envNVVisibleDevices).List()
-}
-
-func getDevicesFromMounts(mounts []Mount) []string {
-	var devices []string
-	for _, m := range mounts {
-		root := filepath.Clean(deviceListAsVolumeMountsRoot)
-		source := filepath.Clean(m.Source)
-		destination := filepath.Clean(m.Destination)
-
-		// Only consider mounts who's host volume is /dev/null
-		if source != "/dev/null" {
-			continue
-		}
-		// Only consider container mount points that begin with 'root'
-		if len(destination) < len(root) {
-			continue
-		}
-		if destination[:len(root)] != root {
-			continue
-		}
-		// Grab the full path beyond 'root' and add it to the list of devices
-		device := destination[len(root):]
-		if len(device) > 0 && device[0] == '/' {
-			device = device[1:]
-		}
-		if len(device) == 0 {
-			continue
-		}
-		devices = append(devices, device)
-	}
-
-	if devices == nil {
-		return nil
-	}
-
-	return devices
+	return containerImage.DevicesFromEnvvars(image.NVIDIAVisibleDevicesEnvVar).List()
 }
 
-func getDevices(hookConfig *HookConfig, image image.CUDA, mounts []Mount, privileged bool) []string {
+func getDevices(hookConfig *HookConfig, image image.CUDA, privileged bool) []string {
 	// If enabled, try and get the device list from volume mounts first
 	if hookConfig.AcceptDeviceListAsVolumeMounts {
-		devices := getDevicesFromMounts(mounts)
+		devices := image.VisibleDevicesFromMounts()
 		if len(devices) > 0 {
 			return devices
 		}
@@ -294,10 +244,10 @@ func (c *HookConfig) getDriverCapabilities(cudaImage image.CUDA, legacyImage boo
 	return capabilities
 }
 
-func getNvidiaConfig(hookConfig *HookConfig, image image.CUDA, mounts []Mount, privileged bool) *nvidiaConfig {
+func getNvidiaConfig(hookConfig *HookConfig, image image.CUDA, privileged bool) *nvidiaConfig {
 	legacyImage := image.IsLegacy()
 
-	devices := getDevices(hookConfig, image, mounts, privileged)
+	devices := getDevices(hookConfig, image, privileged)
 	if len(devices) == 0 {
 		// empty devices means this is not a GPU container.
 		return nil
@@ -357,6 +307,7 @@ func getContainerConfig(hook HookConfig) (config containerConfig) {
 
 	image, err := image.New(
 		image.WithEnv(s.Process.Env),
+		image.WithMounts(s.Mounts),
 		image.WithDisableRequire(hook.DisableRequire),
 	)
 	if err != nil {
@@ -368,6 +319,6 @@ func getContainerConfig(hook HookConfig) (config containerConfig) {
 		Pid:    h.Pid,
 		Rootfs: s.Root.Path,
 		Image:  image,
-		Nvidia: getNvidiaConfig(&hook, image, s.Mounts, privileged),
+		Nvidia: getNvidiaConfig(&hook, image, privileged),
 	}
 }