diff --git a/charts/hami/templates/scheduler/configmap.yaml b/charts/hami/templates/scheduler/configmap.yaml index 960e2b0d87..112f77f5ec 100644 --- a/charts/hami/templates/scheduler/configmap.yaml +++ b/charts/hami/templates/scheduler/configmap.yaml @@ -35,6 +35,7 @@ data: enableHTTPS: false {{- end }} filterVerb: filter + prioritizeVerb: prioritize bindVerb: bind nodeCacheCapable: true weight: 1 @@ -60,6 +61,7 @@ data: "enableHttps": false, {{- end }} "filterVerb": "filter", + "prioritizeVerb": "prioritize", "bindVerb": "bind", "weight": 1, "nodeCacheCapable": true, diff --git a/cmd/scheduler/main.go b/cmd/scheduler/main.go index f7c68d1585..da8be44d08 100644 --- a/cmd/scheduler/main.go +++ b/cmd/scheduler/main.go @@ -144,6 +144,7 @@ func start() error { // start http server router := httprouter.New() router.POST("/filter", routes.PredicateRoute(sher)) + router.POST("/prioritize", routes.PrioritizeRoute(sher)) router.POST("/bind", routes.Bind(sher)) router.POST("/webhook", routes.WebHookRoute()) router.GET("/healthz", routes.HealthzRoute()) diff --git a/pkg/device/pod_test.go b/pkg/device/pod_test.go index a805024c12..273344610d 100644 --- a/pkg/device/pod_test.go +++ b/pkg/device/pod_test.go @@ -229,6 +229,44 @@ func TestGetPod(t *testing.T) { } } +func TestGetPodReturnsConcurrentSnapshot(t *testing.T) { + manager := NewPodManager() + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} + first := PodDevices{"device": {{{UUID: "device-0"}}}} + second := PodDevices{"device": {{{UUID: "device-1"}}}} + manager.AddPod(pod, "node-a", first) + snapshot, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + snapshot.NodeID = "mutated-node" + snapshot.Devices["device"][0][0].UUID = "mutated-device" + stored, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, "node-a", stored.NodeID) + assert.Equal(t, "device-0", stored.Devices["device"][0][0].UUID) + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + for range 10000 { + if allocation, ok := manager.GetPod(pod); ok { + _ = allocation.DeepCopy() + } + } + }() + go func() { + defer wg.Done() + for i := range 10000 { + if i%2 == 0 { + manager.AddPod(pod, "node-a", first) + } else { + manager.AddPod(pod, "node-a", second) + } + } + }() + wg.Wait() +} + func TestAddPod(t *testing.T) { podManager := NewPodManager() podManager.pods["uid1"] = &PodInfo{ @@ -299,6 +337,186 @@ func TestAddPod(t *testing.T) { } } +func TestPodReservationOwnership(t *testing.T) { + manager := NewPodManager() + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} + first := PodDevices{"device": {{{UUID: "device-0"}}}} + second := PodDevices{"device": {{{UUID: "device-1"}}}} + + reserved := manager.ReservePodIfAbsent(pod, "node-a", first) + assert.Equal(t, true, reserved) + reserved = manager.ReservePodIfAbsent(pod, "node-b", second) + assert.Equal(t, false, reserved) + assert.Equal(t, false, manager.AddPod(pod.DeepCopy(), "node-b", second)) + allocation, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, "node-a", allocation.NodeID) + assert.Equal(t, first, allocation.Devices) + _, owned := manager.reservations[pod.UID] + assert.Equal(t, true, owned) +} + +func TestInformerObservationTransfersPodReservationOwnership(t *testing.T) { + manager := NewPodManager() + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} + devices := PodDevices{"device": {{{UUID: "device-0"}}}} + + reserved := manager.ReservePodIfAbsent(pod, "node-a", devices) + assert.Equal(t, true, reserved) + observed := pod.DeepCopy() + observed.ResourceVersion = "2" + assert.Equal(t, false, manager.AddPod(observed, "node-a", devices.DeepCopy())) + _, owned := manager.reservations[pod.UID] + assert.Equal(t, false, owned) + allocation, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, observed, allocation.Pod) + assert.Equal(t, devices, allocation.Devices) +} + +func TestEncodedInformerObservationTransfersPodReservationOwnership(t *testing.T) { + manager := NewPodManager() + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} + devices := PodDevices{"device": {{{UUID: "device-0", Type: "device", Usedmem: 1, Usedcores: 1}}}} + annotations := EncodePodDevices(map[string]string{"device": "hami.io/device-allocated"}, devices) + observed, err := DecodePodDevices(map[string]string{"device": "hami.io/device-allocated"}, annotations) + assert.NoError(t, err) + assert.NotEqual(t, devices, observed) + + reserved := manager.ReservePodIfAbsent(pod, "node-a", devices) + assert.Equal(t, true, reserved) + assert.Equal(t, false, manager.AddPod(pod.DeepCopy(), "node-a", observed)) + _, owned := manager.reservations[pod.UID] + assert.Equal(t, false, owned) + allocation, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, devices, allocation.Devices) +} + +func TestTakeAndDeleteClearsReservationOwnership(t *testing.T) { + manager := NewPodManager() + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} + first := PodDevices{"device": {{{UUID: "device-0"}}}} + second := PodDevices{"device": {{{UUID: "device-1"}}}} + + reserved := manager.ReservePodIfAbsent(pod, "node-a", first) + assert.Equal(t, true, reserved) + _, removed := manager.TakeAndDeletePod(pod) + assert.Equal(t, true, removed) + _, owned := manager.reservations[pod.UID] + assert.Equal(t, false, owned) + reserved = manager.ReservePodIfAbsent(pod, "node-b", second) + assert.Equal(t, true, reserved) + allocation, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, "node-b", allocation.NodeID) + assert.Equal(t, second, allocation.Devices) +} + +func TestStaleInformerUpdateDoesNotOverwritePodReservation(t *testing.T) { + manager := NewPodManager() + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} + reservedDevices := PodDevices{"device": {{{UUID: "device-new"}}}} + staleDevices := PodDevices{"device": {{{UUID: "device-old"}}}} + + reserved := manager.ReservePodIfAbsent(pod, "node-new", reservedDevices) + assert.Equal(t, true, reserved) + assert.Equal(t, false, manager.AddPod(pod.DeepCopy(), "node-old", staleDevices)) + + allocation, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, "node-new", allocation.NodeID) + assert.Equal(t, reservedDevices, allocation.Devices) + + assert.Equal(t, false, manager.AddPod(pod.DeepCopy(), "node-new", reservedDevices.DeepCopy())) + allocation, ok = manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, "node-new", allocation.NodeID) + assert.Equal(t, reservedDevices, allocation.Devices) +} + +func TestReplacementReservationRetainsOwnershipUntilObserved(t *testing.T) { + manager := NewPodManager() + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} + previousDevices := PodDevices{"device": {{{UUID: "device-old"}}}} + replacementDevices := PodDevices{"device": {{{UUID: "device-new"}}}} + manager.AddPod(pod, "node-old", previousDevices) + previous, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + replaced := manager.ReplacePodReservation(pod, previous, "node-new", replacementDevices) + assert.Equal(t, true, replaced) + assert.Equal(t, false, manager.AddPod(pod.DeepCopy(), "node-old", previousDevices.DeepCopy())) + allocation, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, "node-new", allocation.NodeID) + assert.Equal(t, replacementDevices, allocation.Devices) + assert.Equal(t, false, manager.AddPod(pod.DeepCopy(), "node-new", replacementDevices.DeepCopy())) + _, owned := manager.reservations[pod.UID] + assert.Equal(t, false, owned) +} + +func TestReplacePodReservationRejectsStaleExpectedAllocation(t *testing.T) { + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} + reservedDevices := PodDevices{"device": {{{UUID: "device-current"}}}} + replacementDevices := PodDevices{"device": {{{UUID: "device-replacement"}}}} + tests := []struct { + name string + expected *PodInfo + }{ + { + name: "stale node", + expected: &PodInfo{ + Pod: pod.DeepCopy(), + NodeID: "node-stale", + Devices: reservedDevices.DeepCopy(), + }, + }, + { + name: "missing expected allocation", + expected: nil, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + manager := NewPodManager() + assert.Equal(t, true, manager.ReservePodIfAbsent(pod, "node-current", reservedDevices)) + + replaced := manager.ReplacePodReservation(pod, test.expected, "node-replacement", replacementDevices) + assert.Equal(t, false, replaced) + allocation, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, "node-current", allocation.NodeID) + assert.Equal(t, reservedDevices, allocation.Devices) + _, owned := manager.reservations[pod.UID] + assert.Equal(t, true, owned) + }) + } +} + +func TestReplacePodReservationMatchesDecodedInformerDevices(t *testing.T) { + manager := NewPodManager() + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} + previousDevices := PodDevices{"device": {{{UUID: "device-old", Type: "device", Usedmem: 1, Usedcores: 1}}}} + replacementDevices := PodDevices{"device": {{{UUID: "device-new", Type: "device", Usedmem: 1, Usedcores: 1}}}} + manager.AddPod(pod, "node-old", previousDevices) + expected, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + + annotations := EncodePodDevices(map[string]string{"device": "hami.io/device-allocated"}, previousDevices) + observed, err := DecodePodDevices(map[string]string{"device": "hami.io/device-allocated"}, annotations) + assert.NoError(t, err) + assert.NotEqual(t, previousDevices, observed) + manager.AddPod(pod.DeepCopy(), "node-old", observed) + + replaced := manager.ReplacePodReservation(pod, expected, "node-new", replacementDevices) + assert.Equal(t, true, replaced) + allocation, ok := manager.GetPod(pod) + assert.Equal(t, true, ok) + assert.Equal(t, "node-new", allocation.NodeID) + assert.Equal(t, replacementDevices, allocation.Devices) +} + func TestUpdatePod(t *testing.T) { podManager := NewPodManager() diff --git a/pkg/device/pods.go b/pkg/device/pods.go index a549bc4c5d..9669206f57 100644 --- a/pkg/device/pods.go +++ b/pkg/device/pods.go @@ -18,6 +18,7 @@ package device import ( "maps" + "reflect" "sync" corev1 "k8s.io/api/core/v1" @@ -39,13 +40,15 @@ type PodUseDeviceStat struct { } type PodManager struct { - pods map[k8stypes.UID]*PodInfo - mutex sync.RWMutex + pods map[k8stypes.UID]*PodInfo + reservations map[k8stypes.UID]struct{} + mutex sync.RWMutex } func NewPodManager() *PodManager { pm := &PodManager{ - pods: make(map[k8stypes.UID]*PodInfo), + pods: make(map[k8stypes.UID]*PodInfo), + reservations: make(map[k8stypes.UID]struct{}), } klog.InfoS("Pod manager initialized", "podCount", len(pm.pods)) return pm @@ -55,9 +58,9 @@ func (m *PodManager) AddPod(pod *corev1.Pod, nodeID string, devices PodDevices) m.mutex.Lock() defer m.mutex.Unlock() - _, exists := m.pods[pod.UID] + pi, exists := m.pods[pod.UID] if !exists { - pi := &PodInfo{ + pi = &PodInfo{ Pod: pod, NodeID: nodeID, Devices: devices, @@ -68,8 +71,25 @@ func (m *PodManager) AddPod(pod *corev1.Pod, nodeID string, devices PodDevices) "nodeID", nodeID, "devices", devices, ) + } else if _, reserved := m.reservations[pod.UID]; reserved { + if pi.NodeID == nodeID && reservationMatchesInformerDevices(pi.Devices, devices) { + // The informer observed the allocation written by Bind. Clear the + // reservation marker so later informer updates reconcile normally. + pi.Pod = pod + delete(m.reservations, pod.UID) + klog.V(5).InfoS("Pod reservation observed by informer", + "pod", klog.KRef(pod.Namespace, pod.Name), + "nodeID", nodeID, + ) + } else { + klog.V(5).InfoS("Ignoring informer update for bind-owned reservation", + "pod", klog.KRef(pod.Namespace, pod.Name), + "reservedNodeID", pi.NodeID, + "informerNodeID", nodeID, + ) + } } else { - m.pods[pod.UID].Devices = devices + pi.Devices = devices klog.V(5).InfoS("Pod devices updated", "pod", klog.KRef(pod.Namespace, pod.Name), "devices", devices, @@ -79,6 +99,76 @@ func (m *PodManager) AddPod(pod *corev1.Pod, nodeID string, devices PodDevices) return !exists } +func reservationMatchesInformerDevices(reserved, observed PodDevices) bool { + if reflect.DeepEqual(reserved, observed) { + return true + } + + // DecodePodDevices preserves the trailing annotation separator as an empty + // container entry. Remove only that decoder artifact before comparing with + // the in-memory allocation produced by Bind. + normalized := observed.DeepCopy() + for deviceType, containers := range normalized { + if len(containers) > 0 && len(containers[len(containers)-1]) == 0 { + normalized[deviceType] = containers[:len(containers)-1] + } + } + return reflect.DeepEqual(reserved, normalized) +} + +// ReservePodIfAbsent records a bind-owned allocation without replacing an +// allocation installed concurrently by an informer or another binding attempt. +func (m *PodManager) ReservePodIfAbsent(pod *corev1.Pod, nodeID string, devices PodDevices) bool { + m.mutex.Lock() + defer m.mutex.Unlock() + + if _, exists := m.pods[pod.UID]; exists { + return false + } + m.pods[pod.UID] = &PodInfo{ + Pod: pod, + NodeID: nodeID, + Devices: devices, + } + m.reservations[pod.UID] = struct{}{} + klog.InfoS("Pod allocation reserved", + "pod", klog.KRef(pod.Namespace, pod.Name), + "nodeID", nodeID, + "devices", devices, + ) + return true +} + +// ReplacePodReservation atomically replaces the expected allocation with a +// bind-owned reservation. +func (m *PodManager) ReplacePodReservation(pod *corev1.Pod, expected *PodInfo, nodeID string, devices PodDevices) bool { + m.mutex.Lock() + defer m.mutex.Unlock() + + current, exists := m.pods[pod.UID] + currentNodeID := "" + if exists { + currentNodeID = current.NodeID + } + expectedNodeID := "" + if expected != nil { + expectedNodeID = expected.NodeID + } + if !exists || expected == nil || current.NodeID != expected.NodeID || + !reservationMatchesInformerDevices(expected.Devices, current.Devices) { + klog.V(5).InfoS("Pod reservation replacement rejected", + "pod", klog.KRef(pod.Namespace, pod.Name), + "exists", exists, + "currentNodeID", currentNodeID, + "expectedNodeID", expectedNodeID, + ) + return false + } + m.pods[pod.UID] = &PodInfo{Pod: pod, NodeID: nodeID, Devices: devices} + m.reservations[pod.UID] = struct{}{} + return true +} + func (m *PodManager) UpdatePod(pod *corev1.Pod) { m.mutex.Lock() defer m.mutex.Unlock() @@ -103,6 +193,7 @@ func (m *PodManager) DelPod(pod *corev1.Pod) { "nodeID", pi.NodeID, ) delete(m.pods, pod.UID) + delete(m.reservations, pod.UID) } else { klog.InfoS("Pod not found for deletion", "pod", klog.KRef(pod.Namespace, pod.Name), @@ -115,7 +206,10 @@ func (m *PodManager) GetPod(pod *corev1.Pod) (*PodInfo, bool) { defer m.mutex.RUnlock() pi, ok := m.pods[pod.UID] - return pi, ok + if !ok { + return nil, false + } + return pi.DeepCopy(), true } func (m *PodManager) TakeAndDeletePod(pod *corev1.Pod) (*PodInfo, bool) { @@ -125,6 +219,7 @@ func (m *PodManager) TakeAndDeletePod(pod *corev1.Pod) (*PodInfo, bool) { pi, ok := m.pods[pod.UID] if ok { delete(m.pods, pod.UID) + delete(m.reservations, pod.UID) klog.InfoS("Pod taken and deleted", "pod", klog.KRef(pod.Namespace, pod.Name), "nodeID", pi.NodeID) } return pi, ok diff --git a/pkg/scheduler/routes/route.go b/pkg/scheduler/routes/route.go index 38dd0a7127..daf423d7d2 100644 --- a/pkg/scheduler/routes/route.go +++ b/pkg/scheduler/routes/route.go @@ -18,6 +18,7 @@ package routes import ( "bytes" + "context" "encoding/json" "fmt" "io" @@ -95,6 +96,48 @@ func PredicateRoute(s *scheduler.Scheduler) httprouter.Handle { } } +func PrioritizeRoute(s *scheduler.Scheduler) httprouter.Handle { + klog.Infoln("Initializing Prioritize Route") + return prioritizeRoute(s.WaitForCacheSync, s.Prioritize) +} + +func prioritizeRoute( + waitForCacheSync func(context.Context) bool, + prioritize func(extenderv1.ExtenderArgs) (*extenderv1.HostPriorityList, error), +) httprouter.Handle { + return func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) { + limitedReader := io.LimitReader(r.Body, maxRequestSize) + var extenderArgs extenderv1.ExtenderArgs + if err := json.NewDecoder(limitedReader).Decode(&extenderArgs); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + if extenderArgs.Pod == nil { + http.Error(w, "extender args must contain a pod", http.StatusBadRequest) + return + } + if !waitForCacheSync(r.Context()) { + http.Error(w, "context cancelled", http.StatusServiceUnavailable) + return + } + + priorities, err := prioritize(extenderArgs) + if err != nil { + klog.ErrorS(err, "Prioritize error for pod", "pod", extenderArgs.Pod.Name) + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + response, err := json.Marshal(priorities) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(response) + } +} + func Bind(s *scheduler.Scheduler) httprouter.Handle { return func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) { klog.V(5).Infoln("Entering Bind handler") diff --git a/pkg/scheduler/routes/route_test.go b/pkg/scheduler/routes/route_test.go index 015c8f9f9b..15f25cfd84 100644 --- a/pkg/scheduler/routes/route_test.go +++ b/pkg/scheduler/routes/route_test.go @@ -20,6 +20,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "net/http/httptest" "strings" "testing" @@ -76,6 +77,22 @@ func TestMaxRequestSizeBind(t *testing.T) { } } +func TestMaxRequestSizePrioritize(t *testing.T) { + hugePayload := `{"Pod":{"metadata":{"name":"` + strings.Repeat("a", maxRequestSize+100) + `"}}}` + req := httptest.NewRequest("POST", "/prioritize", strings.NewReader(hugePayload)) + w := httptest.NewRecorder() + handler := PrioritizeRoute(&scheduler.Scheduler{}) + + handler(w, req, nil) + + if w.Code != 400 { + t.Fatalf("expected invalid oversized JSON to return 400, got %d", w.Code) + } + if !strings.Contains(w.Body.String(), "unexpected EOF") { + t.Fatalf("expected truncation error, got %q", w.Body.String()) + } +} + func TestWebHookRoute(t *testing.T) { handler := WebHookRoute() if handler == nil { @@ -212,3 +229,103 @@ func TestBind_DecodeError(t *testing.T) { t.Error("expected a decode error to be reported in the bind result") } } + +func TestPrioritizeRoute_DecodeError(t *testing.T) { + req := httptest.NewRequest("POST", "/prioritize", strings.NewReader("{not-json")) + w := httptest.NewRecorder() + + handler := PrioritizeRoute(&scheduler.Scheduler{}) + handler(w, req, nil) + + if w.Code != 400 { + t.Fatalf("expected 400 for invalid prioritize request, got %d", w.Code) + } + if !strings.Contains(w.Body.String(), "invalid character") { + t.Fatalf("expected decode error in response, got %q", w.Body.String()) + } +} + +func TestPrioritizeRoute_MissingPod(t *testing.T) { + req := httptest.NewRequest("POST", "/prioritize", strings.NewReader("{}")) + w := httptest.NewRecorder() + + handler := PrioritizeRoute(&scheduler.Scheduler{}) + handler(w, req, nil) + + if w.Code != 400 { + t.Fatalf("expected 400 for request without a pod, got %d", w.Code) + } + if !strings.Contains(w.Body.String(), "extender args must contain a pod") { + t.Fatalf("expected missing pod error in response, got %q", w.Body.String()) + } +} + +func TestPrioritizeRoute_CacheNotSynced(t *testing.T) { + prioritizeCalled := false + handler := prioritizeRoute( + func(context.Context) bool { return false }, + func(extenderv1.ExtenderArgs) (*extenderv1.HostPriorityList, error) { + prioritizeCalled = true + return nil, nil + }, + ) + req := httptest.NewRequest("POST", "/prioritize", strings.NewReader(`{"Pod":{"metadata":{"name":"test"}}}`)) + w := httptest.NewRecorder() + + handler(w, req, nil) + + if w.Code != 503 { + t.Fatalf("expected 503 while cache is not synced, got %d", w.Code) + } + if prioritizeCalled { + t.Fatal("prioritize should not be called before cache sync") + } +} + +func TestPrioritizeRoute_Success(t *testing.T) { + want := extenderv1.HostPriorityList{{Host: "node-a", Score: 7}} + handler := prioritizeRoute( + func(context.Context) bool { return true }, + func(args extenderv1.ExtenderArgs) (*extenderv1.HostPriorityList, error) { + if args.Pod == nil || args.Pod.Name != "test" { + t.Fatalf("unexpected pod in prioritize request: %#v", args.Pod) + } + return &want, nil + }, + ) + req := httptest.NewRequest("POST", "/prioritize", strings.NewReader(`{"Pod":{"metadata":{"name":"test"}}}`)) + w := httptest.NewRecorder() + + handler(w, req, nil) + + if w.Code != 200 { + t.Fatalf("expected 200 for successful prioritize request, got %d: %s", w.Code, w.Body.String()) + } + var got extenderv1.HostPriorityList + if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil { + t.Fatalf("decode prioritize response: %v", err) + } + if len(got) != 1 || got[0] != want[0] { + t.Fatalf("unexpected prioritize response: %#v", got) + } +} + +func TestPrioritizeRoute_SchedulerError(t *testing.T) { + handler := prioritizeRoute( + func(context.Context) bool { return true }, + func(extenderv1.ExtenderArgs) (*extenderv1.HostPriorityList, error) { + return nil, errors.New("scoring failed") + }, + ) + req := httptest.NewRequest("POST", "/prioritize", strings.NewReader(`{"Pod":{"metadata":{"name":"test"}}}`)) + w := httptest.NewRecorder() + + handler(w, req, nil) + + if w.Code != 500 { + t.Fatalf("expected 500 for prioritize failure, got %d", w.Code) + } + if !strings.Contains(w.Body.String(), "scoring failed") { + t.Fatalf("expected scoring error in response, got %q", w.Body.String()) + } +} diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index eded68aa6d..2cc15b2b36 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -20,8 +20,8 @@ import ( "context" "fmt" "maps" + "math" "os" - "sort" "strconv" "strings" "sync" @@ -31,6 +31,7 @@ import ( corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/labels" + k8stypes "k8s.io/apimachinery/pkg/types" "k8s.io/apimachinery/pkg/util/wait" "k8s.io/client-go/informers" "k8s.io/client-go/kubernetes" @@ -191,10 +192,6 @@ func (s *Scheduler) onDelPod(obj any) { return } - _, ok = pod.Annotations[util.AssignedNodeAnnotations] - if !ok { - return - } if pi, ok := s.podManager.TakeAndDeletePod(pod); ok { s.quotaManager.RmUsage(pod, pi.Devices) } @@ -621,6 +618,10 @@ func nodeListLen(nodes *corev1.NodeList) int { // returns all nodes and its device memory usage, and we filter it with nodeSelector, taints, nodeAffinity // unschedulerable and nodeName. func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[string]*NodeUsage, *map[string]*NodeUsage, map[string]string, error) { + return s.getNodesUsageIgnoringPod(nodes, task, nil) +} + +func (s *Scheduler) getNodesUsageIgnoringPod(nodes *[]string, task *corev1.Pod, ignoredPodUID *k8stypes.UID) (*map[string]*NodeUsage, *map[string]*NodeUsage, map[string]string, error) { overallnodeMap := make(map[string]*NodeUsage) cachenodeMap := make(map[string]*NodeUsage) failedNodes := make(map[string]string) @@ -635,6 +636,9 @@ func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[strin podsInfo := s.podManager.ListPodsInfo() for _, p := range podsInfo { + if ignoredPodUID != nil && p.UID == *ignoredPodUID { + continue + } node, ok := overallnodeMap[p.NodeID] if !ok { klog.V(5).InfoS("pod allocated unknown node resources", @@ -839,6 +843,19 @@ func (s *Scheduler) Bind(args extenderv1.ExtenderBindingArgs) (*extenderv1.Exten }) return &extenderv1.ExtenderBindingResult{Error: err.Error()}, err } + if current.UID != args.PodUID { + err := fmt.Errorf("pod UID mismatch for %s/%s: cached %s, binding request %s", args.PodNamespace, args.PodName, current.UID, args.PodUID) + // The Pod name was reused before this binding request completed. Remove + // only the stale UID's reservation; never mutate the replacement Pod. + stalePod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + Name: args.PodName, + Namespace: args.PodNamespace, + UID: args.PodUID, + }} + s.cleanupStalePodAllocation(stalePod) + s.recordScheduleBindingResultEvent(stalePod, EventReasonBindingFailed, []string{}, err) + return &extenderv1.ExtenderBindingResult{Error: err.Error()}, nil + } klog.InfoS("Trying to get the target node for pod", "pod", args.PodName, "namespace", args.PodNamespace, "node", args.Node) @@ -851,11 +868,6 @@ func (s *Scheduler) Bind(args extenderv1.ExtenderBindingArgs) (*extenderv1.Exten return res, nil } - tmppatch := map[string]string{ - util.DeviceBindPhase: "allocating", - util.BindTimeAnnotations: strconv.FormatInt(time.Now().Unix(), 10), - } - fail := func(e error) (*extenderv1.ExtenderBindingResult, error) { klog.InfoS("Release node locks", "node", args.Node) s.releaseAllDevices(node, current) @@ -872,11 +884,66 @@ func (s *Scheduler) Bind(args extenderv1.ExtenderBindingArgs) (*extenderv1.Exten return fail(err) } - if err = util.PatchPodAnnotations(current, tmppatch); err != nil { + // Filter and Prioritize do not commit a device allocation. Revalidate the + // node selected by kube-scheduler while holding the device lock, then commit + // the concrete device allocation immediately before binding. + resourceReqs := device.Resourcereqs(current) + annotations := map[string]string{ + util.DeviceBindPhase: util.DeviceBindAllocating, + util.BindTimeAnnotations: strconv.FormatInt(time.Now().Unix(), 10), + } + var allocation *policy.NodeScore + var previousAllocation *device.PodInfo + if countDeviceRequests(resourceReqs) > 0 { + if existing, ok := s.podManager.GetPod(current); ok { + previousAllocation = existing + } + nodeNames := []string{args.Node} + ignoredPodUID := current.UID + nodeUsage, _, failedNodes, err := s.getNodesUsageIgnoringPod(&nodeNames, current, &ignoredPodUID) + if err != nil { + return fail(err) + } + nodeScores, err := s.calcScoreWithOptions(nodeUsage, resourceReqs, current, failedNodes, false, false) + if err != nil { + return fail(fmt.Errorf("failed to revalidate node %s: %w", args.Node, err)) + } + if len(nodeScores.NodeList) != 1 { + reason := failedNodes[args.Node] + if reason == "" { + reason = "node is no longer feasible" + } + return fail(fmt.Errorf("failed to revalidate node %s: %s", args.Node, reason)) + } + + allocation = nodeScores.NodeList[0] + annotations[util.AssignedNodeAnnotations] = allocation.NodeID + annotations[util.AssignedTimeAnnotations] = strconv.FormatInt(time.Now().Unix(), 10) + for _, dev := range device.GetDevices() { + dev.PatchAnnotations(current, &annotations, allocation.Devices) + } + + var reserved bool + if previousAllocation == nil { + reserved = s.podManager.ReservePodIfAbsent(current, allocation.NodeID, allocation.Devices) + } else { + reserved = s.podManager.ReplacePodReservation(current, previousAllocation, allocation.NodeID, allocation.Devices) + } + if !reserved { + return fail(fmt.Errorf("pod allocation already exists for UID %s", current.UID)) + } + if previousAllocation != nil { + s.quotaManager.RmUsage(current, previousAllocation.Devices) + } + s.quotaManager.AddUsage(current, allocation.Devices) + } + if err = util.PatchPodAnnotations(current, annotations); err != nil { klog.ErrorS(err, "Failed to patch pod annotations", "pod", klog.KObj(current)) + // A patch error is ambiguous: the API server may have persisted the + // allocation before the response was lost. Keep the reservation until an + // informer update confirms it or the next scheduling attempt replaces it. return fail(err) } - if err = s.kubeClient.CoreV1().Pods(args.PodNamespace).Bind(context.Background(), binding, metav1.CreateOptions{}); err != nil { klog.ErrorS(err, "Failed to bind pod", "pod", args.PodName, "namespace", args.PodNamespace, "node", args.Node) return fail(err) @@ -890,13 +957,7 @@ 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) - } - } - if resourceReqTotal == 0 { + if countDeviceRequests(resourceReqs) == 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")) @@ -952,39 +1013,55 @@ func (s *Scheduler) Filter(args extenderv1.ExtenderArgs) (*extenderv1.ExtenderFi }, nil } klog.V(4).Infoln("nodeScores_len=", len((*nodeScores).NodeList)) - sort.Sort(nodeScores) - m := (*nodeScores).NodeList[len((*nodeScores).NodeList)-1] - klog.InfoS("Scheduling pod to node", - "podNamespace", args.Pod.Namespace, - "podName", args.Pod.Name, - "nodeID", m.NodeID, - "devices", m.Devices) - annotations := make(map[string]string) - annotations[util.AssignedNodeAnnotations] = m.NodeID - annotations[util.AssignedTimeAnnotations] = strconv.FormatInt(time.Now().Unix(), 10) - - for _, val := range device.GetDevices() { - val.PatchAnnotations(args.Pod, &annotations, m.Devices) + fit := make(map[string]struct{}, len(nodeScores.NodeList)) + for _, nodeScore := range nodeScores.NodeList { + fit[nodeScore.NodeID] = struct{}{} + } + feasibleNodes := make([]string, 0, len(fit)) + for _, nodeName := range *args.NodeNames { + if _, ok := fit[nodeName]; ok { + feasibleNodes = append(feasibleNodes, nodeName) + } } + successMsg := genFilterSuccessMsg(len(*args.NodeNames), nodeScores.NodeList) + s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringSucceed, successMsg, nil) + res := extenderv1.ExtenderFilterResult{NodeNames: &feasibleNodes} + if len(failedNodes) > 0 { + res.FailedNodes = failedNodes + } + return &res, nil +} - added := s.podManager.AddPod(args.Pod, m.NodeID, m.Devices) - if added { - s.quotaManager.AddUsage(args.Pod, m.Devices) +// Prioritize exposes HAMi's node policy to kube-scheduler without reserving +// devices or changing the Pod. Higher extender scores are always preferred, +// so spread policy reverses HAMi's internal utilization score before the +// values are normalized to the extender range. +func (s *Scheduler) Prioritize(args extenderv1.ExtenderArgs) (*extenderv1.HostPriorityList, error) { + resourceReqs := device.Resourcereqs(args.Pod) + if countDeviceRequests(resourceReqs) == 0 { + return zeroHostPriorities(args), nil } - err = util.PatchPodAnnotations(args.Pod, annotations) - if err != nil { - s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, "", err) - if added { - s.quotaManager.RmUsage(args.Pod, m.Devices) + var failedNodes map[string]string + var nodeScores *policy.NodeScoreList + var err error + if args.Nodes != nil { + var nodeUsage *map[string]*NodeUsage + nodeUsage, failedNodes, err = s.getSimulationNodesUsage(args.Nodes, args.Pod) + if err == nil { + nodeScores, err = s.calcScoreWithOptions(nodeUsage, resourceReqs, args.Pod, failedNodes, false, true) + } + } else { + var nodeUsage *map[string]*NodeUsage + nodeUsage, _, failedNodes, err = s.getNodesUsage(args.NodeNames, args.Pod) + if err == nil { + nodeScores, err = s.calcScoreWithOptions(nodeUsage, resourceReqs, args.Pod, failedNodes, false, false) } - 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}} - return &res, nil + if err != nil { + return nil, fmt.Errorf("calcScore failed %v for pod %v", err, args.Pod.Name) + } + return normalizeHostPriorities(nodeScores), nil } func (s *Scheduler) filterSimulation(args extenderv1.ExtenderArgs, resourceReqs device.PodDeviceRequests) (*extenderv1.ExtenderFilterResult, error) { @@ -1011,18 +1088,18 @@ func (s *Scheduler) filterSimulation(args extenderv1.ExtenderArgs, resourceReqs FailedNodes: failedNodes, }, nil } - sort.Sort(nodeScores) - bestNodeID := nodeScores.NodeList[len(nodeScores.NodeList)-1].NodeID - filteredNodes := make([]corev1.Node, 0, 1) + fit := make(map[string]struct{}, len(nodeScores.NodeList)) + for _, nodeScore := range nodeScores.NodeList { + fit[nodeScore.NodeID] = struct{}{} + } + filteredNodes := make([]corev1.Node, 0, len(fit)) for i := range args.Nodes.Items { - if args.Nodes.Items[i].Name == bestNodeID { + if _, ok := fit[args.Nodes.Items[i].Name]; ok { filteredNodes = append(filteredNodes, *args.Nodes.Items[i].DeepCopy()) - break } } - klog.V(2).InfoS("Simulation filter selected best node", + klog.V(2).InfoS("Simulation filter returned feasible nodes", "pod", klog.KObj(args.Pod), - "selectedNode", bestNodeID, "filteredNodesLen", len(filteredNodes)) return &extenderv1.ExtenderFilterResult{ Nodes: &corev1.NodeList{ @@ -1032,12 +1109,72 @@ func (s *Scheduler) filterSimulation(args extenderv1.ExtenderArgs, resourceReqs }, nil } -func genSuccessMsg(totalNodes int, target string, nodes []*policy.NodeScore) string { - successMsg := "find fit node(%s), %d nodes not fit, %d nodes fit(%s)" +func countDeviceRequests(resourceReqs device.PodDeviceRequests) int { + total := 0 + for _, containerReqs := range resourceReqs { + for _, req := range containerReqs { + total += int(req.Nums) + } + } + return total +} + +func zeroHostPriorities(args extenderv1.ExtenderArgs) *extenderv1.HostPriorityList { + priorities := make(extenderv1.HostPriorityList, 0) + if args.NodeNames != nil { + priorities = make(extenderv1.HostPriorityList, 0, len(*args.NodeNames)) + for _, nodeName := range *args.NodeNames { + priorities = append(priorities, extenderv1.HostPriority{Host: nodeName}) + } + return &priorities + } + if args.Nodes != nil { + priorities = make(extenderv1.HostPriorityList, 0, len(args.Nodes.Items)) + for i := range args.Nodes.Items { + priorities = append(priorities, extenderv1.HostPriority{Host: args.Nodes.Items[i].Name}) + } + } + return &priorities +} + +func normalizeHostPriorities(nodeScores *policy.NodeScoreList) *extenderv1.HostPriorityList { + priorities := make(extenderv1.HostPriorityList, 0, len(nodeScores.NodeList)) + if len(nodeScores.NodeList) == 0 { + return &priorities + } + + desirability := make([]float64, len(nodeScores.NodeList)) + minScore, maxScore := float64(0), float64(0) + for i, nodeScore := range nodeScores.NodeList { + score := float64(nodeScore.Score) + if nodeScores.Policy == util.NodeSchedulerPolicySpread.String() { + score = -score + } + desirability[i] = score + if i == 0 || score < minScore { + minScore = score + } + if i == 0 || score > maxScore { + maxScore = score + } + } + + for i, nodeScore := range nodeScores.NodeList { + score := int64(0) + if maxScore > minScore { + score = int64(math.Round((desirability[i] - minScore) * float64(extenderv1.MaxExtenderPriority) / (maxScore - minScore))) + } + priorities = append(priorities, extenderv1.HostPriority{Host: nodeScore.NodeID, Score: score}) + } + return &priorities +} + +func genFilterSuccessMsg(totalNodes int, nodes []*policy.NodeScore) string { + successMsg := "%d nodes not fit, %d nodes fit(%s)" var scores []string for _, no := range nodes { scores = append(scores, fmt.Sprintf("%s:%.2f", no.NodeID, no.Score)) } score := strings.Join(scores, ",") - return fmt.Sprintf(successMsg, target, totalNodes-len(nodes), len(nodes), score) + return fmt.Sprintf(successMsg, totalNodes-len(nodes), len(nodes), score) } diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go index 2687cfef22..4ebcfb93f0 100644 --- a/pkg/scheduler/scheduler_test.go +++ b/pkg/scheduler/scheduler_test.go @@ -36,6 +36,7 @@ import ( "k8s.io/apimachinery/pkg/util/wait" "k8s.io/client-go/informers" "k8s.io/client-go/kubernetes/fake" + k8stesting "k8s.io/client-go/testing" "k8s.io/client-go/tools/cache" "k8s.io/client-go/tools/record" "k8s.io/klog/v2" @@ -45,6 +46,7 @@ import ( "github.com/Project-HAMi/HAMi/pkg/device/common" "github.com/Project-HAMi/HAMi/pkg/device/nvidia" "github.com/Project-HAMi/HAMi/pkg/scheduler/config" + "github.com/Project-HAMi/HAMi/pkg/scheduler/policy" "github.com/Project-HAMi/HAMi/pkg/util" "github.com/Project-HAMi/HAMi/pkg/util/client" nodelockutil "github.com/Project-HAMi/HAMi/pkg/util/nodelock" @@ -685,11 +687,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 + wantPriorityNode string + wantErr error }{ { name: "node use binpack gpu use binpack policy", @@ -724,9 +726,9 @@ func Test_Filter(t *testing.T) { }, wantErr: nil, want: &extenderv1.ExtenderFilterResult{ - NodeNames: &[]string{"node2"}, + NodeNames: &[]string{"node1", "node2"}, }, - wantPodAnnotationDeviceID: "device4", + wantPriorityNode: "node2", }, { name: "node use binpack gpu use spread policy", @@ -761,9 +763,9 @@ func Test_Filter(t *testing.T) { }, wantErr: nil, want: &extenderv1.ExtenderFilterResult{ - NodeNames: &[]string{"node2"}, + NodeNames: &[]string{"node1", "node2"}, }, - wantPodAnnotationDeviceID: "device3", + wantPriorityNode: "node2", }, { name: "node use spread gpu use binpack policy", @@ -798,9 +800,9 @@ func Test_Filter(t *testing.T) { }, wantErr: nil, want: &extenderv1.ExtenderFilterResult{ - NodeNames: &[]string{"node1"}, + NodeNames: &[]string{"node1", "node2"}, }, - wantPodAnnotationDeviceID: "device1", + wantPriorityNode: "node1", }, { name: "node use spread gpu use spread policy", @@ -835,9 +837,9 @@ func Test_Filter(t *testing.T) { }, wantErr: nil, want: &extenderv1.ExtenderFilterResult{ - NodeNames: &[]string{"node1"}, + NodeNames: &[]string{"node1", "node2"}, }, - wantPodAnnotationDeviceID: "device2", + wantPriorityNode: "node1", }, } @@ -849,10 +851,460 @@ func Test_Filter(t *testing.T) { assert.DeepEqual(t, test.wantErr, gotErr) 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) + assert.Equal(t, "", getPod.Annotations[util.AssignedNodeAnnotations]) + _, reserved := s.podManager.GetPod(test.args.Pod) + assert.Equal(t, false, reserved) + + priorities, err := s.Prioritize(test.args) + require.NoError(t, err) + require.Len(t, *priorities, 2) + best := (*priorities)[0] + for _, priority := range (*priorities)[1:] { + if priority.Score > best.Score { + best = priority + } + } + assert.Equal(t, test.wantPriorityNode, best.Host) + assert.Equal(t, extenderv1.MaxExtenderPriority, best.Score) + }) + } +} + +func TestNormalizeHostPriorities(t *testing.T) { + tests := []struct { + name string + policy string + want map[string]int64 + }{ + {name: "binpack prefers higher utilization", policy: util.NodeSchedulerPolicyBinpack.String(), want: map[string]int64{"node-low": 0, "node-high": 10}}, + {name: "spread prefers lower utilization", policy: util.NodeSchedulerPolicySpread.String(), want: map[string]int64{"node-low": 10, "node-high": 0}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + priorities := normalizeHostPriorities(&policy.NodeScoreList{ + Policy: test.policy, + NodeList: []*policy.NodeScore{ + {NodeID: "node-low", Score: 2}, + {NodeID: "node-high", Score: 8}, + }, + }) + got := make(map[string]int64, len(*priorities)) + for _, priority := range *priorities { + got[priority.Host] = priority.Score + } + assert.DeepEqual(t, test.want, got) }) } + + equal := normalizeHostPriorities(&policy.NodeScoreList{ + Policy: util.NodeSchedulerPolicyBinpack.String(), + NodeList: []*policy.NodeScore{ + {NodeID: "node-a", Score: 3}, + {NodeID: "node-b", Score: 3}, + }, + }) + require.Len(t, *equal, 2) + assert.Equal(t, int64(0), (*equal)[0].Score) + assert.Equal(t, int64(0), (*equal)[1].Score) +} + +func TestPrioritizeWithoutDeviceRequests(t *testing.T) { + s := &Scheduler{} + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "no-devices"}} + + t.Run("node names", func(t *testing.T) { + nodeNames := []string{"node-a", "node-b"} + priorities, err := s.Prioritize(extenderv1.ExtenderArgs{Pod: pod, NodeNames: &nodeNames}) + require.NoError(t, err) + require.Len(t, *priorities, 2) + assert.Equal(t, "node-a", (*priorities)[0].Host) + assert.Equal(t, int64(0), (*priorities)[0].Score) + assert.Equal(t, "node-b", (*priorities)[1].Host) + assert.Equal(t, int64(0), (*priorities)[1].Score) + }) + + t.Run("full nodes", func(t *testing.T) { + nodes := &corev1.NodeList{Items: []corev1.Node{ + {ObjectMeta: metav1.ObjectMeta{Name: "node-a"}}, + {ObjectMeta: metav1.ObjectMeta{Name: "node-b"}}, + }} + priorities, err := s.Prioritize(extenderv1.ExtenderArgs{Pod: pod, Nodes: nodes}) + require.NoError(t, err) + require.Len(t, *priorities, 2) + assert.Equal(t, "node-a", (*priorities)[0].Host) + assert.Equal(t, int64(0), (*priorities)[0].Score) + assert.Equal(t, "node-b", (*priorities)[1].Host) + assert.Equal(t, int64(0), (*priorities)[1].Score) + }) +} + +type bindAllocationMockDevice struct { + registerMockDevice + fit bool +} + +const bindAllocationMockAnnotation = "hami.io/mock-devices-allocated" + +func (m *bindAllocationMockDevice) CommonWord() string { return "bind-allocation-mock" } +func (m *bindAllocationMockDevice) GetResourceNames() device.ResourceNames { + return device.ResourceNames{ + ResourceCountName: "example.com/mock", + ResourceMemoryName: "example.com/mock-memory", + ResourceCoreName: "example.com/mock-cores", + } +} +func (m *bindAllocationMockDevice) GenerateResourceRequests(ctr *corev1.Container) device.ContainerDeviceRequest { + if ctr.Resources.Limits.Name("example.com/mock", resource.DecimalSI).Value() == 0 { + return device.ContainerDeviceRequest{} + } + return device.ContainerDeviceRequest{Nums: 1, Type: m.CommonWord(), Memreq: 1, Coresreq: 1} +} +func (m *bindAllocationMockDevice) Fit(_ []*device.DeviceUsage, request device.ContainerDeviceRequest, _ *corev1.Pod, _ *device.NodeInfo, _ *device.PodDevices) (bool, map[string]device.ContainerDevices, string) { + if !m.fit { + return false, nil, "mock capacity exhausted" + } + return true, map[string]device.ContainerDevices{ + m.CommonWord(): {{UUID: "mock-device-0", Type: m.CommonWord(), Usedmem: request.Memreq, Usedcores: request.Coresreq}}, + }, "" +} +func (m *bindAllocationMockDevice) PatchAnnotations(_ *corev1.Pod, annotations *map[string]string, devices device.PodDevices) map[string]string { + (*annotations)[bindAllocationMockAnnotation] = device.EncodePodSingleDevice(devices[m.CommonWord()]) + return *annotations +} + +func setupBindAllocationTest(t *testing.T, mock *bindAllocationMockDevice) (*Scheduler, *fake.Clientset, *corev1.Pod, *corev1.Node) { + t.Helper() + oldDevicesMap := device.DevicesMap + device.DevicesMap = map[string]device.Devices{"bind-allocation-mock": mock} + t.Cleanup(func() { device.DevicesMap = oldDevicesMap }) + oldSupportDevices := device.SupportDevices + device.SupportDevices = maps.Clone(device.SupportDevices) + device.SupportDevices[mock.CommonWord()] = bindAllocationMockAnnotation + t.Cleanup(func() { device.SupportDevices = oldSupportDevices }) + + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "bind-allocation", Namespace: "bind-allocation-test", UID: "bind-allocation-uid"}, + Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Name: "worker", + Resources: corev1.ResourceRequirements{Limits: corev1.ResourceList{ + "example.com/mock": *resource.NewQuantity(1, resource.DecimalSI), + }}, + }}}, + } + node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node-selected"}} + fakeClient := fake.NewSimpleClientset(pod, node) + fakeClient.PrependReactor("create", "pods", func(action k8stesting.Action) (bool, runtime.Object, error) { + if action.GetSubresource() == "binding" { + return true, &corev1.Binding{}, nil + } + return false, nil, nil + }) + oldKubeClient := client.KubeClient + t.Cleanup(func() { client.KubeClient = oldKubeClient }) + client.KubeClient = fakeClient + + s := NewScheduler() + s.kubeClient = fakeClient + t.Cleanup(func() { + if allocation, ok := s.podManager.TakeAndDeletePod(pod); ok { + s.quotaManager.RmUsage(pod, allocation.Devices) + } + }) + scheme := runtime.NewScheme() + require.NoError(t, corev1.AddToScheme(scheme)) + s.eventRecorder = record.NewBroadcaster().NewRecorder(scheme, corev1.EventSource{}) + informerFactory := informers.NewSharedInformerFactoryWithOptions(fakeClient, time.Hour) + require.NoError(t, informerFactory.Core().V1().Pods().Informer().GetIndexer().Add(pod)) + require.NoError(t, informerFactory.Core().V1().Nodes().Informer().GetIndexer().Add(node)) + s.podLister = informerFactory.Core().V1().Pods().Lister() + s.nodeLister = informerFactory.Core().V1().Nodes().Lister() + s.addNode(node.Name, &device.NodeInfo{ + ID: node.Name, + Node: node, + Devices: map[string][]device.DeviceInfo{ + "bind-allocation-mock": {{ + ID: "mock-device-0", Type: "bind-allocation-mock", DeviceVendor: "bind-allocation-mock", + Count: 10, Devmem: 100, Devcore: 100, Health: true, + }}, + }, + }) + return s, fakeClient, pod, node +} + +func TestBindRevalidatesAndAllocatesSelectedNode(t *testing.T) { + s, fakeClient, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: true}) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: pod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Empty(t, res.Error) + updated, err := fakeClient.CoreV1().Pods(pod.Namespace).Get(context.Background(), pod.Name, metav1.GetOptions{}) + require.NoError(t, err) + assert.Equal(t, node.Name, updated.Annotations[util.AssignedNodeAnnotations]) + assert.Equal(t, "mock-device-0,bind-allocation-mock,1,1:;", updated.Annotations[bindAllocationMockAnnotation]) + allocation, ok := s.podManager.GetPod(pod) + require.True(t, ok) + assert.Equal(t, node.Name, allocation.NodeID) +} + +func TestBindRejectsNodeThatIsNoLongerFeasible(t *testing.T) { + s, fakeClient, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: false}) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: pod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Contains(t, res.Error, "failed to revalidate node") + updated, err := fakeClient.CoreV1().Pods(pod.Namespace).Get(context.Background(), pod.Name, metav1.GetOptions{}) + require.NoError(t, err) + assert.Equal(t, "", updated.Annotations[util.AssignedNodeAnnotations]) + _, allocated := s.podManager.GetPod(pod) + assert.Equal(t, false, allocated) +} + +func TestBindRevalidationFailureRestoresPreviousAllocation(t *testing.T) { + s, _, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: false}) + previousDevices := device.PodDevices{ + "bind-allocation-mock": {{{ + UUID: "previous-device", Type: "bind-allocation-mock", Usedmem: 1, Usedcores: 1, + }}}, + } + reserved := s.podManager.ReservePodIfAbsent(pod, node.Name, previousDevices) + require.True(t, reserved) + s.quotaManager.AddUsage(pod, previousDevices) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: pod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Contains(t, res.Error, "failed to revalidate node") + allocation, allocated := s.podManager.GetPod(pod) + require.True(t, allocated) + assert.Equal(t, node.Name, allocation.NodeID) + assert.DeepEqual(t, previousDevices, allocation.Devices) + quota := s.quotaManager.GetResourceQuota()[pod.Namespace] + assert.Assert(t, quota != nil) + assert.Equal(t, int64(1), (*quota)["example.com/mock-memory"].Used) + assert.Equal(t, int64(1), (*quota)["example.com/mock-cores"].Used) + stale := pod.DeepCopy() + stale.Annotations = map[string]string{ + util.AssignedNodeAnnotations: "stale-node", + bindAllocationMockAnnotation: "stale-device,bind-allocation-mock,1,1:;", + } + s.onAddPod(stale) + allocation, allocated = s.podManager.GetPod(pod) + require.True(t, allocated) + assert.Equal(t, node.Name, allocation.NodeID) + assert.DeepEqual(t, previousDevices, allocation.Devices) +} + +func TestBindRejectsMismatchedPodUIDBeforeMutation(t *testing.T) { + s, fakeClient, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: true}) + stalePod := pod.DeepCopy() + stalePod.UID = "stale-binding-uid" + staleDevices := device.PodDevices{ + "bind-allocation-mock": {{{ + UUID: "mock-device-0", Type: "bind-allocation-mock", Usedmem: 1, Usedcores: 1, + }}}, + } + reserved := s.podManager.ReservePodIfAbsent(stalePod, node.Name, staleDevices) + require.True(t, reserved) + s.quotaManager.AddUsage(stalePod, staleDevices) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: stalePod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Contains(t, res.Error, "pod UID mismatch") + updated, err := fakeClient.CoreV1().Pods(pod.Namespace).Get(context.Background(), pod.Name, metav1.GetOptions{}) + require.NoError(t, err) + assert.Equal(t, "", updated.Annotations[util.AssignedNodeAnnotations]) + _, staleAllocated := s.podManager.GetPod(stalePod) + assert.Equal(t, false, staleAllocated) + _, replacementAllocated := s.podManager.GetPod(pod) + assert.Equal(t, false, replacementAllocated) + quota := s.quotaManager.GetResourceQuota()[pod.Namespace] + assert.Assert(t, quota != nil) + assert.Equal(t, int64(0), (*quota)["example.com/mock-memory"].Used) + assert.Equal(t, int64(0), (*quota)["example.com/mock-cores"].Used) + for _, action := range fakeClient.Actions() { + assert.Assert(t, action.GetVerb() != "patch", "UID mismatch must not patch the replacement Pod") + assert.Assert(t, action.GetSubresource() != "binding", "UID mismatch must not call the binding API") + } +} + +func TestBindPatchFailurePreservesOwnedAllocation(t *testing.T) { + s, fakeClient, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: true}) + fakeClient.PrependReactor("patch", "pods", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, fmt.Errorf("injected patch failure") + }) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: pod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Contains(t, res.Error, "injected patch failure") + allocation, allocated := s.podManager.GetPod(pod) + require.True(t, allocated) + assert.Equal(t, node.Name, allocation.NodeID) + quota := s.quotaManager.GetResourceQuota()[pod.Namespace] + assert.Assert(t, quota != nil) + assert.Equal(t, int64(1), (*quota)["example.com/mock-memory"].Used) + assert.Equal(t, int64(1), (*quota)["example.com/mock-cores"].Used) + for _, action := range fakeClient.Actions() { + assert.Assert(t, action.GetSubresource() != "binding", "annotation failure must stop before the binding API") + } +} + +func TestBindPatchFailurePreservesReplacementAllocation(t *testing.T) { + s, fakeClient, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: true}) + previousDevices := device.PodDevices{ + "bind-allocation-mock": {{{ + UUID: "previous-device", Type: "bind-allocation-mock", Usedmem: 1, Usedcores: 1, + }}}, + } + s.podManager.AddPod(pod, "previous-node", previousDevices) + s.quotaManager.AddUsage(pod, previousDevices) + fakeClient.PrependReactor("patch", "pods", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, fmt.Errorf("injected patch failure") + }) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: pod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Contains(t, res.Error, "injected patch failure") + allocation, allocated := s.podManager.GetPod(pod) + require.True(t, allocated) + assert.Equal(t, node.Name, allocation.NodeID) + assert.Equal(t, "mock-device-0", allocation.Devices["bind-allocation-mock"][0][0].UUID) + quota := s.quotaManager.GetResourceQuota()[pod.Namespace] + assert.Assert(t, quota != nil) + assert.Equal(t, int64(1), (*quota)["example.com/mock-memory"].Used) + assert.Equal(t, int64(1), (*quota)["example.com/mock-cores"].Used) +} + +func TestBindPatchFailureAcceptsDelayedInformerObservation(t *testing.T) { + s, fakeClient, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: true}) + previousDevices := device.PodDevices{ + "bind-allocation-mock": {{{ + UUID: "previous-device", Type: "bind-allocation-mock", Usedmem: 1, Usedcores: 1, + }}}, + } + s.podManager.AddPod(pod, "previous-node", previousDevices) + s.quotaManager.AddUsage(pod, previousDevices) + fakeClient.PrependReactor("patch", "pods", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, fmt.Errorf("injected ambiguous patch failure") + }) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: pod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Contains(t, res.Error, "injected ambiguous patch failure") + + observedDevices := device.PodDevices{ + "bind-allocation-mock": {{{ + UUID: "mock-device-0", Type: "bind-allocation-mock", Usedmem: 1, Usedcores: 1, + }}}, + } + observed := pod.DeepCopy() + observed.Annotations = map[string]string{ + util.AssignedNodeAnnotations: node.Name, + bindAllocationMockAnnotation: device.EncodePodSingleDevice(observedDevices["bind-allocation-mock"]), + } + s.onAddPod(observed) + + allocation, allocated := s.podManager.GetPod(pod) + require.True(t, allocated) + assert.Equal(t, node.Name, allocation.NodeID) + assert.DeepEqual(t, observedDevices, allocation.Devices) + quota := s.quotaManager.GetResourceQuota()[pod.Namespace] + assert.Assert(t, quota != nil) + assert.Equal(t, int64(1), (*quota)["example.com/mock-memory"].Used) + assert.Equal(t, int64(1), (*quota)["example.com/mock-cores"].Used) +} + +func TestOnDelPodCleansUnconfirmedBindReservation(t *testing.T) { + s, fakeClient, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: true}) + fakeClient.PrependReactor("patch", "pods", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, fmt.Errorf("injected patch failure") + }) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: pod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Contains(t, res.Error, "injected patch failure") + s.onDelPod(pod) + + _, allocated := s.podManager.GetPod(pod) + assert.Equal(t, false, allocated) + quota := s.quotaManager.GetResourceQuota()[pod.Namespace] + assert.Assert(t, quota != nil) + assert.Equal(t, int64(0), (*quota)["example.com/mock-memory"].Used) + assert.Equal(t, int64(0), (*quota)["example.com/mock-cores"].Used) +} + +func TestBindPatchFailurePreservesInformerObservedAllocation(t *testing.T) { + s, fakeClient, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: true}) + devices := device.PodDevices{ + "bind-allocation-mock": {{{ + UUID: "mock-device-0", Type: "bind-allocation-mock", Usedmem: 1, Usedcores: 1, + }}}, + } + fakeClient.PrependReactor("patch", "pods", func(k8stesting.Action) (bool, runtime.Object, error) { + observed := pod.DeepCopy() + observed.Annotations = map[string]string{ + util.AssignedNodeAnnotations: node.Name, + bindAllocationMockAnnotation: device.EncodePodSingleDevice(devices["bind-allocation-mock"]), + } + s.onAddPod(observed) + return true, nil, fmt.Errorf("injected ambiguous patch failure") + }) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: pod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Contains(t, res.Error, "injected ambiguous patch failure") + allocation, allocated := s.podManager.GetPod(pod) + require.True(t, allocated) + assert.Equal(t, node.Name, allocation.NodeID) + assert.DeepEqual(t, devices, allocation.Devices) + quota := s.quotaManager.GetResourceQuota()[pod.Namespace] + assert.Assert(t, quota != nil) + assert.Equal(t, int64(1), (*quota)["example.com/mock-memory"].Used) + assert.Equal(t, int64(1), (*quota)["example.com/mock-cores"].Used) +} + +func TestBindFailurePreservesAllocationForRetry(t *testing.T) { + s, fakeClient, pod, node := setupBindAllocationTest(t, &bindAllocationMockDevice{fit: true}) + fakeClient.PrependReactor("create", "pods", func(action k8stesting.Action) (bool, runtime.Object, error) { + if action.GetSubresource() == "binding" { + return true, nil, fmt.Errorf("injected bind failure") + } + return false, nil, nil + }) + + res, err := s.Bind(extenderv1.ExtenderBindingArgs{ + PodName: pod.Name, PodNamespace: pod.Namespace, PodUID: pod.UID, Node: node.Name, + }) + require.NoError(t, err) + require.Contains(t, res.Error, "injected bind failure") + allocation, allocated := s.podManager.GetPod(pod) + require.True(t, allocated) + assert.Equal(t, node.Name, allocation.NodeID) + quota := s.quotaManager.GetResourceQuota()[pod.Namespace] + assert.Assert(t, quota != nil) + assert.Equal(t, int64(1), (*quota)["example.com/mock-memory"].Used) + assert.Equal(t, int64(1), (*quota)["example.com/mock-cores"].Used) + updated, err := fakeClient.CoreV1().Pods(pod.Namespace).Get(context.Background(), pod.Name, metav1.GetOptions{}) + require.NoError(t, err) + assert.Equal(t, node.Name, updated.Annotations[util.AssignedNodeAnnotations]) + assert.Equal(t, "mock-device-0,bind-allocation-mock,1,1:;", updated.Annotations[bindAllocationMockAnnotation]) + assert.Equal(t, node.Name, updated.Labels[util.AssignedNodeAnnotations]) } func TestSchedulerOnDelNodeCleansLockDirectNode(t *testing.T) {