diff --git a/charts/hami/templates/scheduler/clusterrole.yaml b/charts/hami/templates/scheduler/clusterrole.yaml index 666ce2d86..08f1d8c47 100644 --- a/charts/hami/templates/scheduler/clusterrole.yaml +++ b/charts/hami/templates/scheduler/clusterrole.yaml @@ -22,6 +22,9 @@ rules: - apiGroups: [""] resources: ["resourcequotas"] verbs: ["get", "list", "watch"] + - apiGroups: ["policy"] + resources: ["poddisruptionbudgets"] + verbs: ["get", "list", "watch"] - apiGroups: ["resource.k8s.io"] resources: ["deviceclasses", "resourceclaims", "resourceclaimtemplates", "resourceslices"] verbs: ["get", "list", "watch"] diff --git a/charts/hami/templates/scheduler/configmap.yaml b/charts/hami/templates/scheduler/configmap.yaml index 934f852fd..286bfe86f 100644 --- a/charts/hami/templates/scheduler/configmap.yaml +++ b/charts/hami/templates/scheduler/configmap.yaml @@ -36,6 +36,7 @@ data: {{- end }} filterVerb: filter bindVerb: bind + preemptVerb: preempt nodeCacheCapable: true weight: 1 httpTimeout: 30s @@ -61,6 +62,7 @@ data: {{- end }} "filterVerb": "filter", "bindVerb": "bind", + "preemptVerb": "preempt", "weight": 1, "nodeCacheCapable": true, "httpTimeout": 30000000000, diff --git a/cmd/scheduler/main.go b/cmd/scheduler/main.go index d90dd272a..5c3a7b9a0 100644 --- a/cmd/scheduler/main.go +++ b/cmd/scheduler/main.go @@ -33,6 +33,7 @@ import ( "github.com/Project-HAMi/HAMi/pkg/device" "github.com/Project-HAMi/HAMi/pkg/scheduler" "github.com/Project-HAMi/HAMi/pkg/scheduler/config" + "github.com/Project-HAMi/HAMi/pkg/scheduler/preempt" "github.com/Project-HAMi/HAMi/pkg/scheduler/routes" "github.com/Project-HAMi/HAMi/pkg/util" "github.com/Project-HAMi/HAMi/pkg/util/client" @@ -45,6 +46,7 @@ import ( var ( sher *scheduler.Scheduler + vgpu *preempt.VgpuPreempt tlsKeyFile string tlsCertFile string enableProfiling bool @@ -139,6 +141,10 @@ func start() error { return err } defer sher.Stop() + vgpu, err = preempt.New(sher.GetInformerFactory(), sher.GetEventRecorder()) + if err != nil { + return fmt.Errorf("failed to initialize vgpu preemptor: %w", err) + } // start monitor metrics go initMetrics(config.MetricsBindAddress, legacyMetrics) @@ -150,8 +156,9 @@ func start() error { router.POST("/webhook", routes.WebHookRoute()) router.GET("/healthz", routes.HealthzRoute()) router.GET("/readyz", routes.ReadyzRoute(sher)) - klog.Info("listen on ", config.HTTPBind) + router.POST("/preempt", routes.PreemptRoute(vgpu)) + klog.Info("listen on ", config.HTTPBind) if enableProfiling { injectProfilingRoute(router) klog.Infof("Profiling enabled, visit %s/debug/pprof/ to view profiles", config.HTTPBind) diff --git a/pkg/scheduler/preempt/preempt.go b/pkg/scheduler/preempt/preempt.go new file mode 100644 index 000000000..d246202f1 --- /dev/null +++ b/pkg/scheduler/preempt/preempt.go @@ -0,0 +1,963 @@ +/* +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 preempt + +import ( + "context" + "errors" + "fmt" + "maps" + "sort" + "strings" + "sync/atomic" + + corev1 "k8s.io/api/core/v1" + policyv1 "k8s.io/api/policy/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/informers" + corelisters "k8s.io/client-go/listers/core/v1" + policylisters "k8s.io/client-go/listers/policy/v1" + "k8s.io/client-go/tools/cache" + "k8s.io/client-go/tools/record" + "k8s.io/klog/v2" + extenderv1 "k8s.io/kube-scheduler/extender/v1" + + "github.com/Project-HAMi/HAMi/pkg/device" +) + +const Name = "PreemptPredicate" + +// criticalPriorityThreshold is the numeric value of the system-cluster-critical +// PriorityClass. Any pod whose Spec.Priority is at or above this value is +// treated as a critical system pod and is protected from preemption. +const criticalPriorityThreshold = int32(2000000000) + +// Sentinel error for distinguishing "pod has no native GPU request" from "insufficient native GPUs". +var errNoNativeGPURequested = errors.New("no native whole GPU resources requested") + +// VgpuPreempt refines victim selection proposed by kube-scheduler. +type VgpuPreempt struct { + nodeLister corelisters.NodeLister + podLister corelisters.PodLister + pdbLister policylisters.PodDisruptionBudgetLister + recorder record.EventRecorder + hasSyncFunc func(ctx context.Context) bool +} + +// dummyPDBLister implements PodDisruptionBudgetLister without requiring the policy API. +type dummyPDBLister struct{} + +func (d *dummyPDBLister) List(selector labels.Selector) ([]*policyv1.PodDisruptionBudget, error) { + return nil, nil +} + +func (d *dummyPDBLister) GetPodPodDisruptionBudgets(pod *corev1.Pod) ([]*policyv1.PodDisruptionBudget, error) { + return nil, nil +} + +func (d *dummyPDBLister) PodDisruptionBudgets(namespace string) policylisters.PodDisruptionBudgetNamespaceLister { + return &dummyPDBNamespaceLister{} +} + +type dummyPDBNamespaceLister struct{} + +func (d *dummyPDBNamespaceLister) List(selector labels.Selector) ([]*policyv1.PodDisruptionBudget, error) { + return nil, nil +} + +func (d *dummyPDBNamespaceLister) Get(name string) (*policyv1.PodDisruptionBudget, error) { + return nil, nil +} + +// New creates a new VgpuPreempt plugin. +func New( + factory informers.SharedInformerFactory, + recorder record.EventRecorder, +) (*VgpuPreempt, error) { + if factory == nil { + return nil, fmt.Errorf("informerFactory cannot be nil") + } + + podInformer := factory.Core().V1().Pods().Informer() + nodeInformer := factory.Core().V1().Nodes().Informer() + + nodeLister := factory.Core().V1().Nodes().Lister() + podLister := factory.Core().V1().Pods().Lister() + + var pdbInformer cache.SharedIndexInformer + var pdbLister policylisters.PodDisruptionBudgetLister + + // Access policy/v1 only if available. On older clusters the factory may return nil. + if policyV1 := factory.Policy().V1(); policyV1 != nil { + pdbInformer = policyV1.PodDisruptionBudgets().Informer() + pdbLister = policyV1.PodDisruptionBudgets().Lister() + } else { + klog.Warning("policy/v1 API not available; PDB-aware preemption is disabled") + // Use a dummy lister so violatesPDB() always returns false. + pdbLister = &dummyPDBLister{} + } + + hasSyncFunc := func(ctx context.Context) bool { + syncs := []cache.InformerSynced{ + podInformer.HasSynced, + nodeInformer.HasSynced, + } + if pdbInformer != nil { + syncs = append(syncs, pdbInformer.HasSynced) + } + return cache.WaitForCacheSync(ctx.Done(), syncs...) + } + + return &VgpuPreempt{ + nodeLister: nodeLister, + podLister: podLister, + pdbLister: pdbLister, + recorder: recorder, + hasSyncFunc: hasSyncFunc, + }, nil +} + +func (p *VgpuPreempt) Name() string { return Name } + +func (p *VgpuPreempt) IsReady(ctx context.Context) bool { return p.hasSyncFunc(ctx) } + +// getDevicesSnapshot returns a thread-safe snapshot of registered device handlers. +func (p *VgpuPreempt) getDevicesSnapshot() map[string]device.Devices { + // GetDevices() returns the live global map. To avoid data races, + // create a snapshot at the start of preemption evaluation. + source := device.GetDevices() + snapshot := make(map[string]device.Devices, len(source)) + maps.Copy(snapshot, source) + return snapshot +} + +var ( + preemptionAttempts atomic.Int64 + preemptionSucceeded atomic.Int64 + preemptionDropped atomic.Int64 + victimsFittedWOEvict atomic.Int64 +) + +// IncPreemptionCounter increments the appropriate preemption counter in a thread-safe manner. +func IncPreemptionCounter(counterType string) { + switch counterType { + case "attempts": + preemptionAttempts.Add(1) + case "succeeded": + preemptionSucceeded.Add(1) + case "dropped": + preemptionDropped.Add(1) + case "fitted_no_evict": + victimsFittedWOEvict.Add(1) + } +} + +// getNativeWholeGPUResources returns resource names that represent entire physical GPUs. +// Dynamically built from common patterns HAMi supports. +func getNativeWholeGPUResources() map[string]bool { + return map[string]bool{ + // NVIDIA + "nvidia.com/gpu": true, + // AMD + "amd.com/gpu": true, + // Intel + "intel.com/gpu": true, + // Huawei Ascend (all variants) + "huawei.com/Ascend310": true, + "huawei.com/Ascend310B": true, + "huawei.com/Ascend310P": true, + "huawei.com/Ascend910": true, + "huawei.com/Ascend910B": true, + "huawei.com/Ascend910C": true, + // Cambricon MLU + "cambricon.com/mlu": true, + // Hygon DCU + "hygon.com/dcu": true, + // Iluvatar + "iluvatar.ai/gpu": true, + "metax-tech.com/gpu": true, + "mthreads.com/gpu": true, + "kunlun.com/gpu": true, + "aws.amazon.com/neuron": true, + "vastai.com/gpu": true, + "enflame.com/gpu": true, + } +} + +// GetPreemptionMetrics returns current preemption metrics for monitoring. +func GetPreemptionMetrics() map[string]int64 { + return map[string]int64{ + "preemption_attempts": preemptionAttempts.Load(), + "preemption_succeeded": preemptionSucceeded.Load(), + "preemption_dropped": preemptionDropped.Load(), + "victims_fitted_no_evict": victimsFittedWOEvict.Load(), + } +} + +func (p *VgpuPreempt) Preempt( + ctx context.Context, + args extenderv1.ExtenderPreemptionArgs, +) *extenderv1.ExtenderPreemptionResult { + result := &extenderv1.ExtenderPreemptionResult{ + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{}, + } + + IncPreemptionCounter("attempts") + + pod := args.Pod + if pod == nil { + klog.V(4).InfoS("Preempt called with nil pod, passing input through") + return passthrough(args) + } + + if !isVGPUResourcePod(pod) { + klog.V(5).InfoS("Preempt: pod is not a GPU/device pod, passing input through", + "pod", klog.KObj(pod)) + return passthrough(args) + } + + victimsMap, err := p.resolveVictimsMap(args) + if err != nil { + klog.ErrorS(err, "Preempt: failed to resolve victims, cannot preempt") + return &extenderv1.ExtenderPreemptionResult{} + } + if len(victimsMap) == 0 { + return result + } + + // Single full cluster scan to avoid O(N²) lookups. + allPods, err := p.podLister.List(labels.Everything()) + if err != nil { + klog.ErrorS(err, "PodLister list pods failed in preempt, cannot evaluate preemption") + return &extenderv1.ExtenderPreemptionResult{} + } + + podsByNode := make(map[string][]*corev1.Pod) + for _, pObj := range allPods { + if pObj.Spec.NodeName != "" && isVGPUResourcePod(pObj) { + podsByNode[pObj.Spec.NodeName] = append(podsByNode[pObj.Spec.NodeName], pObj) + } + } + + // Capture device snapshot once per preemption evaluation for thread safety. + devMap := p.getDevicesSnapshot() + + for nodeName, victims := range victimsMap { + if victims == nil { + continue + } + + nodeVGPUPods := podsByNode[nodeName] + + refined, pdbViolations, ok := p.refineForNode(pod, nodeName, victims, nodeVGPUPods, devMap) + if !ok { + klog.V(3).InfoS("Preempt: node cannot fit pod even after preemption, dropping", + "pod", klog.KObj(pod), "node", nodeName) + IncPreemptionCounter("dropped") + continue + } + + // Always include the node if ok is true, even if refined is empty. + meta := &extenderv1.MetaVictims{ + Pods: make([]*extenderv1.MetaPod, 0, len(refined)), + NumPDBViolations: pdbViolations, + } + for _, vp := range refined { + meta.Pods = append(meta.Pods, &extenderv1.MetaPod{UID: string(vp.UID)}) + } + result.NodeNameToMetaVictims[nodeName] = meta + IncPreemptionCounter("succeeded") + + // Emit event on the preemptor to aid observability. + if len(refined) > 0 { + p.recorder.Eventf(pod, corev1.EventTypeNormal, "Preempting", + "Preempting %d pod(s) on node %s to accommodate %s/%s", + len(refined), nodeName, pod.Namespace, pod.Name) + } else { + IncPreemptionCounter("fitted_no_evict") + p.recorder.Eventf(pod, corev1.EventTypeNormal, "SchedulableFitNoEviction", + "Pod fits on node %s without evicting any pods", + nodeName) + } + } + return result +} + +// resolveVictimsMap converts MetaVictims format to Victims format by looking up pod UIDs. +func (p *VgpuPreempt) resolveVictimsMap( + args extenderv1.ExtenderPreemptionArgs, +) (map[string]*extenderv1.Victims, error) { + if len(args.NodeNameToVictims) > 0 { + return args.NodeNameToVictims, nil + } + if len(args.NodeNameToMetaVictims) == 0 { + return nil, nil + } + + allPods, err := p.podLister.List(labels.Everything()) + if err != nil { + return nil, fmt.Errorf("failed to list pods: %w", err) + } + byUID := make(map[types.UID]*corev1.Pod, len(allPods)) + for _, pod := range allPods { + byUID[pod.UID] = pod + } + + resolved := make(map[string]*extenderv1.Victims, len(args.NodeNameToMetaVictims)) + for nodeName, meta := range args.NodeNameToMetaVictims { + if meta == nil { + continue + } + pods := make([]*corev1.Pod, 0, len(meta.Pods)) + for _, mp := range meta.Pods { + if mp == nil { + continue + } + pod, ok := byUID[types.UID(mp.UID)] + if !ok { + klog.V(3).InfoS("Preempt: victim UID not found in cache, skipping", + "uid", mp.UID, "node", nodeName) + continue + } + pods = append(pods, pod) + } + // Always include node, even with empty pods, so refinement can check fit without victims. + resolved[nodeName] = &extenderv1.Victims{ + Pods: pods, + NumPDBViolations: meta.NumPDBViolations, + } + } + return resolved, nil +} + +// refineForNode evaluates and refines victim selection for a single node. +func (p *VgpuPreempt) refineForNode( + preemptor *corev1.Pod, + nodeName string, + victims *extenderv1.Victims, + nodeVGPUPods []*corev1.Pod, + devMap map[string]device.Devices, +) ([]*corev1.Pod, int64, bool) { + node, err := p.nodeLister.Get(nodeName) + if err != nil { + klog.V(3).ErrorS(err, "Preempt: get node failed", "node", nodeName) + return nil, 0, false + } + + // Track whether the scheduler originally proposed any victims for this node. + hadProposedVictims := len(victims.Pods) > 0 + + keep := make([]*corev1.Pod, 0, len(victims.Pods)) + excluded := make(map[types.UID]struct{}, len(victims.Pods)) + preemptorPriority := getPodPriority(preemptor) + + for _, v := range victims.Pods { + if v == nil { + continue + } + if isProtectedFromPreemption(v) { + klog.V(4).InfoS("Preempt: refusing to evict protected pod", + "pod", klog.KObj(v), "node", nodeName) + continue + } + // Reject victims that have equal or higher priority than the preemptor. + if getPodPriority(v) >= preemptorPriority { + klog.V(4).InfoS("Preempt: refusing to evict pod with equal or higher priority", + "pod", klog.KObj(v), "node", nodeName) + continue + } + keep = append(keep, v) + excluded[v.UID] = struct{}{} + } + keptFromInput := len(keep) + + if p.podFitsAfterPreemption(preemptor, *node, nodeVGPUPods, excluded, devMap) { + return keep, pdbViolationsUpperBound(victims.NumPDBViolations, keptFromInput), true + } + + // Only search for additional victims when the scheduler originally proposed at least one victim. + if !hadProposedVictims { + klog.V(4).InfoS("Preempt: pod does not fit and scheduler proposed no victims; dropping node", + "pod", klog.KObj(preemptor), "node", nodeName) + return nil, 0, false + } + + additional := p.findAdditionalVictims(preemptor, node, nodeVGPUPods, excluded) + for _, cand := range additional { + excluded[cand.UID] = struct{}{} + keep = append(keep, cand) + if p.podFitsAfterPreemption(preemptor, *node, nodeVGPUPods, excluded, devMap) { + // added := len(keep) - keptFromInput + return keep, pdbViolationsUpperBound(victims.NumPDBViolations, len(keep)), true + } + } + return nil, 0, false +} + +// podFitsAfterPreemption simulates resource allocation after evicting the excluded pods. +func (p *VgpuPreempt) podFitsAfterPreemption( + preemptor *corev1.Pod, + node corev1.Node, + nodeVGPUPods []*corev1.Pod, + excluded map[types.UID]struct{}, + devMap map[string]device.Devices, +) bool { + nativeGPUResources := getNativeWholeGPUResources() + + ni := &device.NodeInfo{ + ID: node.Name, + Node: &node, + Devices: make(map[string][]device.DeviceInfo, len(devMap)), + } + + // Determine which device types the preemptor actually requires. + preemptorReqs := device.Resourcereqs(preemptor) + requestedDevTypes := make(map[string]bool) + for _, cReqs := range preemptorReqs { + for devType, req := range cReqs { + if req.Nums > 0 || req.Memreq > 0 || req.Coresreq > 0 { + requestedDevTypes[devType] = true + } + } + } + + // Also add native whole GPU requests to requestedDevTypes + for _, c := range append(preemptor.Spec.Containers, preemptor.Spec.InitContainers...) { + for rName := range c.Resources.Requests { + if nativeGPUResources[string(rName)] { + devType := extractDeviceTypeFromResourceName(string(rName), devMap) + if devType != "" { + requestedDevTypes[devType] = true + } + } + } + } + + // Populate device state from node. + for devType, dev := range devMap { + infos, err := dev.GetNodeDevices(node) + if err != nil { + if requestedDevTypes[devType] { + klog.V(3).ErrorS(err, "Failed to get node devices for required type", + "node", node.Name, "type", devType) + return false + } + klog.V(5).InfoS("Ignoring device plugin error for unrequested type", + "node", node.Name, "type", devType) + continue + } + ni.Devices[devType] = make([]device.DeviceInfo, len(infos)) + for i, info := range infos { + if info != nil { + ni.Devices[devType][i] = *info + } + } + } + + state := make(map[string][]*device.DeviceUsage) + for devType, devInfos := range ni.Devices { + usages := make([]*device.DeviceUsage, len(devInfos)) + for i, info := range devInfos { + usages[i] = &device.DeviceUsage{ + ID: info.ID, + Index: info.Index, + Count: info.Count, + Totalmem: info.Devmem, + Totalcore: info.Devcore, + Type: info.Type, + Mode: info.Mode, + Health: info.Health, + } + } + state[devType] = usages + } + + // Group co-located pods to process them in a deterministic, optimal order. + var annotatedPods []*corev1.Pod + var nativePods []*corev1.Pod + var vgpuPods []*corev1.Pod + + hasNativeGPU := func(pod *corev1.Pod) bool { + for _, c := range append(pod.Spec.Containers, pod.Spec.InitContainers...) { + for rName := range c.Resources.Requests { + if nativeGPUResources[string(rName)] { + return true + } + } + } + return false + } + + for _, vPod := range nodeVGPUPods { + if _, skip := excluded[vPod.UID]; skip { + continue + } + podDevs, err := device.DecodePodDevices(device.SupportDevices, vPod.Annotations) + if err == nil && len(podDevs) > 0 { + annotatedPods = append(annotatedPods, vPod) + } else if hasNativeGPU(vPod) { + nativePods = append(nativePods, vPod) + } else { + vgpuPods = append(vgpuPods, vPod) + } + } + + // 1. Process annotated pods first (fixed UUIDs) + for _, vPod := range annotatedPods { + podDevs, _ := device.DecodePodDevices(device.SupportDevices, vPod.Annotations) + for devType, pds := range podDevs { + usages := state[devType] + if usages == nil { + continue + } + for _, ctrDevs := range pds { + for _, cd := range ctrDevs { + var du *device.DeviceUsage + for _, u := range usages { + if u.ID == cd.UUID { + du = u + break + } + } + if du == nil { + continue + } + du.Used++ + du.Usedmem += cd.Usedmem + du.Usedcores += cd.Usedcores + } + } + } + } + + // 2. Process native GPU pods (require whole unused devices) + for _, vPod := range nativePods { + if err := p.accountNativeWholeGPUs(state, vPod, devMap); err != nil { + klog.V(3).ErrorS(err, "Failed to account native pod resource usage", "pod", klog.KObj(vPod)) + return false + } + } + + // 3. Process unannotated vGPU pods (can share remaining capacity) + for _, vPod := range vgpuPods { + if err := p.accountVGPURequests(state, vPod, devMap); err != nil { + klog.V(3).ErrorS(err, "Failed to account pod resource usage", "pod", klog.KObj(vPod)) + return false + } + } + + // Evaluate the preemptor's native whole-GPU requests against the remaining state. + errNative := p.accountNativeWholeGPUs(state, preemptor, devMap) + if errNative != nil && !errors.Is(errNative, errNoNativeGPURequested) { + return false + } + + if len(preemptorReqs) == 0 { + return true + } + + // Attempt to allocate for each container request. + for _, cReqs := range preemptorReqs { + if len(cReqs) == 0 { + continue + } + for devType, cReq := range cReqs { + if cReq.Nums == 0 { + continue + } + if state[devType] == nil { + return false + } + curCopy := deepCopyUsageSlice(state[devType]) + + var alloc device.PodDevices + fitOK, _, _ := devMap[devType].Fit(curCopy, cReq, preemptor, ni, &alloc) + if !fitOK { + return false + } + state[devType] = curCopy + } + } + return true +} + +// accountNativeWholeGPUs reserves full GPU devices for pods that request native whole-GPU resources. +func (p *VgpuPreempt) accountNativeWholeGPUs( + state map[string][]*device.DeviceUsage, + pod *corev1.Pod, + devMap map[string]device.Devices, +) error { + wholeGPUResources := getNativeWholeGPUResources() + wholeRequested := map[string]int64{} + + for _, c := range append(pod.Spec.Containers, pod.Spec.InitContainers...) { + for rName, rQuant := range c.Resources.Requests { + if !wholeGPUResources[string(rName)] { + continue + } + devType := extractDeviceTypeFromResourceName(string(rName), devMap) + if devType == "" { + continue + } + wholeRequested[devType] += rQuant.Value() + } + } + + if len(wholeRequested) == 0 { + return fmt.Errorf("%w: %s/%s", errNoNativeGPURequested, pod.Namespace, pod.Name) + } + + for devType, count := range wholeRequested { + usages := state[devType] + if usages == nil { + return fmt.Errorf("no device state for type %q", devType) + } + used := int64(0) + for _, du := range usages { + if du.Used > 0 { + continue + } + du.Used = du.Count + du.Usedmem = du.Totalmem + du.Usedcores = du.Totalcore + used++ + if used >= count { + break + } + } + if used < count { + return fmt.Errorf("insufficient free devices for %d whole GPU(s) of type %q", count, devType) + } + } + return nil +} + +// accountVGPURequests accounts for pods with HAMi resource requests but without device annotations. +func (p *VgpuPreempt) accountVGPURequests( + state map[string][]*device.DeviceUsage, + pod *corev1.Pod, + devMap map[string]device.Devices, +) error { + reqs := device.Resourcereqs(pod) + for _, cReqs := range reqs { + for devType, req := range cReqs { + usages := state[devType] + if usages == nil { + if req.Nums > 0 || req.Memreq > 0 || req.Coresreq > 0 { + return fmt.Errorf("no device state available for type %q", devType) + } + continue + } + satisfied := false + for _, du := range usages { + if du.Count-du.Used >= 1 && + du.Totalmem-du.Usedmem >= req.Memreq && + du.Totalcore-du.Usedcores >= req.Coresreq { + du.Used++ + du.Usedmem += req.Memreq + du.Usedcores += req.Coresreq + satisfied = true + break + } + } + if !satisfied { + return fmt.Errorf( + "no device could satisfy vGPU request: pod=%s/%s, devType=%s, request=%+v", + pod.Namespace, pod.Name, devType, req, + ) + } + } + } + return nil +} + +// deepCopyUsageSlice creates a deep copy of the slice with cloned DeviceUsage elements. +func deepCopyUsageSlice(src []*device.DeviceUsage) []*device.DeviceUsage { + if src == nil { + return nil + } + dst := make([]*device.DeviceUsage, len(src)) + for i, u := range src { + if u == nil { + continue + } + // Manual deep copy of the struct + copy := *u + dst[i] = © + } + return dst +} + +// findAdditionalVictims searches for lower-priority pods that are not PDB-protected. +func (p *VgpuPreempt) findAdditionalVictims( + preemptor *corev1.Pod, + node *corev1.Node, + nodeVGPUPods []*corev1.Pod, + excluded map[types.UID]struct{}, +) []*corev1.Pod { + preemptorPriority := getPodPriority(preemptor) + candidates := make([]*corev1.Pod, 0) + + for _, candidate := range nodeVGPUPods { + if candidate.UID == preemptor.UID { + continue + } + if _, dup := excluded[candidate.UID]; dup { + continue + } + if getPodPriority(candidate) >= preemptorPriority { + continue + } + if isProtectedFromPreemption(candidate) { + continue + } + + if p.violatesPDB(candidate) { + klog.V(4).InfoS("Preempt: candidate would violate PDB, skipping", + "pod", klog.KObj(candidate), "node", node.Name) + continue + } + + candidates = append(candidates, candidate) + } + + sortVictimsByPreference(candidates) + return candidates +} + +// violatesPDB checks if evicting the pod would violate any applicable PodDisruptionBudget. +func (p *VgpuPreempt) violatesPDB(pod *corev1.Pod) bool { + if p.pdbLister == nil { + return false + } + + if pod == nil || pod.Namespace == "" { + return false + } + + pdbs, err := p.pdbLister.PodDisruptionBudgets(pod.Namespace).List(labels.Everything()) + if err != nil { + klog.V(4).ErrorS(err, "Failed to list PDBs; assuming no PDB violation", "pod", klog.KObj(pod)) + return false + } + + for _, pdb := range pdbs { + if pdb == nil { + continue + } + + selector, err := metav1.LabelSelectorAsSelector(pdb.Spec.Selector) + if err != nil { + klog.V(4).ErrorS(err, "Invalid PDB selector", "pdb", klog.KObj(pdb)) + continue + } + + if !selector.Matches(labels.Set(pod.Labels)) { + continue + } + + if pdb.Status.DisruptionsAllowed == 0 { + klog.V(4).InfoS("Preempt: pod matches PDB with zero disruptions allowed", + "pod", klog.KObj(pod), "pdb", klog.KObj(pdb)) + return true + } + } + + return false +} + +func extractDeviceTypeFromResourceName(rName string, devMap map[string]device.Devices) string { + resourceStr := string(rName) + + for registryKey, dev := range devMap { + resourceNames := dev.GetResourceNames() + + // Check if the requested resource matches any of the three field values + // that this device declares it supports + if resourceStr == resourceNames.ResourceCountName || + resourceStr == resourceNames.ResourceMemoryName || + resourceStr == resourceNames.ResourceCoreName { + return registryKey + } + } + + return "" +} + +func isVGPUResourcePod(pod *corev1.Pod) bool { + if pod == nil { + return false + } + // Check for annotation-based device allocation + if _, ok := pod.Annotations["hami.io/container-devices"]; ok { + return true + } + + vendorPrefixes := []string{ + "hami.io/", + "nvidia.com/", + "amd.com/", + "intel.com/", + "huawei.com/", + "cambricon.com/", + "hygon.com/", + "iluvatar.ai/", + "metax-tech.com/", + "mthreads.com/", + "kunlun.com/", + "aws.amazon.com/", + "vastai.com/", + "enflame.com/", + } + + checkContainers := func(containers []corev1.Container) bool { + for _, c := range containers { + for rName := range c.Resources.Requests { + sName := string(rName) + for _, prefix := range vendorPrefixes { + if strings.HasPrefix(sName, prefix) { + return true + } + } + } + } + return false + } + return checkContainers(pod.Spec.Containers) || checkContainers(pod.Spec.InitContainers) +} + +// isProtectedFromPreemption returns true if the pod cannot be preempted. +func isProtectedFromPreemption(pod *corev1.Pod) bool { + if pod == nil { + return true + } + if pod.DeletionTimestamp != nil { + return true + } + if pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed { + return true + } + if isCriticalPod(pod) { + return true + } + if owner := metav1.GetControllerOf(pod); owner != nil && owner.Kind == "DaemonSet" { + return true + } + return false +} + +// isCriticalPod checks if a pod has critical priority or is in kube-system. +func isCriticalPod(pod *corev1.Pod) bool { + if pod.Spec.PriorityClassName == "system-node-critical" || + pod.Spec.PriorityClassName == "system-cluster-critical" { + return true + } + // Guard by numeric value so pods with high-priority values (>= system-cluster-critical) + // are always protected, regardless of whether the PriorityClass name is set. + if pod.Spec.Priority != nil && *pod.Spec.Priority >= criticalPriorityThreshold { + return true + } + return pod.Namespace == "kube-system" +} + +// getPodPriority returns the pod's priority value. +func getPodPriority(pod *corev1.Pod) int32 { + if pod.Spec.Priority != nil { + return *pod.Spec.Priority + } + return 0 +} + +// sortVictimsByPreference sorts victims by priority (ascending), then by creation time +// (newer first), then by total resource request (descending). +func sortVictimsByPreference(pods []*corev1.Pod) { + sort.SliceStable(pods, func(i, j int) bool { + a, b := pods[i], pods[j] + ap, bp := getPodPriority(a), getPodPriority(b) + if ap != bp { + return ap < bp + } + if !a.CreationTimestamp.Equal(&b.CreationTimestamp) { + return a.CreationTimestamp.After(b.CreationTimestamp.Time) + } + aMem, aCores, aCount := getPodVGPURequest(a) + bMem, bCores, bCount := getPodVGPURequest(b) + return aMem+aCores+aCount > bMem+bCores+bCount + }) +} + +// getPodVGPURequest extracts total vGPU-related resource requests from a pod. +func getPodVGPURequest(pod *corev1.Pod) (mem, cores, count int64) { + for _, c := range append(pod.Spec.Containers, pod.Spec.InitContainers...) { + for rName, rQuant := range c.Resources.Requests { + sName := string(rName) + if strings.Contains(sName, "mem") || strings.Contains(sName, "memory") { + mem += rQuant.Value() + } else if strings.Contains(sName, "cores") { + cores += rQuant.Value() + } else if strings.Contains(sName, "vgpu") || strings.Contains(sName, "gpu") { + count += rQuant.Value() + } + } + } + if mem == 0 && cores == 0 && count == 0 { + if pd, err := device.DecodePodDevices(device.SupportDevices, pod.Annotations); err == nil { + for _, single := range pd { + for _, ctrDevs := range single { + for _, cd := range ctrDevs { + mem += int64(cd.Usedmem) + cores += int64(cd.Usedcores) + count++ + } + } + } + } + } + return +} + +// pdbViolationsUpperBound computes a conservative upper bound on PDB violations. +func pdbViolationsUpperBound(originalCount int64, keptLen int) int64 { + return min(originalCount, int64(keptLen)) +} + +// passthrough returns the original victim list in MetaVictims format. +func passthrough(args extenderv1.ExtenderPreemptionArgs) *extenderv1.ExtenderPreemptionResult { + out := &extenderv1.ExtenderPreemptionResult{ + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{}, + } + if len(args.NodeNameToMetaVictims) > 0 { + maps.Copy(out.NodeNameToMetaVictims, args.NodeNameToMetaVictims) + return out + } + for nodeName, v := range args.NodeNameToVictims { + if v == nil { + continue + } + meta := &extenderv1.MetaVictims{ + Pods: make([]*extenderv1.MetaPod, 0, len(v.Pods)), + NumPDBViolations: v.NumPDBViolations, + } + for _, pod := range v.Pods { + if pod == nil { + continue + } + meta.Pods = append(meta.Pods, &extenderv1.MetaPod{UID: string(pod.UID)}) + } + out.NodeNameToMetaVictims[nodeName] = meta + } + return out +} diff --git a/pkg/scheduler/preempt/preempt_test.go b/pkg/scheduler/preempt/preempt_test.go new file mode 100644 index 000000000..1929344ad --- /dev/null +++ b/pkg/scheduler/preempt/preempt_test.go @@ -0,0 +1,1073 @@ +/* +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 preempt + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + corev1 "k8s.io/api/core/v1" + policyv1 "k8s.io/api/policy/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/intstr" + "k8s.io/client-go/informers" + "k8s.io/client-go/kubernetes/fake" + "k8s.io/client-go/kubernetes/scheme" + "k8s.io/client-go/tools/record" + extenderv1 "k8s.io/kube-scheduler/extender/v1" + + "github.com/Project-HAMi/HAMi/pkg/device" +) + +func init() { + if device.DevicesMap == nil { + device.DevicesMap = make(map[string]device.Devices) + } +} + +type mockDevice struct { + devices []*device.DeviceInfo + failGetNodeDevices bool + fitReturnsTrue bool +} + +func (m *mockDevice) CommonWord() string { return "nvidia" } + +func (m *mockDevice) MutateAdmission(ctr *corev1.Container, pod *corev1.Pod) (bool, error) { + return false, nil +} + +func (m *mockDevice) CheckHealth(devType string, n *corev1.Node) (bool, bool) { return true, false } +func (m *mockDevice) NodeCleanUp(nn string) error { return nil } + +func (m *mockDevice) GetResourceNames() device.ResourceNames { + return device.ResourceNames{ + ResourceCountName: "hami.io/vgpu", + ResourceMemoryName: "hami.io/vgpu-memory", + ResourceCoreName: "hami.io/vgpu-cores", + } +} + +func (m *mockDevice) GetNodeDevices(n corev1.Node) ([]*device.DeviceInfo, error) { + if m.failGetNodeDevices { + return nil, errors.New("simulated hardware driver error") + } + return m.devices, nil +} + +func (m *mockDevice) LockNode(n *corev1.Node, p *corev1.Pod) error { return nil } +func (m *mockDevice) ReleaseNodeLock(n *corev1.Node, p *corev1.Pod) error { return nil } + +func (m *mockDevice) GenerateResourceRequests(ctr *corev1.Container) device.ContainerDeviceRequest { + req := device.ContainerDeviceRequest{} + if qty, ok := ctr.Resources.Requests["nvidia.com/gpu"]; ok { + req.Nums = int32(qty.Value()) + } + if qty, ok := ctr.Resources.Requests["hami.io/vgpu"]; ok { + req.Nums = int32(qty.Value()) + } + if qty, ok := ctr.Resources.Requests["hami.io/vgpu-memory"]; ok { + req.Memreq = int32(qty.Value()) + } + if qty, ok := ctr.Resources.Requests["hami.io/vgpu-cores"]; ok { + req.Coresreq = int32(qty.Value()) + } + return req +} + +func (m *mockDevice) PatchAnnotations(pod *corev1.Pod, annoinput *map[string]string, pd device.PodDevices) map[string]string { + return nil +} + +func (m *mockDevice) ScoreNode(node *corev1.Node, podDevices device.PodSingleDevice, previous []*device.DeviceUsage, policy string) float32 { + return 0 +} + +func (m *mockDevice) AddResourceUsage(pod *corev1.Pod, n *device.DeviceUsage, ctr *device.ContainerDevice) error { + return nil +} + +func (m *mockDevice) Fit( + usages []*device.DeviceUsage, + req device.ContainerDeviceRequest, + pod *corev1.Pod, + nodeInfo *device.NodeInfo, + allocated *device.PodDevices, +) (bool, map[string]device.ContainerDevices, string) { + if !m.fitReturnsTrue { + return false, nil, "insufficient mock device resources" + } + for _, du := range usages { + if du.Count-du.Used >= 1 && + du.Totalmem-du.Usedmem >= req.Memreq && + du.Totalcore-du.Usedcores >= req.Coresreq { + alloc := device.ContainerDevices{{ + UUID: du.ID, + Type: "nvidia", + Usedmem: req.Memreq, + Usedcores: req.Coresreq, + }} + return true, map[string]device.ContainerDevices{"compute-container": alloc}, "" + } + } + return false, nil, "no device with enough resources" +} + +func newMockDevice(count, mem, cores int32) *mockDevice { + return &mockDevice{ + fitReturnsTrue: true, + devices: []*device.DeviceInfo{ + { + ID: "mock-gpu-uuid-1", + Index: 0, + Count: count, + Devmem: mem, + Devcore: cores, + Type: "nvidia", + Health: true, + }, + }, + } +} + +func newDefaultMockDevice() *mockDevice { + return newMockDevice(10, 20000, 100) +} + +func registerMock(m *mockDevice) func() { + devs := device.GetDevices() + old, existed := devs["nvidia"] + devs["nvidia"] = m + return func() { + if existed { + devs["nvidia"] = old + } else { + delete(devs, "nvidia") + } + } +} + +func newTestNode(name string) *corev1.Node { + return &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{Name: name}, + Status: corev1.NodeStatus{ + Capacity: corev1.ResourceList{ + corev1.ResourceName("hami.io/vgpu"): resource.MustParse("2"), + corev1.ResourceName("hami.io/vgpu-memory"): resource.MustParse("20000"), + corev1.ResourceName("hami.io/vgpu-cores"): resource.MustParse("200"), + }, + Allocatable: corev1.ResourceList{ + corev1.ResourceName("hami.io/vgpu"): resource.MustParse("2"), + corev1.ResourceName("hami.io/vgpu-memory"): resource.MustParse("20000"), + corev1.ResourceName("hami.io/vgpu-cores"): resource.MustParse("200"), + }, + }, + } +} + +type podOpt func(*corev1.Pod) + +func withPriority(p int32) podOpt { return func(pod *corev1.Pod) { pod.Spec.Priority = &p } } +func withNodeName(node string) podOpt { return func(pod *corev1.Pod) { pod.Spec.NodeName = node } } +func withDeletionTimestamp() podOpt { + return func(pod *corev1.Pod) { + now := metav1.NewTime(time.Now()) + pod.DeletionTimestamp = &now + } +} +func withOwner(kind string) podOpt { + return func(pod *corev1.Pod) { + ctrl := true + pod.OwnerReferences = []metav1.OwnerReference{{ + Kind: kind, Name: "test-owner", UID: types.UID(uuid.NewString()), APIVersion: "apps/v1", Controller: &ctrl, + }} + } +} +func withNamespace(ns string) podOpt { return func(pod *corev1.Pod) { pod.Namespace = ns } } +func withLabels(labels map[string]string) podOpt { + return func(pod *corev1.Pod) { pod.Labels = labels } +} +func withAnnotations(annot map[string]string) podOpt { + return func(pod *corev1.Pod) { pod.Annotations = annot } +} + +type vgpuRequest struct{ vgpu, mem, cores int64 } + +func newVGPUPod(name string, req vgpuRequest, opts ...podOpt) *corev1.Pod { + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: "default", UID: types.UID(uuid.NewString())}, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{ + Name: "compute-container", + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceName("hami.io/vgpu"): resource.MustParse(fmt.Sprintf("%d", req.vgpu)), + corev1.ResourceName("hami.io/vgpu-memory"): resource.MustParse(fmt.Sprintf("%d", req.mem)), + corev1.ResourceName("hami.io/vgpu-cores"): resource.MustParse(fmt.Sprintf("%d", req.cores)), + }, + }, + }}, + }, + Status: corev1.PodStatus{Phase: corev1.PodRunning}, + } + for _, o := range opts { + o(pod) + } + return pod +} + +func newNativeGPUPod(name string, gpuCount int, opts ...podOpt) *corev1.Pod { + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: "default", UID: types.UID(uuid.NewString())}, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{ + Name: "gpu-container", + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceName("nvidia.com/gpu"): resource.MustParse(fmt.Sprintf("%d", gpuCount)), + }, + }, + }}, + }, + Status: corev1.PodStatus{Phase: corev1.PodRunning}, + } + for _, o := range opts { + o(pod) + } + return pod +} + +func newMultiContainerPod(name string, containers []corev1.Container, opts ...podOpt) *corev1.Pod { + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: "default", UID: types.UID(uuid.NewString())}, + Spec: corev1.PodSpec{Containers: containers}, + Status: corev1.PodStatus{Phase: corev1.PodRunning}, + } + for _, o := range opts { + o(pod) + } + return pod +} + +func newPDB(name, namespace string, selector *metav1.LabelSelector, minAvailable int32) *policyv1.PodDisruptionBudget { + return &policyv1.PodDisruptionBudget{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: namespace}, + Spec: policyv1.PodDisruptionBudgetSpec{ + MinAvailable: &intstr.IntOrString{Type: intstr.Int, IntVal: minAvailable}, + Selector: selector, + }, + Status: policyv1.PodDisruptionBudgetStatus{DisruptionsAllowed: 0}, + } +} + +func newPreemptPluginWithSync(t *testing.T, pods []*corev1.Pod, nodes []*corev1.Node, pdbs ...*policyv1.PodDisruptionBudget) (*VgpuPreempt, context.CancelFunc) { + t.Helper() + objs := make([]runtime.Object, 0, len(pods)+len(nodes)+len(pdbs)) + for _, p := range pods { + objs = append(objs, p) + } + for _, n := range nodes { + objs = append(objs, n) + } + for _, pdb := range pdbs { + objs = append(objs, pdb) + } + k8sClient := fake.NewClientset(objs...) + factory := informers.NewSharedInformerFactory(k8sClient, 0) + broadcaster := record.NewBroadcaster() + recorder := broadcaster.NewRecorder(scheme.Scheme, corev1.EventSource{Component: "test-scheduler-extender"}) + + plugin, err := New(factory, recorder) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + factory.Start(ctx.Done()) + factory.WaitForCacheSync(ctx.Done()) + + cleanup := func() { + cancel() + broadcaster.Shutdown() + } + return plugin, cleanup +} + +func TestNewWithNilFactory(t *testing.T) { + plugin, err := New(nil, nil) + assert.Nil(t, plugin) + assert.NotNil(t, err) + assert.Contains(t, err.Error(), "informerFactory cannot be nil") +} + +func TestNewWithValidFactory(t *testing.T) { + k8sClient := fake.NewClientset() + factory := informers.NewSharedInformerFactory(k8sClient, 0) + broadcaster := record.NewBroadcaster() + recorder := broadcaster.NewRecorder(scheme.Scheme, corev1.EventSource{Component: "test"}) + + plugin, err := New(factory, recorder) + assert.NoError(t, err) + assert.NotNil(t, plugin) + broadcaster.Shutdown() +} + +func TestIsProtectedFromPreemption(t *testing.T) { + tests := []struct { + name string + pod *corev1.Pod + want bool + }{ + {"normal pod", newVGPUPod("p", vgpuRequest{1, 1000, 10}, withPriority(10), withNodeName("n1")), false}, + {"terminating", newVGPUPod("p", vgpuRequest{1, 1000, 10}, withPriority(10), withNodeName("n1"), withDeletionTimestamp()), true}, + {"DaemonSet owned", newVGPUPod("p", vgpuRequest{1, 1000, 10}, withPriority(10), withNodeName("n1"), withOwner("DaemonSet")), true}, + {"ReplicaSet owned", newVGPUPod("p", vgpuRequest{1, 1000, 10}, withPriority(10), withNodeName("n1"), withOwner("ReplicaSet")), false}, + {"critical priority", newVGPUPod("p", vgpuRequest{1, 1000, 10}, withPriority(2000000001), withNodeName("n1")), true}, + {"kube-system namespace", newVGPUPod("p", vgpuRequest{1, 1000, 10}, withNamespace("kube-system"), withNodeName("n1")), true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, isProtectedFromPreemption(tt.pod)) + }) + } +} + +func TestIsCriticalPod(t *testing.T) { + tests := []struct { + name string + pod *corev1.Pod + expected bool + }{ + { + name: "kube-system namespace protected", + pod: newVGPUPod("critical", vgpuRequest{1, 1000, 10}, withNamespace("kube-system")), + expected: true, + }, + { + name: "kube-public namespace NOT protected", + pod: newVGPUPod("info", vgpuRequest{1, 1000, 10}, withNamespace("kube-public")), + expected: false, + }, + { + name: "system-cluster-critical PriorityClass protected", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "critical", Namespace: "my-app"}, + Spec: corev1.PodSpec{ + PriorityClassName: "system-cluster-critical", + Containers: []corev1.Container{{Name: "c"}}, + }, + }, + expected: true, + }, + { + name: "system-node-critical PriorityClass protected", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "critical", Namespace: "my-app"}, + Spec: corev1.PodSpec{ + PriorityClassName: "system-node-critical", + Containers: []corev1.Container{{Name: "c"}}, + }, + }, + expected: true, + }, + { + name: "high priority value protected", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "critical", Namespace: "my-app"}, + Spec: corev1.PodSpec{ + Priority: &[]int32{criticalPriorityThreshold + 100}[0], + Containers: []corev1.Container{{Name: "c"}}, + }, + }, + expected: true, + }, + { + name: "regular pod NOT protected", + pod: newVGPUPod("regular", vgpuRequest{1, 1000, 10}, withNamespace("default"), withPriority(5)), + expected: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.expected, isCriticalPod(tt.pod)) + }) + } +} + +func TestSortVictimsByPreference(t *testing.T) { + now := time.Now() + mkPod := func(name string, prio int32, createdAgo time.Duration) *corev1.Pod { + p := newVGPUPod(name, vgpuRequest{1, 1000, 10}, withPriority(prio), withNodeName("n1")) + p.CreationTimestamp = metav1.NewTime(now.Add(-createdAgo)) + return p + } + + t.Run("priority order ascending", func(t *testing.T) { + pods := []*corev1.Pod{mkPod("high", 100, time.Minute), mkPod("low", 10, time.Minute), mkPod("mid", 50, time.Minute)} + sortVictimsByPreference(pods) + assert.Equal(t, "low", pods[0].Name) + assert.Equal(t, "mid", pods[1].Name) + assert.Equal(t, "high", pods[2].Name) + }) + + t.Run("timestamp tiebreak - newer first", func(t *testing.T) { + older := mkPod("old", 10, time.Hour) + newer := mkPod("new", 10, time.Second) + pods := []*corev1.Pod{older, newer} + sortVictimsByPreference(pods) + assert.Equal(t, "new", pods[0].Name) + }) +} + +func TestExtractDeviceTypeFromResourceName(t *testing.T) { + cleanupMock := registerMock(newDefaultMockDevice()) + defer cleanupMock() + + devMap := device.GetDevices() + + tests := []struct { + name string + resourceName string + expected string + }{ + {"hami vgpu", "hami.io/vgpu", "nvidia"}, + {"hami vgpu memory", "hami.io/vgpu-memory", "nvidia"}, + {"hami vgpu cores", "hami.io/vgpu-cores", "nvidia"}, + {"unknown resource", "unknown.com/device", ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := extractDeviceTypeFromResourceName(tt.resourceName, devMap) + assert.Equal(t, tt.expected, got, "Failed to extract device type for %s", tt.resourceName) + }) + } +} + +func TestIsVGPUResourcePod(t *testing.T) { + tests := []struct { + name string + pod *corev1.Pod + expected bool + }{ + { + name: "hami annotation", + pod: newVGPUPod("p", vgpuRequest{1, 1000, 10}, withAnnotations(map[string]string{"hami.io/container-devices": "{}"})), + expected: true, + }, + { + name: "hami vgpu resource", + pod: newVGPUPod("p", vgpuRequest{1, 1000, 10}), + expected: true, + }, + { + name: "native nvidia gpu", + pod: newNativeGPUPod("p", 1), + expected: true, + }, + { + name: "no gpu resources", + pod: &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "p", Namespace: "default"}, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{ + Name: "c", + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("1"), + }, + }, + }}, + }, + }, + expected: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.expected, isVGPUResourcePod(tt.pod)) + }) + } +} + +func TestViolatePDB(t *testing.T) { + cleanupMock := registerMock(newDefaultMockDevice()) + defer cleanupMock() + + pod := newVGPUPod("victim", vgpuRequest{1, 1000, 10}, withLabels(map[string]string{"app": "critical"}), withPriority(5), withNodeName("n1")) + + pdbSelector := &metav1.LabelSelector{ + MatchLabels: map[string]string{"app": "critical"}, + } + + t.Run("PDB with zero disruptions allowed - violates", func(t *testing.T) { + pdbRestricted := newPDB("critical-pdb", "default", pdbSelector, 1) + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{pod}, nil, pdbRestricted) + defer cleanup() + assert.True(t, plugin.violatesPDB(pod)) + }) + + t.Run("PDB with disruptions allowed - does not violate", func(t *testing.T) { + pdbPermissive := newPDB("critical-pdb", "default", pdbSelector, 1) + pdbPermissive.Status.DisruptionsAllowed = 1 + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{pod}, nil, pdbPermissive) + defer cleanup() + assert.False(t, plugin.violatesPDB(pod)) + }) +} + +func TestViolatePDBWithDummyLister(t *testing.T) { + plugin := &VgpuPreempt{pdbLister: &dummyPDBLister{}} + pod := newVGPUPod("p", vgpuRequest{1, 1000, 10}) + assert.False(t, plugin.violatesPDB(pod)) +} + +func TestPDBViolationsUpperBound(t *testing.T) { + tests := []struct { + name string + originalCount int64 + keptLen int + expected int64 + }{ + {"original smaller", 10, 100, 10}, + {"kept smaller", 100, 10, 10}, + {"equal", 50, 50, 50}, + {"zero original", 0, 50, 0}, + {"zero kept", 50, 0, 0}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := pdbViolationsUpperBound(tt.originalCount, tt.keptLen) + assert.Equal(t, tt.expected, result) + }) + } +} + +func TestGetNativeWholeGPUResourcesCompleteness(t *testing.T) { + resources := getNativeWholeGPUResources() + expectedVendors := []string{ + "nvidia.com/gpu", + "amd.com/gpu", + "intel.com/gpu", + "huawei.com/Ascend310", + "huawei.com/Ascend910", + "cambricon.com/mlu", + "hygon.com/dcu", + "iluvatar.ai/gpu", + "metax-tech.com/gpu", + "mthreads.com/gpu", + "kunlun.com/gpu", + "aws.amazon.com/neuron", + "vastai.com/gpu", + "enflame.com/gpu", + } + + for _, vendor := range expectedVendors { + assert.True(t, resources[vendor], "Missing vendor: %s", vendor) + } +} + +func TestResolveVictimsMapLegacyFormat(t *testing.T) { + cleanupMock := registerMock(newDefaultMockDevice()) + defer cleanupMock() + + pod := newVGPUPod("victim", vgpuRequest{1, 1000, 10}, withPriority(5), withNodeName("n1")) + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{pod}, nil) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + NodeNameToVictims: map[string]*extenderv1.Victims{ + "n1": {Pods: []*corev1.Pod{pod}}, + }, + } + got, err := plugin.resolveVictimsMap(args) + assert.NoError(t, err) + require.Len(t, got, 1) + assert.Equal(t, pod.UID, got["n1"].Pods[0].UID) +} + +func TestResolveVictimsMapEmptyAndMinimal(t *testing.T) { + cleanupMock := registerMock(newDefaultMockDevice()) + defer cleanupMock() + + plugin, cleanup := newPreemptPluginWithSync(t, nil, nil) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{} + got, err := plugin.resolveVictimsMap(args) + assert.NoError(t, err) + assert.Empty(t, got) +} + +func TestPreemptHappyPath(t *testing.T) { + cleanupMock := registerMock(newMockDevice(2, 20000, 200)) + defer cleanupMock() + + node := newTestNode("node1") + lowA := newVGPUPod("low-a", vgpuRequest{1, 10000, 100}, withPriority(10), withNodeName(node.Name)) + lowB := newVGPUPod("low-b", vgpuRequest{1, 10000, 100}, withPriority(10), withNodeName(node.Name)) + preemptor := newVGPUPod("high-preemptor", vgpuRequest{1, 10000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{lowA, lowB, preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(lowB.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + require.Contains(t, res.NodeNameToMetaVictims, node.Name) + assert.Len(t, res.NodeNameToMetaVictims[node.Name].Pods, 1) + assert.Equal(t, string(lowB.UID), res.NodeNameToMetaVictims[node.Name].Pods[0].UID) +} + +func TestPreemptNodeWithZeroVictims(t *testing.T) { + cleanupMock := registerMock(newMockDevice(1, 10000, 100)) + defer cleanupMock() + + node := newTestNode("node1") + preemptor := newVGPUPod("fit", vgpuRequest{1, 10000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{}}, + }, + } + res := plugin.Preempt(context.Background(), args) + require.Contains(t, res.NodeNameToMetaVictims, node.Name) + assert.Empty(t, res.NodeNameToMetaVictims[node.Name].Pods) +} + +func TestPreemptProtectedPodDropsNode(t *testing.T) { + cleanupMock := registerMock(newMockDevice(1, 10000, 100)) + defer cleanupMock() + + node := newTestNode("node1") + dsA := newVGPUPod("ds", vgpuRequest{1, 10000, 100}, withPriority(10), withNodeName(node.Name), withOwner("DaemonSet")) + preemptor := newVGPUPod("high", vgpuRequest{1, 10000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{dsA, preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(dsA.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + assert.NotContains(t, res.NodeNameToMetaVictims, node.Name) +} + +func TestPreemptAddsAdditionalVictim(t *testing.T) { + cleanupMock := registerMock(newMockDevice(2, 20000, 200)) + defer cleanupMock() + + node := newTestNode("node1") + protected := newVGPUPod("protected", vgpuRequest{1, 10000, 100}, withPriority(10), withNodeName(node.Name), withOwner("DaemonSet")) + evictable := newVGPUPod("evict", vgpuRequest{1, 10000, 100}, withPriority(5), withNodeName(node.Name)) + preemptor := newVGPUPod("high", vgpuRequest{1, 10000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{protected, evictable, preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(protected.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + require.Contains(t, res.NodeNameToMetaVictims, node.Name) + meta := res.NodeNameToMetaVictims[node.Name] + require.Len(t, meta.Pods, 1) + assert.Equal(t, string(evictable.UID), meta.Pods[0].UID) +} + +func TestPreemptNativeGPUAccounting(t *testing.T) { + cleanupMock := registerMock(newMockDevice(1, 10000, 100)) + defer cleanupMock() + + node := newTestNode("node1") + nativePod := newNativeGPUPod("native", 1, withPriority(5), withNodeName(node.Name)) + preemptor := newVGPUPod("vgpu", vgpuRequest{1, 10000, 50}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{nativePod, preemptor}, []*corev1.Node{node}) + defer cleanup() + + t.Run("native pod proposed as victim - preemptor fits", func(t *testing.T) { + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(nativePod.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + require.Contains(t, res.NodeNameToMetaVictims, node.Name) + assert.Len(t, res.NodeNameToMetaVictims[node.Name].Pods, 1) + }) + + t.Run("no victims proposed - preemptor cannot fit", func(t *testing.T) { + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{}}, + }, + } + res := plugin.Preempt(context.Background(), args) + assert.NotContains(t, res.NodeNameToMetaVictims, node.Name) + }) +} + +func TestPreemptMultiContainerPod(t *testing.T) { + cleanupMock := registerMock(newMockDevice(3, 30000, 300)) + defer cleanupMock() + + node := newTestNode("node1") + + victim := newMultiContainerPod("victim", []corev1.Container{ + { + Name: "gpu-1", + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceName("hami.io/vgpu"): resource.MustParse("1"), + corev1.ResourceName("hami.io/vgpu-memory"): resource.MustParse("10000"), + corev1.ResourceName("hami.io/vgpu-cores"): resource.MustParse("50"), + }, + }, + }, + { + Name: "gpu-2", + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceName("hami.io/vgpu"): resource.MustParse("1"), + corev1.ResourceName("hami.io/vgpu-memory"): resource.MustParse("10000"), + corev1.ResourceName("hami.io/vgpu-cores"): resource.MustParse("50"), + }, + }, + }, + }, withPriority(5), withNodeName(node.Name)) + + preemptor := newVGPUPod("preemptor", vgpuRequest{2, 20000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{victim, preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(victim.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + require.Contains(t, res.NodeNameToMetaVictims, node.Name) + assert.Len(t, res.NodeNameToMetaVictims[node.Name].Pods, 1) +} + +func TestPreemptPDBViolation(t *testing.T) { + cleanupMock := registerMock(newMockDevice(2, 20000, 200)) + defer cleanupMock() + + node := newTestNode("node1") + + protectedPod := newVGPUPod("protected", vgpuRequest{1, 10000, 100}, + withLabels(map[string]string{"app": "critical"}), + withPriority(5), + withNodeName(node.Name)) + + evictablePod := newVGPUPod("evictable", vgpuRequest{1, 10000, 100}, + withPriority(5), + withNodeName(node.Name)) + + preemptor := newVGPUPod("preemptor", vgpuRequest{1, 10000, 100}, withPriority(100)) + + pdb := newPDB("critical-pdb", "default", &metav1.LabelSelector{ + MatchLabels: map[string]string{"app": "critical"}, + }, 1) + pdb.Status.DisruptionsAllowed = 0 + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{protectedPod, evictablePod, preemptor}, []*corev1.Node{node}, pdb) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(evictablePod.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + require.Contains(t, res.NodeNameToMetaVictims, node.Name) + assert.Len(t, res.NodeNameToMetaVictims[node.Name].Pods, 1) + assert.Equal(t, string(evictablePod.UID), res.NodeNameToMetaVictims[node.Name].Pods[0].UID) +} + +func TestPreemptInitContainerRecognition(t *testing.T) { + cleanupMock := registerMock(newMockDevice(2, 20000, 200)) + defer cleanupMock() + + node := newTestNode("node1") + victim := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "victim", Namespace: "default", UID: types.UID(uuid.NewString())}, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{{ + Name: "init", + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceName("hami.io/vgpu"): resource.MustParse("1"), + corev1.ResourceName("hami.io/vgpu-memory"): resource.MustParse("10000"), + corev1.ResourceName("hami.io/vgpu-cores"): resource.MustParse("50"), + }, + }, + }}, + }, + Status: corev1.PodStatus{Phase: corev1.PodRunning}, + } + victim.Spec.NodeName = node.Name + victim.Spec.Priority = &[]int32{5}[0] + + preemptor := newVGPUPod("high", vgpuRequest{1, 10000, 50}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{victim, preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(victim.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + require.Contains(t, res.NodeNameToMetaVictims, node.Name) + assert.Len(t, res.NodeNameToMetaVictims[node.Name].Pods, 1) +} + +func TestPassthroughMetaVictims(t *testing.T) { + cleanupMock := registerMock(newDefaultMockDevice()) + defer cleanupMock() + + plugin, cleanup := newPreemptPluginWithSync(t, nil, nil) + defer cleanup() + + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "cpu-only", Namespace: "default", UID: types.UID(uuid.NewString())}, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{ + Name: "c", + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1")}, + }, + }}, + }, + } + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: pod, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + "node1": {Pods: []*extenderv1.MetaPod{{UID: "some-uid"}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + require.Contains(t, res.NodeNameToMetaVictims, "node1") + assert.Len(t, res.NodeNameToMetaVictims["node1"].Pods, 1) +} + +func TestPreemptHardwareDriverError(t *testing.T) { + mockDev := newDefaultMockDevice() + mockDev.failGetNodeDevices = true + cleanupMock := registerMock(mockDev) + defer cleanupMock() + + node := newTestNode("node1") + lowPod := newVGPUPod("low", vgpuRequest{1, 10000, 100}, withPriority(10), withNodeName(node.Name)) + preemptor := newVGPUPod("high", vgpuRequest{1, 10000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{lowPod, preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(lowPod.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + assert.NotContains(t, res.NodeNameToMetaVictims, node.Name) +} + +func TestPreemptDeviceFitFails(t *testing.T) { + mockDev := newDefaultMockDevice() + mockDev.fitReturnsTrue = false + cleanupMock := registerMock(mockDev) + defer cleanupMock() + + node := newTestNode("node1") + lowPod := newVGPUPod("low", vgpuRequest{1, 10000, 100}, withPriority(10), withNodeName(node.Name)) + preemptor := newVGPUPod("high", vgpuRequest{1, 10000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{lowPod, preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(lowPod.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + assert.NotContains(t, res.NodeNameToMetaVictims, node.Name) +} + +func TestPreemptNodeNotFoundInCache(t *testing.T) { + cleanupMock := registerMock(newDefaultMockDevice()) + defer cleanupMock() + + preemptor := newVGPUPod("high", vgpuRequest{1, 10000, 100}, withPriority(100)) + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{preemptor}, nil) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + "ghost-node": {Pods: []*extenderv1.MetaPod{{UID: "any-uid"}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + assert.NotContains(t, res.NodeNameToMetaVictims, "ghost-node") +} + +func TestPreemptVictimNotFoundInCache(t *testing.T) { + cleanupMock := registerMock(newDefaultMockDevice()) + defer cleanupMock() + + node := newTestNode("node1") + preemptor := newVGPUPod("high", vgpuRequest{1, 10000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: "ghost-pod-uid"}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + require.Contains(t, res.NodeNameToMetaVictims, node.Name) + assert.Empty(t, res.NodeNameToMetaVictims[node.Name].Pods) +} + +func TestAccountVGPURequestsReturnErrorOnFailure(t *testing.T) { + cleanupMock := registerMock(newDefaultMockDevice()) + defer cleanupMock() + + plugin, cleanup := newPreemptPluginWithSync(t, nil, nil) + defer cleanup() + + state := make(map[string][]*device.DeviceUsage) + pod := newVGPUPod("test", vgpuRequest{1, 1000, 10}) + + err := plugin.accountVGPURequests(state, pod, nil) + assert.NotNil(t, err) + assert.Contains(t, err.Error(), "no device state available") +} + +func TestAccountVGPURequestsNoSatisfyingDevice(t *testing.T) { + cleanupMock := registerMock(newDefaultMockDevice()) + defer cleanupMock() + + plugin, cleanup := newPreemptPluginWithSync(t, nil, nil) + defer cleanup() + + state := map[string][]*device.DeviceUsage{ + "nvidia": { + { + ID: "gpu-0", + Count: 10, + Used: 10, + Totalmem: 10000, + Usedmem: 10000, + Totalcore: 100, + Usedcores: 100, + }, + }, + } + + pod := newVGPUPod("test", vgpuRequest{1, 1000, 10}) + err := plugin.accountVGPURequests(state, pod, nil) + assert.NotNil(t, err) + assert.Contains(t, err.Error(), "no device could satisfy vGPU request") +} + +func TestPreemptMetricsIncrement(t *testing.T) { + cleanupMock := registerMock(newMockDevice(2, 20000, 200)) + defer cleanupMock() + + node := newTestNode("node1") + lowPod := newVGPUPod("low", vgpuRequest{1, 10000, 100}, withPriority(10), withNodeName(node.Name)) + preemptor := newVGPUPod("high", vgpuRequest{1, 10000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{lowPod, preemptor}, []*corev1.Node{node}) + defer cleanup() + + initialMetrics := GetPreemptionMetrics() + initialAttempts := initialMetrics["preemption_attempts"] + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(lowPod.UID)}}}, + }, + } + plugin.Preempt(context.Background(), args) + + finalMetrics := GetPreemptionMetrics() + assert.Greater(t, finalMetrics["preemption_attempts"], initialAttempts) +} + +func TestRefineForNodeWithAccountingError(t *testing.T) { + cleanupMock := registerMock(newMockDevice(1, 5000, 50)) + defer cleanupMock() + + node := newTestNode("node1") + + highRequest := newVGPUPod("demanding", vgpuRequest{1, 20000, 100}, + withPriority(10), withNodeName(node.Name)) + + preemptor := newVGPUPod("high", vgpuRequest{1, 20000, 100}, withPriority(100)) + + plugin, cleanup := newPreemptPluginWithSync(t, []*corev1.Pod{highRequest, preemptor}, []*corev1.Node{node}) + defer cleanup() + + args := extenderv1.ExtenderPreemptionArgs{ + Pod: preemptor, + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{ + node.Name: {Pods: []*extenderv1.MetaPod{{UID: string(highRequest.UID)}}}, + }, + } + res := plugin.Preempt(context.Background(), args) + assert.NotContains(t, res.NodeNameToMetaVictims, node.Name) +} diff --git a/pkg/scheduler/routes/route.go b/pkg/scheduler/routes/route.go index 38dd0a712..afb16b952 100644 --- a/pkg/scheduler/routes/route.go +++ b/pkg/scheduler/routes/route.go @@ -18,16 +18,19 @@ package routes import ( "bytes" + "context" "encoding/json" "fmt" "io" "net/http" + "time" "github.com/julienschmidt/httprouter" "k8s.io/klog/v2" extenderv1 "k8s.io/kube-scheduler/extender/v1" "github.com/Project-HAMi/HAMi/pkg/scheduler" + "github.com/Project-HAMi/HAMi/pkg/scheduler/preempt" ) const maxRequestSize = 1024 * 1024 // 1MB limit @@ -168,3 +171,50 @@ func ReadyzRoute(s *scheduler.Scheduler) httprouter.Handle { w.WriteHeader(http.StatusOK) } } + +func PreemptRoute(p *preempt.VgpuPreempt) httprouter.Handle { + return func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) { + klog.V(5).Infoln("Entering Preempt Route handler") + + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + + checkBody(w, r) + + var buf bytes.Buffer + limitedReader := io.LimitReader(r.Body, maxRequestSize) + body := io.TeeReader(limitedReader, &buf) + + var args extenderv1.ExtenderPreemptionArgs + var result *extenderv1.ExtenderPreemptionResult + + if err := json.NewDecoder(body).Decode(&args); err != nil { + klog.ErrorS(err, "Failed to decode extender preemption arguments") + + result = &extenderv1.ExtenderPreemptionResult{ + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{}, + } + } else { + ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second) + defer cancel() + + result = p.Preempt(ctx, args) + + if result == nil { + result = &extenderv1.ExtenderPreemptionResult{ + NodeNameToMetaVictims: map[string]*extenderv1.MetaVictims{}, + } + } + } + + w.Header().Set("Content-Type", "application/json") + + if err := json.NewEncoder(w).Encode(result); err != nil { + klog.ErrorS(err, "Failed to encode preemption result") + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + } +} diff --git a/pkg/scheduler/routes/route_test.go b/pkg/scheduler/routes/route_test.go index 246a39c47..ca566722b 100644 --- a/pkg/scheduler/routes/route_test.go +++ b/pkg/scheduler/routes/route_test.go @@ -17,6 +17,7 @@ limitations under the License. package routes import ( + "net/http" "net/http/httptest" "strings" "testing" @@ -135,3 +136,39 @@ func TestCheckBodyNil(t *testing.T) { t.Errorf("Expected status 400 for nil body, got %d", w.Code) } } + +func TestPreemptRouteMethodNotAllowed(t *testing.T) { + req := httptest.NewRequest("GET", "/preempt", nil) + w := httptest.NewRecorder() + + // We can pass nil for *preempt.VgpuPreempt since the method check happens + // before the handler attempts to use the preemptor. + handler := PreemptRoute(nil) + handler(w, req, nil) + + if w.Code != http.StatusMethodNotAllowed { + t.Errorf("Expected status 405 Method Not Allowed for GET request, got %d", w.Code) + } +} + +func TestMaxRequestSizePreempt(t *testing.T) { + hugePayload := strings.Repeat(" ", maxRequestSize+100) + req := httptest.NewRequest("POST", "/preempt", strings.NewReader(hugePayload)) + w := httptest.NewRecorder() + + // Passing nil for *preempt.VgpuPreempt is safe here. The LimitReader will trigger + // an EOF error during json.Decode, which skips the p.Preempt() call entirely. + handler := PreemptRoute(nil) + handler(w, req, nil) + + respBody := w.Body.String() + + // Since ExtenderPreemptionResult doesn't have an Error string field like the others, + // the handler defaults to returning an empty NodeNameToMetaVictims map upon decode failure. + if w.Code != http.StatusOK { + t.Errorf("Expected status 200 on handled parse error, got %d", w.Code) + } + if !strings.Contains(respBody, "NodeNameToMetaVictims") { + t.Errorf("LimitReader failed in PreemptRoute or returned unexpected body. Response: %s", respBody) + } +} diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index 711674d76..7133c6283 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -73,9 +73,10 @@ type Scheduler struct { //Node status returned by filter cachedstatus map[string]*NodeUsage //Node Overview - overviewstatus map[string]*NodeUsage - eventRecorder record.EventRecorder - started uint32 // 0 = false, 1 = true + overviewstatus map[string]*NodeUsage + eventRecorder record.EventRecorder + informerFactory informers.SharedInformerFactory + started uint32 // 0 = false, 1 = true lock sync.RWMutex synced bool @@ -133,6 +134,14 @@ func (s *Scheduler) doNodeNotify() { } } +func (s *Scheduler) GetEventRecorder() record.EventRecorder { + return s.eventRecorder +} + +func (s *Scheduler) GetInformerFactory() informers.SharedInformerFactory { + return s.informerFactory +} + func (s *Scheduler) onAddPod(obj any) { pod, ok := obj.(*corev1.Pod) if !ok { @@ -270,6 +279,7 @@ func (s *Scheduler) Start() error { klog.InfoS("Starting HAMi scheduler components") s.kubeClient = client.GetClient() informerFactory := informers.NewSharedInformerFactoryWithOptions(s.kubeClient, defaultResync) + s.informerFactory = informerFactory s.podLister = informerFactory.Core().V1().Pods().Lister() s.nodeLister = informerFactory.Core().V1().Nodes().Lister() s.quotaLister = informerFactory.Core().V1().ResourceQuotas().Lister()