Configure containerd config based on specified annotation prefixes

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

internal/config/engine/containerd/config_v1.go

@@ -50,12 +50,15 @@ func (c *ConfigV1) AddRuntime(name string, path string, setAsDefault bool) error
 		config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "runtime_engine"}, "")
 		config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "privileged_without_host_devices"}, false)
 	}
-	cdiAnnotations := []interface{}{"cdi.k8s.io/*"}
+
-	containerAnnotations, ok := config.GetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "container_annotations"}).([]interface{})
+	if len(c.ContainerAnnotations) > 0 {
-	if ok && containerAnnotations != nil {
+		annotations, err := (*Config)(c).getRuntimeAnnotations([]string{"plugins", "cri", "containerd", "runtimes", name, "container_annotations"})
-		cdiAnnotations = append(containerAnnotations, cdiAnnotations...)
+		if err != nil {
+			return err
+		}
+		annotations = append(c.ContainerAnnotations, annotations...)
+		config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "container_annotations"}, annotations)
 	}
-	config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "container_annotations"}, cdiAnnotations)
 
 	config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "options", "BinaryName"}, path)
 	config.SetPath([]string{"plugins", "cri", "containerd", "runtimes", name, "options", "Runtime"}, path)

internal/config/engine/containerd/config_v2.go

@@ -45,12 +45,14 @@ func (c *Config) AddRuntime(name string, path string, setAsDefault bool) error {
 		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "privileged_without_host_devices"}, false)
 	}
 
-	cdiAnnotations := []interface{}{"cdi.k8s.io/*"}
+	if len(c.ContainerAnnotations) > 0 {
-	containerAnnotations, ok := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "container_annotations"}).([]interface{})
+		annotations, err := c.getRuntimeAnnotations([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "container_annotations"})
-	if ok && containerAnnotations != nil {
+		if err != nil {
-		cdiAnnotations = append(containerAnnotations, cdiAnnotations...)
+			return err
+		}
+		annotations = append(c.ContainerAnnotations, annotations...)
+		config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "container_annotations"}, annotations)
 	}
-	config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "container_annotations"}, cdiAnnotations)
 
 	config.SetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "runtimes", name, "options", "BinaryName"}, path)
 
@@ -62,6 +64,32 @@ func (c *Config) AddRuntime(name string, path string, setAsDefault bool) error {
 	return nil
 }
 
+func (c *Config) getRuntimeAnnotations(path []string) ([]string, error) {
+	if c == nil || c.Tree == nil {
+		return nil, nil
+	}
+
+	config := *c.Tree
+	if !config.HasPath(path) {
+		return nil, nil
+	}
+	annotationsI, ok := config.GetPath(path).([]interface{})
+	if !ok {
+		return nil, fmt.Errorf("invalid annotations: %v", annotationsI)
+	}
+
+	var annotations []string
+	for _, annotation := range annotationsI {
+		a, ok := annotation.(string)
+		if !ok {
+			return nil, fmt.Errorf("invalid annotation: %v", annotation)
+		}
+		annotations = append(annotations, a)
+	}
+
+	return annotations, 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 {

internal/config/engine/containerd/containerd.go

@@ -26,6 +26,7 @@ type Config struct {
 	*toml.Tree
 	RuntimeType           string
 	UseDefaultRuntimeName bool
+	ContainerAnnotations  []string
 }
 
 // New creates a containerd config with the specified options

internal/config/engine/containerd/option.go

@@ -30,9 +30,10 @@ const (
 )
 
 type builder struct {
-	path            string
+	path                 string
-	runtimeType     string
+	runtimeType          string
-	useLegacyConfig bool
+	useLegacyConfig      bool
+	containerAnnotations []string
 }
 
 // Option defines a function that can be used to configure the config builder
@@ -59,6 +60,13 @@ func WithUseLegacyConfig(useLegacyConfig bool) Option {
 	}
 }
 
+// WithContainerAnnotations sets the container annotations for the config builder
+func WithContainerAnnotations(containerAnnotations ...string) Option {
+	return func(b *builder) {
+		b.containerAnnotations = containerAnnotations
+	}
+}
+
 func (b *builder) build() (engine.Interface, error) {
 	if b.path == "" {
 		return nil, fmt.Errorf("config path is empty")
@@ -74,6 +82,7 @@ func (b *builder) build() (engine.Interface, error) {
 	}
 	config.RuntimeType = b.runtimeType
 	config.UseDefaultRuntimeName = !b.useLegacyConfig
+	config.ContainerAnnotations = b.containerAnnotations
 
 	version, err := config.parseVersion(b.useLegacyConfig)
 	if err != nil {

tools/container/containerd/config_v1_test.go

@@ -332,6 +332,7 @@ func TestUpdateV1Config(t *testing.T) {
 				Tree:                  config,
 				UseDefaultRuntimeName: true,
 				RuntimeType:           runtimeType,
+				ContainerAnnotations:  []string{"cdi.k8s.io/*"},
 			}
 
 			err = UpdateConfig(v1, o)
@@ -585,6 +586,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
 				Tree:                  config,
 				UseDefaultRuntimeName: true,
 				RuntimeType:           runtimeType,
+				ContainerAnnotations:  []string{"cdi.k8s.io/*"},
 			}
 
 			err = UpdateConfig(v1, o)

tools/container/containerd/config_v2_test.go

@@ -279,8 +279,9 @@ func TestUpdateV2Config(t *testing.T) {
 			require.NoError(t, err)
 
 			v2 := &containerd.Config{
-				Tree:        config,
+				Tree:                 config,
-				RuntimeType: runtimeType,
+				RuntimeType:          runtimeType,
+				ContainerAnnotations: []string{"cdi.k8s.io/*"},
 			}
 
 			err = UpdateConfig(v2, o)
@@ -520,8 +521,9 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
 			require.NoError(t, err)
 
 			v2 := &containerd.Config{
-				Tree:        config,
+				Tree:                 config,
-				RuntimeType: runtimeType,
+				RuntimeType:          runtimeType,
+				ContainerAnnotations: []string{"cdi.k8s.io/*"},
 			}
 
 			err = UpdateConfig(v2, o)

tools/container/containerd/containerd.go

@@ -72,6 +72,8 @@ type options struct {
 	hostRootMount   string
 	runtimeDir      string
 	useLegacyConfig bool
+
+	ContainerRuntimeModesCDIAnnotationPrefixes cli.StringSlice
 }
 
 func main() {
@@ -173,6 +175,11 @@ func main() {
 			Destination: &options.useLegacyConfig,
 			EnvVars:     []string{"CONTAINERD_USE_LEGACY_CONFIG"},
 		},
+		&cli.StringSliceFlag{
+			Name:        "nvidia-container-runtime-modes.cdi.annotation-prefixes",
+			Destination: &options.ContainerRuntimeModesCDIAnnotationPrefixes,
+			EnvVars:     []string{"NVIDIA_CONTAINER_RUNTIME_MODES_CDI_ANNOTATION_PREFIXES"},
+		},
 	}
 
 	// Update the subcommand flags with the common subcommand flags
@@ -199,6 +206,7 @@ func Setup(c *cli.Context, o *options) error {
 		containerd.WithPath(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)
@@ -241,6 +249,7 @@ func Cleanup(c *cli.Context, o *options) error {
 		containerd.WithPath(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)
@@ -434,3 +443,13 @@ func RestartContainerdSystemd(hostRootMount string) error {
 
 	return nil
 }
+
+// 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
+}