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

Skip injection of nvidia-persistenced socket by default

Browse Source 
This changes skips the injection of the nvidia-persistenced socket by
default.

An include-persistenced-socket feature flag is added to allow the
injection of this socket to be explicitly requested.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2024-09-18 22:10:09 +02:00
parent c490baab63
commit ba1ed3232f
3 changed files with 45 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -89,6 +89,12 @@ func doPrestart() {
rootfs := getRootfsPath(container) rootfs := getRootfsPath(container)
args := []string{getCLIPath(cli)} args := []string{getCLIPath(cli)}
// Only include the nvidia-persistenced socket if it is explicitly enabled.
if !hook.Features.IncludePersistencedSocket.IsEnabled() {
args = append(args, "--no-persistenced")
}
if cli.Root != "" { if cli.Root != "" {
args = append(args, fmt.Sprintf("--root=%s", cli.Root)) args = append(args, fmt.Sprintf("--root=%s", cli.Root))
} }

65
internal/config/features.go
View File

@ -19,10 +19,11 @@ package config
type featureName string type featureName string
const ( const (
FeatureGDS = featureName("gds") FeatureGDS = featureName("gds")
FeatureMOFED = featureName("mofed") FeatureMOFED = featureName("mofed")
FeatureNVSWITCH = featureName("nvswitch") FeatureNVSWITCH = featureName("nvswitch")
FeatureGDRCopy = featureName("gdrcopy") FeatureGDRCopy = featureName("gdrcopy")
FeatureIncludePersistencedSocket = featureName("include-persistenced-socket")
) )
// features specifies a set of named features. // features specifies a set of named features.
@ -31,53 +32,57 @@ type features struct {
MOFED *feature `toml:"mofed,omitempty"` MOFED *feature `toml:"mofed,omitempty"`
NVSWITCH *feature `toml:"nvswitch,omitempty"` NVSWITCH *feature `toml:"nvswitch,omitempty"`
GDRCopy *feature `toml:"gdrcopy,omitempty"` GDRCopy *feature `toml:"gdrcopy,omitempty"`
// IncludePersistencedSocket enables the injection of the nvidia-persistenced
// socket into containers.
IncludePersistencedSocket *feature `toml:"include-persistenced-socket,omitempty"`
} }
type feature bool type feature bool
// IsEnabled checks whether a specified named feature is enabled. // IsEnabledInEnvironment checks whether a specified named feature is enabled.
// An optional list of environments to check for feature-specific environment // An optional list of environments to check for feature-specific environment
// variables can also be supplied. // variables can also be supplied.
func (fs features) IsEnabled(n featureName, in ...getenver) bool { func (fs features) IsEnabledInEnvironment(n featureName, in ...getenver) bool {
featureEnvvars := map[featureName]string{
FeatureGDS: "NVIDIA_GDS",
FeatureMOFED: "NVIDIA_MOFED",
FeatureNVSWITCH: "NVIDIA_NVSWITCH",
FeatureGDRCopy: "NVIDIA_GDRCOPY",
}
envvar := featureEnvvars[n]
switch n { switch n {
// Features with envvar overrides
case FeatureGDS: case FeatureGDS:
return fs.GDS.isEnabled(envvar, in...) return fs.GDS.isEnabledWithEnvvarOverride("NVIDIA_GDS", in...)
case FeatureMOFED: case FeatureMOFED:
return fs.MOFED.isEnabled(envvar, in...) return fs.MOFED.isEnabledWithEnvvarOverride("NVIDIA_MOFED", in...)
case FeatureNVSWITCH: case FeatureNVSWITCH:
return fs.NVSWITCH.isEnabled(envvar, in...) return fs.NVSWITCH.isEnabledWithEnvvarOverride("NVIDIA_NVSWITCH", in...)
case FeatureGDRCopy: case FeatureGDRCopy:
return fs.GDRCopy.isEnabled(envvar, in...) return fs.GDRCopy.isEnabledWithEnvvarOverride("NVIDIA_GDRCOPY", in...)
// Features without envvar overrides
case FeatureIncludePersistencedSocket:
return fs.IncludePersistencedSocket.IsEnabled()
default: default:
return false return false
} }
} }
// isEnabled checks whether a feature is enabled. // IsEnabled checks whether a feature is enabled.
// If the enabled value is explicitly set, this is returned, otherwise the func (f *feature) IsEnabled() bool {
// associated envvar is checked in the specified getenver for the string "enabled"
// A CUDA container / image can be passed here.
func (f *feature) isEnabled(envvar string, ins ...getenver) bool {
if f != nil { if f != nil {
return bool(*f) return bool(*f)
} }
if envvar == "" { return false
return false }
}
for _, in := range ins { // isEnabledWithEnvvarOverride checks whether a feature is enabled and allows an envvar to overide the feature.
if in.Getenv(envvar) == "enabled" { // If the enabled value is explicitly set, this is returned, otherwise the
return true // associated envvar is checked in the specified getenver for the string "enabled"
// A CUDA container / image can be passed here.
func (f *feature) isEnabledWithEnvvarOverride(envvar string, ins ...getenver) bool {
if envvar != "" {
for _, in := range ins {
if in.Getenv(envvar) == "enabled" {
return true
}
} }
} }
return false
return f.IsEnabled()
} }
type getenver interface { type getenver interface {

8
internal/modifier/gated.go
View File

@ -46,7 +46,7 @@ func NewFeatureGatedModifier(logger logger.Interface, cfg *config.Config, image
driverRoot := cfg.NVIDIAContainerCLIConfig.Root driverRoot := cfg.NVIDIAContainerCLIConfig.Root
devRoot := cfg.NVIDIAContainerCLIConfig.Root devRoot := cfg.NVIDIAContainerCLIConfig.Root
if cfg.Features.IsEnabled(config.FeatureGDS, image) { if cfg.Features.IsEnabledInEnvironment(config.FeatureGDS, image) {
d, err := discover.NewGDSDiscoverer(logger, driverRoot, devRoot) d, err := discover.NewGDSDiscoverer(logger, driverRoot, devRoot)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to construct discoverer for GDS devices: %w", err) return nil, fmt.Errorf("failed to construct discoverer for GDS devices: %w", err)
@ -54,7 +54,7 @@ func NewFeatureGatedModifier(logger logger.Interface, cfg *config.Config, image
discoverers = append(discoverers, d) discoverers = append(discoverers, d)
} }
if cfg.Features.IsEnabled(config.FeatureMOFED, image) { if cfg.Features.IsEnabledInEnvironment(config.FeatureMOFED, image) {
d, err := discover.NewMOFEDDiscoverer(logger, devRoot) d, err := discover.NewMOFEDDiscoverer(logger, devRoot)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to construct discoverer for MOFED devices: %w", err) return nil, fmt.Errorf("failed to construct discoverer for MOFED devices: %w", err)
@ -62,7 +62,7 @@ func NewFeatureGatedModifier(logger logger.Interface, cfg *config.Config, image
discoverers = append(discoverers, d) discoverers = append(discoverers, d)
} }
if cfg.Features.IsEnabled(config.FeatureNVSWITCH, image) { if cfg.Features.IsEnabledInEnvironment(config.FeatureNVSWITCH, image) {
d, err := discover.NewNvSwitchDiscoverer(logger, devRoot) d, err := discover.NewNvSwitchDiscoverer(logger, devRoot)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to construct discoverer for NVSWITCH devices: %w", err) return nil, fmt.Errorf("failed to construct discoverer for NVSWITCH devices: %w", err)
@ -70,7 +70,7 @@ func NewFeatureGatedModifier(logger logger.Interface, cfg *config.Config, image
discoverers = append(discoverers, d) discoverers = append(discoverers, d)
} }
if cfg.Features.IsEnabled(config.FeatureGDRCopy, image) { if cfg.Features.IsEnabledInEnvironment(config.FeatureGDRCopy, image) {
d, err := discover.NewGDRCopyDiscoverer(logger, devRoot) d, err := discover.NewGDRCopyDiscoverer(logger, devRoot)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to construct discoverer for GDRCopy devices: %w", err) return nil, fmt.Errorf("failed to construct discoverer for GDRCopy devices: %w", err)