diff --git a/internal/config/image/cuda_image.go b/internal/config/image/cuda_image.go index 51ba2f891..2b9254cfa 100644 --- a/internal/config/image/cuda_image.go +++ b/internal/config/image/cuda_image.go @@ -180,14 +180,7 @@ func (i CUDA) devicesFromEnvvars(envVars ...string) []string { // GetDriverCapabilities returns the requested driver capabilities. func (i CUDA) GetDriverCapabilities() DriverCapabilities { - env := i.env[EnvVarNvidiaDriverCapabilities] - - capabilities := make(DriverCapabilities) - for c := range strings.SplitSeq(env, ",") { - capabilities[DriverCapability(c)] = true - } - - return capabilities + return NewDriverCapabilities(i.env[EnvVarNvidiaDriverCapabilities]) } func (i CUDA) legacyVersion() (string, error) { diff --git a/internal/config/image/cuda_image_test.go b/internal/config/image/cuda_image_test.go index 0000e7129..3b8a58f4f 100644 --- a/internal/config/image/cuda_image_test.go +++ b/internal/config/image/cuda_image_test.go @@ -216,6 +216,56 @@ func TestGetRequirements(t *testing.T) { } } +func TestGetDriverCapabilities(t *testing.T) { + testCases := []struct { + description string + env map[string]string + expected []string + }{ + { + description: "envvar unset", + expected: nil, + }, + { + description: "envvar empty", + env: map[string]string{ + EnvVarNvidiaDriverCapabilities: "", + }, + expected: nil, + }, + { + description: "comma-separated capabilities", + env: map[string]string{ + EnvVarNvidiaDriverCapabilities: "compute,graphics", + }, + expected: []string{"compute", "graphics"}, + }, + { + description: "whitespace around capabilities is trimmed", + env: map[string]string{ + EnvVarNvidiaDriverCapabilities: " compute , graphics ", + }, + expected: []string{"compute", "graphics"}, + }, + { + description: "empty entries are ignored", + env: map[string]string{ + EnvVarNvidiaDriverCapabilities: "compute,,utility,", + }, + expected: []string{"compute", "utility"}, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + image, err := New(WithEnvMap(tc.env)) + require.NoError(t, err) + + require.EqualValues(t, tc.expected, image.GetDriverCapabilities().List()) + }) + } +} + func TestGetDevicesFromEnvvar(t *testing.T) { envDockerResourceGPUs := "DOCKER_RESOURCE_GPUS" gpuID := "GPU-12345" diff --git a/internal/modifier/graphics_test.go b/internal/modifier/graphics_test.go index 540fee5a3..dccc16c5d 100644 --- a/internal/modifier/graphics_test.go +++ b/internal/modifier/graphics_test.go @@ -70,6 +70,14 @@ func TestGraphicsModifier(t *testing.T) { }, expectedDevices: []string{"all"}, }, + { + description: "devices with whitespace around graphics capability creates modifier", + envmap: map[string]string{ + "NVIDIA_VISIBLE_DEVICES": "all", + "NVIDIA_DRIVER_CAPABILITIES": "compute, graphics", + }, + expectedDevices: []string{"all"}, + }, { description: "devices with display capability creates modifier", envmap: map[string]string{