Fix NVIDIA_IMEX_CHANNELS handling on legacy images

For legacy images (images with a CUDA_VERSION set but no CUDA_REQUIRES set), the
default behaviour for device envvars is to treat non-existence as all.

This change ensures that the NVIDIA_IMEX_CHANNELS envvar is not treated in the same
way, instead returning no devices if the envvar is not set.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2024-11-14 13:26:27 -07:00
parent 1467f3f339
commit f774ceeedd
No known key found for this signature in database
2 changed files with 36 additions and 1 deletions

View File

@ -292,7 +292,11 @@ func (i CUDA) CDIDevicesFromMounts() []string {
// ImexChannelsFromEnvVar returns the list of IMEX channels requested for the image.
func (i CUDA) ImexChannelsFromEnvVar() []string {
return i.DevicesFromEnvvars(EnvVarNvidiaImexChannels).List()
imexChannels := i.DevicesFromEnvvars(EnvVarNvidiaImexChannels).List()
if len(imexChannels) == 1 && imexChannels[0] == "all" {
return nil
}
return imexChannels
}
// ImexChannelsFromMounts returns the list of IMEX channels requested for the image.

View File

@ -203,6 +203,37 @@ func TestGetVisibleDevicesFromMounts(t *testing.T) {
}
}
func TestImexChannelsFromEnvVar(t *testing.T) {
testCases := []struct {
description string
env []string
expected []string
}{
{
description: "no imex channels specified",
},
{
description: "imex channel specified",
env: []string{
"NVIDIA_IMEX_CHANNELS=3,4",
},
expected: []string{"3", "4"},
},
}
for _, tc := range testCases {
for id, baseEnvvars := range map[string][]string{"": nil, "legacy": {"CUDA_VERSION=1.2.3"}} {
t.Run(tc.description+id, func(t *testing.T) {
i, err := NewCUDAImageFromEnv(append(baseEnvvars, tc.env...))
require.NoError(t, err)
channels := i.ImexChannelsFromEnvVar()
require.EqualValues(t, tc.expected, channels)
})
}
}
}
func makeTestMounts(paths ...string) []specs.Mount {
var mounts []specs.Mount
for _, path := range paths {