diff --git a/pkg/device/common/common.go b/pkg/device/common/common.go index 6405dc8dc7..174be2bf1c 100644 --- a/pkg/device/common/common.go +++ b/pkg/device/common/common.go @@ -18,6 +18,7 @@ package common import ( "fmt" + "sort" "strings" ) @@ -45,6 +46,7 @@ func GenReason(reasons map[string]int, cards int) string { for r, cnt := range reasons { reason = append(reason, fmt.Sprintf("%d/%d %s", cnt, cards, r)) } + sort.Strings(reason) return strings.Join(reason, ", ") } diff --git a/pkg/device/initContainer.go b/pkg/device/initContainer.go new file mode 100644 index 0000000000..ed7342954d --- /dev/null +++ b/pkg/device/initContainer.go @@ -0,0 +1,138 @@ +/* +Copyright 2026 The HAMi Authors. + +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 device + +import ( + corev1 "k8s.io/api/core/v1" +) + +type deviceKey struct { + devType string + uuid string +} +type usage struct { + mem int32 + cores int32 +} + +// CollapseInitContainerUsage returns the effective device usage for a pod. +func CollapseInitContainerUsage(pod *corev1.Pod, raw PodDevices) PodDevices { + if raw == nil { + return nil + } + numInit := len(pod.Spec.InitContainers) + initPeak := make(map[deviceKey]usage) + appSum := make(map[deviceKey]usage) + + for devType, podSingle := range raw { + for cidx, ctrDevs := range podSingle { + for _, dev := range ctrDevs { + key := deviceKey{devType: devType, uuid: dev.UUID} + if cidx < numInit { + cur := initPeak[key] + if dev.Usedmem > cur.mem { + cur.mem = dev.Usedmem + } + if dev.Usedcores > cur.cores { + cur.cores = dev.Usedcores + } + initPeak[key] = cur + } else { + cur := appSum[key] + cur.mem += dev.Usedmem + cur.cores += dev.Usedcores + appSum[key] = cur + } + } + } + } + + collapsed := make(PodDevices) + for devType := range raw { + uuidSet := make(map[string]struct{}) + for k := range initPeak { + if k.devType == devType { + uuidSet[k.uuid] = struct{}{} + } + } + for k := range appSum { + if k.devType == devType { + uuidSet[k.uuid] = struct{}{} + } + } + + collapsedSingle := make(PodSingleDevice, 1) + var containerDevs ContainerDevices + + for uuid := range uuidSet { + initU := initPeak[deviceKey{devType, uuid}] + appU := appSum[deviceKey{devType, uuid}] + + effMem := max(initU.mem, appU.mem) + effCores := max(initU.cores, appU.cores) + + containerDevs = append(containerDevs, ContainerDevice{ + UUID: uuid, + Type: devType, + Usedmem: effMem, + Usedcores: effCores, + }) + } + + collapsedSingle[0] = containerDevs + collapsed[devType] = collapsedSingle + } + return collapsed +} + +// AppContainersOnlyDeviceUsage returns the device usage considering only app containers. +// Used when init containers have finished and we want to shrink to app-only footprint. +func AppContainersOnlyDeviceUsage(pod *corev1.Pod, raw PodDevices) PodDevices { + if raw == nil { + return nil + } + numInit := len(pod.Spec.InitContainers) + + collapsed := make(PodDevices) + for devType, podSingle := range raw { + sums := make(map[string]usage) + for cidx, ctrDevs := range podSingle { + if cidx < numInit { + continue // skip init containers + } + for _, dev := range ctrDevs { + s := sums[dev.UUID] + s.mem += dev.Usedmem + s.cores += dev.Usedcores + sums[dev.UUID] = s + } + } + collapsedSingle := make(PodSingleDevice, 1) + var containerDevs ContainerDevices + for uuid, s := range sums { + containerDevs = append(containerDevs, ContainerDevice{ + UUID: uuid, + Type: devType, + Usedmem: s.mem, + Usedcores: s.cores, + }) + } + collapsedSingle[0] = containerDevs + collapsed[devType] = collapsedSingle + } + return collapsed +} diff --git a/pkg/device/initContainer_test.go b/pkg/device/initContainer_test.go new file mode 100644 index 0000000000..6bf021777a --- /dev/null +++ b/pkg/device/initContainer_test.go @@ -0,0 +1,326 @@ +/* +Copyright 2026 The HAMi Authors. + +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 device + +import ( + "testing" + + "gotest.tools/v3/assert" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) + +func makePod(name string, numInit, numApp int) *corev1.Pod { + initContainers := make([]corev1.Container, numInit) + for i := range initContainers { + initContainers[i] = corev1.Container{Name: "init"} + } + appContainers := make([]corev1.Container, numApp) + for i := range appContainers { + appContainers[i] = corev1.Container{Name: "app"} + } + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: "default"}, + Spec: corev1.PodSpec{ + InitContainers: initContainers, + Containers: appContainers, + }, + } +} + +func TestCollapseInitContainerUsage_NilInput(t *testing.T) { + result := CollapseInitContainerUsage(nil, nil) + assert.Assert(t, result == nil) + + pod := makePod("test", 1, 1) + result = CollapseInitContainerUsage(pod, nil) + assert.Assert(t, result == nil) +} + +func TestCollapseInitContainerUsage_NoInitContainers(t *testing.T) { + pod := makePod("test", 0, 2) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}}, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + { + ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 300, Usedcores: 30}, + }, + }, + } + result := CollapseInitContainerUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestCollapseInitContainerUsage_OnlyInitContainers(t *testing.T) { + pod := makePod("test", 2, 0) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 30}}, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + { + ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 30}, + }, + }, + } + result := CollapseInitContainerUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestCollapseInitContainerUsage_MixedInitAndApp(t *testing.T) { + pod := makePod("test", 1, 1) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 500, Usedcores: 50}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}}, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + { + ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 500, Usedcores: 50}, + }, + }, + } + result := CollapseInitContainerUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestCollapseInitContainerUsage_InitLargerThanApp(t *testing.T) { + pod := makePod("test", 1, 1) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 1000, Usedcores: 80}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}}, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 1000, Usedcores: 80}}, + }, + } + result := CollapseInitContainerUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestCollapseInitContainerUsage_AppLargerThanInit(t *testing.T) { + pod := makePod("test", 1, 2) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}}, // init + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}}, // app0 + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 300, Usedcores: 30}}, // app1 + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 500, Usedcores: 50}}, + }, + } + result := CollapseInitContainerUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestCollapseInitContainerUsage_MultipleInitContainersPeak(t *testing.T) { + pod := makePod("test", 3, 1) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 300, Usedcores: 30}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 50, Usedcores: 5}}, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 300, Usedcores: 30}}, + }, + } + result := CollapseInitContainerUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestCollapseInitContainerUsage_MultipleDeviceTypes(t *testing.T) { + pod := makePod("test", 1, 1) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 500, Usedcores: 50}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}}, + }, + "kunlun": PodSingleDevice{ + {ContainerDevice{UUID: "xpu0", Type: "kunlun", Usedmem: 200, Usedcores: 20}}, + {ContainerDevice{UUID: "xpu0", Type: "kunlun", Usedmem: 300, Usedcores: 30}}, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 500, Usedcores: 50}}, + }, + "kunlun": PodSingleDevice{ + {ContainerDevice{UUID: "xpu0", Type: "kunlun", Usedmem: 300, Usedcores: 30}}, + }, + } + result := CollapseInitContainerUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestCollapseInitContainerUsage_MultipleDevicesSameContainer(t *testing.T) { + pod := makePod("test", 1, 1) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + { + ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}, + ContainerDevice{UUID: "gpu1", Type: "NVIDIA", Usedmem: 300, Usedcores: 30}, + }, + { + ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}, + ContainerDevice{UUID: "gpu1", Type: "NVIDIA", Usedmem: 150, Usedcores: 15}, + }, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + { + ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}, + ContainerDevice{UUID: "gpu1", Type: "NVIDIA", Usedmem: 300, Usedcores: 30}, + }, + }, + } + result := CollapseInitContainerUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestCollapseInitContainerUsage_EmptyContainerDeviceList(t *testing.T) { + pod := makePod("test", 2, 2) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {}, // init0 empty + {}, // init1 empty + {}, // app0 empty + {}, // app1 empty + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + nil, // no devices => nil slice + }, + } + result := CollapseInitContainerUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestAppContainersOnlyDeviceUsage_NilInput(t *testing.T) { + result := AppContainersOnlyDeviceUsage(nil, nil) + assert.Assert(t, result == nil) + + pod := makePod("test", 1, 1) + result = AppContainersOnlyDeviceUsage(pod, nil) + assert.Assert(t, result == nil) +} + +func TestAppContainersOnlyDeviceUsage_OnlyAppContainers(t *testing.T) { + pod := makePod("test", 0, 2) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}}, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 300, Usedcores: 30}}, + }, + } + result := AppContainersOnlyDeviceUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestAppContainersOnlyDeviceUsage_IgnoresInitContainers(t *testing.T) { + pod := makePod("test", 1, 2) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 999, Usedcores: 99}}, // init + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}}, // app0 + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}}, // app1 + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 300, Usedcores: 30}}, + }, + } + result := AppContainersOnlyDeviceUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestAppContainersOnlyDeviceUsage_MultipleDeviceTypes(t *testing.T) { + pod := makePod("test", 1, 1) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 500, Usedcores: 50}}, + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}}, + }, + "kunlun": PodSingleDevice{ + {ContainerDevice{UUID: "xpu0", Type: "kunlun", Usedmem: 300, Usedcores: 30}}, + {ContainerDevice{UUID: "xpu0", Type: "kunlun", Usedmem: 400, Usedcores: 40}}, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + {ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}}, + }, + "kunlun": PodSingleDevice{ + {ContainerDevice{UUID: "xpu0", Type: "kunlun", Usedmem: 400, Usedcores: 40}}, + }, + } + result := AppContainersOnlyDeviceUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} + +func TestAppContainersOnlyDeviceUsage_MultipleDevicesPerContainer(t *testing.T) { + pod := makePod("test", 1, 1) + raw := PodDevices{ + "NVIDIA": PodSingleDevice{ + { // init (ignored) + ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 1000, Usedcores: 100}, + ContainerDevice{UUID: "gpu1", Type: "NVIDIA", Usedmem: 2000, Usedcores: 200}, + }, + { // app + ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}, + ContainerDevice{UUID: "gpu1", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}, + }, + }, + } + expected := PodDevices{ + "NVIDIA": PodSingleDevice{ + { + ContainerDevice{UUID: "gpu0", Type: "NVIDIA", Usedmem: 100, Usedcores: 10}, + ContainerDevice{UUID: "gpu1", Type: "NVIDIA", Usedmem: 200, Usedcores: 20}, + }, + }, + } + result := AppContainersOnlyDeviceUsage(pod, raw) + assert.DeepEqual(t, expected, result) +} diff --git a/pkg/device/pods.go b/pkg/device/pods.go index a549bc4c5d..60e8ada7dc 100644 --- a/pkg/device/pods.go +++ b/pkg/device/pods.go @@ -28,8 +28,9 @@ import ( type PodInfo struct { *corev1.Pod - NodeID string - Devices PodDevices + NodeID string + Devices PodDevices + InitContainerResourceReleased bool } // PodUseDeviceStat counts pod use device info. @@ -51,6 +52,8 @@ func NewPodManager() *PodManager { return pm } +// AddPod stores the effective (collapsed) device usage for the pod. +// The devices parameter must already be collapsed (caller's responsibility). func (m *PodManager) AddPod(pod *corev1.Pod, nodeID string, devices PodDevices) bool { m.mutex.Lock() defer m.mutex.Unlock() @@ -58,9 +61,10 @@ func (m *PodManager) AddPod(pod *corev1.Pod, nodeID string, devices PodDevices) _, exists := m.pods[pod.UID] if !exists { pi := &PodInfo{ - Pod: pod, - NodeID: nodeID, - Devices: devices, + Pod: pod, + NodeID: nodeID, + Devices: devices, + InitContainerResourceReleased: false, } m.pods[pod.UID] = pi klog.InfoS("Pod added", @@ -69,29 +73,44 @@ func (m *PodManager) AddPod(pod *corev1.Pod, nodeID string, devices PodDevices) "devices", devices, ) } else { - m.pods[pod.UID].Devices = devices - klog.V(5).InfoS("Pod devices updated", + // If the pod already exists (e.g., update), we only update the pod object + // and keep the existing devices, unless we are in a shrink path which is + // handled by ShrinkUsage. + m.pods[pod.UID].Pod = pod + klog.V(5).InfoS("Pod already exists; only pod object updated", "pod", klog.KRef(pod.Namespace, pod.Name), - "devices", devices, ) } return !exists } +// UpdatePod updates only the pod object (used for termination state). func (m *PodManager) UpdatePod(pod *corev1.Pod) { m.mutex.Lock() defer m.mutex.Unlock() if pi, exists := m.pods[pod.UID]; exists { pi.Pod = pod - klog.V(5).InfoS("Pod object updated in cache (terminating state)", + klog.V(5).InfoS("Pod object updated in cache", "pod", klog.KRef(pod.Namespace, pod.Name), - "deletionTimestamp", pod.DeletionTimestamp, ) } } +// DeepCopy must include the new field. +func (p *PodInfo) DeepCopy() *PodInfo { + if p == nil { + return nil + } + return &PodInfo{ + Pod: p.Pod.DeepCopy(), + NodeID: p.NodeID, + Devices: p.Devices.DeepCopy(), + InitContainerResourceReleased: p.InitContainerResourceReleased, + } +} + func (m *PodManager) DelPod(pod *corev1.Pod) { m.mutex.Lock() defer m.mutex.Unlock() @@ -129,6 +148,22 @@ func (m *PodManager) TakeAndDeletePod(pod *corev1.Pod) (*PodInfo, bool) { } return pi, ok } +func (m *PodManager) ShrinkUsage(pod *corev1.Pod, newDevices PodDevices) (oldDevices PodDevices, ok bool) { + m.mutex.Lock() + defer m.mutex.Unlock() + + pi, exists := m.pods[pod.UID] + if !exists { + return nil, false + } + oldDevices = pi.Devices + pi.Devices = newDevices + pi.InitContainerResourceReleased = true + klog.V(4).InfoS("Init container resources released", + "pod", klog.KRef(pod.Namespace, pod.Name), + ) + return oldDevices, true +} func (m *PodManager) ListPodsUID() ([]*corev1.Pod, error) { m.mutex.RLock() @@ -167,17 +202,6 @@ 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(), - } -} - func (pd PodDevices) DeepCopy() PodDevices { if pd == nil { return nil diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index eded68aa6d..be499b995d 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -93,8 +93,6 @@ func NewScheduler() *Scheduler { s.nodeManager = newNodeManager() s.podManager = device.NewPodManager() s.quotaManager = device.NewQuotaManager() - // Use dummy leader manager when leaderElect is disabled - // This ensures IsLeader() always returns true and synced will not be set to false s.leaderManager = leaderelection.NewDummyLeaderManager(true) if config.LeaderElect { callbacks := leaderelection.LeaderCallbacks{ @@ -135,6 +133,18 @@ func (s *Scheduler) doNodeNotify() { } } +func allInitContainersSucceeded(pod *corev1.Pod) bool { + if len(pod.Status.InitContainerStatuses) == 0 { + return false + } + for _, s := range pod.Status.InitContainerStatuses { + if s.State.Terminated == nil || s.State.Terminated.ExitCode != 0 { + return false + } + } + return true +} + func (s *Scheduler) onAddPod(obj any) { pod, ok := obj.(*corev1.Pod) if !ok { @@ -157,18 +167,72 @@ func (s *Scheduler) onAddPod(obj any) { s.podManager.UpdatePod(pod) return } - podDev, err := device.DecodePodDevices(device.SupportDevices, pod.Annotations) + + rawDevices, err := device.DecodePodDevices(device.SupportDevices, pod.Annotations) if err != nil { klog.ErrorS(err, "failed to decode pod devices", "pod", klog.KObj(pod)) return } - if s.podManager.AddPod(pod, nodeID, podDev) { - s.quotaManager.AddUsage(pod, podDev) + + effectiveDevices := device.CollapseInitContainerUsage(pod, rawDevices) + + if s.podManager.AddPod(pod, nodeID, effectiveDevices) { + s.quotaManager.AddUsage(pod, effectiveDevices) } } -func (s *Scheduler) onUpdatePod(_, newObj any) { - s.onAddPod(newObj) +func (s *Scheduler) onUpdatePod(oldObj, newObj any) { + newPod, ok := newObj.(*corev1.Pod) + if !ok { + return + } + + klog.V(5).InfoS("Pod updated", "pod", klog.KObj(newPod)) + + if _, ok := newPod.Annotations[util.AssignedNodeAnnotations]; !ok { + return + } + + if util.IsPodInTerminatedState(newPod) { + if pi, ok := s.podManager.TakeAndDeletePod(newPod); ok { + s.quotaManager.RmUsage(newPod, pi.Devices) + } + return + } + + if util.IsPodTerminating(newPod) { + s.podManager.UpdatePod(newPod) + return + } + + pi, exists := s.podManager.GetPod(newPod) + if !exists { + s.onAddPod(newPod) + return + } + + s.podManager.UpdatePod(newPod) + + if !pi.InitContainerResourceReleased && allInitContainersSucceeded(newPod) { + rawDevices, err := device.DecodePodDevices(device.SupportDevices, newPod.Annotations) + if err != nil { + klog.ErrorS(err, "failed to decode pod devices during shrink", "pod", klog.KObj(newPod)) + return + } + + appOnlyDevices := device.AppContainersOnlyDeviceUsage(newPod, rawDevices) + + oldDevices, ok := s.podManager.ShrinkUsage(newPod, appOnlyDevices) + if ok { + s.quotaManager.RmUsage(newPod, oldDevices) + s.quotaManager.AddUsage(newPod, appOnlyDevices) + klog.InfoS("Init containers completed, shrunk usage", + "pod", klog.KObj(newPod), + "oldUsage", oldDevices, + "newUsage", appOnlyDevices, + ) + } + } } func (s *Scheduler) onDelPod(obj any) { @@ -401,7 +465,14 @@ func (s *Scheduler) register(labelSelector labels.Selector, printedLog map[strin klog.V(5).InfoS("Skipping device cleanup for vendor not present in scheduler cache", "nodeName", val.Name, "deviceVendor", devhandsk) continue } - klog.Warning("Device is unhealthy, cleaning up node", "nodeName", val.Name, "deviceVendor", devhandsk) + // klog.Warning does plain fmt.Print-style concatenation of its arguments - + // klog v2 has no structured WarningS variant. Passing alternating + // "key", value pairs to it (as if it were InfoS/ErrorS) produces a garbled, + // unstructured log line instead of the intended structured fields. Use + // ErrorS (nil error is fine here; this is a detected condition, not a Go + // error) to match the structured logging used throughout the rest of this + // file. + klog.ErrorS(nil, "Device is unhealthy, cleaning up node", "nodeName", val.Name, "deviceVendor", devhandsk) err := devInstance.NodeCleanUp(val.Name) if err != nil { klog.ErrorS(err, "Node cleanup failed", "nodeName", val.Name, "deviceVendor", devhandsk) @@ -890,23 +961,18 @@ func (s *Scheduler) Bind(args extenderv1.ExtenderBindingArgs) (*extenderv1.Exten func (s *Scheduler) Filter(args extenderv1.ExtenderArgs) (*extenderv1.ExtenderFilterResult, error) { klog.InfoS("Starting schedule filter process", "pod", args.Pod.Name, "uuid", args.Pod.UID, "namespace", args.Pod.Namespace) resourceReqs := device.Resourcereqs(args.Pod) - resourceReqTotal := 0 - for _, n := range resourceReqs { - for _, k := range n { - resourceReqTotal += int(k.Nums) + + hasHAMiResource := false + + for _, reqMap := range resourceReqs { + if len(reqMap) > 0 { + hasHAMiResource = true + break } } - if resourceReqTotal == 0 { - klog.V(1).InfoS("Pod does not request any resources", - "pod", args.Pod.Name) - s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, "", fmt.Errorf("does not request any resource")) - if args.Nodes != nil { - return &extenderv1.ExtenderFilterResult{ - Nodes: args.Nodes, - FailedNodes: nil, - Error: "", - }, nil - } + + if !hasHAMiResource { + klog.V(1).InfoS("Pod does not request any resources", "pod", args.Pod.Name) return &extenderv1.ExtenderFilterResult{ NodeNames: args.NodeNames, FailedNodes: nil, @@ -914,17 +980,9 @@ func (s *Scheduler) Filter(args extenderv1.ExtenderArgs) (*extenderv1.ExtenderFi }, nil } if args.Nodes != nil { - klog.V(2).InfoS("Choosing simulation filter path", - "pod", klog.KObj(args.Pod), - "reason", "request contains full nodes", - "nodesLen", nodeListLen(args.Nodes), - "nodeNamesLen", nodeNamesLen(args.NodeNames)) return s.filterSimulation(args, resourceReqs) } - klog.V(2).InfoS("Choosing live filter path", - "pod", klog.KObj(args.Pod), - "reason", "request does not contain full nodes", - "nodeNamesLen", nodeNamesLen(args.NodeNames)) + if pi, ok := s.podManager.TakeAndDeletePod(args.Pod); ok { s.quotaManager.RmUsage(args.Pod, pi.Devices) } @@ -934,8 +992,7 @@ func (s *Scheduler) Filter(args extenderv1.ExtenderArgs) (*extenderv1.ExtenderFi return nil, err } if len(failedNodes) != 0 { - klog.V(5).InfoS("Nodes failed during usage retrieval", - "nodes", failedNodes) + klog.V(5).InfoS("Nodes failed during usage retrieval", "nodes", failedNodes) } nodeScores, err := s.calcScore(nodeUsage, resourceReqs, args.Pod, failedNodes) if err != nil { @@ -944,8 +1001,7 @@ func (s *Scheduler) Filter(args extenderv1.ExtenderArgs) (*extenderv1.ExtenderFi return nil, err } if len((*nodeScores).NodeList) == 0 { - klog.V(4).InfoS("No available nodes meet the required scores", - "pod", args.Pod.Name) + klog.V(4).InfoS("No available nodes meet the required scores", "pod", args.Pod.Name) s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, "", fmt.Errorf("no available node, %d nodes do not meet", len(*args.NodeNames))) return &extenderv1.ExtenderFilterResult{ FailedNodes: failedNodes, @@ -967,20 +1023,24 @@ func (s *Scheduler) Filter(args extenderv1.ExtenderArgs) (*extenderv1.ExtenderFi val.PatchAnnotations(args.Pod, &annotations, m.Devices) } - added := s.podManager.AddPod(args.Pod, m.NodeID, m.Devices) - if added { - s.quotaManager.AddUsage(args.Pod, m.Devices) - } - - err = util.PatchPodAnnotations(args.Pod, annotations) - if err != nil { - s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, "", err) + rawDevices := m.Devices + effectiveDevices := device.CollapseInitContainerUsage(args.Pod, rawDevices) + if args.Nodes == nil { + added := s.podManager.AddPod(args.Pod, m.NodeID, effectiveDevices) if added { - s.quotaManager.RmUsage(args.Pod, m.Devices) + s.quotaManager.AddUsage(args.Pod, effectiveDevices) // use collapsed + } + err = util.PatchPodAnnotations(args.Pod, annotations) + if err != nil { + s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, "", err) + if added { + s.quotaManager.RmUsage(args.Pod, effectiveDevices) + } + s.podManager.DelPod(args.Pod) + return nil, err } - s.podManager.DelPod(args.Pod) - return nil, err } + successMsg := genSuccessMsg(len(*args.NodeNames), m.NodeID, nodeScores.NodeList) s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringSucceed, successMsg, nil) res := extenderv1.ExtenderFilterResult{NodeNames: &[]string{m.NodeID}} diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go index 2687cfef22..dfb1a0a42e 100644 --- a/pkg/scheduler/scheduler_test.go +++ b/pkg/scheduler/scheduler_test.go @@ -20,6 +20,7 @@ import ( "context" "fmt" "maps" + "slices" "strings" "sync/atomic" "testing" @@ -685,11 +686,11 @@ func Test_Filter(t *testing.T) { } tests := []struct { - name string - args extenderv1.ExtenderArgs - want *extenderv1.ExtenderFilterResult - wantPodAnnotationDeviceID string - wantErr error + name string + args extenderv1.ExtenderArgs + want *extenderv1.ExtenderFilterResult + wantPodAnnotationDeviceIDs []string + wantErr error }{ { name: "node use binpack gpu use binpack policy", @@ -726,7 +727,56 @@ func Test_Filter(t *testing.T) { want: &extenderv1.ExtenderFilterResult{ NodeNames: &[]string{"node2"}, }, - wantPodAnnotationDeviceID: "device4", + wantPodAnnotationDeviceIDs: []string{"device4"}, + }, + { + name: "pod with init containers fits correctly using max resource logic (Binpack)", + args: extenderv1.ExtenderArgs{ + Pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-init-containers", + UID: "test-init-uid", + Annotations: map[string]string{ + util.GPUSchedulerPolicyAnnotationKey: util.GPUSchedulerPolicyBinpack.String(), + util.NodeSchedulerPolicyAnnotationKey: util.NodeSchedulerPolicyBinpack.String(), + }, + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-1", + Image: "busybox", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(20, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(5000, resource.BinarySI), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-1", + Image: "chrstnhntschl/gpu_burn", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(20, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(4000, resource.BinarySI), + }, + }, + }, + }, + }, + }, + NodeNames: &[]string{"node1", "node2"}, + }, + wantErr: nil, + want: &extenderv1.ExtenderFilterResult{ + NodeNames: &[]string{"node2"}, + }, + wantPodAnnotationDeviceIDs: []string{"device3"}, }, { name: "node use binpack gpu use spread policy", @@ -763,7 +813,7 @@ func Test_Filter(t *testing.T) { want: &extenderv1.ExtenderFilterResult{ NodeNames: &[]string{"node2"}, }, - wantPodAnnotationDeviceID: "device3", + wantPodAnnotationDeviceIDs: []string{"device3", "device4"}, // Both are acceptable due to tie }, { name: "node use spread gpu use binpack policy", @@ -800,7 +850,7 @@ func Test_Filter(t *testing.T) { want: &extenderv1.ExtenderFilterResult{ NodeNames: &[]string{"node1"}, }, - wantPodAnnotationDeviceID: "device1", + wantPodAnnotationDeviceIDs: []string{"device1"}, }, { name: "node use spread gpu use spread policy", @@ -837,7 +887,7 @@ func Test_Filter(t *testing.T) { want: &extenderv1.ExtenderFilterResult{ NodeNames: &[]string{"node1"}, }, - wantPodAnnotationDeviceID: "device2", + wantPodAnnotationDeviceIDs: []string{"device2"}, }, } @@ -850,7 +900,12 @@ func Test_Filter(t *testing.T) { assert.DeepEqual(t, test.want, got) getPod, _ := client.KubeClient.CoreV1().Pods(test.args.Pod.Namespace).Get(context.Background(), test.args.Pod.Name, metav1.GetOptions{}) podDevices, _ := device.DecodePodDevices(device.SupportDevices, getPod.Annotations) - assert.DeepEqual(t, test.wantPodAnnotationDeviceID, podDevices["NVIDIA"][0][0].UUID) + + actualUUID := podDevices["NVIDIA"][0][0].UUID + + if !slices.Contains(test.wantPodAnnotationDeviceIDs, actualUUID) { + t.Errorf("expected one of %v, got %s", test.wantPodAnnotationDeviceIDs, actualUUID) + } }) } } @@ -1458,7 +1513,7 @@ func Test_ResourceQuota(t *testing.T) { wantErr: nil, want: &extenderv1.ExtenderFilterResult{ FailedNodes: map[string]string{ - "node1": "NodeUnfitPod", + "node1": "1/4 AllocatedCardsInsufficientRequest, 3/4 ResourceQuotaNotFit", }, }, }, @@ -1546,7 +1601,7 @@ func Test_ResourceQuota(t *testing.T) { wantErr: nil, want: &extenderv1.ExtenderFilterResult{ FailedNodes: map[string]string{ - "node1": "NodeUnfitPod", + "node1": "4/4 ResourceQuotaNotFit", }, }, }, @@ -1905,7 +1960,7 @@ func TestFilterTemplateNodesMissingRegisterAnnotation(t *testing.T) { require.Equal(t, "node unregistered", res.FailedNodes["template-node-cold-zero"]) } -func Test_Scheduler_Issue1368_TerminatingPodRetainsCache(t *testing.T) { +func Test_SchedulerTerminatingPodRetainsCache(t *testing.T) { s := NewScheduler() podDevces := device.PodDevices{ @@ -1952,6 +2007,155 @@ func Test_Scheduler_Issue1368_TerminatingPodRetainsCache(t *testing.T) { _, ok = s.podManager.GetPod(terminatedPod) assert.Equal(t, false, ok, "Pod should be removed from cache after reaching a terminal phase (Succeeded/Failed)") } +func Test_onUpdatePod_InitContainerShrink(t *testing.T) { + s := NewScheduler() + sConfig := &config.Config{ + NvidiaConfig: nvidia.NvidiaConfig{ + ResourceCountName: "hami.io/gpu", + ResourceMemoryName: "hami.io/gpumem", + ResourceMemoryPercentageName: "hami.io/gpumem-percentage", + ResourceCoreName: "hami.io/gpucores", + DefaultMemory: 0, + DefaultCores: 0, + DefaultGPUNum: 1, + }, + } + if err := config.InitDevicesWithConfig(sConfig); err != nil { + t.Fatalf("Failed to initialize devices with config: %v", err) + } + client.KubeClient = fake.NewClientset() + s.kubeClient = client.KubeClient + informerFactory := informers.NewSharedInformerFactoryWithOptions(client.KubeClient, time.Hour) + s.podLister = informerFactory.Core().V1().Pods().Lister() + informer := informerFactory.Core().V1().Pods().Informer() + informer.AddEventHandler(cache.ResourceEventHandlerFuncs{ + AddFunc: s.onAddPod, + UpdateFunc: s.onUpdatePod, + DeleteFunc: s.onDelPod, + }) + informerFactory.Start(s.stopCh) + informerFactory.WaitForCacheSync(s.stopCh) + s.addAllEventHandlers() + + // Setup a pod with one init container (20Gi mem, 10 cores) and one app container (10Gi mem, 5 cores). + // Effective = max(20,10)=20Gi mem, max(10,5)=10 cores. + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + UID: "shrink-uid", + Name: "test-shrink", + Namespace: "default", + Annotations: map[string]string{ + util.AssignedNodeAnnotations: "node1", + }, + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{{ + Name: "init", + Image: "busybox", + }}, + Containers: []corev1.Container{{ + Name: "app", + Image: "app", + }}, + }, + Status: corev1.PodStatus{ + Phase: corev1.PodRunning, + }, + } + + rawDevices := device.PodDevices{ + nvidia.NvidiaGPUDevice: device.PodSingleDevice{ + // Init container (index 0) + {{ + Idx: 0, UUID: "GPU0", Usedmem: 20000, Usedcores: 10, + }}, + // App container (index 1) + {{ + Idx: 0, UUID: "GPU0", Usedmem: 10000, Usedcores: 5, + }}, + }, + } + encoded := device.EncodePodDevices(device.SupportDevices, rawDevices) + maps.Copy(pod.Annotations, encoded) + + s.onAddPod(pod) + + pi, ok := s.podManager.GetPod(pod) + require.True(t, ok) + require.Len(t, pi.Devices[nvidia.NvidiaGPUDevice], 1) + require.Len(t, pi.Devices[nvidia.NvidiaGPUDevice][0], 1) + assert.Equal(t, int32(20000), pi.Devices[nvidia.NvidiaGPUDevice][0][0].Usedmem) + assert.Equal(t, int32(10), pi.Devices[nvidia.NvidiaGPUDevice][0][0].Usedcores) + + updatedPod := pod.DeepCopy() + updatedPod.Status.InitContainerStatuses = []corev1.ContainerStatus{ + { + Name: "init", + State: corev1.ContainerState{ + Terminated: &corev1.ContainerStateTerminated{ExitCode: 0}, + }, + }, + } + + s.onUpdatePod(pod, updatedPod) + + piAfter, ok := s.podManager.GetPod(updatedPod) + require.True(t, ok) + assert.Assert(t, piAfter.InitContainerResourceReleased) + assert.Equal(t, int32(10000), piAfter.Devices[nvidia.NvidiaGPUDevice][0][0].Usedmem) + assert.Equal(t, int32(5), piAfter.Devices[nvidia.NvidiaGPUDevice][0][0].Usedcores) + + quotas := s.quotaManager.GetResourceQuota() + require.Contains(t, quotas, "default") + memQuota := (*quotas["default"])["hami.io/gpumem"] + coresQuota := (*quotas["default"])["hami.io/gpucores"] + require.NotNil(t, memQuota) + require.NotNil(t, coresQuota) + assert.Equal(t, int64(10000), memQuota.Used) + assert.Equal(t, int64(5), coresQuota.Used) +} + +func Test_getNodesUsage_WithInitContainers(t *testing.T) { + nodeMage := newNodeManager() + nodeMage.addNode("node1", &device.NodeInfo{ + ID: "node1", Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}, + Devices: map[string][]device.DeviceInfo{ + nvidia.NvidiaGPUDevice: {{ + ID: "GPU0", Index: 0, Count: 10, Devmem: 102400, Devcore: 100, + Numa: 1, Mode: "hami", Health: true, + }}, + }, + }) + podMap := device.NewPodManager() + // Create a pod with init container (mem 20000, cores 10) and app (mem 10000, cores 5) + // Collapsed usage = peak (20000,10) + collapsed := device.CollapseInitContainerUsage( + &corev1.Pod{ + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{{Name: "init"}}, + Containers: []corev1.Container{{Name: "app"}}, + }, + }, + device.PodDevices{ + nvidia.NvidiaGPUDevice: device.PodSingleDevice{ + {{Idx: 0, UUID: "GPU0", Usedmem: 20000, Usedcores: 10}}, + {{Idx: 0, UUID: "GPU0", Usedmem: 10000, Usedcores: 5}}, + }, + }, + ) + podMap.AddPod(&corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{UID: "initpod", Name: "initpod", Namespace: "default"}, + }, "node1", collapsed) + + s := Scheduler{nodeManager: nodeMage, podManager: podMap} + nodes := []string{"node1"} + cachenodeMap, _, _, err := s.getNodesUsage(&nodes, nil) + require.NoError(t, err) + v := (*cachenodeMap)["node1"] + assert.Equal(t, int32(1), v.Devices.DeviceLists[0].Device.Used) + assert.Equal(t, int32(20000), v.Devices.DeviceLists[0].Device.Usedmem) // peak, not sum + assert.Equal(t, int32(10), v.Devices.DeviceLists[0].Device.Usedcores) +} func Test_onAddPod_BadDeviceAnnotation(t *testing.T) { device.SupportDevices["TEST"] = "hami.io/test-allocated" diff --git a/pkg/scheduler/score.go b/pkg/scheduler/score.go index bae2518212..c83eca7f5d 100644 --- a/pkg/scheduler/score.go +++ b/pkg/scheduler/score.go @@ -42,61 +42,158 @@ func viewStatus(usage NodeUsage) { func getNodeResources(list NodeUsage, t string) []*device.DeviceUsage { l := []*device.DeviceUsage{} for _, val := range list.Devices.DeviceLists { - if strings.Contains(val.Device.Type, t) { + if val.Device == nil { + continue + } + if getDeviceBaseType(val.Device.Type) == t { l = append(l, val.Device) } } return l } +// getDeviceBaseType maps a device model string (e.g., "NVIDIA A100-SXM4-40GB") to the base type key used in the device registry (e.g., "NVIDIA"). +func getDeviceBaseType(model string) string { + bestName := "" + bestWordLen := -1 + for name, dev := range device.GetDevices() { + word := dev.CommonWord() + if !strings.HasPrefix(model, word) { + continue + } + if len(word) > bestWordLen || (len(word) == bestWordLen && name < bestName) { + bestName = name + bestWordLen = len(word) + } + } + if bestWordLen == -1 { + return model + } + return bestName +} + +// nodeDeviceBaseTypes returns the set of base device types physically present on the node, computed once so that callers don't repeatedly re-derive it (which, combined with any residual +// ambiguity in getDeviceBaseType, previously risked producing different results across calls within the same scheduling attempt). +func nodeDeviceBaseTypes(list policy.DeviceUsageList) map[string]struct{} { + types := make(map[string]struct{}) + for _, dl := range list.DeviceLists { + if dl.Device != nil { + types[getDeviceBaseType(dl.Device.Type)] = struct{}{} + } + } + return types +} + +// fitInDevices tries to allocate a single container's device requests on the given node. func fitInDevices(node *NodeUsage, requests device.ContainerDeviceRequests, pod *corev1.Pod, nodeInfo *device.NodeInfo, devinput *device.PodDevices) (bool, string) { - //devmap := make(map[string]device.ContainerDevices) - total, totalCore, totalMem := int32(0), int32(0), int32(0) - free, freeCore, freeMem := int32(0), int32(0), int32(0) - sums := 0 - // computer all device score for one node + // Snapshot the entire node state before any modifications. + type devSnapshot struct { + dev *device.DeviceUsage + used int32 + usedcores int32 + usedmem int32 + } + saved := make([]devSnapshot, len(node.Devices.DeviceLists)) + for i := range node.Devices.DeviceLists { + d := node.Devices.DeviceLists[i].Device + saved[i] = devSnapshot{ + dev: d, + used: d.Used, + usedcores: d.Usedcores, + usedmem: d.Usedmem, + } + } + rollbackAll := func() { + for _, s := range saved { + s.dev.Used = s.used + s.dev.Usedcores = s.usedcores + s.dev.Usedmem = s.usedmem + } + } + for index := range node.Devices.DeviceLists { node.Devices.DeviceLists[index].ComputeScore(requests) } - //This loop is for requests for different devices + + // Process each device type in the request. for _, k := range requests { - sums += int(k.Nums) - if int(k.Nums) > len(node.Devices.DeviceLists) { - klog.V(5).InfoS(common.NodeInsufficientDevice, "pod", klog.KObj(pod), "request devices nums", k.Nums, "node device nums", len(node.Devices.DeviceLists)) - return false, common.NodeInsufficientDevice - } + // Sort devices by score (best fit first). sort.Sort(node.Devices) - _, ok := device.GetDevices()[k.Type] + + devPlugin, ok := device.GetDevices()[k.Type] if !ok { - return false, "Device type not found" + rollbackAll() + errMsg := "Device type not found" + klog.ErrorS(nil, errMsg, "pod", klog.KObj(pod), "type", k.Type, "node", node.Node.Name) + return false, errMsg } - fit, tmpDevs, reason := device.GetDevices()[k.Type].Fit(getNodeResources(*node, k.Type), k, pod, nodeInfo, devinput) - if fit { - for idx, val := range tmpDevs[k.Type] { - for nidx, v := range node.Devices.DeviceLists { - //bc node.Devices has been sorted, so we should find out the correct device - if v.Device.ID != val.UUID { - continue - } - total += v.Device.Count - totalCore += v.Device.Totalcore - totalMem += v.Device.Totalmem - free += v.Device.Count - v.Device.Used - freeCore += v.Device.Totalcore - v.Device.Usedcores - freeMem += v.Device.Totalmem - v.Device.Usedmem - err := device.GetDevices()[k.Type].AddResourceUsage(pod, node.Devices.DeviceLists[nidx].Device, &tmpDevs[k.Type][idx]) - if err != nil { - klog.Errorf("AddResourceUsage failed:%s", err.Error()) - return false, "AddResourceUsage failed" - } - klog.V(5).Infoln("After AddResourceUsage:", node.Devices.DeviceLists[nidx].Device) + + typeDevices := getNodeResources(*node, k.Type) + if int(k.Nums) > len(typeDevices) { + klog.V(5).InfoS(common.NodeInsufficientDevice, "pod", klog.KObj(pod), + "request devices nums", k.Nums, "node device nums (type)", len(typeDevices), "type", k.Type) + rollbackAll() + return false, common.NodeInsufficientDevice + } + + fit, tmpDevs, reason := devPlugin.Fit(typeDevices, k, pod, nodeInfo, devinput) + if !fit { + rollbackAll() + return false, reason + } + + type usageSnapshot struct { + idx int + used int32 + usedcores int32 + usedmem int32 + } + modified := make([]usageSnapshot, 0, len(tmpDevs[k.Type])) + + for idx, val := range tmpDevs[k.Type] { + targetIdx := -1 + for nidx, v := range node.Devices.DeviceLists { + if v.Device.ID == val.UUID { + targetIdx = nidx + break } } - } else { - return false, reason + if targetIdx == -1 { + rollbackAll() + errMsg := fmt.Sprintf("Device with UUID %q not found on node %s after Fit", val.UUID, node.Node.Name) + klog.ErrorS(nil, errMsg, "pod", klog.KObj(pod)) + return false, errMsg + } + d := node.Devices.DeviceLists[targetIdx].Device + + snap := usageSnapshot{ + idx: targetIdx, + used: d.Used, + usedcores: d.Usedcores, + usedmem: d.Usedmem, + } + + err := devPlugin.AddResourceUsage(pod, d, &tmpDevs[k.Type][idx]) + if err != nil { + klog.Errorf("AddResourceUsage failed for device %s: %v, rolling back all changes", d.ID, err) + for _, m := range modified { + dev := node.Devices.DeviceLists[m.idx].Device + dev.Used = m.used + dev.Usedcores = m.usedcores + dev.Usedmem = m.usedmem + } + rollbackAll() + errMsg := fmt.Sprintf("AddResourceUsage failed for device %s: %v", d.ID, err) + return false, errMsg + } + modified = append(modified, snap) + klog.V(5).Infof("Allocated device %s: used=%d, cores=%d, mem=%d", + d.ID, d.Used, d.Usedcores, d.Usedmem) } + (*devinput)[k.Type] = append((*devinput)[k.Type], tmpDevs[k.Type]) } + return true, "" } @@ -121,85 +218,228 @@ func (s *Scheduler) calcScoreWithOptions(nodes *map[string]*NodeUsage, resourceR failedNodesMutex := sync.Mutex{} failureReason := make(map[string][]string) errCh := make(chan error, len(*nodes)) + + numInitContainers := len(task.Spec.InitContainers) + for nodeID, node := range *nodes { wg.Add(1) go func(nodeID string, node *NodeUsage) { defer wg.Done() viewStatus(*node) - score := policy.NodeScore{NodeID: nodeID, Node: node.Node, Devices: make(device.PodDevices), Score: 0} - score.ComputeDefaultScore(node.Devices) - snapshot := score.SnapshotDevice(node.Devices) - nodeInfo := node.NodeInfo if nodeInfo == nil { var err error nodeInfo, err = s.GetNode(nodeID) if err != nil { klog.ErrorS(err, "Failed to get node", "nodeID", nodeID) + failedNodesMutex.Lock() + failedNodes[nodeID] = fmt.Sprintf("failed to fetch node info: %v", err) + failedNodesMutex.Unlock() errCh <- err return } } - // Assume the node is a fit by default. This handles pods with no device - // requests, which should be schedulable on any node. + baseTypes := nodeDeviceBaseTypes(node.Devices) + + type peakUsageSnapshot struct { + used int32 + usedcores int32 + usedmem int32 + } + peakUsage := make(map[string]peakUsageSnapshot) + for _, dl := range node.Devices.DeviceLists { + id := dl.Device.ID + peakUsage[id] = peakUsageSnapshot{ + used: dl.Device.Used, + usedcores: dl.Device.Usedcores, + usedmem: dl.Device.Usedmem, + } + } + + // 1) Check init containers (they run sequentially, each on a fresh copy) + var initAllocs device.PodDevices + if numInitContainers > 0 { + initAllocs = make(device.PodDevices) + + initFit := true + for i, req := range resourceReqs { + if i >= numInitContainers { + break + } + // Pad previous slots for all types to maintain index alignment + for typ := range baseTypes { + for len(initAllocs[typ]) < i { + initAllocs[typ] = append(initAllocs[typ], device.ContainerDevices{}) + } + } + if len(req) == 0 { + for typ := range baseTypes { + initAllocs[typ] = append(initAllocs[typ], device.ContainerDevices{}) + } + continue + } + nodeCopy := node.DeepCopy() + fit, reason := fitInDevices(nodeCopy, req, task, nodeInfo, &initAllocs) + if !fit { + klog.V(4).InfoS("Init container does not fit", + "pod", klog.KObj(task), "node", nodeID, "containerIndex", i, "reason", reason) + failedNodesMutex.Lock() + failedNodes[nodeID] = reason + for reasonType := range common.ParseReason(reason) { + failureReason[reasonType] = append(failureReason[reasonType], nodeID) + } + failedNodesMutex.Unlock() + initFit = false + break + } + // Record peak usage for this init container per device (by device ID). + for _, dl := range nodeCopy.Devices.DeviceLists { + id := dl.Device.ID + if p, ok := peakUsage[id]; ok { + if dl.Device.Used > p.used { + p.used = dl.Device.Used + } + if dl.Device.Usedcores > p.usedcores { + p.usedcores = dl.Device.Usedcores + } + if dl.Device.Usedmem > p.usedmem { + p.usedmem = dl.Device.Usedmem + } + peakUsage[id] = p + } else { + peakUsage[id] = peakUsageSnapshot{ + used: dl.Device.Used, + usedcores: dl.Device.Usedcores, + usedmem: dl.Device.Usedmem, + } + } + } + // Ensure every type has an entry for this container (even if empty) + for typ := range baseTypes { + if len(initAllocs[typ]) == i { + initAllocs[typ] = append(initAllocs[typ], device.ContainerDevices{}) + } + } + } + if !initFit { + return + } + } + + // 2) Allocate app containers (they run concurrently, cumulative) + appNodeCopy := node.DeepCopy() + score := policy.NodeScore{ + NodeID: nodeID, + Node: node.Node, + Devices: make(device.PodDevices), + Score: 0, + } + score.ComputeDefaultScore(appNodeCopy.Devices) + snapshot := score.SnapshotDevice(appNodeCopy.Devices) + + appIndex := 0 ctrfit := true - deviceType := "" - //This loop is for different container request + for ctrid, n := range resourceReqs { + if ctrid < numInitContainers { + continue + } + for typ := range baseTypes { + for len(score.Devices[typ]) < appIndex { + score.Devices[typ] = append(score.Devices[typ], device.ContainerDevices{}) + } + } sums := 0 for _, k := range n { sums += int(k.Nums) } - - // container need no device and we have got certain deviceType - if sums == 0 && deviceType != "" { - score.Devices[deviceType] = append(score.Devices[deviceType], device.ContainerDevices{}) - continue - } - klog.V(5).InfoS("fitInDevices", "pod", klog.KObj(task), "node", nodeID) - fit, reason := fitInDevices(node, n, task, nodeInfo, &score.Devices) - // found certain deviceType, fill missing empty allocation for containers before this - for idx := range score.Devices { - deviceType = idx - for len(score.Devices[idx]) <= ctrid { - emptyContainerDevices := device.ContainerDevices{} - emptyPodSingleDevice := device.PodSingleDevice{} - emptyPodSingleDevice = append(emptyPodSingleDevice, emptyContainerDevices) - score.Devices[idx] = append(emptyPodSingleDevice, score.Devices[idx]...) + if sums == 0 { + for typ := range baseTypes { + score.Devices[typ] = append(score.Devices[typ], device.ContainerDevices{}) } + appIndex++ + continue } + fit, reason := fitInDevices(appNodeCopy, n, task, nodeInfo, &score.Devices) ctrfit = fit if !fit { klog.V(4).InfoS(common.NodeUnfitPod, "pod", klog.KObj(task), "node", nodeID, "reason", reason) failedNodesMutex.Lock() - failedNodes[nodeID] = common.NodeUnfitPod - if detailedFailureReason { - failedNodes[nodeID] = reason - } + failedNodes[nodeID] = reason for reasonType := range common.ParseReason(reason) { failureReason[reasonType] = append(failureReason[reasonType], nodeID) } failedNodesMutex.Unlock() break } + // Ensure every type has an entry for this container + for typ := range baseTypes { + if len(score.Devices[typ]) == appIndex { + score.Devices[typ] = append(score.Devices[typ], device.ContainerDevices{}) + } + } + appIndex++ + } + + if !ctrfit { + return + } + + // 3) Prepend init container allocations + if numInitContainers > 0 && initAllocs != nil { + for devType, initConList := range initAllocs { + score.Devices[devType] = append(initConList, score.Devices[devType]...) + } } - if ctrfit { - fitNodesMutex.Lock() - res.NodeList = append(res.NodeList, &score) - fitNodesMutex.Unlock() - score.OverrideScore(snapshot, userNodePolicy) - klog.V(4).InfoS(common.NodeFitPod, "pod", klog.KObj(task), "node", nodeID, "score", score.Score) + for _, dl := range node.Devices.DeviceLists { + id := dl.Device.ID + // Find the corresponding device in the app copy + var appDev *device.DeviceUsage + for _, adl := range appNodeCopy.Devices.DeviceLists { + if adl.Device.ID == id { + appDev = adl.Device + break + } + } + finalUsed := int32(0) + finalCores := int32(0) + finalMem := int32(0) + if appDev != nil { + finalUsed = appDev.Used + finalCores = appDev.Usedcores + finalMem = appDev.Usedmem + } + if p, ok := peakUsage[id]; ok { + if p.used > finalUsed { + finalUsed = p.used + } + if p.usedcores > finalCores { + finalCores = p.usedcores + } + if p.usedmem > finalMem { + finalMem = p.usedmem + } + } + dl.Device.Used = finalUsed + dl.Device.Usedcores = finalCores + dl.Device.Usedmem = finalMem } + + score.OverrideScore(snapshot, userNodePolicy) + fitNodesMutex.Lock() + res.NodeList = append(res.NodeList, &score) + fitNodesMutex.Unlock() + + klog.V(4).InfoS(common.NodeFitPod, "pod", klog.KObj(task), "node", nodeID, "score", score.Score) }(nodeID, node) } wg.Wait() close(errCh) - // only pod scheduler failure will record failure event - if recordEvents && len(res.NodeList) == 0 { + if len(res.NodeList) == 0 { for reasonType, failureNodes := range failureReason { sort.Strings(failureNodes) reason := fmt.Errorf("%d nodes %s(%s)", len(failureNodes), reasonType, strings.Join(failureNodes, ",")) diff --git a/pkg/scheduler/score_test.go b/pkg/scheduler/score_test.go index fd48e61284..e35191e6cd 100644 --- a/pkg/scheduler/score_test.go +++ b/pkg/scheduler/score_test.go @@ -402,6 +402,124 @@ func Test_calcScore(t *testing.T) { err: nil, }, }, + { + name: "init container requests more memory than any single device has (should filter node)", + args: struct { + nodes *map[string]*NodeUsage + nums device.PodDeviceRequests + annos map[string]string + task *corev1.Pod + }{ + nodes: &map[string]*NodeUsage{ + "node1": { + Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}, + Devices: policy.DeviceUsageList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + DeviceLists: []*policy.DeviceListsScore{ + { + Device: &device.DeviceUsage{ + ID: "uuid1", + Index: 0, + Used: 0, + Count: 10, + Usedmem: 7500, // Only 500 available (8000-7500) + Totalmem: 8000, + Totalcore: 100, + Usedcores: 0, + Numa: 0, + Type: nvidia.NvidiaGPUDevice, + Health: true, + }, + }, + { + Device: &device.DeviceUsage{ + ID: "uuid2", + Index: 1, + Used: 0, + Count: 10, + Usedmem: 7500, // Only 500 available (8000-7500) + Totalmem: 8000, + Totalcore: 100, + Usedcores: 0, + Numa: 0, + Type: nvidia.NvidiaGPUDevice, + Health: true, + }, + }, + }, + }, + }, + }, + nums: device.PodDeviceRequests{ + // Index 0: InitContainer requests 4000 (more than any single device has) + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 1, + Type: nvidia.NvidiaGPUDevice, + Memreq: 4000, // Each device only has 500 available + Coresreq: 30, + }, + }, + // Index 1: App container requests only 100 + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 1, + Type: nvidia.NvidiaGPUDevice, + Memreq: 100, + Coresreq: 30, + }, + }, + }, + annos: make(map[string]string), + task: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-init-large-mem", + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-large", + Image: "busybox", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(30, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(4000, resource.BinarySI), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-small", + Image: "busybox", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(30, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(100, resource.BinarySI), + }, + }, + }, + }, + }, + }, + }, + wants: struct { + want *policy.NodeScoreList + failedNodes map[string]string + err error + }{ + want: &policy.NodeScoreList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + NodeList: []*policy.NodeScore{}, // Node should be filtered out + }, + failedNodes: map[string]string{ + "node1": "2/2 CardInsufficientMemory", + }, + err: nil, + }, + }, { name: "one node two device one pod one container use one device,but having use 50%", args: struct { @@ -1574,7 +1692,7 @@ func Test_calcScore(t *testing.T) { NodeList: []*policy.NodeScore{}, }, failedNodes: map[string]string{ - "node1": common.NodeUnfitPod, + "node1": "1/1 CardInsufficientMemory", }, err: nil, }, @@ -1680,8 +1798,146 @@ func Test_calcScore(t *testing.T) { NodeList: []*policy.NodeScore{}, }, failedNodes: map[string]string{ - "node1": common.NodeUnfitPod, - "node2": common.NodeUnfitPod, + "node1": "1/1 CardInsufficientMemory", + "node2": "1/1 CardInsufficientMemory", + }, + err: nil, + }, + }, + { + name: "two init containers, first has no device request, second does - tests slot alignment (issue #1667)", + args: struct { + nodes *map[string]*NodeUsage + nums device.PodDeviceRequests + annos map[string]string + task *corev1.Pod + }{ + nodes: &map[string]*NodeUsage{ + "node1": { + Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}, + Devices: policy.DeviceUsageList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + DeviceLists: []*policy.DeviceListsScore{ + { + Device: &device.DeviceUsage{ + ID: "uuid1", + Index: 0, + Used: 0, + Count: 10, + Usedmem: 0, + Totalmem: 8000, + Totalcore: 100, + Usedcores: 0, + Numa: 0, + Type: nvidia.NvidiaGPUDevice, + Health: true, + }, + }, + }, + }, + }, + }, + nums: device.PodDeviceRequests{ + // Index 0: InitContainer 0 - no device request (e.g. a setup/download step) + {}, + // Index 1: InitContainer 1 - requests a device + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 1, + Type: nvidia.NvidiaGPUDevice, + Memreq: 1000, + Coresreq: 30, + }, + }, + // Index 2: App container - requests a device + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 1, + Type: nvidia.NvidiaGPUDevice, + Memreq: 1000, + Coresreq: 30, + }, + }, + }, + annos: make(map[string]string), + task: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-init-heterogeneous", + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-no-device", + Image: "busybox", + Resources: corev1.ResourceRequirements{}, + }, + { + Name: "init-gpu-check", + Image: "chrstnhntschl/gpu_burn", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(30, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(1000, resource.BinarySI), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-gpu-burn", + Image: "chrstnhntschl/gpu_burn", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(30, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(1000, resource.BinarySI), + }, + }, + }, + }, + }, + }, + }, + wants: struct { + want *policy.NodeScoreList + failedNodes map[string]string + err error + }{ + want: &policy.NodeScoreList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + NodeList: []*policy.NodeScore{ + { + NodeID: "node1", + Devices: device.PodDevices{ + "NVIDIA": device.PodSingleDevice{ + // Index 0: init-no-device -> untouched placeholder slot + {}, + // Index 1: init-gpu-check -> real allocation + { + { + Idx: 0, + UUID: "uuid1", + Type: nvidia.NvidiaGPUDevice, + Usedcores: 30, + Usedmem: 1000, + }, + }, + // Index 2: app-gpu-burn -> real allocation + { + { + Idx: 0, + UUID: "uuid1", + Type: nvidia.NvidiaGPUDevice, + Usedcores: 30, + Usedmem: 1000, + }, + }, + }, + }, + Score: 0, + }, + }, }, err: nil, }, @@ -2444,6 +2700,380 @@ func Test_calcScore(t *testing.T) { err: nil, }, }, + { + name: "one node two devices, pod has one init container (uses 1 device) and one regular container (uses 1 device) - tests InitContainer capacity reset", + args: struct { + nodes *map[string]*NodeUsage + nums device.PodDeviceRequests + annos map[string]string + task *corev1.Pod + }{ + nodes: &map[string]*NodeUsage{ + "node1": { + Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}, + Devices: policy.DeviceUsageList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + DeviceLists: []*policy.DeviceListsScore{ + { + Device: &device.DeviceUsage{ + ID: "uuid1", + Index: 0, + Used: 0, + Count: 10, + Usedmem: 0, + Totalmem: 8000, + Totalcore: 100, + Usedcores: 0, + Numa: 0, + Type: nvidia.NvidiaGPUDevice, + Health: true, + }, + }, + { + Device: &device.DeviceUsage{ + ID: "uuid2", + Index: 0, + Used: 0, + Count: 10, + Usedmem: 0, + Totalmem: 8000, + Totalcore: 100, + Usedcores: 0, + Numa: 0, + Type: nvidia.NvidiaGPUDevice, + Health: true, + }, + }, + }, + }, + }, + }, + nums: device.PodDeviceRequests{ + // Index 0: InitContainer Request + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 1, + Type: nvidia.NvidiaGPUDevice, + Memreq: 1000, + Coresreq: 30, + }, + }, + // Index 1: Regular Container Request + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 1, + Type: nvidia.NvidiaGPUDevice, + Memreq: 1000, + Coresreq: 30, + }, + }, + }, + annos: make(map[string]string), + task: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-init-container", + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-gpu-burn", + Image: "chrstnhntschl/gpu_burn", + Args: []string{"6000"}, + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(30, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(1000, resource.BinarySI), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-gpu-burn", + Image: "chrstnhntschl/gpu_burn", + Args: []string{"6000"}, + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(30, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(1000, resource.BinarySI), + }, + }, + }, + }, + }, + }, + }, + wants: struct { + want *policy.NodeScoreList + failedNodes map[string]string + err error + }{ + want: &policy.NodeScoreList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + NodeList: []*policy.NodeScore{ + { + NodeID: "node1", + Devices: device.PodDevices{ + "NVIDIA": device.PodSingleDevice{ + { + { + Idx: 0, + UUID: "uuid2", + Type: nvidia.NvidiaGPUDevice, + Usedcores: 30, + Usedmem: 1000, + }, + }, + // Index 1: Regular Container allocated normally. + { + { + Idx: 0, + UUID: "uuid2", + Type: nvidia.NvidiaGPUDevice, + Usedcores: 30, + Usedmem: 1000, + }, + }, + }, + }, + Score: 0, + }, + }, + }, + err: nil, + }, + }, + { + name: "init container requests MORE than app container - needsInitClone=true path", + args: struct { + nodes *map[string]*NodeUsage + nums device.PodDeviceRequests + annos map[string]string + task *corev1.Pod + }{ + nodes: &map[string]*NodeUsage{ + "node1": { + Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}, + Devices: policy.DeviceUsageList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + DeviceLists: []*policy.DeviceListsScore{ + { + Device: &device.DeviceUsage{ + ID: "uuid1", Index: 0, Used: 0, Count: 10, + Usedmem: 0, Totalmem: 8000, Totalcore: 100, + Usedcores: 0, Numa: 0, + Type: nvidia.NvidiaGPUDevice, Health: true, + }, + }, + { + Device: &device.DeviceUsage{ + ID: "uuid2", Index: 0, Used: 0, Count: 10, + Usedmem: 0, Totalmem: 8000, Totalcore: 100, + Usedcores: 0, Numa: 0, + Type: nvidia.NvidiaGPUDevice, Health: true, + }, + }, + }, + }, + }, + }, + nums: device.PodDeviceRequests{ + // Index 0: InitContainer requests 2 GPUs (MORE than app) + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 2, Type: nvidia.NvidiaGPUDevice, + Memreq: 1000, Coresreq: 30, + }, + }, + // Index 1: App container requests only 1 GPU (LESS than init) + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 1, Type: nvidia.NvidiaGPUDevice, + Memreq: 1000, Coresreq: 30, + }, + }, + }, + annos: make(map[string]string), + task: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "test-init-bigger"}, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-heavy", + Image: "chrstnhntschl/gpu_burn", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(2, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(30, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(1000, resource.BinarySI), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-light", + Image: "chrstnhntschl/gpu_burn", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(30, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(1000, resource.BinarySI), + }, + }, + }, + }, + }, + }, + }, + wants: struct { + want *policy.NodeScoreList + failedNodes map[string]string + err error + }{ + want: &policy.NodeScoreList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + NodeList: []*policy.NodeScore{ + { + NodeID: "node1", + Devices: device.PodDevices{ + "NVIDIA": device.PodSingleDevice{ + // Init container allocated 2 GPUs from fresh snapshot + { + {Idx: 0, UUID: "uuid2", Type: nvidia.NvidiaGPUDevice, Usedcores: 30, Usedmem: 1000}, + {Idx: 0, UUID: "uuid1", Type: nvidia.NvidiaGPUDevice, Usedcores: 30, Usedmem: 1000}, + }, + // App container allocated 1 GPU from its own fresh pool + { + {Idx: 0, UUID: "uuid2", Type: nvidia.NvidiaGPUDevice, Usedcores: 30, Usedmem: 1000}, + }, + }, + }, + Score: 0, + }, + }, + }, + err: nil, + }, + }, + { + name: "init container requests more cores than any single device has (should filter node)", + args: struct { + nodes *map[string]*NodeUsage + nums device.PodDeviceRequests + annos map[string]string + task *corev1.Pod + }{ + nodes: &map[string]*NodeUsage{ + "node1": { + Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}, + Devices: policy.DeviceUsageList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + DeviceLists: []*policy.DeviceListsScore{ + { + Device: &device.DeviceUsage{ + ID: "uuid1", + Index: 0, + Used: 0, + Count: 10, + Usedmem: 0, + Totalmem: 8000, + Totalcore: 100, + Usedcores: 80, // only 20 cores free + Numa: 0, + Type: nvidia.NvidiaGPUDevice, + Health: true, + }, + }, + { + Device: &device.DeviceUsage{ + ID: "uuid2", + Index: 1, + Used: 0, + Count: 10, + Usedmem: 0, + Totalmem: 8000, + Totalcore: 100, + Usedcores: 80, // same, only 20 cores free + Numa: 0, + Type: nvidia.NvidiaGPUDevice, + Health: true, + }, + }, + }, + }, + }, + }, + nums: device.PodDeviceRequests{ + // Index 0: InitContainer requests 30 cores (more than any single device's 20 free) + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 1, + Type: nvidia.NvidiaGPUDevice, + Memreq: 1000, + Coresreq: 30, + }, + }, + // Index 1: App container requests 10 cores (fits, but init should block the node) + { + nvidia.NvidiaGPUDevice: device.ContainerDeviceRequest{ + Nums: 1, + Type: nvidia.NvidiaGPUDevice, + Memreq: 1000, + Coresreq: 10, + }, + }, + }, + annos: make(map[string]string), + task: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "test-init-large-cores"}, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-core-heavy", + Image: "busybox", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(30, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(1000, resource.BinarySI), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-light", + Image: "busybox", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "hami.io/gpu": *resource.NewQuantity(1, resource.BinarySI), + "hami.io/gpucores": *resource.NewQuantity(10, resource.BinarySI), + "hami.io/gpumem": *resource.NewQuantity(1000, resource.BinarySI), + }, + }, + }, + }, + }, + }, + }, + wants: struct { + want *policy.NodeScoreList + failedNodes map[string]string + err error + }{ + want: &policy.NodeScoreList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + NodeList: []*policy.NodeScore{}, + }, + failedNodes: map[string]string{ + "node1": "2/2 CardInsufficientCore", + }, + err: nil, + }, + }, { name: "two node per node having one device one pod two container use one device", args: struct { diff --git a/pkg/scheduler/webhook.go b/pkg/scheduler/webhook.go index 3f749e8f46..5c3a368f9e 100644 --- a/pkg/scheduler/webhook.go +++ b/pkg/scheduler/webhook.go @@ -71,6 +71,27 @@ func (h *webhook) Handle(_ context.Context, req admission.Request) admission.Res klog.V(5).Infof(template, pod.Namespace, pod.Name, pod.UID) privilegedName, hasPrivileged := privilegedContainerName(pod) hasResource := false + + // 1. Process InitContainers + for idx, ctr := range pod.Spec.InitContainers { + c := &pod.Spec.InitContainers[idx] + if ctr.SecurityContext != nil { + if ctr.SecurityContext.Privileged != nil && *ctr.SecurityContext.Privileged { + klog.Warningf(template+" - Denying admission as init container %s is privileged", pod.Namespace, pod.Name, pod.UID, c.Name) + continue + } + } + for _, val := range device.GetDevices() { + found, err := val.MutateAdmission(c, pod) + if err != nil { + klog.Errorf("validating pod failed:%s", err.Error()) + return admission.Errored(http.StatusInternalServerError, err) + } + hasResource = hasResource || found + } + } + + // 2. Process Regular Containers (Keep your existing loop here) for idx := range pod.Spec.Containers { c := &pod.Spec.Containers[idx] for _, val := range device.GetDevices() { @@ -130,7 +151,6 @@ func isPrivilegedContainer(ctr *corev1.Container) bool { func fitResourceQuota(pod *corev1.Pod) bool { for deviceName, dev := range device.GetDevices() { - // Only supports NVIDIA if deviceName != nvidia.NvidiaGPUDevice { continue } @@ -139,8 +159,7 @@ func fitResourceQuota(pod *corev1.Pod) bool { resourceName := corev1.ResourceName(resourceNames.ResourceCountName) memResourceName := corev1.ResourceName(resourceNames.ResourceMemoryName) coreResourceName := corev1.ResourceName(resourceNames.ResourceCoreName) - var memoryReq int64 = 0 - var coresReq int64 = 0 + getRequest := func(ctr *corev1.Container, resName corev1.ResourceName) (int64, bool) { v, ok := ctr.Resources.Limits[resName] if !ok { @@ -153,24 +172,49 @@ func fitResourceQuota(pod *corev1.Pod) bool { } return 0, false } + + var initMemoryReq, initCoresReq, initCountReq int64 + var appMemoryReq, appCoresReq, appCountReq int64 + + for _, ctr := range pod.Spec.InitContainers { + req, ok := getRequest(&ctr, resourceName) + if ok { + initCountReq = max(initCountReq, req) + if memReq, ok := getRequest(&ctr, memResourceName); ok { + initMemoryReq = max(initMemoryReq, memReq*req) + } + if coreReq, ok := getRequest(&ctr, coreResourceName); ok { + initCoresReq = max(initCoresReq, coreReq*req) + } + } + } + for _, ctr := range pod.Spec.Containers { req, ok := getRequest(&ctr, resourceName) if ok { + appCountReq += req if memReq, ok := getRequest(&ctr, memResourceName); ok { - memoryReq += memReq * req + appMemoryReq += memReq * req } if coreReq, ok := getRequest(&ctr, coreResourceName); ok { - coresReq += coreReq * req + appCoresReq += coreReq * req } } } + + memoryReq := max(appMemoryReq, initMemoryReq) + coresReq := max(appCoresReq, initCoresReq) + effectiveCountReq := max(appCountReq, initCountReq) + _ = effectiveCountReq + if memoryFactor > 1 { oriMemReq := memoryReq memoryReq = memoryReq * int64(memoryFactor) klog.V(5).Infof("Adjusting memory request for quota check: oriMemReq %d, memoryReq %d, factor %d", oriMemReq, memoryReq, memoryFactor) } + if !device.GetLocalCache().FitQuota(pod.Namespace, memoryReq, memoryFactor, coresReq, deviceName) { - klog.Infof(template+" - Denying admission", pod.Namespace, pod.Name, pod.UID) + klog.Infof("Namespace %s, Pod %s, UID %s - Denying admission", pod.Namespace, pod.Name, pod.UID) return false } } diff --git a/pkg/scheduler/webhook_test.go b/pkg/scheduler/webhook_test.go index a03e488310..f742caf555 100644 --- a/pkg/scheduler/webhook_test.go +++ b/pkg/scheduler/webhook_test.go @@ -18,7 +18,6 @@ package scheduler import ( "context" - "strings" "testing" admissionv1 "k8s.io/api/admission/v1" @@ -263,12 +262,15 @@ func TestFitResourceQuota(t *testing.T) { klog.Fatalf("Failed to initialize devices with config: %v", err) } - qm := device.NewQuotaManager() ns := "default" memName := "nvidia.com/gpumem" coreName := "nvidia.com/gpucores" - qm.Quotas[ns] = &device.DeviceQuota{ + cache := device.GetLocalCache() + if cache.Quotas == nil { + cache.Quotas = make(map[string]*device.DeviceQuota) + } + cache.Quotas[ns] = &device.DeviceQuota{ memName: &device.Quota{Used: 1000, Limit: 2000}, coreName: &device.Quota{Used: 200, Limit: 400}, } @@ -386,6 +388,89 @@ func TestFitResourceQuota(t *testing.T) { }, fit: true, }, + { + name: "InitContainers run sequentially: max init fits quota, but simple sum would exceed", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-pod-init-fit", + Namespace: "default", + }, + Spec: corev1.PodSpec{ + SchedulerName: "hami-scheduler", + // Two init containers each asking for 800. Max = 800. + InitContainers: []corev1.Container{ + { + Name: "init1", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("800"), + }, + }, + }, + { + Name: "init2", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("800"), + }, + }, + }, + }, + // App container asking for 500. Total effective requirement = max(800, 500) = 800. + // 800 is less than the available 1000 limit, so it should fit. + // If the quota manager wrongly sums everything (800+800+500=2100), this test will catch it by failing. + Containers: []corev1.Container{ + { + Name: "app1", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("500"), + }, + }, + }, + }, + }, + }, + fit: true, + }, + { + name: "InitContainer request exceeds quota directly", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-pod-init-fail", + Namespace: "default", + }, + Spec: corev1.PodSpec{ + SchedulerName: "hami-scheduler", + InitContainers: []corev1.Container{ + { + Name: "init-massive", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("1500"), // 1500 > 1000 available limit + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app1", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("100"), + }, + }, + }, + }, + }, + }, + fit: false, + }, { name: "request ascend", pod: &corev1.Pod{ @@ -644,170 +729,102 @@ func TestSchedulerNameEmptyNoOverwrite(t *testing.T) { } } -func TestPrivilegedContainerDenied(t *testing.T) { - prevSchedulerName := config.SchedulerName - prevDevicesMap := device.DevicesMap - prevDevicesToHandle := device.DevicesToHandle - t.Cleanup(func() { - config.SchedulerName = prevSchedulerName - device.DevicesMap = prevDevicesMap - device.DevicesToHandle = prevDevicesToHandle - }) - +func TestFitResourceQuota_InitContainerPeakSequence(t *testing.T) { config.SchedulerName = "hami-scheduler" sConfig := &config.Config{ NvidiaConfig: nvidia.NvidiaConfig{ - ResourceCountName: "hami.io/gpu", - ResourceMemoryName: "hami.io/gpumem", - ResourceMemoryPercentageName: "hami.io/gpumem-percentage", - ResourceCoreName: "hami.io/gpucores", + ResourceCountName: "nvidia.com/gpu", + ResourceMemoryName: "nvidia.com/gpumem", + ResourceMemoryPercentageName: "nvidia.com/gpumem-percentage", + ResourceCoreName: "nvidia.com/gpucores", DefaultMemory: 0, DefaultCores: 0, DefaultGPUNum: 1, + MemoryFactor: 1, }, } if err := config.InitDevicesWithConfig(sConfig); err != nil { t.Fatalf("Failed to initialize devices with config: %v", err) } - privileged := true - testCases := []struct { - name string - pod *corev1.Pod - allowed bool - }{ - { - name: "privileged container only without gpu", - pod: &corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{Name: "privileged-pod", Namespace: "default"}, - Spec: corev1.PodSpec{ - Containers: []corev1.Container{ - { - Name: "privileged", - SecurityContext: &corev1.SecurityContext{ - Privileged: &privileged, - }, - }, - }, - }, + ns := "default" + cache := device.GetLocalCache() + if cache.Quotas == nil { + cache.Quotas = make(map[string]*device.DeviceQuota) + } + + // initial quota: 30000 limit, 0 used + cache.Quotas[ns] = &device.DeviceQuota{ + "nvidia.com/gpumem": &device.Quota{Used: 0, Limit: 30000}, + "nvidia.com/gpucores": &device.Quota{Used: 0, Limit: 0}, + } + + makePod := func(name string, initMem, appMem int64) *corev1.Pod { + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: name, + Namespace: ns, }, - allowed: true, - }, - { - name: "privileged sidecar with gpu workload", - pod: &corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{Name: "mixed-pod", Namespace: "default"}, - Spec: corev1.PodSpec{ - Containers: []corev1.Container{ - { - Name: "privileged-sidecar", - SecurityContext: &corev1.SecurityContext{ - Privileged: &privileged, - }, - }, - { - Name: "gpu-workload", - Resources: corev1.ResourceRequirements{ - Limits: corev1.ResourceList{ - "hami.io/gpu": resource.MustParse("1"), - }, + Spec: corev1.PodSpec{ + SchedulerName: "hami-scheduler", + InitContainers: []corev1.Container{ + { + Name: "init", + Image: "busybox", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": *resource.NewQuantity(initMem, resource.DecimalSI), }, }, }, }, - }, - allowed: false, - }, - { - name: "privileged init container with gpu workload", - pod: &corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{Name: "init-privileged-pod", Namespace: "default"}, - Spec: corev1.PodSpec{ - InitContainers: []corev1.Container{ - { - Name: "privileged-init", - SecurityContext: &corev1.SecurityContext{ - Privileged: &privileged, - }, - }, - }, - Containers: []corev1.Container{ - { - Name: "gpu-workload", - Resources: corev1.ResourceRequirements{ - Limits: corev1.ResourceList{ - "hami.io/gpu": resource.MustParse("1"), - }, - }, - }, - }, - }, - }, - allowed: false, - }, - { - name: "privileged pod with different scheduler", - pod: &corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{Name: "other-scheduler-pod", Namespace: "default"}, - Spec: corev1.PodSpec{ - SchedulerName: "other-scheduler", - Containers: []corev1.Container{ - { - Name: "privileged", - SecurityContext: &corev1.SecurityContext{ - Privileged: &privileged, + Containers: []corev1.Container{ + { + Name: "app", + Image: "busybox", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": *resource.NewQuantity(appMem, resource.DecimalSI), }, }, }, }, }, - allowed: true, - }, + } } - wh, err := NewWebHook() - if err != nil { - t.Fatalf("Error creating WebHook: %v", err) + // Step 1: Pod1 (init 20000, app 10000) should be allowed + pod1 := makePod("pod1", 20000, 10000) + if !fitResourceQuota(pod1) { + t.Fatal("Step 1 failed: pod1 should be allowed (peak 20000 ≤ 30000)") } - scheme := runtime.NewScheme() - corev1.AddToScheme(scheme) - codec := serializer.NewCodecFactory(scheme).LegacyCodec(corev1.SchemeGroupVersion) + // Simulate pod1 scheduled → record its peak usage (20000) + if dq, ok := cache.Quotas[ns]; ok { + if q, ok := (*dq)["nvidia.com/gpumem"]; ok { + q.Used = 20000 + } else { + (*dq)["nvidia.com/gpumem"] = &device.Quota{Used: 20000, Limit: 30000} + } + } - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - podBytes, err := runtime.Encode(codec, tc.pod) - if err != nil { - t.Fatalf("Error encoding pod: %v", err) - } + // Step 2: Pod2 (same) must be DENIED + pod2 := makePod("pod2", 20000, 10000) + if fitResourceQuota(pod2) { + t.Fatal("Step 2 failed: pod2 should be denied (total used 20000 + request 20000 > 30000)") + } - req := admission.Request{ - AdmissionRequest: admissionv1.AdmissionRequest{ - UID: "test-uid", - Namespace: tc.pod.Namespace, - Name: tc.pod.Name, - Object: runtime.RawExtension{ - Raw: podBytes, - }, - }, - } + // Step 3: Pod1 init finished → usage drops to app only (10000) + if dq, ok := cache.Quotas[ns]; ok { + if q, ok := (*dq)["nvidia.com/gpumem"]; ok { + q.Used = 10000 + } + } - resp := wh.Handle(context.Background(), req) - if tc.allowed { - if !resp.Allowed { - t.Fatalf("Expected allowed response, but got denied: %+v", resp.Result) - } - return - } - if resp.Allowed { - t.Fatalf("Expected denied response for privileged pod, but got allowed with %d patches", len(resp.Patches)) - } - if len(resp.Patches) != 0 { - t.Fatalf("Expected no patches for privileged pod, got %d", len(resp.Patches)) - } - if resp.Result == nil || !strings.Contains(resp.Result.Message, "is privileged") { - t.Fatalf("Expected privilege denial message, got: %+v", resp.Result) - } - }) + // Now pod2 should be allowed + if !fitResourceQuota(pod2) { + t.Fatal("Step 3 failed: pod2 should be allowed after pod1 init finished (total 10000+20000=30000)") } }