From 61595aa0fa37ba6e82d2ff737ebbd4c85dfe01a3 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Tue, 31 Oct 2023 15:24:27 +0100
Subject: [PATCH] Add cdi.enabled option to runtime configure

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 cmd/nvidia-ctk/runtime/configure/configure.go | 23 +++++++++++++++++++
 pkg/config/engine/api.go                      |  1 +
 pkg/config/engine/containerd/config_v1.go     |  8 +++++++
 pkg/config/engine/containerd/config_v2.go     |  8 +++++++
 pkg/config/engine/crio/crio.go                |  5 ++++
 pkg/config/engine/docker/docker.go            |  5 ++++
 6 files changed, 50 insertions(+)

diff --git a/cmd/nvidia-ctk/runtime/configure/configure.go b/cmd/nvidia-ctk/runtime/configure/configure.go
index 90af0cb6..11f37387 100644
--- a/cmd/nvidia-ctk/runtime/configure/configure.go
+++ b/cmd/nvidia-ctk/runtime/configure/configure.go
@@ -71,6 +71,11 @@ type config struct {
 		hookPath     string
 		setAsDefault bool
 	}
+
+	// cdi-specific options
+	cdi struct {
+		enabled bool
+	}
 }
 
 func (m command) build() *cli.Command {
@@ -141,6 +146,11 @@ func (m command) build() *cli.Command {
 			Usage:       "set the NVIDIA runtime as the default runtime",
 			Destination: &config.nvidiaRuntime.setAsDefault,
 		},
+		&cli.BoolFlag{
+			Name:        "cdi.enabled",
+			Usage:       "Enable CDI in the configured runtime",
+			Destination: &config.cdi.enabled,
+		},
 	}
 
 	return &configure
@@ -175,6 +185,13 @@ func (m command) validateFlags(c *cli.Context, config *config) error {
 		}
 	}
 
+	if config.runtime != "containerd" {
+		if config.cdi.enabled {
+			m.logger.Warningf("Ignoring cdi.enabled flag for %v", config.runtime)
+		}
+		config.cdi.enabled = false
+	}
+
 	return nil
 }
 
@@ -227,6 +244,12 @@ func (m command) configureConfigFile(c *cli.Context, config *config) error {
 		return fmt.Errorf("unable to update config: %v", err)
 	}
 
+	if config.cdi.enabled {
+		if err := cfg.Set("enable_cdi", true); err != nil {
+			return fmt.Errorf("failed enable CDI in containerd: %w", err)
+		}
+	}
+
 	outputPath := config.getOuputConfigPath()
 	n, err := cfg.Save(outputPath)
 	if err != nil {
diff --git a/pkg/config/engine/api.go b/pkg/config/engine/api.go
index 5cd5fb52..3e6d1035 100644
--- a/pkg/config/engine/api.go
+++ b/pkg/config/engine/api.go
@@ -20,6 +20,7 @@ package engine
 type Interface interface {
 	DefaultRuntime() string
 	AddRuntime(string, string, bool) error
+	Set(string, interface{}) error
 	RemoveRuntime(string) error
 	Save(string) (int64, error)
 }
diff --git a/pkg/config/engine/containerd/config_v1.go b/pkg/config/engine/containerd/config_v1.go
index c2343941..d058d21c 100644
--- a/pkg/config/engine/containerd/config_v1.go
+++ b/pkg/config/engine/containerd/config_v1.go
@@ -133,6 +133,14 @@ func (c *ConfigV1) RemoveRuntime(name string) error {
 	return nil
 }
 
+// SetOption sets the specified containerd option.
+func (c *ConfigV1) Set(key string, value interface{}) error {
+	config := *c.Tree
+	config.SetPath([]string{"plugins", "cri", "containerd", key}, value)
+	*c.Tree = config
+	return nil
+}
+
 // Save wrotes the config to a file
 func (c ConfigV1) Save(path string) (int64, error) {
 	return (Config)(c).Save(path)
diff --git a/pkg/config/engine/containerd/config_v2.go b/pkg/config/engine/containerd/config_v2.go
index 35205446..6d3b2d80 100644
--- a/pkg/config/engine/containerd/config_v2.go
+++ b/pkg/config/engine/containerd/config_v2.go
@@ -89,6 +89,14 @@ func (c *Config) getRuntimeAnnotations(path []string) ([]string, error) {
 	return annotations, nil
 }
 
+// Set sets the specified containerd option.
+func (c *Config) Set(key string, value interface{}) error {
+	config := *c.Tree
+	config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", key}, value)
+	*c.Tree = config
+	return nil
+}
+
 // DefaultRuntime returns the default runtime for the cri-o config
 func (c Config) DefaultRuntime() string {
 	if runtime, ok := c.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"}).(string); ok {
diff --git a/pkg/config/engine/crio/crio.go b/pkg/config/engine/crio/crio.go
index dcf8572a..6832d3a1 100644
--- a/pkg/config/engine/crio/crio.go
+++ b/pkg/config/engine/crio/crio.go
@@ -98,6 +98,11 @@ func (c *Config) RemoveRuntime(name string) error {
 	return nil
 }
 
+// Set is not supported for cri-o configs.
+func (c *Config) Set(key string, value interface{}) error {
+	return fmt.Errorf("Set is not supported for crio configs")
+}
+
 // Save writes the config to the specified path
 func (c Config) Save(path string) (int64, error) {
 	config := (toml.Tree)(c)
diff --git a/pkg/config/engine/docker/docker.go b/pkg/config/engine/docker/docker.go
index 92594cc5..965dc2e4 100644
--- a/pkg/config/engine/docker/docker.go
+++ b/pkg/config/engine/docker/docker.go
@@ -114,6 +114,11 @@ func (c *Config) RemoveRuntime(name string) error {
 	return nil
 }
 
+// Set is not supported for docker configs.
+func (c *Config) Set(key string, value interface{}) error {
+	return fmt.Errorf("Set is not supported for crio configs")
+}
+
 // Save writes the config to the specified path
 func (c Config) Save(path string) (int64, error) {
 	output, err := json.MarshalIndent(c, "", "    ")