Skip to content

Add --disable-hook flag to cdi generate command #1077

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Add --disable-hook flag to cdi generate command #1077

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
96ecfd0
Add --disable-hook flag to cdi generate command
ArangoGutierrez May 12, 2025
d657e0a
rename flag to disable-hooks
ArangoGutierrez May 14, 2025
191f11e
Add all option to disable all hooks
ArangoGutierrez May 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions cmd/nvidia-ctk/cdi/generate/generate.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ type options struct {

configSearchPaths cli.StringSlice
librarySearchPaths cli.StringSlice
disabledHooks cli.StringSlice

csv struct {
files cli.StringSlice
Expand Down Expand Up @@ -176,6 +177,13 @@ func (m command) build() *cli.Command {
Usage: "Specify a pattern the CSV mount specifications.",
Destination: &opts.csv.ignorePatterns,
},
&cli.StringSliceFlag{
Name: "disable-hook",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we also add a disable-all-hooks flag so people don't have to enumerate all of them to disable them.

Also, should we add a list-hooks command so people can discover what hooks are available for disabling

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@klueska my thoughts on how to handle the disabling of all hooks is to also accept all as a valid hook name. So that --disable-hook(s)=all can be specified.

Copy link
Member

@elezar elezar May 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, should we add a list-hooks command so people can discover what hooks are available for disabling

Do you mean something like:

nvidia-ctk cdi list-hooks

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have added an alias of disable-hooks and a better Usage description.

Aliases: []string{"disable-hooks"},
Usage: "Hook to skip when generating the CDI specification. Can be specified multiple times. Can be a comma-separated list of hooks or a single hook name.",
Value: cli.NewStringSlice(),
Destination: &opts.disabledHooks,
},
}

return &c
Expand Down Expand Up @@ -262,7 +270,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
deviceNamers = append(deviceNamers, deviceNamer)
}

cdilib, err := nvcdi.New(
initOpts := []nvcdi.Option{
nvcdi.WithLogger(m.logger),
nvcdi.WithDriverRoot(opts.driverRoot),
nvcdi.WithDevRoot(opts.devRoot),
Expand All @@ -276,7 +284,13 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
nvcdi.WithCSVIgnorePatterns(opts.csv.ignorePatterns.Value()),
// We set the following to allow for dependency injection:
nvcdi.WithNvmlLib(opts.nvmllib),
)
}

for _, hook := range opts.disabledHooks.Value() {
initOpts = append(initOpts, nvcdi.WithDisabledHook(hook))
}

cdilib, err := nvcdi.New(initOpts...)
if err != nil {
return nil, fmt.Errorf("failed to create CDI library: %v", err)
}
Expand Down
166 changes: 166 additions & 0 deletions cmd/nvidia-ctk/cdi/generate/generate_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (
"github.com/NVIDIA/go-nvml/pkg/nvml/mock/dgxa100"
testlog "github.com/sirupsen/logrus/hooks/test"
"github.com/stretchr/testify/require"
"github.com/urfave/cli/v2"

"github.com/NVIDIA/nvidia-container-toolkit/internal/test"
)
Expand All @@ -36,6 +37,9 @@ func TestGenerateSpec(t *testing.T) {
require.NoError(t, err)

driverRoot := filepath.Join(moduleRoot, "testdata", "lookup", "rootfs-1")
disableHook1 := cli.NewStringSlice("enable-cuda-compat")
disableHook2 := cli.NewStringSlice("enable-cuda-compat", "update-ldcache")
disableHook3 := cli.NewStringSlice("all")

logger, _ := testlog.NewNullLogger()
testCases := []struct {
Expand Down Expand Up @@ -112,6 +116,168 @@ containerEdits:
- nosuid
- nodev
- bind
`,
},
{
description: "disableHooks1",
options: options{
format: "yaml",
mode: "nvml",
vendor: "example.com",
class: "device",
driverRoot: driverRoot,
disabledHooks: *disableHook1,
},
expectedOptions: options{
format: "yaml",
mode: "nvml",
vendor: "example.com",
class: "device",
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
driverRoot: driverRoot,
disabledHooks: *disableHook1,
},
expectedSpec: `---
cdiVersion: 0.5.0
kind: example.com/device
devices:
- name: "0"
containerEdits:
deviceNodes:
- path: /dev/nvidia0
hostPath: {{ .driverRoot }}/dev/nvidia0
- name: all
containerEdits:
deviceNodes:
- path: /dev/nvidia0
hostPath: {{ .driverRoot }}/dev/nvidia0
containerEdits:
env:
- NVIDIA_VISIBLE_DEVICES=void
deviceNodes:
- path: /dev/nvidiactl
hostPath: {{ .driverRoot }}/dev/nvidiactl
hooks:
- hookName: createContainer
path: /usr/bin/nvidia-cdi-hook
args:
- nvidia-cdi-hook
- create-symlinks
- --link
- libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
- hookName: createContainer
path: /usr/bin/nvidia-cdi-hook
args:
- nvidia-cdi-hook
- update-ldcache
- --folder
- /lib/x86_64-linux-gnu
mounts:
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
options:
- ro
- nosuid
- nodev
- bind
`,
},
{
description: "disableHooks2",
options: options{
format: "yaml",
mode: "nvml",
vendor: "example.com",
class: "device",
driverRoot: driverRoot,
disabledHooks: *disableHook2,
},
expectedOptions: options{
format: "yaml",
mode: "nvml",
vendor: "example.com",
class: "device",
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
driverRoot: driverRoot,
disabledHooks: *disableHook2,
},
expectedSpec: `---
cdiVersion: 0.5.0
kind: example.com/device
devices:
- name: "0"
containerEdits:
deviceNodes:
- path: /dev/nvidia0
hostPath: {{ .driverRoot }}/dev/nvidia0
- name: all
containerEdits:
deviceNodes:
- path: /dev/nvidia0
hostPath: {{ .driverRoot }}/dev/nvidia0
containerEdits:
env:
- NVIDIA_VISIBLE_DEVICES=void
deviceNodes:
- path: /dev/nvidiactl
hostPath: {{ .driverRoot }}/dev/nvidiactl
hooks:
- hookName: createContainer
path: /usr/bin/nvidia-cdi-hook
args:
- nvidia-cdi-hook
- create-symlinks
- --link
- libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
mounts:
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
options:
- ro
- nosuid
- nodev
- bind
`,
},
{
description: "disableHooksAll",
options: options{
format: "yaml",
mode: "nvml",
vendor: "example.com",
class: "device",
driverRoot: driverRoot,
disabledHooks: *disableHook3,
},
expectedOptions: options{
format: "yaml",
mode: "nvml",
vendor: "example.com",
class: "device",
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
driverRoot: driverRoot,
disabledHooks: *disableHook3,
},
expectedSpec: `---
cdiVersion: 0.5.0
kind: example.com/device
devices:
- name: "0"
containerEdits:
deviceNodes:
- path: /dev/nvidia0
hostPath: {{ .driverRoot }}/dev/nvidia0
- name: all
containerEdits:
deviceNodes:
- path: /dev/nvidia0
hostPath: {{ .driverRoot }}/dev/nvidia0
containerEdits:
env:
- NVIDIA_VISIBLE_DEVICES=void
deviceNodes:
- path: /dev/nvidiactl
hostPath: {{ .driverRoot }}/dev/nvidiactl
`,
},
}
Expand Down
13 changes: 13 additions & 0 deletions pkg/nvcdi/api.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,4 +44,17 @@ 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")
// HookCreateSymlinks refers to the hook used create symlinks inside the
// 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,
}
31 changes: 22 additions & 9 deletions pkg/nvcdi/driver-nvml.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,21 +99,34 @@ func (l *nvcdilib) NewDriverLibraryDiscoverer(version string) (discover.Discover

var discoverers []discover.Discover

driverDotSoSymlinksDiscoverer := discover.WithDriverDotSoSymlinks(
libraries,
version,
l.nvidiaCDIHookPath,
)
discoverers = append(discoverers, driverDotSoSymlinksDiscoverer)
if l.HookIsSupported(HookCreateSymlinks) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For this hook specifically (and the update-ldcache hook too), there are many places in the code where we are injecting this hook. As such this is not going to be sufficient to prevent their injection. This was why I suggested that handling these hooks is out of scope for this PR.

driverDotSoSymlinksDiscoverer := discover.WithDriverDotSoSymlinks(
libraries,
version,
l.nvidiaCDIHookPath,
)
discoverers = append(discoverers, driverDotSoSymlinksDiscoverer)
}

if l.HookIsSupported(HookEnableCudaCompat) {
// TODO: The following should use the version directly.
cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(l.logger, l.nvidiaCDIHookPath, l.driver)
cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(
l.logger,
l.nvidiaCDIHookPath,
l.driver,
)
discoverers = append(discoverers, cudaCompatLibHookDiscoverer)
}

updateLDCache, _ := discover.NewLDCacheUpdateHook(l.logger, libraries, l.nvidiaCDIHookPath, l.ldconfigPath)
discoverers = append(discoverers, updateLDCache)
if l.HookIsSupported(HookUpdateLDCache) {
updateLDCache, _ := discover.NewLDCacheUpdateHook(
l.logger,
libraries,
l.nvidiaCDIHookPath,
l.ldconfigPath,
)
discoverers = append(discoverers, updateLDCache)
}

d := discover.Merge(discoverers...)

Expand Down
11 changes: 11 additions & 0 deletions pkg/nvcdi/hooks.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,3 +28,14 @@ func (l *nvcdilib) HookIsSupported(h HookName) bool {
}
return !l.disabledHooks[h]
}

func (d disabledHooks) Set(value HookName) {
if value == "all" {
for _, hook := range AllHooks {
d[hook] = true
}
return
}

d[value] = true
}
4 changes: 2 additions & 2 deletions pkg/nvcdi/options.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,11 +158,11 @@ 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.Set(HookName(hook))
}
}