Merge branch 'CNT-4075' into 'release-1.13'

Allow same envars for all runtime configs

See merge request nvidia/container-toolkit/container-toolkit!419

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## v1.13.3
 
+* [toolkit-container] Allow same envars for all runtime configs
+
 ## v1.13.2
 
 * Add `nvidia-container-runtime-hook.path` config option to specify NVIDIA Container Runtime Hook path explicitly.

test/container/crio_test.sh

@@ -27,7 +27,7 @@ testing::crio::hook_created() {
 }
 
 testing::crio::hook_cleanup() {
-	testing::docker_run::toolkit::shell 'crio cleanup'
+	testing::docker_run::toolkit::shell 'crio cleanup --nvidia-runtime-dir /run/nvidia/toolkit/'
 
 	test -z "$(ls -A "${shared_dir}${CRIO_HOOKS_DIR}")"
 }

tools/container/container.go

@@ -0,0 +1,169 @@
+/**
+# Copyright (c) NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+**/
+
+package container
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator"
+	"github.com/sirupsen/logrus"
+	"github.com/urfave/cli/v2"
+)
+
+const (
+	restartModeNone    = "none"
+	restartModeSignal  = "signal"
+	restartModeSystemd = "systemd"
+)
+
+// Options defines the shared options for the CLIs to configure containers runtimes.
+type Options struct {
+	Config        string
+	Socket        string
+	RuntimeName   string
+	RuntimeDir    string
+	SetAsDefault  bool
+	RestartMode   string
+	HostRootMount string
+}
+
+// ParseArgs parses the command line arguments to the CLI
+func ParseArgs(c *cli.Context, o *Options) error {
+	if o.RuntimeDir != "" {
+		logrus.Debug("Runtime directory already set; ignoring arguments")
+		return nil
+	}
+
+	args := c.Args()
+
+	logrus.Infof("Parsing arguments: %v", args.Slice())
+	if c.NArg() != 1 {
+		return fmt.Errorf("incorrect number of arguments")
+	}
+
+	o.RuntimeDir = args.Get(0)
+
+	logrus.Infof("Successfully parsed arguments")
+
+	return nil
+}
+
+// Configure applies the options to the specified config
+func (o Options) Configure(cfg engine.Interface) error {
+	err := o.UpdateConfig(cfg)
+	if err != nil {
+		return fmt.Errorf("unable to update config: %v", err)
+	}
+	return o.flush(cfg)
+}
+
+// Unconfigure removes the options from the specified config
+func (o Options) Unconfigure(cfg engine.Interface) error {
+	err := o.RevertConfig(cfg)
+	if err != nil {
+		return fmt.Errorf("unable to update config: %v", err)
+	}
+	return o.flush(cfg)
+}
+
+// flush flushes the specified config to disk
+func (o Options) flush(cfg engine.Interface) error {
+	logrus.Infof("Flushing config to %v", o.Config)
+	n, err := cfg.Save(o.Config)
+	if err != nil {
+		return fmt.Errorf("unable to flush config: %v", err)
+	}
+	if n == 0 {
+		logrus.Infof("Config file is empty, removed")
+	}
+	return nil
+}
+
+// UpdateConfig updates the specified config to include the nvidia runtimes
+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)
+		if err != nil {
+			return fmt.Errorf("failed to update runtime %q: %v", name, err)
+		}
+	}
+
+	return nil
+}
+
+// RevertConfig reverts the specified config to remove the nvidia runtimes
+func (o Options) RevertConfig(cfg engine.Interface) error {
+	runtimes := operator.GetRuntimes(
+		operator.WithNvidiaRuntimeName(o.RuntimeName),
+		operator.WithSetAsDefault(o.SetAsDefault),
+		operator.WithRoot(o.RuntimeDir),
+	)
+	for name := range runtimes {
+		err := cfg.RemoveRuntime(name)
+		if err != nil {
+			return fmt.Errorf("failed to remove runtime %q: %v", name, err)
+		}
+	}
+
+	return nil
+}
+
+// Restart restarts the specified service
+func (o Options) Restart(service string, withSignal func(string) error) error {
+	switch o.RestartMode {
+	case restartModeNone:
+		logrus.Warnf("Skipping restart of %v due to --restart-mode=%v", service, o.RestartMode)
+		return nil
+	case restartModeSignal:
+		return withSignal(o.Socket)
+	case restartModeSystemd:
+		return o.SystemdRestart(service)
+	}
+
+	return fmt.Errorf("invalid restart mode specified: %v", o.RestartMode)
+}
+
+// SystemdRestart restarts the specified service using systemd
+func (o Options) SystemdRestart(service string) error {
+	var args []string
+	var msg string
+	if o.HostRootMount != "" {
+		msg = " on host"
+		args = append(args, "chroot", o.HostRootMount)
+	}
+	args = append(args, "systemctl", "restart", service)
+
+	logrus.Infof("Restarting %v%v using systemd: %v", service, msg, args)
+
+	cmd := exec.Command(args[0], args[1:]...)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	err := cmd.Run()
+	if err != nil {
+		return fmt.Errorf("error restarting %v using systemd: %v", service, err)
+	}
+
+	return nil
+}

tools/container/containerd/config_v1_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/containerd"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 	"github.com/pelletier/go-toml"
 	"github.com/stretchr/testify/require"
 )
@@ -31,7 +32,7 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 	testCases := []struct {
 		legacyConfig                 bool
 		setAsDefault                 bool
-		runtimeClass                 string
+		runtimeName                  string
 		expectedDefaultRuntimeName   interface{}
 		expectedDefaultRuntimeBinary interface{}
 	}{
@@ -51,14 +52,14 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 		{
 			legacyConfig:                 true,
 			setAsDefault:                 true,
-			runtimeClass:                 "NAME",
+			runtimeName:                  "NAME",
 			expectedDefaultRuntimeName:   nil,
 			expectedDefaultRuntimeBinary: "/test/runtime/dir/nvidia-container-runtime",
 		},
 		{
 			legacyConfig:                 true,
 			setAsDefault:                 true,
-			runtimeClass:                 "nvidia-experimental",
+			runtimeName:                  "nvidia-experimental",
 			expectedDefaultRuntimeName:   nil,
 			expectedDefaultRuntimeBinary: "/test/runtime/dir/nvidia-container-runtime.experimental",
 		},
@@ -77,14 +78,14 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 		{
 			legacyConfig:                 false,
 			setAsDefault:                 true,
-			runtimeClass:                 "NAME",
+			runtimeName:                  "NAME",
 			expectedDefaultRuntimeName:   "NAME",
 			expectedDefaultRuntimeBinary: nil,
 		},
 		{
 			legacyConfig:                 false,
 			setAsDefault:                 true,
-			runtimeClass:                 "nvidia-experimental",
+			runtimeName:                  "nvidia-experimental",
 			expectedDefaultRuntimeName:   "nvidia-experimental",
 			expectedDefaultRuntimeBinary: nil,
 		},
@@ -93,11 +94,13 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
-				useLegacyConfig: tc.legacyConfig,
+				Options: container.Options{
-				setAsDefault:    tc.setAsDefault,
+					RuntimeName:  tc.runtimeName,
-				runtimeClass:    tc.runtimeClass,
+					RuntimeDir:   runtimeDir,
+					SetAsDefault: tc.setAsDefault,
+				},
 				runtimeType:     runtimeType,
-				runtimeDir:      runtimeDir,
+				useLegacyConfig: tc.legacyConfig,
 			}
 
 			config, err := toml.TreeFromMap(map[string]interface{}{})
@@ -109,7 +112,7 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
 				RuntimeType:           runtimeType,
 			}
 
-			err = UpdateConfig(v1, o)
+			err = o.UpdateConfig(v1)
 			require.NoError(t, err, "%d: %v", i, tc)
 
 			defaultRuntimeName := v1.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"})
@@ -138,11 +141,11 @@ func TestUpdateV1Config(t *testing.T) {
 	const runtimeDir = "/test/runtime/dir"
 
 	testCases := []struct {
-		runtimeClass   string
+		runtimeName    string
 		expectedConfig map[string]interface{}
 	}{
 		{
-			runtimeClass: "nvidia",
+			runtimeName: "nvidia",
 			expectedConfig: map[string]interface{}{
 				"version": int64(1),
 				"plugins": map[string]interface{}{
@@ -200,7 +203,7 @@ func TestUpdateV1Config(t *testing.T) {
 			},
 		},
 		{
-			runtimeClass: "NAME",
+			runtimeName: "NAME",
 			expectedConfig: map[string]interface{}{
 				"version": int64(1),
 				"plugins": map[string]interface{}{
@@ -258,7 +261,7 @@ func TestUpdateV1Config(t *testing.T) {
 			},
 		},
 		{
-			runtimeClass: "nvidia-experimental",
+			runtimeName: "nvidia-experimental",
 			expectedConfig: map[string]interface{}{
 				"version": int64(1),
 				"plugins": map[string]interface{}{
@@ -320,9 +323,10 @@ func TestUpdateV1Config(t *testing.T) {
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
-				runtimeClass: tc.runtimeClass,
+				Options: container.Options{
-				runtimeType:  runtimeType,
+					RuntimeName: tc.runtimeName,
-				runtimeDir:   runtimeDir,
+					RuntimeDir:  runtimeDir,
+				},
 			}
 
 			config, err := toml.TreeFromMap(map[string]interface{}{})
@@ -335,7 +339,7 @@ func TestUpdateV1Config(t *testing.T) {
 				ContainerAnnotations:  []string{"cdi.k8s.io/*"},
 			}
 
-			err = UpdateConfig(v1, o)
+			err = o.UpdateConfig(v1)
 			require.NoError(t, err)
 
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
@@ -350,11 +354,11 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 	const runtimeDir = "/test/runtime/dir"
 
 	testCases := []struct {
-		runtimeClass   string
+		runtimeName    string
 		expectedConfig map[string]interface{}
 	}{
 		{
-			runtimeClass: "nvidia",
+			runtimeName: "nvidia",
 			expectedConfig: map[string]interface{}{
 				"version": int64(1),
 				"plugins": map[string]interface{}{
@@ -426,7 +430,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 			},
 		},
 		{
-			runtimeClass: "NAME",
+			runtimeName: "NAME",
 			expectedConfig: map[string]interface{}{
 				"version": int64(1),
 				"plugins": map[string]interface{}{
@@ -498,7 +502,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 			},
 		},
 		{
-			runtimeClass: "nvidia-experimental",
+			runtimeName: "nvidia-experimental",
 			expectedConfig: map[string]interface{}{
 				"version": int64(1),
 				"plugins": map[string]interface{}{
@@ -574,9 +578,10 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
-				runtimeClass: tc.runtimeClass,
+				Options: container.Options{
-				runtimeType:  runtimeType,
+					RuntimeName: tc.runtimeName,
-				runtimeDir:   runtimeDir,
+					RuntimeDir:  runtimeDir,
+				},
 			}
 
 			config, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary"))
@@ -589,7 +594,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 				ContainerAnnotations:  []string{"cdi.k8s.io/*"},
 			}
 
-			err = UpdateConfig(v1, o)
+			err = o.UpdateConfig(v1)
 			require.NoError(t, err)
 
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
@@ -653,7 +658,9 @@ func TestRevertV1Config(t *testing.T) {
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
-				runtimeClass: "nvidia",
+				Options: container.Options{
+					RuntimeName: "nvidia",
+				},
 			}
 
 			config, err := toml.TreeFromMap(tc.config)
@@ -668,7 +675,7 @@ func TestRevertV1Config(t *testing.T) {
 				RuntimeType:           runtimeType,
 			}
 
-			err = RevertConfig(v1, o)
+			err = o.RevertConfig(v1)
 			require.NoError(t, err, "%d: %v", i, tc)
 
 			configContents, _ := toml.Marshal(config)

tools/container/containerd/config_v2_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/containerd"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 	"github.com/pelletier/go-toml"
 	"github.com/stretchr/testify/require"
 )
@@ -34,38 +35,38 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
 
 	testCases := []struct {
 		setAsDefault               bool
-		runtimeClass               string
+		runtimeName                string
 		expectedDefaultRuntimeName interface{}
 	}{
 		{},
 		{
 			setAsDefault:               false,
-			runtimeClass:               "nvidia",
+			runtimeName:                "nvidia",
 			expectedDefaultRuntimeName: nil,
 		},
 		{
 			setAsDefault:               false,
-			runtimeClass:               "NAME",
+			runtimeName:                "NAME",
 			expectedDefaultRuntimeName: nil,
 		},
 		{
 			setAsDefault:               false,
-			runtimeClass:               "nvidia-experimental",
+			runtimeName:                "nvidia-experimental",
 			expectedDefaultRuntimeName: nil,
 		},
 		{
 			setAsDefault:               true,
-			runtimeClass:               "nvidia",
+			runtimeName:                "nvidia",
 			expectedDefaultRuntimeName: "nvidia",
 		},
 		{
 			setAsDefault:               true,
-			runtimeClass:               "NAME",
+			runtimeName:                "NAME",
 			expectedDefaultRuntimeName: "NAME",
 		},
 		{
 			setAsDefault:               true,
-			runtimeClass:               "nvidia-experimental",
+			runtimeName:                "nvidia-experimental",
 			expectedDefaultRuntimeName: "nvidia-experimental",
 		},
 	}
@@ -73,9 +74,11 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
-				setAsDefault: tc.setAsDefault,
+				Options: container.Options{
-				runtimeClass: tc.runtimeClass,
+					RuntimeName:  tc.runtimeName,
-				runtimeDir:   runtimeDir,
+					RuntimeDir:   runtimeDir,
+					SetAsDefault: tc.setAsDefault,
+				},
 			}
 
 			config, err := toml.TreeFromMap(map[string]interface{}{})
@@ -86,7 +89,7 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
 				RuntimeType: runtimeType,
 			}
 
-			err = UpdateConfig(v2, o)
+			err = o.UpdateConfig(v2)
 			require.NoError(t, err)
 
 			defaultRuntimeName := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"})
@@ -100,11 +103,11 @@ func TestUpdateV2Config(t *testing.T) {
 	const expectedVersion = int64(2)
 
 	testCases := []struct {
-		runtimeClass   string
+		runtimeName    string
 		expectedConfig map[string]interface{}
 	}{
 		{
-			runtimeClass: "nvidia",
+			runtimeName: "nvidia",
 			expectedConfig: map[string]interface{}{
 				"version": int64(2),
 				"plugins": map[string]interface{}{
@@ -158,7 +161,7 @@ func TestUpdateV2Config(t *testing.T) {
 			},
 		},
 		{
-			runtimeClass: "NAME",
+			runtimeName: "NAME",
 			expectedConfig: map[string]interface{}{
 				"version": int64(2),
 				"plugins": map[string]interface{}{
@@ -212,7 +215,7 @@ func TestUpdateV2Config(t *testing.T) {
 			},
 		},
 		{
-			runtimeClass: "nvidia-experimental",
+			runtimeName: "nvidia-experimental",
 			expectedConfig: map[string]interface{}{
 				"version": int64(2),
 				"plugins": map[string]interface{}{
@@ -270,9 +273,11 @@ func TestUpdateV2Config(t *testing.T) {
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
-				runtimeClass: tc.runtimeClass,
+				Options: container.Options{
-				runtimeType:  runtimeType,
+					RuntimeName: tc.runtimeName,
-				runtimeDir:   runtimeDir,
+					RuntimeDir:  runtimeDir,
+				},
+				runtimeType: runtimeType,
 			}
 
 			config, err := toml.TreeFromMap(map[string]interface{}{})
@@ -284,7 +289,7 @@ func TestUpdateV2Config(t *testing.T) {
 				ContainerAnnotations: []string{"cdi.k8s.io/*"},
 			}
 
-			err = UpdateConfig(v2, o)
+			err = o.UpdateConfig(v2)
 			require.NoError(t, err)
 
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
@@ -300,11 +305,11 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 	const runtimeDir = "/test/runtime/dir"
 
 	testCases := []struct {
-		runtimeClass   string
+		runtimeName    string
 		expectedConfig map[string]interface{}
 	}{
 		{
-			runtimeClass: "nvidia",
+			runtimeName: "nvidia",
 			expectedConfig: map[string]interface{}{
 				"version": int64(2),
 				"plugins": map[string]interface{}{
@@ -372,7 +377,7 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 			},
 		},
 		{
-			runtimeClass: "NAME",
+			runtimeName: "NAME",
 			expectedConfig: map[string]interface{}{
 				"version": int64(2),
 				"plugins": map[string]interface{}{
@@ -440,7 +445,7 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 			},
 		},
 		{
-			runtimeClass: "nvidia-experimental",
+			runtimeName: "nvidia-experimental",
 			expectedConfig: map[string]interface{}{
 				"version": int64(2),
 				"plugins": map[string]interface{}{
@@ -512,9 +517,10 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
-				runtimeClass: tc.runtimeClass,
+				Options: container.Options{
-				runtimeType:  runtimeType,
+					RuntimeName: tc.runtimeName,
-				runtimeDir:   runtimeDir,
+					RuntimeDir:  runtimeDir,
+				},
 			}
 
 			config, err := toml.TreeFromMap(runcConfigMapV2("/runc-binary"))
@@ -526,7 +532,7 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 				ContainerAnnotations: []string{"cdi.k8s.io/*"},
 			}
 
-			err = UpdateConfig(v2, o)
+			err = o.UpdateConfig(v2)
 			require.NoError(t, err)
 
 			expected, err := toml.TreeFromMap(tc.expectedConfig)
@@ -585,7 +591,9 @@ func TestRevertV2Config(t *testing.T) {
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			o := &options{
-				runtimeClass: "nvidia",
+				Options: container.Options{
+					RuntimeName: "nvidia",
+				},
 			}
 
 			config, err := toml.TreeFromMap(tc.config)
@@ -599,7 +607,7 @@ func TestRevertV2Config(t *testing.T) {
 				RuntimeType: runtimeType,
 			}
 
-			err = RevertConfig(v2, o)
+			err = o.RevertConfig(v2)
 			require.NoError(t, err)
 
 			configContents, _ := toml.Marshal(config)

tools/container/containerd/containerd.go

@@ -20,33 +20,23 @@ import (
 	"fmt"
 	"net"
 	"os"
-	"os/exec"
 	"syscall"
 	"time"
 
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/containerd"
-	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/info"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 	log "github.com/sirupsen/logrus"
 	cli "github.com/urfave/cli/v2"
 )
 
 const (
-	restartModeSignal  = "signal"
-	restartModeSystemd = "systemd"
-	restartModeNone    = "none"
-
-	nvidiaRuntimeName               = "nvidia"
-	nvidiaRuntimeBinary             = "nvidia-container-runtime"
-	nvidiaExperimentalRuntimeName   = "nvidia-experimental"
-	nvidiaExperimentalRuntimeBinary = "nvidia-container-runtime.experimental"
-
 	defaultConfig        = "/etc/containerd/config.toml"
 	defaultSocket        = "/run/containerd/containerd.sock"
 	defaultRuntimeClass  = "nvidia"
 	defaultRuntmeType    = "io.containerd.runc.v2"
 	defaultSetAsDefault  = true
-	defaultRestartMode   = restartModeSignal
+	defaultRestartMode   = "signal"
 	defaultHostRootMount = "/host"
 
 	reloadBackoff     = 5 * time.Second
@@ -55,23 +45,13 @@ const (
 	socketMessageToGetPID = ""
 )
 
-// nvidiaRuntimeBinaries defines a map of runtime names to binary names
-var nvidiaRuntimeBinaries = map[string]string{
-	nvidiaRuntimeName:             nvidiaRuntimeBinary,
-	nvidiaExperimentalRuntimeName: nvidiaExperimentalRuntimeBinary,
-}
-
 // options stores the configuration from the command line or environment variables
 type options struct {
-	config          string
+	container.Options
-	socket          string
+
-	runtimeClass    string
+	// containerd-specific options
-	runtimeType     string
-	setAsDefault    bool
-	restartMode     string
-	hostRootMount   string
-	runtimeDir      string
 	useLegacyConfig bool
+	runtimeType     string
 
 	ContainerRuntimeModesCDIAnnotationPrefixes cli.StringSlice
 }
@@ -83,7 +63,7 @@ func main() {
 	c := cli.NewApp()
 	c.Name = "containerd"
 	c.Usage = "Update a containerd config with the nvidia-container-runtime"
-	c.Version = "0.1.0"
+	c.Version = info.GetVersionString()
 
 	// Create the 'setup' subcommand
 	setup := cli.Command{}
@@ -93,6 +73,9 @@ func main() {
 	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{}
@@ -102,6 +85,9 @@ func main() {
 	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{
@@ -116,57 +102,53 @@ func main() {
 	commonFlags := []cli.Flag{
 		&cli.StringFlag{
 			Name:        "config",
-			Aliases:     []string{"c"},
 			Usage:       "Path to the containerd config file",
 			Value:       defaultConfig,
-			Destination: &options.config,
+			Destination: &options.Config,
-			EnvVars:     []string{"CONTAINERD_CONFIG"},
+			EnvVars:     []string{"RUNTIME_CONFIG", "CONTAINERD_CONFIG"},
 		},
 		&cli.StringFlag{
 			Name:        "socket",
-			Aliases:     []string{"s"},
 			Usage:       "Path to the containerd socket file",
 			Value:       defaultSocket,
-			Destination: &options.socket,
+			Destination: &options.Socket,
-			EnvVars:     []string{"CONTAINERD_SOCKET"},
+			EnvVars:     []string{"RUNTIME_SOCKET", "CONTAINERD_SOCKET"},
-		},
-		&cli.StringFlag{
-			Name:        "runtime-class",
-			Aliases:     []string{"r"},
-			Usage:       "The name of the runtime class to set for the nvidia-container-runtime",
-			Value:       defaultRuntimeClass,
-			Destination: &options.runtimeClass,
-			EnvVars:     []string{"CONTAINERD_RUNTIME_CLASS"},
-		},
-		&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"},
-		},
-		// The flags below are only used by the 'setup' command.
-		&cli.BoolFlag{
-			Name:        "set-as-default",
-			Aliases:     []string{"d"},
-			Usage:       "Set nvidia-container-runtime as the default runtime",
-			Value:       defaultSetAsDefault,
-			Destination: &options.setAsDefault,
-			EnvVars:     []string{"CONTAINERD_SET_AS_DEFAULT"},
-			Hidden:      true,
 		},
 		&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,
+			Destination: &options.RestartMode,
-			EnvVars:     []string{"CONTAINERD_RESTART_MODE"},
+			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,
+			Destination: &options.HostRootMount,
 			EnvVars:     []string{"HOST_ROOT_MOUNT"},
 		},
 		&cli.BoolFlag{
@@ -175,6 +157,13 @@ func main() {
 			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,
@@ -196,14 +185,8 @@ func main() {
 func Setup(c *cli.Context, o *options) error {
 	log.Infof("Starting 'setup' for %v", c.App.Name)
 
-	runtimeDir, err := ParseArgs(c)
-	if err != nil {
-		return fmt.Errorf("unable to parse args: %v", err)
-	}
-	o.runtimeDir = runtimeDir
-
 	cfg, err := containerd.New(
-		containerd.WithPath(o.config),
+		containerd.WithPath(o.Config),
 		containerd.WithRuntimeType(o.runtimeType),
 		containerd.WithUseLegacyConfig(o.useLegacyConfig),
 		containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...),
@@ -212,18 +195,9 @@ func Setup(c *cli.Context, o *options) error {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	err = UpdateConfig(cfg, o)
+	err = o.Configure(cfg)
 	if err != nil {
-		return fmt.Errorf("unable to update config: %v", err)
+		return fmt.Errorf("unable to configure containerd: %v", err)
-	}
-
-	log.Infof("Flushing containerd config to %v", o.config)
-	n, err := cfg.Save(o.config)
-	if err != nil {
-		return fmt.Errorf("unable to flush config: %v", err)
-	}
-	if n == 0 {
-		log.Infof("Config file is empty, removed")
 	}
 
 	err = RestartContainerd(o)
@@ -240,13 +214,8 @@ func Setup(c *cli.Context, o *options) error {
 func Cleanup(c *cli.Context, o *options) error {
 	log.Infof("Starting 'cleanup' for %v", c.App.Name)
 
-	_, err := ParseArgs(c)
-	if err != nil {
-		return fmt.Errorf("unable to parse args: %v", err)
-	}
-
 	cfg, err := containerd.New(
-		containerd.WithPath(o.config),
+		containerd.WithPath(o.Config),
 		containerd.WithRuntimeType(o.runtimeType),
 		containerd.WithUseLegacyConfig(o.useLegacyConfig),
 		containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...),
@@ -255,18 +224,9 @@ func Cleanup(c *cli.Context, o *options) error {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	err = RevertConfig(cfg, o)
+	err = o.Unconfigure(cfg)
 	if err != nil {
-		return fmt.Errorf("unable to update config: %v", err)
+		return fmt.Errorf("unable to unconfigure containerd: %v", err)
-	}
-
-	log.Infof("Flushing containerd config to %v", o.config)
-	n, err := cfg.Save(o.config)
-	if err != nil {
-		return fmt.Errorf("unable to flush config: %v", err)
-	}
-	if n == 0 {
-		log.Infof("Config file is empty, removed")
 	}
 
 	err = RestartContainerd(o)
@@ -279,80 +239,18 @@ func Cleanup(c *cli.Context, o *options) error {
 	return nil
 }
 
-// ParseArgs parses the command line arguments to the CLI
-func ParseArgs(c *cli.Context) (string, error) {
-	args := c.Args()
-
-	log.Infof("Parsing arguments: %v", args.Slice())
-	if args.Len() != 1 {
-		return "", fmt.Errorf("incorrect number of arguments")
-	}
-	runtimeDir := args.Get(0)
-	log.Infof("Successfully parsed arguments")
-
-	return runtimeDir, nil
-}
-
-// UpdateConfig updates the containerd config to include the nvidia-container-runtime
-func UpdateConfig(cfg engine.Interface, o *options) error {
-	runtimes := operator.GetRuntimes(
-		operator.WithNvidiaRuntimeName(o.runtimeClass),
-		operator.WithSetAsDefault(o.setAsDefault),
-		operator.WithRoot(o.runtimeDir),
-	)
-	for class, runtime := range runtimes {
-		err := cfg.AddRuntime(class, runtime.Path, runtime.SetAsDefault)
-		if err != nil {
-			return fmt.Errorf("unable to update config for runtime class '%v': %v", class, err)
-		}
-	}
-
-	return nil
-}
-
-// RevertConfig reverts the containerd config to remove the nvidia-container-runtime
-func RevertConfig(cfg engine.Interface, o *options) error {
-	runtimes := operator.GetRuntimes(
-		operator.WithNvidiaRuntimeName(o.runtimeClass),
-		operator.WithSetAsDefault(o.setAsDefault),
-		operator.WithRoot(o.runtimeDir),
-	)
-	for class := range runtimes {
-		err := cfg.RemoveRuntime(class)
-		if err != nil {
-			return fmt.Errorf("unable to revert config for runtime class '%v': %v", class, err)
-		}
-	}
-	return nil
-}
-
 // RestartContainerd restarts containerd depending on the value of restartModeFlag
 func RestartContainerd(o *options) error {
-	switch o.restartMode {
+	return o.Restart("containerd", SignalContainerd)
-	case restartModeNone:
-		log.Warnf("Skipping sending signal to containerd due to --restart-mode=%v", o.restartMode)
-		return nil
-	case restartModeSignal:
-		err := SignalContainerd(o)
-		if err != nil {
-			return fmt.Errorf("unable to signal containerd: %v", err)
-		}
-	case restartModeSystemd:
-		return RestartContainerdSystemd(o.hostRootMount)
-	default:
-		return fmt.Errorf("Invalid restart mode specified: %v", o.restartMode)
-	}
-
-	return nil
 }
 
 // SignalContainerd sends a SIGHUP signal to the containerd daemon
-func SignalContainerd(o *options) error {
+func SignalContainerd(socket string) error {
 	log.Infof("Sending SIGHUP signal to containerd")
 
 	// Wrap the logic to perform the SIGHUP in a function so we can retry it on failure
 	retriable := func() error {
-		conn, err := net.Dial("unix", o.socket)
+		conn, err := net.Dial("unix", socket)
 		if err != nil {
 			return fmt.Errorf("unable to dial: %v", err)
 		}
@@ -426,24 +324,6 @@ func SignalContainerd(o *options) error {
 	return nil
 }
 
-// RestartContainerdSystemd restarts containerd using systemctl
-func RestartContainerdSystemd(hostRootMount string) error {
-	log.Infof("Restarting containerd using systemd and host root mounted at %v", hostRootMount)
-
-	command := "chroot"
-	args := []string{hostRootMount, "systemctl", "restart", "containerd"}
-
-	cmd := exec.Command(command, args...)
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-	err := cmd.Run()
-	if err != nil {
-		return fmt.Errorf("error restarting containerd using systemd: %v", err)
-	}
-
-	return nil
-}
-
 // containerAnnotationsFromCDIPrefixes returns the container annotations to set for the given CDI prefixes.
 func (o *options) containerAnnotationsFromCDIPrefixes() []string {
 	var annotations []string

tools/container/crio/crio.go

@@ -20,21 +20,17 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
-	"os/exec"
 	"path/filepath"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/crio"
-	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/info"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 	log "github.com/sirupsen/logrus"
 	cli "github.com/urfave/cli/v2"
 )
 
 const (
-	restartModeSystemd = "systemd"
-	restartModeNone    = "none"
-
 	defaultConfigMode = "hook"
 
 	// Hook-based settings
@@ -43,25 +39,22 @@ const (
 
 	// Config-based settings
 	defaultConfig        = "/etc/crio/crio.conf"
+	defaultSocket        = "/var/run/crio/crio.sock"
 	defaultRuntimeClass  = "nvidia"
 	defaultSetAsDefault  = true
-	defaultRestartMode   = restartModeSystemd
+	defaultRestartMode   = "systemd"
 	defaultHostRootMount = "/host"
 )
 
 // options stores the configuration from the command linek or environment variables
 type options struct {
+	container.Options
+
 	configMode string
 
+	// hook-specific options
 	hooksDir     string
 	hookFilename string
-	runtimeDir   string
-
-	config        string
-	runtimeClass  string
-	setAsDefault  bool
-	restartMode   string
-	hostRootMount string
 }
 
 func main() {
@@ -71,8 +64,7 @@ func main() {
 	c := cli.NewApp()
 	c.Name = "crio"
 	c.Usage = "Update cri-o hooks to include the NVIDIA runtime hook"
-	c.ArgsUsage = "<toolkit_dirname>"
+	c.Version = info.GetVersionString()
-	c.Version = "0.1.0"
 
 	// Create the 'setup' subcommand
 	setup := cli.Command{}
@@ -83,7 +75,7 @@ func main() {
 		return Setup(c, &options)
 	}
 	setup.Before = func(c *cli.Context) error {
-		return ParseArgs(c, &options)
+		return container.ParseArgs(c, &options.Options)
 	}
 
 	// Create the 'cleanup' subcommand
@@ -93,6 +85,10 @@ func main() {
 	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,
@@ -104,9 +100,61 @@ func main() {
 	// 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 cri-o config file",
+			Value:       defaultConfig,
+			Destination: &options.Config,
+			EnvVars:     []string{"RUNTIME_CONFIG", "CRIO_CONFIG"},
+		},
+		&cli.StringFlag{
+			Name:        "socket",
+			Usage:       "Path to the crio socket file",
+			Value:       "",
+			Destination: &options.Socket,
+			EnvVars:     []string{"RUNTIME_SOCKET", "CRIO_SOCKET"},
+			// Note: We hide this option since restarting cri-o via a socket is not supported.
+			Hidden: true,
+		},
+		&cli.StringFlag{
+			Name:        "restart-mode",
+			Usage:       "Specify how cri-o should be restarted;  If 'none' is selected, it will not be restarted [systemd | none]",
+			Value:       defaultRestartMode,
+			Destination: &options.RestartMode,
+			EnvVars:     []string{"RUNTIME_RESTART_MODE", "CRIO_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", "CRIO_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", "CRIO_SET_AS_DEFAULT"},
+			Hidden:      true,
+		},
+		&cli.StringFlag{
+			Name:        "host-root",
+			Usage:       "Specify the path to the host root to be used when restarting crio using systemd",
+			Value:       defaultHostRootMount,
+			Destination: &options.HostRootMount,
+			EnvVars:     []string{"HOST_ROOT_MOUNT"},
+		},
 		&cli.StringFlag{
 			Name:        "hooks-dir",
-			Aliases:     []string{"d"},
 			Usage:       "path to the cri-o hooks directory",
 			Value:       defaultHooksDir,
 			Destination: &options.hooksDir,
@@ -115,7 +163,6 @@ func main() {
 		},
 		&cli.StringFlag{
 			Name:        "hook-filename",
-			Aliases:     []string{"f"},
 			Usage:       "filename of the cri-o hook that will be created / removed in the hooks directory",
 			Value:       defaultHookFilename,
 			Destination: &options.hookFilename,
@@ -129,43 +176,6 @@ func main() {
 			Destination: &options.configMode,
 			EnvVars:     []string{"CRIO_CONFIG_MODE"},
 		},
-		&cli.StringFlag{
-			Name:        "config",
-			Usage:       "Path to the cri-o config file",
-			Value:       defaultConfig,
-			Destination: &options.config,
-			EnvVars:     []string{"CRIO_CONFIG"},
-		},
-		&cli.StringFlag{
-			Name:        "runtime-class",
-			Usage:       "The name of the runtime class to set for the nvidia-container-runtime",
-			Value:       defaultRuntimeClass,
-			Destination: &options.runtimeClass,
-			EnvVars:     []string{"CRIO_RUNTIME_CLASS"},
-		},
-		// The flags below are only used by the 'setup' command.
-		&cli.BoolFlag{
-			Name:        "set-as-default",
-			Usage:       "Set nvidia-container-runtime as the default runtime",
-			Value:       defaultSetAsDefault,
-			Destination: &options.setAsDefault,
-			EnvVars:     []string{"CRIO_SET_AS_DEFAULT"},
-			Hidden:      true,
-		},
-		&cli.StringFlag{
-			Name:        "restart-mode",
-			Usage:       "Specify how cri-o should be restarted;  If 'none' is selected, it will not be restarted [systemd | none]",
-			Value:       defaultRestartMode,
-			Destination: &options.restartMode,
-			EnvVars:     []string{"CRIO_RESTART_MODE"},
-		},
-		&cli.StringFlag{
-			Name:        "host-root",
-			Usage:       "Specify the path to the host root to be used when restarting crio using systemd",
-			Value:       defaultHostRootMount,
-			Destination: &options.hostRootMount,
-			EnvVars:     []string{"HOST_ROOT_MOUNT"},
-		},
 	}
 
 	// Update the subcommand flags with the common subcommand flags
@@ -202,7 +212,7 @@ func setupHook(o *options) error {
 	}
 
 	hookPath := getHookPath(o.hooksDir, o.hookFilename)
-	err = createHook(o.runtimeDir, hookPath)
+	err = createHook(o.RuntimeDir, hookPath)
 	if err != nil {
 		return fmt.Errorf("error creating hook: %v", err)
 	}
@@ -215,24 +225,15 @@ func setupConfig(o *options) error {
 	log.Infof("Updating config file")
 
 	cfg, err := crio.New(
-		crio.WithPath(o.config),
+		crio.WithPath(o.Config),
 	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	err = UpdateConfig(cfg, o)
+	err = o.Configure(cfg)
 	if err != nil {
-		return fmt.Errorf("unable to update config: %v", err)
+		return fmt.Errorf("unable to configure cri-o: %v", err)
-	}
-
-	log.Infof("Flushing cri-o config to %v", o.config)
-	n, err := cfg.Save(o.config)
-	if err != nil {
-		return fmt.Errorf("unable to flush config: %v", err)
-	}
-	if n == 0 {
-		log.Infof("Config file is empty, removed")
 	}
 
 	err = RestartCrio(o)
@@ -275,24 +276,15 @@ func cleanupConfig(o *options) error {
 	log.Infof("Reverting config file modifications")
 
 	cfg, err := crio.New(
-		crio.WithPath(o.config),
+		crio.WithPath(o.Config),
 	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	err = RevertConfig(cfg, o)
+	err = o.Unconfigure(cfg)
 	if err != nil {
-		return fmt.Errorf("unable to update config: %v", err)
+		return fmt.Errorf("unable to unconfigure cri-o: %v", err)
-	}
-
-	log.Infof("Flushing cri-o config to %v", o.config)
-	n, err := cfg.Save(o.config)
-	if err != nil {
-		return fmt.Errorf("unable to flush config: %v", err)
-	}
-	if n == 0 {
-		log.Infof("Config file is empty, removed")
 	}
 
 	err = RestartCrio(o)
@@ -303,20 +295,6 @@ func cleanupConfig(o *options) error {
 	return nil
 }
 
-// ParseArgs parses the command line arguments to the CLI
-func ParseArgs(c *cli.Context, o *options) error {
-	args := c.Args()
-
-	log.Infof("Parsing arguments: %v", args.Slice())
-	if c.NArg() != 1 {
-		return fmt.Errorf("incorrect number of arguments")
-	}
-	o.runtimeDir = args.Get(0)
-	log.Infof("Successfully parsed arguments")
-
-	return nil
-}
-
 func createHook(toolkitDir string, hookPath string) error {
 	hook, err := os.Create(hookPath)
 	if err != nil {
@@ -357,66 +335,7 @@ func generateOciHook(toolkitDir string) podmanHook {
 	return hook
 }
 
-// UpdateConfig updates the cri-o config to include the NVIDIA Container Runtime
-func UpdateConfig(cfg engine.Interface, o *options) error {
-	runtimes := operator.GetRuntimes(
-		operator.WithNvidiaRuntimeName(o.runtimeClass),
-		operator.WithSetAsDefault(o.setAsDefault),
-		operator.WithRoot(o.runtimeDir),
-	)
-	for class, runtime := range runtimes {
-		err := cfg.AddRuntime(class, runtime.Path, runtime.SetAsDefault)
-		if err != nil {
-			return fmt.Errorf("unable to update config for runtime class '%v': %v", class, err)
-		}
-	}
-
-	return nil
-}
-
-// RevertConfig reverts the cri-o config to remove the NVIDIA Container Runtime
-func RevertConfig(cfg engine.Interface, o *options) error {
-	runtimes := operator.GetRuntimes(
-		operator.WithNvidiaRuntimeName(o.runtimeClass),
-		operator.WithSetAsDefault(o.setAsDefault),
-		operator.WithRoot(o.runtimeDir),
-	)
-	for class := range runtimes {
-		err := cfg.RemoveRuntime(class)
-		if err != nil {
-			return fmt.Errorf("unable to revert config for runtime class '%v': %v", class, err)
-		}
-	}
-	return nil
-}
-
 // RestartCrio restarts crio depending on the value of restartModeFlag
 func RestartCrio(o *options) error {
-	switch o.restartMode {
+	return o.Restart("crio", func(string) error { return fmt.Errorf("supporting crio via signal is unsupported") })
-	case restartModeNone:
-		log.Warnf("Skipping restart of crio due to --restart-mode=%v", o.restartMode)
-		return nil
-	case restartModeSystemd:
-		return RestartCrioSystemd(o.hostRootMount)
-	default:
-		return fmt.Errorf("invalid restart mode specified: %v", o.restartMode)
-	}
-}
-
-// RestartCrioSystemd restarts cri-o using systemctl
-func RestartCrioSystemd(hostRootMount string) error {
-	log.Infof("Restarting cri-o using systemd and host root mounted at %v", hostRootMount)
-
-	command := "chroot"
-	args := []string{hostRootMount, "systemctl", "restart", "crio"}
-
-	cmd := exec.Command(command, args...)
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-	err := cmd.Run()
-	if err != nil {
-		return fmt.Errorf("error restarting crio using systemd: %v", err)
-	}
-
-	return nil
 }

tools/container/docker/docker.go

@@ -23,50 +23,31 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/docker"
-	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/info"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 	log "github.com/sirupsen/logrus"
 	cli "github.com/urfave/cli/v2"
 )
 
 const (
-	restartModeSignal = "signal"
-	restartModeNone   = "none"
-
-	nvidiaRuntimeName               = "nvidia"
-	nvidiaRuntimeBinary             = "nvidia-container-runtime"
-	nvidiaExperimentalRuntimeName   = "nvidia-experimental"
-	nvidiaExperimentalRuntimeBinary = "nvidia-container-runtime.experimental"
-
 	defaultConfig       = "/etc/docker/daemon.json"
 	defaultSocket       = "/var/run/docker.sock"
 	defaultSetAsDefault = true
 	// defaultRuntimeName specifies the NVIDIA runtime to be use as the default runtime if setting the default runtime is enabled
-	defaultRuntimeName = nvidiaRuntimeName
+	defaultRuntimeName   = "nvidia"
-	defaultRestartMode = restartModeSignal
+	defaultRestartMode   = "signal"
+	defaultHostRootMount = "/host"
 
 	reloadBackoff     = 5 * time.Second
 	maxReloadAttempts = 6
 
-	defaultDockerRuntime  = "runc"
 	socketMessageToGetPID = "GET /info HTTP/1.0\r\n\r\n"
 )
 
-// nvidiaRuntimeBinaries defines a map of runtime names to binary names
-var nvidiaRuntimeBinaries = map[string]string{
-	nvidiaRuntimeName:             nvidiaRuntimeBinary,
-	nvidiaExperimentalRuntimeName: nvidiaExperimentalRuntimeBinary,
-}
-
 // options stores the configuration from the command line or environment variables
 type options struct {
-	config       string
+	container.Options
-	socket       string
-	runtimeName  string
-	setAsDefault bool
-	runtimeDir   string
-	restartMode  string
 }
 
 func main() {
@@ -76,7 +57,7 @@ func main() {
 	c := cli.NewApp()
 	c.Name = "docker"
 	c.Usage = "Update docker config with the nvidia runtime"
-	c.Version = "0.1.0"
+	c.Version = info.GetVersionString()
 
 	// Create the 'setup' subcommand
 	setup := cli.Command{}
@@ -86,6 +67,9 @@ func main() {
 	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{}
@@ -95,6 +79,9 @@ func main() {
 	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{
@@ -109,44 +96,57 @@ func main() {
 	commonFlags := []cli.Flag{
 		&cli.StringFlag{
 			Name:        "config",
-			Aliases:     []string{"c"},
 			Usage:       "Path to docker config file",
 			Value:       defaultConfig,
-			Destination: &options.config,
+			Destination: &options.Config,
-			EnvVars:     []string{"DOCKER_CONFIG"},
+			EnvVars:     []string{"RUNTIME_CONFIG", "DOCKER_CONFIG"},
 		},
 		&cli.StringFlag{
 			Name:        "socket",
-			Aliases:     []string{"s"},
 			Usage:       "Path to the docker socket file",
 			Value:       defaultSocket,
-			Destination: &options.socket,
+			Destination: &options.Socket,
-			EnvVars:     []string{"DOCKER_SOCKET"},
+			EnvVars:     []string{"RUNTIME_SOCKET", "DOCKER_SOCKET"},
-		},
-		// The flags below are only used by the 'setup' command.
-		&cli.StringFlag{
-			Name:        "runtime-name",
-			Aliases:     []string{"r"},
-			Usage:       "Specify the name of the `nvidia` runtime. If set-as-default is selected, the runtime is used as the default runtime.",
-			Value:       defaultRuntimeName,
-			Destination: &options.runtimeName,
-			EnvVars:     []string{"DOCKER_RUNTIME_NAME"},
-		},
-		&cli.BoolFlag{
-			Name:        "set-as-default",
-			Aliases:     []string{"d"},
-			Usage:       "Set the `nvidia` runtime as the default runtime. If --runtime-name is specified as `nvidia-experimental` the experimental runtime is set as the default runtime instead",
-			Value:       defaultSetAsDefault,
-			Destination: &options.setAsDefault,
-			EnvVars:     []string{"DOCKER_SET_AS_DEFAULT"},
-			Hidden:      true,
 		},
 		&cli.StringFlag{
 			Name:        "restart-mode",
-			Usage:       "Specify how docker should be restarted; If 'none' is selected it will not be restarted [signal | none]",
+			Usage:       "Specify how docker should be restarted; If 'none' is selected it will not be restarted [signal | systemd | none ]",
 			Value:       defaultRestartMode,
-			Destination: &options.restartMode,
+			Destination: &options.RestartMode,
-			EnvVars:     []string{"DOCKER_RESTART_MODE"},
+			EnvVars:     []string{"RUNTIME_RESTART_MODE", "DOCKER_RESTART_MODE"},
+		},
+		&cli.StringFlag{
+			Name:        "host-root",
+			Usage:       "Specify the path to the host root to be used when restarting docker using systemd",
+			Value:       defaultHostRootMount,
+			Destination: &options.HostRootMount,
+			EnvVars:     []string{"HOST_ROOT_MOUNT"},
+			// Restart using systemd is currently not supported.
+			// We hide this option for the time being.
+			Hidden: true,
+		},
+		&cli.StringFlag{
+			Name:        "runtime-name",
+			Aliases:     []string{"nvidia-runtime-name", "runtime-class"},
+			Usage:       "Specify the name of the `nvidia` runtime. If set-as-default is selected, the runtime is used as the default runtime.",
+			Value:       defaultRuntimeName,
+			Destination: &options.RuntimeName,
+			EnvVars:     []string{"NVIDIA_RUNTIME_NAME", "DOCKER_RUNTIME_NAME"},
+		},
+		&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 the `nvidia` runtime as the default runtime. If --runtime-name is specified as `nvidia-experimental` the experimental runtime is set as the default runtime instead",
+			Value:       defaultSetAsDefault,
+			Destination: &options.SetAsDefault,
+			EnvVars:     []string{"NVIDIA_RUNTIME_SET_AS_DEFAULT", "DOCKER_SET_AS_DEFAULT"},
+			Hidden:      true,
 		},
 	}
 
@@ -165,28 +165,16 @@ func main() {
 func Setup(c *cli.Context, o *options) error {
 	log.Infof("Starting 'setup' for %v", c.App.Name)
 
-	runtimeDir, err := ParseArgs(c)
-	if err != nil {
-		return fmt.Errorf("unable to parse args: %v", err)
-	}
-	o.runtimeDir = runtimeDir
-
 	cfg, err := docker.New(
-		docker.WithPath(o.config),
+		docker.WithPath(o.Config),
 	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	err = UpdateConfig(cfg, o)
+	err = o.Configure(cfg)
 	if err != nil {
-		return fmt.Errorf("unable to update config: %v", err)
+		return fmt.Errorf("unable to configure docker: %v", err)
-	}
-
-	log.Infof("Flushing docker config to %v", o.config)
-	_, err = cfg.Save(o.config)
-	if err != nil {
-		return fmt.Errorf("unable to flush config: %v", err)
 	}
 
 	err = RestartDocker(o)
@@ -203,30 +191,16 @@ func Setup(c *cli.Context, o *options) error {
 func Cleanup(c *cli.Context, o *options) error {
 	log.Infof("Starting 'cleanup' for %v", c.App.Name)
 
-	_, err := ParseArgs(c)
-	if err != nil {
-		return fmt.Errorf("unable to parse args: %v", err)
-	}
-
 	cfg, err := docker.New(
-		docker.WithPath(o.config),
+		docker.WithPath(o.Config),
 	)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
-	err = RevertConfig(cfg, o)
+	err = o.Unconfigure(cfg)
 	if err != nil {
-		return fmt.Errorf("unable to update config: %v", err)
+		return fmt.Errorf("unable to unconfigure docker: %v", err)
-	}
-
-	log.Infof("Flushing docker config to %v", o.config)
-	n, err := cfg.Save(o.config)
-	if err != nil {
-		return fmt.Errorf("unable to flush config: %v", err)
-	}
-	if n == 0 {
-		log.Infof("Config file is empty, removed")
 	}
 
 	err = RestartDocker(o)
@@ -239,69 +213,9 @@ func Cleanup(c *cli.Context, o *options) error {
 	return nil
 }
 
-// ParseArgs parses the command line arguments to the CLI
-func ParseArgs(c *cli.Context) (string, error) {
-	args := c.Args()
-
-	log.Infof("Parsing arguments: %v", args.Slice())
-	if args.Len() != 1 {
-		return "", fmt.Errorf("incorrect number of arguments")
-	}
-	runtimeDir := args.Get(0)
-	log.Infof("Successfully parsed arguments")
-
-	return runtimeDir, nil
-}
-
-// UpdateConfig updates the docker config to include the nvidia runtimes
-func UpdateConfig(cfg engine.Interface, o *options) 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)
-		if err != nil {
-			return fmt.Errorf("failed to update runtime %q: %v", name, err)
-		}
-	}
-
-	return nil
-}
-
-// RevertConfig reverts the docker config to remove the nvidia runtime
-func RevertConfig(cfg engine.Interface, o *options) error {
-	runtimes := operator.GetRuntimes(
-		operator.WithNvidiaRuntimeName(o.runtimeName),
-		operator.WithSetAsDefault(o.setAsDefault),
-		operator.WithRoot(o.runtimeDir),
-	)
-	for name := range runtimes {
-		err := cfg.RemoveRuntime(name)
-		if err != nil {
-			return fmt.Errorf("failed to remove runtime %q: %v", name, err)
-		}
-	}
-
-	return nil
-}
-
 // RestartDocker restarts docker depending on the value of restartModeFlag
 func RestartDocker(o *options) error {
-	switch o.restartMode {
+	return o.Restart("docker", SignalDocker)
-	case restartModeNone:
-		log.Warnf("Skipping sending signal to docker due to --restart-mode=%v", o.restartMode)
-	case restartModeSignal:
-		err := SignalDocker(o.socket)
-		if err != nil {
-			return fmt.Errorf("unable to signal docker: %v", err)
-		}
-	default:
-		return fmt.Errorf("invalid restart mode specified: %v", o.restartMode)
-	}
-
-	return nil
 }
 
 // SignalDocker sends a SIGHUP signal to docker daemon

tools/container/docker/docker_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/docker"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
 	"github.com/stretchr/testify/require"
 )
 
@@ -56,14 +57,16 @@ func TestUpdateConfigDefaultRuntime(t *testing.T) {
 
 	for i, tc := range testCases {
 		o := &options{
-			setAsDefault: tc.setAsDefault,
+			Options: container.Options{
-			runtimeName:  tc.runtimeName,
+				RuntimeName:  tc.runtimeName,
-			runtimeDir:   runtimeDir,
+				RuntimeDir:   runtimeDir,
+				SetAsDefault: tc.setAsDefault,
+			},
 		}
 
 		config := docker.Config(map[string]interface{}{})
 
-		err := UpdateConfig(&config, o)
+		err := o.UpdateConfig(&config)
 		require.NoError(t, err, "%d: %v", i, tc)
 
 		defaultRuntimeName := config["default-runtime"]
@@ -314,13 +317,15 @@ func TestUpdateConfig(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
-		options := &options{
+		o := &options{
-			setAsDefault: tc.setAsDefault,
+			Options: container.Options{
-			runtimeName:  tc.runtimeName,
+				RuntimeName:  tc.runtimeName,
-			runtimeDir:   runtimeDir,
+				RuntimeDir:   runtimeDir,
+				SetAsDefault: tc.setAsDefault,
+			},
 		}
 
-		err := UpdateConfig(&tc.config, options)
+		err := o.UpdateConfig(&tc.config)
 		require.NoError(t, err, "%d: %v", i, tc)
 
 		configContent, err := json.MarshalIndent(tc.config, "", "    ")
@@ -457,7 +462,8 @@ func TestRevertConfig(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
-		err := RevertConfig(&tc.config, &options{})
+		o := &options{}
+		err := o.RevertConfig(&tc.config)
 
 		require.NoError(t, err, "%d: %v", i, tc)