mirror of
				https://github.com/NVIDIA/nvidia-container-toolkit
				synced 2025-06-26 18:18:24 +00:00 
			
		
		
		
	Add config.Toml type to handle config files
This change introduced a config.Toml type that is used as the base for config file processing and manipulation. This ensures that configs -- including commented values -- can be handled consistently. Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
		
							parent
							
								
									c2d4de54b0
								
							
						
					
					
						commit
						a69657dde7
					
				| @ -43,13 +43,15 @@ func loadConfig() (*config.Config, error) { | ||||
| 	} | ||||
| 
 | ||||
| 	for _, p := range configPaths { | ||||
| 		cfg, err := config.Load(p) | ||||
| 		cfg, err := config.New( | ||||
| 			config.WithConfigFile(p), | ||||
| 		) | ||||
| 		if err == nil { | ||||
| 			return cfg, nil | ||||
| 			return cfg.Config() | ||||
| 		} else if os.IsNotExist(err) && !required { | ||||
| 			continue | ||||
| 		} | ||||
| 		return nil, fmt.Errorf("couldn't open configuration file: %v", err) | ||||
| 		return nil, fmt.Errorf("couldn't open required configuration file: %v", err) | ||||
| 	} | ||||
| 
 | ||||
| 	return config.GetDefault() | ||||
|  | ||||
| @ -17,12 +17,10 @@ | ||||
| package defaultsubcommand | ||||
| 
 | ||||
| import ( | ||||
| 	"bytes" | ||||
| 	"fmt" | ||||
| 	"io" | ||||
| 	"os" | ||||
| 	"path/filepath" | ||||
| 	"regexp" | ||||
| 
 | ||||
| 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config" | ||||
| 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger" | ||||
| @ -102,57 +100,23 @@ func (m command) run(c *cli.Context, opts *options) error { | ||||
| 		return fmt.Errorf("unable to create output directory: %v", err) | ||||
| 	} | ||||
| 
 | ||||
| 	contents, err := opts.getFormattedConfig() | ||||
| 	cfgToml, err := opts.getConfig() | ||||
| 	if err != nil { | ||||
| 		return fmt.Errorf("unable to fix comments: %v", err) | ||||
| 		return fmt.Errorf("failed to load config: %v", err) | ||||
| 	} | ||||
| 
 | ||||
| 	if _, err := opts.Write(contents); err != nil { | ||||
| 	if _, err := opts.Write(cfgToml); err != nil { | ||||
| 		return fmt.Errorf("unable to write to output: %v", err) | ||||
| 	} | ||||
| 
 | ||||
| 	return nil | ||||
| } | ||||
| 
 | ||||
| // getFormattedConfig returns the default config formatted as required from the specified config file.
 | ||||
| // The config is then formatted as required.
 | ||||
| // No indentation is used and comments are modified so that there is no space
 | ||||
| // after the '#' character.
 | ||||
| func (opts options) getFormattedConfig() ([]byte, error) { | ||||
| 	cfg, err := config.Load(opts.config) | ||||
| 	if err != nil { | ||||
| 		return nil, fmt.Errorf("unable to load or create config: %v", err) | ||||
| 	} | ||||
| 
 | ||||
| 	buffer := bytes.NewBuffer(nil) | ||||
| 
 | ||||
| 	if _, err := cfg.Save(buffer); err != nil { | ||||
| 		return nil, fmt.Errorf("unable to save config: %v", err) | ||||
| 	} | ||||
| 	return fixComments(buffer.Bytes()) | ||||
| } | ||||
| 
 | ||||
| func fixComments(contents []byte) ([]byte, error) { | ||||
| 	r, err := regexp.Compile(`(\n*)\s*?#\s*(\S.*)`) | ||||
| 	if err != nil { | ||||
| 		return nil, fmt.Errorf("unable to compile regexp: %v", err) | ||||
| 	} | ||||
| 	replaced := r.ReplaceAll(contents, []byte("$1#$2")) | ||||
| 
 | ||||
| 	return replaced, nil | ||||
| } | ||||
| 
 | ||||
| func (opts options) outputExists() (bool, error) { | ||||
| 	if opts.output == "" { | ||||
| 		return false, nil | ||||
| 	} | ||||
| 	_, err := os.Stat(opts.output) | ||||
| 	if err == nil { | ||||
| 		return true, nil | ||||
| 	} else if !os.IsNotExist(err) { | ||||
| 		return false, fmt.Errorf("unable to stat output file: %v", err) | ||||
| 	} | ||||
| 	return false, nil | ||||
| // getConfig returns the TOML config for the specified options.
 | ||||
| func (opts options) getConfig() (*config.Toml, error) { | ||||
| 	return config.New( | ||||
| 		config.WithConfigFile(opts.config), | ||||
| 	) | ||||
| } | ||||
| 
 | ||||
| func (opts options) ensureOutputFolder() error { | ||||
| @ -166,7 +130,7 @@ func (opts options) ensureOutputFolder() error { | ||||
| } | ||||
| 
 | ||||
| // Write writes the contents to the output file specified in the options.
 | ||||
| func (opts options) Write(contents []byte) (int, error) { | ||||
| func (opts options) Write(cfg *config.Toml) (int, error) { | ||||
| 	var output io.Writer | ||||
| 	if opts.output == "" { | ||||
| 		output = os.Stdout | ||||
| @ -179,5 +143,6 @@ func (opts options) Write(contents []byte) (int, error) { | ||||
| 		output = outputFile | ||||
| 	} | ||||
| 
 | ||||
| 	return output.Write(contents) | ||||
| 	n, err := cfg.Save(output) | ||||
| 	return int(n), err | ||||
| } | ||||
|  | ||||
| @ -1,82 +0,0 @@ | ||||
| /** | ||||
| # Copyright (c) NVIDIA CORPORATION.  All rights reserved. | ||||
| # | ||||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||||
| # you may not use this file except in compliance with the License. | ||||
| # You may obtain a copy of the License at | ||||
| # | ||||
| #     http://www.apache.org/licenses/LICENSE-2.0
 | ||||
| # | ||||
| # Unless required by applicable law or agreed to in writing, software | ||||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||||
| # See the License for the specific language governing permissions and | ||||
| # limitations under the License. | ||||
| **/ | ||||
| 
 | ||||
| package defaultsubcommand | ||||
| 
 | ||||
| import ( | ||||
| 	"strings" | ||||
| 	"testing" | ||||
| 
 | ||||
| 	"github.com/stretchr/testify/require" | ||||
| ) | ||||
| 
 | ||||
| func TestFixComment(t *testing.T) { | ||||
| 	testCases := []struct { | ||||
| 		input    string | ||||
| 		expected string | ||||
| 	}{ | ||||
| 		{ | ||||
| 			input:    "# comment", | ||||
| 			expected: "#comment", | ||||
| 		}, | ||||
| 		{ | ||||
| 			input:    " #comment", | ||||
| 			expected: "#comment", | ||||
| 		}, | ||||
| 		{ | ||||
| 			input:    " # comment", | ||||
| 			expected: "#comment", | ||||
| 		}, | ||||
| 		{ | ||||
| 			input: strings.Join([]string{ | ||||
| 				"some", | ||||
| 				"# comment", | ||||
| 				" # comment", | ||||
| 				" #comment", | ||||
| 				"other"}, "\n"), | ||||
| 			expected: strings.Join([]string{ | ||||
| 				"some", | ||||
| 				"#comment", | ||||
| 				"#comment", | ||||
| 				"#comment", | ||||
| 				"other"}, "\n"), | ||||
| 		}, | ||||
| 	} | ||||
| 
 | ||||
| 	for _, tc := range testCases { | ||||
| 		t.Run(tc.input, func(t *testing.T) { | ||||
| 			actual, _ := fixComments([]byte(tc.input)) | ||||
| 			require.Equal(t, tc.expected, string(actual)) | ||||
| 		}) | ||||
| 	} | ||||
| } | ||||
| 
 | ||||
| func TestGetFormattedConfig(t *testing.T) { | ||||
| 	expectedLines := []string{ | ||||
| 		"#no-cgroups = false", | ||||
| 		"#debug = \"/var/log/nvidia-container-toolkit.log\"", | ||||
| 		"#debug = \"/var/log/nvidia-container-runtime.log\"", | ||||
| 	} | ||||
| 
 | ||||
| 	opts := &options{} | ||||
| 	contents, err := opts.getFormattedConfig() | ||||
| 	require.NoError(t, err) | ||||
| 	lines := strings.Split(string(contents), "\n") | ||||
| 
 | ||||
| 	for _, line := range expectedLines { | ||||
| 		require.Contains(t, lines, line) | ||||
| 	} | ||||
| } | ||||
| @ -18,8 +18,6 @@ package config | ||||
| 
 | ||||
| import ( | ||||
| 	"bufio" | ||||
| 	"fmt" | ||||
| 	"io" | ||||
| 	"os" | ||||
| 	"path/filepath" | ||||
| 	"strings" | ||||
| @ -28,7 +26,6 @@ import ( | ||||
| 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger" | ||||
| 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup" | ||||
| 	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi" | ||||
| 	"github.com/pelletier/go-toml" | ||||
| ) | ||||
| 
 | ||||
| const ( | ||||
| @ -79,58 +76,14 @@ func GetConfigFilePath() string { | ||||
| // GetConfig sets up the config struct. Values are read from a toml file
 | ||||
| // or set via the environment.
 | ||||
| func GetConfig() (*Config, error) { | ||||
| 	return Load(GetConfigFilePath()) | ||||
| } | ||||
| 
 | ||||
| // Load loads the config from the specified file path.
 | ||||
| func Load(configFilePath string) (*Config, error) { | ||||
| 	if configFilePath == "" { | ||||
| 		return GetDefault() | ||||
| 	} | ||||
| 
 | ||||
| 	tomlFile, err := os.Open(configFilePath) | ||||
| 	if err != nil { | ||||
| 		return GetDefault() | ||||
| 	} | ||||
| 	defer tomlFile.Close() | ||||
| 
 | ||||
| 	cfg, err := LoadFrom(tomlFile) | ||||
| 	if err != nil { | ||||
| 		return nil, fmt.Errorf("failed to read config values: %v", err) | ||||
| 	} | ||||
| 
 | ||||
| 	return cfg, nil | ||||
| } | ||||
| 
 | ||||
| // LoadFrom reads the config from the specified Reader
 | ||||
| func LoadFrom(reader io.Reader) (*Config, error) { | ||||
| 	var tree *toml.Tree | ||||
| 	if reader != nil { | ||||
| 		toml, err := toml.LoadReader(reader) | ||||
| 		if err != nil { | ||||
| 			return nil, err | ||||
| 		} | ||||
| 		tree = toml | ||||
| 	} | ||||
| 
 | ||||
| 	return getFromTree(tree) | ||||
| } | ||||
| 
 | ||||
| // getFromTree reads the nvidia container runtime config from the specified toml Tree.
 | ||||
| func getFromTree(toml *toml.Tree) (*Config, error) { | ||||
| 	cfg, err := GetDefault() | ||||
| 	cfg, err := New( | ||||
| 		WithConfigFile(GetConfigFilePath()), | ||||
| 	) | ||||
| 	if err != nil { | ||||
| 		return nil, err | ||||
| 	} | ||||
| 	if toml == nil { | ||||
| 		return cfg, nil | ||||
| 	} | ||||
| 
 | ||||
| 	if err := toml.Unmarshal(cfg); err != nil { | ||||
| 		return nil, fmt.Errorf("failed to unmarshal config: %v", err) | ||||
| 	} | ||||
| 
 | ||||
| 	return cfg, nil | ||||
| 	return cfg.Config() | ||||
| } | ||||
| 
 | ||||
| // GetDefault defines the default values for the config
 | ||||
| @ -260,64 +213,3 @@ func resolveWithDefault(logger logger.Interface, label string, path string, defa | ||||
| 
 | ||||
| 	return resolvedPath | ||||
| } | ||||
| 
 | ||||
| func (c Config) asCommentedToml() (*toml.Tree, error) { | ||||
| 	contents, err := toml.Marshal(c) | ||||
| 	if err != nil { | ||||
| 		return nil, err | ||||
| 	} | ||||
| 	asToml, err := toml.LoadBytes(contents) | ||||
| 	if err != nil { | ||||
| 		return nil, err | ||||
| 	} | ||||
| 
 | ||||
| 	commentedDefaults := map[string]interface{}{ | ||||
| 		"swarm-resource": "DOCKER_RESOURCE_GPU", | ||||
| 		"accept-nvidia-visible-devices-envvar-when-unprivileged": true, | ||||
| 		"accept-nvidia-visible-devices-as-volume-mounts":         false, | ||||
| 		"nvidia-container-cli.root":                              "/run/nvidia/driver", | ||||
| 		"nvidia-container-cli.path":                              "/usr/bin/nvidia-container-cli", | ||||
| 		"nvidia-container-cli.debug":                             "/var/log/nvidia-container-toolkit.log", | ||||
| 		"nvidia-container-cli.ldcache":                           "/etc/ld.so.cache", | ||||
| 		"nvidia-container-cli.no-cgroups":                        false, | ||||
| 		"nvidia-container-cli.user":                              "root:video", | ||||
| 		"nvidia-container-runtime.debug":                         "/var/log/nvidia-container-runtime.log", | ||||
| 	} | ||||
| 	for k, v := range commentedDefaults { | ||||
| 		set := asToml.Get(k) | ||||
| 		if !shouldComment(k, v, set) { | ||||
| 			continue | ||||
| 		} | ||||
| 		asToml.SetWithComment(k, "", true, v) | ||||
| 	} | ||||
| 
 | ||||
| 	return asToml, nil | ||||
| } | ||||
| 
 | ||||
| func shouldComment(key string, defaultValue interface{}, setTo interface{}) bool { | ||||
| 	if key == "nvidia-container-cli.user" && !getCommentedUserGroup() { | ||||
| 		return false | ||||
| 	} | ||||
| 	if key == "nvidia-container-runtime.debug" && setTo == "/dev/null" { | ||||
| 		return true | ||||
| 	} | ||||
| 	if setTo == nil || defaultValue == setTo || setTo == "" { | ||||
| 		return true | ||||
| 	} | ||||
| 
 | ||||
| 	return false | ||||
| } | ||||
| 
 | ||||
| // Save writes the config to the specified writer.
 | ||||
| func (c Config) Save(w io.Writer) (int64, error) { | ||||
| 	asToml, err := c.asCommentedToml() | ||||
| 	if err != nil { | ||||
| 		return 0, err | ||||
| 	} | ||||
| 
 | ||||
| 	enc := toml.NewEncoder(w).Indentation("") | ||||
| 	if err := enc.Encode(asToml); err != nil { | ||||
| 		return 0, fmt.Errorf("invalid config: %v", err) | ||||
| 	} | ||||
| 	return 0, nil | ||||
| } | ||||
|  | ||||
| @ -17,8 +17,6 @@ | ||||
| package config | ||||
| 
 | ||||
| import ( | ||||
| 	"bytes" | ||||
| 	"io/ioutil" | ||||
| 	"os" | ||||
| 	"path/filepath" | ||||
| 	"strings" | ||||
| @ -28,24 +26,20 @@ import ( | ||||
| ) | ||||
| 
 | ||||
| func TestGetConfigWithCustomConfig(t *testing.T) { | ||||
| 	wd, err := os.Getwd() | ||||
| 	require.NoError(t, err) | ||||
| 	testDir := t.TempDir() | ||||
| 	t.Setenv(configOverride, testDir) | ||||
| 
 | ||||
| 	filename := filepath.Join(testDir, configFilePath) | ||||
| 
 | ||||
| 	// By default debug is disabled
 | ||||
| 	contents := []byte("[nvidia-container-runtime]\ndebug = \"/nvidia-container-toolkit.log\"") | ||||
| 	testDir := filepath.Join(wd, "test") | ||||
| 	filename := filepath.Join(testDir, configFilePath) | ||||
| 
 | ||||
| 	os.Setenv(configOverride, testDir) | ||||
| 
 | ||||
| 	require.NoError(t, os.MkdirAll(filepath.Dir(filename), 0766)) | ||||
| 	require.NoError(t, ioutil.WriteFile(filename, contents, 0766)) | ||||
| 
 | ||||
| 	defer func() { require.NoError(t, os.RemoveAll(testDir)) }() | ||||
| 	require.NoError(t, os.WriteFile(filename, contents, 0766)) | ||||
| 
 | ||||
| 	cfg, err := GetConfig() | ||||
| 	require.NoError(t, err) | ||||
| 	require.Equal(t, cfg.NVIDIAContainerRuntimeConfig.DebugFilePath, "/nvidia-container-toolkit.log") | ||||
| 	require.Equal(t, "/nvidia-container-toolkit.log", cfg.NVIDIAContainerRuntimeConfig.DebugFilePath) | ||||
| } | ||||
| 
 | ||||
| func TestGetConfig(t *testing.T) { | ||||
| @ -219,12 +213,14 @@ func TestGetConfig(t *testing.T) { | ||||
| 		t.Run(tc.description, func(t *testing.T) { | ||||
| 			reader := strings.NewReader(strings.Join(tc.contents, "\n")) | ||||
| 
 | ||||
| 			cfg, err := LoadFrom(reader) | ||||
| 			tomlCfg, err := loadConfigTomlFrom(reader) | ||||
| 			if tc.expectedError != nil { | ||||
| 				require.Error(t, err) | ||||
| 			} else { | ||||
| 				require.NoError(t, err) | ||||
| 			} | ||||
| 			cfg, err := tomlCfg.Config() | ||||
| 			require.NoError(t, err) | ||||
| 
 | ||||
| 			// We first handle the ldconfig path since this is currently system-dependent.
 | ||||
| 			if tc.inspectLdconfig { | ||||
| @ -240,105 +236,3 @@ func TestGetConfig(t *testing.T) { | ||||
| 		}) | ||||
| 	} | ||||
| } | ||||
| 
 | ||||
| func TestConfigDefault(t *testing.T) { | ||||
| 	config, err := GetDefault() | ||||
| 	require.NoError(t, err) | ||||
| 
 | ||||
| 	buffer := new(bytes.Buffer) | ||||
| 	_, err = config.Save(buffer) | ||||
| 	require.NoError(t, err) | ||||
| 
 | ||||
| 	var lines []string | ||||
| 	for _, l := range strings.Split(buffer.String(), "\n") { | ||||
| 		l = strings.TrimSpace(l) | ||||
| 		if strings.HasPrefix(l, "# ") { | ||||
| 			l = "#" + strings.TrimPrefix(l, "# ") | ||||
| 		} | ||||
| 		lines = append(lines, l) | ||||
| 	} | ||||
| 
 | ||||
| 	// We take the lines from the config that was included in previous packages.
 | ||||
| 	expectedLines := []string{ | ||||
| 		"disable-require = false", | ||||
| 		"#swarm-resource = \"DOCKER_RESOURCE_GPU\"", | ||||
| 		"#accept-nvidia-visible-devices-envvar-when-unprivileged = true", | ||||
| 		"#accept-nvidia-visible-devices-as-volume-mounts = false", | ||||
| 
 | ||||
| 		"#root = \"/run/nvidia/driver\"", | ||||
| 		"#path = \"/usr/bin/nvidia-container-cli\"", | ||||
| 		"environment = []", | ||||
| 		"#debug = \"/var/log/nvidia-container-toolkit.log\"", | ||||
| 		"#ldcache = \"/etc/ld.so.cache\"", | ||||
| 		"load-kmods = true", | ||||
| 		"#no-cgroups = false", | ||||
| 		"#user = \"root:video\"", | ||||
| 
 | ||||
| 		"[nvidia-container-runtime]", | ||||
| 		"#debug = \"/var/log/nvidia-container-runtime.log\"", | ||||
| 		"log-level = \"info\"", | ||||
| 		"mode = \"auto\"", | ||||
| 
 | ||||
| 		"mount-spec-path = \"/etc/nvidia-container-runtime/host-files-for-container.d\"", | ||||
| 	} | ||||
| 
 | ||||
| 	require.Subset(t, lines, expectedLines) | ||||
| } | ||||
| 
 | ||||
| func TestFormat(t *testing.T) { | ||||
| 	testCases := []struct { | ||||
| 		input    string | ||||
| 		expected string | ||||
| 	}{ | ||||
| 		{ | ||||
| 			input:    "# comment", | ||||
| 			expected: "#comment", | ||||
| 		}, | ||||
| 		{ | ||||
| 			input:    " #comment", | ||||
| 			expected: "#comment", | ||||
| 		}, | ||||
| 		{ | ||||
| 			input:    " # comment", | ||||
| 			expected: "#comment", | ||||
| 		}, | ||||
| 		{ | ||||
| 			input: strings.Join([]string{ | ||||
| 				"some", | ||||
| 				"# comment", | ||||
| 				" # comment", | ||||
| 				" #comment", | ||||
| 				"other"}, "\n"), | ||||
| 			expected: strings.Join([]string{ | ||||
| 				"some", | ||||
| 				"#comment", | ||||
| 				"#comment", | ||||
| 				"#comment", | ||||
| 				"other"}, "\n"), | ||||
| 		}, | ||||
| 	} | ||||
| 
 | ||||
| 	for _, tc := range testCases { | ||||
| 		t.Run(tc.input, func(t *testing.T) { | ||||
| 			actual, _ := Config{}.format([]byte(tc.input)) | ||||
| 			require.Equal(t, tc.expected, string(actual)) | ||||
| 		}) | ||||
| 	} | ||||
| } | ||||
| 
 | ||||
| func TestGetFormattedConfig(t *testing.T) { | ||||
| 	expectedLines := []string{ | ||||
| 		"#no-cgroups = false", | ||||
| 		"#debug = \"/var/log/nvidia-container-toolkit.log\"", | ||||
| 		"#debug = \"/var/log/nvidia-container-runtime.log\"", | ||||
| 	} | ||||
| 
 | ||||
| 	config, _ := GetDefault() | ||||
| 	contents, err := config.contents() | ||||
| 	require.NoError(t, err) | ||||
| 	lines := strings.Split(string(contents), "\n") | ||||
| 
 | ||||
| 	for _, line := range expectedLines { | ||||
| 		require.Contains(t, lines, line) | ||||
| 	} | ||||
| } | ||||
|  | ||||
							
								
								
									
										188
									
								
								internal/config/toml.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										188
									
								
								internal/config/toml.go
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,188 @@ | ||||
| /** | ||||
| # Copyright (c) NVIDIA CORPORATION.  All rights reserved. | ||||
| # | ||||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||||
| # you may not use this file except in compliance with the License. | ||||
| # You may obtain a copy of the License at | ||||
| # | ||||
| #     http://www.apache.org/licenses/LICENSE-2.0
 | ||||
| # | ||||
| # Unless required by applicable law or agreed to in writing, software | ||||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||||
| # See the License for the specific language governing permissions and | ||||
| # limitations under the License. | ||||
| **/ | ||||
| 
 | ||||
| package config | ||||
| 
 | ||||
| import ( | ||||
| 	"bytes" | ||||
| 	"fmt" | ||||
| 	"io" | ||||
| 	"os" | ||||
| 	"regexp" | ||||
| 
 | ||||
| 	"github.com/pelletier/go-toml" | ||||
| ) | ||||
| 
 | ||||
| // Toml is a type for the TOML representation of a config.
 | ||||
| type Toml toml.Tree | ||||
| 
 | ||||
| type options struct { | ||||
| 	configFile string | ||||
| } | ||||
| 
 | ||||
| // Option is a functional option for loading TOML config files.
 | ||||
| type Option func(*options) | ||||
| 
 | ||||
| // WithConfigFile sets the config file option.
 | ||||
| func WithConfigFile(configFile string) Option { | ||||
| 	return func(o *options) { | ||||
| 		o.configFile = configFile | ||||
| 	} | ||||
| } | ||||
| 
 | ||||
| // New creates a new toml tree based on the provided options
 | ||||
| func New(opts ...Option) (*Toml, error) { | ||||
| 	o := &options{} | ||||
| 	for _, opt := range opts { | ||||
| 		opt(o) | ||||
| 	} | ||||
| 
 | ||||
| 	return loadConfigToml(o.configFile) | ||||
| } | ||||
| 
 | ||||
| func loadConfigToml(filename string) (*Toml, error) { | ||||
| 	if filename == "" { | ||||
| 		return defaultToml() | ||||
| 	} | ||||
| 
 | ||||
| 	tomlFile, err := os.Open(filename) | ||||
| 	if os.IsNotExist(err) { | ||||
| 		return defaultToml() | ||||
| 	} else if err != nil { | ||||
| 		return nil, fmt.Errorf("failed to load specified config file: %v", err) | ||||
| 	} | ||||
| 	defer tomlFile.Close() | ||||
| 
 | ||||
| 	return loadConfigTomlFrom(tomlFile) | ||||
| 
 | ||||
| } | ||||
| 
 | ||||
| func defaultToml() (*Toml, error) { | ||||
| 	cfg, err := GetDefault() | ||||
| 	if err != nil { | ||||
| 		return nil, err | ||||
| 	} | ||||
| 	contents, err := toml.Marshal(cfg) | ||||
| 	if err != nil { | ||||
| 		return nil, err | ||||
| 	} | ||||
| 
 | ||||
| 	return loadConfigTomlFrom(bytes.NewReader(contents)) | ||||
| } | ||||
| 
 | ||||
| func loadConfigTomlFrom(reader io.Reader) (*Toml, error) { | ||||
| 	tree, err := toml.LoadReader(reader) | ||||
| 	if err != nil { | ||||
| 		return nil, err | ||||
| 	} | ||||
| 	return (*Toml)(tree), nil | ||||
| } | ||||
| 
 | ||||
| // Config returns the typed config associated with the toml tree.
 | ||||
| func (t *Toml) Config() (*Config, error) { | ||||
| 	cfg, err := GetDefault() | ||||
| 	if err != nil { | ||||
| 		return nil, err | ||||
| 	} | ||||
| 	if t == nil { | ||||
| 		return cfg, nil | ||||
| 	} | ||||
| 	if err := t.Unmarshal(cfg); err != nil { | ||||
| 		return nil, fmt.Errorf("failed to unmarshal config: %v", err) | ||||
| 	} | ||||
| 	return cfg, nil | ||||
| } | ||||
| 
 | ||||
| // Unmarshal wraps the toml.Tree Unmarshal function.
 | ||||
| func (t *Toml) Unmarshal(v interface{}) error { | ||||
| 	return (*toml.Tree)(t).Unmarshal(v) | ||||
| } | ||||
| 
 | ||||
| // Save saves the config to the specified Writer.
 | ||||
| func (t *Toml) Save(w io.Writer) (int64, error) { | ||||
| 	contents, err := t.contents() | ||||
| 	if err != nil { | ||||
| 		return 0, err | ||||
| 	} | ||||
| 
 | ||||
| 	n, err := w.Write(contents) | ||||
| 	return int64(n), err | ||||
| } | ||||
| 
 | ||||
| // contents returns the config TOML as a byte slice.
 | ||||
| // Any required formatting is applied.
 | ||||
| func (t Toml) contents() ([]byte, error) { | ||||
| 	commented := t.commentDefaults() | ||||
| 
 | ||||
| 	buffer := bytes.NewBuffer(nil) | ||||
| 
 | ||||
| 	enc := toml.NewEncoder(buffer).Indentation("") | ||||
| 	if err := enc.Encode((*toml.Tree)(commented)); err != nil { | ||||
| 		return nil, fmt.Errorf("invalid config: %v", err) | ||||
| 	} | ||||
| 	return t.format(buffer.Bytes()) | ||||
| } | ||||
| 
 | ||||
| // format fixes the comments for the config to ensure that they start in column
 | ||||
| // 1 and are not followed by a space.
 | ||||
| func (t Toml) format(contents []byte) ([]byte, error) { | ||||
| 	r, err := regexp.Compile(`(\n*)\s*?#\s*(\S.*)`) | ||||
| 	if err != nil { | ||||
| 		return nil, fmt.Errorf("unable to compile regexp: %v", err) | ||||
| 	} | ||||
| 	replaced := r.ReplaceAll(contents, []byte("$1#$2")) | ||||
| 
 | ||||
| 	return replaced, nil | ||||
| } | ||||
| 
 | ||||
| // commentDefaults applies the required comments for default values to the Toml.
 | ||||
| func (t *Toml) commentDefaults() *Toml { | ||||
| 	asToml := (*toml.Tree)(t) | ||||
| 	commentedDefaults := map[string]interface{}{ | ||||
| 		"swarm-resource": "DOCKER_RESOURCE_GPU", | ||||
| 		"accept-nvidia-visible-devices-envvar-when-unprivileged": true, | ||||
| 		"accept-nvidia-visible-devices-as-volume-mounts":         false, | ||||
| 		"nvidia-container-cli.root":                              "/run/nvidia/driver", | ||||
| 		"nvidia-container-cli.path":                              "/usr/bin/nvidia-container-cli", | ||||
| 		"nvidia-container-cli.debug":                             "/var/log/nvidia-container-toolkit.log", | ||||
| 		"nvidia-container-cli.ldcache":                           "/etc/ld.so.cache", | ||||
| 		"nvidia-container-cli.no-cgroups":                        false, | ||||
| 		"nvidia-container-cli.user":                              "root:video", | ||||
| 		"nvidia-container-runtime.debug":                         "/var/log/nvidia-container-runtime.log", | ||||
| 	} | ||||
| 	for k, v := range commentedDefaults { | ||||
| 		set := asToml.Get(k) | ||||
| 		if !shouldComment(k, v, set) { | ||||
| 			continue | ||||
| 		} | ||||
| 		asToml.SetWithComment(k, "", true, v) | ||||
| 	} | ||||
| 	return (*Toml)(asToml) | ||||
| } | ||||
| 
 | ||||
| func shouldComment(key string, defaultValue interface{}, setTo interface{}) bool { | ||||
| 	if key == "nvidia-container-cli.user" && !getCommentedUserGroup() { | ||||
| 		return false | ||||
| 	} | ||||
| 	if key == "nvidia-container-runtime.debug" && setTo == "/dev/null" { | ||||
| 		return true | ||||
| 	} | ||||
| 	if setTo == nil || defaultValue == setTo || setTo == "" { | ||||
| 		return true | ||||
| 	} | ||||
| 
 | ||||
| 	return false | ||||
| } | ||||
							
								
								
									
										248
									
								
								internal/config/toml_test.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										248
									
								
								internal/config/toml_test.go
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,248 @@ | ||||
| /** | ||||
| # Copyright (c) NVIDIA CORPORATION.  All rights reserved. | ||||
| # | ||||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||||
| # you may not use this file except in compliance with the License. | ||||
| # You may obtain a copy of the License at | ||||
| # | ||||
| #     http://www.apache.org/licenses/LICENSE-2.0
 | ||||
| # | ||||
| # Unless required by applicable law or agreed to in writing, software | ||||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||||
| # See the License for the specific language governing permissions and | ||||
| # limitations under the License. | ||||
| **/ | ||||
| 
 | ||||
| package config | ||||
| 
 | ||||
| import ( | ||||
| 	"bytes" | ||||
| 	"strings" | ||||
| 	"testing" | ||||
| 
 | ||||
| 	"github.com/pelletier/go-toml" | ||||
| 	"github.com/stretchr/testify/require" | ||||
| ) | ||||
| 
 | ||||
| func TestTomlSave(t *testing.T) { | ||||
| 	testCases := []struct { | ||||
| 		description string | ||||
| 		config      *Toml | ||||
| 		expected    string | ||||
| 	}{ | ||||
| 		{ | ||||
| 			description: "defaultConfig", | ||||
| 			config: func() *Toml { | ||||
| 				t, _ := defaultToml() | ||||
| 				// TODO: We handle the ldconfig path specifically, since this is platform
 | ||||
| 				// dependent.
 | ||||
| 				(*toml.Tree)(t).Set("nvidia-container-cli.ldconfig", "OVERRIDDEN") | ||||
| 				return t | ||||
| 			}(), | ||||
| 			expected: ` | ||||
| #accept-nvidia-visible-devices-as-volume-mounts = false | ||||
| #accept-nvidia-visible-devices-envvar-when-unprivileged = true | ||||
| disable-require = false | ||||
| supported-driver-capabilities = "compat32,compute,display,graphics,ngx,utility,video" | ||||
| #swarm-resource = "DOCKER_RESOURCE_GPU" | ||||
| 
 | ||||
| [nvidia-container-cli] | ||||
| #debug = "/var/log/nvidia-container-toolkit.log" | ||||
| environment = [] | ||||
| #ldcache = "/etc/ld.so.cache" | ||||
| ldconfig = "OVERRIDDEN" | ||||
| load-kmods = true | ||||
| #no-cgroups = false | ||||
| #path = "/usr/bin/nvidia-container-cli" | ||||
| #root = "/run/nvidia/driver" | ||||
| #user = "root:video" | ||||
| 
 | ||||
| [nvidia-container-runtime] | ||||
| #debug = "/var/log/nvidia-container-runtime.log" | ||||
| log-level = "info" | ||||
| mode = "auto" | ||||
| runtimes = ["docker-runc", "runc"] | ||||
| 
 | ||||
| [nvidia-container-runtime.modes] | ||||
| 
 | ||||
| [nvidia-container-runtime.modes.cdi] | ||||
| annotation-prefixes = ["cdi.k8s.io/"] | ||||
| default-kind = "nvidia.com/gpu" | ||||
| spec-dirs = ["/etc/cdi", "/var/run/cdi"] | ||||
| 
 | ||||
| [nvidia-container-runtime.modes.csv] | ||||
| mount-spec-path = "/etc/nvidia-container-runtime/host-files-for-container.d" | ||||
| 
 | ||||
| [nvidia-container-runtime-hook] | ||||
| path = "nvidia-container-runtime-hook" | ||||
| skip-mode-detection = false | ||||
| 
 | ||||
| [nvidia-ctk] | ||||
| path = "nvidia-ctk" | ||||
| `, | ||||
| 		}, | ||||
| 	} | ||||
| 
 | ||||
| 	for _, tc := range testCases { | ||||
| 		t.Run(tc.description, func(t *testing.T) { | ||||
| 			buffer := new(bytes.Buffer) | ||||
| 			_, err := tc.config.Save(buffer) | ||||
| 			require.NoError(t, err) | ||||
| 
 | ||||
| 			require.EqualValues(t, | ||||
| 				strings.TrimSpace(tc.expected), | ||||
| 				strings.TrimSpace(buffer.String()), | ||||
| 			) | ||||
| 		}) | ||||
| 	} | ||||
| } | ||||
| 
 | ||||
| func TestFormat(t *testing.T) { | ||||
| 	testCases := []struct { | ||||
| 		input    string | ||||
| 		expected string | ||||
| 	}{ | ||||
| 		{ | ||||
| 			input:    "# comment", | ||||
| 			expected: "#comment", | ||||
| 		}, | ||||
| 		{ | ||||
| 			input:    " #comment", | ||||
| 			expected: "#comment", | ||||
| 		}, | ||||
| 		{ | ||||
| 			input:    " # comment", | ||||
| 			expected: "#comment", | ||||
| 		}, | ||||
| 		{ | ||||
| 			input: strings.Join([]string{ | ||||
| 				"some", | ||||
| 				"# comment", | ||||
| 				" # comment", | ||||
| 				" #comment", | ||||
| 				"other"}, "\n"), | ||||
| 			expected: strings.Join([]string{ | ||||
| 				"some", | ||||
| 				"#comment", | ||||
| 				"#comment", | ||||
| 				"#comment", | ||||
| 				"other"}, "\n"), | ||||
| 		}, | ||||
| 	} | ||||
| 
 | ||||
| 	for _, tc := range testCases { | ||||
| 		t.Run(tc.input, func(t *testing.T) { | ||||
| 			actual, _ := (Toml{}).format([]byte(tc.input)) | ||||
| 			require.Equal(t, tc.expected, string(actual)) | ||||
| 		}) | ||||
| 	} | ||||
| } | ||||
| 
 | ||||
| func TestGetFormattedConfig(t *testing.T) { | ||||
| 	expectedLines := []string{ | ||||
| 		"#no-cgroups = false", | ||||
| 		"#debug = \"/var/log/nvidia-container-toolkit.log\"", | ||||
| 		"#debug = \"/var/log/nvidia-container-runtime.log\"", | ||||
| 	} | ||||
| 
 | ||||
| 	contents, err := createEmpty().contents() | ||||
| 	require.NoError(t, err) | ||||
| 	lines := strings.Split(string(contents), "\n") | ||||
| 
 | ||||
| 	for _, line := range expectedLines { | ||||
| 		require.Contains(t, lines, line) | ||||
| 	} | ||||
| } | ||||
| 
 | ||||
| func TestTomlContents(t *testing.T) { | ||||
| 	testCases := []struct { | ||||
| 		description string | ||||
| 		contents    map[string]interface{} | ||||
| 		expected    string | ||||
| 	}{ | ||||
| 		{ | ||||
| 			description: "empty config returns commented defaults", | ||||
| 			expected: ` | ||||
| #accept-nvidia-visible-devices-as-volume-mounts = false | ||||
| #accept-nvidia-visible-devices-envvar-when-unprivileged = true | ||||
| #swarm-resource = "DOCKER_RESOURCE_GPU" | ||||
| 
 | ||||
| [nvidia-container-cli] | ||||
| #debug = "/var/log/nvidia-container-toolkit.log" | ||||
| #ldcache = "/etc/ld.so.cache" | ||||
| #no-cgroups = false | ||||
| #path = "/usr/bin/nvidia-container-cli" | ||||
| #root = "/run/nvidia/driver" | ||||
| #user = "root:video" | ||||
| 
 | ||||
| [nvidia-container-runtime] | ||||
| #debug = "/var/log/nvidia-container-runtime.log"`, | ||||
| 		}, | ||||
| 	} | ||||
| 
 | ||||
| 	for _, tc := range testCases { | ||||
| 		t.Run(tc.description, func(t *testing.T) { | ||||
| 			tree, err := toml.TreeFromMap(tc.contents) | ||||
| 			require.NoError(t, err) | ||||
| 			cfg := (*Toml)(tree) | ||||
| 			contents, err := cfg.contents() | ||||
| 			require.NoError(t, err) | ||||
| 
 | ||||
| 			require.EqualValues(t, | ||||
| 				strings.TrimSpace(tc.expected), | ||||
| 				strings.TrimSpace(string(contents)), | ||||
| 			) | ||||
| 		}) | ||||
| 	} | ||||
| } | ||||
| 
 | ||||
| func TestConfigFromToml(t *testing.T) { | ||||
| 	testCases := []struct { | ||||
| 		description    string | ||||
| 		contents       map[string]interface{} | ||||
| 		expectedConfig *Config | ||||
| 	}{ | ||||
| 		{ | ||||
| 			description: "empty config returns default config", | ||||
| 			contents:    nil, | ||||
| 			expectedConfig: func() *Config { | ||||
| 				c, _ := GetDefault() | ||||
| 				return c | ||||
| 			}(), | ||||
| 		}, | ||||
| 		{ | ||||
| 			description: "contents overrides default", | ||||
| 			contents: map[string]interface{}{ | ||||
| 				"nvidia-container-runtime": map[string]interface{}{ | ||||
| 					"debug": "/some/log/file.log", | ||||
| 					"mode":  "csv", | ||||
| 				}, | ||||
| 			}, | ||||
| 			expectedConfig: func() *Config { | ||||
| 				c, _ := GetDefault() | ||||
| 				c.NVIDIAContainerRuntimeConfig.DebugFilePath = "/some/log/file.log" | ||||
| 				c.NVIDIAContainerRuntimeConfig.Mode = "csv" | ||||
| 				return c | ||||
| 			}(), | ||||
| 		}, | ||||
| 	} | ||||
| 
 | ||||
| 	for _, tc := range testCases { | ||||
| 		t.Run(tc.description, func(t *testing.T) { | ||||
| 			tomlCfg := fromMap(tc.contents) | ||||
| 			config, err := tomlCfg.Config() | ||||
| 			require.NoError(t, err) | ||||
| 			require.EqualValues(t, tc.expectedConfig, config) | ||||
| 		}) | ||||
| 	} | ||||
| } | ||||
| 
 | ||||
| func fromMap(c map[string]interface{}) *Toml { | ||||
| 	tree, _ := toml.TreeFromMap(c) | ||||
| 	return (*Toml)(tree) | ||||
| } | ||||
| 
 | ||||
| func createEmpty() *Toml { | ||||
| 	return fromMap(nil) | ||||
| } | ||||
		Loading…
	
		Reference in New Issue
	
	Block a user