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

Allow IMEX channels to be requested as volume mounts

Browse Source 
This change allows IMEX channels to be requested using the
volume mount mechanism.

A mount from /dev/null to /var/run/nvidia-container-devices/imex/{{ .ChannelID }}
is equivalent to including {{ .ChannelID }} in the NVIDIA_IMEX_CHANNELS
envvironment variables.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar
2024-10-14 14:32:06 +02:00
parent 92df542f2f
commit 2e6712d2bc
4 changed files with 46 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -30,7 +30,8 @@ import (
const (
DeviceListAsVolumeMountsRoot = "/var/run/nvidia-container-devices"
volumeMountDevicePrefixCDI = "cdi/"
volumeMountDevicePrefixCDI = "cdi/"
volumeMountDevicePrefixImex = "imex/"
)
// CUDA represents a CUDA image that can be used for GPU computing. This wraps
@@ -225,7 +226,10 @@ func (i CUDA) VisibleDevicesFromEnvVar() []string {
func (i CUDA) VisibleDevicesFromMounts() []string {
var devices []string
for _, device := range i.DevicesFromMounts() {
if strings.HasPrefix(device, volumeMountDevicePrefixCDI) {
switch {
case strings.HasPrefix(device, volumeMountDevicePrefixCDI):
continue
case strings.HasPrefix(device, volumeMountDevicePrefixImex):
continue
}
devices = append(devices, device)
@@ -286,6 +290,19 @@ func (i CUDA) CDIDevicesFromMounts() []string {
return devices
}
func (i CUDA) IsEnabled(envvar string) bool {
return i.Getenv(envvar) == "enabled"
// ImexChannelsFromEnvVar returns the list of IMEX channels requested for the image.
func (i CUDA) ImexChannelsFromEnvVar() []string {
return i.DevicesFromEnvvars(EnvVarNvidiaImexChannels).List()
}
// ImexChannelsFromMounts returns the list of IMEX channels requested for the image.
func (i CUDA) ImexChannelsFromMounts() []string {
var channels []string
for _, mountDevice := range i.DevicesFromMounts() {
if !strings.HasPrefix(mountDevice, volumeMountDevicePrefixImex) {
continue
}
channels = append(channels, strings.TrimPrefix(mountDevice, volumeMountDevicePrefixImex))
}
return channels
}

5
internal/config/image/cuda_image_test.go
View File

@@ -189,6 +189,11 @@ func TestGetVisibleDevicesFromMounts(t *testing.T) {
mounts: makeTestMounts("GPU0", "cdi/nvidia.com/gpu=all", "GPU1"),
expectedDevices: []string{"GPU0", "GPU1"},
},
{
description: "imex devices are ignored",
mounts: makeTestMounts("GPU0", "imex/0", "GPU1"),
expectedDevices: []string{"GPU0", "GPU1"},
},
}
for _, tc := range tests {
t.Run(tc.description, func(t *testing.T) {