Add ability to merge envars of the form NVIDIA_VISIBLE_DEVICES_*

This allows someone to (for example) pass the following environment
variables:

NVIDIA_VISIBLE_DEVICES_0="0,1"
NVIDIA_VISIBLE_DEVICES_1="2,3"
NVIDIA_VISIBLE_DEVICES_WHATEVER="4,5"

and have the nvidia-container-toolkit automatically merge these into:

NVIDIA_VISIBLE_DEVICES="0,1,2,3,4,5"

This is useful (for example) if the full list of devices comes
from multiple, disparate sources.

Note: This will override whatever the original value of
NVIDIA_VISIBLE_DEVICES was (*excluding* its original value) if it also
exists as an environment variable already. We exclude the original value
to ensure that we have a way to override the default value of
NVIDIA_VISIBLE_DEVICES set to "all" inside a container image.

Signed-off-by: Kevin Klues <kklues@nvidia.com>

pkg/container_config.go

@@ -82,7 +82,7 @@ func parseCudaVersion(cudaVersion string) (vmaj, vmin, vpatch uint32) {
 	return
 }
 
-func getEnvMap(e []string) (m map[string]string) {
+func getEnvMap(e []string, config CLIConfig) (m map[string]string) {
 	m = make(map[string]string)
 	for _, s := range e {
 		p := strings.SplitN(s, "=", 2)
@@ -91,6 +91,17 @@ func getEnvMap(e []string) (m map[string]string) {
 		}
 		m[p[0]] = p[1]
 	}
+	if config.AlphaMergeVisibleDevicesEnvvars {
+		var mergable []string
+		for k, v := range m {
+			if strings.HasPrefix(k, envNVVisibleDevices+"_") {
+				mergable = append(mergable, v)
+			}
+		}
+		if len(mergable) > 0 {
+			m[envNVVisibleDevices] = strings.Join(mergable, ",")
+		}
+	}
 	return
 }
 
@@ -254,7 +265,7 @@ func getContainerConfig(hook HookConfig) (config containerConfig) {
 
 	s := loadSpec(path.Join(b, "config.json"))
 
-	env := getEnvMap(s.Process.Env)
+	env := getEnvMap(s.Process.Env, hook.NvidiaContainerCLI)
 	envSwarmGPU = hook.SwarmResource
 	return containerConfig{
 		Pid:    h.Pid,

pkg/hook_config.go

@@ -20,16 +20,17 @@ var defaultPaths = [...]string{
 
 // CLIConfig: options for nvidia-container-cli.
 type CLIConfig struct {
-	Root        *string  `toml:"root"`
-	Path        *string  `toml:"path"`
-	Environment []string `toml:"environment"`
-	Debug       *string  `toml:"debug"`
-	Ldcache     *string  `toml:"ldcache"`
-	LoadKmods   bool     `toml:"load-kmods"`
-	NoPivot     bool     `toml:"no-pivot"`
-	NoCgroups   bool     `toml:"no-cgroups"`
-	User        *string  `toml:"user"`
-	Ldconfig    *string  `toml:"ldconfig"`
+	Root                            *string  `toml:"root"`
+	Path                            *string  `toml:"path"`
+	Environment                     []string `toml:"environment"`
+	Debug                           *string  `toml:"debug"`
+	Ldcache                         *string  `toml:"ldcache"`
+	LoadKmods                       bool     `toml:"load-kmods"`
+	NoPivot                         bool     `toml:"no-pivot"`
+	NoCgroups                       bool     `toml:"no-cgroups"`
+	User                            *string  `toml:"user"`
+	Ldconfig                        *string  `toml:"ldconfig"`
+	AlphaMergeVisibleDevicesEnvvars bool     `toml:"alpha-merge-visible-devices-envvars"`
 }
 
 type HookConfig struct {
@@ -44,16 +45,17 @@ func getDefaultHookConfig() (config HookConfig) {
 		DisableRequire: false,
 		SwarmResource:  nil,
 		NvidiaContainerCLI: CLIConfig{
-			Root:        nil,
-			Path:        nil,
-			Environment: []string{},
-			Debug:       nil,
-			Ldcache:     nil,
-			LoadKmods:   true,
-			NoPivot:     false,
-			NoCgroups:   false,
-			User:        nil,
-			Ldconfig:    nil,
+			Root:                            nil,
+			Path:                            nil,
+			Environment:                     []string{},
+			Debug:                           nil,
+			Ldcache:                         nil,
+			LoadKmods:                       true,
+			NoPivot:                         false,
+			NoCgroups:                       false,
+			User:                            nil,
+			Ldconfig:                        nil,
+			AlphaMergeVisibleDevicesEnvvars: false,
 		},
 	}
 }