Convert containerd to runtime package

Signed-off-by: Evan Lezar <elezar@nvidia.com>

tools/container/containerd/containerd.go

@@ -1,260 +0,0 @@
-/**
-# Copyright (c) 2020-2021, 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 main
-
-import (
-	"fmt"
-	"os"
-
-	log "github.com/sirupsen/logrus"
-	cli "github.com/urfave/cli/v2"
-
-	"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"
-)
-
-const (
-	defaultConfig        = "/etc/containerd/config.toml"
-	defaultSocket        = "/run/containerd/containerd.sock"
-	defaultRuntimeClass  = "nvidia"
-	defaultRuntmeType    = "io.containerd.runc.v2"
-	defaultSetAsDefault  = true
-	defaultRestartMode   = "signal"
-	defaultHostRootMount = "/host"
-)
-
-// options stores the configuration from the command line or environment variables
-type options struct {
-	container.Options
-
-	// containerd-specific options
-	useLegacyConfig bool
-	runtimeType     string
-
-	ContainerRuntimeModesCDIAnnotationPrefixes cli.StringSlice
-
-	runtimeConfigOverrideJSON string
-}
-
-func main() {
-	options := options{}
-
-	// Create the top-level CLI
-	c := cli.NewApp()
-	c.Name = "containerd"
-	c.Usage = "Update a containerd config with the nvidia-container-runtime"
-	c.Version = info.GetVersionString()
-
-	// Create the 'setup' subcommand
-	setup := cli.Command{}
-	setup.Name = "setup"
-	setup.Usage = "Trigger a containerd config to be updated"
-	setup.ArgsUsage = "<runtime_dirname>"
-	setup.Action = func(c *cli.Context) error {
-		return Setup(c, &options)
-	}
-	setup.Before = func(c *cli.Context) error {
-		return container.ParseArgs(c, &options.Options)
-	}
-
-	// Create the 'cleanup' subcommand
-	cleanup := cli.Command{}
-	cleanup.Name = "cleanup"
-	cleanup.Usage = "Trigger any updates made to a containerd config to be undone"
-	cleanup.ArgsUsage = "<runtime_dirname>"
-	cleanup.Action = func(c *cli.Context) error {
-		return Cleanup(c, &options)
-	}
-	cleanup.Before = func(c *cli.Context) error {
-		return container.ParseArgs(c, &options.Options)
-	}
-
-	// Register the subcommands with the top-level CLI
-	c.Commands = []*cli.Command{
-		&setup,
-		&cleanup,
-	}
-
-	// Setup common flags across both subcommands. All subcommands get the same
-	// set of flags even if they don't use some of them. This is so that we
-	// only require the user to specify one set of flags for both 'startup'
-	// and 'cleanup' to simplify things.
-	commonFlags := []cli.Flag{
-		&cli.StringFlag{
-			Name:        "config",
-			Usage:       "Path to the containerd config file",
-			Value:       defaultConfig,
-			Destination: &options.Config,
-			EnvVars:     []string{"RUNTIME_CONFIG", "CONTAINERD_CONFIG"},
-		},
-		&cli.StringFlag{
-			Name:        "socket",
-			Usage:       "Path to the containerd socket file",
-			Value:       defaultSocket,
-			Destination: &options.Socket,
-			EnvVars:     []string{"RUNTIME_SOCKET", "CONTAINERD_SOCKET"},
-		},
-		&cli.StringFlag{
-			Name:        "restart-mode",
-			Usage:       "Specify how containerd should be restarted;  If 'none' is selected, it will not be restarted [signal | systemd | none]",
-			Value:       defaultRestartMode,
-			Destination: &options.RestartMode,
-			EnvVars:     []string{"RUNTIME_RESTART_MODE", "CONTAINERD_RESTART_MODE"},
-		},
-		&cli.StringFlag{
-			Name:        "runtime-name",
-			Aliases:     []string{"nvidia-runtime-name", "runtime-class"},
-			Usage:       "The name of the runtime class to set for the nvidia-container-runtime",
-			Value:       defaultRuntimeClass,
-			Destination: &options.RuntimeName,
-			EnvVars:     []string{"NVIDIA_RUNTIME_NAME", "CONTAINERD_RUNTIME_CLASS"},
-		},
-		&cli.StringFlag{
-			Name:        "nvidia-runtime-dir",
-			Aliases:     []string{"runtime-dir"},
-			Usage:       "The path where the nvidia-container-runtime binaries are located. If this is not specified, the first argument will be used instead",
-			Destination: &options.RuntimeDir,
-			EnvVars:     []string{"NVIDIA_RUNTIME_DIR"},
-		},
-		&cli.BoolFlag{
-			Name:        "set-as-default",
-			Usage:       "Set nvidia-container-runtime as the default runtime",
-			Value:       defaultSetAsDefault,
-			Destination: &options.SetAsDefault,
-			EnvVars:     []string{"NVIDIA_RUNTIME_SET_AS_DEFAULT", "CONTAINERD_SET_AS_DEFAULT"},
-			Hidden:      true,
-		},
-		&cli.StringFlag{
-			Name:        "host-root",
-			Usage:       "Specify the path to the host root to be used when restarting containerd using systemd",
-			Value:       defaultHostRootMount,
-			Destination: &options.HostRootMount,
-			EnvVars:     []string{"HOST_ROOT_MOUNT"},
-		},
-		&cli.BoolFlag{
-			Name:        "use-legacy-config",
-			Usage:       "Specify whether a legacy (pre v1.3) config should be used",
-			Destination: &options.useLegacyConfig,
-			EnvVars:     []string{"CONTAINERD_USE_LEGACY_CONFIG"},
-		},
-		&cli.StringFlag{
-			Name:        "runtime-type",
-			Usage:       "The runtime_type to use for the configured runtime classes",
-			Value:       defaultRuntmeType,
-			Destination: &options.runtimeType,
-			EnvVars:     []string{"CONTAINERD_RUNTIME_TYPE"},
-		},
-		&cli.StringSliceFlag{
-			Name:        "nvidia-container-runtime-modes.cdi.annotation-prefixes",
-			Destination: &options.ContainerRuntimeModesCDIAnnotationPrefixes,
-			EnvVars:     []string{"NVIDIA_CONTAINER_RUNTIME_MODES_CDI_ANNOTATION_PREFIXES"},
-		},
-		&cli.StringFlag{
-			Name:        "runtime-config-override",
-			Destination: &options.runtimeConfigOverrideJSON,
-			Usage:       "specify additional runtime options as a JSON string. The paths are relative to the runtime config.",
-			Value:       "{}",
-			EnvVars:     []string{"RUNTIME_CONFIG_OVERRIDE", "CONTAINERD_RUNTIME_CONFIG_OVERRIDE"},
-		},
-	}
-
-	// Update the subcommand flags with the common subcommand flags
-	setup.Flags = append([]cli.Flag{}, commonFlags...)
-	cleanup.Flags = append([]cli.Flag{}, commonFlags...)
-
-	// Run the top-level CLI
-	if err := c.Run(os.Args); err != nil {
-		log.Fatal(fmt.Errorf("error: %v", err))
-	}
-}
-
-// Setup updates a containerd configuration to include the nvidia-containerd-runtime and reloads it
-func Setup(c *cli.Context, o *options) error {
-	log.Infof("Starting 'setup' for %v", c.App.Name)
-
-	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()...),
-	)
-	if err != nil {
-		return fmt.Errorf("unable to load config: %v", err)
-	}
-
-	err = o.Configure(cfg)
-	if err != nil {
-		return fmt.Errorf("unable to configure containerd: %v", err)
-	}
-
-	err = RestartContainerd(o)
-	if err != nil {
-		return fmt.Errorf("unable to restart containerd: %v", err)
-	}
-
-	log.Infof("Completed 'setup' for %v", c.App.Name)
-
-	return nil
-}
-
-// Cleanup reverts a containerd configuration to remove the nvidia-containerd-runtime and reloads it
-func Cleanup(c *cli.Context, o *options) error {
-	log.Infof("Starting 'cleanup' for %v", c.App.Name)
-
-	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()...),
-	)
-	if err != nil {
-		return fmt.Errorf("unable to load config: %v", err)
-	}
-
-	err = o.Unconfigure(cfg)
-	if err != nil {
-		return fmt.Errorf("unable to unconfigure containerd: %v", err)
-	}
-
-	err = RestartContainerd(o)
-	if err != nil {
-		return fmt.Errorf("unable to restart containerd: %v", err)
-	}
-
-	log.Infof("Completed 'cleanup' for %v", c.App.Name)
-
-	return nil
-}
-
-// RestartContainerd restarts containerd depending on the value of restartModeFlag
-func RestartContainerd(o *options) error {
-	return o.Restart("containerd", SignalContainerd)
-}
-
-// containerAnnotationsFromCDIPrefixes returns the container annotations to set for the given CDI prefixes.
-func (o *options) containerAnnotationsFromCDIPrefixes() []string {
-	var annotations []string
-	for _, prefix := range o.ContainerRuntimeModesCDIAnnotationPrefixes.Value() {
-		annotations = append(annotations, prefix+"*")
-	}
-
-	return annotations
-}

tools/container/runtime/containerd/config_v1_test.go

@@ -14,7 +14,7 @@
 # limitations under the License.
 */
 
-package main
+package containerd
 
 import (
 	"fmt"
@@ -82,14 +82,10 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
-			o := &options{
+			o := &container.Options{
-				Options: container.Options{
+				RuntimeName:  tc.runtimeName,
-					RuntimeName:  tc.runtimeName,
+				RuntimeDir:   runtimeDir,
-					RuntimeDir:   runtimeDir,
+				SetAsDefault: tc.setAsDefault,
-					SetAsDefault: tc.setAsDefault,
-				},
-				runtimeType:     runtimeType,
-				useLegacyConfig: tc.legacyConfig,
 			}
 
 			cfg, err := toml.Empty.Load()
@@ -233,11 +229,9 @@ func TestUpdateV1Config(t *testing.T) {
 
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
-			o := &options{
+			o := &container.Options{
-				Options: container.Options{
+				RuntimeName: tc.runtimeName,
-					RuntimeName: tc.runtimeName,
+				RuntimeDir:  runtimeDir,
-					RuntimeDir:  runtimeDir,
-				},
 			}
 
 			cfg, err := toml.Empty.Load()
@@ -394,11 +388,9 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
-			o := &options{
+			o := &container.Options{
-				Options: container.Options{
+				RuntimeName: tc.runtimeName,
-					RuntimeName: tc.runtimeName,
+				RuntimeDir:  runtimeDir,
-					RuntimeDir:  runtimeDir,
-				},
 			}
 
 			cfg, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary"))
@@ -473,10 +465,8 @@ func TestRevertV1Config(t *testing.T) {
 
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
-			o := &options{
+			o := &container.Options{
-				Options: container.Options{
+				RuntimeName: "nvidia",
-					RuntimeName: "nvidia",
-				},
 			}
 
 			cfg, err := toml.LoadMap(tc.config)

tools/container/runtime/containerd/config_v2_test.go

@@ -14,7 +14,7 @@
 # limitations under the License.
 */
 
-package main
+package containerd
 
 import (
 	"fmt"
@@ -66,12 +66,10 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
 
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
-			o := &options{
+			o := &container.Options{
-				Options: container.Options{
+				RuntimeName:  tc.runtimeName,
-					RuntimeName:  tc.runtimeName,
+				RuntimeDir:   runtimeDir,
-					RuntimeDir:   runtimeDir,
+				SetAsDefault: tc.setAsDefault,
-					SetAsDefault: tc.setAsDefault,
-				},
 			}
 
 			cfg, err := toml.LoadMap(map[string]interface{}{})
@@ -192,12 +190,9 @@ func TestUpdateV2Config(t *testing.T) {
 
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
-			o := &options{
+			o := &container.Options{
-				Options: container.Options{
+				RuntimeName: tc.runtimeName,
-					RuntimeName: tc.runtimeName,
+				RuntimeDir:  runtimeDir,
-					RuntimeDir:  runtimeDir,
-				},
-				runtimeType: runtimeType,
 			}
 
 			cfg, err := toml.LoadMap(map[string]interface{}{})
@@ -206,7 +201,7 @@ func TestUpdateV2Config(t *testing.T) {
 			v2 := &containerd.Config{
 				Logger:               logger,
 				Tree:                 cfg,
-				RuntimeType:          o.runtimeType,
+				RuntimeType:          runtimeType,
 				ContainerAnnotations: []string{"cdi.k8s.io/*"},
 			}
 
@@ -348,11 +343,9 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
-			o := &options{
+			o := &container.Options{
-				Options: container.Options{
+				RuntimeName: tc.runtimeName,
-					RuntimeName: tc.runtimeName,
+				RuntimeDir:  runtimeDir,
-					RuntimeDir:  runtimeDir,
-				},
 			}
 
 			cfg, err := toml.LoadMap(runcConfigMapV2("/runc-binary"))
@@ -421,10 +414,8 @@ func TestRevertV2Config(t *testing.T) {
 
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
-			o := &options{
+			o := &container.Options{
-				Options: container.Options{
+				RuntimeName: "nvidia",
-					RuntimeName: "nvidia",
-				},
 			}
 
 			cfg, err := toml.LoadMap(tc.config)

tools/container/runtime/containerd/containerd.go

@@ -0,0 +1,166 @@
+/**
+# Copyright (c) 2020-2021, 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 containerd
+
+import (
+	"encoding/json"
+	"fmt"
+
+	log "github.com/sirupsen/logrus"
+	cli "github.com/urfave/cli/v2"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
+)
+
+const (
+	Name = "containerd"
+
+	DefaultConfig      = "/etc/containerd/config.toml"
+	DefaultSocket      = "/run/containerd/containerd.sock"
+	DefaultRestartMode = "signal"
+
+	defaultRuntmeType = "io.containerd.runc.v2"
+)
+
+// Options stores the containerd-specific options
+type Options struct {
+	useLegacyConfig bool
+	runtimeType     string
+
+	ContainerRuntimeModesCDIAnnotationPrefixes cli.StringSlice
+
+	runtimeConfigOverrideJSON string
+}
+
+func Flags(opts *Options) []cli.Flag {
+	flags := []cli.Flag{
+		&cli.BoolFlag{
+			Name:        "use-legacy-config",
+			Usage:       "Specify whether a legacy (pre v1.3) config should be used",
+			Destination: &opts.useLegacyConfig,
+			EnvVars:     []string{"CONTAINERD_USE_LEGACY_CONFIG"},
+		},
+		&cli.StringFlag{
+			Name:        "runtime-type",
+			Usage:       "The runtime_type to use for the configured runtime classes",
+			Value:       defaultRuntmeType,
+			Destination: &opts.runtimeType,
+			EnvVars:     []string{"CONTAINERD_RUNTIME_TYPE"},
+		},
+		&cli.StringSliceFlag{
+			Name:        "nvidia-container-runtime-modes.cdi.annotation-prefixes",
+			Destination: &opts.ContainerRuntimeModesCDIAnnotationPrefixes,
+			EnvVars:     []string{"NVIDIA_CONTAINER_RUNTIME_MODES_CDI_ANNOTATION_PREFIXES"},
+		},
+		&cli.StringFlag{
+			Name:        "runtime-config-override",
+			Destination: &opts.runtimeConfigOverrideJSON,
+			Usage:       "specify additional runtime options as a JSON string. The paths are relative to the runtime config.",
+			Value:       "{}",
+			EnvVars:     []string{"RUNTIME_CONFIG_OVERRIDE", "CONTAINERD_RUNTIME_CONFIG_OVERRIDE"},
+		},
+	}
+
+	return flags
+}
+
+// Setup updates a containerd configuration to include the nvidia-containerd-runtime and reloads it
+func Setup(c *cli.Context, o *container.Options, co *Options) error {
+	log.Infof("Starting 'setup' for %v", c.App.Name)
+
+	cfg, err := containerd.New(
+		containerd.WithPath(o.Config),
+		containerd.WithRuntimeType(co.runtimeType),
+		containerd.WithUseLegacyConfig(co.useLegacyConfig),
+		containerd.WithContainerAnnotations(co.containerAnnotationsFromCDIPrefixes()...),
+	)
+	if err != nil {
+		return fmt.Errorf("unable to load config: %v", err)
+	}
+
+	err = o.Configure(cfg)
+	if err != nil {
+		return fmt.Errorf("unable to configure containerd: %v", err)
+	}
+
+	err = RestartContainerd(o)
+	if err != nil {
+		return fmt.Errorf("unable to restart containerd: %v", err)
+	}
+
+	log.Infof("Completed 'setup' for %v", c.App.Name)
+
+	return nil
+}
+
+// Cleanup reverts a containerd configuration to remove the nvidia-containerd-runtime and reloads it
+func Cleanup(c *cli.Context, o *container.Options, co *Options) error {
+	log.Infof("Starting 'cleanup' for %v", c.App.Name)
+
+	cfg, err := containerd.New(
+		containerd.WithPath(o.Config),
+		containerd.WithRuntimeType(co.runtimeType),
+		containerd.WithUseLegacyConfig(co.useLegacyConfig),
+		containerd.WithContainerAnnotations(co.containerAnnotationsFromCDIPrefixes()...),
+	)
+	if err != nil {
+		return fmt.Errorf("unable to load config: %v", err)
+	}
+
+	err = o.Unconfigure(cfg)
+	if err != nil {
+		return fmt.Errorf("unable to unconfigure containerd: %v", err)
+	}
+
+	err = RestartContainerd(o)
+	if err != nil {
+		return fmt.Errorf("unable to restart containerd: %v", err)
+	}
+
+	log.Infof("Completed 'cleanup' for %v", c.App.Name)
+
+	return nil
+}
+
+// RestartContainerd restarts containerd depending on the value of restartModeFlag
+func RestartContainerd(o *container.Options) error {
+	return o.Restart("containerd", SignalContainerd)
+}
+
+// containerAnnotationsFromCDIPrefixes returns the container annotations to set for the given CDI prefixes.
+func (o *Options) containerAnnotationsFromCDIPrefixes() []string {
+	var annotations []string
+	for _, prefix := range o.ContainerRuntimeModesCDIAnnotationPrefixes.Value() {
+		annotations = append(annotations, prefix+"*")
+	}
+
+	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/runtime/containerd/containerd_linux.go

@@ -14,7 +14,7 @@
 # limitations under the License.
 **/
 
-package main
+package containerd
 
 import (
 	"fmt"

tools/container/runtime/containerd/containerd_other.go

@@ -17,7 +17,7 @@
 # limitations under the License.
 **/
 
-package main
+package containerd
 
 import (
 	"errors"

tools/container/runtime/containerd/containerd_test.go

@@ -0,0 +1,72 @@
+/**
+# 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 (
+	"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/runtime/docker/docker.go

@@ -34,6 +34,12 @@ const (
 	DefaultRestartMode = "signal"
 )
 
+type Options struct{}
+
+func Flags(opts *Options) []cli.Flag {
+	return nil
+}
+
 // Setup updates docker configuration to include the nvidia runtime and reloads it
 func Setup(c *cli.Context, o *container.Options) error {
 	log.Infof("Starting 'setup' for %v", c.App.Name)

tools/container/runtime/runtime.go

@@ -22,6 +22,7 @@ import (
 	"github.com/urfave/cli/v2"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/runtime/containerd"
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/runtime/docker"
 )
 
@@ -36,6 +37,8 @@ const (
 
 type Options struct {
 	container.Options
+
+	containerdOptions containerd.Options
 }
 
 func Flags(opts *Options) []cli.Flag {
@@ -86,6 +89,8 @@ func Flags(opts *Options) []cli.Flag {
 		},
 	}
 
+	flags = append(flags, containerd.Flags(&opts.containerdOptions)...)
+
 	return flags
 }
 
@@ -96,6 +101,16 @@ func ValidateOptions(opts *Options, runtime string, toolkitRoot string) error {
 
 	// Apply the runtime-specific config changes.
 	switch runtime {
+	case containerd.Name:
+		if opts.Config == runtimeSpecificDefault {
+			opts.Config = containerd.DefaultConfig
+		}
+		if opts.Socket == runtimeSpecificDefault {
+			opts.Socket = containerd.DefaultSocket
+		}
+		if opts.RestartMode == runtimeSpecificDefault {
+			opts.RestartMode = containerd.DefaultRestartMode
+		}
 	case docker.Name:
 		if opts.Config == runtimeSpecificDefault {
 			opts.Config = docker.DefaultConfig
@@ -115,6 +130,8 @@ func ValidateOptions(opts *Options, runtime string, toolkitRoot string) error {
 
 func Setup(c *cli.Context, opts *Options, runtime string) error {
 	switch runtime {
+	case containerd.Name:
+		return containerd.Setup(c, &opts.Options, &opts.containerdOptions)
 	case docker.Name:
 		return docker.Setup(c, &opts.Options)
 	default:
@@ -124,6 +141,8 @@ func Setup(c *cli.Context, opts *Options, runtime string) error {
 
 func Cleanup(c *cli.Context, opts *Options, runtime string) error {
 	switch runtime {
+	case containerd.Name:
+		return containerd.Cleanup(c, &opts.Options, &opts.containerdOptions)
 	case docker.Name:
 		return docker.Cleanup(c, &opts.Options)
 	default: