Add WSL2 discovery and spec generation

These changes add a wsl discovery mode to the nvidia-ctk cdi generate command.

If wsl mode is enabled, the driver store for the available devices is used as
the source for discovered entities.

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

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

@@ -37,9 +37,12 @@ import (
 
 const (
 	discoveryModeNVML = "nvml"
+	discoveryModeWSL  = "wsl"
 
 	formatJSON = "json"
 	formatYAML = "yaml"
+
+	allDeviceName = "all"
 )
 
 type command struct {
@@ -93,7 +96,7 @@ func (m command) build() *cli.Command {
 		},
 		&cli.StringFlag{
 			Name:        "discovery-mode",
-			Usage:       "The mode to use when discovering the available entities. One of [nvml]",
+			Usage:       "The mode to use when discovering the available entities. One of [nvml | wsl]",
 			Value:       discoveryModeNVML,
 			Destination: &cfg.discoveryMode,
 		},
@@ -130,6 +133,7 @@ func (m command) validateFlags(r *cli.Context, cfg *config) error {
 	cfg.discoveryMode = strings.ToLower(cfg.discoveryMode)
 	switch cfg.discoveryMode {
 	case discoveryModeNVML:
+	case discoveryModeWSL:
 	default:
 		return fmt.Errorf("invalid discovery mode: %v", cfg.discoveryMode)
 	}
@@ -252,10 +256,20 @@ func (m command) generateSpec(cfg *config) (*specs.Spec, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to create device CDI specs: %v", err)
 	}
-
-	allDevice := createAllDevice(deviceSpecs)
-
-	deviceSpecs = append(deviceSpecs, allDevice)
+	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 {
@@ -287,22 +301,32 @@ func (m command) generateSpec(cfg *config) (*specs.Spec, error) {
 	return &spec, nil
 }
 
-// createAllDevice creates an 'all' device which combines the edits from the previous devices
-func createAllDevice(deviceSpecs []specs.Device) specs.Device {
-	edits := edits.NewContainerEdits()
+// 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,
 		}
-		edits.Append(&edit)
+		mergedEdits.Append(&edit)
 	}
 
-	all := specs.Device{
-		Name:           "all",
-		ContainerEdits: *edits.ContainerEdits,
+	merged := specs.Device{
+		Name:           mergedDeviceName,
+		ContainerEdits: *mergedEdits.ContainerEdits,
 	}
-	return all
+	return merged, nil
 }
 
 // createParentDirsIfRequired creates the parent folders of the specified path if requried.
@@ -315,5 +339,3 @@ func createParentDirsIfRequired(filename string) error {
 	}
 	return os.MkdirAll(dir, 0755)
 }
-
-type discoveryMode string

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

@@ -0,0 +1,117 @@
+/**
+# 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)
+		})
+	}
+}