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-k8shim.git
The following commit(s) were added to refs/heads/master by this push:
new e666a782 [YUNIKORN-3316] Handle AssumePod failures in shim (#1052)
e666a782 is described below
commit e666a782fa62098ae767528b9900e385068d1844
Author: Aditya Maheshwari <[email protected]>
AuthorDate: Mon Aug 10 22:41:43 2026 +1000
[YUNIKORN-3316] Handle AssumePod failures in shim (#1052)
During AssumePod cache update if there are any failures in the API call
the task goes to failed state. The application is completed if the failed
task is the last task for the application.
Yunikorn should retry the task and possibly select a different node that
does not fail the AssumePod.
As part of these changes following changes are done -
* Shim will send a SCHEDULING_FAILED_ON_RM termination type to core after
exhausting all retries on assume pod failures.
* Task will be kept in scheduling state so that it can be retries again.
* Task will clear node assignment, clear pvc bindings from cache and will
remove itself from node (forgetPod).
Core will rollback based on the core changes apache/yunikorn-core#1104
Closes: #1052
Signed-off-by: Wilfred Spiegelenburg <[email protected]>
---
go.mod | 4 +-
go.sum | 8 +--
pkg/cache/context.go | 31 ++++++++++++
pkg/cache/context_test.go | 97 ++++++++++++++++++++++++++++++++++++
pkg/cache/scheduler_callback.go | 11 +++-
pkg/cache/scheduler_callback_test.go | 58 ++++++++++++++++++++-
pkg/cache/task.go | 49 ++++++++++++++++++
pkg/cache/task_test.go | 93 ++++++++++++++++++++++++++++++++++
pkg/common/test/volumebinder_mock.go | 20 ++++++--
pkg/shim/scheduler_test.go | 31 +++++++++---
10 files changed, 383 insertions(+), 19 deletions(-)
diff --git a/go.mod b/go.mod
index 414d6156..349efd5e 100644
--- a/go.mod
+++ b/go.mod
@@ -21,8 +21,8 @@ module github.com/apache/yunikorn-k8shim
go 1.26.0
require (
- github.com/apache/yunikorn-core v0.0.0-20260605065500-fae6e256cd5c
- github.com/apache/yunikorn-scheduler-interface
v0.0.0-20260528033204-c474acff6d53
+ github.com/apache/yunikorn-core v0.0.0-20260728123950-900bd603df82
+ github.com/apache/yunikorn-scheduler-interface
v0.0.0-20260727092410-674338955bdf
github.com/google/go-cmp v0.7.0
github.com/google/uuid v1.6.0
github.com/looplab/fsm v1.0.3
diff --git a/go.sum b/go.sum
index 08489887..974ceb67 100644
--- a/go.sum
+++ b/go.sum
@@ -10,10 +10,10 @@ github.com/alexbrainman/sspi
v0.0.0-20250919150558-7d374ff0d59e h1:4dAU9FXIyQktp
github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e/go.mod
h1:cEWa1LVoE5KvSD9ONXsZrj0z6KqySlCCNKHlLzbqAt4=
github.com/antlr4-go/antlr/v4 v4.13.1
h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ=
github.com/antlr4-go/antlr/v4 v4.13.1/go.mod
h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw=
-github.com/apache/yunikorn-core v0.0.0-20260605065500-fae6e256cd5c
h1:iuShQ3op12fL+jQ420ZQLsgyDp/1L1/ygi9ZG8awvko=
-github.com/apache/yunikorn-core v0.0.0-20260605065500-fae6e256cd5c/go.mod
h1:yxn4stqkqgiSNYqOyulBHf9IgyTBM+rFrzWcyicueKU=
-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-core v0.0.0-20260728123950-900bd603df82
h1:uvaD7WchXd9nLkXArGjdul1OQfxXO+ukvI0VJJ0hAwk=
+github.com/apache/yunikorn-core v0.0.0-20260728123950-900bd603df82/go.mod
h1:Z/8pYK5VeDk9Q74vTCm3FnCVm1DtciaGYjl5/P7kO+M=
+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/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5
h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio=
github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod
h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs=
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
diff --git a/pkg/cache/context.go b/pkg/cache/context.go
index e353f5bc..9b99be02 100644
--- a/pkg/cache/context.go
+++ b/pkg/cache/context.go
@@ -897,6 +897,37 @@ func (ctx *Context) ForgetPod(name string) {
log.Log(log.ShimContext).Debug("unable to forget pod: not found in
cache", zap.String("pod", name))
}
+// RevertPodVolumeAssumptions undoes any PV/PVC assumptions made by the volume
binder
+// for the given pod on the given node. This is idempotent and safe to call
even if
+// AssumePodVolumes was never called or already reverted internally.
+func (ctx *Context) RevertPodVolumeAssumptions(podName, nodeID string) {
+ ctx.lock.Lock()
+ defer ctx.lock.Unlock()
+ pod := ctx.schedulerCache.GetPod(podName)
+ if pod == nil {
+ return
+ }
+ node := ctx.schedulerCache.GetNode(nodeID)
+ if node == nil {
+ return
+ }
+ podVolumeClaims, err :=
ctx.apiProvider.GetAPIs().VolumeBinder.GetPodVolumeClaims(ctx.klogger, pod)
+ if err != nil {
+ log.Log(log.ShimContext).Error("RevertPodVolumeAssumptions:
failed to get pod volume claims",
+ zap.String("pod", podName), zap.Error(err))
+ return
+ }
+ podVolumes, _, err :=
ctx.apiProvider.GetAPIs().VolumeBinder.FindPodVolumes(ctx.klogger, pod,
podVolumeClaims, node.Node())
+ if err != nil || podVolumes == nil {
+ log.Log(log.ShimContext).Error("RevertPodVolumeAssumptions:
failed to find pod volumes",
+ zap.String("pod", podName), zap.Error(err))
+ return
+ }
+
ctx.apiProvider.GetAPIs().VolumeBinder.RevertAssumedPodVolumes(podVolumes)
+ log.Log(log.ShimContext).Info("reverted assumed pod volumes",
+ zap.String("pod", podName), zap.String("node", nodeID))
+}
+
func (ctx *Context) notifyTaskComplete(app *Application, taskID string) {
if app == nil {
log.Log(log.ShimContext).Debug("In notifyTaskComplete but app
is nil",
diff --git a/pkg/cache/context_test.go b/pkg/cache/context_test.go
index d4a7d669..226dee53 100644
--- a/pkg/cache/context_test.go
+++ b/pkg/cache/context_test.go
@@ -2325,6 +2325,35 @@ func TestAssumePod_PodNotFound(t *testing.T) {
assert.Equal(t, podInCache.Spec.NodeName, "", "NodeName in pod spec was
set unexpectedly")
}
+func TestAssumePod_VolumesNotFullyBound(t *testing.T) {
+ binder := test.NewVolumeBinderMock()
+ binder.SetAllBound(false)
+ context := initAssumePodTest(binder)
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+
+ err := context.AssumePod(pod1UID, fakeNodeName)
+ assert.NilError(t, err)
+ assert.Assert(t,
!context.schedulerCache.ArePodVolumesAllBound(pod1UID), "volumes should not be
reported as fully bound")
+ assumedPod := context.schedulerCache.GetPod(pod1UID)
+ assert.Assert(t, assumedPod != nil, "pod not found in cache")
+ assert.Equal(t, assumedPod.Spec.NodeName, fakeNodeName)
+ assert.Assert(t, context.schedulerCache.IsAssumedPod(pod1UID))
+}
+
+func TestAssumePod_NodeNotFound(t *testing.T) {
+ context := initAssumePodTest(test.NewVolumeBinderMock())
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+
+ err := context.AssumePod(pod1UID, "nonexistent-node")
+ assert.NilError(t, err)
+ assert.Assert(t, !context.schedulerCache.IsAssumedPod(pod1UID))
+ podInCache := context.schedulerCache.GetPod(pod1UID)
+ assert.Assert(t, podInCache != nil, "pod not found in cache")
+ assert.Equal(t, podInCache.Spec.NodeName, "", "NodeName in pod spec was
set unexpectedly")
+}
+
// TestOriginatorPodAfterRestart Test to ensure originator pod remains same
even after restart. After restart, ordering of pods may change which can lead to
// incorrect originator pod selection. Instead of doing actual restart, create
a situation where in pods are being processed in any random order.
// For example, placeholders are processed first and then real driver pod.
@@ -2452,6 +2481,74 @@ func initAssumePodTest(binder *test.VolumeBinderMock)
*Context {
return context
}
+func TestRevertPodVolumeAssumptions(t *testing.T) {
+ binder := test.NewVolumeBinderMock()
+ binder.SetPodVolumes(&volumebinding.PodVolumes{})
+ context := initAssumePodTest(binder)
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+
+ context.RevertPodVolumeAssumptions(pod1UID, fakeNodeName)
+ assert.Equal(t, binder.RevertCalledCount(), 1)
+}
+
+func TestRevertPodVolumeAssumptions_PodNotFound(t *testing.T) {
+ binder := test.NewVolumeBinderMock()
+ binder.SetPodVolumes(&volumebinding.PodVolumes{})
+ context := initAssumePodTest(binder)
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+
+ context.RevertPodVolumeAssumptions("nonexistent-pod", fakeNodeName)
+ assert.Equal(t, binder.RevertCalledCount(), 0)
+}
+
+func TestRevertPodVolumeAssumptions_NodeNotFound(t *testing.T) {
+ binder := test.NewVolumeBinderMock()
+ binder.SetPodVolumes(&volumebinding.PodVolumes{})
+ context := initAssumePodTest(binder)
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+
+ context.RevertPodVolumeAssumptions(pod1UID, "nonexistent-node")
+ assert.Equal(t, binder.RevertCalledCount(), 0)
+}
+
+func TestRevertPodVolumeAssumptions_GetPodVolumeClaimsError(t *testing.T) {
+ binder := test.NewVolumeBinderMock()
+ binder.SetPodVolumes(&volumebinding.PodVolumes{})
+ binder.EnableVolumeClaimsError("volume claims error")
+ context := initAssumePodTest(binder)
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+
+ context.RevertPodVolumeAssumptions(pod1UID, fakeNodeName)
+ assert.Equal(t, binder.RevertCalledCount(), 0)
+}
+
+func TestRevertPodVolumeAssumptions_FindPodVolumesError(t *testing.T) {
+ binder := test.NewVolumeBinderMock()
+ binder.SetPodVolumes(&volumebinding.PodVolumes{})
+ binder.EnableFindPodVolumesError("find volumes error")
+ context := initAssumePodTest(binder)
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+
+ context.RevertPodVolumeAssumptions(pod1UID, fakeNodeName)
+ assert.Equal(t, binder.RevertCalledCount(), 0)
+}
+
+func TestRevertPodVolumeAssumptions_NilPodVolumes(t *testing.T) {
+ // By default the mock returns nil podVolumes; RevertAssumedPodVolumes
must not be called.
+ binder := test.NewVolumeBinderMock()
+ context := initAssumePodTest(binder)
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+
+ context.RevertPodVolumeAssumptions(pod1UID, fakeNodeName)
+ assert.Equal(t, binder.RevertCalledCount(), 0)
+}
+
func waitForNodeAcceptedEvent(recorder *k8sEvents.FakeRecorder) error {
// fetch the "node accepted" event
err := utils.WaitForCondition(func() bool {
diff --git a/pkg/cache/scheduler_callback.go b/pkg/cache/scheduler_callback.go
index ab93e6f9..a2cdaf0a 100644
--- a/pkg/cache/scheduler_callback.go
+++ b/pkg/cache/scheduler_callback.go
@@ -77,8 +77,15 @@ func (callback *AsyncRMCallback) UpdateAllocation(response
*si.AllocationRespons
return callback.context.AssumePod(alloc.AllocationKey,
alloc.NodeID)
})
if err != nil {
- task.FailWithEvent(err.Error(), "AssumePodError")
- return err
+ if task.IsPlaceholder() {
+ // Placeholder tasks do not have volume
bindings, so AssumePod failure
+ // is unexpected and unrecoverable; wrap the
error with context.
+ wrappedErr := fmt.Errorf("placeholder task does
not have volume bindings, AssumePod failed: %w", err)
+ task.FailWithEvent(wrappedErr.Error(),
"AssumePodError")
+ return wrappedErr
+ }
+ task.rollbackOnAssumePodFailure(alloc.AllocationKey,
alloc.NodeID)
+ continue
}
if utils.IsAssignedPod(task.GetTaskPod()) {
diff --git a/pkg/cache/scheduler_callback_test.go
b/pkg/cache/scheduler_callback_test.go
index 22b8c37e..cd4d7980 100644
--- a/pkg/cache/scheduler_callback_test.go
+++ b/pkg/cache/scheduler_callback_test.go
@@ -94,6 +94,49 @@ func TestUpdateAllocation_NewTask_AssumePodFails(t
*testing.T) {
binder.SetAssumePodVolumesError(errMsg)
setVolumeBinder(context, binder)
+ var rollbackSent atomic.Bool
+ //nolint:errcheck
+
context.apiProvider.(*client.MockedAPIProvider).MockSchedulerAPIUpdateAllocationFn(func(request
*si.AllocationRequest) error {
+ if request.Releases != nil {
+ for _, rel := range
request.Releases.AllocationsToRelease {
+ if rel.TerminationType ==
si.TerminationType_SCHEDULING_FAILED_ON_RM {
+ rollbackSent.Store(true)
+ }
+ }
+ }
+ return nil
+ })
+
+ err := callback.UpdateAllocation(&si.AllocationResponse{
+ New: []*si.Allocation{
+ {
+ ApplicationID: appID,
+ AllocationKey: taskUID1,
+ NodeID: fakeNodeName,
+ },
+ },
+ })
+ assert.NilError(t, err, "error updating allocation")
+ assert.Assert(t, !context.schedulerCache.IsAssumedPod(taskUID1))
+ assert.Assert(t, rollbackSent.Load(), "rollback release was not sent to
scheduler")
+ task := context.getTask(appID, taskUID1)
+ assert.Equal(t, "", task.GetAllocationKey(), "allocation key should be
cleared after rollback")
+ assert.Equal(t, TaskStates().Scheduling, task.GetTaskState(), "task
should remain in Scheduling state")
+}
+
+func TestUpdateAllocation_PlaceholderTask_AssumePodFails(t *testing.T) {
+ callback, context := initCallbackTest(t, false, true)
+ defer dispatcher.UnregisterAllEventHandlers()
+ defer dispatcher.Stop()
+ binder := test.NewVolumeBinderMock()
+ const errMsg = "error assuming pod volumes"
+ binder.SetAssumePodVolumesError(errMsg)
+ setVolumeBinder(context, binder)
+ NewPlaceholderManager(context.apiProvider.GetAPIs())
+ recorder := k8sEvents.NewFakeRecorder(1024)
+ events.SetRecorder(recorder)
+ defer events.SetRecorder(events.NewMockedRecorder())
+
err := callback.UpdateAllocation(&si.AllocationResponse{
New: []*si.Allocation{
{
@@ -103,13 +146,24 @@ func TestUpdateAllocation_NewTask_AssumePodFails(t
*testing.T) {
},
},
})
- assert.Error(t, err, errMsg)
+ assert.ErrorContains(t, err, "placeholder task does not have volume
bindings, AssumePod failed", "error should contain placeholder context")
+ assert.ErrorContains(t, err, errMsg, "error should wrap the original
AssumePod error")
assert.Assert(t, !context.schedulerCache.IsAssumedPod(taskUID1))
task := context.getTask(appID, taskUID1)
err = utils.WaitForCondition(func() bool {
return task.GetTaskState() == TaskStates().Failed
}, 10*time.Millisecond, time.Second)
- assert.NilError(t, err, "task has not transitioned to Failed state")
+ assert.NilError(t, err, "placeholder task has not transitioned to
Failed state")
+ // FailWithEvent emits a Warning/"AssumePodError" event; beforeTaskFail
emits a Normal/"TaskFailed" event.
+ assert.Equal(t, 2, len(recorder.Events), "expected two K8s events to be
recorded")
+ assumePodErrorFound := false
+ for i := 0; i < 2; i++ {
+ event := <-recorder.Events
+ if strings.Contains(event, "AssumePodError") {
+ assumePodErrorFound = true
+ }
+ }
+ assert.Assert(t, assumePodErrorFound, "no event with reason
'AssumePodError' was recorded")
}
func TestUpdateAllocation_NewTask_PodAlreadyAssigned(t *testing.T) {
diff --git a/pkg/cache/task.go b/pkg/cache/task.go
index e68b7b54..595b8c0c 100644
--- a/pkg/cache/task.go
+++ b/pkg/cache/task.go
@@ -620,6 +620,55 @@ func (task *Task) failWithEvent(errorMessage, actionReason
string) {
dispatcher.Dispatch(NewFailTaskEvent(task.applicationID, task.taskID,
errorMessage))
}
+// rollbackOnAssumePodFailure is called when AssumePod fails after all retries.
+// It resets task state and notifies the core to move the allocation back to a
+// pending ask so it can be re-scheduled on a different node.
+// Must be called without holding the task lock.
+func (task *Task) rollbackOnAssumePodFailure(allocationKey, nodeID string) {
+ // Read fields needed for event posting and release request.
+ // Clear stale node assignment under write lock so the task is clean
for the next allocation.
+ task.lock.Lock()
+ podCopy := task.pod.DeepCopy()
+ alias := task.alias
+ appID := task.applicationID
+ partition := task.application.partition
+ task.allocationKey = ""
+ task.nodeName = ""
+ task.lock.Unlock()
+
+ // Post a warning event so operators can see the retry via kubectl
describe pod.
+ events.GetRecorder().Eventf(podCopy, nil,
+ v1.EventTypeWarning, "AssumePodFailed", "AssumePodFailed",
+ "Node assignment failed for %s on node %s, it will be retried",
alias, nodeID)
+
+ // Revert any PV/PVC assumptions made by the volume binder. Idempotent:
safe to call
+ // even if AssumePodVolumes was never reached or already cleaned up
internally.
+ task.context.RevertPodVolumeAssumptions(allocationKey, nodeID)
+
+ // ForgetPod is idempotent: removes the pod from the assumed-pods cache.
+ task.context.ForgetPod(allocationKey)
+
+ // Notify the core to roll back the allocation to a pending ask.
+ if schedulerAPI := task.context.apiProvider.GetAPIs().SchedulerAPI;
schedulerAPI != nil {
+ releaseRequest := common.CreateReleaseRequestForTask(
+ appID,
+ allocationKey,
+ partition,
+ si.TerminationType_SCHEDULING_FAILED_ON_RM,
+ )
+ if err := schedulerAPI.UpdateAllocation(releaseRequest); err !=
nil {
+ log.Log(log.ShimCacheTask).Error("failed to send
rollback request to scheduler",
+ zap.String("appID", appID),
+ zap.String("allocationKey", allocationKey),
+ zap.Error(err))
+ }
+ }
+
+ log.Log(log.ShimCacheTask).Info("task allocation rolled back, will
retry on a different node",
+ zap.String("appID", appID),
+ zap.String("allocationKey", allocationKey))
+}
+
func (task *Task) SetTaskPod(pod *v1.Pod) {
task.lock.Lock()
defer task.lock.Unlock()
diff --git a/pkg/cache/task_test.go b/pkg/cache/task_test.go
index c8a51609..98184b2b 100644
--- a/pkg/cache/task_test.go
+++ b/pkg/cache/task_test.go
@@ -19,6 +19,9 @@
package cache
import (
+ "fmt"
+ "strings"
+ "sync/atomic"
"testing"
"time"
@@ -27,6 +30,7 @@ import (
schedulingv1 "k8s.io/api/scheduling/v1"
"k8s.io/apimachinery/pkg/api/resource"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+ k8sEvents "k8s.io/client-go/tools/events"
"github.com/apache/yunikorn-k8shim/pkg/client"
"github.com/apache/yunikorn-k8shim/pkg/common/constants"
@@ -930,3 +934,92 @@ func TestCheckPodMetadataBeforeScheduling(t *testing.T) {
})
}
}
+
+// newRollbackTask creates a task in Scheduling state with allocationKey and
nodeName set,
+// ready to exercise rollbackOnAssumePodFailure.
+func newRollbackTask(ctx *Context, allocationKey, nodeID string) *Task {
+ app := NewApplication(appID1, queueNameA, testUser, testGroups,
map[string]string{},
+ ctx.apiProvider.GetAPIs().SchedulerAPI)
+ pod := &v1.Pod{
+ TypeMeta: metav1.TypeMeta{Kind: "Pod", APIVersion: "v1"},
+ ObjectMeta: metav1.ObjectMeta{Name: "rollback-pod", UID:
"rollback-uid"},
+ }
+ task := NewTask(allocationKey, app, ctx, pod)
+ task.sm.SetState(TaskStates().Scheduling)
+ task.allocationKey = allocationKey
+ task.nodeName = nodeID
+ return task
+}
+
+// TestRollbackOnAssumePodFailure_ClearsStateAndSendsRelease verifies the
happy-path:
+// allocationKey and nodeName are cleared, a Warning "AssumePodFailed" event
is posted,
+// and a SCHEDULING_FAILED_ON_RM release request is sent to the scheduler core.
+func TestRollbackOnAssumePodFailure_ClearsStateAndSendsRelease(t *testing.T) {
+ mockedContext, apiProvider := initContextAndAPIProviderForTest()
+
+ recorder := k8sEvents.NewFakeRecorder(1024)
+ events.SetRecorder(recorder)
+ defer events.SetRecorder(events.NewMockedRecorder())
+
+ var rollbackSent atomic.Bool
+ apiProvider.MockSchedulerAPIUpdateAllocationFn(func(request
*si.AllocationRequest) error {
+ if request.Releases != nil {
+ for _, rel := range
request.Releases.AllocationsToRelease {
+ if rel.TerminationType ==
si.TerminationType_SCHEDULING_FAILED_ON_RM &&
+ rel.AllocationKey == taskUID1 {
+ rollbackSent.Store(true)
+ }
+ }
+ }
+ return nil
+ })
+
+ task := newRollbackTask(mockedContext, taskUID1, fakeNodeName)
+ task.rollbackOnAssumePodFailure(taskUID1, fakeNodeName)
+
+ assert.Equal(t, "", task.GetAllocationKey(), "allocationKey should be
cleared after rollback")
+ assert.Equal(t, "", task.GetNodeName(), "nodeName should be cleared
after rollback")
+ assert.Assert(t, rollbackSent.Load(), "SCHEDULING_FAILED_ON_RM release
request should be sent to scheduler")
+
+ assert.Equal(t, 1, len(recorder.Events), "expected one AssumePodFailed
event")
+ event := <-recorder.Events
+ assert.Assert(t, strings.Contains(event, "AssumePodFailed"), "event
should contain AssumePodFailed reason")
+}
+
+// TestRollbackOnAssumePodFailure_UpdateAllocationError verifies that if the
scheduler
+// API returns an error, the task state (allocationKey, nodeName) is still
cleaned up
+// and the function does not panic.
+func TestRollbackOnAssumePodFailure_UpdateAllocationError(t *testing.T) {
+ mockedContext, apiProvider := initContextAndAPIProviderForTest()
+ events.SetRecorder(events.NewMockedRecorder())
+ defer events.SetRecorder(events.NewMockedRecorder())
+
+ apiProvider.MockSchedulerAPIUpdateAllocationFn(func(request
*si.AllocationRequest) error {
+ return fmt.Errorf("scheduler unavailable")
+ })
+
+ task := newRollbackTask(mockedContext, taskUID1, fakeNodeName)
+ // Must not panic even when UpdateAllocation fails.
+ task.rollbackOnAssumePodFailure(taskUID1, fakeNodeName)
+
+ assert.Equal(t, "", task.GetAllocationKey(), "allocationKey should be
cleared even when UpdateAllocation errors")
+ assert.Equal(t, "", task.GetNodeName(), "nodeName should be cleared
even when UpdateAllocation errors")
+}
+
+// TestRollbackOnAssumePodFailure_NilSchedulerAPI verifies that
rollbackOnAssumePodFailure
+// does not panic when the scheduler API is nil, and still clears the
allocation state.
+func TestRollbackOnAssumePodFailure_NilSchedulerAPI(t *testing.T) {
+ mockedContext, _ := initContextAndAPIProviderForTest()
+ events.SetRecorder(events.NewMockedRecorder())
+ defer events.SetRecorder(events.NewMockedRecorder())
+
+ // Simulate a context where the scheduler API is not available.
+ mockedContext.apiProvider.GetAPIs().SchedulerAPI = nil
+
+ task := newRollbackTask(mockedContext, taskUID1, fakeNodeName)
+ // Must not panic with a nil SchedulerAPI.
+ task.rollbackOnAssumePodFailure(taskUID1, fakeNodeName)
+
+ assert.Equal(t, "", task.GetAllocationKey(), "allocationKey should be
cleared even with nil SchedulerAPI")
+ assert.Equal(t, "", task.GetNodeName(), "nodeName should be cleared
even with nil SchedulerAPI")
+}
diff --git a/pkg/common/test/volumebinder_mock.go
b/pkg/common/test/volumebinder_mock.go
index 3bc38276..149d60a6 100644
--- a/pkg/common/test/volumebinder_mock.go
+++ b/pkg/common/test/volumebinder_mock.go
@@ -37,9 +37,10 @@ type VolumeBinderMock struct {
bindError error
conflictReasons volumebinding.ConflictReasons
- podVolumeClaim *volumebinding.PodVolumeClaims
- podVolumes *volumebinding.PodVolumes
- allBound bool
+ podVolumeClaim *volumebinding.PodVolumeClaims
+ podVolumes *volumebinding.PodVolumes
+ allBound bool
+ revertCalledCount int
}
func NewVolumeBinderMock() *VolumeBinderMock {
@@ -81,6 +82,15 @@ func (v *VolumeBinderMock) AssumePodVolumes(_ klog.Logger, _
*v1.Pod, _ string,
}
func (v *VolumeBinderMock) RevertAssumedPodVolumes(_
*volumebinding.PodVolumes) {
+ v.revertCalledCount++
+}
+
+func (v *VolumeBinderMock) RevertCalledCount() int {
+ return v.revertCalledCount
+}
+
+func (v *VolumeBinderMock) SetPodVolumes(podVolumes *volumebinding.PodVolumes)
{
+ v.podVolumes = podVolumes
}
func (v *VolumeBinderMock) BindPodVolumes(_ context.Context, _ *v1.Pod, _
*volumebinding.PodVolumes) error {
@@ -106,3 +116,7 @@ func (v *VolumeBinderMock) SetConflictReasons(reasons
...string) {
func (v *VolumeBinderMock) SetAssumePodVolumesError(message string) {
v.assumeVolumeError = errors.New(message)
}
+
+func (v *VolumeBinderMock) SetAllBound(allBound bool) {
+ v.allBound = allBound
+}
diff --git a/pkg/shim/scheduler_test.go b/pkg/shim/scheduler_test.go
index afdcf92c..ecfb8d4b 100644
--- a/pkg/shim/scheduler_test.go
+++ b/pkg/shim/scheduler_test.go
@@ -34,6 +34,7 @@ import (
"github.com/apache/yunikorn-k8shim/pkg/common"
"github.com/apache/yunikorn-k8shim/pkg/common/constants"
"github.com/apache/yunikorn-k8shim/pkg/common/test"
+ "github.com/apache/yunikorn-k8shim/pkg/common/utils"
"github.com/apache/yunikorn-scheduler-interface/lib/go/api"
siCommon "github.com/apache/yunikorn-scheduler-interface/lib/go/common"
"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
@@ -248,12 +249,30 @@ func TestAssumePodError(t *testing.T) {
pod1 := createTestPod("root.a", "app0001", "task0001", taskResource)
cluster.AddPod(pod1)
- // expect app to enter Completing state with allocation+ask removed
- err = cluster.waitForApplicationStateInCore("app0001", partitionName,
"Completing")
- assert.NilError(t, err)
- app := cluster.getApplicationFromCore("app0001", partitionName)
- assert.Equal(t, 0, len(app.GetAllRequests()), "asks were not removed
from the application")
- assert.Equal(t, 0, len(app.GetAllAllocations()), "allocations were not
removed from the application")
+ // wait for core to deliver the first allocation to the shim; the
allocationKey is
+ // set in the callback before the AssumePod retry loop starts, so it
becomes
+ // visible immediately — well before the long exponential backoff runs
to completion.
+ err = utils.WaitForCondition(func() bool {
+ shimApp := cluster.context.GetApplication("app0001")
+ if shimApp == nil {
+ return false
+ }
+ task := shimApp.GetTask("task0001")
+ return task != nil && task.GetAllocationKey() != ""
+ }, time.Second, 30*time.Second)
+ assert.NilError(t, err, "task allocation key was never set; core did
not allocate the ask")
+
+ // with the rollback-on-AssumePod-failure implementation the task must
NOT fail;
+ // instead, each failed AssumePod attempt rolls the allocation back to
a pending
+ // ask in core so the scheduler can retry on a different node.
+ shimApp := cluster.context.GetApplication("app0001")
+ assert.Assert(t, shimApp != nil, "application should still be present
in the shim")
+ task := shimApp.GetTask("task0001")
+ assert.Assert(t, task != nil, "task should still be present in the
shim")
+ assert.Check(t, task.GetTaskState() == cache.TaskStates().Scheduling,
+ "task must be in Scheduling state while AssumePod retries are
in flight")
+ assert.Equal(t, 0, len(cluster.GetBoundPods(false)),
+ "no pods should be bound when AssumePod always fails")
}
func TestForeignPodTracking(t *testing.T) {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
