UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2024-11-22 00:08:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add support for an NVIDIA_IMEX_CHANNELS envvar

Browse Source 
Signed-off-by: Kevin Klues <kklues@nvidia.com>
This commit is contained in:
Kevin Klues 2024-01-17 22:38:10 +00:00 committed by Evan Lezar
parent d57d83405d
commit 2fa37973e0
2 changed files with 20 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
cmd/nvidia-container-runtime-hook/container_config.go
View File

@ -9,9 +9,10 @@ import (
"path/filepath" "path/filepath"
"strings" "strings"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
"github.com/opencontainers/runtime-spec/specs-go" "github.com/opencontainers/runtime-spec/specs-go"
"golang.org/x/mod/semver" "golang.org/x/mod/semver"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
) )
const ( const (
@ -22,6 +23,7 @@ const (
envNVVisibleDevices = "NVIDIA_VISIBLE_DEVICES" envNVVisibleDevices = "NVIDIA_VISIBLE_DEVICES"
envNVMigConfigDevices = "NVIDIA_MIG_CONFIG_DEVICES" envNVMigConfigDevices = "NVIDIA_MIG_CONFIG_DEVICES"
envNVMigMonitorDevices = "NVIDIA_MIG_MONITOR_DEVICES" envNVMigMonitorDevices = "NVIDIA_MIG_MONITOR_DEVICES"
envNVImexChannels = "NVIDIA_IMEX_CHANNELS"
envNVDriverCapabilities = "NVIDIA_DRIVER_CAPABILITIES" envNVDriverCapabilities = "NVIDIA_DRIVER_CAPABILITIES"
) )
@ -37,6 +39,7 @@ type nvidiaConfig struct {
Devices string Devices string
MigConfigDevices string MigConfigDevices string
MigMonitorDevices string MigMonitorDevices string
ImexChannels string
DriverCapabilities string DriverCapabilities string
// Requirements defines the requirements DSL for the container to run. // Requirements defines the requirements DSL for the container to run.
// This is empty if no specific requirements are needed, or if requirements are // This is empty if no specific requirements are needed, or if requirements are
@ -271,6 +274,13 @@ func getMigMonitorDevices(env map[string]string) *string {
return nil return nil
} }
func getImexChannels(env map[string]string) *string {
if chans, ok := env[envNVImexChannels]; ok {
return &chans
}
return nil
}
func (c *HookConfig) getDriverCapabilities(env map[string]string, legacyImage bool) image.DriverCapabilities { func (c *HookConfig) getDriverCapabilities(env map[string]string, legacyImage bool) image.DriverCapabilities {
// We use the default driver capabilities by default. This is filtered to only include the // We use the default driver capabilities by default. This is filtered to only include the
// supported capabilities // supported capabilities
@ -324,6 +334,11 @@ func getNvidiaConfig(hookConfig *HookConfig, image image.CUDA, mounts []Mount, p
log.Panicln("cannot set MIG_MONITOR_DEVICES in non privileged container") log.Panicln("cannot set MIG_MONITOR_DEVICES in non privileged container")
} }
var imexChannels string
if c := getImexChannels(image); c != nil {
imexChannels = *c
}
driverCapabilities := hookConfig.getDriverCapabilities(image, legacyImage).String() driverCapabilities := hookConfig.getDriverCapabilities(image, legacyImage).String()
requirements, err := image.GetRequirements() requirements, err := image.GetRequirements()
@ -335,6 +350,7 @@ func getNvidiaConfig(hookConfig *HookConfig, image image.CUDA, mounts []Mount, p
Devices: devices, Devices: devices,
MigConfigDevices: migConfigDevices, MigConfigDevices: migConfigDevices,
MigMonitorDevices: migMonitorDevices, MigMonitorDevices: migMonitorDevices,
ImexChannels: imexChannels,
DriverCapabilities: driverCapabilities, DriverCapabilities: driverCapabilities,
Requirements: requirements, Requirements: requirements,
} }

3
cmd/nvidia-container-runtime-hook/main.go
View File

@ -126,6 +126,9 @@ func doPrestart() {
if len(nvidia.MigMonitorDevices) > 0 { if len(nvidia.MigMonitorDevices) > 0 {
args = append(args, fmt.Sprintf("--mig-monitor=%s", nvidia.MigMonitorDevices)) args = append(args, fmt.Sprintf("--mig-monitor=%s", nvidia.MigMonitorDevices))
} }
if len(nvidia.ImexChannels) > 0 {
args = append(args, fmt.Sprintf("--imex-channel=%s", nvidia.ImexChannels))
}
for _, cap := range strings.Split(nvidia.DriverCapabilities, ",") { for _, cap := range strings.Split(nvidia.DriverCapabilities, ",") {
if len(cap) == 0 { if len(cap) == 0 {