Skip to content

Commit aa696ef

Browse files
Consolidate HookName functionality on internal/discover pkg
Signed-off-by: Carlos Eduardo Arango Gutierrez <[email protected]>
1 parent f93d96a commit aa696ef

File tree

6 files changed

+37
-16
lines changed

6 files changed

+37
-16
lines changed

internal/discover/hooks.go

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,29 @@ func (h *Hook) Hooks() ([]Hook, error) {
4444
return []Hook{*h}, nil
4545
}
4646

47+
type HookName string
48+
49+
// DisabledHooks allows individual hooks to be disabled.
50+
type DisabledHooks map[HookName]bool
51+
52+
const (
53+
// HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility.
54+
// This was added with v1.17.5 of the NVIDIA Container Toolkit.
55+
HookEnableCudaCompat = HookName("enable-cuda-compat")
56+
// directory path to be mounted into a container.
57+
HookCreateSymlinks = HookName("create-symlinks")
58+
// HookUpdateLDCache refers to the hook used to Update the dynamic linker
59+
// cache inside the directory path to be mounted into a container.
60+
HookUpdateLDCache = HookName("update-ldcache")
61+
)
62+
63+
// AllHooks maintains a future-proof list of all defined hooks.
64+
var AllHooks = []HookName{
65+
HookEnableCudaCompat,
66+
HookCreateSymlinks,
67+
HookUpdateLDCache,
68+
}
69+
4770
// Option is a function that configures the nvcdilib
4871
type Option func(*CDIHook)
4972

@@ -52,7 +75,7 @@ type CDIHook struct {
5275
}
5376

5477
type HookCreator interface {
55-
Create(string, ...string) *Hook
78+
Create(HookName, ...string) *Hook
5679
}
5780

5881
func NewHookCreator(nvidiaCDIHookPath string) HookCreator {
@@ -63,7 +86,7 @@ func NewHookCreator(nvidiaCDIHookPath string) HookCreator {
6386
return CDIHook
6487
}
6588

66-
func (c CDIHook) Create(name string, args ...string) *Hook {
89+
func (c CDIHook) Create(name HookName, args ...string) *Hook {
6790
if name == "create-symlinks" {
6891
if len(args) == 0 {
6992
return nil
@@ -79,7 +102,7 @@ func (c CDIHook) Create(name string, args ...string) *Hook {
79102
return &Hook{
80103
Lifecycle: cdi.CreateContainerHook,
81104
Path: c.nvidiaCDIHookPath,
82-
Args: append(c.requiredArgs(name), args...),
105+
Args: append(c.requiredArgs(string(name)), args...),
83106
}
84107
}
85108

internal/discover/ldconfig.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ func createLDCacheUpdateHook(hookCreator HookCreator, ldconfig string, libraries
7272
args = append(args, "--folder", f)
7373
}
7474

75-
return hookCreator.Create("update-ldcache", args...)
75+
return hookCreator.Create(HookUpdateLDCache, args...)
7676
}
7777

7878
// getLibraryPaths extracts the library dirs from the specified mounts

pkg/nvcdi/api.go

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"tags.cncf.io/container-device-interface/pkg/cdi"
2222
"tags.cncf.io/container-device-interface/specs-go"
2323

24+
"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
2425
"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
2526
)
2627

@@ -36,12 +37,12 @@ type Interface interface {
3637
GetDeviceSpecsByID(...string) ([]specs.Device, error)
3738
}
3839

39-
// A HookName refers to one of the predefined set of CDI hooks that may be
40-
// included in the generated CDI specification.
41-
type HookName string
40+
// HookName is an alias for the discover.HookName type.
41+
type HookName = discover.HookName
4242

43+
// Aliases for the discover.HookName constants.
4344
const (
44-
// HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility.
45-
// This was added with v1.17.5 of the NVIDIA Container Toolkit.
46-
HookEnableCudaCompat = HookName("enable-cuda-compat")
45+
HookEnableCudaCompat = discover.HookEnableCudaCompat
46+
HookCreateSymlinks = discover.HookCreateSymlinks
47+
HookUpdateLDCache = discover.HookUpdateLDCache
4748
)

pkg/nvcdi/driver-wsl.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ func (m nvidiaSMISimlinkHook) Hooks() ([]discover.Hook, error) {
135135
}
136136
link := "/usr/bin/nvidia-smi"
137137
links := []string{fmt.Sprintf("%s::%s", target, link)}
138-
symlinkHook := m.hookCreator.Create("create-symlinks", links...)
138+
symlinkHook := m.hookCreator.Create(HookCreateSymlinks, links...)
139139

140140
return symlinkHook.Hooks()
141141
}

pkg/nvcdi/hooks.go

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,6 @@
1616

1717
package nvcdi
1818

19-
// disabledHooks allows individual hooks to be disabled.
20-
type disabledHooks map[HookName]bool
21-
2219
// HookIsSupported checks whether a hook of the specified name is supported.
2320
// Hooks must be explicitly disabled, meaning that if no disabled hooks are
2421
// all hooks are supported.

pkg/nvcdi/lib.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,14 +56,14 @@ type nvcdilib struct {
5656

5757
mergedDeviceOptions []transform.MergedDeviceOption
5858

59-
disabledHooks disabledHooks
59+
disabledHooks discover.DisabledHooks
6060
hookCreator discover.HookCreator
6161
}
6262

6363
// New creates a new nvcdi library
6464
func New(opts ...Option) (Interface, error) {
6565
l := &nvcdilib{
66-
disabledHooks: make(disabledHooks),
66+
disabledHooks: make(discover.DisabledHooks),
6767
}
6868
for _, opt := range opts {
6969
opt(l)

0 commit comments

Comments
 (0)