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

Add FeatureFlags to the nvcdi API

Browse Source 
This change adds support for feature flags to the nvcdi API.

A feature flag to disable nvsandboxutils is also added to allow
more flexibility in cases where this library causes issue.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2025-05-22 16:35:01 +02:00
parent 872aa2fe1c
commit 7bd65da91e
No known key found for this signature in database
3 changed files with 27 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
pkg/nvcdi/api.go
View File

@ -45,3 +45,13 @@ const (
// This was added with v1.17.5 of the NVIDIA Container Toolkit.
HookEnableCudaCompat = HookName("enable-cuda-compat")
)
// A FeatureFlag refers to a specific feature that can be toggled in the CDI api.
// All features are off by default.
type FeatureFlag string
const (
// FeatureDisableNvsandboxUtils disables the use of nvsandboxutils when
// querying devices.
FeatureDisableNvsandboxUtils = FeatureFlag("disable-nvsandbox-utils")
)

6
pkg/nvcdi/lib.go
View File

@ -56,6 +56,8 @@ type nvcdilib struct {
mergedDeviceOptions []transform.MergedDeviceOption
featureFlags map[FeatureFlag]bool
disabledHooks disabledHooks
hookCreator discover.HookCreator
}
@ -64,6 +66,7 @@ type nvcdilib struct {
func New(opts ...Option) (Interface, error) {
l := &nvcdilib{
disabledHooks: make(disabledHooks),
featureFlags: make(map[FeatureFlag]bool),
}
for _, opt := range opts {
opt(l)
@ -218,6 +221,9 @@ func (l *nvcdilib) getCudaVersionNvsandboxutils() (string, error) {
// getNvsandboxUtilsLib returns the nvsandboxutilslib to use for CDI spec
// generation.
func (l *nvcdilib) getNvsandboxUtilsLib() nvsandboxutils.Interface {
if l.featureFlags[FeatureDisableNvsandboxUtils] {
return nil
}
if l.nvsandboxutilslib != nil {
return l.nvsandboxutilslib
}

11
pkg/nvcdi/options.go
View File

@ -166,3 +166,14 @@ func WithDisabledHook(hook HookName) Option {
o.disabledHooks[hook] = true
}
}
// WithFeatureFlag allows specified features to be toggled on.
// This option can be specified multiple times for each feature flag.
func WithFeatureFlag(featureFlag FeatureFlag) Option {
return func(o *nvcdilib) {
if o.featureFlags == nil {
o.featureFlags = make(map[FeatureFlag]bool)
}
o.featureFlags[featureFlag] = true
}
}