This is an automated email from the ASF dual-hosted git repository.
wilfred-s pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/yunikorn-core.git
The following commit(s) were added to refs/heads/master by this push:
new 900bd603 [YUNIKORN-3316] Handle AssumePod failures in core (#1104)
900bd603 is described below
commit 900bd603df82609152d30de2f4811daf109380ed
Author: Aditya Maheshwari <[email protected]>
AuthorDate: Tue Jul 28 22:39:50 2026 +1000
[YUNIKORN-3316] Handle AssumePod failures in core (#1104)
During UpdateAllocation, shim tries to call AssumePod which updates
in-memory cache of k8s for pod volume bindings (not the actual binding).
We have retries in place to retry any failure during this step. But if
all retries get exhausted we fail the task and ask core to release the
allocation.
Instead of failing the task, we should retry again the task after all
the cleanup.
Core changes to handle updates of tasks with termination type
SCHEDULING_FAILED_ON_RM
Remove allocation from node, application, update and clear resources on
queue, update partition allocation count.
Closes: #1104
Signed-off-by: Wilfred Spiegelenburg <[email protected]>
---
go.mod | 2 +-
go.sum | 4 +-
pkg/scheduler/objects/application.go | 52 +++++++
pkg/scheduler/objects/application_test.go | 245 ++++++++++++++++++++++++++++++
pkg/scheduler/partition.go | 109 +++++++++----
pkg/scheduler/partition_test.go | 136 +++++++++++++++++
6 files changed, 517 insertions(+), 31 deletions(-)
diff --git a/go.mod b/go.mod
index c83ea670..76556ee2 100644
--- a/go.mod
+++ b/go.mod
@@ -22,7 +22,7 @@ module github.com/apache/yunikorn-core
go 1.25.0
require (
- github.com/apache/yunikorn-scheduler-interface
v0.0.0-20260528033204-c474acff6d53
+ github.com/apache/yunikorn-scheduler-interface
v0.0.0-20260727092410-674338955bdf
github.com/go-ldap/ldap/v3 v3.4.13
github.com/google/btree v1.1.3
github.com/google/go-cmp v0.7.0
diff --git a/go.sum b/go.sum
index d8100c43..10e3c159 100644
--- a/go.sum
+++ b/go.sum
@@ -2,8 +2,8 @@ github.com/Azure/go-ntlmssp v0.1.1
h1:l+FM/EEMb0U9QZE7mKNEDw5Mu3mFiaa2GKOoTSsNDP
github.com/Azure/go-ntlmssp v0.1.1/go.mod
h1:NYqdhxd/8aAct/s4qSYZEerdPuH1liG2/X9DiVTbhpk=
github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e
h1:4dAU9FXIyQktpoUAgOJK3OTFc/xug0PCXYCqU0FgDKI=
github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e/go.mod
h1:cEWa1LVoE5KvSD9ONXsZrj0z6KqySlCCNKHlLzbqAt4=
-github.com/apache/yunikorn-scheduler-interface
v0.0.0-20260528033204-c474acff6d53
h1:zodKoODatR57zSGoQGdQ8EZjJgQqy9nh5/yxmI9hWpI=
-github.com/apache/yunikorn-scheduler-interface
v0.0.0-20260528033204-c474acff6d53/go.mod
h1:XygReHrRd3TtfNS3uVtRQI0sCOfprcLUfhlmv/LdCmk=
+github.com/apache/yunikorn-scheduler-interface
v0.0.0-20260727092410-674338955bdf
h1:IXEpAeqZgCXODJtvB6Ib5CCcpRyvF0mqcicz80rA3Cc=
+github.com/apache/yunikorn-scheduler-interface
v0.0.0-20260727092410-674338955bdf/go.mod
h1:qb739Bdm82PH7gsfEYabulGF90xKGNQ1hWmf197rDfw=
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
github.com/beorn7/perks v1.0.1/go.mod
h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
github.com/cespare/xxhash/v2 v2.3.0
h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
diff --git a/pkg/scheduler/objects/application.go
b/pkg/scheduler/objects/application.go
index 694924a4..d22f094e 100644
--- a/pkg/scheduler/objects/application.go
+++ b/pkg/scheduler/objects/application.go
@@ -514,6 +514,17 @@ func (sa *Application) GetAllocationAsk(allocationKey
string) *Allocation {
return sa.requests[allocationKey]
}
+// GetAllocationNodeID returns the node ID for a confirmed allocation
identified by allocationKey.
+// Returns an empty string if the allocation does not exist.
+func (sa *Application) GetAllocationNodeID(allocationKey string) string {
+ sa.RLock()
+ defer sa.RUnlock()
+ if alloc := sa.allocations[allocationKey]; alloc != nil {
+ return alloc.GetNodeID()
+ }
+ return ""
+}
+
// GetAllocatedResource returns the currently allocated resources for this
application
func (sa *Application) GetAllocatedResource() *resources.Resource {
sa.RLock()
@@ -799,6 +810,47 @@ func (sa *Application) DeallocateAsk(allocKey string)
(*resources.Resource, erro
return nil, fmt.Errorf("failed to locate ask with key %s", allocKey)
}
+// RollbackAllocation atomically reverts an allocated ask back to pending
state.
+// 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) {
+ sa.Lock()
+ defer sa.Unlock()
+
+ if !sa.IsAccepted() && !sa.IsRunning() {
+ return nil, fmt.Errorf("cannot rollback allocation %s:
application %s is in state %s", allocKey, sa.ApplicationID, sa.CurrentState())
+ }
+
+ ask := sa.allocations[allocKey]
+ if ask == nil {
+ return nil, fmt.Errorf("failed to locate allocation with key %s
for rollback", allocKey)
+ }
+
+ // Restore pending FIRST. This ensures hasZeroAllocations() sees
pending > 0 when
+ // we later decrement allocatedResource, 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)
+ }
+
+ // Clear stale node assignment on the ask so it is re-schedulable
cleanly.
+ ask.SetNodeID("")
+
+ res := ask.GetAllocatedResource()
+ sa.allocatedResource = resources.Sub(sa.allocatedResource, res)
+ sa.allocatedResource.Prune()
+ sa.decUserResourceUsage(res, false)
+ delete(sa.allocations, allocKey)
+
+ sa.appEvents.SendRemoveAllocationEvent(sa.ApplicationID, allocKey, res,
si.TerminationType_SCHEDULING_FAILED_ON_RM)
+
+ log.Log(log.SchedApplication).Info("allocation rolled back to pending
ask",
+ zap.String("appID", sa.ApplicationID),
+ zap.String("allocationKey", allocKey))
+
+ return res, nil
+}
+
func (sa *Application) allocateAsk(ask *Allocation) (*resources.Resource,
error) {
if !ask.allocate() {
return nil, fmt.Errorf("unable to allocate previously allocated
ask %s on app %s", ask.GetAllocationKey(), sa.ApplicationID)
diff --git a/pkg/scheduler/objects/application_test.go
b/pkg/scheduler/objects/application_test.go
index 7233788d..2d21f250 100644
--- a/pkg/scheduler/objects/application_test.go
+++ b/pkg/scheduler/objects/application_test.go
@@ -4012,6 +4012,251 @@ func TestAppSubmissionTime(t *testing.T) {
assert.Equal(t, app.submissionTime, time.Unix(0, 30), "app submission
time is not set properly")
}
+// TestGetAllocationNodeID verifies the node ID lookup for confirmed
allocations.
+func TestGetAllocationNodeID(t *testing.T) {
+ app := newApplication(appID1, "default", "root.unknown")
+
+ // Unknown key returns empty string.
+ assert.Equal(t, "", app.GetAllocationNodeID("nonexistent"), "unknown
key should return empty string")
+
+ // Empty key returns empty string.
+ assert.Equal(t, "", app.GetAllocationNodeID(""), "empty key should
return empty string")
+
+ // Existing allocation returns the correct node ID.
+ res :=
resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5})
+ alloc := newAllocationWithKey(aKey, appID1, nodeID1, res)
+ app.allocations[aKey] = alloc
+ assert.Equal(t, nodeID1, app.GetAllocationNodeID(aKey), "existing
allocation should return its node ID")
+
+ // A second allocation with a different node returns its own node ID
independently.
+ alloc2 := newAllocationWithKey(aKey2, appID1, nodeID2, res)
+ app.allocations[aKey2] = alloc2
+ assert.Equal(t, nodeID2, app.GetAllocationNodeID(aKey2), "second
allocation should return its own node ID")
+ assert.Equal(t, nodeID1, app.GetAllocationNodeID(aKey), "first
allocation node ID should be unchanged")
+
+ // After removing an allocation, the key returns empty string again.
+ delete(app.allocations, aKey)
+ assert.Equal(t, "", app.GetAllocationNodeID(aKey), "removed allocation
should return empty string")
+}
+
+// TestRollbackAllocationInvalidState verifies that RollbackAllocation returns
an error
+// when the application is not in Accepted or Running state.
+func TestRollbackAllocationInvalidState(t *testing.T) {
+ setupUGM()
+ queue, err := createRootQueue(nil)
+ assert.NilError(t, err, "queue create failed")
+ res :=
resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5})
+
+ for _, state := range []string{New.String(), Completing.String(),
Completed.String(), Rejected.String(), Failed.String()} {
+ app := newApplication(appID1, "default", "root.unknown")
+ app.queue = queue
+ app.SetState(state)
+ // inject a fake ask and allocation so the state check is the
only failure path
+ ask := newAllocationAsk(aKey, appID1, res)
+ app.requests[aKey] = ask
+ alloc := newAllocationWithKey(aKey, appID1, nodeID1, res)
+ app.allocations[aKey] = alloc
+
+ _, rollbackErr := app.RollbackAllocation(aKey)
+ assert.Assert(t, rollbackErr != nil, "expected error for state
%s", state)
+ assert.Assert(t, strings.Contains(rollbackErr.Error(), state),
"error should mention the invalid state %q, got: %s", state,
rollbackErr.Error())
+ }
+}
+
+// TestRollbackAllocationAskNotFound verifies that RollbackAllocation returns
an error
+// when the ask key does not exist in the application's request map.
+func TestRollbackAllocationAskNotFound(t *testing.T) {
+ setupUGM()
+ queue, err := createRootQueue(nil)
+ assert.NilError(t, err, "queue create failed")
+ app := newApplication(appID1, "default", "root.unknown")
+ app.queue = queue
+ app.SetState(Accepted.String())
+
+ _, err = app.RollbackAllocation("nonexistent-key")
+ assert.Assert(t, err != nil, "expected error for missing ask key")
+ assert.Assert(t, strings.Contains(err.Error(), "nonexistent-key"),
"error should mention the missing key")
+}
+
+// TestRollbackAllocationAllocationNotFound verifies that RollbackAllocation
returns an
+// error when the ask exists in requests but no corresponding allocation has
been recorded.
+func TestRollbackAllocationAllocationNotFound(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})
+
+ ask := newAllocationAsk(aKey, appID1, res)
+ err = app.AddAllocationAsk(ask)
+ assert.NilError(t, err, "ask should have been added to app")
+ assert.Assert(t, app.IsAccepted(), "app should be in Accepted state")
+ // ask is in requests but no allocation has been added
+
+ _, err = app.RollbackAllocation(aKey)
+ assert.Assert(t, err != nil, "expected error when no allocation exists
for the ask")
+ assert.Assert(t, strings.Contains(err.Error(), aKey), "error should
mention the missing allocation key")
+}
+
+// TestRollbackAllocationFromAccepted verifies a successful rollback when the
application
+// is in Accepted state. The allocation is injected directly into internal
state to avoid
+// the Running-state transition triggered by AddAllocation.
+func TestRollbackAllocationFromAccepted(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})
+
+ // Add ask → app moves to Accepted, ask enters requests.
+ ask := newAllocationAsk(aKey, appID1, res)
+ err = app.AddAllocationAsk(ask)
+ assert.NilError(t, err, "ask should have been added to app")
+ assert.Assert(t, app.IsAccepted(), "app should be in Accepted state")
+
+ // Simulate the core allocating the ask: mark it allocated and record
the allocation
+ // in internal state without triggering the Running state transition.
+ _, err = app.AllocateAsk(aKey)
+ assert.NilError(t, err, "AllocateAsk should succeed")
+ ask.SetNodeID(nodeID1)
+
+ app.Lock()
+ app.allocations[aKey] = ask
+ app.allocatedResource = resources.Add(app.allocatedResource, res)
+ app.incUserResourceUsage(res)
+ app.Unlock()
+
+ assert.Assert(t, app.IsAccepted(), "app should still be in Accepted
state before rollback")
+
+ // 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)
+
+ // Allocation is removed from the map.
+ assert.Assert(t, app.allocations[aKey] == nil, "allocation should have
been removed from map")
+
+ // Pending resource is restored.
+ assert.Assert(t, resources.Equals(app.GetPendingResource(), res),
"pending resource should be restored, got %v", app.GetPendingResource())
+
+ // Ask still exists in requests (rolled back to pending, not deleted).
+ assert.Assert(t, app.GetAllocationAsk(aKey) != nil, "ask should still
exist in requests after rollback")
+
+ // allocatedResource is back to zero.
+ assert.Assert(t, resources.IsZero(app.GetAllocatedResource()),
"allocated resource should be zero after rollback, got %v",
app.GetAllocatedResource())
+
+ // Node ID is cleared on the ask so it can be rescheduled on a
different node.
+ assert.Equal(t, app.GetAllocationAsk(aKey).GetNodeID(), "", "ask node
ID should be cleared after rollback")
+
+ // App remains in Accepted state (pending > 0 prevents Completing
transition).
+ assert.Assert(t, app.IsAccepted(), "app should remain in Accepted state
after rollback")
+}
+
+// TestRollbackAllocationFromRunning verifies a successful rollback from
Running state,
+// including full resource accounting and UGM tracking validation.
+func TestRollbackAllocationFromRunning(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})
+
+ // Standard allocation flow: ask → Accepted, allocate → Running.
+ ask := newAllocationAsk(aKey, appID1, res)
+ err = app.AddAllocationAsk(ask)
+ assert.NilError(t, err, "ask should have been added to app")
+ assert.Assert(t, app.IsAccepted(), "app should be in Accepted state
after adding ask")
+
+ delta, err := app.AllocateAsk(aKey)
+ assert.NilError(t, err, "AllocateAsk should succeed")
+ assert.Assert(t, resources.Equals(delta, res), "AllocateAsk delta
should equal res, got %v", delta)
+ ask.SetNodeID(nodeID1)
+
+ app.AddAllocation(ask)
+ assert.Assert(t, app.IsRunning(), "app should be in Running state after
AddAllocation")
+ assert.Assert(t, resources.IsZero(app.GetPendingResource()), "pending
should be zero before rollback")
+ assert.Assert(t, resources.Equals(app.GetAllocatedResource(), res),
"allocated should equal res before rollback")
+ assertUserGroupResource(t, getTestUserGroup(), res)
+
+ // 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)
+
+ // Allocation is removed.
+ assert.Assert(t, app.allocations[aKey] == nil, "allocation should have
been removed from map")
+
+ // Pending resource is restored.
+ assert.Assert(t, resources.Equals(app.GetPendingResource(), res),
"pending resource should be restored after rollback, got %v",
app.GetPendingResource())
+
+ // Ask still exists in requests.
+ assert.Assert(t, app.GetAllocationAsk(aKey) != nil, "ask should still
exist in requests after rollback")
+
+ // allocatedResource is back to zero.
+ assert.Assert(t, resources.IsZero(app.GetAllocatedResource()),
"allocated resource should be zero after rollback, got %v",
app.GetAllocatedResource())
+
+ // Node ID is cleared so the ask can land on a different node.
+ assert.Equal(t, app.GetAllocationAsk(aKey).GetNodeID(), "", "ask node
ID should be cleared after rollback")
+
+ // UGM resource usage is decremented.
+ assertUserGroupResource(t, getTestUserGroup(), nil)
+}
+
+// TestRollbackAllocationPartial verifies that rolling back one of several
allocations
+// leaves the remaining allocations and their accounting untouched.
+func TestRollbackAllocationPartial(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})
+
+ // Allocate first ask.
+ ask1 := newAllocationAsk(aKey, appID1, res)
+ err = app.AddAllocationAsk(ask1)
+ assert.NilError(t, err, "ask1 should have been added")
+ _, err = app.AllocateAsk(aKey)
+ assert.NilError(t, err, "AllocateAsk for ask1 should succeed")
+ ask1.SetNodeID(nodeID1)
+ app.AddAllocation(ask1)
+
+ // Allocate second ask.
+ ask2 := newAllocationAsk(aKey2, appID1, res)
+ err = app.AddAllocationAsk(ask2)
+ assert.NilError(t, err, "ask2 should have been added")
+ _, err = app.AllocateAsk(aKey2)
+ assert.NilError(t, err, "AllocateAsk for ask2 should succeed")
+ ask2.SetNodeID(nodeID2)
+ app.AddAllocation(ask2)
+
+ doubleRes := resources.Multiply(res, 2)
+ assert.Assert(t, app.IsRunning(), "app should be Running with two
allocations")
+ assert.Assert(t, resources.Equals(app.GetAllocatedResource(),
doubleRes), "allocated should be 2x res, got %v", app.GetAllocatedResource())
+ assertUserGroupResource(t, getTestUserGroup(), doubleRes)
+
+ // 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)
+
+ // Second allocation is untouched.
+ assert.Assert(t, app.allocations[aKey2] != nil, "second allocation
should still exist")
+ assert.Assert(t, resources.Equals(app.GetAllocatedResource(), res),
"allocated should be 1x res after partial rollback, got %v",
app.GetAllocatedResource())
+
+ // First ask is back to pending; second ask remains allocated.
+ assert.Assert(t, resources.Equals(app.GetPendingResource(), res),
"pending should be restored to res for the rolled-back ask, got %v",
app.GetPendingResource())
+
+ // UGM reflects only the remaining allocation.
+ assertUserGroupResource(t, getTestUserGroup(), res)
+
+ // App remains Running because the second allocation is still active.
+ assert.Assert(t, app.IsRunning(), "app should remain Running after
partial 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 28a4859f..3a8d2ddc 100644
--- a/pkg/scheduler/partition.go
+++ b/pkg/scheduler/partition.go
@@ -1504,6 +1504,14 @@ func (pc *PartitionContext) removeAllocation(release
*si.AllocationRelease) ([]*
return nil, nil
}
+ // Handle shim-initiated scheduling failure: roll the allocation back
to a pending ask
+ // so the core can re-schedule it on a different node. This bypasses
the normal
+ // remove-and-destroy path entirely. We return nil,nil to suppress the
echo back to
+ // the shim (the shim initiated this release; no confirmation is
needed).
+ if release.TerminationType ==
si.TerminationType_SCHEDULING_FAILED_ON_RM {
+ return pc.rollbackAllocation(appID, allocationKey, app)
+ }
+
// **** DO NOT MOVE **** this must be called before any allocations are
released.
// Processing a removal while in the Completing state could race with
the state change. The race occurs between
// removing the allocation and updating the queue after node
processing. If the state change removes the queue link
@@ -1528,34 +1536,7 @@ func (pc *PartitionContext) removeAllocation(release
*si.AllocationRelease) ([]*
continue
}
if release.TerminationType ==
si.TerminationType_PLACEHOLDER_REPLACED {
- confirmed = alloc.GetRelease()
- // we need to check the resources equality
- delta :=
resources.Sub(confirmed.GetAllocatedResource(), alloc.GetAllocatedResource())
- // Any negative value in the delta means that at least
one of the requested resource in the
- // placeholder is larger than the real allocation. The
node and queue need adjusting.
- // The reverse case is handled during allocation.
- if delta.HasNegativeValue() {
- // This looks incorrect but the delta is
negative and the result will be an increase of the
- // total tracked. The total will later be
deducted from the queue usage.
- total.SubFrom(delta)
- log.Log(log.SchedPartition).Warn("replacing
placeholder: placeholder is larger than real allocation",
- zap.String("allocationKey",
confirmed.GetAllocationKey()),
- zap.Stringer("requested resource",
confirmed.GetAllocatedResource()),
- zap.String("placeholderKey",
alloc.GetAllocationKey()),
- zap.Stringer("placeholder resource",
alloc.GetAllocatedResource()))
- }
- // replacements could be on a different node and
different size handle all cases
- if confirmed.GetNodeID() == alloc.GetNodeID() {
- // this is the real swap on the node, adjust
usage if needed
-
node.ReplaceAllocation(alloc.GetAllocationKey(), confirmed, delta)
- } else {
- // we have already added the real allocation to
the new node, just remove the placeholder
- node.RemoveAllocation(alloc.GetAllocationKey())
- }
- log.Log(log.SchedPartition).Info("replacing placeholder
allocation on node",
- zap.String("nodeID", alloc.GetNodeID()),
- zap.String("allocationKey",
alloc.GetAllocationKey()),
- zap.String("allocation nodeID",
confirmed.GetNodeID()))
+ confirmed = pc.replacePlaceholderOnNode(alloc, node,
total)
} else if node.RemoveAllocation(alloc.GetAllocationKey()) !=
nil {
// all non replacement are real removes: must update
the queue usage
total.AddTo(alloc.GetAllocatedResource())
@@ -1605,6 +1586,78 @@ func (pc *PartitionContext) removeAllocation(release
*si.AllocationRelease) ([]*
return released, confirmed
}
+// replacePlaceholderOnNode handles a placeholder replacement on a node. It
adjusts the total resource tracking
+// when the placeholder is larger than the real allocation, updates the node
accordingly, and returns the confirmed
+// real allocation.
+func (pc *PartitionContext) replacePlaceholderOnNode(alloc
*objects.Allocation, node *objects.Node, total *resources.Resource)
*objects.Allocation {
+ confirmed := alloc.GetRelease()
+ // we need to check the resources equality
+ delta := resources.Sub(confirmed.GetAllocatedResource(),
alloc.GetAllocatedResource())
+ // Any negative value in the delta means that at least one of the
requested resource in the
+ // placeholder is larger than the real allocation. The node and queue
need adjusting.
+ // The reverse case is handled during allocation.
+ if delta.HasNegativeValue() {
+ // This looks incorrect but the delta is negative and the
result will be an increase of the
+ // total tracked. The total will later be deducted from the
queue usage.
+ total.SubFrom(delta)
+ log.Log(log.SchedPartition).Warn("replacing placeholder:
placeholder is larger than real allocation",
+ zap.String("allocationKey",
confirmed.GetAllocationKey()),
+ zap.Stringer("requested resource",
confirmed.GetAllocatedResource()),
+ zap.String("placeholderKey", alloc.GetAllocationKey()),
+ zap.Stringer("placeholder resource",
alloc.GetAllocatedResource()))
+ }
+ // replacements could be on a different node and different size handle
all cases
+ if confirmed.GetNodeID() == alloc.GetNodeID() {
+ // this is the real swap on the node, adjust usage if needed
+ node.ReplaceAllocation(alloc.GetAllocationKey(), confirmed,
delta)
+ } else {
+ // we have already added the real allocation to the new node,
just remove the placeholder
+ node.RemoveAllocation(alloc.GetAllocationKey())
+ }
+ log.Log(log.SchedPartition).Info("replacing placeholder allocation on
node",
+ zap.String("nodeID", alloc.GetNodeID()),
+ zap.String("allocationKey", alloc.GetAllocationKey()),
+ zap.String("allocation nodeID", confirmed.GetNodeID()))
+ return confirmed
+}
+
+// rollbackAllocation handles a shim-initiated scheduling failure by rolling
the allocation back to a pending ask
+// so the core can re-schedule it on a different node.
+// NOTE: this is a lock free call. It must NOT be called holding the
PartitionContext lock.
+func (pc *PartitionContext) rollbackAllocation(appID, allocationKey string,
app *objects.Application) ([]*objects.Allocation, *objects.Allocation) {
+ queue := app.GetQueue()
+ // Retrieve node ID before rolling back (RollbackAllocation clears it
on the ask).
+ nodeID := app.GetAllocationNodeID(allocationKey)
+ res, err := app.RollbackAllocation(allocationKey)
+ if err != nil {
+ log.Log(log.SchedPartition).Warn("failed to rollback
allocation",
+ zap.String("appID", appID),
+ zap.String("allocationKey", allocationKey),
+ zap.Error(err))
+ return nil, nil
+ }
+ if node := pc.GetNode(nodeID); node != nil {
+ node.RemoveAllocation(allocationKey)
+ } else {
+ log.Log(log.SchedPartition).Warn("node not found while rolling
back allocation",
+ zap.String("appID", appID),
+ zap.String("allocationKey", allocationKey),
+ zap.String("nodeID", nodeID))
+ }
+ if err := queue.DecAllocatedResource(res); 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)
+ log.Log(log.SchedPartition).Info("allocation rolled back to pending
ask",
+ zap.String("appID", appID),
+ zap.String("allocationKey", allocationKey),
+ zap.String("nodeID", nodeID))
+ return nil, nil
+}
+
// updatePhAllocationCount checks the released allocations and updates the
partition context counter of allocated
// placeholders.
func (pc *PartitionContext) updatePhAllocationCount(released
[]*objects.Allocation) {
diff --git a/pkg/scheduler/partition_test.go b/pkg/scheduler/partition_test.go
index f832f2b4..3368ae0c 100644
--- a/pkg/scheduler/partition_test.go
+++ b/pkg/scheduler/partition_test.go
@@ -5426,3 +5426,139 @@ func TestUpdateResolver(t *testing.T) {
assert.NilError(t, err, "unable to update partition config")
assert.Equal(t, partition.GetUserGroupResolverType(), "test", "resolver
cannot be updated on running partition")
}
+
+// TestRemoveAllocationSchedulingFailedOnRM verifies that a
SCHEDULING_FAILED_ON_RM release rolls the
+// allocation back to a pending ask: allocation count decrements, node
allocation is removed, and the
+// queue allocated resource is decremented. The function must return nil, nil
(no echo to the shim).
+func TestRemoveAllocationSchedulingFailedOnRM(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})
+ node := setupNode(t, nodeID1, partition, nodeRes)
+
+ app := newApplication(appID1, "default", defQueue)
+ err = partition.AddApplication(app)
+ assert.NilError(t, err, "add application failed")
+
+ askRes :=
resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 1})
+
+ // Add a pending ask, then transition it to allocated on nodeID1.
+ ask := newAllocationAsk(allocKey, appID1, askRes)
+ _, _, err = partition.UpdateAllocation(ask)
+ assert.NilError(t, err, "add ask failed")
+
+ alloc := newAllocation(allocKey, appID1, nodeID1, askRes)
+ _, allocCreated, err := partition.UpdateAllocation(alloc)
+ assert.NilError(t, err, "transition to allocated failed")
+ assert.Check(t, allocCreated, "allocation should have been created")
+
+ assert.Equal(t, 1, partition.GetTotalAllocationCount(), "allocation
count should be 1 before rollback")
+ assert.Assert(t, node.GetAllocation(allocKey) != nil, "node should have
the allocation before rollback")
+ queueAllocBefore :=
partition.GetQueue(defQueue).GetAllocatedResource().Clone()
+ assert.Assert(t, resources.StrictlyGreaterThanZero(queueAllocBefore),
"queue should have allocated resources before rollback")
+
+ release := &si.AllocationRelease{
+ PartitionName: "test",
+ ApplicationID: appID1,
+ AllocationKey: allocKey,
+ TerminationType: si.TerminationType_SCHEDULING_FAILED_ON_RM,
+ }
+ released, confirmed := partition.removeAllocation(release)
+
+ 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.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")
+}
+
+// TestRemoveAllocationSchedulingFailedOnRMRollbackError verifies that when
RollbackAllocation returns an
+// error (allocation key not found in app), removeAllocation logs a warning
and returns nil, nil without
+// modifying the allocation count.
+func TestRemoveAllocationSchedulingFailedOnRMRollbackError(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)
+
+ app := newApplication(appID1, "default", defQueue)
+ err = partition.AddApplication(app)
+ assert.NilError(t, err, "add application failed")
+
+ askRes :=
resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 1})
+ ask := newAllocationAsk(allocKey, appID1, askRes)
+ _, _, err = partition.UpdateAllocation(ask)
+ assert.NilError(t, err, "add ask failed")
+
+ alloc := newAllocation(allocKey, appID1, nodeID1, askRes)
+ _, allocCreated, err := partition.UpdateAllocation(alloc)
+ assert.NilError(t, err, "transition to allocated failed")
+ assert.Check(t, allocCreated)
+
+ // Use a non-existent allocation key to force RollbackAllocation to
fail.
+ release := &si.AllocationRelease{
+ PartitionName: "test",
+ ApplicationID: appID1,
+ AllocationKey: "does-not-exist",
+ TerminationType: si.TerminationType_SCHEDULING_FAILED_ON_RM,
+ }
+ released, confirmed := partition.removeAllocation(release)
+
+ assert.Assert(t, released == nil, "should return nil on rollback error")
+ assert.Assert(t, confirmed == nil, "should return nil on rollback
error")
+ // Allocation count must be unchanged since the rollback failed.
+ assert.Equal(t, 1, partition.GetTotalAllocationCount(), "allocation
count should be unchanged after failed rollback")
+}
+
+// TestRemoveAllocationSchedulingFailedOnRMNodeNotFound verifies that when the
node referenced by the
+// allocation has already been removed, the rollback still succeeds: the
allocation is removed from the
+// app and queue, the allocation count is decremented, and the function
returns nil, nil.
+func TestRemoveAllocationSchedulingFailedOnRMNodeNotFound(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)
+
+ app := newApplication(appID1, "default", defQueue)
+ err = partition.AddApplication(app)
+ assert.NilError(t, err, "add application failed")
+
+ askRes :=
resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 1})
+ ask := newAllocationAsk(allocKey, appID1, askRes)
+ _, _, err = partition.UpdateAllocation(ask)
+ assert.NilError(t, err, "add ask failed")
+
+ alloc := newAllocation(allocKey, appID1, nodeID1, askRes)
+ _, allocCreated, err := partition.UpdateAllocation(alloc)
+ assert.NilError(t, err, "transition to allocated failed")
+ assert.Check(t, allocCreated)
+
+ assert.Equal(t, 1, partition.GetTotalAllocationCount(), "allocation
count should be 1 before rollback")
+
+ // Remove the node to simulate it disappearing before the shim sends
the failure release.
+ partition.nodes.RemoveNode(nodeID1)
+
+ release := &si.AllocationRelease{
+ PartitionName: "test",
+ ApplicationID: appID1,
+ AllocationKey: allocKey,
+ TerminationType: si.TerminationType_SCHEDULING_FAILED_ON_RM,
+ }
+ released, confirmed := partition.removeAllocation(release)
+
+ 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 decremented even when node is missing")
+ assert.Assert(t,
resources.IsZero(partition.GetQueue(defQueue).GetAllocatedResource()), "queue
resource should be zero after rollback")
+ assert.Assert(t,
resources.StrictlyGreaterThanZero(app.GetPendingResource()), "ask should be
pending again after rollback")
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
