mirror of
https://github.com/NVIDIA/nvidia-container-toolkit
synced 2025-06-11 00:57:15 +00:00
Consolidate HookName functionality on internal/discover pkg
Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>
This commit is contained in:
parent
890db82b46
commit
b4787511d2
@ -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
|
||||
|
@ -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
|
||||
|
@ -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.
|
||||
|
@ -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()
|
||||
}
|
||||
|
@ -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.
|
||||
|
@ -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)
|
||||
|
@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user