UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-06-26 18:18:24 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge f882eb56da into 4f9c860a37

Browse Source
This commit is contained in:
Carlos Eduardo Arango Gutierrez
2025-06-13 17:05:45 +02:00
committed by GitHub
parent 4f9c860a37 f882eb56da
commit 7399061b17
7 changed files with 697 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
cmd/nvidia-ctk/cdi/generate/generate.go
View File

@@ -68,6 +68,14 @@ type options struct {
nvmllib nvml.Interface
}
// Add setter methods for csv.files and csv.ignorePatterns
func (o *options) SetCSVFiles(files []string) {
o.csv.files = *cli.NewStringSlice(files...)
}
func (o *options) SetCSVIgnorePatterns(patterns []string) {
o.csv.ignorePatterns = *cli.NewStringSlice(patterns...)
}
// NewCommand constructs a generate-cdi command with the specified logger
func NewCommand(logger logger.Interface) *cli.Command {
c := command{
@@ -207,6 +215,16 @@ func (m command) build() *cli.Command {
}
func (m command) validateFlags(c *cli.Context, opts *options) error {
// Load config file as base configuration
cfg, err := config.GetConfig()
if err != nil {
return fmt.Errorf("failed to load config: %v", err)
}
// Use centralized flag resolution (CLI > config file > default)
config.ResolveCDIGenerateOptions(c, cfg, opts)
// Additional validation (format, mode, etc.) can remain here if needed
opts.format = strings.ToLower(opts.format)
switch opts.format {
case spec.FormatJSON:

56
cmd/nvidia-ctk/cdi/generate/generate_test.go
View File

@@ -18,6 +18,7 @@ package generate
import (
"bytes"
"os"
"path/filepath"
"strings"
"testing"
@@ -32,6 +33,29 @@ import (
)
func TestGenerateSpec(t *testing.T) {
// Create a temporary directory for config
tmpDir, err := os.MkdirTemp("", "nvidia-container-toolkit-test-*")
require.NoError(t, err)
defer os.RemoveAll(tmpDir)
// Create a temporary config file
configContent := `
[nvidia-container-runtime]
mode = "nvml"
[[nvidia-container-runtime.modes.cdi]]
spec-dirs = ["/etc/cdi", "/usr/local/cdi"]
[nvidia-container-runtime.modes.csv]
mount-spec-path = "/etc/nvidia-container-runtime/host-files-for-container.d"
`
configPath := filepath.Join(tmpDir, "config.toml")
err = os.WriteFile(configPath, []byte(configContent), 0600)
require.NoError(t, err)
// Set XDG_CONFIG_HOME to point to our temporary directory
oldXDGConfigHome := os.Getenv("XDG_CONFIG_HOME")
os.Setenv("XDG_CONFIG_HOME", tmpDir)
defer os.Setenv("XDG_CONFIG_HOME", oldXDGConfigHome)
t.Setenv("__NVCT_TESTING_DEVICES_ARE_FILES", "true")
moduleRoot, err := test.GetModuleRoot()
require.NoError(t, err)
@@ -63,6 +87,13 @@ func TestGenerateSpec(t *testing.T) {
class: "device",
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
driverRoot: driverRoot,
csv: struct {
files cli.StringSlice
ignorePatterns cli.StringSlice
}{
files: *cli.NewStringSlice("/etc/nvidia-container-runtime/host-files-for-container.d"),
ignorePatterns: *cli.NewStringSlice(),
},
},
expectedSpec: `---
cdiVersion: 0.5.0
@@ -147,6 +178,13 @@ containerEdits:
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
driverRoot: driverRoot,
disabledHooks: valueOf(cli.NewStringSlice("enable-cuda-compat")),
csv: struct {
files cli.StringSlice
ignorePatterns cli.StringSlice
}{
files: *cli.NewStringSlice("/etc/nvidia-container-runtime/host-files-for-container.d"),
ignorePatterns: *cli.NewStringSlice(),
},
},
expectedSpec: `---
cdiVersion: 0.5.0
@@ -223,6 +261,13 @@ containerEdits:
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
driverRoot: driverRoot,
disabledHooks: valueOf(cli.NewStringSlice("enable-cuda-compat", "update-ldcache")),
csv: struct {
files cli.StringSlice
ignorePatterns cli.StringSlice
}{
files: *cli.NewStringSlice("/etc/nvidia-container-runtime/host-files-for-container.d"),
ignorePatterns: *cli.NewStringSlice(),
},
},
expectedSpec: `---
cdiVersion: 0.5.0
@@ -290,6 +335,13 @@ containerEdits:
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
driverRoot: driverRoot,
disabledHooks: valueOf(cli.NewStringSlice("all")),
csv: struct {
files cli.StringSlice
ignorePatterns cli.StringSlice
}{
files: *cli.NewStringSlice("/etc/nvidia-container-runtime/host-files-for-container.d"),
ignorePatterns: *cli.NewStringSlice(),
},
},
expectedSpec: `---
cdiVersion: 0.5.0
@@ -332,6 +384,10 @@ containerEdits:
err := c.validateFlags(nil, &tc.options)
require.ErrorIs(t, err, tc.expectedValidateError)
// Set the ldconfig path to empty.
// This is required during test because config.GetConfig() returns
// the default ldconfig path, even if it is not set in the config file.
tc.options.ldconfigPath = ""
require.EqualValues(t, tc.expectedOptions, tc.options)
// Set up a mock server, reusing the DGX A100 mock.

18
cmd/nvidia-ctk/cdi/list/list.go
View File

@@ -23,6 +23,7 @@ import (
"github.com/urfave/cli/v2"
"tags.cncf.io/container-device-interface/pkg/cdi"
ctkconfig "github.com/NVIDIA/nvidia-container-toolkit/internal/config"
"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
)
@@ -34,6 +35,11 @@ type config struct {
cdiSpecDirs cli.StringSlice
}
// SetCDISpecDirs sets the cdiSpecDirs field from a []string
func (c *config) SetCDISpecDirs(dirs []string) {
c.cdiSpecDirs = *cli.NewStringSlice(dirs...)
}
// NewCommand constructs a cdi list command with the specified logger
func NewCommand(logger logger.Interface) *cli.Command {
c := command{
@@ -71,10 +77,20 @@ func (m command) build() *cli.Command {
return &c
}
func (m command) validateFlags(c *cli.Context, cfg *config) error {
func (m command) validateFlags(ctx *cli.Context, cfg *config) error {
// Load config file as base configuration
c, err := ctkconfig.GetConfig()
if err != nil {
return fmt.Errorf("failed to load config: %v", err)
}
// Use centralized normalization
ctkconfig.ResolveCDIListConfig(ctx, c, cfg)
if len(cfg.cdiSpecDirs.Value()) == 0 {
return errors.New("at least one CDI specification directory must be specified")
}
return nil
}

95
cmd/nvidia-ctk/cdi/list/list_test.go Normal file
View File

@@ -0,0 +1,95 @@
package list
import (
"os"
"path/filepath"
"testing"
testlog "github.com/sirupsen/logrus/hooks/test"
"github.com/stretchr/testify/require"
"github.com/urfave/cli/v2"
)
func TestValidateFlags(t *testing.T) {
logger, _ := testlog.NewNullLogger()
// Create a temporary directory for config
tmpDir, err := os.MkdirTemp("", "nvidia-container-toolkit-test-*")
require.NoError(t, err)
defer os.RemoveAll(tmpDir)
// Create a temporary config file
configContent := `
[nvidia-container-runtime]
mode = "cdi"
[[nvidia-container-runtime.modes.cdi]]
spec-dirs = ["/etc/cdi", "/usr/local/cdi"]
`
configPath := filepath.Join(tmpDir, "config.toml")
err = os.WriteFile(configPath, []byte(configContent), 0600)
require.NoError(t, err)
// Set XDG_CONFIG_HOME to point to our temporary directory
oldXDGConfigHome := os.Getenv("XDG_CONFIG_HOME")
os.Setenv("XDG_CONFIG_HOME", tmpDir)
defer os.Setenv("XDG_CONFIG_HOME", oldXDGConfigHome)
tests := []struct {
name string
cliArgs []string
envVars map[string]string
expectedDirs []string
expectError bool
errorContains string
}{
{
name: "command line takes precedence",
cliArgs: []string{"--spec-dir=/custom/dir1", "--spec-dir=/custom/dir2"},
expectedDirs: []string{"/custom/dir1", "/custom/dir2"},
},
{
name: "environment variable takes precedence over config",
envVars: map[string]string{"NVIDIA_CTK_CDI_SPEC_DIRS": "/env/dir1:/env/dir2"},
expectedDirs: []string{"/env/dir1", "/env/dir2"},
},
{
name: "config file used as fallback",
expectedDirs: []string{"/etc/cdi", "/usr/local/cdi"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Set up environment variables
for k, v := range tt.envVars {
old := os.Getenv(k)
os.Setenv(k, v)
defer os.Setenv(k, old)
}
// Create command
cmd := NewCommand(logger)
// Create a new context with the command
app := &cli.App{
Commands: []*cli.Command{
{
Name: "cdi",
Subcommands: []*cli.Command{cmd},
},
},
}
// Run command
args := append([]string{"nvidia-ctk", "cdi", "list"}, tt.cliArgs...)
err := app.Run(args)
if tt.expectError {
require.Error(t, err)
require.Contains(t, err.Error(), tt.errorContains)
return
}
require.NoError(t, err)
})
}
}