Consolidate HookName functionality on internal/discover pkg

Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>
This commit is contained in:
Carlos Eduardo Arango Gutierrez 2025-05-22 13:03:53 +02:00
parent 890db82b46
commit b4787511d2
No known key found for this signature in database
GPG Key ID: 42D9CB42F300A852
7 changed files with 40 additions and 24 deletions

View File

@ -45,6 +45,29 @@ func (h *Hook) Hooks() ([]Hook, error) {
return []Hook{*h}, nil
}
type HookName string
// DisabledHooks allows individual hooks to be disabled.
type DisabledHooks map[HookName]bool
const (
// HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility.
// This was added with v1.17.5 of the NVIDIA Container Toolkit.
HookEnableCudaCompat = HookName("enable-cuda-compat")
// directory path to be mounted into a container.
HookCreateSymlinks = HookName("create-symlinks")
// HookUpdateLDCache refers to the hook used to Update the dynamic linker
// cache inside the directory path to be mounted into a container.
HookUpdateLDCache = HookName("update-ldcache")
)
// AllHooks maintains a future-proof list of all defined hooks.
var AllHooks = []HookName{
HookEnableCudaCompat,
HookCreateSymlinks,
HookUpdateLDCache,
}
// Option is a function that configures the nvcdilib
type Option func(*CDIHook)
@ -54,7 +77,7 @@ type CDIHook struct {
}
type HookCreator interface {
Create(string, ...string) *Hook
Create(HookName, ...string) *Hook
}
func NewHookCreator(nvidiaCDIHookPath string, debugLogging bool) HookCreator {
@ -66,7 +89,7 @@ func NewHookCreator(nvidiaCDIHookPath string, debugLogging bool) HookCreator {
return CDIHook
}
func (c CDIHook) Create(name string, args ...string) *Hook {
func (c CDIHook) Create(name HookName, args ...string) *Hook {
if name == "create-symlinks" {
if len(args) == 0 {
return nil

View File

@ -72,7 +72,7 @@ func createLDCacheUpdateHook(hookCreator HookCreator, ldconfig string, libraries
args = append(args, "--folder", f)
}
return hookCreator.Create("update-ldcache", args...)
return hookCreator.Create(HookUpdateLDCache, args...)
}
// getLibraryPaths extracts the library dirs from the specified mounts

View File

@ -21,6 +21,7 @@ import (
"tags.cncf.io/container-device-interface/pkg/cdi"
"tags.cncf.io/container-device-interface/specs-go"
"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
)
@ -36,14 +37,14 @@ type Interface interface {
GetDeviceSpecsByID(...string) ([]specs.Device, error)
}
// A HookName refers to one of the predefined set of CDI hooks that may be
// included in the generated CDI specification.
type HookName string
// HookName is an alias for the discover.HookName type.
type HookName = discover.HookName
// Aliases for the discover.HookName constants.
const (
// HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility.
// This was added with v1.17.5 of the NVIDIA Container Toolkit.
HookEnableCudaCompat = HookName("enable-cuda-compat")
HookEnableCudaCompat = discover.HookEnableCudaCompat
HookCreateSymlinks = discover.HookCreateSymlinks
HookUpdateLDCache = discover.HookUpdateLDCache
)
// A FeatureFlag refers to a specific feature that can be toggled in the CDI api.

View File

@ -135,7 +135,7 @@ func (m nvidiaSMISimlinkHook) Hooks() ([]discover.Hook, error) {
}
link := "/usr/bin/nvidia-smi"
links := []string{fmt.Sprintf("%s::%s", target, link)}
symlinkHook := m.hookCreator.Create("create-symlinks", links...)
symlinkHook := m.hookCreator.Create(HookCreateSymlinks, links...)
return symlinkHook.Hooks()
}

View File

@ -16,9 +16,6 @@
package nvcdi
// disabledHooks allows individual hooks to be disabled.
type disabledHooks map[HookName]bool
// HookIsSupported checks whether a hook of the specified name is supported.
// Hooks must be explicitly disabled, meaning that if no disabled hooks are
// all hooks are supported.

View File

@ -58,16 +58,13 @@ type nvcdilib struct {
featureFlags map[FeatureFlag]bool
disabledHooks disabledHooks
disabledHooks []discover.HookName
hookCreator discover.HookCreator
}
// New creates a new nvcdi library
func New(opts ...Option) (Interface, error) {
l := &nvcdilib{
disabledHooks: make(disabledHooks),
featureFlags: make(map[FeatureFlag]bool),
}
l := &nvcdilib{}
for _, opt := range opts {
opt(l)
}
@ -136,7 +133,7 @@ func New(opts ...Option) (Interface, error) {
l.vendor = "management.nvidia.com"
}
// Management containers in general do not require CUDA Forward compatibility.
l.disabledHooks[HookEnableCudaCompat] = true
l.disabledHooks = append(l.disabledHooks, discover.HookEnableCudaCompat)
lib = (*managementlib)(l)
case ModeNvml:
lib = (*nvmllib)(l)

View File

@ -21,6 +21,7 @@ import (
"github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
"github.com/NVIDIA/go-nvml/pkg/nvml"
"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/transform"
)
@ -158,12 +159,9 @@ func WithLibrarySearchPaths(paths []string) Option {
// WithDisabledHook allows specific hooks to the disabled.
// This option can be specified multiple times for each hook.
func WithDisabledHook(hook HookName) Option {
func WithDisabledHook[T string | HookName](hook T) Option {
return func(o *nvcdilib) {
if o.disabledHooks == nil {
o.disabledHooks = make(map[HookName]bool)
}
o.disabledHooks[hook] = true
o.disabledHooks = append(o.disabledHooks, discover.HookName(hook))
}
}