fix(scheduler): add prioritize extender and defer allocation to bind

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,

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())

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()