Refactor toolking to setup and cleanup configs

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

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,
-				setAsDefault:    tc.setAsDefault,
-				runtimeClass:    tc.runtimeClass,
+				Options: container.Options{
+					RuntimeName:  tc.runtimeName,
+					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,
-				runtimeType:  runtimeType,
-				runtimeDir:   runtimeDir,
+				Options: container.Options{
+					RuntimeName: tc.runtimeName,
+					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,
-				runtimeType:  runtimeType,
-				runtimeDir:   runtimeDir,
+				Options: container.Options{
+					RuntimeName: tc.runtimeName,
+					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,
-				runtimeClass: tc.runtimeClass,
-				runtimeDir:   runtimeDir,
+				Options: container.Options{
+					RuntimeName:  tc.runtimeName,
+					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,
-				runtimeType:  runtimeType,
-				runtimeDir:   runtimeDir,
+				Options: container.Options{
+					RuntimeName: tc.runtimeName,
+					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,
-				runtimeType:  runtimeType,
-				runtimeDir:   runtimeDir,
+				Options: container.Options{
+					RuntimeName: tc.runtimeName,
+					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

@@ -24,9 +24,8 @@ import (
 	"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/tools/container"
 	log "github.com/sirupsen/logrus"
 	cli "github.com/urfave/cli/v2"
 )
@@ -63,13 +62,7 @@ var nvidiaRuntimeBinaries = map[string]string{
 
 // options stores the configuration from the command line or environment variables
 type options struct {
-	config        string
-	socket        string
-	runtimeClass  string
-	runtimeDir    string
-	setAsDefault  bool
-	restartMode   string
-	hostRootMount string
+	container.Options
 
 	// containerd-specific options
 	useLegacyConfig bool
@@ -96,7 +89,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
@@ -108,7 +101,7 @@ func main() {
 		return Cleanup(c, &options)
 	}
 	cleanup.Before = func(c *cli.Context) error {
-		return ParseArgs(c, &options)
+		return container.ParseArgs(c, &options.Options)
 	}
 
 	// Register the subcommands with the top-level CLI
@@ -126,21 +119,21 @@ func main() {
 			Name:        "config",
 			Usage:       "Path to the containerd config file",
 			Value:       defaultConfig,
-			Destination: &options.config,
+			Destination: &options.Config,
 			EnvVars:     []string{"CONTAINERD_CONFIG", "RUNTIME_CONFIG"},
 		},
 		&cli.StringFlag{
 			Name:        "socket",
 			Usage:       "Path to the containerd socket file",
 			Value:       defaultSocket,
-			Destination: &options.socket,
+			Destination: &options.Socket,
 			EnvVars:     []string{"CONTAINERD_SOCKET", "RUNTIME_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,
+			Destination: &options.RestartMode,
 			EnvVars:     []string{"CONTAINERD_RESTART_MODE", "RUNTIME_RESTART_MODE"},
 		},
 		&cli.StringFlag{
@@ -148,21 +141,21 @@ func main() {
 			Aliases:     []string{"runtime-name", "nvidia-runtime-name"},
 			Usage:       "The name of the runtime class to set for the nvidia-container-runtime",
 			Value:       defaultRuntimeClass,
-			Destination: &options.runtimeClass,
+			Destination: &options.RuntimeName,
 			EnvVars:     []string{"CONTAINERD_RUNTIME_CLASS", "NVIDIA_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,
+			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,
+			Destination: &options.SetAsDefault,
 			EnvVars:     []string{"CONTAINERD_SET_AS_DEFAULT", "NVIDIA_RUNTIME_SET_AS_DEFAULT"},
 			Hidden:      true,
 		},
@@ -170,7 +163,7 @@ func main() {
 			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{
@@ -208,7 +201,7 @@ 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.WithPath(o.Config),
 		containerd.WithRuntimeType(o.runtimeType),
 		containerd.WithUseLegacyConfig(o.useLegacyConfig),
 		containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...),
@@ -217,18 +210,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)
-	}
-
-	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")
+		return fmt.Errorf("unable to configure containerd: %v", err)
 	}
 
 	err = RestartContainerd(o)
@@ -246,7 +230,7 @@ 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.WithPath(o.Config),
 		containerd.WithRuntimeType(o.runtimeType),
 		containerd.WithUseLegacyConfig(o.useLegacyConfig),
 		containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...),
@@ -255,18 +239,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)
-	}
-
-	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")
+		return fmt.Errorf("unable to unconfigure containerd: %v", err)
 	}
 
 	err = RestartContainerd(o)
@@ -279,65 +254,11 @@ func Cleanup(c *cli.Context, o *options) error {
 	return nil
 }
 
-// ParseArgs parses the command line arguments to the CLI
-func ParseArgs(c *cli.Context, o *options) error {
-	if o.runtimeDir != "" {
-		log.Debug("Runtime directory already set")
-		return nil
-	}
-
-	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
-}
-
-// 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 {
+	switch o.RestartMode {
 	case restartModeNone:
-		log.Warnf("Skipping sending signal to containerd due to --restart-mode=%v", o.restartMode)
+		log.Warnf("Skipping sending signal to containerd due to --restart-mode=%v", o.RestartMode)
 		return nil
 	case restartModeSignal:
 		err := SignalContainerd(o)
@@ -345,9 +266,9 @@ func RestartContainerd(o *options) error {
 			return fmt.Errorf("unable to signal containerd: %v", err)
 		}
 	case restartModeSystemd:
-		return RestartContainerdSystemd(o.hostRootMount)
+		return RestartContainerdSystemd(o.HostRootMount)
 	default:
-		return fmt.Errorf("Invalid restart mode specified: %v", o.restartMode)
+		return fmt.Errorf("Invalid restart mode specified: %v", o.RestartMode)
 	}
 
 	return nil
@@ -359,7 +280,7 @@ func SignalContainerd(o *options) error {
 
 	// 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", o.Socket)
 		if err != nil {
 			return fmt.Errorf("unable to dial: %v", err)
 		}