Merge pull request #93 from mresvanis/add-mig-configuration

Add MIG config support when MIG-backed vGPU type

Author: cdesiniotis

api/spec/v1/spec.go

@@ -20,6 +20,9 @@ import (
 	"encoding/json"
 	"fmt"
 
+	migpartedv1 "github.com/NVIDIA/mig-parted/api/spec/v1"
+	migtypes "github.com/NVIDIA/mig-parted/pkg/types"
+
 	"github.com/NVIDIA/vgpu-device-manager/pkg/types"
 )
 
@@ -160,6 +163,44 @@ func (s *VGPUConfigSpec) UnmarshalJSON(b []byte) error {
 	return nil
 }
 
+func (s VGPUConfigSpecSlice) ToMigConfigSpecSlice() (migpartedv1.MigConfigSpecSlice, error) {
+	var migConfigSpecs migpartedv1.MigConfigSpecSlice
+
+	for _, vgpuSpec := range s {
+		migSpec := migpartedv1.MigConfigSpec{
+			DeviceFilter: vgpuSpec.DeviceFilter,
+			Devices:      vgpuSpec.Devices,
+			MigDevices:   make(migtypes.MigConfig),
+		}
+
+		migEnabled := false
+		for vgpuType := range vgpuSpec.VGPUDevices {
+			vgpu, err := types.ParseVGPUType(vgpuType)
+			if err != nil {
+				return nil, fmt.Errorf("failed to parse vGPU type %s: %w", vgpuType, err)
+			}
+
+			if vgpu.G > 0 {
+				migEnabled = true
+				migProfile := fmt.Sprintf("%dg.%dgb", vgpu.G, vgpu.GB)
+				for _, attr := range vgpu.Attr {
+					if attr == types.AttributeMediaExtensions {
+						migProfile += ".me"
+						break
+					}
+				}
+				migSpec.MigDevices[migProfile] = vgpuSpec.VGPUDevices[vgpuType]
+			}
+		}
+
+		migSpec.MigEnabled = migEnabled
+
+		migConfigSpecs = append(migConfigSpecs, migSpec)
+	}
+
+	return migConfigSpecs, nil
+}
+
 func containsKey(m map[string]json.RawMessage, s string) bool {
 	_, exists := m[s]
 	return exists

api/spec/v1/spec_test.go

@@ -21,6 +21,11 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"sigs.k8s.io/yaml"
+
+	migpartedv1 "github.com/NVIDIA/mig-parted/api/spec/v1"
+	migtypes "github.com/NVIDIA/mig-parted/pkg/types"
+
+	"github.com/NVIDIA/vgpu-device-manager/pkg/types"
 )
 
 func TestSpec(t *testing.T) {
@@ -230,3 +235,190 @@ func TestVGPUConfigSpec(t *testing.T) {
 	}
 
 }
+
+func TestVGPUConfigSpecSliceToMigConfigSpecSlice(t *testing.T) {
+	testCases := []struct {
+		Description           string
+		VGPUConfigSpecSlice   VGPUConfigSpecSlice
+		ExpectedMigConfigSpec migpartedv1.MigConfigSpecSlice
+		ExpectedError         string
+	}{
+		{
+			"Empty slice",
+			VGPUConfigSpecSlice{},
+			nil,
+			"",
+		},
+		{
+			"Single MIG-backed vGPU type",
+			VGPUConfigSpecSlice{
+				{
+					DeviceFilter: "MODEL",
+					Devices:      "all",
+					VGPUDevices: types.VGPUConfig{
+						"A100-1-5C": 4,
+					},
+				},
+			},
+			migpartedv1.MigConfigSpecSlice{
+				{
+					DeviceFilter: "MODEL",
+					Devices:      "all",
+					MigEnabled:   true,
+					MigDevices: migtypes.MigConfig{
+						"1g.5gb": 4,
+					},
+				},
+			},
+			"",
+		},
+		{
+			"Multiple MIG-backed vGPU types",
+			VGPUConfigSpecSlice{
+				{
+					DeviceFilter: []string{"MODEL1", "MODEL2"},
+					Devices:      []int{0, 1},
+					VGPUDevices: types.VGPUConfig{
+						"A100-1-5C":  2,
+						"A100-2-10C": 1,
+					},
+				},
+			},
+			migpartedv1.MigConfigSpecSlice{
+				{
+					DeviceFilter: []string{"MODEL1", "MODEL2"},
+					Devices:      []int{0, 1},
+					MigEnabled:   true,
+					MigDevices: migtypes.MigConfig{
+						"1g.5gb":  2,
+						"2g.10gb": 1,
+					},
+				},
+			},
+			"",
+		},
+		{
+			"MIG-backed vGPU type with media extensions",
+			VGPUConfigSpecSlice{
+				{
+					Devices: "all",
+					VGPUDevices: types.VGPUConfig{
+						"A100-1-5CME": 2,
+					},
+				},
+			},
+			migpartedv1.MigConfigSpecSlice{
+				{
+					Devices:    "all",
+					MigEnabled: true,
+					MigDevices: migtypes.MigConfig{
+						"1g.5gb.me": 2,
+					},
+				},
+			},
+			"",
+		},
+		{
+			"Non-MIG vGPU type",
+			VGPUConfigSpecSlice{
+				{
+					Devices: "all",
+					VGPUDevices: types.VGPUConfig{
+						"A100-40C": 2,
+					},
+				},
+			},
+			migpartedv1.MigConfigSpecSlice{
+				{
+					Devices:    "all",
+					MigEnabled: false,
+					MigDevices: migtypes.MigConfig{},
+				},
+			},
+			"",
+		},
+		{
+			"Mixed MIG and non-MIG vGPU types",
+			VGPUConfigSpecSlice{
+				{
+					Devices: "all",
+					VGPUDevices: types.VGPUConfig{
+						"A100-40C":  1,
+						"A100-1-5C": 2,
+					},
+				},
+			},
+			migpartedv1.MigConfigSpecSlice{
+				{
+					Devices:    "all",
+					MigEnabled: true,
+					MigDevices: migtypes.MigConfig{
+						"1g.5gb": 2,
+					},
+				},
+			},
+			"",
+		},
+		{
+			"Multiple specs with different configurations",
+			VGPUConfigSpecSlice{
+				{
+					Devices: "all",
+					VGPUDevices: types.VGPUConfig{
+						"A100-1-5C": 4,
+					},
+				},
+				{
+					DeviceFilter: "MODEL",
+					Devices:      []int{0, 1},
+					VGPUDevices: types.VGPUConfig{
+						"A100-40C": 1,
+					},
+				},
+			},
+			migpartedv1.MigConfigSpecSlice{
+				{
+					Devices:    "all",
+					MigEnabled: true,
+					MigDevices: migtypes.MigConfig{
+						"1g.5gb": 4,
+					},
+				},
+				{
+					DeviceFilter: "MODEL",
+					Devices:      []int{0, 1},
+					MigEnabled:   false,
+					MigDevices:   migtypes.MigConfig{},
+				},
+			},
+			"",
+		},
+		{
+			"Invalid vGPU type",
+			VGPUConfigSpecSlice{
+				{
+					Devices: "all",
+					VGPUDevices: types.VGPUConfig{
+						"InvalidType": 1,
+					},
+				},
+			},
+			nil,
+			"failed to parse vGPU type InvalidType:",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Description, func(t *testing.T) {
+			result, err := tc.VGPUConfigSpecSlice.ToMigConfigSpecSlice()
+			if tc.ExpectedError != "" {
+				require.NotNil(t, err, "Expected failure but got success")
+				require.Nil(t, result, "Expected nil result on failure")
+				require.ErrorContains(t, err, tc.ExpectedError)
+			} else {
+				require.Nil(t, err, "Unexpected failure: %v", err)
+				require.Equal(t, tc.ExpectedMigConfigSpec, result, "Unexpected result")
+			}
+		})
+	}
+}

cmd/nvidia-k8s-vgpu-dm/find.go

@@ -0,0 +1,71 @@
+/*
+# Copyright 2025 NVIDIA CORPORATION
+#
+# 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 main
+
+import (
+	"fmt"
+	"path/filepath"
+)
+
+type root string
+
+// getDriverLibraryPath returns path to `libnvidia-ml.so.1` in the driver root.
+// The folder for this file is also expected to be the location of other driver files.
+func (r root) getDriverLibraryPath() (string, error) {
+	librarySearchPaths := []string{
+		"/usr/lib64",
+		"/usr/lib/x86_64-linux-gnu",
+		"/usr/lib/aarch64-linux-gnu",
+		"/lib64",
+		"/lib/x86_64-linux-gnu",
+		"/lib/aarch64-linux-gnu",
+	}
+
+	libraryPath, err := r.findFile("libnvidia-ml.so.1", librarySearchPaths...)
+	if err != nil {
+		return "", err
+	}
+
+	return libraryPath, nil
+}
+
+// findFile searches the root for a specified file.
+// A number of folders can be specified to search in addition to the root itself.
+// If the file represents a symlink, this is resolved and the final path is returned.
+func (r root) findFile(name string, searchIn ...string) (string, error) {
+	for _, d := range append([]string{"/"}, searchIn...) {
+		l := filepath.Join(string(r), d, name)
+		candidate, err := resolveLink(l)
+		if err != nil {
+			continue
+		}
+		return candidate, nil
+	}
+
+	return "", fmt.Errorf("error locating %q", name)
+}
+
+// resolveLink finds the target of a symlink or the file itself in the
+// case of a regular file.
+// This is equivalent to running `readlink -f ${l}`.
+func resolveLink(l string) (string, error) {
+	resolved, err := filepath.EvalSymlinks(l)
+	if err != nil {
+		return "", fmt.Errorf("error resolving link '%s': %w", l, err)
+	}
+	return resolved, nil
+}