UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-06-26 18:18:24 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

[no-relnote] new helper func visibleEnvVars

Browse Source 
Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Carlos Eduardo Arango Gutierrez 2025-06-05 13:21:25 +02:00 committed by Evan Lezar
parent 8149be09ac
commit 5fe7b06514
No known key found for this signature in database
1 changed files with 22 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
internal/config/image/cuda_image.go
View File

@ -232,6 +232,22 @@ func (i CUDA) OnlyFullyQualifiedCDIDevices() bool {
return hasCDIdevice
}
// visibleEnvVars returns the environment variables that are used to determine device visibility.
// It returns the preferred environment variables that are set, or NVIDIA_VISIBLE_DEVICES if none are set.
func (i CUDA) visibleEnvVars() []string {
var envVars []string
for _, envVar := range i.preferredVisibleDeviceEnvVars {
if !i.HasEnvvar(envVar) {
continue
}
envVars = append(envVars, envVar)
}
if len(envVars) > 0 {
return envVars
}
return []string{EnvVarNvidiaVisibleDevices}
}
// VisibleDevices returns a list of devices requested in the container image.
// If volume mount requests are enabled these are returned if requested,
// otherwise device requests through environment variables are considered.
@ -265,7 +281,10 @@ func (i CUDA) VisibleDevices() []string {
}
// We log a warning if we are ignoring the environment variable requests.
i.logger.Warningf("Ignoring devices specified in NVIDIA_VISIBLE_DEVICES in unprivileged container")
envVars := i.visibleEnvVars()
if len(envVars) > 0 {
i.logger.Warningf("Ignoring devices requested by environment variable(s) in unprivileged container: %v", envVars)
}
return nil
}
@ -300,12 +319,8 @@ func (i CUDA) cdiDeviceRequestsFromAnnotations() []string {
// are used to determine the visible devices. If this is not the case, the
// NVIDIA_VISIBLE_DEVICES environment variable is used.
func (i CUDA) VisibleDevicesFromEnvVar() []string {
for _, envVar := range i.preferredVisibleDeviceEnvVars {
if i.HasEnvvar(envVar) {
return i.DevicesFromEnvvars(i.preferredVisibleDeviceEnvVars...).List()
}
}
return i.DevicesFromEnvvars(EnvVarNvidiaVisibleDevices).List()
envVars := i.visibleEnvVars()
return i.DevicesFromEnvvars(envVars...).List()
}
// visibleDevicesFromMounts returns the set of visible devices requested as mounts.