mirror of
https://github.com/NVIDIA/nvidia-container-toolkit
synced 2024-11-22 08:18:32 +00:00
01b4381282
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>
132 lines
2.6 KiB
Go
132 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"github.com/stretchr/testify/require"
|
|
"sort"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestMergeVisibleDevicesEnvvars(t *testing.T) {
|
|
var tests = []struct {
|
|
name string
|
|
input []string
|
|
expected string
|
|
enableMerge bool
|
|
}{
|
|
{
|
|
"Simple Merge Enabled",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES_0=0,1",
|
|
"NVIDIA_VISIBLE_DEVICES_1=2,3",
|
|
"NVIDIA_VISIBLE_DEVICES_WHATEVER=4,5",
|
|
},
|
|
"0,1,2,3,4,5",
|
|
true,
|
|
},
|
|
{
|
|
"Simple Merge Disabled",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES_0=0,1",
|
|
"NVIDIA_VISIBLE_DEVICES_1=2,3",
|
|
"NVIDIA_VISIBLE_DEVICES_WHATEVER=4,5",
|
|
},
|
|
"",
|
|
false,
|
|
},
|
|
{
|
|
"Merge No Override (Enabled)",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES=all",
|
|
},
|
|
"all",
|
|
true,
|
|
},
|
|
{
|
|
"Merge No Override (Disabled)",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES=all",
|
|
},
|
|
"all",
|
|
false,
|
|
},
|
|
{
|
|
"Merge Override (Enabled, Before)",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES=all",
|
|
"NVIDIA_VISIBLE_DEVICES_0=0,1",
|
|
"NVIDIA_VISIBLE_DEVICES_1=2,3",
|
|
"NVIDIA_VISIBLE_DEVICES_WHATEVER=4,5",
|
|
},
|
|
"0,1,2,3,4,5",
|
|
true,
|
|
},
|
|
{
|
|
"Merge Override (Enabled, After)",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES_0=0,1",
|
|
"NVIDIA_VISIBLE_DEVICES_1=2,3",
|
|
"NVIDIA_VISIBLE_DEVICES_WHATEVER=4,5",
|
|
"NVIDIA_VISIBLE_DEVICES=all",
|
|
},
|
|
"0,1,2,3,4,5",
|
|
true,
|
|
},
|
|
{
|
|
"Merge Override (Enabled, In Between)",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES_0=0,1",
|
|
"NVIDIA_VISIBLE_DEVICES_1=2,3",
|
|
"NVIDIA_VISIBLE_DEVICES=all",
|
|
"NVIDIA_VISIBLE_DEVICES_WHATEVER=4,5",
|
|
},
|
|
"0,1,2,3,4,5",
|
|
true,
|
|
},
|
|
{
|
|
"Merge Override (Disabled, Before)",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES=all",
|
|
"NVIDIA_VISIBLE_DEVICES_0=0,1",
|
|
"NVIDIA_VISIBLE_DEVICES_1=2,3",
|
|
"NVIDIA_VISIBLE_DEVICES_WHATEVER=4,5",
|
|
},
|
|
"all",
|
|
false,
|
|
},
|
|
{
|
|
"Merge Override (Disabled, After)",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES_0=0,1",
|
|
"NVIDIA_VISIBLE_DEVICES_1=2,3",
|
|
"NVIDIA_VISIBLE_DEVICES_WHATEVER=4,5",
|
|
"NVIDIA_VISIBLE_DEVICES=all",
|
|
},
|
|
"all",
|
|
false,
|
|
},
|
|
{
|
|
"Merge Override (Disabled, In Between)",
|
|
[]string{
|
|
"NVIDIA_VISIBLE_DEVICES_0=0,1",
|
|
"NVIDIA_VISIBLE_DEVICES_1=2,3",
|
|
"NVIDIA_VISIBLE_DEVICES=all",
|
|
"NVIDIA_VISIBLE_DEVICES_WHATEVER=4,5",
|
|
},
|
|
"all",
|
|
false,
|
|
},
|
|
}
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
config := CLIConfig{
|
|
AlphaMergeVisibleDevicesEnvvars: tc.enableMerge,
|
|
}
|
|
envvars := getEnvMap(tc.input, config)
|
|
devices := strings.Split(envvars[envNVVisibleDevices], ",")
|
|
sort.Strings(devices)
|
|
require.Equal(t, tc.expected, strings.Join(devices, ","))
|
|
})
|
|
}
|
|
}
|