diff --git a/pkg/device/devices.go b/pkg/device/devices.go index 43ee5c876..12be06aa4 100644 --- a/pkg/device/devices.go +++ b/pkg/device/devices.go @@ -20,6 +20,7 @@ import ( "encoding/json" "errors" "fmt" + "maps" "slices" "strconv" "strings" @@ -177,6 +178,62 @@ func init() { SupportDevices = make(map[string]string) } +func (d *DeviceUsage) DeepCopy() *DeviceUsage { + if d == nil { + return nil + } + dup := &DeviceUsage{ + ID: d.ID, + Index: d.Index, + Used: d.Used, + Count: d.Count, + Usedmem: d.Usedmem, + Totalmem: d.Totalmem, + Totalcore: d.Totalcore, + Usedcores: d.Usedcores, + Mode: d.Mode, + Numa: d.Numa, + Type: d.Type, + Health: d.Health, + } + + if d.MigTemplate != nil { + dup.MigTemplate = make([]Geometry, len(d.MigTemplate)) + for i, g := range d.MigTemplate { + dup.MigTemplate[i] = make(Geometry, len(g)) + copy(dup.MigTemplate[i], g) + } + } + + dup.MigUsage = d.MigUsage.DeepCopy() + + if d.PodInfos != nil { + dup.PodInfos = make([]*PodInfo, len(d.PodInfos)) + for i, pi := range d.PodInfos { + dup.PodInfos[i] = pi.DeepCopy() + } + } + + if d.CustomInfo != nil { + dup.CustomInfo = make(map[string]any, len(d.CustomInfo)) + maps.Copy(dup.CustomInfo, d.CustomInfo) + } + + return dup +} + +func (m MigInUse) DeepCopy() MigInUse { + var usageList MIGS + if m.UsageList != nil { + usageList = make(MIGS, len(m.UsageList)) + copy(usageList, m.UsageList) + } + return MigInUse{ + Index: m.Index, + UsageList: usageList, + } +} + func GetDevices() map[string]Devices { return DevicesMap } diff --git a/pkg/device/devices_test.go b/pkg/device/devices_test.go index 30b659bbd..f5cd99b59 100644 --- a/pkg/device/devices_test.go +++ b/pkg/device/devices_test.go @@ -1190,3 +1190,160 @@ func TestCheckUUID(t *testing.T) { }) } } + +func TestDeviceUsageDeepCopy(t *testing.T) { + tests := []struct { + name string + original *DeviceUsage + }{ + { + name: "nil input", + original: nil, + }, + { + name: "empty struct", + original: &DeviceUsage{}, + }, + { + name: "fully populated", + original: &DeviceUsage{ + ID: "GPU-0", + Index: 1, + Used: 2, + Count: 10, + Usedmem: 1024, + Totalmem: 8192, + Totalcore: 100, + Usedcores: 10, + Mode: "hami-core", + MigTemplate: []Geometry{ + { + {Name: "1g.5gb", Core: 1, Memory: 5, Count: 1}, + }, + }, + MigUsage: MigInUse{ + Index: 0, + UsageList: MIGS{ + {Name: "1g.5gb", Core: 1, Memory: 5, InUse: false}, + }, + }, + Numa: 0, + Type: "NVIDIA", + Health: true, + PodInfos: []*PodInfo{ + { + Pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "pod1", + Namespace: "default", + }, + }, + NodeID: "node1", + Devices: PodDevices{ + "NVIDIA": { + { + {UUID: "GPU-0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}, + }, + }, + }, + CtrIDs: []string{"ctr1"}, + }, + }, + CustomInfo: map[string]any{ + "key1": "value1", + "key2": 42, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + copy := tt.original.DeepCopy() + + if tt.original == nil { + if copy != nil { + t.Fatalf("expected nil, got %v", copy) + } + return + } + + // 1. Copy must be deeply equal to original. + assert.DeepEqual(t, tt.original, copy) + + // 2. Mutating the copy must not affect the original. + originalID := tt.original.ID + copy.ID = "mutated-id" + assert.Equal(t, tt.original.ID, originalID) + + if len(copy.MigTemplate) > 0 && len(copy.MigTemplate[0]) > 0 { + originalTemplateName := tt.original.MigTemplate[0][0].Name + copy.MigTemplate[0][0].Name = "mutated-template" + assert.Equal(t, tt.original.MigTemplate[0][0].Name, originalTemplateName) + } + + if len(copy.MigUsage.UsageList) > 0 { + originalUsageName := tt.original.MigUsage.UsageList[0].Name + copy.MigUsage.UsageList[0].Name = "mutated-usage" + assert.Equal(t, tt.original.MigUsage.UsageList[0].Name, originalUsageName) + } + + if len(copy.PodInfos) > 0 { + originalNodeID := tt.original.PodInfos[0].NodeID + originalCtrIDsLen := len(tt.original.PodInfos[0].CtrIDs) + copy.PodInfos[0].NodeID = "mutated-node" + copy.PodInfos[0].CtrIDs = append(copy.PodInfos[0].CtrIDs, "ctr2") + assert.Equal(t, tt.original.PodInfos[0].NodeID, originalNodeID) + assert.Equal(t, len(tt.original.PodInfos[0].CtrIDs), originalCtrIDsLen) + } + + if copy.CustomInfo != nil { + copy.CustomInfo["newkey"] = "newvalue" + _, exists := tt.original.CustomInfo["newkey"] + assert.Assert(t, !exists, "original CustomInfo should not have newkey") + } + }) + } +} + +func TestMigInUseDeepCopy(t *testing.T) { + tests := []struct { + name string + original MigInUse + }{ + { + name: "empty", + original: MigInUse{}, + }, + { + name: "with data", + original: MigInUse{ + Index: 1, + UsageList: MIGS{ + {Name: "1g.5gb", Core: 1, Memory: 5, InUse: true}, + {Name: "2g.10gb", Core: 2, Memory: 10, InUse: false}, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + copy := tt.original.DeepCopy() + + // 1. Copy must be deeply equal to original. + assert.DeepEqual(t, tt.original, copy) + + // 2. Mutating the copy must not affect the original. + originalIndex := tt.original.Index + copy.Index = 99 + assert.Equal(t, tt.original.Index, originalIndex) + + if len(copy.UsageList) > 0 { + originalName := tt.original.UsageList[0].Name + copy.UsageList[0].Name = "mutated" + assert.Equal(t, tt.original.UsageList[0].Name, originalName) + } + }) + } +} diff --git a/pkg/device/pod_test.go b/pkg/device/pod_test.go index b01a57df1..e6072531b 100644 --- a/pkg/device/pod_test.go +++ b/pkg/device/pod_test.go @@ -376,3 +376,169 @@ func TestUpdatePod(t *testing.T) { }) } } + +func TestPodInfoDeepCopy(t *testing.T) { + tests := []struct { + name string + original *PodInfo + }{ + { + name: "nil input", + original: nil, + }, + { + name: "empty struct", + original: &PodInfo{}, + }, + { + name: "fully populated", + original: &PodInfo{ + Pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "my-pod", + UID: k8stypes.UID("12345678"), + }, + }, + NodeID: "node1", + Devices: PodDevices{ + "NVIDIA": { + { + ContainerDevice{UUID: "GPU-0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}, + }, + }, + }, + CtrIDs: []string{"ctr1", "ctr2"}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + copy := tt.original.DeepCopy() + + if tt.original == nil { + if copy != nil { + t.Fatalf("expected nil, got %v", copy) + } + return + } + + // 1. Copy must be deeply equal to original. + assert.Equal(t, tt.original.NodeID, copy.NodeID) + assert.Equal(t, tt.original.CtrIDs, copy.CtrIDs) + assert.Equal(t, tt.original.Devices, copy.Devices) + if tt.original.Pod != nil { + assert.Equal(t, tt.original.Name, copy.Name) + } + + // 2. Mutating the copy must not affect the original. + if copy.Pod != nil { + originalPodName := tt.original.Name + copy.Name = "mutated-pod" + assert.Equal(t, tt.original.Name, originalPodName) + } + originalNodeID := tt.original.NodeID + copy.NodeID = "mutated-node" + assert.Equal(t, tt.original.NodeID, originalNodeID) + originalCtrIDsLen := len(tt.original.CtrIDs) + copy.CtrIDs = append(copy.CtrIDs, "ctr3") + assert.Equal(t, len(tt.original.CtrIDs), originalCtrIDsLen) + if copy.Devices == nil { + copy.Devices = make(PodDevices) + } + copy.Devices["AMD"] = PodSingleDevice{ + ContainerDevices{{UUID: "AMD-0", Type: "AMD"}}, + } + _, exists := tt.original.Devices["AMD"] + assert.False(t, exists, "original Devices should not have AMD key") + }) + } +} + +func TestPodDevicesDeepCopy(t *testing.T) { + original := PodDevices{ + "NVIDIA": { + { + ContainerDevice{UUID: "GPU-0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}, + }, + }, + } + + copy := original.DeepCopy() + + // 1. Copy must be deeply equal to original. + assert.Equal(t, original, copy) + + // 2. Mutating the copy must not affect the original. + copy["AMD"] = PodSingleDevice{ + ContainerDevices{{UUID: "AMD-0", Type: "AMD"}}, + } + copy["NVIDIA"][0][0].UUID = "mutated-gpu" + + _, exists := original["AMD"] + assert.False(t, exists, "original should not have AMD key") + assert.Equal(t, original["NVIDIA"][0][0].UUID, "GPU-0") +} + +func TestPodSingleDeviceDeepCopy(t *testing.T) { + original := PodSingleDevice{ + { + ContainerDevice{UUID: "GPU-0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}, + }, + { + ContainerDevice{UUID: "GPU-1", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}, + }, + } + + copy := original.DeepCopy() + + // 1. Copy must be deeply equal to original. + assert.Equal(t, original, copy) + + // 2. Mutating the copy must not affect the original. + copy[0][0].UUID = "mutated-gpu" + + assert.Equal(t, original[0][0].UUID, "GPU-0") +} + +func TestContainerDevicesDeepCopy(t *testing.T) { + original := ContainerDevices{ + {UUID: "GPU-0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}, + {UUID: "GPU-1", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}, + } + + copy := original.DeepCopy() + + // 1. Copy must be deeply equal to original. + assert.Equal(t, original, copy) + + // 2. Mutating the copy must not affect the original. + copy[0].UUID = "mutated-gpu" + + assert.Equal(t, original[0].UUID, "GPU-0") +} + +func TestContainerDeviceDeepCopy(t *testing.T) { + original := ContainerDevice{ + Idx: 0, + UUID: "GPU-0", + Type: "NVIDIA", + Usedmem: 100, + Usedcores: 10, + CustomInfo: map[string]any{"key1": "value1"}, + } + + copy := original.DeepCopy() + + // 1. Copy must be deeply equal to original. + assert.Equal(t, original, copy) + + // 2. Mutating the copy must not affect the original. + copy.UUID = "mutated-gpu" + copy.CustomInfo["key2"] = "value2" + + assert.Equal(t, original.UUID, "GPU-0") + _, exists := original.CustomInfo["key2"] + assert.False(t, exists, "original CustomInfo should not have key2") +} diff --git a/pkg/device/pods.go b/pkg/device/pods.go index a65b5afdc..5b0d916ad 100644 --- a/pkg/device/pods.go +++ b/pkg/device/pods.go @@ -156,6 +156,66 @@ func (m *PodManager) ListPodsInfo() []*PodInfo { return pods } +func (p *PodInfo) DeepCopy() *PodInfo { + if p == nil { + return nil + } + return &PodInfo{ + Pod: p.Pod.DeepCopy(), + NodeID: p.NodeID, + Devices: p.Devices.DeepCopy(), + CtrIDs: append([]string(nil), p.CtrIDs...), + } +} + +func (pd PodDevices) DeepCopy() PodDevices { + if pd == nil { + return nil + } + dup := make(PodDevices, len(pd)) + for k, v := range pd { + dup[k] = v.DeepCopy() + } + return dup +} + +func (psd PodSingleDevice) DeepCopy() PodSingleDevice { + if psd == nil { + return nil + } + dup := make(PodSingleDevice, len(psd)) + for i, cd := range psd { + dup[i] = cd.DeepCopy() + } + return dup +} + +func (cd ContainerDevices) DeepCopy() ContainerDevices { + if cd == nil { + return nil + } + dup := make(ContainerDevices, len(cd)) + for i, c := range cd { + dup[i] = c.DeepCopy() + } + return dup +} + +func (c ContainerDevice) DeepCopy() ContainerDevice { + dup := ContainerDevice{ + Idx: c.Idx, + UUID: c.UUID, + Type: c.Type, + Usedmem: c.Usedmem, + Usedcores: c.Usedcores, + } + if c.CustomInfo != nil { + dup.CustomInfo = make(map[string]any, len(c.CustomInfo)) + maps.Copy(dup.CustomInfo, c.CustomInfo) + } + return dup +} + func (m *PodManager) GetScheduledPods() (map[k8stypes.UID]*PodInfo, error) { m.mutex.RLock() defer m.mutex.RUnlock() diff --git a/pkg/scheduler/nodes.go b/pkg/scheduler/nodes.go index bae2b1a6a..ecac94be5 100644 --- a/pkg/scheduler/nodes.go +++ b/pkg/scheduler/nodes.go @@ -32,6 +32,16 @@ type NodeUsage struct { Devices policy.DeviceUsageList } +func (n *NodeUsage) DeepCopy() *NodeUsage { + if n == nil { + return nil + } + return &NodeUsage{ + Node: n.Node.DeepCopy(), + Devices: n.Devices.DeepCopy(), + } +} + type nodeManager struct { nodes map[string]*device.NodeInfo mutex sync.RWMutex diff --git a/pkg/scheduler/nodes_test.go b/pkg/scheduler/nodes_test.go index 81d497b99..f777ae510 100644 --- a/pkg/scheduler/nodes_test.go +++ b/pkg/scheduler/nodes_test.go @@ -22,9 +22,11 @@ import ( "github.com/Project-HAMi/HAMi/pkg/device" "github.com/Project-HAMi/HAMi/pkg/scheduler/config" + "github.com/Project-HAMi/HAMi/pkg/scheduler/policy" "gotest.tools/v3/assert" corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) func Test_addNode_ListNodes(t *testing.T) { @@ -213,6 +215,90 @@ func Test_GetNode(t *testing.T) { } } +func TestNodeUsageDeepCopy(t *testing.T) { + tests := []struct { + name string + original *NodeUsage + }{ + { + name: "nil input", + original: nil, + }, + { + name: "empty fields", + original: &NodeUsage{ + Node: nil, + Devices: policy.DeviceUsageList{}, + }, + }, + { + name: "fully populated", + original: &NodeUsage{ + Node: &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-node", + }, + }, + Devices: policy.DeviceUsageList{ + Policy: "binpack", + DeviceLists: []*policy.DeviceListsScore{ + { + Device: &device.DeviceUsage{ + ID: "GPU-0", + Index: 0, + Used: 1, + Count: 10, + Usedmem: 1024, + Totalmem: 8192, + Totalcore: 100, + Usedcores: 10, + Numa: 0, + Type: "NVIDIA", + Health: true, + }, + Score: 1.5, + }, + }, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + copy := tt.original.DeepCopy() + + if tt.original == nil { + if copy != nil { + t.Fatalf("expected nil, got %v", copy) + } + return + } + + // 1. Copy must be deeply equal to original. + assert.DeepEqual(t, tt.original, copy) + + // 2. Mutating the copy must not affect the original. + if copy.Node != nil { + originalNodeName := tt.original.Node.Name + copy.Node.Name = "mutated-node" + assert.Equal(t, tt.original.Node.Name, originalNodeName) + } + originalPolicy := tt.original.Devices.Policy + copy.Devices.Policy = "spread" + assert.Equal(t, tt.original.Devices.Policy, originalPolicy) + if len(copy.Devices.DeviceLists) > 0 { + originalScore := tt.original.Devices.DeviceLists[0].Score + originalDeviceID := tt.original.Devices.DeviceLists[0].Device.ID + copy.Devices.DeviceLists[0].Score = 99.9 + copy.Devices.DeviceLists[0].Device.ID = "mutated-gpu" + assert.Equal(t, tt.original.Devices.DeviceLists[0].Score, originalScore) + assert.Equal(t, tt.original.Devices.DeviceLists[0].Device.ID, originalDeviceID) + } + }) + } +} + func Test_rmNodeDevices(t *testing.T) { tests := []struct { name string diff --git a/pkg/scheduler/policy/gpu_policy.go b/pkg/scheduler/policy/gpu_policy.go index 272e199b4..af5d91d10 100644 --- a/pkg/scheduler/policy/gpu_policy.go +++ b/pkg/scheduler/policy/gpu_policy.go @@ -56,6 +56,30 @@ func (l DeviceUsageList) Less(i, j int) bool { return l.DeviceLists[i].Device.Numa < l.DeviceLists[j].Device.Numa } +func (l DeviceUsageList) DeepCopy() DeviceUsageList { + var deviceLists []*DeviceListsScore + if l.DeviceLists != nil { + deviceLists = make([]*DeviceListsScore, len(l.DeviceLists)) + for i, ds := range l.DeviceLists { + deviceLists[i] = ds.DeepCopy() + } + } + return DeviceUsageList{ + DeviceLists: deviceLists, + Policy: l.Policy, + } +} + +func (ds *DeviceListsScore) DeepCopy() *DeviceListsScore { + if ds == nil { + return nil + } + return &DeviceListsScore{ + Device: ds.Device.DeepCopy(), + Score: ds.Score, + } +} + func (ds *DeviceListsScore) ComputeScore(requests device.ContainerDeviceRequests) { request, core, mem := int32(0), int32(0), int32(0) // Here we are required to use the same type device diff --git a/pkg/scheduler/policy/gpu_policy_test.go b/pkg/scheduler/policy/gpu_policy_test.go index c17b4fd26..f291d1189 100644 --- a/pkg/scheduler/policy/gpu_policy_test.go +++ b/pkg/scheduler/policy/gpu_policy_test.go @@ -20,6 +20,8 @@ import ( "testing" "github.com/Project-HAMi/HAMi/pkg/device" + + "gotest.tools/v3/assert" ) func TestDeviceUsageListLen(t *testing.T) { @@ -196,6 +198,110 @@ func TestDeviceUsageList_Less(t *testing.T) { } } +func TestDeviceUsageListDeepCopy(t *testing.T) { + tests := []struct { + name string + original DeviceUsageList + }{ + { + name: "empty", + original: DeviceUsageList{}, + }, + { + name: "fully populated", + original: DeviceUsageList{ + Policy: "binpack", + DeviceLists: []*DeviceListsScore{ + { + Device: &device.DeviceUsage{ + ID: "GPU-0", + Type: "NVIDIA", + Health: true, + }, + Score: 1.0, + }, + { + Device: &device.DeviceUsage{ + ID: "GPU-1", + Type: "NVIDIA", + Health: false, + }, + Score: 2.0, + }, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + copy := tt.original.DeepCopy() + + // 1. Copy must be deeply equal to original. + assert.DeepEqual(t, tt.original, copy) + + // 2. Mutating the copy must not affect the original. + originalPolicy := tt.original.Policy + copy.Policy = "spread" + assert.Equal(t, tt.original.Policy, originalPolicy) + if len(copy.DeviceLists) > 0 { + originalScore := tt.original.DeviceLists[0].Score + originalDeviceID := tt.original.DeviceLists[0].Device.ID + copy.DeviceLists[0].Score = 99.9 + copy.DeviceLists[0].Device.ID = "mutated-gpu" + assert.Equal(t, tt.original.DeviceLists[0].Score, originalScore) + assert.Equal(t, tt.original.DeviceLists[0].Device.ID, originalDeviceID) + } + }) + } +} + +func TestDeviceListsScoreDeepCopy(t *testing.T) { + tests := []struct { + name string + original *DeviceListsScore + }{ + { + name: "nil", + original: nil, + }, + { + name: "fully populated", + original: &DeviceListsScore{ + Device: &device.DeviceUsage{ + ID: "GPU-0", + Type: "NVIDIA", + Health: true, + }, + Score: 3.5, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + copy := tt.original.DeepCopy() + + if tt.original == nil { + if copy != nil { + t.Fatalf("expected nil, got %v", copy) + } + return + } + + // 1. Copy must be deeply equal to original. + assert.DeepEqual(t, tt.original, copy) + + // 2. Mutating the copy must not affect the original. + copy.Score = 99.9 + copy.Device.ID = "mutated-gpu" + + assert.Equal(t, tt.original.Score, float32(3.5)) + assert.Equal(t, tt.original.Device.ID, "GPU-0") + }) + } +} + func TestComputeScore(t *testing.T) { tests := []struct { name string