Add MergedDevice transform to generate all device

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

cmd/nvidia-ctk/cdi/generate/generate.go

@@ -23,11 +23,10 @@ import (
 	"strings"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/transform"
 	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi"
-	specs "github.com/container-orchestrated-devices/container-device-interface/specs-go"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v2"
 )
@@ -225,27 +224,13 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to create device CDI specs: %v", err)
 	}
-	var hasAll bool
-	for _, deviceSpec := range deviceSpecs {
-		if deviceSpec.Name == allDeviceName {
-			hasAll = true
-			break
-		}
-	}
-	if !hasAll {
-		allDevice, err := MergeDeviceSpecs(deviceSpecs, allDeviceName)
-		if err != nil {
-			return nil, fmt.Errorf("failed to create CDI specification for %q device: %v", allDeviceName, err)
-		}
-		deviceSpecs = append(deviceSpecs, allDevice)
-	}
 
 	commonEdits, err := cdilib.GetCommonEdits()
 	if err != nil {
 		return nil, fmt.Errorf("failed to create edits common for entities: %v", err)
 	}
 
-	return spec.New(
+	spec, err := spec.New(
 		spec.WithVendor(opts.vendor),
 		spec.WithClass(opts.class),
 		spec.WithDeviceSpecs(deviceSpecs),
@@ -253,32 +238,20 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		spec.WithFormat(opts.format),
 		spec.WithPermissions(0644),
 	)
-}
-
-// MergeDeviceSpecs creates a device with the specified name which combines the edits from the previous devices.
-// If a device of the specified name already exists, an error is returned.
-func MergeDeviceSpecs(deviceSpecs []specs.Device, mergedDeviceName string) (specs.Device, error) {
-	if err := cdi.ValidateDeviceName(mergedDeviceName); err != nil {
-		return specs.Device{}, fmt.Errorf("invalid device name %q: %v", mergedDeviceName, err)
-	}
-	for _, d := range deviceSpecs {
-		if d.Name == mergedDeviceName {
-			return specs.Device{}, fmt.Errorf("device %q already exists", mergedDeviceName)
-		}
-	}
-
-	mergedEdits := edits.NewContainerEdits()
-
-	for _, d := range deviceSpecs {
-		edit := cdi.ContainerEdits{
-			ContainerEdits: &d.ContainerEdits,
-		}
-		mergedEdits.Append(&edit)
-	}
-
-	merged := specs.Device{
-		Name:           mergedDeviceName,
-		ContainerEdits: *mergedEdits.ContainerEdits,
-	}
-	return merged, nil
+	if err != nil {
+		return nil, fmt.Errorf("failed to create CDI spec: %v", err)
+	}
+
+	addAllDevice, err := transform.NewMergedDevice(
+		transform.WithName(allDeviceName),
+		transform.WithSkipIfExists(true),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create merged device: %v", err)
+	}
+	if err := addAllDevice.Transform(spec.Raw()); err != nil {
+		return nil, fmt.Errorf("failed to add merged device: %v", err)
+	}
+
+	return spec, nil
 }

cmd/nvidia-ctk/cdi/generate/generate_test.go

@@ -1,117 +0,0 @@
-/**
-# 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 generate
-
-import (
-	"fmt"
-	"testing"
-
-	"github.com/container-orchestrated-devices/container-device-interface/specs-go"
-	"github.com/stretchr/testify/require"
-)
-
-func TestMergeDeviceSpecs(t *testing.T) {
-	testCases := []struct {
-		description      string
-		deviceSpecs      []specs.Device
-		mergedDeviceName string
-		expectedError    error
-		expected         specs.Device
-	}{
-		{
-			description:      "no devices",
-			mergedDeviceName: "all",
-			expected: specs.Device{
-				Name: "all",
-			},
-		},
-		{
-			description:      "one device",
-			mergedDeviceName: "all",
-			deviceSpecs: []specs.Device{
-				{
-					Name: "gpu0",
-					ContainerEdits: specs.ContainerEdits{
-						Env: []string{"GPU=0"},
-					},
-				},
-			},
-			expected: specs.Device{
-				Name: "all",
-				ContainerEdits: specs.ContainerEdits{
-					Env: []string{"GPU=0"},
-				},
-			},
-		},
-		{
-			description:      "two devices",
-			mergedDeviceName: "all",
-			deviceSpecs: []specs.Device{
-				{
-					Name: "gpu0",
-					ContainerEdits: specs.ContainerEdits{
-						Env: []string{"GPU=0"},
-					},
-				},
-				{
-					Name: "gpu1",
-					ContainerEdits: specs.ContainerEdits{
-						Env: []string{"GPU=1"},
-					},
-				},
-			},
-			expected: specs.Device{
-				Name: "all",
-				ContainerEdits: specs.ContainerEdits{
-					Env: []string{"GPU=0", "GPU=1"},
-				},
-			},
-		},
-		{
-			description:      "has merged device",
-			mergedDeviceName: "gpu0",
-			deviceSpecs: []specs.Device{
-				{
-					Name: "gpu0",
-					ContainerEdits: specs.ContainerEdits{
-						Env: []string{"GPU=0"},
-					},
-				},
-			},
-			expectedError: fmt.Errorf("device %q already exists", "gpu0"),
-		},
-		{
-			description:      "invalid merged device name",
-			mergedDeviceName: ".-not-valid",
-			expectedError:    fmt.Errorf("invalid device name %q", ".-not-valid"),
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.description, func(t *testing.T) {
-			mergedDevice, err := MergeDeviceSpecs(tc.deviceSpecs, tc.mergedDeviceName)
-
-			if tc.expectedError != nil {
-				require.Error(t, err)
-				return
-			}
-
-			require.NoError(t, err)
-			require.EqualValues(t, tc.expected, mergedDevice)
-		})
-	}
-}

pkg/nvcdi/transform/merged-device.go

@@ -0,0 +1,123 @@
+/**
+# 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 transform
+
+import (
+	"fmt"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
+	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi"
+	"github.com/container-orchestrated-devices/container-device-interface/specs-go"
+)
+
+const (
+	allDeviceName = "all"
+)
+
+type mergedDevice struct {
+	name         string
+	skipIfExists bool
+	simplifier   Transformer
+}
+
+var _ Transformer = (*mergedDevice)(nil)
+
+// MergedDeviceOption is a function that configures a merged device
+type MergedDeviceOption func(*mergedDevice)
+
+// WithName sets the name of the merged device
+func WithName(name string) MergedDeviceOption {
+	return func(m *mergedDevice) {
+		m.name = name
+	}
+}
+
+// WithSkipIfExists sets whether to skip adding the merged device if it already exists
+func WithSkipIfExists(skipIfExists bool) MergedDeviceOption {
+	return func(m *mergedDevice) {
+		m.skipIfExists = skipIfExists
+	}
+}
+
+// NewMergedDevice creates a transformer with the specified options
+func NewMergedDevice(opts ...MergedDeviceOption) (Transformer, error) {
+	m := &mergedDevice{}
+	for _, opt := range opts {
+		opt(m)
+	}
+	if m.name == "" {
+		m.name = allDeviceName
+	}
+	m.simplifier = NewSimplifier()
+
+	if err := cdi.ValidateDeviceName(m.name); err != nil {
+		return nil, fmt.Errorf("invalid device name %q: %v", m.name, err)
+	}
+
+	return m, nil
+}
+
+// Transform adds a merged device to the spec
+func (m mergedDevice) Transform(spec *specs.Spec) error {
+	if spec == nil {
+		return nil
+	}
+
+	mergedDevice, err := mergeDeviceSpecs(spec.Devices, m.name)
+	if err != nil {
+		return fmt.Errorf("failed to generate merged device %q: %v", m.name, err)
+	}
+	if mergedDevice == nil {
+		if m.skipIfExists {
+			return nil
+		}
+		return fmt.Errorf("device %q already exists", m.name)
+	}
+
+	spec.Devices = append(spec.Devices, *mergedDevice)
+
+	if err := m.simplifier.Transform(spec); err != nil {
+		return fmt.Errorf("failed to simplify spec after merging device %q: %v", m.name, err)
+	}
+
+	return nil
+}
+
+// mergeDeviceSpecs creates a device with the specified name which combines the edits from the previous devices.
+// If a device of the specified name already exists, no device is created and nil is returned.
+func mergeDeviceSpecs(deviceSpecs []specs.Device, mergedDeviceName string) (*specs.Device, error) {
+	for _, d := range deviceSpecs {
+		if d.Name == mergedDeviceName {
+			return nil, nil
+		}
+	}
+
+	mergedEdits := edits.NewContainerEdits()
+
+	for _, d := range deviceSpecs {
+		edit := cdi.ContainerEdits{
+			ContainerEdits: &d.ContainerEdits,
+		}
+		mergedEdits.Append(&edit)
+	}
+
+	merged := specs.Device{
+		Name:           mergedDeviceName,
+		ContainerEdits: *mergedEdits.ContainerEdits,
+	}
+	return &merged, nil
+}

pkg/nvcdi/transform/merged-device_test.go

@@ -0,0 +1,222 @@
+/**
+# 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 transform
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/container-orchestrated-devices/container-device-interface/specs-go"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMergeDeviceSpecs(t *testing.T) {
+	testCases := []struct {
+		description      string
+		deviceSpecs      []specs.Device
+		mergedDeviceName string
+		createError      error
+		expectedError    error
+		expected         *specs.Device
+	}{
+		{
+			description:      "no devices",
+			mergedDeviceName: "all",
+			expected: &specs.Device{
+				Name: "all",
+			},
+		},
+		{
+			description:      "one device",
+			mergedDeviceName: "all",
+			deviceSpecs: []specs.Device{
+				{
+					Name: "gpu0",
+					ContainerEdits: specs.ContainerEdits{
+						Env: []string{"GPU=0"},
+					},
+				},
+			},
+			expected: &specs.Device{
+				Name: "all",
+				ContainerEdits: specs.ContainerEdits{
+					Env: []string{"GPU=0"},
+				},
+			},
+		},
+		{
+			description:      "two devices",
+			mergedDeviceName: "all",
+			deviceSpecs: []specs.Device{
+				{
+					Name: "gpu0",
+					ContainerEdits: specs.ContainerEdits{
+						Env: []string{"GPU=0"},
+					},
+				},
+				{
+					Name: "gpu1",
+					ContainerEdits: specs.ContainerEdits{
+						Env: []string{"GPU=1"},
+					},
+				},
+			},
+			expected: &specs.Device{
+				Name: "all",
+				ContainerEdits: specs.ContainerEdits{
+					Env: []string{"GPU=0", "GPU=1"},
+				},
+			},
+		},
+		{
+			description:      "has merged device",
+			mergedDeviceName: "gpu0",
+			deviceSpecs: []specs.Device{
+				{
+					Name: "gpu0",
+					ContainerEdits: specs.ContainerEdits{
+						Env: []string{"GPU=0"},
+					},
+				},
+			},
+		},
+		{
+			description:      "invalid merged device name",
+			mergedDeviceName: ".-not-valid",
+			createError:      fmt.Errorf("invalid device name %q", ".-not-valid"),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			_, err := NewMergedDevice(
+				WithName(tc.mergedDeviceName),
+			)
+			if tc.createError != nil {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+
+			device, err := mergeDeviceSpecs(tc.deviceSpecs, tc.mergedDeviceName)
+			if tc.expectedError != nil {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+
+			require.EqualValues(t, tc.expected, device)
+		})
+	}
+}
+
+func TestMergedDevice(t *testing.T) {
+	testCases := []struct {
+		description   string
+		spec          *specs.Spec
+		expectedError error
+		expectedSpec  *specs.Spec
+	}{
+		{
+			description: "duplicate hooks are removed",
+			spec: &specs.Spec{
+				Devices: []specs.Device{
+					{
+						Name: "gpu0",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+					{
+						Name: "gpu1",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+				},
+			},
+			expectedSpec: &specs.Spec{
+				Devices: []specs.Device{
+					{
+						Name: "gpu0",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+					{
+						Name: "gpu1",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+					{
+						Name: "all",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			m, err := NewMergedDevice()
+			require.NoError(t, err)
+
+			err = m.Transform(tc.spec)
+			if tc.expectedError != nil {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+
+			require.EqualValues(t, tc.expectedSpec, tc.spec)
+		})
+	}
+}

pkg/nvcdi/transform/simplify_test.go

@@ -105,6 +105,94 @@ func TestSimplify(t *testing.T) {
 				},
 			},
 		},
+		{
+			description: "simplify removes hooks",
+			spec: &specs.Spec{
+				Devices: []specs.Device{
+					{
+						Name: "gpu0",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+					{
+						Name: "gpu1",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+					{
+						Name: "all",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+				},
+			},
+			expectedSpec: &specs.Spec{
+				Devices: []specs.Device{
+					{
+						Name: "gpu0",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+					{
+						Name: "gpu1",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+					{
+						Name: "all",
+						ContainerEdits: specs.ContainerEdits{
+							Hooks: []*specs.Hook{
+								{
+									HookName: "createContainer",
+									Path:     "/usr/bin/nvidia-ctk",
+									Args:     []string{"nvidia-ctk", "hook", "chmod", "--mode", "755", "--path", "/dev/dri"},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, tc := range testCases {