From 5dd0350cf114cc2494deaf32625f60546a8d0492 Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Fri, 9 Oct 2026 18:13:01 +0800 Subject: [PATCH 1/5] [YUNIKORN-3509] Fix placeholder accounting on allocation rollback A placeholder rolled back after a shim bind failure was subtracted from the application's allocated resource instead of its placeholder resource, and the partition kept counting it as a placeholder. The application could then report negative usage, move to Running before its gang was complete, and get stuck in Completing. --- pkg/scheduler/objects/application.go | 21 +++++++--- pkg/scheduler/objects/application_test.go | 40 ++++++++++++++++-- pkg/scheduler/partition.go | 7 +++- pkg/scheduler/partition_test.go | 49 +++++++++++++++++++++++ 4 files changed, 107 insertions(+), 10 deletions(-) diff --git a/pkg/scheduler/objects/application.go b/pkg/scheduler/objects/application.go index 4b0c1a948..6363cbbdd 100644 --- a/pkg/scheduler/objects/application.go +++ b/pkg/scheduler/objects/application.go @@ -962,7 +962,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 +981,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 +990,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() + // a timer left running without allocated placeholders would release the real allocations on timeout + if resources.IsZero(sa.allocatedPlaceholder) { + sa.clearPlaceholderTimer() + sa.hasPlaceholderAlloc = false + } + } else { + sa.allocatedResource = resources.Sub(sa.allocatedResource, res) + sa.allocatedResource.Prune() + } sa.decUserResourceUsage(res, false) delete(sa.allocations, allocKey) @@ -1000,7 +1011,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) { diff --git a/pkg/scheduler/objects/application_test.go b/pkg/scheduler/objects/application_test.go index 42dc28011..1eb18cb4e 100644 --- a/pkg/scheduler/objects/application_test.go +++ b/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,40 @@ 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, including the placeholder timer, and leaves the real allocation tracking untouched. +func TestRollbackAllocationPlaceholder(t *testing.T) { + setupUGM() + queue, err := createRootQueue(nil) + assert.NilError(t, err, "queue create failed") + app := newApplication(appID1, "default", "root.unknown") + app.queue = queue + res := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}) + + ph := newAllocationAskTG(aKey, appID1, tg1, res) + err = app.AddAllocationAsk(ph) + assert.NilError(t, err, "placeholder ask should have been added") + _, err = app.AllocateAsk(aKey) + assert.NilError(t, err, "AllocateAsk should succeed") + ph.SetNodeID(nodeID1) + err = app.AddAllocation(ph) + assert.NilError(t, err, "placeholder allocation should have been added") + assert.Assert(t, resources.Equals(app.GetPlaceholderResource(), res), "placeholder resource should equal res before rollback, got %v", app.GetPlaceholderResource()) + assert.Assert(t, app.getPlaceholderTimer() != nil, "placeholder timer should be started by the placeholder allocation") + assertUserGroupResource(t, getTestUserGroup(), res) + + 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.IsZero(app.GetPlaceholderResource()), "placeholder resource should be zero after rollback, got %v", app.GetPlaceholderResource()) + assert.Assert(t, app.getPlaceholderTimer() == nil, "placeholder timer should be cleared when no placeholder is allocated") + 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(), res), "pending resource should be restored after rollback, got %v", app.GetPendingResource()) + assertUserGroupResource(t, getTestUserGroup(), nil) + assert.Assert(t, app.IsAccepted(), "app should remain in Accepted state after rollback") +} + func TestApplicationBackoff(t *testing.T) { rootQ, err := createRootQueue(map[string]string{"first": "10"}) assert.NilError(t, err, "unexpected error when creating root queue") diff --git a/pkg/scheduler/partition.go b/pkg/scheduler/partition.go index c9ada617f..b3bf22b6b 100644 --- a/pkg/scheduler/partition.go +++ b/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), diff --git a/pkg/scheduler/partition_test.go b/pkg/scheduler/partition_test.go index d4f621b22..db4042643 100644 --- a/pkg/scheduler/partition_test.go +++ b/pkg/scheduler/partition_test.go @@ -5682,6 +5682,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) From 666984709f72fe1930dbbfc59abdadcb708faa03 Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Fri, 9 Oct 2026 19:21:49 +0800 Subject: [PATCH 2/5] [YUNIKORN-3509] Cover partial placeholder rollback and the normal rollback count The tests only rolled back the last allocated placeholder, so clearing the placeholder timer while another placeholder is still allocated went unnoticed, as did decrementing the partition placeholder count for a normal allocation. --- pkg/scheduler/objects/application_test.go | 47 +++++++++++++++-------- pkg/scheduler/partition_test.go | 1 + 2 files changed, 32 insertions(+), 16 deletions(-) diff --git a/pkg/scheduler/objects/application_test.go b/pkg/scheduler/objects/application_test.go index 1eb18cb4e..4c39a1554 100644 --- a/pkg/scheduler/objects/application_test.go +++ b/pkg/scheduler/objects/application_test.go @@ -4828,38 +4828,53 @@ 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, including the placeholder timer, and leaves the real allocation tracking untouched. +// TestRollbackAllocationPlaceholder verifies that a placeholder rollback reverts the placeholder tracking and +// leaves the real allocation tracking untouched. The placeholder timer keeps running while a placeholder is +// still allocated and is cleared with the last one. func TestRollbackAllocationPlaceholder(t *testing.T) { setupUGM() queue, err := createRootQueue(nil) assert.NilError(t, err, "queue create failed") - app := newApplication(appID1, "default", "root.unknown") - app.queue = queue 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 - ph := newAllocationAskTG(aKey, appID1, tg1, res) - err = app.AddAllocationAsk(ph) - assert.NilError(t, err, "placeholder ask should have been added") - _, err = app.AllocateAsk(aKey) - assert.NilError(t, err, "AllocateAsk should succeed") - ph.SetNodeID(nodeID1) - err = app.AddAllocation(ph) - assert.NilError(t, err, "placeholder allocation should have been added") - assert.Assert(t, resources.Equals(app.GetPlaceholderResource(), res), "placeholder resource should equal res before rollback, got %v", app.GetPlaceholderResource()) + 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(), res) + 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 be cleared when no placeholder is allocated") 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(), res), "pending resource should be restored after rollback, got %v", app.GetPendingResource()) + 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.IsAccepted(), "app should remain in Accepted state after rollback") + assert.Assert(t, app.IsRunning(), "app should remain in Running state after rollback") } func TestApplicationBackoff(t *testing.T) { diff --git a/pkg/scheduler/partition_test.go b/pkg/scheduler/partition_test.go index db4042643..6dbb80909 100644 --- a/pkg/scheduler/partition_test.go +++ b/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") From 79fb1ed43c15edfde5089ecb05287066641cc377 Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Sun, 11 Oct 2026 19:57:14 +0800 Subject: [PATCH 3/5] [YUNIKORN-3509] Keep the placeholder timer running after a rollback Clearing the timer once a rollback left no placeholder allocated reset the gang deadline in Accepted, and left a placeholder allocated again in Running without any timeout. The timer was only cleared because a started application without allocated placeholders took the timeout path that releases every allocation; choosing that path by application state alone removes the need to clear it. --- pkg/scheduler/objects/application.go | 7 +-- pkg/scheduler/objects/application_test.go | 65 +++++++++++++++++++++-- 2 files changed, 66 insertions(+), 6 deletions(-) diff --git a/pkg/scheduler/objects/application.go b/pkg/scheduler/objects/application.go index 6363cbbdd..9612bee0e 100644 --- a/pkg/scheduler/objects/application.go +++ b/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 @@ -993,9 +994,9 @@ func (sa *Application) RollbackAllocation(allocKey string) (*Allocation, error) if ask.IsPlaceholder() { sa.allocatedPlaceholder = resources.Sub(sa.allocatedPlaceholder, res) sa.allocatedPlaceholder.Prune() - // a timer left running without allocated placeholders would release the real allocations on timeout + // 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.clearPlaceholderTimer() sa.hasPlaceholderAlloc = false } } else { diff --git a/pkg/scheduler/objects/application_test.go b/pkg/scheduler/objects/application_test.go index 4c39a1554..4c4fca07a 100644 --- a/pkg/scheduler/objects/application_test.go +++ b/pkg/scheduler/objects/application_test.go @@ -4829,8 +4829,8 @@ func TestRollbackAllocationPartial(t *testing.T) { } // TestRollbackAllocationPlaceholder verifies that a placeholder rollback reverts the placeholder tracking and -// leaves the real allocation tracking untouched. The placeholder timer keeps running while a placeholder is -// still allocated and is cleared with the last one. +// 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) @@ -4869,7 +4869,7 @@ func TestRollbackAllocationPlaceholder(t *testing.T) { _, 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 be cleared when no placeholder is allocated") + 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()) @@ -4877,6 +4877,65 @@ func TestRollbackAllocationPlaceholder(t *testing.T) { 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") From c80a775b72fa094d83868f01cf9d946182d672e8 Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Thu, 15 Oct 2026 18:15:03 +0800 Subject: [PATCH 4/5] [YUNIKORN-3509] Roll back the gang state with a placeholder rollback A placeholder is only rolled back while its gang is forming, so the rollback has to undo what that allocation progressed. Left in Running after its last placeholder was rolled back, an application whose gang timed out only released its placeholders and was never resumed or failed, which left the shim waiting for the gang. With no placeholder left allocated nothing was replaced yet, so the timer starts over, and the peak usage no longer counts a placeholder that never ran. --- pkg/scheduler/objects/application.go | 18 ++- pkg/scheduler/objects/application_state.go | 12 +- .../objects/application_state_test.go | 30 +++++ pkg/scheduler/objects/application_test.go | 121 +++++++----------- pkg/scheduler/objects/utilities_test.go | 23 ++++ 5 files changed, 125 insertions(+), 79 deletions(-) diff --git a/pkg/scheduler/objects/application.go b/pkg/scheduler/objects/application.go index 9612bee0e..1033b9979 100644 --- a/pkg/scheduler/objects/application.go +++ b/pkg/scheduler/objects/application.go @@ -472,8 +472,7 @@ func (sa *Application) clearPlaceholderTimer() { func (sa *Application) timeoutPlaceholderProcessing() { sa.Lock() defer sa.Unlock() - // a rollback can leave a started application without allocated placeholders: Case 2 must not release its real allocations - if sa.IsRunning() || sa.IsCompleting() { + if (sa.IsRunning() || sa.IsCompleting()) && !resources.IsZero(sa.allocatedPlaceholder) { // 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 @@ -994,11 +993,22 @@ func (sa *Application) RollbackAllocation(allocKey string) (*Allocation, error) 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 + // placeholders are only rolled back while the gang forms and placeholder usage only grows, so the peak + // includes this placeholder + sa.maxAllocatedResource = resources.ComponentWiseMax(resources.Sub(sa.maxAllocatedResource, res), sa.allocatedPlaceholder) + sa.maxAllocatedResource.Prune() if resources.IsZero(sa.allocatedPlaceholder) { + sa.clearPlaceholderTimer() sa.hasPlaceholderAlloc = false } + // the application moved to Running when its last placeholder was allocated: the gang is incomplete again + if sa.IsRunning() && resources.IsZero(sa.allocatedResource) { + if err := sa.HandleApplicationEvent(RollbackApplication); err != nil { + log.Log(log.SchedApplication).Warn("Unexpected app state change failure while rolling back a placeholder", + zap.String("currentState", sa.stateMachine.Current()), + zap.Error(err)) + } + } } else { sa.allocatedResource = resources.Sub(sa.allocatedResource, res) sa.allocatedResource.Prune() diff --git a/pkg/scheduler/objects/application_state.go b/pkg/scheduler/objects/application_state.go index 708a4f3eb..5a7260e2a 100644 --- a/pkg/scheduler/objects/application_state.go +++ b/pkg/scheduler/objects/application_state.go @@ -45,6 +45,7 @@ const ( FailApplication ExpireApplication ResumeApplication + RollbackApplication ) const ( @@ -52,7 +53,7 @@ const ( ) func (ae applicationEvent) String() string { - return [...]string{"runApplication", "rejectApplication", "completeApplication", "failApplication", "expireApplication", "resumeApplication"}[ae] + return [...]string{"runApplication", "rejectApplication", "completeApplication", "failApplication", "expireApplication", "resumeApplication", "rollbackApplication"}[ae] } // ---------------------------------- @@ -130,6 +131,10 @@ func eventDesc() fsm.Events { Name: ResumeApplication.String(), Src: []string{New.String(), Accepted.String()}, Dst: Resuming.String(), + }, { + Name: RollbackApplication.String(), + Src: []string{Running.String()}, + Dst: Accepted.String(), }, { Name: ExpireApplication.String(), Src: []string{Completed.String(), Failed.String(), Rejected.String()}, @@ -183,7 +188,10 @@ func callbacks() fsm.Callbacks { fmt.Sprintf("enter_%s", Accepted.String()): func(_ context.Context, event *fsm.Event) { app := event.Args[0].(*Application) //nolint:errcheck metrics.GetQueueMetrics(app.queuePath).IncQueueApplicationsAccepted() - metrics.GetSchedulerMetrics().IncTotalApplicationsAccepted() + // a rollback returns an application that was already counted as accepted + if event.Event != RollbackApplication.String() { + metrics.GetSchedulerMetrics().IncTotalApplicationsAccepted() + } }, fmt.Sprintf("leave_%s", Accepted.String()): func(_ context.Context, event *fsm.Event) { app := event.Args[0].(*Application) //nolint:errcheck diff --git a/pkg/scheduler/objects/application_state_test.go b/pkg/scheduler/objects/application_state_test.go index 3a1bd38c8..95f5c53a0 100644 --- a/pkg/scheduler/objects/application_state_test.go +++ b/pkg/scheduler/objects/application_state_test.go @@ -210,6 +210,36 @@ func TestFailedStateTransition(t *testing.T) { assert.NilError(t, err, "App should be in Failed state") } +func TestRollbackStateTransition(t *testing.T) { + queue := createQueue(t, "rollback") + metrics.GetSchedulerMetrics().Reset() + metrics.GetQueueMetrics("root.rollback").Reset() + app := newApplication("app-00001", "default", "root.rollback") + app.SetQueue(queue) + + err := app.HandleApplicationEvent(RunApplication) + assertState(t, app, err, Accepted.String()) + err = app.HandleApplicationEvent(RollbackApplication) + assert.Assert(t, err != nil, "error expected accepted to accepted (rollback test)") + assert.Equal(t, app.CurrentState(), Accepted.String()) + + err = app.HandleApplicationEvent(RunApplication) + assertState(t, app, err, Running.String()) + err = app.HandleApplicationEvent(RollbackApplication) + assertState(t, app, err, Accepted.String()) + assertTotalAppsAcceptedMetrics(t, 1) + assertTotalAppsRunningMetrics(t, 0) + assertQueueRunningApps(t, app, 0) + assertQueueApplicationsRunningMetrics(t, app, 0) + assertQueueApplicationsAcceptedMetrics(t, app, 1) + + err = app.HandleApplicationEvent(RunApplication) + assertState(t, app, err, Running.String()) + assertTotalAppsRunningMetrics(t, 1) + assertQueueRunningApps(t, app, 1) + assertQueueApplicationsAcceptedMetrics(t, app, 0) +} + func TestAppStateTransitionEvents(t *testing.T) { events.Init() eventSystem := events.GetEventSystem().(*events.EventSystemImpl) //nolint:errcheck diff --git a/pkg/scheduler/objects/application_test.go b/pkg/scheduler/objects/application_test.go index 4c4fca07a..6a88b138f 100644 --- a/pkg/scheduler/objects/application_test.go +++ b/pkg/scheduler/objects/application_test.go @@ -4828,32 +4828,18 @@ 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. +// TestRollbackAllocationPlaceholder verifies that a placeholder rollback reverts the placeholder tracking and peak, +// moves the application back to Accepted, and clears the placeholder timer once no placeholder is left allocated. 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 := newGangApplication(res, 2) 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) + allocatePlaceholder(t, app, key, res) } 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") @@ -4862,78 +4848,67 @@ func TestRollbackAllocationPlaceholder(t *testing.T) { 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, app.IsAccepted(), "app should move back to Accepted state, got state %s", app.CurrentState()) assert.Assert(t, resources.Equals(app.GetPlaceholderResource(), res), "placeholder resource should only include the remaining placeholder, got %v", app.GetPlaceholderResource()) + assert.Assert(t, resources.Equals(app.GetMaxAllocatedResource(), res), "peak should not include the rolled back placeholder, got %v", app.GetMaxAllocatedResource()) 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, app.IsAccepted(), "app should remain in Accepted state, got state %s", app.CurrentState()) 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, resources.IsZero(app.GetMaxAllocatedResource()), "peak should be zero after rollback, got %v", app.GetMaxAllocatedResource()) + assert.Assert(t, app.getPlaceholderTimer() == nil, "placeholder timer should be cleared when no placeholder is allocated") 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) { +// TestRollbackAllocationPlaceholderWithRealAllocation verifies that a placeholder rollback leaves the state of an +// application that runs because of a real allocation alone. +func TestRollbackAllocationPlaceholderWithRealAllocation(t *testing.T) { + setupUGM() + queue, err := createRootQueue(nil) + assert.NilError(t, err, "queue create failed") 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 := newGangApplication(res, 2) + app.queue = queue - 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") - }) + real := newAllocationAsk(aKey3, appID1, res) + err = app.AddAllocationAsk(real) + assert.NilError(t, err, "real ask should have been added") + _, err = app.AllocateAsk(aKey3) + assert.NilError(t, err, "AllocateAsk should succeed") + real.SetNodeID(nodeID1) + err = app.AddAllocation(real) + assert.NilError(t, err, "real allocation should have been added") + allocatePlaceholder(t, app, aKey, res) + assert.Assert(t, app.IsRunning(), "app should be in Running state with a real allocation") - 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()) + _, err = app.RollbackAllocation(aKey) + assert.NilError(t, err, "rollback should succeed") + assert.Assert(t, app.IsRunning(), "app should remain in Running state, got state %s", app.CurrentState()) +} - 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") - }) +// TestTimeoutPlaceholderAfterRollback verifies that a gang that is incomplete again after a rollback times out as +// an incomplete gang. +func TestTimeoutPlaceholderAfterRollback(t *testing.T) { + setupUGM() + queue, err := createRootQueue(nil) + assert.NilError(t, err, "queue create failed") + res := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}) + app := newGangApplication(res, 2) + app.queue = queue + allocatePlaceholder(t, app, aKey, res) + allocatePlaceholder(t, app, aKey2, res) + _, err = app.RollbackAllocation(aKey2) + assert.NilError(t, err, "rollback should succeed") + + 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") } func TestApplicationBackoff(t *testing.T) { diff --git a/pkg/scheduler/objects/utilities_test.go b/pkg/scheduler/objects/utilities_test.go index 17122b8f3..9c008e7da 100644 --- a/pkg/scheduler/objects/utilities_test.go +++ b/pkg/scheduler/objects/utilities_test.go @@ -175,6 +175,29 @@ func newApplicationWithPlaceholderTimeout(appID, partition, queueName string, ph return NewApplication(siApp, getTestUserGroup(), mockEventHandler, ""), mockEventHandler } +// newGangApplication creates an application that asks for count placeholders of res +func newGangApplication(res *resources.Resource, count int64) *Application { + siApp := &si.AddApplicationRequest{ + ApplicationID: appID1, + QueueName: "root.unknown", + PartitionName: "default", + PlaceholderAsk: resources.Multiply(res, count).ToProto(), + } + return NewApplication(siApp, getTestUserGroup(), nil, "") +} + +// allocatePlaceholder adds a placeholder ask for key and allocates it on nodeID1 +func allocatePlaceholder(t *testing.T, app *Application, key string, res *resources.Resource) { + 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) +} + // Create node with minimal info func newNode(nodeID string, totalMap map[string]resources.Quantity) *Node { total := resources.NewResourceFromMap(totalMap) From 34af2ec9ab22a8f7bf08d8557bdf4edac5d6822a Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Thu, 15 Oct 2026 19:05:42 +0800 Subject: [PATCH 5/5] [YUNIKORN-3509] Keep the running slot of a gang rolled back into Accepted An application back in Accepted that still held placeholders was no longer counted against the queue's max running applications. Another application could take the slot, after which the gang could not get its rolled back placeholder again and held the rest until the gang timeout. The peak rollback assumes a forming gang, which has no real allocations. Applied with real allocations it could drop the peak below the current usage, so it now shares the state rollback's precondition. --- pkg/scheduler/objects/application.go | 22 ++++---- pkg/scheduler/objects/application_state.go | 6 ++ pkg/scheduler/objects/application_test.go | 66 +++++++++++++++------- pkg/scheduler/objects/utilities_test.go | 31 +++++----- 4 files changed, 80 insertions(+), 45 deletions(-) diff --git a/pkg/scheduler/objects/application.go b/pkg/scheduler/objects/application.go index 1033b9979..1f86385b0 100644 --- a/pkg/scheduler/objects/application.go +++ b/pkg/scheduler/objects/application.go @@ -993,20 +993,22 @@ func (sa *Application) RollbackAllocation(allocKey string) (*Allocation, error) if ask.IsPlaceholder() { sa.allocatedPlaceholder = resources.Sub(sa.allocatedPlaceholder, res) sa.allocatedPlaceholder.Prune() - // placeholders are only rolled back while the gang forms and placeholder usage only grows, so the peak - // includes this placeholder - sa.maxAllocatedResource = resources.ComponentWiseMax(resources.Sub(sa.maxAllocatedResource, res), sa.allocatedPlaceholder) - sa.maxAllocatedResource.Prune() if resources.IsZero(sa.allocatedPlaceholder) { sa.clearPlaceholderTimer() sa.hasPlaceholderAlloc = false } - // the application moved to Running when its last placeholder was allocated: the gang is incomplete again - if sa.IsRunning() && resources.IsZero(sa.allocatedResource) { - if err := sa.HandleApplicationEvent(RollbackApplication); err != nil { - log.Log(log.SchedApplication).Warn("Unexpected app state change failure while rolling back a placeholder", - zap.String("currentState", sa.stateMachine.Current()), - zap.Error(err)) + // without real allocations the gang is still forming: placeholder usage only grew, so the peak includes this + // placeholder, and a Running application got there when its last placeholder was allocated + if resources.IsZero(sa.allocatedResource) { + sa.maxAllocatedResource = resources.ComponentWiseMax(resources.Sub(sa.maxAllocatedResource, res), sa.allocatedPlaceholder) + sa.maxAllocatedResource.Prune() + if sa.IsRunning() { + if err := sa.HandleApplicationEvent(RollbackApplication); err != nil { + log.Log(log.SchedApplication).Warn("Application state not changed to Accepted while rolling back a placeholder allocation", + zap.String("currentState", sa.CurrentState()), + zap.Stringer("event", RollbackApplication), + zap.Error(err)) + } } } } else { diff --git a/pkg/scheduler/objects/application_state.go b/pkg/scheduler/objects/application_state.go index 5a7260e2a..999022197 100644 --- a/pkg/scheduler/objects/application_state.go +++ b/pkg/scheduler/objects/application_state.go @@ -26,6 +26,7 @@ import ( "github.com/looplab/fsm" "go.uber.org/zap" + "github.com/apache/yunikorn-core/pkg/common/resources" "github.com/apache/yunikorn-core/pkg/log" "github.com/apache/yunikorn-core/pkg/metrics" "github.com/apache/yunikorn-scheduler-interface/lib/go/si" @@ -224,6 +225,11 @@ func callbacks() fsm.Callbacks { fmt.Sprintf("leave_%s", Running.String()): func(_ context.Context, event *fsm.Event) { if event.Dst != Running.String() { app := event.Args[0].(*Application) //nolint:errcheck + // an application back in Accepted that still holds placeholders is counted as running, like before it + // went Running: add it before the running count drops so that its slot is never free in between + if event.Dst == Accepted.String() && !resources.IsZero(app.allocatedPlaceholder) { + app.queue.setAllocatingAccepted(app.ApplicationID) + } app.queue.decRunningApps() metrics.GetQueueMetrics(app.queuePath).DecQueueApplicationsRunning() metrics.GetSchedulerMetrics().DecTotalApplicationsRunning() diff --git a/pkg/scheduler/objects/application_test.go b/pkg/scheduler/objects/application_test.go index 6a88b138f..0d8aeae6c 100644 --- a/pkg/scheduler/objects/application_test.go +++ b/pkg/scheduler/objects/application_test.go @@ -4835,12 +4835,11 @@ func TestRollbackAllocationPlaceholder(t *testing.T) { queue, err := createRootQueue(nil) assert.NilError(t, err, "queue create failed") res := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}) - app := newGangApplication(res, 2) + app := newApplicationTG(appID1, "default", "root.unknown", resources.Multiply(res, 2)) app.queue = queue - for _, key := range []string{aKey, aKey2} { - allocatePlaceholder(t, app, key, res) - } + allocateAndAdd(t, app, newAllocationAskTG(aKey, appID1, tg1, res)) + allocateAndAdd(t, app, newAllocationAskTG(aKey2, appID1, tg1, res)) 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)) @@ -4866,30 +4865,59 @@ func TestRollbackAllocationPlaceholder(t *testing.T) { assertUserGroupResource(t, getTestUserGroup(), nil) } -// TestRollbackAllocationPlaceholderWithRealAllocation verifies that a placeholder rollback leaves the state of an -// application that runs because of a real allocation alone. +// TestRollbackAllocationPlaceholderWithRealAllocation verifies that a placeholder rollback leaves the state and the +// peak of an application with a real allocation alone. func TestRollbackAllocationPlaceholderWithRealAllocation(t *testing.T) { setupUGM() queue, err := createRootQueue(nil) assert.NilError(t, err, "queue create failed") res := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}) - app := newGangApplication(res, 2) + app := newApplicationTG(appID1, "default", "root.unknown", resources.Multiply(res, 2)) app.queue = queue - real := newAllocationAsk(aKey3, appID1, res) - err = app.AddAllocationAsk(real) - assert.NilError(t, err, "real ask should have been added") - _, err = app.AllocateAsk(aKey3) - assert.NilError(t, err, "AllocateAsk should succeed") - real.SetNodeID(nodeID1) - err = app.AddAllocation(real) - assert.NilError(t, err, "real allocation should have been added") - allocatePlaceholder(t, app, aKey, res) + allocateAndAdd(t, app, newAllocationAsk(aKey3, appID1, res)) + allocateAndAdd(t, app, newAllocationAskTG(aKey, appID1, tg1, res)) assert.Assert(t, app.IsRunning(), "app should be in Running state with a real allocation") _, err = app.RollbackAllocation(aKey) assert.NilError(t, err, "rollback should succeed") assert.Assert(t, app.IsRunning(), "app should remain in Running state, got state %s", app.CurrentState()) + assert.Assert(t, resources.Equals(app.GetMaxAllocatedResource(), res), "peak should be unchanged, got %v", app.GetMaxAllocatedResource()) +} + +// TestRollbackAllocationPlaceholderRunningSlot verifies that an application moved back to Accepted keeps its running +// application slot in the queue only while it still holds placeholders. +func TestRollbackAllocationPlaceholderRunningSlot(t *testing.T) { + res := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}) + tests := []struct { + name string + placeholders int + slotKept bool + }{ + {"placeholder left", 2, true}, + {"no placeholder left", 1, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + setupUGM() + queue, err := createRootQueue(nil) + assert.NilError(t, err, "queue create failed") + queue.maxRunningApps = 1 + app := newApplicationTG(appID1, "default", "root.unknown", resources.Multiply(res, int64(tt.placeholders))) + app.queue = queue + keys := []string{aKey, aKey2}[:tt.placeholders] + for _, key := range keys { + allocateAndAdd(t, app, newAllocationAskTG(key, appID1, tg1, res)) + } + assert.Assert(t, app.IsRunning(), "app should be in Running state once all placeholders are allocated") + assert.Assert(t, !queue.canRunApp(appID2), "running app should take the only slot") + + _, err = app.RollbackAllocation(keys[len(keys)-1]) + assert.NilError(t, err, "rollback should succeed") + assert.Assert(t, app.IsAccepted(), "app should move back to Accepted state, got state %s", app.CurrentState()) + assert.Equal(t, queue.canRunApp(appID2), !tt.slotKept, "unexpected slot for another application") + }) + } } // TestTimeoutPlaceholderAfterRollback verifies that a gang that is incomplete again after a rollback times out as @@ -4899,10 +4927,10 @@ func TestTimeoutPlaceholderAfterRollback(t *testing.T) { queue, err := createRootQueue(nil) assert.NilError(t, err, "queue create failed") res := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5}) - app := newGangApplication(res, 2) + app := newApplicationTG(appID1, "default", "root.unknown", resources.Multiply(res, 2)) app.queue = queue - allocatePlaceholder(t, app, aKey, res) - allocatePlaceholder(t, app, aKey2, res) + allocateAndAdd(t, app, newAllocationAskTG(aKey, appID1, tg1, res)) + allocateAndAdd(t, app, newAllocationAskTG(aKey2, appID1, tg1, res)) _, err = app.RollbackAllocation(aKey2) assert.NilError(t, err, "rollback should succeed") diff --git a/pkg/scheduler/objects/utilities_test.go b/pkg/scheduler/objects/utilities_test.go index 9c008e7da..7ef27335f 100644 --- a/pkg/scheduler/objects/utilities_test.go +++ b/pkg/scheduler/objects/utilities_test.go @@ -175,27 +175,26 @@ func newApplicationWithPlaceholderTimeout(appID, partition, queueName string, ph return NewApplication(siApp, getTestUserGroup(), mockEventHandler, ""), mockEventHandler } -// newGangApplication creates an application that asks for count placeholders of res -func newGangApplication(res *resources.Resource, count int64) *Application { +// Create application with a placeholder ask for the task groups +func newApplicationTG(appID, partition, queueName string, task *resources.Resource) *Application { siApp := &si.AddApplicationRequest{ - ApplicationID: appID1, - QueueName: "root.unknown", - PartitionName: "default", - PlaceholderAsk: resources.Multiply(res, count).ToProto(), + ApplicationID: appID, + QueueName: queueName, + PartitionName: partition, + PlaceholderAsk: task.ToProto(), } return NewApplication(siApp, getTestUserGroup(), nil, "") } -// allocatePlaceholder adds a placeholder ask for key and allocates it on nodeID1 -func allocatePlaceholder(t *testing.T, app *Application, key string, res *resources.Resource) { - 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) +// allocateAndAdd adds the ask to the application, allocates it on nodeID1 and adds the allocation +func allocateAndAdd(t *testing.T, app *Application, ask *Allocation) { + err := app.AddAllocationAsk(ask) + assert.NilError(t, err, "ask %s should have been added", ask.GetAllocationKey()) + _, err = app.AllocateAsk(ask.GetAllocationKey()) + assert.NilError(t, err, "AllocateAsk %s should succeed", ask.GetAllocationKey()) + ask.SetNodeID(nodeID1) + err = app.AddAllocation(ask) + assert.NilError(t, err, "allocation %s should have been added", ask.GetAllocationKey()) } // Create node with minimal info