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 f93d96a0de
commit aa696ef919
No known key found for this signature in database
GPG Key ID: 42D9CB42F300A852
6 changed files with 37 additions and 16 deletions

View File

@ -44,6 +44,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)
@ -52,7 +75,7 @@ type CDIHook struct {
}
type HookCreator interface {
Create(string, ...string) *Hook
Create(HookName, ...string) *Hook
}
func NewHookCreator(nvidiaCDIHookPath string) HookCreator {
@ -63,7 +86,7 @@ func NewHookCreator(nvidiaCDIHookPath string) 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
@ -79,7 +102,7 @@ func (c CDIHook) Create(name string, args ...string) *Hook {
return &Hook{
Lifecycle: cdi.CreateContainerHook,
Path: c.nvidiaCDIHookPath,
Args: append(c.requiredArgs(name), args...),
Args: append(c.requiredArgs(string(name)), args...),
}
}

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,12 +37,12 @@ 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
)

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

@ -56,14 +56,14 @@ type nvcdilib struct {
mergedDeviceOptions []transform.MergedDeviceOption
disabledHooks disabledHooks
disabledHooks discover.DisabledHooks
hookCreator discover.HookCreator
}
// New creates a new nvcdi library
func New(opts ...Option) (Interface, error) {
l := &nvcdilib{
disabledHooks: make(disabledHooks),
disabledHooks: make(discover.DisabledHooks),
}
for _, opt := range opts {
opt(l)