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 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 // Option is a function that configures the nvcdilib
type Option func(*CDIHook) type Option func(*CDIHook)
@ -52,7 +75,7 @@ type CDIHook struct {
} }
type HookCreator interface { type HookCreator interface {
Create(string, ...string) *Hook Create(HookName, ...string) *Hook
} }
func NewHookCreator(nvidiaCDIHookPath string) HookCreator { func NewHookCreator(nvidiaCDIHookPath string) HookCreator {
@ -63,7 +86,7 @@ func NewHookCreator(nvidiaCDIHookPath string) HookCreator {
return CDIHook 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 name == "create-symlinks" {
if len(args) == 0 { if len(args) == 0 {
return nil return nil
@ -79,7 +102,7 @@ func (c CDIHook) Create(name string, args ...string) *Hook {
return &Hook{ return &Hook{
Lifecycle: cdi.CreateContainerHook, Lifecycle: cdi.CreateContainerHook,
Path: c.nvidiaCDIHookPath, 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) 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 // 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/pkg/cdi"
"tags.cncf.io/container-device-interface/specs-go" "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" "github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
) )
@ -36,12 +37,12 @@ type Interface interface {
GetDeviceSpecsByID(...string) ([]specs.Device, error) GetDeviceSpecsByID(...string) ([]specs.Device, error)
} }
// A HookName refers to one of the predefined set of CDI hooks that may be // HookName is an alias for the discover.HookName type.
// included in the generated CDI specification. type HookName = discover.HookName
type HookName string
// Aliases for the discover.HookName constants.
const ( const (
// HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility. HookEnableCudaCompat = discover.HookEnableCudaCompat
// This was added with v1.17.5 of the NVIDIA Container Toolkit. HookCreateSymlinks = discover.HookCreateSymlinks
HookEnableCudaCompat = HookName("enable-cuda-compat") HookUpdateLDCache = discover.HookUpdateLDCache
) )

View File

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

View File

@ -16,9 +16,6 @@
package nvcdi 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. // HookIsSupported checks whether a hook of the specified name is supported.
// Hooks must be explicitly disabled, meaning that if no disabled hooks are // Hooks must be explicitly disabled, meaning that if no disabled hooks are
// all hooks are supported. // all hooks are supported.

View File

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