Add support for containerd version 3 config

This change adds support for containerd configs with version=3.
From the perspective of the runtime configuration the contents of the
config are the same. This means that we just have to load the new
version and ensure that this is propagated to the generated config.

Note that v3 config also requires a switch to the 'io.containerd.cri.v1.runtime'
CRI runtime plugin. See:
https://github.com/containerd/containerd/blob/v2.0.0/docs/PLUGINS.md
https://github.com/containerd/containerd/issues/10132

Note that we still use a default config of version=2 since we need to
ensure compatibility with older containerd versions (1.6.x and 1.7.x).

Signed-off-by: Sam Lockart <sam.lockart@zendesk.com>
Signed-off-by: Evan Lezar <elezar@nvidia.com>
Signed-off-by: Christopher Desiniotis <cdesiniotis@nvidia.com>

pkg/config/engine/containerd/config.go

@@ -30,40 +30,40 @@ func (c *Config) AddRuntime(name string, path string, setAsDefault bool) error {
 	}
 	config := *c.Tree
 
-	config.Set("version", int64(2))
+	config.Set("version", c.Version)
 
 	runtimeNamesForConfig := engine.GetLowLevelRuntimes(c)
 	for _, r := range runtimeNamesForConfig {
-		options := config.GetSubtreeByPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", r})
+		options := config.GetSubtreeByPath([]string{"plugins", c.CRIRuntimePluginName, "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())
+		config.SetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name}, options.Copy())
 		break
 	}
 
-	if config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name}) == nil {
+	if config.GetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name}) == nil {
 		c.Logger.Warningf("could not infer options from runtimes %v; using defaults", runtimeNamesForConfig)
-		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "runtime_type"}, c.RuntimeType)
+		config.SetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name, "runtime_type"}, c.RuntimeType)
-		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "runtime_root"}, "")
+		config.SetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name, "runtime_root"}, "")
-		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "runtime_engine"}, "")
+		config.SetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name, "runtime_engine"}, "")
-		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "privileged_without_host_devices"}, false)
+		config.SetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name, "privileged_without_host_devices"}, false)
 	}
 
 	if len(c.ContainerAnnotations) > 0 {
-		annotations, err := c.getRuntimeAnnotations([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "container_annotations"})
+		annotations, err := c.getRuntimeAnnotations([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name, "container_annotations"})
 		if err != nil {
 			return err
 		}
 		annotations = append(c.ContainerAnnotations, annotations...)
-		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "container_annotations"}, annotations)
+		config.SetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name, "container_annotations"}, annotations)
 	}
 
-	config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "options", "BinaryName"}, path)
+	config.SetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name, "options", "BinaryName"}, path)
 
 	if setAsDefault {
-		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"}, name)
+		config.SetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "default_runtime_name"}, name)
 	}
 
 	*c.Tree = config
@@ -99,13 +99,13 @@ func (c *Config) getRuntimeAnnotations(path []string) ([]string, error) {
 // Set sets the specified containerd option.
 func (c *Config) Set(key string, value interface{}) {
 	config := *c.Tree
-	config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", key}, value)
+	config.SetPath([]string{"plugins", c.CRIRuntimePluginName, key}, value)
 	*c.Tree = config
 }
 
 // DefaultRuntime returns the default runtime for the cri-o config
 func (c Config) DefaultRuntime() string {
-	if runtime, ok := c.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"}).(string); ok {
+	if runtime, ok := c.GetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "default_runtime_name"}).(string); ok {
 		return runtime
 	}
 	return ""
@@ -119,14 +119,14 @@ func (c *Config) RemoveRuntime(name string) error {
 
 	config := *c.Tree
 
-	config.DeletePath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name})
+	config.DeletePath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name})
-	if runtime, ok := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"}).(string); ok {
+	if runtime, ok := config.GetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "default_runtime_name"}).(string); ok {
 		if runtime == name {
-			config.DeletePath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"})
+			config.DeletePath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "default_runtime_name"})
 		}
 	}
 
-	runtimePath := []string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name}
+	runtimePath := []string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name}
 	for i := 0; i < len(runtimePath); i++ {
 		if runtimes, ok := config.GetPath(runtimePath[:len(runtimePath)-i]).(*toml.Tree); ok {
 			if len(runtimes.Keys()) == 0 {

pkg/config/engine/containerd/config_test.go

@@ -195,6 +195,68 @@ func TestAddRuntime(t *testing.T) {
 						SystemdCgroup = false
 				`,
 		},
+		{
+			description: "empty v3 spec is supported",
+			config: `
+			version = 3
+			`,
+			expectedConfig: `
+			version = 3
+			[plugins]
+			[plugins."io.containerd.cri.v1.runtime"]
+				[plugins."io.containerd.cri.v1.runtime".containerd]
+				[plugins."io.containerd.cri.v1.runtime".containerd.runtimes]
+					[plugins."io.containerd.cri.v1.runtime".containerd.runtimes.test]
+					privileged_without_host_devices = false
+					runtime_engine = ""
+					runtime_root = ""
+					runtime_type = "io.containerd.runc.v2"
+					[plugins."io.containerd.cri.v1.runtime".containerd.runtimes.test.options]
+						BinaryName = "/usr/bin/test"
+			`,
+			expectedError: nil,
+		},
+		{
+			description: "v3 spec is supported",
+			config: `
+			version = 3
+			[plugins]
+			[plugins."io.containerd.cri.v1.runtime"]
+				[plugins."io.containerd.cri.v1.runtime".containerd]
+				[plugins."io.containerd.cri.v1.runtime".containerd.runtimes]
+					[plugins."io.containerd.cri.v1.runtime".containerd.runtimes.runc]
+					privileged_without_host_devices = true
+					runtime_engine = "engine"
+					runtime_root = "root"
+					runtime_type = "type"
+					[plugins."io.containerd.cri.v1.runtime".containerd.runtimes.runc.options]
+						BinaryName = "/usr/bin/runc"
+						SystemdCgroup = true
+			`,
+			expectedConfig: `
+			version = 3
+			[plugins]
+			[plugins."io.containerd.cri.v1.runtime"]
+				[plugins."io.containerd.cri.v1.runtime".containerd]
+				[plugins."io.containerd.cri.v1.runtime".containerd.runtimes]
+					[plugins."io.containerd.cri.v1.runtime".containerd.runtimes.runc]
+					privileged_without_host_devices = true
+					runtime_engine = "engine"
+					runtime_root = "root"
+					runtime_type = "type"
+					[plugins."io.containerd.cri.v1.runtime".containerd.runtimes.runc.options]
+						BinaryName = "/usr/bin/runc"
+						SystemdCgroup = true
+					[plugins."io.containerd.cri.v1.runtime".containerd.runtimes.test]
+					privileged_without_host_devices = true
+					runtime_engine = "engine"
+					runtime_root = "root"
+					runtime_type = "type"
+					[plugins."io.containerd.cri.v1.runtime".containerd.runtimes.test.options]
+						BinaryName = "/usr/bin/test"
+						SystemdCgroup = true
+				`,
+		},
 	}
 
 	for _, tc := range testCases {

pkg/config/engine/containerd/containerd.go

@@ -24,9 +24,15 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
+const (
+	defaultConfigVersion = 2
+	defaultRuntimeType   = "io.containerd.runc.v2"
+)
+
 // Config represents the containerd config
 type Config struct {
 	*toml.Tree
+	Version              int64
 	Logger               logger.Interface
 	RuntimeType          string
 	ContainerAnnotations []string
@@ -36,6 +42,10 @@ type Config struct {
 	// This was deprecated in v1.4 in favour of containerd.default_runtime_name.
 	// Support for this section has been removed in v2.0.
 	UseLegacyConfig bool
+	// CRIRuntimePluginName represents the fully qualified name of the containerd plugin
+	// for the CRI runtime service. The name of this plugin was changed in v3 of the
+	// containerd configuration file.
+	CRIRuntimePluginName string
 }
 
 var _ engine.Interface = (*Config)(nil)
@@ -60,6 +70,7 @@ func (c *containerdCfgRuntime) GetBinaryPath() string {
 // New creates a containerd config with the specified options
 func New(opts ...Option) (engine.Interface, error) {
 	b := &builder{
+		configVersion: defaultConfigVersion,
 		runtimeType:   defaultRuntimeType,
 	}
 	for _, opt := range opts {
@@ -77,56 +88,74 @@ func New(opts ...Option) (engine.Interface, error) {
 		return nil, fmt.Errorf("failed to load config: %v", err)
 	}
 
+	configVersion, err := b.parseVersion(tomlConfig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse config version: %w", err)
+	}
+	b.logger.Infof("Using config version %v", configVersion)
+
+	criRuntimePluginName, err := b.criRuntimePluginName(configVersion)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get CRI runtime plugin name: %w", err)
+	}
+	b.logger.Infof("Using CRI runtime plugin name %q", criRuntimePluginName)
+
 	cfg := &Config{
 		Tree:                 tomlConfig,
+		Version:              configVersion,
+		CRIRuntimePluginName: criRuntimePluginName,
 		Logger:               b.logger,
 		RuntimeType:          b.runtimeType,
-		ContainerAnnotations: b.containerAnnotations,
 		UseLegacyConfig:      b.useLegacyConfig,
+		ContainerAnnotations: b.containerAnnotations,
 	}
 
-	version, err := cfg.parseVersion()
+	switch configVersion {
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse config version: %v", err)
-	}
-	switch version {
 	case 1:
 		return (*ConfigV1)(cfg), nil
-	case 2:
+	default:
 		return cfg, nil
 	}
-
-	return nil, fmt.Errorf("unsupported config version: %v", version)
 }
 
 // parseVersion returns the version of the config
-func (c *Config) parseVersion() (int, error) {
+func (b *builder) parseVersion(c *toml.Tree) (int64, error) {
-	defaultVersion := 2
+	if c == nil || len(c.Keys()) == 0 {
-	// For legacy configs, we default to v1 configs.
+		// No config exists, or the config file is empty.
-	if c.UseLegacyConfig {
+		if b.useLegacyConfig {
-		defaultVersion = 1
+			// If a legacy config is explicitly requested, we default to a v1 config.
+			return 1, nil
+		}
+		// Use the requested version.
+		return int64(b.configVersion), nil
 	}
 
 	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
+		return v, nil
 	default:
 		return -1, fmt.Errorf("unsupported type for version field: %v", v)
 	}
 }
 
+func (b *builder) criRuntimePluginName(configVersion int64) (string, error) {
+	switch configVersion {
+	case 1:
+		return "cri", nil
+	case 2:
+		return "io.containerd.grpc.v1.cri", nil
+	default:
+		return "io.containerd.cri.v1.runtime", nil
+	}
+}
+
 func (c *Config) GetRuntimeConfig(name string) (engine.RuntimeConfig, error) {
 	if c == nil || c.Tree == nil {
 		return nil, fmt.Errorf("config is nil")
 	}
-	runtimeData := c.GetSubtreeByPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name})
+	runtimeData := c.GetSubtreeByPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "runtimes", name})
 	return &containerdCfgRuntime{
 		tree: runtimeData,
 	}, nil

pkg/config/engine/containerd/option.go

@@ -21,16 +21,13 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/toml"
 )
 
-const (
-	defaultRuntimeType = "io.containerd.runc.v2"
-)
-
 type builder struct {
 	logger               logger.Interface
 	configSource         toml.Loader
+	configVersion        int
+	useLegacyConfig      bool
 	path                 string
 	runtimeType          string
-	useLegacyConfig      bool
 	containerAnnotations []string
 }
 
@@ -65,13 +62,20 @@ func WithRuntimeType(runtimeType string) Option {
 	}
 }
 
-// WithUseLegacyConfig sets the useLegacyConfig flag for the config builder
+// WithUseLegacyConfig sets the useLegacyConfig flag for the config builder.
 func WithUseLegacyConfig(useLegacyConfig bool) Option {
 	return func(b *builder) {
 		b.useLegacyConfig = useLegacyConfig
 	}
 }
 
+// WithConfigVersion sets the config version for the config builder
+func WithConfigVersion(configVersion int) Option {
+	return func(b *builder) {
+		b.configVersion = configVersion
+	}
+}
+
 // WithContainerAnnotations sets the container annotations for the config builder
 func WithContainerAnnotations(containerAnnotations ...string) Option {
 	return func(b *builder) {

pkg/config/toml/source-map.go

@@ -0,0 +1,26 @@
+/**
+# 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
+
+type tomlMap map[string]interface{}
+
+var _ Loader = (*tomlFile)(nil)
+
+// Load loads the contents of the specified TOML file as a map.
+func (l tomlMap) Load() (*Tree, error) {
+	return LoadMap(l)
+}

pkg/config/toml/source.go

@@ -25,15 +25,6 @@ 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)
-}
-
 // FromCommandLine creates a TOML source from the output of a shell command and its corresponding args.
 // If the command is empty, an empty config is returned.
 func FromCommandLine(cmds ...string) Loader {
@@ -46,6 +37,24 @@ func FromCommandLine(cmds ...string) Loader {
 	}
 }
 
+// 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)
+}
+
+// FromMap creates a TOML source for the specified map.
+// If an empty map is passed and empty tomly config is used.
+func FromMap(m map[string]interface{}) Loader {
+	if m == nil {
+		return Empty
+	}
+	return tomlMap(m)
+}
+
 // FromString creates a TOML source for the specified contents.
 // If an empty string is passed an empty toml config is used.
 func FromString(contents string) Loader {

tools/container/runtime/containerd/config_v1_test.go

@@ -88,37 +88,36 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 				SetAsDefault: tc.setAsDefault,
 			}
 
-			cfg, err := toml.Empty.Load()
+			v1, err := containerd.New(
-			require.NoError(t, err, "%d: %v", i, tc)
+				containerd.WithLogger(logger),
-
+				containerd.WithConfigSource(toml.Empty),
-			v1 := &containerd.ConfigV1{
+				containerd.WithRuntimeType(runtimeType),
-				Logger:          logger,
+				containerd.WithUseLegacyConfig(tc.legacyConfig),
-				Tree:            cfg,
+				containerd.WithConfigVersion(1),
-				UseLegacyConfig: tc.legacyConfig,
+			)
-				RuntimeType:     runtimeType,
+			require.NoError(t, err)
-			}
 
 			err = o.UpdateConfig(v1)
-			require.NoError(t, err, "%d: %v", i, tc)
+			require.NoError(t, err)
 
-			defaultRuntimeName := v1.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"})
+			cfg := v1.(*containerd.ConfigV1)
-			require.EqualValues(t, tc.expectedDefaultRuntimeName, defaultRuntimeName, "%d: %v", i, tc)
+			defaultRuntimeName := cfg.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"})
+			require.EqualValues(t, tc.expectedDefaultRuntimeName, defaultRuntimeName)
 
-			defaultRuntime := v1.GetPath([]string{"plugins", "cri", "containerd", "default_runtime"})
+			defaultRuntime := cfg.GetPath([]string{"plugins", "cri", "containerd", "default_runtime"})
 			if tc.expectedDefaultRuntimeBinary == nil {
-				require.Nil(t, defaultRuntime, "%d: %v", i, tc)
+				require.Nil(t, defaultRuntime)
 			} else {
 				require.NotNil(t, defaultRuntime)
 
 				expected, err := defaultRuntimeTomlConfigV1(tc.expectedDefaultRuntimeBinary.(string))
-				require.NoError(t, err, "%d: %v", i, tc)
+				require.NoError(t, err)
 
 				configContents, _ := toml.Marshal(defaultRuntime.(*toml.Tree))
 				expectedContents, _ := toml.Marshal(expected)
 
 				require.Equal(t, string(expectedContents), string(configContents), "%d: %v: %v", i, tc)
 			}
-
 		})
 	}
 }
@@ -234,24 +233,22 @@ func TestUpdateV1Config(t *testing.T) {
 				RuntimeDir:  runtimeDir,
 			}
 
-			cfg, err := toml.Empty.Load()
+			v1, err := containerd.New(
+				containerd.WithLogger(logger),
+				containerd.WithConfigSource(toml.Empty),
+				containerd.WithRuntimeType(runtimeType),
+				containerd.WithConfigVersion(1),
+				containerd.WithContainerAnnotations("cdi.k8s.io/*"),
+			)
 			require.NoError(t, err)
 
-			v1 := &containerd.ConfigV1{
-				Logger:               logger,
-				Tree:                 cfg,
-				UseLegacyConfig:      true,
-				RuntimeType:          runtimeType,
-				ContainerAnnotations: []string{"cdi.k8s.io/*"},
-			}
-
 			err = o.UpdateConfig(v1)
 			require.NoError(t, err)
 
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
 			require.NoError(t, err)
 
-			require.Equal(t, expected.String(), cfg.String())
+			require.Equal(t, expected.String(), v1.String())
 		})
 	}
 }
@@ -393,29 +390,28 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 				RuntimeDir:  runtimeDir,
 			}
 
-			cfg, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary"))
+			v1, err := containerd.New(
+				containerd.WithLogger(logger),
+				containerd.WithConfigSource(toml.FromMap(runcConfigMapV1("/runc-binary"))),
+				containerd.WithRuntimeType(runtimeType),
+				containerd.WithConfigVersion(1),
+				containerd.WithContainerAnnotations("cdi.k8s.io/*"),
+			)
 			require.NoError(t, err)
 
-			v1 := &containerd.ConfigV1{
-				Logger:               logger,
-				Tree:                 cfg,
-				UseLegacyConfig:      true,
-				RuntimeType:          runtimeType,
-				ContainerAnnotations: []string{"cdi.k8s.io/*"},
-			}
-
 			err = o.UpdateConfig(v1)
 			require.NoError(t, err)
 
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
 			require.NoError(t, err)
 
-			require.Equal(t, expected.String(), cfg.String())
+			require.Equal(t, expected.String(), v1.String())
 		})
 	}
 }
 
 func TestRevertV1Config(t *testing.T) {
+	logger, _ := testlog.NewNullLogger()
 	testCases := []struct {
 		config map[string]interface {
 		}
@@ -469,25 +465,22 @@ func TestRevertV1Config(t *testing.T) {
 				RuntimeName: "nvidia",
 			}
 
-			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)
+			require.NoError(t, err)
 
-			v1 := &containerd.ConfigV1{
+			v1, err := containerd.New(
-				Tree:            cfg,
+				containerd.WithLogger(logger),
-				UseLegacyConfig: true,
+				containerd.WithConfigSource(toml.FromMap(tc.config)),
-				RuntimeType:     runtimeType,
+				containerd.WithRuntimeType(runtimeType),
-			}
+
+				containerd.WithContainerAnnotations("cdi.k8s.io/*"),
+			)
+			require.NoError(t, err)
 
 			err = o.RevertConfig(v1)
-			require.NoError(t, err, "%d: %v", i, tc)
+			require.NoError(t, err)
 
-			configContents, _ := toml.Marshal(cfg)
+			require.Equal(t, expected.String(), v1.String())
-			expectedContents, _ := toml.Marshal(expected)
-
-			require.Equal(t, string(expectedContents), string(configContents), "%d: %v", i, tc)
 		})
 	}
 }

tools/container/runtime/containerd/config_v2_test.go

@@ -72,18 +72,19 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
 				SetAsDefault: tc.setAsDefault,
 			}
 
-			cfg, err := toml.LoadMap(map[string]interface{}{})
+			v2, err := containerd.New(
+				containerd.WithLogger(logger),
+				containerd.WithConfigSource(toml.Empty),
+				containerd.WithRuntimeType(runtimeType),
+				containerd.WithContainerAnnotations("cdi.k8s.io/*"),
+			)
 			require.NoError(t, err)
 
-			v2 := &containerd.Config{
-				Logger:      logger,
-				Tree:        cfg,
-				RuntimeType: runtimeType,
-			}
-
 			err = o.UpdateConfig(v2)
 			require.NoError(t, err)
 
+			cfg := v2.(*containerd.Config)
+
 			defaultRuntimeName := cfg.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"})
 			require.EqualValues(t, tc.expectedDefaultRuntimeName, defaultRuntimeName)
 		})
@@ -195,23 +196,21 @@ func TestUpdateV2Config(t *testing.T) {
 				RuntimeDir:  runtimeDir,
 			}
 
-			cfg, err := toml.LoadMap(map[string]interface{}{})
+			v2, err := containerd.New(
+				containerd.WithLogger(logger),
+				containerd.WithConfigSource(toml.Empty),
+				containerd.WithRuntimeType(runtimeType),
+				containerd.WithContainerAnnotations("cdi.k8s.io/*"),
+			)
 			require.NoError(t, err)
 
-			v2 := &containerd.Config{
-				Logger:               logger,
-				Tree:                 cfg,
-				RuntimeType:          runtimeType,
-				ContainerAnnotations: []string{"cdi.k8s.io/*"},
-			}
-
 			err = o.UpdateConfig(v2)
 			require.NoError(t, err)
 
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
 			require.NoError(t, err)
 
-			require.Equal(t, expected.String(), cfg.String())
+			require.Equal(t, expected.String(), v2.String())
 		})
 	}
 
@@ -348,28 +347,28 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 				RuntimeDir:  runtimeDir,
 			}
 
-			cfg, err := toml.LoadMap(runcConfigMapV2("/runc-binary"))
+			v2, err := containerd.New(
+				containerd.WithLogger(logger),
+				containerd.WithConfigSource(toml.FromMap(runcConfigMapV2("/runc-binary"))),
+				containerd.WithRuntimeType(runtimeType),
+				containerd.WithContainerAnnotations("cdi.k8s.io/*"),
+			)
 			require.NoError(t, err)
 
-			v2 := &containerd.Config{
-				Logger:               logger,
-				Tree:                 cfg,
-				RuntimeType:          runtimeType,
-				ContainerAnnotations: []string{"cdi.k8s.io/*"},
-			}
-
 			err = o.UpdateConfig(v2)
 			require.NoError(t, err)
 
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
 			require.NoError(t, err)
 
-			require.Equal(t, expected.String(), cfg.String())
+			require.Equal(t, expected.String(), v2.String())
 		})
 	}
 }
 
 func TestRevertV2Config(t *testing.T) {
+	logger, _ := testlog.NewNullLogger()
+
 	testCases := []struct {
 		config map[string]interface {
 		}
@@ -418,24 +417,21 @@ func TestRevertV2Config(t *testing.T) {
 				RuntimeName: "nvidia",
 			}
 
-			cfg, err := toml.LoadMap(tc.config)
-			require.NoError(t, err)
-
 			expected, err := toml.TreeFromMap(tc.expected)
 			require.NoError(t, err)
 
-			v2 := &containerd.Config{
+			v2, err := containerd.New(
-				Tree:        cfg,
+				containerd.WithLogger(logger),
-				RuntimeType: runtimeType,
+				containerd.WithConfigSource(toml.FromMap(tc.config)),
-			}
+				containerd.WithRuntimeType(runtimeType),
+				containerd.WithContainerAnnotations("cdi.k8s.io/*"),
+			)
+			require.NoError(t, err)
 
 			err = o.RevertConfig(v2)
 			require.NoError(t, err)
 
-			configContents, _ := toml.Marshal(cfg)
+			require.Equal(t, expected.String(), v2.String())
-			expectedContents, _ := toml.Marshal(expected)
-
-			require.Equal(t, string(expectedContents), string(configContents))
 		})
 	}
 }
@@ -454,6 +450,7 @@ func runtimeMapV2(binary string) map[string]interface{} {
 
 func runcConfigMapV2(binary string) map[string]interface{} {
 	return map[string]interface{}{
+		"version": 2,
 		"plugins": map[string]interface{}{
 			"io.containerd.grpc.v1.cri": map[string]interface{}{
 				"containerd": map[string]interface{}{