[YUNIKORN-2978] Fix handling of reserved allocations where node differs

pkg/scheduler/objects/node.go

@@ -344,6 +344,12 @@ func (sn *Node) RemoveAllocation(allocationKey string) *Allocation {
 		}
 		sn.availableResource.AddTo(alloc.GetAllocatedResource())
 		sn.nodeEvents.SendAllocationRemovedEvent(sn.NodeID, alloc.allocationKey, alloc.GetAllocatedResource())
+		log.Log(log.SchedNode).Info("node allocation removed",
+			zap.String("appID", alloc.GetApplicationID()),
+			zap.String("allocationKey", alloc.GetAllocationKey()),
+			zap.Stringer("allocatedResource", alloc.GetAllocatedResource()),
+			zap.Bool("placeholder", alloc.IsPlaceholder()),
+			zap.String("targetNode", sn.NodeID))
 		return alloc
 	}
 
@@ -382,6 +388,10 @@ func (sn *Node) UpdateForeignAllocation(alloc *Allocation) *Allocation {
 	delta := resources.Sub(newResource, existingResource)
 	delta.Prune()
 
+	log.Log(log.SchedNode).Info("node foreign allocation updated",
+		zap.String("allocationKey", alloc.GetAllocationKey()),
+		zap.Stringer("deltaResource", delta),
+		zap.String("targetNode", sn.NodeID))
 	sn.occupiedResource.AddTo(delta)
 	sn.occupiedResource.Prune()
 	sn.refreshAvailableResource()
@@ -416,6 +426,12 @@ func (sn *Node) addAllocationInternal(alloc *Allocation, force bool) bool {
 		sn.availableResource.SubFrom(res)
 		sn.availableResource.Prune()
 		sn.nodeEvents.SendAllocationAddedEvent(sn.NodeID, alloc.allocationKey, res)
+		log.Log(log.SchedNode).Info("node allocation processed",
+			zap.String("appID", alloc.GetApplicationID()),
+			zap.String("allocationKey", alloc.GetAllocationKey()),
+			zap.Stringer("allocatedResource", alloc.GetAllocatedResource()),
+			zap.Bool("placeholder", alloc.IsPlaceholder()),
+			zap.String("targetNode", sn.NodeID))
 		result = true
 		return result
 	}
@@ -440,6 +456,12 @@ func (sn *Node) ReplaceAllocation(allocationKey string, replace *Allocation, del
 	sn.allocatedResource.AddTo(delta)
 	sn.availableResource.SubFrom(delta)
 	sn.availableResource.Prune()
+	log.Log(log.SchedNode).Info("node allocation replaced",
+		zap.String("appID", replace.GetApplicationID()),
+		zap.String("allocationKey", replace.GetAllocationKey()),
+		zap.Stringer("allocatedResource", replace.GetAllocatedResource()),
+		zap.String("placeholderKey", allocationKey),
+		zap.String("targetNode", sn.NodeID))
 	if !before.FitIn(sn.allocatedResource) {
 		log.Log(log.SchedNode).Warn("unexpected increase in node usage after placeholder replacement",
 			zap.String("placeholder allocationKey", allocationKey),

pkg/scheduler/partition.go

@@ -869,44 +869,55 @@
 	// find the node make sure it still exists
 	// if the node was passed in use that ID instead of the one from the allocation
 	// the node ID is set when a reservation is allocated on a non-reserved node
-	var nodeID string
 	alloc := result.Request
-	if result.ReservedNodeID == "" {
-		nodeID = result.NodeID
-	} else {
-		nodeID = result.ReservedNodeID
-		log.Log(log.SchedPartition).Debug("Reservation allocated on different node",
-			zap.String("current node", result.NodeID),
-			zap.String("reserved node", nodeID),
-			zap.String("appID", appID))
-	}
-	node := pc.GetNode(nodeID)
-	if node == nil {
-		log.Log(log.SchedPartition).Info("Node was removed while allocating",
-			zap.String("nodeID", nodeID),
+	targetNodeID := result.NodeID
+	targetNode := pc.GetNode(targetNodeID)
+	if targetNode == nil {
+		log.Log(log.SchedPartition).Info("Target node was removed while allocating",
+			zap.String("nodeID", targetNodeID),
 			zap.String("appID", appID))
 
 		// attempt to deallocate
 		if alloc.IsAllocated() {
 			allocKey := alloc.GetAllocationKey()
 			if _, err := app.DeallocateAsk(allocKey); err != nil {
 				log.Log(log.SchedPartition).Warn("Failed to unwind allocation",
-					zap.String("nodeID", nodeID),
+					zap.String("nodeID", targetNodeID),
 					zap.String("appID", appID),
 					zap.String("allocationKey", allocKey),
 					zap.Error(err))
 			}
 		}
 		return nil
 	}
+
 	// reservation
 	if result.ResultType == objects.Reserved {
-		pc.reserve(app, node, result.Request)
+		pc.reserve(app, targetNode, result.Request)
 		return nil
 	}
+
 	// unreserve
 	if result.ResultType == objects.Unreserved || result.ResultType == objects.AllocatedReserved {
-		pc.unReserve(app, node, result.Request)
+		var reservedNodeID string
+		if result.ReservedNodeID == "" {
+			reservedNodeID = result.NodeID
+		} else {
+			reservedNodeID = result.ReservedNodeID
+			log.Log(log.SchedPartition).Debug("Reservation allocated on different node",
+				zap.String("current node", result.NodeID),
+				zap.String("reserved node", reservedNodeID),
+				zap.String("appID", appID))
+		}
+
+		reservedNode := pc.GetNode(reservedNodeID)
+		if reservedNode != nil {
+			pc.unReserve(app, reservedNode, result.Request)
+		} else {
+			log.Log(log.SchedPartition).Info("Reserved node was removed while allocating",
+				zap.String("nodeID", reservedNodeID),
+				zap.String("appID", appID))
+		}
 		if result.ResultType == objects.Unreserved {
 			return nil
 		}
@@ -915,8 +926,8 @@
 	}
 
 	alloc.SetBindTime(time.Now())
-	alloc.SetNodeID(nodeID)
-	alloc.SetInstanceType(node.GetInstanceType())
+	alloc.SetNodeID(targetNodeID)
+	alloc.SetInstanceType(targetNode.GetInstanceType())
 
 	// track the number of allocations
 	pc.updateAllocationCount(1)
@@ -929,7 +940,7 @@
 		zap.String("allocationKey", result.Request.GetAllocationKey()),
 		zap.Stringer("allocatedResource", result.Request.GetAllocatedResource()),
 		zap.Bool("placeholder", result.Request.IsPlaceholder()),
-		zap.String("targetNode", alloc.GetNodeID()))
+		zap.String("targetNode", targetNodeID))
 	// pass the allocation result back to the RM via the cluster context
 	return result
 }

pkg/scheduler/partition_test.go

@@ -2369,6 +2369,9 @@ func TestAllocReserveNewNode(t *testing.T) {
 	assert.Equal(t, 0, len(node1.GetReservationKeys()), "old node should have no more reservations")
 	assert.Equal(t, 0, len(app.GetReservations()), "ask should have been reserved")
 	assertLimits(t, getTestUserGroup(), resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 16000}))
+	alloc2 := node2.GetAllocation("alloc-2")
+	assert.Assert(t, alloc2 != nil, "alloc was nil")
+	assert.Equal(t, nodeID2, alloc2.GetNodeID(), "wrong node id")
 }
 
 func TestTryAllocateReserve(t *testing.T) {