Merge branch 'cherry-picks-v0.4.3' into 'release-0.4'

Cherry-picks of required patches for v0.4.3

See merge request nvidia/cloud-native/mig-parted!106

Author: klueska

deployments/systemd/nvidia-mig-manager.service

@@ -16,7 +16,7 @@
 Description=Configure MIG on NVIDIA GPUs
 DefaultDependencies=no
 After=sysinit.target local-fs.target
-Before=basic.target nvidia-persistenced.service
+Before=basic.target nvidia-persistenced.service systemd-resolved.service
 
 [Service]
 Type=oneshot

internal/nvml/mock.go

@@ -29,6 +29,7 @@ type MockA100Device struct {
 	MigMode            int
 	GpuInstances       map[*MockA100GpuInstance]struct{}
 	GpuInstanceCounter uint32
+	MemoryInfo         Memory
 }
 type MockA100GpuInstance struct {
 	Info                   GpuInstanceInfo
@@ -330,6 +331,7 @@ func NewMockA100Device(index int) Device {
 		Index:              index,
 		GpuInstances:       make(map[*MockA100GpuInstance]struct{}),
 		GpuInstanceCounter: 0,
+		MemoryInfo:         Memory{42949672960, 0, 0},
 	}
 }
 
@@ -387,6 +389,10 @@ func (d *MockA100Device) GetUUID() (string, Return) {
 	return d.UUID, MockReturn(SUCCESS)
 }
 
+func (d *MockA100Device) GetMemoryInfo() (Memory, Return) {
+	return d.MemoryInfo, MockReturn(SUCCESS)
+}
+
 func (d *MockA100Device) GetPciInfo() (PciInfo, Return) {
 	p := PciInfo{
 		PciDeviceId: 0x20B010DE,

internal/nvml/nvml.go

@@ -72,6 +72,11 @@ func (d nvmlDevice) GetUUID() (string, Return) {
 	return u, nvmlReturn(r)
 }
 
+func (d nvmlDevice) GetMemoryInfo() (Memory, Return) {
+	m, r := nvml.Device(d).GetMemoryInfo()
+	return Memory(m), nvmlReturn(r)
+}
+
 func (d nvmlDevice) GetPciInfo() (PciInfo, Return) {
 	p, r := nvml.Device(d).GetPciInfo()
 	return PciInfo(p), nvmlReturn(r)

internal/nvml/types.go

@@ -32,6 +32,7 @@ type Interface interface {
 type Device interface {
 	GetIndex() (int, Return)
 	GetUUID() (string, Return)
+	GetMemoryInfo() (Memory, Return)
 	GetPciInfo() (PciInfo, Return)
 	SetMigMode(Mode int) (Return, Return)
 	GetMigMode() (int, int, Return)
@@ -69,6 +70,7 @@ type ComputeInstanceInfo struct {
 	Placement   ComputeInstancePlacement
 }
 
+type Memory nvml.Memory
 type PciInfo nvml.PciInfo
 type GpuInstanceProfileInfo nvml.GpuInstanceProfileInfo
 type GpuInstancePlacement nvml.GpuInstancePlacement

pkg/mig/config/config.go

@@ -66,6 +66,11 @@ func (m *nvmlMigConfigManager) GetMigConfig(gpu int) (types.MigConfig, error) {
 		return nil, fmt.Errorf("error getting device handle: %v", ret)
 	}
 
+	deviceMemory, ret := device.GetMemoryInfo()
+	if ret.Value() != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting device memory: %v", ret)
+	}
+
 	err := m.nvlib.Mig.Device(device).AssertMigEnabled()
 	if err != nil {
 		return nil, fmt.Errorf("error asserting MIG enabled: %v", err)
@@ -74,7 +79,7 @@ func (m *nvmlMigConfigManager) GetMigConfig(gpu int) (types.MigConfig, error) {
 	migConfig := types.MigConfig{}
 	err = m.nvlib.Mig.Device(device).WalkGpuInstances(func(gi nvml.GpuInstance, giProfileID int, giProfileInfo nvml.GpuInstanceProfileInfo) error {
 		err := m.nvlib.Mig.GpuInstance(gi).WalkComputeInstances(func(ci nvml.ComputeInstance, ciProfileID int, ciEngProfileID int, ciProfileInfo nvml.ComputeInstanceProfileInfo) error {
-			mp := types.NewMigProfile(giProfileID, ciProfileID, ciEngProfileID, &giProfileInfo, &ciProfileInfo)
+			mp := types.NewMigProfile(giProfileID, ciProfileID, ciEngProfileID, &giProfileInfo, &ciProfileInfo, deviceMemory.Total)
 			migConfig[mp.String()]++
 			return nil
 		})
@@ -102,6 +107,11 @@ func (m *nvmlMigConfigManager) SetMigConfig(gpu int, config types.MigConfig) err
 		return fmt.Errorf("error getting device handle: %v", ret)
 	}
 
+	deviceMemory, ret := device.GetMemoryInfo()
+	if ret.Value() != nvml.SUCCESS {
+		return fmt.Errorf("error getting device memory: %v", ret)
+	}
+
 	err := m.nvlib.Mig.Device(device).AssertMigEnabled()
 	if err != nil {
 		return fmt.Errorf("error asserting MIG enabled: %v", err)
@@ -169,7 +179,7 @@ func (m *nvmlMigConfigManager) SetMigConfig(gpu int, config types.MigConfig) err
 					return fmt.Errorf("error creating Compute instance for '%v': %v", mp, ret)
 				}
 
-				valid := types.NewMigProfile(mp.GIProfileID, mp.CIProfileID, mp.CIEngProfileID, &giProfileInfo, &ciProfileInfo)
+				valid := types.NewMigProfile(mp.GIProfileID, mp.CIProfileID, mp.CIEngProfileID, &giProfileInfo, &ciProfileInfo, deviceMemory.Total)
 				if !mp.Equals(valid) {
 					if reuseGI {
 						reuseGI = false

pkg/types/mig_profile.go

@@ -18,6 +18,7 @@ package types
 
 import (
 	"fmt"
+	"math"
 	"strconv"
 	"strings"
 
@@ -41,11 +42,11 @@ type MigProfile struct {
 }
 
 // NewMigProfile constructs a new MigProfile struct using info from the giProfiles and ciProfiles used to create it.
-func NewMigProfile(giProfileID, ciProfileID, ciEngProfileID int, giProfileInfo *nvml.GpuInstanceProfileInfo, ciProfileInfo *nvml.ComputeInstanceProfileInfo) *MigProfile {
+func NewMigProfile(giProfileID, ciProfileID, ciEngProfileID int, giProfileInfo *nvml.GpuInstanceProfileInfo, ciProfileInfo *nvml.ComputeInstanceProfileInfo, totalDeviceMemory uint64) *MigProfile {
 	return &MigProfile{
 		C:              int(ciProfileInfo.SliceCount),
 		G:              int(giProfileInfo.SliceCount),
-		GB:             int((giProfileInfo.MemorySizeMB + 1024 - 1) / 1024),
+		GB:             int(getMigMemorySizeInGB(totalDeviceMemory, giProfileInfo.MemorySizeMB)),
 		GIProfileID:    giProfileID,
 		CIProfileID:    ciProfileID,
 		CIEngProfileID: ciEngProfileID,
@@ -256,3 +257,13 @@ func parseMigProfileAttributes(s string) ([]string, error) {
 	}
 	return attr, nil
 }
+
+func getMigMemorySizeInGB(totalDeviceMemory, migMemorySizeMB uint64) uint64 {
+	const fracDenominator = 8
+	const oneMB = 1024 * 1024
+	const oneGB = 1024 * 1024 * 1024
+	fractionalGpuMem := (float64(migMemorySizeMB) * oneMB) / float64(totalDeviceMemory)
+	fractionalGpuMem = math.Ceil(fractionalGpuMem*fracDenominator) / fracDenominator
+	totalMemGB := float64((totalDeviceMemory + oneGB - 1) / oneGB)
+	return uint64(math.Round(fractionalGpuMem * totalMemGB))
+}

pkg/types/types_test.go

@@ -17,6 +17,7 @@
 package types
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -266,3 +267,48 @@ func TestParseMigProfile(t *testing.T) {
 		})
 	}
 }
+
+func TestGetMigMemorySizeInGB(t *testing.T) {
+	type testCase struct {
+		totalDeviceMemory    uint64
+		migMemorySizeMB      uint64
+		expectedMemorySizeGB uint64
+	}
+
+	const maxMemorySlices = 8
+	const oneMB = uint64(1024 * 1024)
+	const oneGB = uint64(1024 * 1024 * 1024)
+
+	totalDeviceMemory := []uint64{
+		24 * oneGB,
+		40 * oneGB,
+		80 * oneGB,
+	}
+
+	testCases := []testCase{}
+	for _, tdm := range totalDeviceMemory {
+		sliceSize := tdm / maxMemorySlices
+
+		const stepSize = oneGB / 32
+		for i := stepSize; i <= tdm; i += stepSize {
+			tc := testCase{
+				totalDeviceMemory: tdm,
+				migMemorySizeMB:   i / oneMB,
+			}
+			for j := uint64(sliceSize); j <= tdm; j += sliceSize {
+				if i <= j {
+					tc.expectedMemorySizeGB = j / oneGB
+					break
+				}
+			}
+			testCases = append(testCases, tc)
+		}
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("%v", tc.migMemorySizeMB), func(t *testing.T) {
+			memorySizeGB := getMigMemorySizeInGB(tc.totalDeviceMemory, tc.migMemorySizeMB)
+			require.Equal(t, int(tc.expectedMemorySizeGB), int(memorySizeGB))
+		})
+	}
+}