[nvidia-ctk-installer] introduce unconfigure mode to determine cleanup behavior

Signed-off-by: Christopher Desiniotis <[email protected]>

Author: cdesiniotis

cmd/nvidia-ctk-installer/container/container.go

@@ -57,6 +57,8 @@ type Options struct {
 	HostRootMount string
 
 	ConfigSources []string
+
+	UnconfigureMode string
 }
 
 // Configure applies the options to the specified config
@@ -68,6 +70,29 @@ func (o Options) Configure(cfg engine.Interface) error {
 	return o.flush(cfg)
 }
 
+func (o Options) UnsetDefaultRuntime(cfg engine.Interface) error {
+	runtimes := operator.GetRuntimes(
+		operator.WithNvidiaRuntimeName(o.RuntimeName),
+		operator.WithSetAsDefault(o.SetAsDefault),
+		operator.WithRoot(o.RuntimeDir),
+	)
+	defaultRuntimeName := ""
+	for name, runtime := range runtimes {
+		if runtime.SetAsDefault {
+			defaultRuntimeName = name
+		}
+	}
+	if defaultRuntimeName == "" {
+		return nil
+	}
+
+	cfg.UnsetDefaultRuntime(defaultRuntimeName)
+	if err := o.flush(cfg); err != nil {
+		return err
+	}
+	return nil
+}
+
 // Unconfigure removes the options from the specified config
 func (o Options) Unconfigure(cfg engine.Interface) error {
 	err := o.RevertConfig(cfg)

cmd/nvidia-ctk-installer/container/runtime/containerd/containerd.go

@@ -39,6 +39,8 @@ const (
 	DefaultSocket      = "/run/containerd/containerd.sock"
 	DefaultRestartMode = "signal"
 
+	DefaultUnconfigureMode = "all"
+
 	defaultRuntimeType = "io.containerd.runc.v2"
 )
 
@@ -115,11 +117,21 @@ func Setup(c *cli.Command, o *container.Options, co *Options) error {
 func Cleanup(c *cli.Command, o *container.Options, co *Options) error {
 	log.Infof("Starting 'cleanup' for %v", c.Name)
 
+	if o.UnconfigureMode == "none" {
+		log.Infof("Skipping cleanup as unconfigure mode is set to 'none'")
+		return nil
+	}
+
 	cfg, err := getRuntimeConfig(o, co)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
+	if o.UnconfigureMode == "default-runtime-only" {
+		log.Infof("Removing nvidia as the default runtime if configured")
+		return o.UnsetDefaultRuntime(cfg)
+	}
+
 	err = o.Unconfigure(cfg)
 	if err != nil {
 		return fmt.Errorf("unable to unconfigure containerd: %v", err)

cmd/nvidia-ctk-installer/container/runtime/crio/config_test.go

@@ -88,7 +88,6 @@ func TestCrioConfigLifecycle(t *testing.T) {
 			},
 			assertCleanupPostConditions: func(t *testing.T, co *container.Options, _ *Options) error {
 				require.NoFileExists(t, co.TopLevelConfigPath)
-				require.NoFileExists(t, co.DropInConfig)
 				return nil
 			},
 		},
@@ -182,8 +181,6 @@ signature_policy = "/etc/crio/policy.json"
 			assertCleanupPostConditions: func(t *testing.T, co *container.Options, o *Options) error {
 				require.FileExists(t, co.TopLevelConfigPath)
 
-				require.NoFileExists(t, co.DropInConfig)
-
 				actualTopLevel, err := os.ReadFile(co.TopLevelConfigPath)
 				require.NoError(t, err)
 
@@ -480,8 +477,6 @@ plugin_dirs = [
 `
 				require.Equal(t, expected, string(actual))
 
-				require.NoFileExists(t, co.DropInConfig)
-
 				return nil
 			},
 		},

cmd/nvidia-ctk-installer/container/runtime/crio/crio.go

@@ -51,6 +51,8 @@ const (
 
 	DefaultSocket      = "/var/run/crio/crio.sock"
 	DefaultRestartMode = "systemd"
+
+	DefaultUnconfigureMode = "default-runtime-only"
 )
 
 // Options defines the cri-o specific options.
@@ -182,11 +184,21 @@ func cleanupHook(co *Options) error {
 func cleanupConfig(o *container.Options) error {
 	log.Infof("Reverting config file modifications")
 
+	if o.UnconfigureMode == "none" {
+		log.Infof("Skipping cleanup as unconfigure mode is set to 'none'")
+		return nil
+	}
+
 	cfg, err := getRuntimeConfig(o)
 	if err != nil {
 		return fmt.Errorf("unable to load config: %v", err)
 	}
 
+	if o.UnconfigureMode == "default-runtime-only" {
+		log.Infof("Removing nvidia as the default runtime if configured")
+		return o.UnsetDefaultRuntime(cfg)
+	}
+
 	err = o.Unconfigure(cfg)
 	if err != nil {
 		return fmt.Errorf("unable to unconfigure cri-o: %v", err)

cmd/nvidia-ctk-installer/container/runtime/docker/docker.go

@@ -30,9 +30,10 @@ import (
 const (
 	Name = "docker"
 
-	DefaultConfig      = "/etc/docker/daemon.json"
-	DefaultSocket      = "/var/run/docker.sock"
-	DefaultRestartMode = "signal"
+	DefaultConfig          = "/etc/docker/daemon.json"
+	DefaultSocket          = "/var/run/docker.sock"
+	DefaultRestartMode     = "signal"
+	DefaultUnconfigureMode = "all"
 )
 
 type Options struct{}

cmd/nvidia-ctk-installer/container/runtime/runtime.go

@@ -128,6 +128,13 @@ func Flags(opts *Options) []cli.Flag {
 			Destination: &opts.ConfigSources,
 			Sources:     cli.EnvVars("RUNTIME_CONFIG_SOURCES", "RUNTIME_CONFIG_SOURCE"),
 		},
+		&cli.StringFlag{
+			Name:        "unconfigure-mode",
+			Usage:       "Specify how the runtime should be unconfigured; If 'none' is selected nothing will be reverted in the config [ all | default-runtime-only | none ]",
+			Value:       runtimeSpecificDefault,
+			Destination: &opts.UnconfigureMode,
+			Sources:     cli.EnvVars("RUNTIME_UNCONFIGURE_MODE"),
+		},
 	}
 
 	flags = append(flags, containerd.Flags(&opts.containerdOptions)...)
@@ -178,6 +185,9 @@ func (opts *Options) Validate(logger logger.Interface, c *cli.Command, runtime s
 		if opts.RestartMode == runtimeSpecificDefault {
 			opts.RestartMode = containerd.DefaultRestartMode
 		}
+		if opts.UnconfigureMode == runtimeSpecificDefault {
+			opts.UnconfigureMode = containerd.DefaultUnconfigureMode
+		}
 	case crio.Name:
 		if opts.TopLevelConfigPath == runtimeSpecificDefault {
 			opts.TopLevelConfigPath = crio.DefaultConfig
@@ -191,6 +201,9 @@ func (opts *Options) Validate(logger logger.Interface, c *cli.Command, runtime s
 		if opts.RestartMode == runtimeSpecificDefault {
 			opts.RestartMode = crio.DefaultRestartMode
 		}
+		if opts.UnconfigureMode == runtimeSpecificDefault {
+			opts.UnconfigureMode = crio.DefaultUnconfigureMode
+		}
 	case docker.Name:
 		if opts.TopLevelConfigPath == runtimeSpecificDefault {
 			opts.TopLevelConfigPath = docker.DefaultConfig
@@ -205,6 +218,9 @@ func (opts *Options) Validate(logger logger.Interface, c *cli.Command, runtime s
 		if opts.RestartMode == runtimeSpecificDefault {
 			opts.RestartMode = docker.DefaultRestartMode
 		}
+		if opts.UnconfigureMode == runtimeSpecificDefault {
+			opts.UnconfigureMode = docker.DefaultUnconfigureMode
+		}
 	default:
 		return fmt.Errorf("undefined runtime %v", runtime)
 	}

pkg/config/engine/api.go

@@ -29,6 +29,7 @@ type Interface interface {
 	EnableCDI()
 	GetRuntimeConfig(string) (RuntimeConfig, error)
 	RemoveRuntime(string) error
+	UnsetDefaultRuntime(string)
 	Save(string) (int64, error)
 	String() string
 }

pkg/config/engine/config.go

@@ -43,6 +43,7 @@ type RuntimeConfigDestination interface {
 	AddRuntimeWithOptions(string, string, bool, interface{}) error
 	EnableCDI()
 	RemoveRuntime(string) error
+	UnsetDefaultRuntime(string)
 	Save(string) (int64, error)
 	String() string
 }
@@ -60,6 +61,11 @@ func (c *Config) RemoveRuntime(runtime string) error {
 	return c.Destination.RemoveRuntime(runtime)
 }
 
+// UnsetDefaultRuntime removes the default runtime setting from the destination config.
+func (c *Config) UnsetDefaultRuntime(runtime string) {
+	c.Destination.UnsetDefaultRuntime(runtime)
+}
+
 // EnableCDI enables CDI in the destination config.
 func (c *Config) EnableCDI() {
 	c.Destination.EnableCDI()

pkg/config/engine/containerd/config.go

@@ -123,6 +123,20 @@ func (c *Config) EnableCDI() {
 	*c.Tree = config
 }
 
+func (c *Config) UnsetDefaultRuntime(name string) {
+	if c == nil || c.Tree == nil {
+		return
+	}
+
+	config := *c.Tree
+	if runtime, ok := config.GetPath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "default_runtime_name"}).(string); ok {
+		if runtime == name {
+			config.DeletePath([]string{"plugins", c.CRIRuntimePluginName, "containerd", "default_runtime_name"})
+		}
+	}
+	*c.Tree = config
+}
+
 // RemoveRuntime removes a runtime from the docker config
 func (c *Config) RemoveRuntime(name string) error {
 	if c == nil || c.Tree == nil {

pkg/config/engine/containerd/config_drop_in.go

@@ -88,6 +88,10 @@ func (c *ConfigWithDropIn) Save(dropInPath string) (int64, error) {
 	return bytesWritten, nil
 }
 
+func (c *ConfigWithDropIn) UnsetDefaultRuntime(name string) {
+	c.Interface.UnsetDefaultRuntime(name)
+}
+
 // RemoveRuntime removes the runtime from both configs.
 func (c *ConfigWithDropIn) RemoveRuntime(name string) error {
 	if err := c.topLevelConfig.RemoveRuntime(name); err != nil {