[YUNIKORN-3509] Fix placeholder accounting on allocation rollback

pkg/scheduler/objects/application.go

@@ -472,7 +472,8 @@ func (sa *Application) clearPlaceholderTimer() {
 func (sa *Application) timeoutPlaceholderProcessing() {
 	sa.Lock()
 	defer sa.Unlock()
-	if (sa.IsRunning() || sa.IsCompleting()) && !resources.IsZero(sa.allocatedPlaceholder) {
+	// a rollback can leave a started application without allocated placeholders: Case 2 must not release its real allocations
+	if sa.IsRunning() || sa.IsCompleting() {
 		// Case 1: if all app's placeholders are allocated, only part of them gets replaced, just delete the remaining placeholders
 		var toRelease []*Allocation
 		replacing := 0
@@ -962,7 +963,8 @@ func (sa *Application) DeallocateAsk(allocKey string) (*resources.Resource, erro
 // This is used when the shim fails to bind a pod after the core has allocated it,
 // allowing the task to be re-scheduled on a different node.
 // The application must be in Accepted or Running state for rollback to succeed.
-func (sa *Application) RollbackAllocation(allocKey string) (*resources.Resource, error) {
+// Returns the rolled back allocation.
+func (sa *Application) RollbackAllocation(allocKey string) (*Allocation, error) {
 	sa.Lock()
 	defer sa.Unlock()
 
@@ -980,7 +982,7 @@ func (sa *Application) RollbackAllocation(allocKey string) (*resources.Resource,
 	}
 
 	// Restore pending FIRST. This ensures hasZeroAllocations() sees pending > 0 when
-	// we later decrement allocatedResource, preventing a spurious CompleteApplication event.
+	// we later decrement allocatedResource for a non-placeholder, preventing a spurious CompleteApplication event.
 	if _, err := sa.deallocateAsk(ask); err != nil {
 		return nil, fmt.Errorf("failed to deallocate ask %s during rollback on app %s: %w", allocKey, sa.ApplicationID, err)
 	}
@@ -989,8 +991,18 @@ func (sa *Application) RollbackAllocation(allocKey string) (*resources.Resource,
 	ask.SetNodeID("")
 
 	res := ask.GetAllocatedResource()
-	sa.allocatedResource = resources.Sub(sa.allocatedResource, res)
-	sa.allocatedResource.Prune()
+	if ask.IsPlaceholder() {
+		sa.allocatedPlaceholder = resources.Sub(sa.allocatedPlaceholder, res)
+		sa.allocatedPlaceholder.Prune()
+		// keep the placeholder timer: the gang deadline counts from the first placeholder allocation and also
+		// covers this placeholder when it is allocated again
+		if resources.IsZero(sa.allocatedPlaceholder) {
+			sa.hasPlaceholderAlloc = false
+		}
+	} else {
+		sa.allocatedResource = resources.Sub(sa.allocatedResource, res)
+		sa.allocatedResource.Prune()
+	}
 	sa.decUserResourceUsage(res, false)
 	delete(sa.allocations, allocKey)
 
@@ -1000,7 +1012,7 @@ func (sa *Application) RollbackAllocation(allocKey string) (*resources.Resource,
 		zap.String("appID", sa.ApplicationID),
 		zap.String("allocationKey", allocKey))
 
-	return res, nil
+	return ask, nil
 }
 
 func (sa *Application) allocateAsk(ask *Allocation) (*resources.Resource, error) {

pkg/scheduler/objects/application_test.go

@@ -4701,7 +4701,7 @@ func TestRollbackAllocationFromAccepted(t *testing.T) {
 	// Execute rollback.
 	returned, err := app.RollbackAllocation(aKey)
 	assert.NilError(t, err, "rollback should succeed from Accepted state")
-	assert.Assert(t, resources.Equals(returned, res), "returned resource should match allocation resource, got %v", returned)
+	assert.Equal(t, returned, ask, "rollback should return the rolled back allocation")
 
 	// Allocation is removed from the map.
 	assert.Assert(t, app.allocations[aKey] == nil, "allocation should have been removed from map")
@@ -4753,7 +4753,7 @@ func TestRollbackAllocationFromRunning(t *testing.T) {
 	// Execute rollback.
 	returned, err := app.RollbackAllocation(aKey)
 	assert.NilError(t, err, "rollback should succeed from Running state")
-	assert.Assert(t, resources.Equals(returned, res), "returned resource should match allocation resource, got %v", returned)
+	assert.Equal(t, returned, ask, "rollback should return the rolled back allocation")
 
 	// Allocation is removed.
 	assert.Assert(t, app.allocations[aKey] == nil, "allocation should have been removed from map")
@@ -4812,7 +4812,7 @@ func TestRollbackAllocationPartial(t *testing.T) {
 	// Roll back only the first allocation.
 	returned, err := app.RollbackAllocation(aKey)
 	assert.NilError(t, err, "partial rollback should succeed")
-	assert.Assert(t, resources.Equals(returned, res), "returned resource should equal res, got %v", returned)
+	assert.Equal(t, returned, ask1, "rollback should return the rolled back allocation")
 
 	// Second allocation is untouched.
 	assert.Assert(t, app.allocations[aKey2] != nil, "second allocation should still exist")
@@ -4828,6 +4828,114 @@ func TestRollbackAllocationPartial(t *testing.T) {
 	assert.Assert(t, app.IsRunning(), "app should remain Running after partial rollback")
 }
 
+// TestRollbackAllocationPlaceholder verifies that a placeholder rollback reverts the placeholder tracking and
+// leaves the real allocation tracking untouched. The placeholder timer keeps running, also after the last
+// placeholder is rolled back.
+func TestRollbackAllocationPlaceholder(t *testing.T) {
+	setupUGM()
+	queue, err := createRootQueue(nil)
+	assert.NilError(t, err, "queue create failed")
+	res := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5})
+	siApp := &si.AddApplicationRequest{
+		ApplicationID:  appID1,
+		QueueName:      "root.unknown",
+		PartitionName:  "default",
+		PlaceholderAsk: resources.Multiply(res, 2).ToProto(),
+	}
+	app := NewApplication(siApp, getTestUserGroup(), nil, "")
+	app.queue = queue
+
+	for _, key := range []string{aKey, aKey2} {
+		ph := newAllocationAskTG(key, appID1, tg1, res)
+		err = app.AddAllocationAsk(ph)
+		assert.NilError(t, err, "placeholder ask %s should have been added", key)
+		_, err = app.AllocateAsk(key)
+		assert.NilError(t, err, "AllocateAsk %s should succeed", key)
+		ph.SetNodeID(nodeID1)
+		err = app.AddAllocation(ph)
+		assert.NilError(t, err, "placeholder allocation %s should have been added", key)
+	}
+	assert.Assert(t, app.IsRunning(), "app should be in Running state once all placeholders are allocated")
+	assert.Assert(t, app.getPlaceholderTimer() != nil, "placeholder timer should be started by the placeholder allocation")
+	assertUserGroupResource(t, getTestUserGroup(), resources.Multiply(res, 2))
+
+	returned, err := app.RollbackAllocation(aKey)
+	assert.NilError(t, err, "rollback should succeed")
+	assert.Assert(t, returned.IsPlaceholder(), "returned allocation should be the placeholder")
+	assert.Assert(t, resources.Equals(app.GetPlaceholderResource(), res), "placeholder resource should only include the remaining placeholder, got %v", app.GetPlaceholderResource())
+	assert.Assert(t, app.getPlaceholderTimer() != nil, "placeholder timer should keep running while a placeholder is allocated")
+	assert.Assert(t, app.HasPlaceholderAllocation(), "app should still report a placeholder allocation")
+
+	_, err = app.RollbackAllocation(aKey2)
+	assert.NilError(t, err, "rollback of the last placeholder should succeed")
+	assert.Assert(t, resources.IsZero(app.GetPlaceholderResource()), "placeholder resource should be zero after rollback, got %v", app.GetPlaceholderResource())
+	assert.Assert(t, app.getPlaceholderTimer() != nil, "placeholder timer should keep running after the last placeholder is rolled back")
+	assert.Assert(t, !app.HasPlaceholderAllocation(), "app should not report a placeholder allocation after rollback")
+	assert.Assert(t, resources.IsZero(app.GetAllocatedResource()), "allocated resource should be zero after rollback, got %v", app.GetAllocatedResource())
+	assert.Assert(t, resources.Equals(app.GetPendingResource(), resources.Multiply(res, 2)), "pending resource should be restored after rollback, got %v", app.GetPendingResource())
+	assertUserGroupResource(t, getTestUserGroup(), nil)
+	assert.Assert(t, app.IsRunning(), "app should remain in Running state after rollback")
+}
+
+// TestTimeoutPlaceholderAfterRollback verifies the placeholder timeout once a rollback left no placeholder allocated:
+// an Accepted application times out as an incomplete gang, a Running one keeps its real allocations.
+func TestTimeoutPlaceholderAfterRollback(t *testing.T) {
+	res := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5})
+	newGangApp := func(t *testing.T) *Application {
+		setupUGM()
+		queue, err := createRootQueue(nil)
+		assert.NilError(t, err, "queue create failed")
+		siApp := &si.AddApplicationRequest{
+			ApplicationID:  appID1,
+			QueueName:      "root.unknown",
+			PartitionName:  "default",
+			PlaceholderAsk: resources.Multiply(res, 2).ToProto(),
+		}
+		app := NewApplication(siApp, getTestUserGroup(), nil, "")
+		app.queue = queue
+		return app
+	}
+	allocate := func(t *testing.T, app *Application, alloc *Allocation) {
+		err := app.AddAllocationAsk(alloc)
+		assert.NilError(t, err, "ask %s should have been added", alloc.GetAllocationKey())
+		_, err = app.AllocateAsk(alloc.GetAllocationKey())
+		assert.NilError(t, err, "AllocateAsk %s should succeed", alloc.GetAllocationKey())
+		alloc.SetNodeID(nodeID1)
+		err = app.AddAllocation(alloc)
+		assert.NilError(t, err, "allocation %s should have been added", alloc.GetAllocationKey())
+	}
+
+	t.Run("Accepted", func(t *testing.T) {
+		app := newGangApp(t)
+		allocate(t, app, newAllocationAskTG(aKey, appID1, tg1, res))
+		err := app.AddAllocationAsk(newAllocationAskTG(aKey2, appID1, tg1, res))
+		assert.NilError(t, err, "placeholder ask should have been added")
+		_, err = app.RollbackAllocation(aKey)
+		assert.NilError(t, err, "rollback should succeed")
+		assert.Assert(t, app.getPlaceholderTimer() != nil, "placeholder timer should keep the gang deadline")
+
+		app.timeoutPlaceholderProcessing()
+		assert.Assert(t, app.IsResuming(), "soft gang should resume on timeout, got state %s", app.CurrentState())
+		assert.Equal(t, len(app.getAllRequestsInternal()), 0, "placeholder asks should be removed on timeout")
+	})
+
+	t.Run("Running", func(t *testing.T) {
+		app := newGangApp(t)
+		allocate(t, app, newAllocationAskTG(aKey, appID1, tg1, res))
+		allocate(t, app, newAllocationAskTG(aKey2, appID1, tg1, res))
+		app.RemoveAllocation(aKey, si.TerminationType_PLACEHOLDER_REPLACED)
+		real := newAllocationAsk(aKey3, appID1, res)
+		allocate(t, app, real)
+		_, err := app.RollbackAllocation(aKey2)
+		assert.NilError(t, err, "rollback should succeed")
+		assert.Assert(t, resources.IsZero(app.GetPlaceholderResource()), "no placeholder should be allocated, got %v", app.GetPlaceholderResource())
+
+		app.timeoutPlaceholderProcessing()
+		assert.Assert(t, app.IsRunning(), "app should remain in Running state, got state %s", app.CurrentState())
+		assert.Assert(t, !real.IsReleased(), "real allocation should not be released on placeholder timeout")
+	})
+}
+
 func TestApplicationBackoff(t *testing.T) {
 	rootQ, err := createRootQueue(map[string]string{"first": "10"})
 	assert.NilError(t, err, "unexpected error when creating root queue")

pkg/scheduler/partition.go

@@ -1795,7 +1795,7 @@ func (pc *PartitionContext) rollbackAllocation(appID, allocationKey string, app
 	queue := app.GetQueue()
 	// Retrieve node ID before rolling back (RollbackAllocation clears it on the ask).
 	nodeID := app.GetAllocationNodeID(allocationKey)
-	res, err := app.RollbackAllocation(allocationKey)
+	alloc, err := app.RollbackAllocation(allocationKey)
 	if err != nil {
 		log.Log(log.SchedPartition).Warn("failed to rollback allocation",
 			zap.String("appID", appID),
@@ -1811,13 +1811,16 @@ func (pc *PartitionContext) rollbackAllocation(appID, allocationKey string, app
 			zap.String("allocationKey", allocationKey),
 			zap.String("nodeID", nodeID))
 	}
-	if err := queue.DecAllocatedResource(res); err != nil {
+	if err := queue.DecAllocatedResource(alloc.GetAllocatedResource()); err != nil {
 		log.Log(log.SchedPartition).Warn("failed to release resources from queue during rollback",
 			zap.String("appID", appID),
 			zap.String("allocationKey", allocationKey),
 			zap.Error(err))
 	}
 	pc.updateAllocationCount(-1)
+	if alloc.IsPlaceholder() {
+		pc.decPhAllocationCount(1)
+	}
 	log.Log(log.SchedPartition).Info("allocation rolled back to pending ask",
 		zap.String("appID", appID),
 		zap.String("allocationKey", allocationKey),

pkg/scheduler/partition_test.go

@@ -5590,6 +5590,7 @@ func TestRemoveAllocationSchedulingFailedOnRM(t *testing.T) {
 	assert.Assert(t, released == nil, "SCHEDULING_FAILED_ON_RM should not echo released allocations to the shim")
 	assert.Assert(t, confirmed == nil, "SCHEDULING_FAILED_ON_RM should not return a confirmed allocation")
 	assert.Equal(t, 0, partition.GetTotalAllocationCount(), "allocation count should be 0 after rollback")
+	assert.Equal(t, 0, partition.getPhAllocationCount(), "rollback of a normal allocation should not change the placeholder count")
 	assert.Assert(t, node.GetAllocation(allocKey) == nil, "node should not have the allocation after rollback")
 	assert.Assert(t, resources.IsZero(partition.GetQueue(defQueue).GetAllocatedResource()), "queue allocated resource should be zero after rollback")
 	assert.Assert(t, resources.StrictlyGreaterThanZero(app.GetPendingResource()), "ask should be pending again after rollback")
@@ -5682,6 +5683,55 @@ func TestRemoveAllocationSchedulingFailedOnRMNodeNotFound(t *testing.T) {
 	assert.Assert(t, resources.StrictlyGreaterThanZero(app.GetPendingResource()), "ask should be pending again after rollback")
 }
 
+// TestRemoveAllocationSchedulingFailedOnRMPlaceholder verifies that a placeholder rolled back after a bind
+// failure is no longer tracked as a placeholder by the app and the partition, and is counted once when it is
+// allocated again.
+func TestRemoveAllocationSchedulingFailedOnRMPlaceholder(t *testing.T) {
+	setupUGM()
+	partition, err := newBasePartition()
+	assert.NilError(t, err, "partition create failed")
+	defer partition.userGroupCache.Stop()
+
+	nodeRes := resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 10})
+	setupNode(t, nodeID1, partition, nodeRes)
+
+	phRes := resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 1})
+	app := newApplicationTG(appID1, "default", defQueue, resources.Multiply(phRes, 2))
+	err = partition.AddApplication(app)
+	assert.NilError(t, err, "add application failed")
+	err = app.AddAllocationAsk(newAllocationAskTG(phID, appID1, taskGroup, phRes, true))
+	assert.NilError(t, err, "failed to add placeholder ask ph-1 to app")
+	err = app.AddAllocationAsk(newAllocationAskTG(phID2, appID1, taskGroup, phRes, true))
+	assert.NilError(t, err, "failed to add placeholder ask ph-2 to app")
+
+	result := partition.tryAllocate()
+	assert.Assert(t, result != nil && result.Request != nil, "placeholder should have been allocated")
+	assert.Equal(t, 1, partition.getPhAllocationCount(), "placeholder allocation should be registered")
+
+	released, confirmed := partition.removeAllocation(&si.AllocationRelease{
+		PartitionName:   "test",
+		ApplicationID:   appID1,
+		AllocationKey:   result.Request.GetAllocationKey(),
+		TerminationType: si.TerminationType_SCHEDULING_FAILED_ON_RM,
+	})
+	assert.Assert(t, released == nil, "SCHEDULING_FAILED_ON_RM should not echo released allocations to the shim")
+	assert.Assert(t, confirmed == nil, "SCHEDULING_FAILED_ON_RM should not return a confirmed allocation")
+	assert.Equal(t, 0, partition.GetTotalAllocationCount(), "allocation count should be 0 after rollback")
+	assert.Equal(t, 0, partition.getPhAllocationCount(), "placeholder count should be 0 after rollback")
+	assert.Assert(t, resources.IsZero(app.GetPlaceholderResource()), "app placeholder resource should be zero after rollback, got %v", app.GetPlaceholderResource())
+	assert.Assert(t, resources.IsZero(app.GetAllocatedResource()), "app allocated resource should be zero after rollback, got %v", app.GetAllocatedResource())
+
+	result = partition.tryAllocate()
+	assert.Assert(t, result != nil && result.Request != nil, "rolled back placeholder should have been allocated again")
+	assert.Assert(t, app.IsAccepted(), "app should wait for the second placeholder, got state %s", app.CurrentState())
+	result = partition.tryAllocate()
+	assert.Assert(t, result != nil && result.Request != nil, "second placeholder should have been allocated")
+	assert.Assert(t, app.IsRunning(), "app should be running once all placeholders are allocated, got state %s", app.CurrentState())
+	assert.Equal(t, 2, partition.getPhAllocationCount(), "both placeholder allocations should be registered once")
+	assert.Assert(t, resources.Equals(app.GetPlaceholderResource(), resources.Multiply(phRes, 2)), "app placeholder resource should count both placeholders once, got %v", app.GetPlaceholderResource())
+	assert.Assert(t, resources.IsZero(app.GetAllocatedResource()), "app allocated resource should be zero, got %v", app.GetAllocatedResource())
+}
+
 func TestRemoveAppWithReservations(t *testing.T) {
 	setupUGM()
 	partition := createQueuesNodes(t)