Merge pull request #643 from elezar/refactor-toml-source

Refactor handling of TOML config files for runtimes

cmd/nvidia-ctk/runtime/configure/configure.go

@@ -17,7 +17,6 @@
 package configure
 
 import (
-	"encoding/json"
 	"fmt"
 	"path/filepath"
 
@@ -29,6 +28,7 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/crio"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/docker"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/ocihook"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
 const (
@@ -156,13 +156,6 @@ func (m command) build() *cli.Command {
 			Usage:       "Enable CDI in the configured runtime",
 			Destination: &config.cdi.enabled,
 		},
-		&cli.StringFlag{
-			Name:        "runtime-config-override",
-			Destination: &config.runtimeConfigOverrideJSON,
-			Usage:       "specify additional runtime options as a JSON string. The paths are relative to the runtime config.",
-			Value:       "",
-			EnvVars:     []string{"RUNTIME_CONFIG_OVERRIDE"},
-		},
 	}
 
 	return &configure
@@ -234,11 +227,13 @@ func (m command) configureConfigFile(c *cli.Context, config *config) error {
 		cfg, err = containerd.New(
 			containerd.WithLogger(m.logger),
 			containerd.WithPath(configFilePath),
+			containerd.WithConfigSource(toml.FromFile(configFilePath)),
 		)
 	case "crio":
 		cfg, err = crio.New(
 			crio.WithLogger(m.logger),
 			crio.WithPath(configFilePath),
+			crio.WithConfigSource(toml.FromFile(configFilePath)),
 		)
 	case "docker":
 		cfg, err = docker.New(
@@ -252,16 +247,10 @@ func (m command) configureConfigFile(c *cli.Context, config *config) error {
 		return fmt.Errorf("unable to load config for runtime %v: %v", config.runtime, err)
 	}
 
-	runtimeConfigOverride, err := config.runtimeConfigOverride()
-	if err != nil {
-		return fmt.Errorf("unable to parse config overrides: %w", err)
-	}
-
 	err = cfg.AddRuntime(
 		config.nvidiaRuntime.name,
 		config.nvidiaRuntime.path,
 		config.nvidiaRuntime.setAsDefault,
-		runtimeConfigOverride,
 	)
 	if err != nil {
 		return fmt.Errorf("unable to update config: %v", err)
@@ -314,20 +303,6 @@ func (c *config) getOuputConfigPath() string {
 	return c.resolveConfigFilePath()
 }
 
-// runtimeConfigOverride converts the specified runtimeConfigOverride JSON string to a map.
-func (o *config) runtimeConfigOverride() (map[string]interface{}, error) {
-	if o.runtimeConfigOverrideJSON == "" {
-		return nil, nil
-	}
-
-	runtimeOptions := make(map[string]interface{})
-	if err := json.Unmarshal([]byte(o.runtimeConfigOverrideJSON), &runtimeOptions); err != nil {
-		return nil, fmt.Errorf("failed to read %v as JSON: %w", o.runtimeConfigOverrideJSON, err)
-	}
-
-	return runtimeOptions, nil
-}
-
 // configureOCIHook creates and configures the OCI hook for the NVIDIA runtime
 func (m *command) configureOCIHook(c *cli.Context, config *config) error {
 	err := ocihook.CreateHook(config.hookFilePath, config.nvidiaRuntime.hookPath)

pkg/config/engine/api.go

@@ -19,7 +19,7 @@ package engine
 // Interface defines the API for a runtime config updater.
 type Interface interface {
 	DefaultRuntime() string
-	AddRuntime(string, string, bool, ...map[string]interface{}) error
+	AddRuntime(string, string, bool) error
 	Set(string, interface{})
 	RemoveRuntime(string) error
 	Save(string) (int64, error)

pkg/config/engine/containerd/config_v1.go

@@ -30,7 +30,7 @@ type ConfigV1 Config
 var _ engine.Interface = (*ConfigV1)(nil)
 
 // AddRuntime adds a runtime to the containerd config
-func (c *ConfigV1) AddRuntime(name string, path string, setAsDefault bool, configOverrides ...map[string]interface{}) error {
+func (c *ConfigV1) AddRuntime(name string, path string, setAsDefault bool) error {
 	if c == nil || c.Tree == nil {
 		return fmt.Errorf("config is nil")
 	}
@@ -45,12 +45,14 @@ func (c *ConfigV1) AddRuntime(name string, path string, setAsDefault bool, confi
 		runtimeNamesForConfig = append(runtimeNamesForConfig, name)
 	}
 	for _, r := range runtimeNamesForConfig {
-		if options, ok := config.GetPath([]string{"plugins", "cri", "containerd", "runtimes", r}).(*toml.Tree); ok {
-			c.Logger.Debugf("using options from runtime %v: %v", r, options.String())
-			options, _ = toml.Load(options.String())
-			config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name}, options)
-			break
+		options := config.GetSubtreeByPath([]string{"plugins", "cri", "containerd", "runtimes", r})
+		if options == nil {
+			continue
 		}
+		c.Logger.Debugf("using options from runtime %v: %v", r, options)
+		config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name}, options.Copy())
+		break
+
 	}
 
 	if config.GetPath([]string{"plugins", "cri", "containerd", "runtimes", name}) == nil {
@@ -85,16 +87,6 @@ func (c *ConfigV1) AddRuntime(name string, path string, setAsDefault bool, confi
 		}
 		config.SetPath([]string{"plugins", "cri", "containerd", "default_runtime", "options", "BinaryName"}, path)
 		config.SetPath([]string{"plugins", "cri", "containerd", "default_runtime", "options", "Runtime"}, path)
-
-		defaultRuntimeSubtree := subtreeAtPath(config, "plugins", "cri", "containerd", "default_runtime")
-		if err := defaultRuntimeSubtree.applyOverrides(configOverrides...); err != nil {
-			return fmt.Errorf("failed to apply config overrides to default_runtime: %w", err)
-		}
-	}
-
-	runtimeSubtree := subtreeAtPath(config, "plugins", "cri", "containerd", "runtimes", name)
-	if err := runtimeSubtree.applyOverrides(configOverrides...); err != nil {
-		return fmt.Errorf("failed to apply config overrides: %w", err)
 	}
 
 	*c.Tree = config

pkg/config/engine/containerd/config_v1_test.go

@@ -19,20 +19,20 @@ package containerd
 import (
 	"testing"
 
-	"github.com/pelletier/go-toml"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
 func TestAddRuntimeV1(t *testing.T) {
 	logger, _ := testlog.NewNullLogger()
 	testCases := []struct {
-		description     string
-		config          string
-		setAsDefault    bool
-		configOverrides []map[string]interface{}
-		expectedConfig  string
-		expectedError   error
+		description    string
+		config         string
+		setAsDefault   bool
+		expectedConfig string
+		expectedError  error
 	}{
 		{
 			description: "empty config not default runtime",
@@ -53,32 +53,6 @@ func TestAddRuntimeV1(t *testing.T) {
 			`,
 			expectedError: nil,
 		},
-		{
-			description: "empty config not default runtime with overrides",
-			configOverrides: []map[string]interface{}{
-				{
-					"options": map[string]interface{}{
-						"SystemdCgroup": true,
-					},
-				},
-			},
-			expectedConfig: `
-			version = 1
-			[plugins]
-			[plugins.cri]
-				[plugins.cri.containerd]
-				[plugins.cri.containerd.runtimes]
-					[plugins.cri.containerd.runtimes.test]
-					privileged_without_host_devices = false
-					runtime_engine = ""
-					runtime_root = ""
-					runtime_type = ""
-					[plugins.cri.containerd.runtimes.test.options]
-						BinaryName = "/usr/bin/test"
-						Runtime = "/usr/bin/test"
-						SystemdCgroup = true
-			`,
-		},
 		{
 			description: "options from runc are imported",
 			config: `
@@ -226,20 +200,20 @@ func TestAddRuntimeV1(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			config, err := toml.Load(tc.config)
+			cfg, err := toml.Load(tc.config)
 			require.NoError(t, err)
 			expectedConfig, err := toml.Load(tc.expectedConfig)
 			require.NoError(t, err)
 
 			c := &ConfigV1{
 				Logger: logger,
-				Tree:   config,
+				Tree:   cfg,
 			}
 
-			err = c.AddRuntime("test", "/usr/bin/test", tc.setAsDefault, tc.configOverrides...)
+			err = c.AddRuntime("test", "/usr/bin/test", tc.setAsDefault)
 			require.NoError(t, err)
 
-			require.EqualValues(t, expectedConfig.String(), config.String())
+			require.EqualValues(t, expectedConfig.String(), cfg.String())
 		})
 	}
 }

pkg/config/engine/containerd/config_v2.go

@@ -19,13 +19,11 @@ package containerd
 import (
 	"fmt"
 
-	"github.com/pelletier/go-toml"
-
-	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
 // AddRuntime adds a runtime to the containerd config
-func (c *Config) AddRuntime(name string, path string, setAsDefault bool, configOverrides ...map[string]interface{}) error {
+func (c *Config) AddRuntime(name string, path string, setAsDefault bool) error {
 	if c == nil || c.Tree == nil {
 		return fmt.Errorf("config is nil")
 	}
@@ -39,12 +37,13 @@ func (c *Config) AddRuntime(name string, path string, setAsDefault bool, configO
 		runtimeNamesForConfig = append(runtimeNamesForConfig, name)
 	}
 	for _, r := range runtimeNamesForConfig {
-		if options, ok := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", r}).(*toml.Tree); ok {
-			c.Logger.Debugf("using options from runtime %v: %v", r, options.String())
-			options, _ = toml.Load(options.String())
-			config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name}, options)
-			break
+		options := config.GetSubtreeByPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", r})
+		if options == nil {
+			continue
 		}
+		c.Logger.Debugf("using options from runtime %v: %v", r, options)
+		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name}, options.Copy())
+		break
 	}
 
 	if config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name}) == nil {
@@ -70,11 +69,6 @@ func (c *Config) AddRuntime(name string, path string, setAsDefault bool, configO
 		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"}, name)
 	}
 
-	runtimeSubtree := subtreeAtPath(config, "plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name)
-	if err := runtimeSubtree.applyOverrides(configOverrides...); err != nil {
-		return fmt.Errorf("failed to apply config overrides: %w", err)
-	}
-
 	*c.Tree = config
 	return nil
 }
@@ -151,15 +145,3 @@ func (c *Config) RemoveRuntime(name string) error {
 	*c.Tree = config
 	return nil
 }
-
-// Save writes the config to the specified path
-func (c Config) Save(path string) (int64, error) {
-	config := c.Tree
-	output, err := config.Marshal()
-	if err != nil {
-		return 0, fmt.Errorf("unable to convert to TOML: %v", err)
-	}
-
-	n, err := engine.Config(path).Write(output)
-	return int64(n), err
-}

pkg/config/engine/containerd/config_v2_test.go

@@ -19,20 +19,20 @@ package containerd
 import (
 	"testing"
 
-	"github.com/pelletier/go-toml"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
 func TestAddRuntime(t *testing.T) {
 	logger, _ := testlog.NewNullLogger()
 	testCases := []struct {
-		description     string
-		config          string
-		setAsDefault    bool
-		configOverrides []map[string]interface{}
-		expectedConfig  string
-		expectedError   error
+		description    string
+		config         string
+		setAsDefault   bool
+		expectedConfig string
+		expectedError  error
 	}{
 		{
 			description: "empty config not default runtime",
@@ -52,31 +52,6 @@ func TestAddRuntime(t *testing.T) {
 			`,
 			expectedError: nil,
 		},
-		{
-			description: "empty config not default runtime with overrides",
-			configOverrides: []map[string]interface{}{
-				{
-					"options": map[string]interface{}{
-						"SystemdCgroup": true,
-					},
-				},
-			},
-			expectedConfig: `
-			version = 2
-			[plugins]
-			[plugins."io.containerd.grpc.v1.cri"]
-				[plugins."io.containerd.grpc.v1.cri".containerd]
-				[plugins."io.containerd.grpc.v1.cri".containerd.runtimes]
-					[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.test]
-					privileged_without_host_devices = false
-					runtime_engine = ""
-					runtime_root = ""
-					runtime_type = ""
-					[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.test.options]
-						BinaryName = "/usr/bin/test"
-						SystemdCgroup = true
-			`,
-		},
 		{
 			description: "options from runc are imported",
 			config: `
@@ -224,20 +199,20 @@ func TestAddRuntime(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			config, err := toml.Load(tc.config)
+			cfg, err := toml.Load(tc.config)
 			require.NoError(t, err)
 			expectedConfig, err := toml.Load(tc.expectedConfig)
 			require.NoError(t, err)
 
 			c := &Config{
 				Logger: logger,
-				Tree:   config,
+				Tree:   cfg,
 			}
 
-			err = c.AddRuntime("test", "/usr/bin/test", tc.setAsDefault, tc.configOverrides...)
+			err = c.AddRuntime("test", "/usr/bin/test", tc.setAsDefault)
 			require.NoError(t, err)
 
-			require.EqualValues(t, expectedConfig.String(), config.String())
+			require.EqualValues(t, expectedConfig.String(), cfg.String())
 		})
 	}
 }

pkg/config/engine/containerd/containerd.go

@@ -17,10 +17,11 @@
 package containerd
 
 import (
-	"github.com/pelletier/go-toml"
+	"fmt"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
 // Config represents the containerd config
@@ -36,14 +37,64 @@ var _ engine.Interface = (*Config)(nil)
 
 // New creates a containerd config with the specified options
 func New(opts ...Option) (engine.Interface, error) {
-	b := &builder{}
+	b := &builder{
+		runtimeType: defaultRuntimeType,
+	}
 	for _, opt := range opts {
 		opt(b)
 	}
-
 	if b.logger == nil {
 		b.logger = logger.New()
 	}
+	if b.configSource == nil {
+		b.configSource = toml.FromFile(b.path)
+	}
 
-	return b.build()
+	tomlConfig, err := b.configSource.Load()
+	if err != nil {
+		return nil, fmt.Errorf("failed to load config: %v", err)
+	}
+
+	cfg := &Config{
+		Tree:                  tomlConfig,
+		Logger:                b.logger,
+		RuntimeType:           b.runtimeType,
+		UseDefaultRuntimeName: b.useLegacyConfig,
+		ContainerAnnotations:  b.containerAnnotations,
+	}
+
+	version, err := cfg.parseVersion(b.useLegacyConfig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse config version: %v", err)
+	}
+	switch version {
+	case 1:
+		return (*ConfigV1)(cfg), nil
+	case 2:
+		return cfg, nil
+	}
+
+	return nil, fmt.Errorf("unsupported config version: %v", version)
+}
+
+// parseVersion returns the version of the config
+func (c *Config) parseVersion(useLegacyConfig bool) (int, error) {
+	defaultVersion := 2
+	if useLegacyConfig {
+		defaultVersion = 1
+	}
+
+	switch v := c.Get("version").(type) {
+	case nil:
+		switch len(c.Keys()) {
+		case 0: // No config exists, or the config file is empty, use version inferred from containerd
+			return defaultVersion, nil
+		default: // A config file exists, has content, and no version is set
+			return 1, nil
+		}
+	case int64:
+		return int(v), nil
+	default:
+		return -1, fmt.Errorf("unsupported type for version field: %v", v)
+	}
 }

pkg/config/engine/containerd/option.go

@@ -17,13 +17,8 @@
 package containerd
 
 import (
-	"fmt"
-	"os"
-
-	"github.com/pelletier/go-toml"
-
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
-	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
 const (
@@ -32,6 +27,7 @@ const (
 
 type builder struct {
 	logger               logger.Interface
+	configSource         toml.Loader
 	path                 string
 	runtimeType          string
 	useLegacyConfig      bool
@@ -55,6 +51,13 @@ func WithPath(path string) Option {
 	}
 }
 
+// WithConfigSource sets the source for the config.
+func WithConfigSource(configSource toml.Loader) Option {
+	return func(b *builder) {
+		b.configSource = configSource
+	}
+}
+
 // WithRuntimeType sets the runtime type for the config builder
 func WithRuntimeType(runtimeType string) Option {
 	return func(b *builder) {
@@ -75,82 +78,3 @@ func WithContainerAnnotations(containerAnnotations ...string) Option {
 		b.containerAnnotations = containerAnnotations
 	}
 }
-
-func (b *builder) build() (engine.Interface, error) {
-	if b.path == "" {
-		return nil, fmt.Errorf("config path is empty")
-	}
-
-	if b.runtimeType == "" {
-		b.runtimeType = defaultRuntimeType
-	}
-
-	config, err := b.loadConfig(b.path)
-	if err != nil {
-		return nil, fmt.Errorf("failed to load config: %v", err)
-	}
-	config.Logger = b.logger
-	config.RuntimeType = b.runtimeType
-	config.UseDefaultRuntimeName = !b.useLegacyConfig
-	config.ContainerAnnotations = b.containerAnnotations
-
-	version, err := config.parseVersion(b.useLegacyConfig)
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse config version: %v", err)
-	}
-	switch version {
-	case 1:
-		return (*ConfigV1)(config), nil
-	case 2:
-		return config, nil
-	}
-
-	return nil, fmt.Errorf("unsupported config version: %v", version)
-}
-
-// loadConfig loads the containerd config from disk
-func (b *builder) loadConfig(config string) (*Config, error) {
-	info, err := os.Stat(config)
-	if os.IsExist(err) && info.IsDir() {
-		return nil, fmt.Errorf("config file is a directory")
-	}
-
-	if os.IsNotExist(err) {
-		b.logger.Infof("Config file does not exist; using empty config")
-		config = "/dev/null"
-	} else {
-		b.logger.Infof("Loading config from %v", config)
-	}
-
-	tomlConfig, err := toml.LoadFile(config)
-	if err != nil {
-		return nil, err
-	}
-
-	cfg := Config{
-		Tree: tomlConfig,
-	}
-	return &cfg, nil
-}
-
-// parseVersion returns the version of the config
-func (c *Config) parseVersion(useLegacyConfig bool) (int, error) {
-	defaultVersion := 2
-	if useLegacyConfig {
-		defaultVersion = 1
-	}
-
-	switch v := c.Get("version").(type) {
-	case nil:
-		switch len(c.Keys()) {
-		case 0: // No config exists, or the config file is empty, use version inferred from containerd
-			return defaultVersion, nil
-		default: // A config file exists, has content, and no version is set
-			return 1, nil
-		}
-	case int64:
-		return int(v), nil
-	default:
-		return -1, fmt.Errorf("unsupported type for version field: %v", v)
-	}
-}

pkg/config/engine/containerd/toml.go

@@ -1,56 +0,0 @@
-/**
-# Copyright 2024 NVIDIA CORPORATION
-#
-# 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 containerd
-
-import (
-	"fmt"
-
-	"github.com/pelletier/go-toml"
-)
-
-// tomlTree is an alias for toml.Tree that allows for extensions.
-type tomlTree toml.Tree
-
-func subtreeAtPath(c toml.Tree, path ...string) *tomlTree {
-	tree := c.GetPath(path).(*toml.Tree)
-	return (*tomlTree)(tree)
-}
-
-func (t *tomlTree) insert(other map[string]interface{}) error {
-
-	for key, value := range other {
-		if insertsubtree, ok := value.(map[string]interface{}); ok {
-			subtree := (*toml.Tree)(t).Get(key).(*toml.Tree)
-			return (*tomlTree)(subtree).insert(insertsubtree)
-		}
-		(*toml.Tree)(t).Set(key, value)
-	}
-	return nil
-}
-
-func (t *tomlTree) applyOverrides(overrides ...map[string]interface{}) error {
-	for _, override := range overrides {
-		subconfig, err := toml.TreeFromMap(override)
-		if err != nil {
-			return fmt.Errorf("invalid toml config: %w", err)
-		}
-		if err := t.insert(subconfig.ToMap()); err != nil {
-			return err
-		}
-	}
-	return nil
-}

pkg/config/engine/crio/crio.go

@@ -19,10 +19,9 @@ package crio
 import (
 	"fmt"
 
-	"github.com/pelletier/go-toml"
-
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
 // Config represents the cri-o config
@@ -39,12 +38,27 @@ func New(opts ...Option) (engine.Interface, error) {
 	for _, opt := range opts {
 		opt(b)
 	}
+	if b.logger == nil {
+		b.logger = logger.New()
+	}
+	if b.configSource == nil {
+		b.configSource = toml.FromFile(b.path)
+	}
 
-	return b.build()
+	tomlConfig, err := b.configSource.Load()
+	if err != nil {
+		return nil, err
+	}
+
+	cfg := Config{
+		Tree:   tomlConfig,
+		Logger: b.logger,
+	}
+	return &cfg, nil
 }
 
 // AddRuntime adds a new runtime to the crio config
-func (c *Config) AddRuntime(name string, path string, setAsDefault bool, _ ...map[string]interface{}) error {
+func (c *Config) AddRuntime(name string, path string, setAsDefault bool) error {
 	if c == nil {
 		return fmt.Errorf("config is nil")
 	}
@@ -115,22 +129,3 @@ func (c *Config) RemoveRuntime(name string) error {
 	*c.Tree = config
 	return nil
 }
-
-// Set sets the specified cri-o option.
-func (c *Config) Set(key string, value interface{}) {
-	config := *c.Tree
-	config.Set(key, value)
-	*c.Tree = config
-}
-
-// Save writes the config to the specified path
-func (c Config) Save(path string) (int64, error) {
-	config := c.Tree
-	output, err := config.Marshal()
-	if err != nil {
-		return 0, fmt.Errorf("unable to convert to TOML: %v", err)
-	}
-
-	n, err := engine.Config(path).Write(output)
-	return int64(n), err
-}

pkg/config/engine/crio/crio_test.go

@@ -19,20 +19,20 @@ package crio
 import (
 	"testing"
 
-	"github.com/pelletier/go-toml"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
 func TestAddRuntime(t *testing.T) {
 	logger, _ := testlog.NewNullLogger()
 	testCases := []struct {
-		description     string
-		config          string
-		setAsDefault    bool
-		configOverrides []map[string]interface{}
-		expectedConfig  string
-		expectedError   error
+		description    string
+		config         string
+		setAsDefault   bool
+		expectedConfig string
+		expectedError  error
 	}{
 		{
 			description: "empty config not default runtime",
@@ -127,20 +127,20 @@ func TestAddRuntime(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			config, err := toml.Load(tc.config)
+			cfg, err := toml.Load(tc.config)
 			require.NoError(t, err)
 			expectedConfig, err := toml.Load(tc.expectedConfig)
 			require.NoError(t, err)
 
 			c := &Config{
 				Logger: logger,
-				Tree:   config,
+				Tree:   cfg,
 			}
 
-			err = c.AddRuntime("test", "/usr/bin/test", tc.setAsDefault, tc.configOverrides...)
+			err = c.AddRuntime("test", "/usr/bin/test", tc.setAsDefault)
 			require.NoError(t, err)
 
-			require.EqualValues(t, expectedConfig.String(), config.String())
+			require.EqualValues(t, expectedConfig.String(), cfg.String())
 		})
 	}
 }

pkg/config/engine/crio/option.go

@@ -17,17 +17,14 @@
 package crio
 
 import (
-	"fmt"
-	"os"
-
-	"github.com/pelletier/go-toml"
-
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
 type builder struct {
-	logger logger.Interface
-	path   string
+	logger       logger.Interface
+	configSource toml.Loader
+	path         string
 }
 
 // Option defines a function that can be used to configure the config builder
@@ -47,48 +44,9 @@ func WithPath(path string) Option {
 	}
 }
 
-func (b *builder) build() (*Config, error) {
-	if b.logger == nil {
-		b.logger = logger.New()
+// WithConfigSource sets the TOML source for the config.
+func WithConfigSource(configSource toml.Loader) Option {
+	return func(b *builder) {
+		b.configSource = configSource
 	}
-	if b.path == "" {
-		empty := toml.Tree{}
-		c := Config{
-			Tree:   &empty,
-			Logger: b.logger,
-		}
-		return &c, nil
-	}
-
-	return b.loadConfig(b.path)
-}
-
-// loadConfig loads the cri-o config from disk
-func (b *builder) loadConfig(config string) (*Config, error) {
-	b.logger.Infof("Loading config: %v", config)
-
-	info, err := os.Stat(config)
-	if os.IsExist(err) && info.IsDir() {
-		return nil, fmt.Errorf("config file is a directory")
-	}
-
-	if os.IsNotExist(err) {
-		b.logger.Infof("Config file does not exist; using empty config")
-		config = "/dev/null"
-	} else {
-		b.logger.Infof("Loading config from %v", config)
-	}
-
-	cfg, err := toml.LoadFile(config)
-	if err != nil {
-		return nil, err
-	}
-
-	b.logger.Infof("Successfully loaded config")
-
-	c := Config{
-		Tree:   cfg,
-		Logger: b.logger,
-	}
-	return &c, nil
 }

pkg/config/engine/docker/docker.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
 )
 
@@ -49,7 +50,7 @@ func New(opts ...Option) (engine.Interface, error) {
 }
 
 // AddRuntime adds a new runtime to the docker config
-func (c *Config) AddRuntime(name string, path string, setAsDefault bool, _ ...map[string]interface{}) error {
+func (c *Config) AddRuntime(name string, path string, setAsDefault bool) error {
 	if c == nil {
 		return fmt.Errorf("config is nil")
 	}
@@ -128,6 +129,6 @@ func (c Config) Save(path string) (int64, error) {
 		return 0, fmt.Errorf("unable to convert to JSON: %v", err)
 	}
 
-	n, err := engine.Config(path).Write(output)
+	n, err := config.Raw(path).Write(output)
 	return int64(n), err
 }

pkg/config/raw.go

@@ -1,5 +1,5 @@
 /**
-# Copyright (c) NVIDIA CORPORATION.  All rights reserved.
+# Copyright 2024 NVIDIA CORPORATION
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
 # limitations under the License.
 **/
 
-package engine
+package config
 
 import (
 	"fmt"
@@ -22,11 +22,11 @@ import (
 	"path/filepath"
 )
 
-// Config represents a runtime config
-type Config string
+// Raw represents a raw config file
+type Raw string
 
 // Write writes the specified contents to a config file.
-func (c Config) Write(output []byte) (int, error) {
+func (c Raw) Write(output []byte) (int, error) {
 	path := string(c)
 	if path == "" {
 		n, err := os.Stdout.Write(output)

pkg/config/toml/source-empty.go

@@ -0,0 +1,35 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# 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 toml
+
+import (
+	"github.com/pelletier/go-toml"
+)
+
+type empty string
+
+var _ Loader = (*empty)(nil)
+
+// Load is a no-op for an empty source.
+func (e empty) Load() (*Tree, error) {
+	return newEmpty(), nil
+}
+
+func newEmpty() *Tree {
+	tomlTree, _ := toml.TreeFromMap(nil)
+	return (*Tree)(tomlTree)
+}

pkg/config/toml/source-file.go

@@ -0,0 +1,40 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# 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 toml
+
+import (
+	"fmt"
+	"os"
+)
+
+type tomlFile string
+
+var _ Loader = (*tomlFile)(nil)
+
+// Load loads the contents of the specified TOML file as a map.
+func (f tomlFile) Load() (*Tree, error) {
+	info, err := os.Stat(string(f))
+	if os.IsExist(err) && info.IsDir() {
+		return nil, fmt.Errorf("config file %s is a directory", string(f))
+	}
+
+	if os.IsNotExist(err) {
+		return Empty.Load()
+	}
+
+	return LoadFile(string(f))
+}

pkg/config/toml/source.go

@@ -0,0 +1,35 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# 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 toml
+
+const (
+	Empty = empty("")
+)
+
+// Loader represents a source for a toml config.
+type Loader interface {
+	Load() (*Tree, error)
+}
+
+// FromFile creates a TOML source from the specified file.
+// If an empty string is passed an empty toml config is used.
+func FromFile(path string) Loader {
+	if path == "" {
+		return Empty
+	}
+	return tomlFile(path)
+}

pkg/config/toml/toml.go

@@ -0,0 +1,159 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# 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 toml
+
+import (
+	"fmt"
+
+	"github.com/pelletier/go-toml"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config"
+)
+
+type Tree toml.Tree
+
+// Copy produces a copy of the contents of the Tree.
+func (t *Tree) Copy() *Tree {
+	copy, _ := Load(t.String())
+	return copy
+}
+
+func (t *Tree) GetSubtreeByPath(keys []string) *Tree {
+	subtree := t.GetPath(keys)
+	if subtree == nil {
+		return nil
+	}
+
+	switch subtree := subtree.(type) {
+	case *toml.Tree:
+		return (*Tree)(subtree)
+	case *Tree:
+		return subtree
+	default:
+		panic(fmt.Errorf("invalid subtree type %T", subtree))
+	}
+}
+
+func (t *Tree) DeletePath(keys []string) error {
+	return (*toml.Tree)(t).DeletePath(keys)
+}
+
+func (t *Tree) HasPath(keys []string) bool {
+	return (*toml.Tree)(t).HasPath(keys)
+}
+
+func (t *Tree) Get(key string) interface{} {
+	return toTreeFromRaw((*toml.Tree)(t).Get(key))
+}
+
+func (t *Tree) GetPath(keys []string) interface{} {
+	return toTreeFromRaw((*toml.Tree)(t).GetPath(keys))
+}
+
+func (t *Tree) SetPath(keys []string, value interface{}) {
+	(*toml.Tree)(t).SetPath(keys, toRawFromTree(value))
+}
+
+func (t *Tree) Set(key string, value interface{}) {
+	(*toml.Tree)(t).Set(key, toRawFromTree(value))
+}
+
+func (t *Tree) Delete(key string) error {
+	return (*toml.Tree)(t).Delete(key)
+}
+
+func (t *Tree) Keys() []string {
+	return (*toml.Tree)(t).Keys()
+}
+
+func (t *Tree) String() string {
+	return (*toml.Tree)(t).String()
+}
+
+func (t *Tree) ToMap() map[string]interface{} {
+	return (*toml.Tree)(t).ToMap()
+}
+
+func (t *Tree) Raw() *toml.Tree {
+	return (*toml.Tree)(t)
+}
+
+func toRawFromTree(value interface{}) interface{} {
+	if tree, ok := value.(*Tree); ok {
+		return (*toml.Tree)(tree)
+	}
+	return value
+}
+
+func toTreeFromRaw(value interface{}) interface{} {
+	if tree, ok := value.(*toml.Tree); ok {
+		return (*Tree)(tree)
+	}
+	return value
+}
+
+func TreeFromMap(m map[string]interface{}) (*Tree, error) {
+	return new(func() (*toml.Tree, error) {
+		return toml.TreeFromMap(m)
+	})
+}
+
+func Load(content string) (*Tree, error) {
+	return new(func() (*toml.Tree, error) {
+		return toml.Load(content)
+	})
+}
+
+func LoadBytes(b []byte) (*Tree, error) {
+	return new(func() (*toml.Tree, error) {
+		return toml.LoadBytes(b)
+	})
+}
+
+func LoadFile(path string) (*Tree, error) {
+	return new(func() (*toml.Tree, error) {
+		return toml.LoadFile(path)
+	})
+}
+
+func LoadMap(m map[string]interface{}) (*Tree, error) {
+	return TreeFromMap(m)
+}
+
+func Marshal(v interface{}) ([]byte, error) {
+	return toml.Marshal(v)
+}
+
+func new(construct func() (*toml.Tree, error)) (*Tree, error) {
+	tomlTree, err := construct()
+	if err != nil {
+		return nil, err
+	}
+	return (*Tree)(tomlTree), nil
+}
+
+// Save writes the config to the specified path
+func (t *Tree) Save(path string) (int64, error) {
+	cfg := (*toml.Tree)(t)
+	output, err := cfg.Marshal()
+	if err != nil {
+		return 0, fmt.Errorf("unable to convert to TOML: %v", err)
+	}
+
+	n, err := config.Raw(path).Write(output)
+	return int64(n), err
+}

tools/container/container.go

@@ -67,8 +67,8 @@ func ParseArgs(c *cli.Context, o *Options) error {
 }
 
 // Configure applies the options to the specified config
-func (o Options) Configure(cfg engine.Interface, configOverrides ...map[string]interface{}) error {
-	err := o.UpdateConfig(cfg, configOverrides...)
+func (o Options) Configure(cfg engine.Interface) error {
+	err := o.UpdateConfig(cfg)
 	if err != nil {
 		return fmt.Errorf("unable to update config: %v", err)
 	}
@@ -98,14 +98,14 @@ func (o Options) flush(cfg engine.Interface) error {
 }
 
 // UpdateConfig updates the specified config to include the nvidia runtimes
-func (o Options) UpdateConfig(cfg engine.Interface, configOverrides ...map[string]interface{}) error {
+func (o Options) UpdateConfig(cfg engine.Interface) error {
 	runtimes := operator.GetRuntimes(
 		operator.WithNvidiaRuntimeName(o.RuntimeName),
 		operator.WithSetAsDefault(o.SetAsDefault),
 		operator.WithRoot(o.RuntimeDir),
 	)
 	for name, runtime := range runtimes {
-		err := cfg.AddRuntime(name, runtime.Path, runtime.SetAsDefault, configOverrides...)
+		err := cfg.AddRuntime(name, runtime.Path, runtime.SetAsDefault)
 		if err != nil {
 			return fmt.Errorf("failed to update runtime %q: %v", name, err)
 		}

tools/container/containerd/config_v1_test.go

@@ -20,11 +20,11 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/pelletier/go-toml"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 )
 
@@ -92,12 +92,12 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 				useLegacyConfig: tc.legacyConfig,
 			}
 
-			config, err := toml.TreeFromMap(map[string]interface{}{})
+			cfg, err := toml.Empty.Load()
 			require.NoError(t, err, "%d: %v", i, tc)
 
 			v1 := &containerd.ConfigV1{
 				Logger:                logger,
-				Tree:                  config,
+				Tree:                  cfg,
 				UseDefaultRuntimeName: !tc.legacyConfig,
 				RuntimeType:           runtimeType,
 			}
@@ -240,12 +240,12 @@ func TestUpdateV1Config(t *testing.T) {
 				},
 			}
 
-			config, err := toml.TreeFromMap(map[string]interface{}{})
+			cfg, err := toml.Empty.Load()
 			require.NoError(t, err)
 
 			v1 := &containerd.ConfigV1{
 				Logger:                logger,
-				Tree:                  config,
+				Tree:                  cfg,
 				UseDefaultRuntimeName: true,
 				RuntimeType:           runtimeType,
 				ContainerAnnotations:  []string{"cdi.k8s.io/*"},
@@ -257,7 +257,7 @@ func TestUpdateV1Config(t *testing.T) {
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
 			require.NoError(t, err)
 
-			require.Equal(t, expected.String(), config.String())
+			require.Equal(t, expected.String(), cfg.String())
 		})
 	}
 }
@@ -401,12 +401,12 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 				},
 			}
 
-			config, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary"))
+			cfg, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary"))
 			require.NoError(t, err)
 
 			v1 := &containerd.ConfigV1{
 				Logger:                logger,
-				Tree:                  config,
+				Tree:                  cfg,
 				UseDefaultRuntimeName: true,
 				RuntimeType:           runtimeType,
 				ContainerAnnotations:  []string{"cdi.k8s.io/*"},
@@ -418,7 +418,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
 			require.NoError(t, err)
 
-			require.Equal(t, expected.String(), config.String())
+			require.Equal(t, expected.String(), cfg.String())
 		})
 	}
 }
@@ -479,14 +479,14 @@ func TestRevertV1Config(t *testing.T) {
 				},
 			}
 
-			config, err := toml.TreeFromMap(tc.config)
+			cfg, err := toml.LoadMap(tc.config)
 			require.NoError(t, err, "%d: %v", i, tc)
 
 			expected, err := toml.TreeFromMap(tc.expected)
 			require.NoError(t, err, "%d: %v", i, tc)
 
 			v1 := &containerd.ConfigV1{
-				Tree:                  config,
+				Tree:                  cfg,
 				UseDefaultRuntimeName: true,
 				RuntimeType:           runtimeType,
 			}
@@ -494,7 +494,7 @@ func TestRevertV1Config(t *testing.T) {
 			err = o.RevertConfig(v1)
 			require.NoError(t, err, "%d: %v", i, tc)
 
-			configContents, _ := toml.Marshal(config)
+			configContents, _ := toml.Marshal(cfg)
 			expectedContents, _ := toml.Marshal(expected)
 
 			require.Equal(t, string(expectedContents), string(configContents), "%d: %v", i, tc)

tools/container/containerd/config_v2_test.go

@@ -20,11 +20,11 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/pelletier/go-toml"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 )
 
@@ -74,19 +74,19 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
 				},
 			}
 
-			config, err := toml.TreeFromMap(map[string]interface{}{})
+			cfg, err := toml.LoadMap(map[string]interface{}{})
 			require.NoError(t, err)
 
 			v2 := &containerd.Config{
 				Logger:      logger,
-				Tree:        config,
+				Tree:        cfg,
 				RuntimeType: runtimeType,
 			}
 
 			err = o.UpdateConfig(v2)
 			require.NoError(t, err)
 
-			defaultRuntimeName := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"})
+			defaultRuntimeName := cfg.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"})
 			require.EqualValues(t, tc.expectedDefaultRuntimeName, defaultRuntimeName)
 		})
 	}
@@ -200,12 +200,12 @@ func TestUpdateV2Config(t *testing.T) {
 				runtimeType: runtimeType,
 			}
 
-			config, err := toml.TreeFromMap(map[string]interface{}{})
+			cfg, err := toml.LoadMap(map[string]interface{}{})
 			require.NoError(t, err)
 
 			v2 := &containerd.Config{
 				Logger:               logger,
-				Tree:                 config,
+				Tree:                 cfg,
 				RuntimeType:          o.runtimeType,
 				ContainerAnnotations: []string{"cdi.k8s.io/*"},
 			}
@@ -216,7 +216,7 @@ func TestUpdateV2Config(t *testing.T) {
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
 			require.NoError(t, err)
 
-			require.Equal(t, expected.String(), config.String())
+			require.Equal(t, expected.String(), cfg.String())
 		})
 	}
 
@@ -355,12 +355,12 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 				},
 			}
 
-			config, err := toml.TreeFromMap(runcConfigMapV2("/runc-binary"))
+			cfg, err := toml.LoadMap(runcConfigMapV2("/runc-binary"))
 			require.NoError(t, err)
 
 			v2 := &containerd.Config{
 				Logger:               logger,
-				Tree:                 config,
+				Tree:                 cfg,
 				RuntimeType:          runtimeType,
 				ContainerAnnotations: []string{"cdi.k8s.io/*"},
 			}
@@ -371,7 +371,7 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
 			require.NoError(t, err)
 
-			require.Equal(t, expected.String(), config.String())
+			require.Equal(t, expected.String(), cfg.String())
 		})
 	}
 }
@@ -427,21 +427,21 @@ func TestRevertV2Config(t *testing.T) {
 				},
 			}
 
-			config, err := toml.TreeFromMap(tc.config)
+			cfg, err := toml.LoadMap(tc.config)
 			require.NoError(t, err)
 
 			expected, err := toml.TreeFromMap(tc.expected)
 			require.NoError(t, err)
 
 			v2 := &containerd.Config{
-				Tree:        config,
+				Tree:        cfg,
 				RuntimeType: runtimeType,
 			}
 
 			err = o.RevertConfig(v2)
 			require.NoError(t, err)
 
-			configContents, _ := toml.Marshal(config)
+			configContents, _ := toml.Marshal(cfg)
 			expectedContents, _ := toml.Marshal(expected)
 
 			require.Equal(t, string(expectedContents), string(configContents))

tools/container/containerd/containerd.go

@@ -17,7 +17,6 @@
 package main
 
 import (
-	"encoding/json"
 	"fmt"
 	"os"
 
@@ -26,6 +25,7 @@ import (
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/info"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 )
 
@@ -190,6 +190,7 @@ func Setup(c *cli.Context, o *options) error {
 
 	cfg, err := containerd.New(
 		containerd.WithPath(o.Config),
+		containerd.WithConfigSource(toml.FromFile(o.Config)),
 		containerd.WithRuntimeType(o.runtimeType),
 		containerd.WithUseLegacyConfig(o.useLegacyConfig),
 		containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...),
@@ -198,11 +199,7 @@ func Setup(c *cli.Context, o *options) error {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	runtimeConfigOverride, err := o.runtimeConfigOverride()
-	if err != nil {
-		return fmt.Errorf("unable to parse config overrides: %w", err)
-	}
-	err = o.Configure(cfg, runtimeConfigOverride)
+	err = o.Configure(cfg)
 	if err != nil {
 		return fmt.Errorf("unable to configure containerd: %v", err)
 	}
@@ -223,6 +220,7 @@ func Cleanup(c *cli.Context, o *options) error {
 
 	cfg, err := containerd.New(
 		containerd.WithPath(o.Config),
+		containerd.WithConfigSource(toml.FromFile(o.Config)),
 		containerd.WithRuntimeType(o.runtimeType),
 		containerd.WithUseLegacyConfig(o.useLegacyConfig),
 		containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...),
@@ -260,16 +258,3 @@ func (o *options) containerAnnotationsFromCDIPrefixes() []string {
 
 	return annotations
 }
-
-func (o *options) runtimeConfigOverride() (map[string]interface{}, error) {
-	if o.runtimeConfigOverrideJSON == "" {
-		return nil, nil
-	}
-
-	runtimeOptions := make(map[string]interface{})
-	if err := json.Unmarshal([]byte(o.runtimeConfigOverrideJSON), &runtimeOptions); err != nil {
-		return nil, fmt.Errorf("failed to read %v as JSON: %w", o.runtimeConfigOverrideJSON, err)
-	}
-
-	return runtimeOptions, nil
-}

tools/container/containerd/containerd_test.go

@@ -1,72 +0,0 @@
-/**
-# Copyright 2024 NVIDIA CORPORATION
-#
-# 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 main
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestRuntimeOptions(t *testing.T) {
-	testCases := []struct {
-		description   string
-		options       options
-		expected      map[string]interface{}
-		expectedError error
-	}{
-		{
-			description: "empty is nil",
-		},
-		{
-			description: "empty json",
-			options: options{
-				runtimeConfigOverrideJSON: "{}",
-			},
-			expected:      map[string]interface{}{},
-			expectedError: nil,
-		},
-		{
-			description: "SystemdCgroup is true",
-			options: options{
-				runtimeConfigOverrideJSON: "{\"SystemdCgroup\": true}",
-			},
-			expected: map[string]interface{}{
-				"SystemdCgroup": true,
-			},
-			expectedError: nil,
-		},
-		{
-			description: "SystemdCgroup is false",
-			options: options{
-				runtimeConfigOverrideJSON: "{\"SystemdCgroup\": false}",
-			},
-			expected: map[string]interface{}{
-				"SystemdCgroup": false,
-			},
-			expectedError: nil,
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.description, func(t *testing.T) {
-			runtimeOptions, err := tc.options.runtimeConfigOverride()
-			require.ErrorIs(t, tc.expectedError, err)
-			require.EqualValues(t, tc.expected, runtimeOptions)
-		})
-	}
-}

tools/container/crio/crio.go

@@ -28,6 +28,7 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/info"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/crio"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/ocihook"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 )
 
@@ -222,6 +223,7 @@ func setupConfig(o *options) error {
 
 	cfg, err := crio.New(
 		crio.WithPath(o.Config),
+		crio.WithConfigSource(toml.FromFile(o.Config)),
 	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
@@ -273,6 +275,7 @@ func cleanupConfig(o *options) error {
 
 	cfg, err := crio.New(
 		crio.WithPath(o.Config),
+		crio.WithConfigSource(toml.FromFile(o.Config)),
 	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)