diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b1a34c1..077c46fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ * Generate a simplified CDI specification by default. This means that entities in the common edits in a spec are not included in device definitions. * Also return an error from the nvcdi.New constructor instead of panicing. * Detect XOrg libraries for injection and CDI spec generation. +* Add `nvidia-container-runtime.modes.cdi.annotation-prefixes` config option that allows the CDI annotation prefixes that are read to be overridden. * [libnvidia-container] Fix segmentation fault when RPC initialization fails. * [libnvidia-container] Build centos variants of the NVIDIA Container Library with static libtirpc v1.3.2. diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 4c378d55..72f03336 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -71,7 +71,8 @@ func TestGetConfig(t *testing.T) { MountSpecPath: "/etc/nvidia-container-runtime/host-files-for-container.d", }, CDI: cdiModeConfig{ - DefaultKind: "nvidia.com/gpu", + DefaultKind: "nvidia.com/gpu", + AnnotationPrefixes: []string{"cdi.k8s.io/"}, }, }, }, @@ -92,6 +93,7 @@ func TestGetConfig(t *testing.T) { "nvidia-container-runtime.runtimes = [\"/some/runtime\",]", "nvidia-container-runtime.mode = \"not-auto\"", "nvidia-container-runtime.modes.cdi.default-kind = \"example.vendor.com/device\"", + "nvidia-container-runtime.modes.cdi.annotation-prefixes = [\"cdi.k8s.io/\", \"example.vendor.com/\",]", "nvidia-container-runtime.modes.csv.mount-spec-path = \"/not/etc/nvidia-container-runtime/host-files-for-container.d\"", "nvidia-ctk.path = \"/foo/bar/nvidia-ctk\"", }, @@ -111,6 +113,10 @@ func TestGetConfig(t *testing.T) { }, CDI: cdiModeConfig{ DefaultKind: "example.vendor.com/device", + AnnotationPrefixes: []string{ + "cdi.k8s.io/", + "example.vendor.com/", + }, }, }, }, @@ -134,6 +140,7 @@ func TestGetConfig(t *testing.T) { "mode = \"not-auto\"", "[nvidia-container-runtime.modes.cdi]", "default-kind = \"example.vendor.com/device\"", + "annotation-prefixes = [\"cdi.k8s.io/\", \"example.vendor.com/\",]", "[nvidia-container-runtime.modes.csv]", "mount-spec-path = \"/not/etc/nvidia-container-runtime/host-files-for-container.d\"", "[nvidia-ctk]", @@ -155,6 +162,10 @@ func TestGetConfig(t *testing.T) { }, CDI: cdiModeConfig{ DefaultKind: "example.vendor.com/device", + AnnotationPrefixes: []string{ + "cdi.k8s.io/", + "example.vendor.com/", + }, }, }, }, diff --git a/internal/config/runtime.go b/internal/config/runtime.go index 0248754e..8ff8c6e8 100644 --- a/internal/config/runtime.go +++ b/internal/config/runtime.go @@ -19,6 +19,7 @@ package config import ( "fmt" + "github.com/container-orchestrated-devices/container-device-interface/pkg/cdi" "github.com/pelletier/go-toml" "github.com/sirupsen/logrus" ) @@ -52,6 +53,8 @@ type cdiModeConfig struct { SpecDirs []string `toml:"spec-dirs"` // DefaultKind sets the default kind to be used when constructing fully-qualified CDI device names DefaultKind string `toml:"default-kind"` + // AnnotationPrefixes sets the allowed prefixes for CDI annotation-based device injection + AnnotationPrefixes []string `toml:"annotation-prefixes"` } type csvModeConfig struct { @@ -98,6 +101,9 @@ func GetDefaultRuntimeConfig() *RuntimeConfig { }, CDI: cdiModeConfig{ DefaultKind: "nvidia.com/gpu", + AnnotationPrefixes: []string{ + cdi.AnnotationPrefix, + }, }, }, } diff --git a/internal/modifier/cdi.go b/internal/modifier/cdi.go index a5f42eb8..879967b9 100644 --- a/internal/modifier/cdi.go +++ b/internal/modifier/cdi.go @@ -18,6 +18,7 @@ package modifier import ( "fmt" + "strings" "github.com/NVIDIA/nvidia-container-toolkit/internal/config" "github.com/NVIDIA/nvidia-container-toolkit/internal/config/image" @@ -67,7 +68,7 @@ func getDevicesFromSpec(logger *logrus.Logger, ociSpec oci.Spec, cfg *config.Con return nil, fmt.Errorf("failed to load OCI spec: %v", err) } - _, annotationDevices, err := cdi.ParseAnnotations(rawSpec.Annotations) + annotationDevices, err := getAnnotationDevices(cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.AnnotationPrefixes, rawSpec.Annotations) if err != nil { return nil, fmt.Errorf("failed to parse container annotations: %v", err) } @@ -107,6 +108,38 @@ func getDevicesFromSpec(logger *logrus.Logger, ociSpec oci.Spec, cfg *config.Con return nil, nil } +// getAnnotationDevices returns a list of devices specified in the annotations. +// Keys starting with the specified prefixes are considered and expected to contain a comma-separated list of +// fully-qualified CDI devices names. If any device name is not fully-quality an error is returned. +// The list of returned devices is deduplicated. +func getAnnotationDevices(prefixes []string, annotations map[string]string) ([]string, error) { + devicesByKey := make(map[string][]string) + for key, value := range annotations { + for _, prefix := range prefixes { + if strings.HasPrefix(key, prefix) { + devicesByKey[key] = strings.Split(value, ",") + } + } + } + + seen := make(map[string]bool) + var annotationDevices []string + for key, devices := range devicesByKey { + for _, device := range devices { + if !cdi.IsQualifiedName(device) { + return nil, fmt.Errorf("invalid device name %q in annotation %q", device, key) + } + if seen[device] { + continue + } + annotationDevices = append(annotationDevices, device) + seen[device] = true + } + } + + return annotationDevices, nil +} + // Modify loads the CDI registry and injects the specified CDI devices into the OCI runtime specification. func (m cdiModifier) Modify(spec *specs.Spec) error { registry := cdi.GetRegistry( diff --git a/internal/modifier/cdi_test.go b/internal/modifier/cdi_test.go new file mode 100644 index 00000000..88ff697a --- /dev/null +++ b/internal/modifier/cdi_test.go @@ -0,0 +1,92 @@ +/** +# 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 modifier + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestGetAnnotationDevices(t *testing.T) { + testCases := []struct { + description string + prefixes []string + annotations map[string]string + expectedDevices []string + expectedError error + }{ + { + description: "no annotations", + }, + { + description: "no matching annotations", + prefixes: []string{"not-prefix/"}, + annotations: map[string]string{ + "prefix/foo": "example.com/device=bar", + }, + }, + { + description: "single matching annotation", + prefixes: []string{"prefix/"}, + annotations: map[string]string{ + "prefix/foo": "example.com/device=bar", + }, + expectedDevices: []string{"example.com/device=bar"}, + }, + { + description: "multiple matching annotations", + prefixes: []string{"prefix/", "another-prefix/"}, + annotations: map[string]string{ + "prefix/foo": "example.com/device=bar", + "another-prefix/bar": "example.com/device=baz", + }, + expectedDevices: []string{"example.com/device=bar", "example.com/device=baz"}, + }, + { + description: "multiple matching annotations with duplicate devices", + prefixes: []string{"prefix/", "another-prefix/"}, + annotations: map[string]string{ + "prefix/foo": "example.com/device=bar", + "another-prefix/bar": "example.com/device=bar", + }, + expectedDevices: []string{"example.com/device=bar"}, + }, + { + description: "invalid devices", + prefixes: []string{"prefix/"}, + annotations: map[string]string{ + "prefix/foo": "example.com/device", + }, + expectedError: fmt.Errorf("invalid device %q", "example.com/device"), + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + devices, err := getAnnotationDevices(tc.prefixes, tc.annotations) + if tc.expectedError != nil { + require.Error(t, err) + return + } + + require.NoError(t, err) + require.ElementsMatch(t, tc.expectedDevices, devices) + }) + } +}